scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(proxy): write key deleted audit logs for cascade and alias key deletions (#42446)

Browse source 
* fix(proxy): write key deleted audit logs for cascade and alias key deletions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): assert persisted key deleted audit rows for cascade paths

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(e2e): route /audit and /v2/login to the control plane in split transport

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-22 11:52:08 -07:00 • committed by GitHub
parent 5af9b0136c
commit f3a0bb7249
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 645 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

69
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -1,7 +1,8 @@
import asyncio
import json
from collections.abc import Sequence
from datetime import datetime, timezone
from typing import Final
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
@ -23,6 +24,9 @@ from litellm.proxy._types import (
from litellm.proxy.utils import _hash_token_if_needed
from litellm.secret_managers.base_secret_manager import BaseSecretManager
if TYPE_CHECKING:
from prisma import models as prisma_models
# NOTE: This is the prefix for all virtual keys stored in AWS Secrets Manager
LITELLM_PREFIX_STORED_VIRTUAL_KEYS: Final = "litellm/"
@ -233,6 +237,19 @@ class KeyManagementEventHooks:
Handles the following:
- Storing Audit Logs for key deletion
"""
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=keys_being_deleted,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager(keys_being_deleted=keys_being_deleted)
@staticmethod
def create_key_deleted_audit_logs(
keys_being_deleted: Sequence["LiteLLM_VerificationToken | prisma_models.LiteLLM_VerificationToken"],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
) -> None:
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
get_audit_log_changed_by,
@ -240,35 +257,33 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if is_audit_logging_enabled() and data.keys is not None:
# make an audit log for each key deleted
for key in keys_being_deleted:
if key.token is None:
continue
_key_row = key.model_dump_json(exclude_none=True)
if not is_audit_logging_enabled():
return
for key in keys_being_deleted:
key_row = LiteLLM_VerificationToken.model_validate(key, from_attributes=True)
if key_row.token is None:
continue
_key_row = key_row.model_dump_json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key.token,
action="deleted",
updated_values="{}",
before_value=_key_row,
)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=get_audit_log_changed_by(
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
),
changed_by_api_key=user_api_key_dict.token,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=key_row.token,
action="deleted",
updated_values="{}",
before_value=_key_row,
)
)
# delete the keys from the secret manager
await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager(keys_being_deleted=keys_being_deleted)
)
@staticmethod
async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str, team_id: str | None = None):

7
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -48,6 +48,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
user_object_permission_id_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
from litellm.proxy.management_endpoints.common_daily_activity import (
@ -2539,6 +2540,12 @@ async def delete_user(
prisma_client=prisma_client,
)
await _verification_token_table(prisma_client).delete_many(where=key_filter)
if keys_to_delete:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=keys_to_delete,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await delete_cache_key_objects(
hashed_tokens=hashed_tokens_to_delete,
user_api_key_cache=user_api_key_cache,

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

@ -118,6 +118,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.model_max_budget_limiter import (
build_model_max_budget_usage,
resolve_model_budget,
@ -3751,6 +3752,13 @@ async def _team_member_delete(
}
)
if keys_to_delete:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=keys_to_delete,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
await delete_cache_team_object(
team_id=data.team_id,
team_alias=existing_team_row.team_alias,
@ -4469,6 +4477,13 @@ async def delete_team(
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
if keys_to_delete:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=keys_to_delete,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _invalidate_deleted_key_cache(
keys=keys_to_delete,
user_api_key_cache=user_api_key_cache,

35
litellm/proxy/management_helpers/bulk_user_deletion.py
View file

@ -31,6 +31,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_checks import delete_cache_key_objects, get_jwt_key_mapping_cache_keys_for_tokens
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem
from litellm.proxy.management_endpoints.common_utils import (
@ -93,20 +94,28 @@ class _TeamRemoval:
team: LiteLLM_TeamTable
removed: frozenset[str]
matched: frozenset[int]
deleted_key_tokens: tuple[str, ...]
deleted_keys: tuple["prisma_models.LiteLLM_VerificationToken", ...]
jwt_mapping_cache_keys: tuple[str, ...]
@property
def deleted_key_tokens(self) -> tuple[str, ...]:
return tuple(k.token for k in self.deleted_keys)
@dataclass(frozen=True, slots=True)
class _UserBatchDeletion:
removals: Mapping[str, _TeamRemoval]
deleted_key_tokens: tuple[str, ...]
deleted_keys: tuple["prisma_models.LiteLLM_VerificationToken", ...]
jwt_mapping_cache_keys: tuple[str, ...]
@property
def deleted_key_tokens(self) -> tuple[str, ...]:
return tuple(k.token for k in self.deleted_keys)
@dataclass(frozen=True, slots=True)
class _DeletedKeys:
tokens: tuple[str, ...]
keys: tuple["prisma_models.LiteLLM_VerificationToken", ...]
jwt_mapping_cache_keys: tuple[str, ...]
@ -276,7 +285,7 @@ async def _remove_members_from_team(
),
removed=cleanup_ids,
matched=matched,
deleted_key_tokens=tuple(k.token for k in keys),
deleted_keys=tuple(keys),
jwt_mapping_cache_keys=jwt_mapping_cache_keys,
)
@ -330,6 +339,12 @@ async def bulk_remove_team_members(
members: Final = tuple(data.members[i] for i in kept_indexes)
async with prisma_client.tx(timeout=_BATCH_TX_TIMEOUT) as tx:
removal: Final = await _remove_members_from_team(prisma_client, tx, team_id, members, user_api_key_dict)
if removal.deleted_keys:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=removal.deleted_keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
await delete_cache_key_objects(
hashed_tokens=removal.deleted_key_tokens,
user_api_key_cache=user_api_key_cache,
@ -407,7 +422,7 @@ async def _delete_user_rows(
await _org_membership_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
await _membership_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
await _user_tx_db(tx).delete_many(where=_in_filter("user_id", user_ids))
return _DeletedKeys(tokens=tuple(k.token for k in keys), jwt_mapping_cache_keys=jwt_mapping_cache_keys)
return _DeletedKeys(keys=tuple(keys), jwt_mapping_cache_keys=jwt_mapping_cache_keys)
async def _delete_users_tx(
@ -446,7 +461,7 @@ async def _delete_users_tx(
)
return _UserBatchDeletion(
removals=removals,
deleted_key_tokens=deleted_keys.tokens + tuple(t for r in removals.values() for t in r.deleted_key_tokens),
deleted_keys=deleted_keys.keys + tuple(k for r in removals.values() for k in r.deleted_keys),
jwt_mapping_cache_keys=deleted_keys.jwt_mapping_cache_keys
+ tuple(k for r in removals.values() for k in r.jwt_mapping_cache_keys),
)
@ -469,6 +484,12 @@ async def _delete_users(
except Exception as e: # noqa: BLE001 # the rolled-back batch is reported per row, not as a request failure
verbose_proxy_logger.error("users/bulk_delete: failed to delete users %s: %s", sorted(user_ids), e)
return _error_message(e)
if deletion.deleted_keys:
KeyManagementEventHooks.create_key_deleted_audit_logs(
keys_being_deleted=deletion.deleted_keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await delete_cache_key_objects(
hashed_tokens=deletion.deleted_key_tokens,
user_api_key_cache=user_api_key_cache,
@ -555,7 +576,7 @@ async def bulk_delete_users(
litellm_changed_by,
)
if candidates
else _UserBatchDeletion(removals=MappingProxyType({}), deleted_key_tokens=(), jwt_mapping_cache_keys=())
else _UserBatchDeletion(removals=MappingProxyType({}), deleted_keys=(), jwt_mapping_cache_keys=())
)
def result(index: int, user_id: str) -> UserDeleteResult:

4
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -94,6 +94,10 @@
- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"}
- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"}
- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"}
- {id: mgmt.key.delete.audit_logged, module: mgmt, tier: P0, surface: api, assertions: [audit_logged], source: "key_management_endpoints.py:3981", rationale: "Every hard key deletion writes one LiteLLM_VerificationToken deleted audit row, whether the key is addressed by key or by alias"}
- {id: mgmt.team.member_delete.audit_logs_keys, module: mgmt, tier: P0, surface: api, assertions: [audit_logs_keys], source: "team_endpoints.py:3563", rationale: "Removing a team member hard-deletes their keys and each deleted key writes a deleted audit row"}
- {id: mgmt.team.delete.audit_logs_keys, module: mgmt, tier: P0, surface: api, assertions: [audit_logs_keys], source: "team_endpoints.py:4344", rationale: "Deleting a team hard-deletes its keys and each deleted key writes a deleted audit row"}
- {id: mgmt.user.delete.audit_logs_keys, module: mgmt, tier: P0, surface: api, assertions: [audit_logs_keys], source: "internal_user_endpoints.py:2369", rationale: "Deleting a user hard-deletes their keys and each deleted key writes a deleted audit row"}
- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"}
- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"}

28
tests/e2e/management/management_client.py
View file

@ -25,6 +25,8 @@ from e2e_http import (
unwrap,
)
from models import (
AuditLogPage,
AuditLogParams,
ChatBody,
ChatMessage,
ConnectionTestBody,
@ -35,6 +37,7 @@ from models import (
CustomerResponse,
KeyBlockBody,
KeyDeleteBody,
KeyDeleteByAliasBody,
KeyGenerateBody,
KeyGenerateResponse,
KeyInfoParams,
@ -159,6 +162,31 @@ class ManagementClient:
def update_key_models(self, key: str, models: list[str]) -> None:
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
def delete_key_by_alias(self, key_alias: str) -> None:
_ = unwrap(
self.proxy.transport.post(
"/key/delete",
headers=self.proxy.management_headers(),
json=KeyDeleteByAliasBody(key_aliases=[key_alias]),
response_type=NoBody,
)
)
def key_deleted_audit_logs(self, token_hash: str) -> AuditLogPage:
return unwrap(
self.proxy.transport.get(
"/audit",
headers=self.proxy.management_headers(),
params=AuditLogParams(
object_id=token_hash,
action="deleted",
table_name="LiteLLM_VerificationToken",
page_size=100,
),
response_type=AuditLogPage,
)
)
def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]:
return self.proxy.transport.get(
"/key/info",

132
tests/e2e/management/test_management_e2e.py
View file

@ -16,8 +16,8 @@ from typing import Final
import pytest
from e2e_config import UI_PASSWORD, UI_USERNAME, unique_marker
from e2e_http import StreamingResponse, Success
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, UI_PASSWORD, UI_USERNAME, unique_marker
from e2e_http import StreamingResponse, Success, unwrap
from lifecycle import ResourceManager
from management_client import (
DASHBOARD_SESSION_TEAM_ID,
@ -26,7 +26,9 @@ from management_client import (
ManagementClient,
)
from models import (
AuditLogPage,
KeyGenerateBody,
KeyGenerateResponse,
KeyUpdateBody,
LiteLLMParamsBody,
ModelInfoEntry,
@ -40,6 +42,7 @@ from models import (
UserNewBody,
UserUpdateBody,
)
from proxy_client import Converged, await_converged
pytestmark = pytest.mark.e2e
@ -801,3 +804,128 @@ class TestCustomer:
assert info.user_id == customer, (
f"/customer/info did not report the created end-user; got {info.user_id!r}"
)
def _await_deleted_audit_rows(client: ManagementClient, token_hash: str) -> AuditLogPage:
outcome = await_converged(
lambda: client.key_deleted_audit_logs(token_hash),
converged=lambda page: page.total >= 1,
timeout=POLL_TIMEOUT,
interval=POLL_INTERVAL,
now=time.monotonic,
sleep=time.sleep,
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
def _assert_single_deleted_row(page: AuditLogPage, token_hash: str) -> None:
assert page.total == 1, page
row = page.audit_logs[0]
assert row.action == "deleted", row
assert row.table_name == "LiteLLM_VerificationToken", row
assert row.object_id == token_hash, row
assert row.changed_by, row
def _assert_key_deleted(client: ManagementClient, key: str) -> None:
def gone() -> bool | None:
match client.key_info_as(key):
case Success(data=response):
return True if response.info.status == "deleted" else None
case _:
return True
_ = _poll(
client,
gone,
"/key/info never reported status 'deleted' for a key whose deletion returned",
)
def _token_of(created: KeyGenerateResponse) -> str:
assert created.token is not None, created
return created.token
def _generate_response(
client: ManagementClient, resources: ResourceManager, body: KeyGenerateBody
) -> KeyGenerateResponse:
created = unwrap(client.generate_key(body))
resources.defer(lambda: client.delete_key_strict(created.key, missing_ok=True))
return created
class TestKeyDeletionAuditLog:
@pytest.mark.covers("mgmt.key.delete.audit_logged")
def test_key_delete_by_key_writes_audit_row(
self, client: ManagementClient, resources: ResourceManager
) -> None:
created = _generate_response(client, resources, KeyGenerateBody(key_alias=f"e2e-audit-{unique_marker()}"))
token = _token_of(created)
client.delete_key_strict(created.key)
_assert_key_deleted(client, created.key)
_assert_single_deleted_row(_await_deleted_audit_rows(client, token), token)
@pytest.mark.covers("mgmt.key.delete.audit_logged")
def test_key_delete_by_alias_writes_audit_row(
self, client: ManagementClient, resources: ResourceManager
) -> None:
alias = f"e2e-audit-{unique_marker()}"
created = _generate_response(client, resources, KeyGenerateBody(key_alias=alias))
token = _token_of(created)
client.delete_key_by_alias(alias)
_assert_key_deleted(client, created.key)
_assert_single_deleted_row(_await_deleted_audit_rows(client, token), token)
@pytest.mark.covers("mgmt.team.member_delete.audit_logs_keys")
def test_team_member_delete_writes_audit_row_for_member_keys(
self, client: ManagementClient, resources: ResourceManager
) -> None:
team_id = _create_team(client, resources, f"e2e-audit-team-{unique_marker()}", [])
user_id = _create_user(
client,
resources,
UserNewBody(user_email=f"e2e-audit-{unique_marker()}@example.com", user_role="internal_user"),
)
client.add_team_member(team_id, user_id)
created = _generate_response(client, resources, KeyGenerateBody(user_id=user_id, team_id=team_id))
token = _token_of(created)
client.delete_team_member(team_id, user_id)
_assert_key_deleted(client, created.key)
_assert_single_deleted_row(_await_deleted_audit_rows(client, token), token)
@pytest.mark.covers("mgmt.team.delete.audit_logs_keys")
def test_team_delete_writes_audit_row_for_team_keys(
self, client: ManagementClient, resources: ResourceManager
) -> None:
team_id = _create_team(client, resources, f"e2e-audit-team-{unique_marker()}", [])
created = _generate_response(client, resources, KeyGenerateBody(team_id=team_id))
token = _token_of(created)
client.delete_team(team_id)
_assert_key_deleted(client, created.key)
_assert_single_deleted_row(_await_deleted_audit_rows(client, token), token)
@pytest.mark.covers("mgmt.user.delete.audit_logs_keys")
def test_user_delete_writes_audit_row_for_user_keys(
self, client: ManagementClient, resources: ResourceManager
) -> None:
user_id = _create_user(
client,
resources,
UserNewBody(user_email=f"e2e-audit-{unique_marker()}@example.com", user_role="internal_user"),
)
created = _generate_response(client, resources, KeyGenerateBody(user_id=user_id))
token = _token_of(created)
client.delete_user_strict(user_id)
_assert_key_deleted(client, created.key)
_assert_single_deleted_row(_await_deleted_audit_rows(client, token), token)

26
tests/e2e/models.py
View file

@ -124,6 +124,32 @@ class KeyDeleteBody(BaseModel):
keys: list[str]
class KeyDeleteByAliasBody(BaseModel):
key_aliases: list[str]
class AuditLogParams(BaseModel):
object_id: str
action: str
table_name: str
page_size: int
class AuditLogEntry(BaseModel):
id: str
changed_by: str | None = None
changed_by_api_key: str | None = None
action: str
table_name: str
object_id: str
before_value: object | None = None
class AuditLogPage(BaseModel):
audit_logs: list[AuditLogEntry]
total: int
class KeyInfoParams(BaseModel):
key: str

2
tests/e2e/transport.py
View file

@ -313,6 +313,8 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
"/config",
"/guardrails",
"/router/settings",
"/audit",
"/v2/login",
"/openapi.json",
)

46
tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py
View file

@ -701,3 +701,49 @@ class TestKeyUpdatedAuditLogObjectId:
assert updated_values["project_id"] is None
assert json.loads(audit_row.before_value)["project_id"] == "project-orbit"
assert updated_values["max_budget"] == 2000.0
@pytest.mark.asyncio
async def test_key_deleted_hook_writes_audit_log_for_alias_deletion():
from litellm.proxy._types import (
KeyRequest,
LiteLLM_AuditLogs,
LiteLLM_VerificationToken,
LitellmTableNames,
UserAPIKeyAuth,
)
captured: Final[list[LiteLLM_AuditLogs]] = []
async def capture_audit_log(request_data: LiteLLM_AuditLogs) -> None:
captured.append(request_data)
with (
patch("litellm.store_audit_logs", True),
patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=capture_audit_log,
),
patch.object(
KeyManagementEventHooks,
"_delete_virtual_keys_from_secret_manager",
new_callable=AsyncMock,
),
):
await KeyManagementEventHooks.async_key_deleted_hook(
data=KeyRequest(key_aliases=["a"]),
keys_being_deleted=[LiteLLM_VerificationToken(token="hashed", key_alias="a")],
response={},
user_api_key_dict=UserAPIKeyAuth(user_id="admin", token="callertok"),
)
for _ in range(100):
if captured:
break
await asyncio.sleep(0.01)
assert len(captured) == 1
audit_row = captured[0]
assert audit_row.action == "deleted"
assert audit_row.object_id == "hashed"
assert audit_row.table_name == LitellmTableNames.KEY_TABLE_NAME
assert audit_row.changed_by == "admin"

60
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -1,3 +1,4 @@
import asyncio
import hashlib
import json
import logging
@ -4779,3 +4780,62 @@ def test_user_routes_answer_503_no_db_connection_when_the_callers_user_read_hits
assert response.status_code == 503, response.text
assert response.json() == _DB_OUTAGE_503_BODY
@pytest.mark.asyncio
async def test_delete_user_writes_deleted_audit_log_for_user_keys(mocker):
from litellm.proxy._types import (
DeleteUserRequest,
LiteLLM_VerificationToken,
LitellmTableNames,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user
mock_prisma_client = mocker.MagicMock()
mock_user_row = mocker.MagicMock()
mock_user_row.user_id = "doomed-user"
mock_user_row.user_email = "doomed@example.com"
mock_user_row.teams = []
mock_user_row.model_dump_json.return_value = "{}"
mock_user_row.model_dump.return_value = {"user_id": "doomed-user", "user_email": "doomed@example.com", "teams": []}
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=mock_user_row)
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
user_key = LiteLLM_VerificationToken(token="hashed-user-key", user_id="doomed-user")
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[user_key])
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=1)
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
mock_prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[])
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch("litellm.store_audit_logs", True)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=_capture,
)
caller = UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
await delete_user(data=DeleteUserRequest(user_ids=["doomed-user"]), user_api_key_dict=caller)
for _ in range(100):
if captured:
break
await asyncio.sleep(0.01)
key_rows: Final = [r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]
assert len(key_rows) == 1
audit_row: Final = key_rows[0]
assert audit_row.action == "deleted"
assert audit_row.object_id == user_key.token
assert audit_row.changed_by
assert json.loads(audit_row.before_value)["token"] == user_key.token

162
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -5129,6 +5129,168 @@ async def test_team_member_delete_cleans_verification_tokens(
)
@pytest.mark.asyncio
async def test_team_member_delete_writes_deleted_audit_log_for_member_keys(
mock_db_client, mock_admin_auth
):
from litellm.proxy._types import (
LiteLLM_VerificationToken,
LitellmTableNames,
TeamMemberDeleteRequest,
)
from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
test_team_id = "team-del-audit-123"
test_user_id = "user-audit@example.com"
member_key = LiteLLM_VerificationToken(token="hashed-member-key", team_id=test_team_id, user_id=test_user_id)
mock_team_row = MagicMock()
mock_team_row.model_dump.return_value = {
"team_id": test_team_id,
"members_with_roles": [
{"user_id": test_user_id, "user_email": None, "role": "user"}
],
"team_member_permissions": [],
"metadata": {},
"models": [],
"spend": 0.0,
}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team_row
)
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row)
mock_user_row = MagicMock()
mock_user_row.user_id = test_user_id
mock_user_row.teams = [test_team_id]
mock_db_client.db.litellm_usertable.find_many = AsyncMock(
return_value=[mock_user_row]
)
mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
mock_db_client.db.litellm_teammembership = MagicMock()
mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(
return_value=MagicMock()
)
mock_db_client.db.litellm_verificationtoken = MagicMock()
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[member_key])
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(
return_value=MagicMock()
)
mock_db_client.db.litellm_jwtkeymapping = MagicMock()
mock_db_client.db.litellm_jwtkeymapping.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_deletedverificationtoken = MagicMock()
mock_db_client.db.litellm_deletedverificationtoken.create_many = AsyncMock(return_value=MagicMock())
_wire_member_delete_tx(mock_db_client)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
with (
patch("litellm.store_audit_logs", True),
patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=_capture,
),
):
await team_member_delete(
data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id),
user_api_key_dict=mock_admin_auth,
)
for _ in range(100):
if captured:
break
await asyncio.sleep(0.01)
key_rows: Final = [r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]
assert len(key_rows) == 1
audit_row: Final = key_rows[0]
assert audit_row.action == "deleted"
assert audit_row.object_id == member_key.token
assert audit_row.changed_by
assert json.loads(audit_row.before_value)["token"] == member_key.token
@pytest.mark.asyncio
async def test_delete_team_writes_deleted_audit_log_for_team_keys(
monkeypatch,
):
from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken, LitellmTableNames
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
team = LiteLLM_TeamTable(
team_id="team-doomed",
team_alias="doomed-team",
members_with_roles=[],
metadata={},
model_max_budget={},
model_spend={},
)
team_key = LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]})
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[team_key])
mock_prisma_client.db.execute_raw = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
mock_prisma_client.get_data = AsyncMock(return_value=None)
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_tx_cm = MagicMock()
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
_wire_team_delete_tx(mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache())
monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
caller = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
monkeypatch.setattr("litellm.store_audit_logs", True)
monkeypatch.setattr(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
_capture,
)
await delete_team(
data=DeleteTeamRequest(team_ids=["team-doomed"]),
http_request=MagicMock(),
user_api_key_dict=caller,
litellm_changed_by="admin-user",
)
for _ in range(100):
if captured:
break
await asyncio.sleep(0.01)
key_rows: Final = [r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]
assert len(key_rows) == 1
audit_row: Final = key_rows[0]
assert audit_row.action == "deleted"
assert audit_row.object_id == team_key.token
assert audit_row.changed_by
assert json.loads(audit_row.before_value)["token"] == team_key.token
@pytest.mark.asyncio
async def test_team_member_delete_reads_on_the_lock_holding_transaction(
mock_db_client, mock_admin_auth

96
tests/test_litellm/proxy/management_helpers/test_bulk_user_deletion.py
View file

@ -1,3 +1,4 @@
import asyncio
import copy
import json
from collections.abc import Callable, Mapping, Sequence
@ -7,7 +8,7 @@ from typing import Final
import pytest
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth
from litellm.proxy._types import LiteLLM_TeamTable, LitellmTableNames, LitellmUserRoles, Member, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.list_api.common import ManagementProblem
@ -748,3 +749,96 @@ def test_request_models_reject_unknown_fields():
BulkTeamMemberDeleteRequest.model_validate({"members": [{"user_id": "u1", "role": "admin"}]})
with pytest.raises(ValidationError, match="dry_run"):
BulkDeleteUserRequest.model_validate({"user_ids": ["u1"], "dry_run": True})
@pytest.mark.asyncio
async def test_bulk_delete_writes_deleted_audit_log_for_deleted_keys(mocker):
prisma = _FakePrisma(
users=[_user("u1", "t1"), _user("keep", "t1")],
teams=[_team("t1", "u1", "keep")],
tokens=[
{"token": "team-key", "user_id": "u1", "team_id": "t1"},
{"token": "personal-key", "user_id": "u1"},
{"token": "keep-key", "user_id": "keep", "team_id": "t1"},
],
)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
mocker.patch("litellm.store_audit_logs", True)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=_capture,
)
await _delete(prisma, ["u1"])
for _ in range(100):
if len([r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]) >= 2:
break
await asyncio.sleep(0.01)
key_rows: Final = [r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]
assert {r.object_id for r in key_rows} == {"team-key", "personal-key"}
assert {r.action for r in key_rows} == {"deleted"}
assert all(r.changed_by for r in key_rows)
assert {json.loads(r.before_value)["token"] for r in key_rows} == {"team-key", "personal-key"}
@pytest.mark.asyncio
async def test_bulk_member_delete_writes_deleted_audit_log_for_removed_team_keys(mocker):
prisma = _FakePrisma(
users=[_user("u1", "t1"), _user("keep", "t1")],
teams=[_team("t1", "u1", "keep")],
tokens=[
{"token": "team-key", "user_id": "u1", "team_id": "t1"},
{"token": "keep-key", "user_id": "keep", "team_id": "t1"},
],
)
captured: Final[list] = []
async def _capture(request_data):
captured.append(request_data)
mocker.patch("litellm.store_audit_logs", True)
mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=_capture,
)
await _remove(prisma, "t1", [{"user_id": "u1"}])
for _ in range(100):
if captured:
break
await asyncio.sleep(0.01)
key_rows: Final = [r for r in captured if r.table_name == LitellmTableNames.KEY_TABLE_NAME]
assert len(key_rows) == 1
audit_row: Final = key_rows[0]
assert audit_row.action == "deleted"
assert audit_row.object_id == "team-key"
assert audit_row.changed_by
assert json.loads(audit_row.before_value)["token"] == "team-key"
@pytest.mark.asyncio
async def test_bulk_delete_skips_the_key_audit_log_when_the_tx_rolls_back(mocker):
prisma = _FakePrisma(
users=[_user("u1", "a-good", "z-bad"), _user("u2", "a-good")],
teams=[_team("a-good", "u1", "u2"), _team("z-bad", "u1")],
tokens=[{"token": "k1", "user_id": "u1", "team_id": "a-good"}],
fail_locks=frozenset({"z-bad"}),
)
cache = _cache_with("k1")
mocker.patch("litellm.store_audit_logs", True)
mock_audit_write = mocker.patch(
"litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update",
new=mocker.AsyncMock(),
)
results = await _delete(prisma, ["u1", "u2"], cache=cache)
assert [(r.user_id, r.success) for r in results] == [("u1", False), ("u2", False)]
assert [t["token"] for t in prisma.db.litellm_verificationtoken.rows] == ["k1"]
mock_audit_write.assert_not_called()