scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172

fix(proxy): track and drain management audit log writes on shutdown (#43594)

Browse source 
* fix(proxy): track and drain management audit log writes on shutdown

* fix(proxy): track hook spawns and defer registration in audit log drain

* fix: satisfy LIT006 on audit log task registry by documenting the required cast

* fix(proxy): satisfy LIT006 and LIT010 in audit task registry

* test(proxy): drive the audit drain tests with events instead of wall-clock waits

* fix(proxy): drain audit writes beside the shutdown flushes and track only premium writes

Run the audit drain concurrently with the gateway and request-error flushes so
the two waits overlap inside the teardown budget, register the running audit
task only after the premium check, drop the four duplicate task-failure
loggers, and cover the shutdown order and a real hook spawn with tests

* fix(proxy): start hook audit writes before their email and secret manager calls

Key and user hooks now spawn their tracked audit write first, so a slow
invite email or secret manager call no longer spends the drain budget
before the write starts. The endpoints stop tracking the hook tasks, so
shutdown never waits on hook I/O when audit logging is off.

Adds regression tests for the hook order, the failed block or unblock
update, and the shutdown drain without a database.

* fix(proxy): track only the audit insert, not the code that awaits it

The shutdown drain now waits on each audit row insert itself, so it never waits on a caller's later work and covers every caller without per-endpoint wrappers. The user/new hook still tracks its own task because it reads the user row before the insert

* test(proxy): drain the audit registry in the config audit tests

The audit insert now runs in its own tracked task, one hop after the caller's task, so a single yield no longer reaches it

* fix(proxy): name dropped audit rows in the drain timeout warning

Each tracked audit task is named after its action, table, object id and destination, so a shutdown that drops a row says which one. Adds tests that the drain waits for slow audit callbacks and that the warning names the dropped row

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Tufail Akram 2026-10-16 14:11:34 +05:30 • committed by GitHub
parent ecbea93381
commit 84eeebd509
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 765 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -57,13 +57,6 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
# Send email notification - non-blocking, independent operation
if data.send_invite_email is True:
try:
await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True))
except Exception as e:
verbose_proxy_logger.warning("Failed to send key created email: %s", e)
if is_audit_logging_enabled():
_updated_values: Final = response.model_dump_json(exclude_none=True)
asyncio.create_task(
@ -86,6 +79,13 @@ class KeyManagementEventHooks:
)
)
# Send email notification - non-blocking, independent operation
if data.send_invite_email is True:
try:
await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True))
except Exception as e:
verbose_proxy_logger.warning("Failed to send key created email: %s", e)
# Store the generated key in the secret manager - non-blocking, independent operation
try:
await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
@ -118,16 +118,6 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
if data.key_alias is not None and data.key_alias != existing_key_row.key_alias:
try:
await KeyManagementEventHooks._rename_virtual_key_in_secret_manager(
current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}",
new_secret_name=data.key_alias,
team_id=existing_key_row.team_id,
)
except Exception as e:
verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e)
if is_audit_logging_enabled():
updated_fields: Final = {
**data.model_dump(exclude_none=True),
@ -156,6 +146,16 @@ class KeyManagementEventHooks:
)
asyncio.create_task(create_audit_log_for_update(request_data=request_data))
if data.key_alias is not None and data.key_alias != existing_key_row.key_alias:
try:
await KeyManagementEventHooks._rename_virtual_key_in_secret_manager(
current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}",
new_secret_name=data.key_alias,
team_id=existing_key_row.team_id,
)
except Exception as e:
verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e)
@staticmethod
async def async_key_rotated_hook(
data: RegenerateKeyRequest | None,
@ -171,6 +171,28 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
# store the audit log
if is_audit_logging_enabled() and existing_key_row.token is not None:
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=existing_key_row.token,
action="rotated",
updated_values=response.model_dump_json(exclude_none=True),
before_value=existing_key_row.model_dump_json(exclude_none=True),
)
)
)
# Store the generated key in the secret manager - non-blocking, independent operation
if response.token_id is not None:
try:
@ -204,28 +226,6 @@ class KeyManagementEventHooks:
except Exception as e:
verbose_proxy_logger.warning("Failed to send key rotated email: %s", e)
# store the audit log
if is_audit_logging_enabled() and existing_key_row.token is not None:
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=existing_key_row.token,
action="rotated",
updated_values=response.model_dump_json(exclude_none=True),
before_value=existing_key_row.model_dump_json(exclude_none=True),
)
)
)
@staticmethod
async def async_key_deleted_hook(
data: KeyRequest,

43
litellm/proxy/hooks/user_management_event_hooks.py
View file

@ -23,6 +23,7 @@ from litellm.proxy._types import (
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
is_audit_logging_enabled,
track_audit_task,
)
from litellm.repositories.user_repository import UserRepository
@ -41,7 +42,18 @@ class UserManagementEventHooks:
- Creating an audit log for the user creation
- Sending a user invitation email to the user
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import litellm_proxy_admin_name
if is_audit_logging_enabled():
track_audit_task(
asyncio.create_task(
UserManagementEventHooks._create_user_created_audit_log(
response=response,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
)
)
#########################################################
########## Send User Invitation Email ################
@ -52,9 +64,14 @@ class UserManagementEventHooks:
user_api_key_dict=user_api_key_dict,
)
#########################################################
########## CREATE AUDIT LOG ################
#########################################################
@staticmethod
async def _create_user_created_audit_log(
response: NewUserResponse,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str | None,
) -> None:
from litellm.proxy.proxy_server import prisma_client
try:
if prisma_client is None:
raise Exception(CommonProxyErrors.db_not_connected_error.value)
@ -63,16 +80,14 @@ class UserManagementEventHooks:
user_row: Final = await UserRepository(prisma_client).find_by_id(response.user_id)
if user_row is None:
raise Exception(f"no user row found for user_id={response.user_id}")
asyncio.create_task(
UserManagementEventHooks.create_internal_user_audit_log(
user_id=user_row.user_id,
action="created",
litellm_changed_by=user_api_key_dict.user_id,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
before_value=None,
after_value=user_row.model_dump_json(exclude_none=True),
)
await UserManagementEventHooks.create_internal_user_audit_log(
user_id=user_row.user_id,
action="created",
litellm_changed_by=user_api_key_dict.user_id,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
before_value=None,
after_value=user_row.model_dump_json(exclude_none=True),
)
except Exception as e:
verbose_proxy_logger.warning("Unable to create audit log for user on `/user/new` - %s", e)

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

@ -7713,6 +7713,11 @@ async def block_key(
code=status.HTTP_404_NOT_FOUND,
)
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
where={"token": hashed_token},
data=with_settings_updated_at({"blocked": True}),
)
if is_audit_logging_enabled():
asyncio.create_task(
create_audit_log_for_update(
@ -7734,11 +7739,6 @@ async def block_key(
)
)
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
where={"token": hashed_token},
data=with_settings_updated_at({"blocked": True}),
)
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await delete_cache_key_object(
hashed_token=hashed_token,
@ -7827,6 +7827,11 @@ async def unblock_key(
code=status.HTTP_404_NOT_FOUND,
)
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
where={"token": hashed_token},
data=with_settings_updated_at({"blocked": False}),
)
if is_audit_logging_enabled():
asyncio.create_task(
create_audit_log_for_update(
@ -7848,11 +7853,6 @@ async def unblock_key(
)
)
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
where={"token": hashed_token},
data=with_settings_updated_at({"blocked": False}),
)
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await delete_cache_key_object(
hashed_token=hashed_token,

83
litellm/proxy/management_helpers/audit_logs.py
View file

@ -5,7 +5,7 @@ Functions to create audit logs for LiteLLM Proxy
import asyncio
import json
from datetime import datetime, timezone
from typing import Final
from typing import Final, TypeVar
import litellm
from litellm._logging import verbose_proxy_logger
@ -134,6 +134,56 @@ def _audit_log_task_done_callback(task: asyncio.Task) -> None:
verbose_proxy_logger.error("Audit log callback task failed: %s", exc, exc_info=exc)
_T: Final = TypeVar("_T")
class AuditTaskRegistry:
__slots__ = ("_pending_tasks",)
def __init__(self) -> None:
self._pending_tasks: Final[set[asyncio.Task[object]]] = set() # mutable-ok: background task tracking
def track(self, task: asyncio.Task[_T]) -> asyncio.Task[_T]:
if task.done() or task in self._pending_tasks:
return task
self._pending_tasks.add(task)
task.add_done_callback(self._on_task_done)
return task
def _on_task_done(self, task: asyncio.Task[object]) -> None:
self._pending_tasks.discard(task)
_audit_log_task_done_callback(task)
async def drain(self, timeout: float = 5.0) -> None:
loop: Final = asyncio.get_running_loop()
deadline: Final = loop.time() + timeout
while self._pending_tasks:
_, pending = await asyncio.wait(tuple(self._pending_tasks), timeout=max(deadline - loop.time(), 0))
if pending:
self._pending_tasks.difference_update(pending)
verbose_proxy_logger.warning(
"Timed out draining %d audit tasks on shutdown, dropping: %s",
len(pending),
", ".join(sorted(task.get_name() for task in pending)),
)
return
_audit_task_registry: Final = AuditTaskRegistry()
def track_audit_task(task: asyncio.Task[_T]) -> asyncio.Task[_T]:
return _audit_task_registry.track(task)
async def drain_audit_tasks(timeout: float = 5.0) -> None:
await _audit_task_registry.drain(timeout=timeout)
def _audit_task_name(request_data: LiteLLM_AuditLogs, destination: str) -> str:
return f"{request_data.action} {request_data.table_name.value} {request_data.object_id} to {destination}"
async def _dispatch_audit_log_to_callbacks(
request_data: LiteLLM_AuditLogs,
) -> None:
@ -153,8 +203,12 @@ async def _dispatch_audit_log_to_callbacks(
continue
if isinstance(resolved, CustomLogger):
task = asyncio.create_task(resolved.async_log_audit_log_event(payload))
task.add_done_callback(_audit_log_task_done_callback)
track_audit_task(
asyncio.create_task(
resolved.async_log_audit_log_event(payload),
name=_audit_task_name(request_data, type(resolved).__name__),
)
)
except Exception as e:
verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e)
@ -204,6 +258,14 @@ async def create_object_audit_log(
)
async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None:
try:
await repository.table.create(data={**request_data.model_dump(exclude_none=True)})
except Exception as e:
# [Non-Blocking Exception. Do not allow blocking LLM API call]
verbose_proxy_logger.error("Failed Creating audit log %s", e)
async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
"""
Create an audit log for an object.
@ -231,14 +293,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
return
_request_data: Final = request_data.model_dump(exclude_none=True)
try:
await AuditLogRepository(prisma_client).table.create(
data={
**_request_data,
}
await track_audit_task(
asyncio.create_task(
_insert_audit_log_row(AuditLogRepository(prisma_client), request_data),
name=_audit_task_name(request_data, "the database"),
)
except Exception as e:
# [Non-Blocking Exception. Do not allow blocking LLM API call]
verbose_proxy_logger.error("Failed Creating audit log %s", e)
)

19
litellm/proxy/proxy_server.py
View file

@ -1200,6 +1200,20 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e)
async def _drain_audit_tasks_on_shutdown() -> None:
try:
from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
await drain_audit_tasks()
except Exception as e: # noqa: BLE001 # shutdown must continue even if the drain fails
verbose_proxy_logger.exception("Error draining audit tasks on shutdown: %s", e)
async def _flush_request_accumulators(prisma: PrismaClient) -> None:
await flush_gateway_requests(prisma, gateway_request_accumulator)
await flush_request_errors(prisma, request_error_accumulator)
async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None:
global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
@ -1210,10 +1224,11 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
# is lost, and a write attempted after disconnect raises
# ClientNotConnectedError rather than persisting anything. Ordering this
# inside the same guard is what keeps the two from drifting apart.
await flush_gateway_requests(prisma_client, gateway_request_accumulator)
await flush_request_errors(prisma_client, request_error_accumulator)
await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client))
verbose_proxy_logger.debug("Disconnecting from Prisma")
await prisma_client.disconnect()
else:
await _drain_audit_tasks_on_shutdown()
await drain_post_call_redis_batches()
if litellm.cache is not None:

37
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -3407,6 +3407,43 @@ async def test_block_key_existing_key_succeeds(monkeypatch):
assert result == mock_updated_record
@pytest.mark.asyncio
@pytest.mark.parametrize("endpoint_name", ["block_key", "unblock_key"])
async def test_block_and_unblock_write_no_audit_row_when_the_key_update_fails(monkeypatch, endpoint_name):
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints import key_management_endpoints
from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
existing_key_record = MagicMock()
existing_key_record.token = hashed_token
existing_key_record.model_dump_json.return_value = f'{{"token": "{hashed_token}"}}'
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_record)
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed"))
audit_writer = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", audit_writer)
monkeypatch.setattr("litellm.store_audit_logs", True)
with pytest.raises(RuntimeError, match="update failed"):
await getattr(key_management_endpoints, endpoint_name)(
data=BlockKeyRequest(key=hashed_token),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
),
litellm_changed_by=None,
)
await drain_audit_tasks(timeout=1.0)
assert audit_writer.call_count == 0
@pytest.mark.asyncio
async def test_validate_key_team_change_with_member_permissions():
"""

496
tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
View file

@ -1,14 +1,18 @@
import os
import traceback
from litellm._uuid import uuid
from datetime import datetime
from datetime import datetime, timezone
from collections.abc import Coroutine
from typing import Any, Final
from dotenv import load_dotenv
from fastapi import Request
from fastapi.routing import APIRoute
import contextlib
import io
import logging
import time
# this file is to test litellm/proxy
@ -38,11 +42,16 @@ from starlette.datastructures import URL
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
drain_audit_tasks,
get_audit_log_changed_by,
track_audit_task,
)
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from unittest.mock import patch, AsyncMock
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardAuditLogPayload
from unittest.mock import AsyncMock, MagicMock, patch
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
import json
@ -174,7 +183,6 @@ async def test_create_audit_log_for_update_premium_user():
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
):
mock_prisma.db.litellm_auditlog.create = AsyncMock()
request_data = LiteLLM_AuditLogs(
@ -215,9 +223,7 @@ def prisma_client():
os.environ["DATABASE_URL"] = modified_url
# Assuming PrismaClient is a class that needs to be instantiated
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
prisma_client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj)
return prisma_client
@ -252,10 +258,482 @@ async def test_create_audit_log_in_db(prisma_client):
await asyncio.sleep(1)
# now read the last log from the db
last_log = await prisma_client.db.litellm_auditlog.find_first(
where={"id": audit_log_id}
)
last_log = await prisma_client.db.litellm_auditlog.find_first(where={"id": audit_log_id})
assert last_log.id == audit_log_id
setattr(litellm, "store_audit_logs", False)
@pytest.mark.asyncio
async def test_track_audit_task_lifecycle():
task_completed: Final = asyncio.Event()
async def _sample_coroutine() -> None:
await asyncio.sleep(0)
task_completed.set()
task: Final = track_audit_task(asyncio.create_task(_sample_coroutine()))
assert not task.done()
await drain_audit_tasks(timeout=1.0)
assert task_completed.is_set()
assert task.done()
completed_task: Final = asyncio.create_task(asyncio.sleep(0))
await completed_task
assert track_audit_task(completed_task) is completed_task
@pytest.mark.asyncio
async def test_drain_audit_tasks_waits_for_all_tasks():
event1: Final = asyncio.Event()
event2: Final = asyncio.Event()
async def _worker(event: asyncio.Event) -> None:
await asyncio.sleep(0)
event.set()
task1: Final = track_audit_task(asyncio.create_task(_worker(event1)))
task2: Final = track_audit_task(asyncio.create_task(_worker(event2)))
assert not task1.done()
assert not task2.done()
await drain_audit_tasks(timeout=1.0)
assert event1.is_set()
assert event2.is_set()
assert task1.done()
assert task2.done()
@pytest.mark.asyncio
async def test_drain_audit_tasks_handles_failing_task_cleanly():
async def _failing_worker() -> None:
await asyncio.sleep(0)
raise RuntimeError("test audit log write failure")
task: Final = track_audit_task(asyncio.create_task(_failing_worker()))
assert not task.done()
await drain_audit_tasks(timeout=1.0)
assert task.done()
assert isinstance(task.exception(), RuntimeError)
@pytest.mark.asyncio
async def test_drain_audit_tasks_timeout_does_not_cancel_pending_tasks():
async def _long_running_worker() -> None:
await asyncio.Event().wait()
task: Final = track_audit_task(asyncio.create_task(_long_running_worker()))
try:
assert not task.done()
await drain_audit_tasks(timeout=0)
assert not task.cancelled()
assert not task.done()
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
@pytest.mark.asyncio
async def test_drain_audit_tasks_persists_management_audit_write_before_db_disconnect():
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
is_disconnected: Final = [False]
async def _mock_create(data: dict[str, Any]) -> None:
if is_disconnected[0]:
raise RuntimeError("Database already disconnected")
await asyncio.sleep(0)
writes.append(data)
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create)
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
request_data: Final = LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by="admin-user",
changed_by_api_key="sk-admin",
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id="token-123",
action="blocked",
updated_values="{}",
before_value='{"key": "test"}',
)
task: Final = track_audit_task(asyncio.create_task(create_audit_log_for_update(request_data=request_data)))
assert not task.done()
await drain_audit_tasks(timeout=2.0)
is_disconnected[0] = True
assert task.done()
assert len(writes) == 1
assert writes[0]["object_id"] == "token-123"
assert writes[0]["action"] == "blocked"
@pytest.mark.asyncio
async def test_create_audit_log_auto_registers_in_drain():
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
async def _mock_create(data: dict[str, Any]) -> None:
await asyncio.sleep(0)
writes.append(data)
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create)
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
request_data: Final = LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by="admin-user",
changed_by_api_key="sk-admin",
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id="token-456",
action="deleted",
updated_values="{}",
before_value='{"key": "test-deleted"}',
)
task: Final = asyncio.create_task(create_audit_log_for_update(request_data=request_data))
await asyncio.sleep(0)
await drain_audit_tasks(timeout=2.0)
assert task.done()
assert len(writes) == 1
assert writes[0]["object_id"] == "token-456"
assert writes[0]["action"] == "deleted"
@pytest.mark.asyncio
async def test_drain_audit_tasks_discards_timed_out_tasks_on_timeout():
async def _hung_worker() -> None:
await asyncio.Event().wait()
hung_task: Final = track_audit_task(asyncio.create_task(_hung_worker()))
try:
await drain_audit_tasks(timeout=0)
assert not hung_task.done()
second_drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
for _ in range(3):
await asyncio.sleep(0)
assert second_drain.done()
finally:
hung_task.cancel()
try:
await hung_task
except asyncio.CancelledError:
pass
@pytest.mark.asyncio
async def test_drain_audit_tasks_captures_tasks_queued_during_drain():
finished_tasks: Final[list[str]] = []
async def _first_worker() -> None:
await asyncio.sleep(0)
track_audit_task(asyncio.create_task(_second_worker()))
finished_tasks.append("first")
async def _second_worker() -> None:
await asyncio.sleep(0)
finished_tasks.append("second")
track_audit_task(asyncio.create_task(_first_worker()))
await drain_audit_tasks(timeout=2.0)
assert finished_tasks == ["first", "second"]
class _AuditCallbackBlockedUntilReleased(CustomLogger):
def __init__(self, release: asyncio.Event) -> None:
super().__init__()
self.release: Final = release
self.logged_object_ids: Final[list[str]] = []
async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None:
await self.release.wait()
self.logged_object_ids.append(audit_log["object_id"])
class _WarningRecorder(logging.Handler):
def __init__(self) -> None:
super().__init__(level=logging.WARNING)
self.messages: Final[list[str]] = []
def emit(self, record: logging.LogRecord) -> None:
self.messages.append(record.getMessage())
def _deleted_key_audit_log(object_id: str) -> LiteLLM_AuditLogs:
return LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by="admin-user",
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=object_id,
action="deleted",
updated_values="{}",
before_value="{}",
)
@pytest.mark.asyncio
async def test_drain_waits_for_a_slow_audit_log_callback():
release: Final = asyncio.Event()
callback: Final = _AuditCallbackBlockedUntilReleased(release)
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.audit_log_callbacks", [callback]),
):
await create_audit_log_for_update(request_data=_deleted_key_audit_log("token-callback"))
drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
for _ in range(3):
await asyncio.sleep(0)
drain_returned_before_the_callback_finished: Final = drain.done()
release.set()
await drain
assert not drain_returned_before_the_callback_finished
assert callback.logged_object_ids == ["token-callback"]
@pytest.mark.asyncio
async def test_drain_timeout_warning_names_each_dropped_audit_row():
mock_prisma: Final = MagicMock()
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_io_that_never_finishes)
recorder: Final = _WarningRecorder()
verbose_proxy_logger.addHandler(recorder)
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
writer: Final = asyncio.create_task(
create_audit_log_for_update(request_data=_deleted_key_audit_log("token-dropped"))
)
try:
for _ in range(3):
await asyncio.sleep(0)
await drain_audit_tasks(timeout=0.05)
finally:
verbose_proxy_logger.removeHandler(recorder)
writer.cancel()
with contextlib.suppress(asyncio.CancelledError):
await writer
assert recorder.messages == [
"Timed out draining 1 audit tasks on shutdown, dropping: "
f"deleted {LitellmTableNames.KEY_TABLE_NAME.value} token-dropped to the database"
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("store_audit_logs", "premium_user", "expected_write_count"),
[(True, True, 1), (False, True, 0), (True, False, 0)],
ids=["row_written", "logging_disabled", "not_premium"],
)
async def test_drain_does_not_wait_for_code_that_runs_after_an_inline_audit_write(
store_audit_logs: bool, premium_user: bool, expected_write_count: int
):
release_caller: Final = asyncio.Event()
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
async def _caller_that_outlives_the_audit_call() -> None:
await create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by="test-user",
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id="test-obj",
action="updated",
updated_values="{}",
before_value="{}",
)
)
await release_caller.wait()
with (
patch("litellm.store_audit_logs", store_audit_logs),
patch("litellm.proxy.proxy_server.premium_user", premium_user),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
caller: Final = asyncio.create_task(_caller_that_outlives_the_audit_call())
for _ in range(3):
await asyncio.sleep(0)
drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
for _ in range(3):
await asyncio.sleep(0)
drain_returned_while_caller_was_running: Final = drain.done()
release_caller.set()
await caller
await drain
assert drain_returned_while_caller_was_running
assert len(writes) == expected_write_count
async def _io_that_never_finishes(*args: object, **kwargs: object) -> None:
await asyncio.Event().wait()
async def _drain_while_hook_io_hangs(hook: Coroutine[object, object, None]) -> None:
hook_task: Final = asyncio.create_task(hook)
for _ in range(3):
await asyncio.sleep(0)
await drain_audit_tasks(timeout=2.0)
hook_task.cancel()
@pytest.mark.asyncio
async def test_key_generated_audit_write_is_drained_while_the_invite_email_hangs():
from litellm.proxy._types import GenerateKeyRequest, GenerateKeyResponse
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(KeyManagementEventHooks, "_send_key_created_email", side_effect=_io_that_never_finishes),
):
await _drain_while_hook_io_hangs(
KeyManagementEventHooks.async_key_generated_hook(
data=GenerateKeyRequest(send_invite_email=True),
response=GenerateKeyResponse(key="sk-new", token_id="hashed-new"),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
)
)
assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-new", "created")]
@pytest.mark.asyncio
async def test_key_updated_audit_write_is_drained_while_the_secret_manager_rename_hangs():
from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(
KeyManagementEventHooks, "_rename_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes
),
):
await _drain_while_hook_io_hangs(
KeyManagementEventHooks.async_key_updated_hook(
data=UpdateKeyRequest(key="sk-existing", key_alias="new-alias"),
existing_key_row=LiteLLM_VerificationToken(token="hashed-existing", key_alias="old-alias"),
response=None,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
)
)
assert [(write["object_id"], write["action"]) for write in writes] == [(hash_token("sk-existing"), "updated")]
@pytest.mark.asyncio
async def test_key_rotated_audit_write_is_drained_while_the_secret_manager_rotation_hangs():
from litellm.proxy._types import GenerateKeyResponse, LiteLLM_VerificationToken
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(
KeyManagementEventHooks, "_rotate_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes
),
):
await _drain_while_hook_io_hangs(
KeyManagementEventHooks.async_key_rotated_hook(
data=None,
existing_key_row=LiteLLM_VerificationToken(token="hashed-old"),
response=GenerateKeyResponse(key="sk-rotated", token_id="hashed-new"),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
)
)
assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-old", "rotated")]
@pytest.mark.asyncio
async def test_user_created_audit_write_is_drained_while_its_user_read_is_still_running():
from litellm.proxy._types import NewUserRequest, NewUserResponse
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
mock_prisma: Final = MagicMock()
writes: Final[list[dict[str, Any]]] = []
release_user_read: Final = asyncio.Event()
async def _slow_user_read(*args: object, **kwargs: object) -> dict[str, str]:
await release_user_read.wait()
return {"user_id": "new-user"}
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=_slow_user_read)
with (
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(UserManagementEventHooks, "async_send_user_invitation_email", side_effect=_io_that_never_finishes),
):
hook: Final = asyncio.create_task(
UserManagementEventHooks.async_user_created_hook(
data=NewUserRequest(user_id="new-user", send_invite_email=True),
response=NewUserResponse(key="sk-new", user_id="new-user"),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
)
)
for _ in range(3):
await asyncio.sleep(0)
drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
for _ in range(3):
await asyncio.sleep(0)
drain_waited_for_the_user_read: Final = not drain.done()
release_user_read.set()
await drain
hook.cancel()
assert drain_waited_for_the_user_read
assert [(write["object_id"], write["action"]) for write in writes] == [("new-user", "created")]

