From c1b36403104224a888b939429921707aced0154c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 9 Mar 2026 09:48:10 -0700 Subject: [PATCH 1/4] feat: add audit log export to external callbacks (S3, Datadog, etc.) Audit logs (CRUD events on keys, teams, users, models) were only stored in the Prisma DB. This adds a pluggable callback system so audit logs can be forwarded to external services like S3 for ingestion into security monitoring tools. New config key `audit_log_callbacks` under `litellm_settings` reuses the existing callback infrastructure. Any CustomLogger subclass can opt in by overriding `async_log_audit_log_event()`. S3Logger (s3_v2) is implemented as the first handler, storing audit logs under `audit_logs/{date}/` prefix. Co-Authored-By: Claude Opus 4.6 --- litellm/__init__.py | 1 + litellm/integrations/custom_logger.py | 7 + litellm/integrations/s3_v2.py | 34 ++- .../proxy/management_helpers/audit_logs.py | 84 +++++- litellm/proxy/proxy_server.py | 13 + litellm/types/utils.py | 14 + .../test_audit_log_callbacks.py | 270 ++++++++++++++++++ 7 files changed, 420 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 4fc71e1270..e66a9b6b08 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -96,6 +96,7 @@ input_callback: List[CALLBACK_TYPES] = [] success_callback: List[CALLBACK_TYPES] = [] failure_callback: List[CALLBACK_TYPES] = [] service_callback: List[CALLBACK_TYPES] = [] +audit_log_callbacks: List[CALLBACK_TYPES] = [] # logging_callback_manager is lazy-loaded via __getattr__ _custom_logger_compatible_callbacks_literal = Literal[ "lago", diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index c244363e38..8867adfe86 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -27,6 +27,7 @@ from litellm.types.utils import ( LLMResponseTypes, ModelResponse, ModelResponseStream, + StandardAuditLogPayload, StandardCallbackDynamicParams, StandardLoggingPayload, ) @@ -177,6 +178,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): pass + async def async_log_audit_log_event( + self, audit_log: "StandardAuditLogPayload" + ): + """Called when an audit log is created. Override in subclasses to handle.""" + pass + #### PROMPT MANAGEMENT HOOKS #### async def async_get_chat_completion_prompt( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index eddc80dbc1..134766aa18 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -22,7 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.integrations.s3_v2 import s3BatchLoggingElement -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger @@ -244,6 +244,38 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) pass + async def async_log_audit_log_event( + self, audit_log: StandardAuditLogPayload + ) -> None: + """Batch audit logs and upload to S3 under audit_logs/ prefix.""" + try: + from datetime import timezone + + now = datetime.now(timezone.utc) + audit_log_id = audit_log.get("id", "unknown") + + s3_path = cast(Optional[str], self.s3_path) or "" + s3_path = s3_path.rstrip("/") + "/" if s3_path else "" + + s3_object_key = ( + f"{s3_path}audit_logs/" + f"{now.strftime('%Y-%m-%d')}/" + f"{now.strftime('%H-%M-%S')}_{audit_log_id}.json" + ) + + element = s3BatchLoggingElement( + payload=dict(audit_log), + s3_object_key=s3_object_key, + s3_object_download_filename=f"audit-{audit_log_id}.json", + ) + + self.log_queue.append(element) + + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() + except Exception as e: + verbose_logger.exception("S3 audit log error: %s", e) + async def _async_log_event_base(self, kwargs, response_obj, start_time, end_time): try: verbose_logger.debug( diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ea082f468a..dee7c82ee5 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -2,12 +2,15 @@ Functions to create audit logs for LiteLLM Proxy """ +import asyncio import json -from litellm._uuid import uuid from datetime import datetime, timezone +from typing import Dict import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( AUDIT_ACTIONS, LiteLLM_AuditLogs, @@ -15,6 +18,82 @@ from litellm.proxy._types import ( Optional, UserAPIKeyAuth, ) +from litellm.types.utils import StandardAuditLogPayload + +_audit_log_callback_cache: Dict[str, CustomLogger] = {} + + +def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]: + """Resolve a string callback name to a CustomLogger instance, with caching.""" + if name in _audit_log_callback_cache: + return _audit_log_callback_cache[name] + + from litellm.litellm_core_utils.litellm_logging import ( + _init_custom_logger_compatible_class, + ) + + instance = _init_custom_logger_compatible_class( + logging_integration=name, # type: ignore + internal_usage_cache=None, + llm_router=None, + ) + + if instance is not None: + _audit_log_callback_cache[name] = instance + return instance + + +def _build_audit_log_payload( + request_data: LiteLLM_AuditLogs, +) -> StandardAuditLogPayload: + """Convert LiteLLM_AuditLogs to StandardAuditLogPayload for callback dispatch.""" + updated_at = "" + if request_data.updated_at is not None: + updated_at = request_data.updated_at.isoformat() + + table_name = request_data.table_name + if isinstance(table_name, LitellmTableNames): + table_name = table_name.value + + return StandardAuditLogPayload( + id=request_data.id, + updated_at=updated_at, + changed_by=request_data.changed_by or "", + changed_by_api_key=request_data.changed_by_api_key or "", + action=request_data.action, + table_name=str(table_name), + object_id=request_data.object_id, + before_value=request_data.before_value, + updated_values=request_data.updated_values, + ) + + +async def _dispatch_audit_log_to_callbacks( + request_data: LiteLLM_AuditLogs, +) -> None: + """Dispatch audit log to all registered audit_log_callbacks.""" + if not litellm.audit_log_callbacks: + return + + payload = _build_audit_log_payload(request_data) + + for callback in litellm.audit_log_callbacks: + try: + resolved = callback + if isinstance(callback, str): + resolved = _resolve_audit_log_callback(callback) + if resolved is None: + verbose_proxy_logger.warning( + "Could not resolve audit log callback: %s", callback + ) + continue + + if isinstance(resolved, CustomLogger): + asyncio.create_task(resolved.async_log_audit_log_event(payload)) + except Exception as e: + verbose_proxy_logger.error( + "Failed dispatching audit log to callback: %s", e + ) async def create_object_audit_log( @@ -104,4 +183,5 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): # [Non-Blocking Exception. Do not allow blocking LLM API call] verbose_proxy_logger.error(f"Failed Creating audit log {e}") - return + # Dispatch to external audit log callbacks + await _dispatch_audit_log_to_callbacks(request_data) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bd7b21c3b5..aa5adc054d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2946,6 +2946,19 @@ class ProxyConfig: print( # noqa f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}" ) # noqa + elif key == "audit_log_callbacks": + litellm.audit_log_callbacks = [] + + for callback in value: + if "." in callback: + litellm.audit_log_callbacks.append( + get_instance_fn(value=callback) + ) + else: + litellm.audit_log_callbacks.append(callback) + print( # noqa + f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" + ) # noqa elif key == "cache_params": # this is set in the cache branch # see usage here: https://docs.litellm.ai/docs/proxy/caching diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8ae0cf2892..dab40ffc6e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2761,6 +2761,20 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False): """ +class StandardAuditLogPayload(TypedDict): + """Payload for audit log events dispatched to external callbacks.""" + + id: str + updated_at: str # ISO-8601 + changed_by: str + changed_by_api_key: str + action: str # "created" | "updated" | "deleted" | "blocked" | "rotated" + table_name: str + object_id: str + before_value: Optional[str] + updated_values: Optional[str] + + class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) diff --git a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py new file mode 100644 index 0000000000..cfdf6ac033 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py @@ -0,0 +1,270 @@ +""" +Tests for audit log callback dispatch. + +Tests the flow: create_audit_log_for_update -> _dispatch_audit_log_to_callbacks -> CustomLogger.async_log_audit_log_event +""" + +import asyncio +import json +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames +from litellm.proxy.management_helpers.audit_logs import ( + _build_audit_log_payload, + _dispatch_audit_log_to_callbacks, + create_audit_log_for_update, +) +from litellm.types.utils import StandardAuditLogPayload + + +@pytest.fixture(autouse=True) +def reset_audit_log_callbacks(): + """Reset audit_log_callbacks before and after each test.""" + original = litellm.audit_log_callbacks + litellm.audit_log_callbacks = [] + yield + litellm.audit_log_callbacks = original + + +def _make_audit_log( + action: str = "created", + table_name: LitellmTableNames = LitellmTableNames.TEAM_TABLE_NAME, +) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id="test-audit-id", + updated_at=datetime(2026, 3, 9, 12, 0, 0, tzinfo=timezone.utc), + changed_by="user-123", + changed_by_api_key="sk-abc", + action=action, + table_name=table_name, + object_id="team-456", + updated_values=json.dumps({"name": "new-team"}), + before_value=json.dumps({"name": "old-team"}), + ) + + +class TestBuildAuditLogPayload: + def test_builds_correct_payload(self): + audit_log = _make_audit_log() + payload = _build_audit_log_payload(audit_log) + + assert payload["id"] == "test-audit-id" + assert payload["updated_at"] == "2026-03-09T12:00:00+00:00" + assert payload["changed_by"] == "user-123" + assert payload["changed_by_api_key"] == "sk-abc" + assert payload["action"] == "created" + assert payload["table_name"] == "LiteLLM_TeamTable" + assert payload["object_id"] == "team-456" + assert payload["updated_values"] == json.dumps({"name": "new-team"}) + assert payload["before_value"] == json.dumps({"name": "old-team"}) + + def test_handles_none_values(self): + audit_log = LiteLLM_AuditLogs( + id="test-id", + updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + changed_by=None, + changed_by_api_key=None, + action="deleted", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="key-789", + updated_values=None, + before_value=None, + ) + payload = _build_audit_log_payload(audit_log) + + assert payload["changed_by"] == "" + assert payload["changed_by_api_key"] == "" + assert payload["before_value"] is None + assert payload["updated_values"] is None + + +class TestDispatchAuditLogToCallbacks: + @pytest.mark.asyncio + async def test_dispatches_to_custom_logger_instance(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + audit_log = _make_audit_log() + await _dispatch_audit_log_to_callbacks(audit_log) + + # Let asyncio.create_task run + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_called_once() + payload = mock_logger.async_log_audit_log_event.call_args[0][0] + assert payload["id"] == "test-audit-id" + assert payload["action"] == "created" + + @pytest.mark.asyncio + async def test_no_dispatch_when_callbacks_empty(self): + litellm.audit_log_callbacks = [] + audit_log = _make_audit_log() + # Should return immediately without error + await _dispatch_audit_log_to_callbacks(audit_log) + + @pytest.mark.asyncio + async def test_resolves_string_callback(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + + litellm.audit_log_callbacks = ["s3_v2"] + + with patch( + "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", + return_value=mock_logger, + ): + audit_log = _make_audit_log() + await _dispatch_audit_log_to_callbacks(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_nonblocking_on_callback_failure(self): + """Callback errors should not propagate.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock( + side_effect=RuntimeError("boom") + ) + litellm.audit_log_callbacks = [mock_logger] + + audit_log = _make_audit_log() + # Should not raise + await _dispatch_audit_log_to_callbacks(audit_log) + await asyncio.sleep(0.1) + + @pytest.mark.asyncio + async def test_skips_unresolvable_string_callback(self): + litellm.audit_log_callbacks = ["nonexistent_callback"] + + with patch( + "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", + return_value=None, + ): + audit_log = _make_audit_log() + # Should not raise + await _dispatch_audit_log_to_callbacks(audit_log) + + +class TestCreateAuditLogForUpdateWithCallbacks: + @pytest.mark.asyncio + async def test_dispatches_to_callbacks_after_db_write(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_auditlog.create = AsyncMock() + + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # DB write should happen + mock_prisma.db.litellm_auditlog.create.assert_called_once() + # Callback should also be called + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_no_dispatch_when_not_premium(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", False), patch( + "litellm.store_audit_logs", True + ): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_not_called() + + @pytest.mark.asyncio + async def test_no_dispatch_when_store_audit_logs_false(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.store_audit_logs", False): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_not_called() + + +class TestS3LoggerAuditLogEvent: + @pytest.mark.asyncio + async def test_queues_audit_log_with_correct_s3_key(self): + with patch( + "litellm.integrations.s3_v2.S3Logger.__init__", return_value=None + ): + from litellm.integrations.s3_v2 import S3Logger + + logger = S3Logger() + logger.s3_path = "my-prefix" + logger.log_queue = [] + logger.batch_size = 100 + + audit_log = StandardAuditLogPayload( + id="audit-123", + updated_at="2026-03-09T12:00:00+00:00", + changed_by="user-1", + changed_by_api_key="sk-abc", + action="created", + table_name="LiteLLM_TeamTable", + object_id="team-1", + before_value=None, + updated_values='{"name": "new"}', + ) + + await logger.async_log_audit_log_event(audit_log) + + assert len(logger.log_queue) == 1 + element = logger.log_queue[0] + assert element.s3_object_key.startswith("my-prefix/audit_logs/") + assert "audit-123" in element.s3_object_key + assert element.s3_object_key.endswith(".json") + assert element.s3_object_download_filename == "audit-audit-123.json" + assert element.payload["id"] == "audit-123" + assert element.payload["action"] == "created" + + @pytest.mark.asyncio + async def test_s3_key_format_no_path(self): + with patch( + "litellm.integrations.s3_v2.S3Logger.__init__", return_value=None + ): + from litellm.integrations.s3_v2 import S3Logger + + logger = S3Logger() + logger.s3_path = None + logger.log_queue = [] + logger.batch_size = 100 + + audit_log = StandardAuditLogPayload( + id="audit-456", + updated_at="2026-03-09T12:00:00+00:00", + changed_by="user-1", + changed_by_api_key="sk-abc", + action="deleted", + table_name="LiteLLM_VerificationToken", + object_id="key-1", + before_value=None, + updated_values=None, + ) + + await logger.async_log_audit_log_event(audit_log) + + assert len(logger.log_queue) == 1 + element = logger.log_queue[0] + assert element.s3_object_key.startswith("audit_logs/") + assert "audit-456" in element.s3_object_key From 364f5c6e03dfe72c2223482d731c81ffd5858b56 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Mar 2026 11:27:48 -0700 Subject: [PATCH 2/4] fix: ensure audit log callbacks fire even when DB is unavailable Two fixes based on PR feedback: 1. Move callback dispatch before the prisma_client check so audit logs still reach S3/Datadog even if the DB is down. Also changed the prisma_client=None case from raising an exception to logging an error and returning gracefully. 2. Attach a done_callback to asyncio tasks created for audit log callbacks so exceptions are logged through verbose_proxy_logger instead of silently swallowed. Co-Authored-By: Claude Opus 4.6 --- .../proxy/management_helpers/audit_logs.py | 32 ++++++-- .../test_audit_log_callbacks.py | 75 +++++++++++++++++++ 2 files changed, 100 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index dee7c82ee5..a26a38bc4c 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -68,6 +68,18 @@ def _build_audit_log_payload( ) +def _audit_log_task_done_callback(task: asyncio.Task) -> None: + """Log exceptions from audit log callback tasks so they don't slip through silently.""" + try: + exc = task.exception() + except asyncio.CancelledError: + return + if exc is not None: + verbose_proxy_logger.error( + "Audit log callback task failed: %s", exc, exc_info=exc + ) + + async def _dispatch_audit_log_to_callbacks( request_data: LiteLLM_AuditLogs, ) -> None: @@ -89,7 +101,10 @@ async def _dispatch_audit_log_to_callbacks( continue if isinstance(resolved, CustomLogger): - asyncio.create_task(resolved.async_log_audit_log_event(payload)) + task = asyncio.create_task( + resolved.async_log_audit_log_event(payload) + ) + task.add_done_callback(_audit_log_task_done_callback) except Exception as e: verbose_proxy_logger.error( "Failed dispatching audit log to callback: %s", e @@ -160,9 +175,6 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): if premium_user is not True: return - if prisma_client is None: - raise Exception("prisma_client is None, no DB connected") - verbose_proxy_logger.debug("creating audit log for %s", request_data) if isinstance(request_data.updated_values, dict): @@ -171,6 +183,15 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): if isinstance(request_data.before_value, dict): request_data.before_value = json.dumps(request_data.before_value) + # Dispatch to external audit log callbacks regardless of DB availability + await _dispatch_audit_log_to_callbacks(request_data) + + if prisma_client is None: + verbose_proxy_logger.error( + "prisma_client is None, cannot write audit log to DB" + ) + return + _request_data = request_data.model_dump(exclude_none=True) try: @@ -182,6 +203,3 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): except Exception as e: # [Non-Blocking Exception. Do not allow blocking LLM API call] verbose_proxy_logger.error(f"Failed Creating audit log {e}") - - # Dispatch to external audit log callbacks - await _dispatch_audit_log_to_callbacks(request_data) diff --git a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py index cfdf6ac033..85eda4368c 100644 --- a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py +++ b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py @@ -15,6 +15,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames from litellm.proxy.management_helpers.audit_logs import ( + _audit_log_task_done_callback, _build_audit_log_payload, _dispatch_audit_log_to_callbacks, create_audit_log_for_update, @@ -201,6 +202,80 @@ class TestCreateAuditLogForUpdateWithCallbacks: mock_logger.async_log_audit_log_event.assert_not_called() + @pytest.mark.asyncio + async def test_dispatches_even_when_prisma_client_is_none(self): + """Callbacks should fire even if DB is unavailable.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client", None): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # Callback should still be called despite no DB + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_dispatches_even_when_db_write_fails(self): + """Callbacks should fire even if the DB write raises.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_auditlog.create = AsyncMock( + side_effect=RuntimeError("DB connection lost") + ) + + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # Callback should still be called despite DB failure + mock_logger.async_log_audit_log_event.assert_called_once() + + +class TestAuditLogTaskDoneCallback: + def test_logs_exception_from_failed_task(self): + """Done callback should log task exceptions.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.return_value = RuntimeError("callback failed") + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_called_once() + assert "callback failed" in str(mock_logger.error.call_args) + + def test_no_log_on_success(self): + """Done callback should not log when task succeeds.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.return_value = None + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_not_called() + + def test_handles_cancelled_task(self): + """Done callback should handle cancelled tasks gracefully.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.side_effect = asyncio.CancelledError() + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_not_called() + class TestS3LoggerAuditLogEvent: @pytest.mark.asyncio From 9c47d50f7496ede7bd75d5eaa86100e04b3c47b2 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 20 Mar 2026 16:50:35 -0700 Subject: [PATCH 3/4] add startup messages --- litellm/proxy/proxy_server.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ade952b3cc..61fb12a855 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2971,9 +2971,19 @@ class ProxyConfig: ) else: litellm.audit_log_callbacks.append(callback) - print( # noqa - f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" - ) # noqa + + _store_audit_logs = litellm_settings.get( + "store_audit_logs", litellm.store_audit_logs + ) + if _store_audit_logs: + print( # noqa + f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" + ) # noqa + else: + verbose_proxy_logger.warning( + "'audit_log_callbacks' is configured but 'store_audit_logs' is not enabled. " + "Audit log callbacks will not fire until 'store_audit_logs: true' is added to litellm_settings." + ) elif key == "cache_params": # this is set in the cache branch # see usage here: https://docs.litellm.ai/docs/proxy/caching From ae21527aeaeeaf74ccc7dcdbaf6883ff5bf9878c Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 20 Mar 2026 17:35:05 -0700 Subject: [PATCH 4/4] fix black formatting Made-with: Cursor --- litellm/integrations/custom_logger.py | 4 +--- litellm/proxy/management_helpers/audit_logs.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0b198ed8d1..cccabf53e5 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -178,9 +178,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): pass - async def async_log_audit_log_event( - self, audit_log: "StandardAuditLogPayload" - ): + async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): """Called when an audit log is created. Override in subclasses to handle.""" pass diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index a26a38bc4c..08d92c3e6a 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -101,9 +101,7 @@ async def _dispatch_audit_log_to_callbacks( continue if isinstance(resolved, CustomLogger): - task = asyncio.create_task( - resolved.async_log_audit_log_event(payload) - ) + task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) task.add_done_callback(_audit_log_task_done_callback) except Exception as e: verbose_proxy_logger.error(