tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 00:23:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18214 from BerriAI/litellm_key_management_fix

Browse Source 
[Fix] Key Delete and Regenerate Permissions Fix
This commit is contained in:
yuneng-jiang
2025-12-19 08:48:21 -08:00
committed by GitHub
parent 388535003e 887aa62ac7
commit e1d25670cd
3 changed files with 845 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/key_management_endpoints.py
+64 -19
View File
@@ -2310,31 +2310,70 @@ async def _team_key_deletion_check(
return False
async def can_delete_verification_token(
async def can_modify_verification_token(
key_info: LiteLLM_VerificationToken,
user_api_key_cache: DualCache,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> bool:
"""
- check if user is proxy admin
- check if user is team admin and key is a team key
- check if key is personal key
Check if user has permission to modify (delete/regenerate) a verification token.
Rules:
- Proxy admin can modify any key
- For team keys: only team admin or key owner can modify
- For personal keys: only key owner can modify
Args:
key_info: The verification token to check
user_api_key_cache: Cache for user API keys
user_api_key_dict: The user making the request
prisma_client: Prisma client for database access
Returns:
True if user can modify the key, False otherwise
"""
is_team_key = _is_team_key(data=key_info)
# 1. Proxy admin can modify any key
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return True
elif is_team_key and key_info.team_id is not None:
return await _team_key_deletion_check(
user_api_key_dict=user_api_key_dict,
key_info=key_info,
# 2. For team keys: only team admin or key owner can modify
if is_team_key and key_info.team_id is not None:
# Get team object to check if user is team admin
team_table = await get_team_object(
team_id=key_info.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
elif key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
return True
else:
if team_table is None:
return False
# Check if user is team admin
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict,
team_obj=team_table,
):
return True
# Check if the key belongs to the user (they own it)
if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
return True
# Not team admin and doesn't own the key
return False
# 3. For personal keys: only key owner can modify
if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
return True
# Default: deny
return False
async def delete_verification_tokens(
@@ -2388,7 +2427,7 @@ async def delete_verification_tokens(
for key in _keys_being_deleted:
async def _delete_key(key: LiteLLM_VerificationToken):
if await can_delete_verification_token(
if await can_modify_verification_token(
key_info=key,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
@@ -2739,6 +2778,18 @@ async def regenerate_key_fn(
user_api_key_cache=user_api_key_cache,
)
# check if user has ownership permission to regenerate key
if not await can_modify_verification_token(
key_info=_key_in_db,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": "You are not authorized to regenerate this key"},
)
verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
new_token = get_new_token(data=data)
@@ -2777,14 +2828,8 @@ async def regenerate_key_fn(
### 3. remove existing key entry from cache
######################################################################
if key:
await _delete_cache_key_object(
hashed_token=hash_token(key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if hashed_api_key:
if hashed_api_key or key:
await _delete_cache_key_object(
hashed_token=hash_token(key),
user_api_key_cache=user_api_key_cache,
tests/otel_tests/test_team_member_permissions.py
+20 -11
View File
@@ -20,11 +20,12 @@
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
- Note: Delete/regenerate require key ownership or team admin status, not just team member permissions
- User tries deleting a key with team_id = team_id -> expect to fail (403) unless user owns the key or is team admin
- User tries regenerating a key with team_id = team_id -> expect to fail (403) unless user owns the key or is team admin
Invalid 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
@@ -303,10 +304,11 @@ async def test_default_member_permissions():
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"
assert "status" in delete_result and delete_result["status"] == 403, "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"
# Delete endpoint now returns 403 with authorization error, not team_member_permission_error
assert "error" in error_data, "Error should contain error field"
# User tries regenerating a key with team_id
print("Regular team member trying to regenerate a key with team_id. Expecting error.")
@@ -318,7 +320,8 @@ async def test_default_member_permissions():
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"
# Regenerate endpoint now returns 403 with authorization error, not team_member_permission_error
assert "error" in error_data, "Error should contain error field"
# Test valid permissions
# User tries calling /key/info with team_id
@@ -378,13 +381,15 @@ async def test_edit_delete_permissions():
)
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
# User tries deleting a key with team_id
# Note: Even with /key/delete permission, users can only delete keys they own or if they're team admin
# The delete endpoint checks ownership/team admin status, not just team member permissions
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"
assert "status" in delete_result and delete_result["status"] == 403, "User should not be able to delete keys they don't own (even with /key/delete permission, ownership is required)"
# Test invalid permissions
# User tries creating a key with team_id
@@ -396,13 +401,14 @@ async def test_edit_delete_permissions():
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
# Note: Even with /key/regenerate permission, users can only regenerate keys they own or if they're team admin
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"
assert "status" in regenerate_result and regenerate_result["status"] == 401, "User should not be able to regenerate keys they don't own (even with /key/regenerate permission, ownership is required)"
@pytest.mark.asyncio()
async def test_create_permissions():
@@ -475,13 +481,16 @@ async def test_create_permissions():
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"
assert "status" in delete_result and delete_result["status"] == 403, "User should not be able to delete keys for team"
# User tries regenerating a key with team_id
# User doesn't have /key/regenerate permission, so should get 401 (team member permission error)
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"
assert "status" in regenerate_result and regenerate_result["status"] == 401, "User should not be able to regenerate keys for team (no /key/regenerate permission)"
error_data = json.loads(regenerate_result["error"])
assert error_data["error"]["type"] == ProxyErrorTypes.team_member_permission_error.value, "Error should be a team member permission error"
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+761
View File
@@ -20,6 +20,7 @@ from litellm.proxy._types import (
LiteLLM_TeamTableCachedObj,
LiteLLM_VerificationToken,
LitellmUserRoles,
Member,
ProxyException,
UpdateKeyRequest,
)
@@ -29,6 +30,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_team_key_limits,
_common_key_generation_helper,
_list_key_helper,
can_modify_verification_token,
check_org_key_model_specific_limits,
check_team_key_model_specific_limits,
generate_key_helper_fn,
@@ -2613,3 +2615,762 @@ def test_check_org_key_model_specific_limits_org_model_tpm_overallocation():
"Allocated TPM limit=17000 + Key TPM limit=4000 is greater than organization TPM limit=20000"
in str(exc_info.value.detail)
)
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch):
"""Test that proxy admin can delete any team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatch):
"""Test that proxy admin can delete any personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_team_admin_own_team(monkeypatch):
"""Test that team admin can delete team keys from their own team."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="team-admin-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="team-admin-user", role="admin"),
Member(user_id="other-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_team_admin_different_team(monkeypatch):
"""Test that team admin cannot delete team keys from a different team."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-456",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="team-admin-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-456",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="different-admin", role="admin"),
Member(user_id="other-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_delete_verification_token_key_owner_team_key(monkeypatch):
"""Test that key owner can delete their own team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="key-owner-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_key_owner_personal_key(monkeypatch):
"""Test that key owner can delete their own personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_other_user_team_key(monkeypatch):
"""Test that other user cannot delete team keys they don't own and aren't admin for."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="other-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="key-owner-user", role="user"),
Member(user_id="other-user", role="user"),
Member(user_id="team-admin-user", role="admin"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_delete_verification_token_other_user_personal_key(monkeypatch):
"""Test that other user cannot delete personal keys they don't own."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="other-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_delete_verification_token_team_key_no_team_found(monkeypatch):
"""Test that deletion fails when team is not found in database."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="non-existent-team",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return None
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_delete_verification_token_personal_key_no_user_id(monkeypatch):
"""Test that deletion fails for personal key when key has no user_id."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id=None,
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="some-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_modify_verification_token_proxy_admin_team_key(monkeypatch):
"""Test that proxy admin can modify any team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_verification_token_proxy_admin_personal_key(monkeypatch):
"""Test that proxy admin can modify any personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_verification_token_team_admin_own_team(monkeypatch):
"""Test that team admin can modify team keys from their own team."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="team-admin-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="team-admin-user", role="admin"),
Member(user_id="other-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_verification_token_team_admin_different_team(monkeypatch):
"""Test that team admin cannot modify team keys from a different team."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-456",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="team-admin-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-456",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="different-admin", role="admin"),
Member(user_id="other-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_modify_verification_token_key_owner_team_key(monkeypatch):
"""Test that key owner can modify their own team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="key-owner-user", role="user"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_verification_token_key_owner_personal_key(monkeypatch):
"""Test that key owner can modify their own personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_modify_verification_token_other_user_team_key(monkeypatch):
"""Test that other user cannot modify team keys they don't own and aren't admin for."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="other-user",
api_key="sk-user",
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-123",
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="key-owner-user", role="user"),
Member(user_id="other-user", role="user"),
Member(user_id="team-admin-user", role="admin"),
],
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_modify_verification_token_other_user_personal_key(monkeypatch):
"""Test that other user cannot modify personal keys they don't own."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="other-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_modify_verification_token_team_key_no_team_found(monkeypatch):
"""Test that modification fails when team is not found in database."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="key-owner-user",
team_id="non-existent-team",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="key-owner-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
async def mock_get_team_object(*args, **kwargs):
return None
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False
@pytest.mark.asyncio
async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch):
"""Test that modification fails for personal key when key has no user_id."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id=None,
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="some-user",
api_key="sk-user",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is False