phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-24 18:24:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Team Member permissions] - Fixes (#9945)

Browse source 
* only load member permissions for non-admins

* run member permission checks on update + regenerate endpoints

* run check for /key/generate

* working test_default_member_permissions

* passing test with permissions on update delete endpoints

* test_create_permissions

* _team_key_generation_check

* fix TeamBase

* fix team endpoints

* fix api docs check
This commit is contained in:
Ishaan Jaff 2025-04-12 11:17:51 -07:00 • committed by GitHub
parent d2a462fc93
commit 4e81b2cab4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 601 additions and 64 deletions
Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/_types.py
View file

@ -435,21 +435,19 @@ class LiteLLMRoutes(enum.Enum):
"/get/litellm_model_cost_map",
] + info_routes
internal_user_routes = [
"/key/generate",
"/key/{token_id}/regenerate",
"/key/update",
"/key/delete",
"/key/health",
"/key/info",
"/global/spend/tags",
"/global/spend/keys",
"/global/spend/models",
"/global/spend/provider",
"/global/spend/end_users",
"/global/activity",
"/global/activity/model",
] + spend_tracking_routes
internal_user_routes = (
[
"/global/spend/tags",
"/global/spend/keys",
"/global/spend/models",
"/global/spend/provider",
"/global/spend/end_users",
"/global/activity",
"/global/activity/model",
]
+ spend_tracking_routes
+ key_management_routes
)
internal_user_view_only_routes = (
spend_tracking_routes + global_spend_tracking_routes
@ -983,6 +981,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
admins: list = []
members: list = []
members_with_roles: List[Member] = []
team_member_permissions: Optional[List[str]] = None
metadata: Optional[dict] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
@ -1138,7 +1137,6 @@ class LiteLLM_TeamTable(TeamBase):
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
model_id: Optional[int] = None
team_member_permissions: Optional[List[str]] = None
litellm_model_table: Optional[LiteLLM_ModelTable] = None
created_at: Optional[datetime] = None

102
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1,7 +1,7 @@
"""
KEY MANAGEMENT
All /key management endpoints
All /key management endpoints
/key/generate
/key/info
@ -42,6 +42,9 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
TeamMemberPermissionChecks,
)
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key
from litellm.proxy.utils import (
@ -104,20 +107,16 @@ def _is_allowed_to_make_key_request(
and user_api_key_dict.team_id == UI_TEAM_ID
):
return True # handle https://github.com/BerriAI/litellm/issues/7482
assert (
user_api_key_dict.team_id == team_id
), "User can only create keys for their own team. Got={}, Your Team ID={}".format(
team_id, user_api_key_dict.team_id
)
return True
def _team_key_generation_team_member_check(
def _team_key_operation_team_member_check(
assigned_user_id: Optional[str],
team_table: LiteLLM_TeamTableCachedObj,
user_api_key_dict: UserAPIKeyAuth,
team_key_generation: TeamUIKeyGenerationConfig,
route: KeyManagementRoutes,
):
if assigned_user_id is not None:
key_assigned_user_in_team = _get_user_in_team(
@ -130,7 +129,7 @@ def _team_key_generation_team_member_check(
detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}",
)
key_creating_user_in_team = _get_user_in_team(
team_member_object = _get_user_in_team(
team_table=team_table, user_id=user_api_key_dict.user_id
)
@ -141,20 +140,26 @@ def _team_key_generation_team_member_check(
if is_admin:
return True
elif key_creating_user_in_team is None:
elif team_member_object is None:
raise HTTPException(
status_code=400,
detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}",
)
elif (
"allowed_team_member_roles" in team_key_generation
and key_creating_user_in_team.role
and team_member_object.role
not in team_key_generation["allowed_team_member_roles"]
):
raise HTTPException(
status_code=400,
detail=f"Team member role {key_creating_user_in_team.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
)
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
team_member_object=team_member_object,
team_table=team_table,
route=route,
)
return True
@ -178,6 +183,7 @@ def _team_key_generation_check(
team_table: LiteLLM_TeamTableCachedObj,
user_api_key_dict: UserAPIKeyAuth,
data: GenerateKeyRequest,
route: KeyManagementRoutes,
):
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return True
@ -191,11 +197,12 @@ def _team_key_generation_check(
allowed_team_member_roles=["admin", "user"],
)
_team_key_generation_team_member_check(
_team_key_operation_team_member_check(
assigned_user_id=data.user_id,
team_table=team_table,
user_api_key_dict=user_api_key_dict,
team_key_generation=_team_key_generation,
route=route,
)
_key_generation_required_param_check(
data,
@ -252,6 +259,7 @@ def key_generation_check(
team_table: Optional[LiteLLM_TeamTableCachedObj],
user_api_key_dict: UserAPIKeyAuth,
data: GenerateKeyRequest,
route: KeyManagementRoutes,
) -> bool:
"""
Check if admin has restricted key creation to certain roles for teams or individuals
@ -271,6 +279,7 @@ def key_generation_check(
team_table=team_table,
user_api_key_dict=user_api_key_dict,
data=data,
route=route,
)
else:
return _personal_key_generation_check(
@ -425,6 +434,7 @@ async def generate_key_fn( # noqa: PLR0915
team_table=team_table,
user_api_key_dict=user_api_key_dict,
data=data,
route=KeyManagementRoutes.KEY_GENERATE,
)
common_key_access_checks(
@ -567,9 +577,9 @@ async def generate_key_fn( # noqa: PLR0915
request_type="key", **data_json, table_name="key"
)
response[
"soft_budget"
] = data.soft_budget # include the user-input soft budget in the response
response["soft_budget"] = (
data.soft_budget
) # include the user-input soft budget in the response
response = GenerateKeyResponse(**response)
@ -762,6 +772,15 @@ async def update_key_fn(
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
# check if user has permission to update key
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
non_default_values = prepare_key_update_data(
data=data, existing_key_row=existing_key_row
)
@ -1049,7 +1068,7 @@ async def info_key_fn(
)
if (
_can_user_query_key_info(
await _can_user_query_key_info(
user_api_key_dict=user_api_key_dict,
key=key,
key_info=key_info,
@ -1376,11 +1395,12 @@ async def _team_key_deletion_check(
)
# check if user is team admin
if team_table is not None:
return _team_key_generation_team_member_check(
return _team_key_operation_team_member_check(
assigned_user_id=user_api_key_dict.user_id,
team_table=team_table,
user_api_key_dict=user_api_key_dict,
team_key_generation=_team_key_generation,
route=KeyManagementRoutes.KEY_DELETE,
)
else:
raise HTTPException(
@ -1447,10 +1467,10 @@ async def delete_verification_tokens(
try:
if prisma_client:
tokens = [_hash_token_if_needed(token=key) for key in tokens]
_keys_being_deleted: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
)
# Assuming 'db' is your Prisma Client instance
@ -1552,9 +1572,9 @@ async def _rotate_master_key(
from litellm.proxy.proxy_server import proxy_config
try:
models: Optional[
List
] = await prisma_client.db.litellm_proxymodeltable.find_many()
models: Optional[List] = (
await prisma_client.db.litellm_proxymodeltable.find_many()
)
except Exception:
models = None
# 2. process model table
@ -1743,6 +1763,15 @@ async def regenerate_key_fn(
detail={"error": f"Key {key} not found."},
)
# check if user has permission to regenerate key
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_REGENERATE,
prisma_client=prisma_client,
existing_key_row=_key_in_db,
user_api_key_cache=user_api_key_cache,
)
verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
@ -1832,11 +1861,11 @@ async def validate_key_list_check(
param="user_id",
code=status.HTTP_403_FORBIDDEN,
)
complete_user_info_db_obj: Optional[
BaseModel
] = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
complete_user_info_db_obj: Optional[BaseModel] = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
)
if complete_user_info_db_obj is None:
@ -1897,10 +1926,10 @@ async def get_admin_team_ids(
if complete_user_info is None:
return []
# Get all teams that user is an admin of
teams: Optional[
List[BaseModel]
] = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
teams: Optional[List[BaseModel]] = (
await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
)
if teams is None:
return []
@ -2450,7 +2479,7 @@ async def key_health(
)
def _can_user_query_key_info(
async def _can_user_query_key_info(
user_api_key_dict: UserAPIKeyAuth,
key: Optional[str],
key_info: LiteLLM_VerificationToken,
@ -2468,6 +2497,11 @@ def _can_user_query_key_info(
# user can query their own key info
elif key_info.user_id == user_api_key_dict.user_id:
return True
elif await TeamMemberPermissionChecks.user_belongs_to_keys_team(
user_api_key_dict=user_api_key_dict,
existing_key_row=key_info,
):
return True
return False

2
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -142,6 +142,7 @@ async def new_team( # noqa: PLR0915
- team_alias: Optional[str] - User defined team alias
- team_id: Optional[str] - The team id of the user. If none passed, we'll generate it.
- members_with_roles: List[{"role": "admin" or "user", "user_id": "<user-id>"}] - A list of users and their roles in the team. Get user_id when making a new user via `/user/new`.
- team_member_permissions: Optional[List[str]] - A list of routes that non-admin team members can access. example: ["/key/generate", "/key/update", "/key/delete"]
- metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"extra_info": "some info"}
- tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
@ -420,6 +421,7 @@ async def update_team(
Parameters:
- team_id: str - The team id of the user. Required param.
- team_alias: Optional[str] - User defined team alias
- team_member_permissions: Optional[List[str]] - A list of routes that non-admin team members can access. example: ["/key/generate", "/key/update", "/key/delete"]
- metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
- tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit

4
litellm/proxy/management_helpers/team_member_permission_checks.py
View file

@ -10,7 +10,6 @@ from litellm.proxy._types import (
Member,
ProxyErrorTypes,
ProxyException,
Span,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import get_team_object
@ -57,7 +56,6 @@ class TeamMemberPermissionChecks:
route: KeyManagementRoutes,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span],
existing_key_row: LiteLLM_VerificationToken,
):
"""
@ -129,7 +127,7 @@ class TeamMemberPermissionChecks:
route=route, allowed_routes=team_member_permissions
):
raise ProxyException(
message=f"Team member does not have permissions for endpoint: {route}. You only have access to the following endpoints: {team_member_permissions}",
message=f"Team member does not have permissions for endpoint: {route}. You only have access to the following endpoints: {team_member_permissions} for team {team_table.team_id}",
type=ProxyErrorTypes.team_member_permission_error,
param=route,
code=401,

488
tests/otel_tests/test_team_member_permissions.py Normal file
View file

@ -0,0 +1,488 @@
"""
1. Default permissions for members in a team - allowed to call /key/info and /key/health
- Create a team, create a member in a team (role = "user")
Invalid Permissions:
- User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
Valid Permissions:
- User tries calling /key/info with team_id, expect to get valid response
2. Permissions - members allowd to edit, delete keys but not allowed to create keys
- Create a team with member_permissions = ["/key/update", "/key/delete", "/key/info"]
- Create a member in the team with role = "user"
Valid Permissions:
- User tries editing a key with team_id = team_id -> expect to pass. Valid Permissions
- User tries deleting a key with team_id = team_id -> expect to pass. Valid Permissions
- User tries creating a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries calling /key/info with team_id, expect to get valid response
3. Permissions - members allowed to create keys but not allowed to edit, delete keys
- Create a team with member_permissions = ["/key/generate"]
- Create a member in the team with role = "user"
Valid Permissions:
- User tries creating a key with team_id = team_id -> expect to pass. Valid Permissions
Invalid Permissions:
- User tries editing a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries deleting a key with team_id = team_id -> expect to fail. Invalid Permissions
- User tries regenerating a key with team_id = team_id -> expect to fail. Invalid Permissions
"""
import pytest
import asyncio
import aiohttp, openai
import uuid
import json
from litellm.proxy._types import ProxyErrorTypes
from typing import Optional
LITELLM_MASTER_KEY = "sk-1234"
async def create_team(session, key, member_permissions=None):
url = "http://0.0.0.0:4000/team/new"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"team_member_permissions": member_permissions
}
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
raise Exception(response_text)
return await response.json()
async def create_user(session, key, user_id, team_id=None):
url = "http://0.0.0.0:4000/user/new"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"user_id": user_id
}
if team_id:
data["team_id"] = team_id
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
raise Exception(response_text)
return await response.json()
async def add_team_member(session, key, team_id, user_id, role="user"):
url = "http://0.0.0.0:4000/team/member_add"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"team_id": team_id,
"member": {
"role": role,
"user_id": user_id
}
}
print("Adding team member with data: ", data)
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
raise Exception(response_text)
return await response.json()
async def generate_key(session, key, team_id=None, user_id=None):
url = "http://0.0.0.0:4000/key/generate"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {}
if team_id:
data["team_id"] = team_id
if user_id:
data["user_id"] = user_id
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
return {"status": status, "error": response_text}
return await response.json()
async def key_info(session, key, key_id):
url = f"http://0.0.0.0:4000/key/info?key={key_id}"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
async with session.get(url, headers=headers) as response:
status = response.status
response_text = await response.text()
if status != 200:
return {"status": status, "error": response_text}
return await response.json()
async def update_key(
session: aiohttp.ClientSession,
key: str,
key_id: str,
team_id: Optional[str] = None,
):
"""
Update a key
Args:
key: key to use for authentication
key_id: key to update
"""
url = "http://0.0.0.0:4000/key/update"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"key": key_id,
"metadata": {"updated": True}
}
if team_id:
data["team_id"] = team_id
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
return {"status": status, "error": response_text}
return await response.json()
async def delete_key(session, key, key_id):
url = "http://0.0.0.0:4000/key/delete"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"keys": [key_id]
}
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
return {"status": status, "error": response_text}
return await response.json()
async def regenerate_key(session, key, key_id, team_id=None):
url = "http://0.0.0.0:4000/key/regenerate"
headers = {
"Authorization": f"Bearer {key}",
"Content-Type": "application/json",
}
data = {
"key": key_id
}
if team_id:
data["team_id"] = team_id
async with session.post(url, headers=headers, json=data) as response:
status = response.status
response_text = await response.text()
if status != 200:
return {"status": status, "error": response_text}
return await response.json()
@pytest.mark.asyncio()
async def test_default_member_permissions():
"""
Test default permissions for members in a team - allowed to call /key/info and /key/health
"""
async with aiohttp.ClientSession() as session:
master_key = LITELLM_MASTER_KEY
# Create a team
team_data = await create_team(
session=session,
key=master_key
)
team_id = team_data["team_id"]
# create a team key
team_key_data = await generate_key(
session=session,
key=master_key,
team_id=team_id
)
team_key = team_key_data["key"]
# create a user
user_data = await create_user(
session=session,
key=master_key,
user_id=f"user_{uuid.uuid4().hex[:8]}",
team_id=team_id
)
user_id = user_data["user_id"]
# Create a user key
print("New user data: ", user_data)
# Create a user key
user_key_data = await generate_key(
session=session,
key=master_key,
user_id=user_id
)
print("new user key: ", user_key_data)
user_key = user_key_data["key"]
# Test invalid permissions
# User tries creating a key with team_id
print("Regular team member trying to create a key with team_id. Expecting error.")
create_result = await generate_key(
session=session,
key=user_key,
team_id=team_id
)
print("result: ", create_result)
assert "status" in create_result and create_result["status"] == 401, "User should not be able to create keys for team"
error_data = json.loads(create_result["error"])
print("error response =", json.dumps(error_data, indent=4))
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
# User tries editing a key with team_id
print("Regular team member trying to edit a key with team_id. Expecting error.")
update_result = await update_key(
session=session,
key=user_key,
key_id=team_key,
team_id="ATTACKER_TEAM_ID"
)
assert "status" in update_result and update_result["status"] == 401, "User should not be able to update keys for team"
error_data = json.loads(update_result["error"])
print("error response =", json.dumps(error_data, indent=4))
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
# User tries deleting a key with team_id
print("Regular team member trying to delete a key with team_id. Expecting error.")
delete_result = await delete_key(
session=session,
key=user_key,
key_id=team_key,
)
assert "status" in delete_result and delete_result["status"] == 401, "User should not be able to delete keys for team"
error_data = json.loads(delete_result["error"])
print("error response =", json.dumps(error_data, indent=4))
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
# User tries regenerating a key with team_id
print("Regular team member trying to regenerate a key with team_id. Expecting error.")
regenerate_result = await regenerate_key(
session=session,
key=user_key,
key_id=team_key,
)
assert "status" in regenerate_result and regenerate_result["status"] == 401, "User should not be able to regenerate keys for team"
error_data = json.loads(regenerate_result["error"])
print("error response =", json.dumps(error_data, indent=4))
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
# Test valid permissions
# User tries calling /key/info with team_id
print("Regular team member trying to get key info with team_id. Expecting success.")
info_result = await key_info(
session=session,
key=user_key,
key_id=team_key,
)
print("info result =", info_result)
assert "status" not in info_result, "Admin should be able to get key info"
@pytest.mark.asyncio()
async def test_edit_delete_permissions():
"""
Test permissions - members allowed to edit, delete keys but not allowed to create keys
"""
async with aiohttp.ClientSession() as session:
master_key = LITELLM_MASTER_KEY
# Create a team with specific member permissions
team_data = await create_team(
session=session,
key=master_key,
member_permissions=["/key/update", "/key/delete", "/key/info"]
)
team_id = team_data["team_id"]
# create a user in team=team_id
user_data = await create_user(
session=session,
key=master_key,
user_id=f"user_{uuid.uuid4().hex[:8]}",
team_id=team_id
)
user_id = user_data["user_id"]
# Generate an admin key for the team
admin_key_data = await generate_key(session, master_key, team_id)
admin_key = admin_key_data["key"]
key_id = admin_key_data["key"]
# Create a user key
user_key_data = await generate_key(
session=session,
key=master_key,
user_id=user_id
)
user_key = user_key_data["key"]
# Test valid permissions
# User tries editing a key with team_id
update_result = await update_key(
session=session,
key=user_key,
key_id=key_id,
team_id=team_id
)
assert "status" not in update_result, "User should be able to update keys for team"
# User tries deleting a key with team_id - test this last
delete_result = await delete_key(
session=session,
key=user_key,
key_id=key_id
)
assert "status" not in delete_result, "User should be able to delete keys for team"
# Test invalid permissions
# User tries creating a key with team_id
create_result = await generate_key(
session=session,
key=user_key,
team_id=team_id
)
assert "status" in create_result and create_result["status"] != 200, "User should not be able to create keys for team"
# User tries regenerating a key with team_id
regenerate_result = await regenerate_key(
session=session,
key=user_key,
key_id=key_id,
team_id=team_id
)
assert "status" in regenerate_result and regenerate_result["status"] != 200, "User should not be able to regenerate keys for team"
@pytest.mark.asyncio()
async def test_create_permissions():
"""
Test permissions - members allowed to create keys but not allowed to edit, delete keys
"""
async with aiohttp.ClientSession() as session:
master_key = LITELLM_MASTER_KEY
# Create a team with specific member permissions
team_data = await create_team(
session=session,
key=master_key,
member_permissions=["/key/generate"]
)
team_id = team_data["team_id"]
# Create a user in the team
user_id = f"user_{uuid.uuid4().hex[:8]}"
await add_team_member(
session=session,
key=master_key,
team_id=team_id,
user_id=user_id,
role="user"
)
# Generate an admin key for the team
admin_key_data = await generate_key(
session=session,
key=master_key,
team_id=team_id
)
admin_key = admin_key_data["key"]
key_id = admin_key_data["key"]
# Create a user key
user_key_data = await generate_key(
session=session,
key=master_key,
user_id=user_id
)
user_key = user_key_data["key"]
# Test valid permissions
# User tries creating a key with team_id
create_result = await generate_key(
session=session,
key=user_key,
team_id=team_id
)
print("success, user created key for team=", create_result)
assert "key" in create_result, "User should be able to create keys for team"
assert create_result["team_id"] == team_id, "User should be able to create keys for team"
assert "status" not in create_result, "User should be able to create keys for team"
# Test invalid permissions
# User tries editing a key with team_id
update_result = await update_key(
session=session,
key=user_key,
key_id=key_id,
team_id=team_id
)
assert "status" in update_result and update_result["status"] != 200, "User should not be able to update keys for team"
# User tries deleting a key with team_id
delete_result = await delete_key(
session=session,
key=user_key,
key_id=key_id
)
assert "status" in delete_result and delete_result["status"] != 200, "User should not be able to delete keys for team"
# User tries regenerating a key with team_id
regenerate_result = await regenerate_key(
session=session,
key=user_key,
key_id=key_id,
team_id=team_id
)
assert "status" in regenerate_result and regenerate_result["status"] != 200, "User should not be able to regenerate keys for team"

21
tests/proxy_admin_ui_tests/test_key_management.py
View file

@ -661,6 +661,7 @@ def test_team_key_generation_team_member_check():
team_member=Member(role="admin", user_id="test_user_id"),
),
data=GenerateKeyRequest(),
route=KeyManagementRoutes.KEY_GENERATE,
)
team_table = LiteLLM_TeamTableCachedObj(
@ -679,6 +680,7 @@ def test_team_key_generation_team_member_check():
team_member=Member(role="user", user_id="test_user_id"),
),
data=GenerateKeyRequest(),
route=KeyManagementRoutes.KEY_GENERATE,
)
@ -739,13 +741,26 @@ def test_key_generation_required_params_check(
if expected_result:
if key_type == "team_key":
assert _team_key_generation_check(team_table, user_api_key_dict, input_data)
assert _team_key_generation_check(
team_table=team_table,
user_api_key_dict=user_api_key_dict,
data=input_data,
route=KeyManagementRoutes.KEY_GENERATE,
)
elif key_type == "personal_key":
assert _personal_key_generation_check(user_api_key_dict, input_data)
assert _personal_key_generation_check(
user_api_key_dict=user_api_key_dict,
data=input_data,
)
else:
if key_type == "team_key":
with pytest.raises(HTTPException):
_team_key_generation_check(team_table, user_api_key_dict, input_data)
_team_key_generation_check(
team_table=team_table,
user_api_key_dict=user_api_key_dict,
data=input_data,
route=KeyManagementRoutes.KEY_GENERATE,
)
elif key_type == "personal_key":
with pytest.raises(HTTPException):
_personal_key_generation_check(user_api_key_dict, input_data)

18
ui/litellm-dashboard/src/components/team/team_info.tsx
View file

@ -32,7 +32,7 @@ import { PencilAltIcon, PlusIcon, TrashIcon } from "@heroicons/react/outline";
import MemberModal from "./edit_membership";
import UserSearchModal from "@/components/common_components/user_search_modal";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
import { isAdminRole } from "@/utils/roles";
export interface TeamData {
team_id: string;
@ -296,13 +296,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
</TabPanel>
{/* Member Permissions Panel */}
<TabPanel>
<MemberPermissions
teamId={teamId}
accessToken={accessToken}
canEditTeam={canEditTeam}
/>
</TabPanel>
{canEditTeam && (
<TabPanel>
<MemberPermissions
teamId={teamId}
accessToken={accessToken}
canEditTeam={canEditTeam}
/>
</TabPanel>
)}
{/* Settings Panel */}
<TabPanel>