From 84eeebd509eccbe495812b30cad12f57e4d20481 Mon Sep 17 00:00:00 2001 From: Tufail Akram <72789060+akram1089@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:11:34 +0530 Subject: [PATCH] fix(proxy): track and drain management audit log writes on shutdown (#43594) * 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> --- .../proxy/hooks/key_management_event_hooks.py | 78 +-- .../hooks/user_management_event_hooks.py | 43 +- .../key_management_endpoints.py | 20 +- .../proxy/management_helpers/audit_logs.py | 83 ++- litellm/proxy/proxy_server.py | 19 +- .../test_key_management_endpoints.py | 37 ++ .../test_audit_logs_proxy.py | 496 +++++++++++++++++- .../unit/proxy/proxy_server/test_lifecycle.py | 72 +++ ...test_proxy_server_endpoints_and_startup.py | 4 + 9 files changed, 765 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ef93809d47b..ef60a800cdd 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -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, diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 6d978929c05..6618af70c0a 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -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) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7877d887409..592c732bfd3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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, diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ecd6abea3c3..ef95e8e2c0f 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -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) + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97d8eda1fad..bbd9d96b84b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 1af9def3afb..57e0a4fc47c 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -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(): """ diff --git a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py index 1ea608cfc17..325275ab391 100644 --- a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py +++ b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py @@ -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")] diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 69f9c72c9c8..71daa3bb873 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -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() diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 0ecebc7560c..e824c5017c6 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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"]