72
tests/unit/proxy/proxy_server/test_lifecycle.py
View file

@ -17,12 +17,14 @@ Pins covered:
from __future__ import annotations
import asyncio
import inspect
import json
import logging
import os
import subprocess
from collections.abc import Awaitable, Callable
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -48,6 +50,8 @@ from litellm.proxy.proxy_server import (
save_worker_config,
)
from litellm.proxy.management_helpers.audit_logs import track_audit_task
from .conftest import normalize
# ---------------------------------------------------------------------------
@ -223,6 +227,46 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
assert calls == ["flush", "disconnect"]
@pytest.mark.asyncio
async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_before_disconnecting(monkeypatch):
"""
A management call returns before its audit row is written, so shutdown waits
for that write while prisma is still connected. It waits beside the gateway
flush rather than before it, so the two waits overlap inside the teardown budget.
"""
calls: list[str] = [] # mutable-ok: records call order, which is the assertion
gateway_flush_started: Final = asyncio.Event()
fake_prisma = MagicMock()
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
async def _record_flush(client, accumulator):
gateway_flush_started.set()
calls.append("flush")
async def _audit_write_that_finishes_after_the_flush_starts() -> None:
await gateway_flush_started.wait()
calls.append("audit")
monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
fake_jwt = MagicMock()
fake_jwt.close = AsyncMock()
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
import litellm
monkeypatch.setattr(litellm, "cache", None, raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts()))
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
assert calls == ["flush", "audit", "disconnect"]
@pytest.mark.asyncio
async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
"""No prisma client means nothing to drain to, and no attempt is made."""
@ -245,6 +289,34 @@ async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch
assert flush.await_count == 0
@pytest.mark.asyncio
async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_the_rest_of_teardown(monkeypatch):
calls: list[str] = [] # mutable-ok: records call order, which is the assertion
monkeypatch.setattr(ps, "prisma_client", None, raising=False)
monkeypatch.setattr(
ps, "drain_post_call_redis_batches", AsyncMock(side_effect=lambda: calls.append("redis")), raising=False
)
fake_jwt = MagicMock()
fake_jwt.close = AsyncMock()
monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
import litellm
monkeypatch.setattr(litellm, "cache", None, raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
async def _audit_callback_still_sending() -> None:
await asyncio.sleep(0.05)
calls.append("audit")
track_audit_task(asyncio.create_task(_audit_callback_still_sending()))
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
assert calls == ["audit", "redis"]
@pytest.mark.asyncio
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
fake_prisma = MagicMock()

4
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -12612,6 +12612,7 @@ def test_dump_redacted_config_serializes_non_json_native_values():
async def test_update_config_general_settings_emits_audit_log(monkeypatch):
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import ConfigFieldUpdate
from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
from litellm.proxy.proxy_server import update_config_general_settings
existing = {"max_parallel_requests": 5, "some_api_key": "sk-stored-secret"}
@ -12635,6 +12636,7 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch):
)
# Audit is scheduled via asyncio.create_task; yield so it runs.
await asyncio.sleep(0)
await drain_audit_tasks()
fake.db.litellm_auditlog.create.assert_awaited_once()
written = fake.db.litellm_auditlog.create.call_args.kwargs["data"]
@ -13004,6 +13006,7 @@ async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch):
async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatch):
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import ConfigFieldDelete
from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
from litellm.proxy.proxy_server import delete_config_general_settings
existing = {"max_parallel_requests": 5}
@ -13023,6 +13026,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc
)
# Audit is scheduled via asyncio.create_task; yield so it runs.
await asyncio.sleep(0)
await drain_audit_tasks()
fake.db.litellm_auditlog.create.assert_awaited_once()
written = fake.db.litellm_auditlog.create.call_args.kwargs["data"]