fix(logging): preserve spend accounting when callbacks time out

litellm/integrations/custom_logger.py

@@ -66,6 +66,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile(
 class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
     # Class variables or attributes
     server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset()
+    is_internal_accounting_callback: ClassVar[bool] = False
 
     enforces_request_content: bool = False
     """

litellm/litellm_core_utils/litellm_logging.py

@@ -19,6 +19,7 @@ from functools import lru_cache
 from types import MappingProxyType, TracebackType
 from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
 
+from anyio import fail_after
 from httpx import Response, Timeout
 from pydantic import BaseModel, JsonValue, TypeAdapter
 
@@ -111,9 +112,11 @@ from litellm.litellm_core_utils.logging_utils import (
     truncate_base64_in_messages,
     truncate_base64_in_messages_async,
 )
+from litellm.litellm_core_utils.logging_worker import optional_callback_budget
 from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
 from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request
 from litellm.litellm_core_utils.redact_messages import (
+    perform_redaction,
     redact_message_input_output_from_custom_logger,
     redact_message_input_output_from_logging,
     redact_streaming_responses_for_custom_logger,
@@ -3526,6 +3529,68 @@ class Logging(LiteLLMLoggingBaseClass):
                 e,
             )
 
+    async def _run_async_success_hooks(
+        self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object
+    ) -> object:
+        for callback in callbacks:
+            try:
+                if isinstance(callback, CustomGuardrail):
+                    from litellm.types.guardrails import GuardrailEventHooks
+
+                    if (
+                        callback.should_run_guardrail(
+                            data=self.model_call_details,
+                            event_type=GuardrailEventHooks.logging_only,
+                        )
+                        is not True
+                    ):
+                        continue
+
+                    self.shadow_eval_request_snapshot = None
+                    self.model_call_details, result = await callback.async_logging_hook(
+                        kwargs=self.model_call_details,
+                        result=result,
+                        call_type=self.call_type,
+                    )
+                elif isinstance(callback, CustomLogger):
+                    result = cast(  # cast-ok: legacy redaction returns heterogeneous response objects
+                        object,
+                        redact_message_input_output_from_custom_logger(
+                            result=result, litellm_logging_obj=self, custom_logger=callback
+                        ),
+                    )
+                    self.model_call_details, result = await callback.async_logging_hook(
+                        kwargs=self.model_call_details,
+                        result=result,
+                        call_type=self.call_type,
+                    )
+            except Exception:  # noqa: BLE001  # one failing hook must not skip later callbacks (slot release)
+                verbose_logger.error(
+                    "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s",
+                    traceback.format_exc(),
+                )
+                self._handle_callback_failure(callback=callback)
+
+        return result
+
+    async def _prepare_accounting_success_hooks(
+        self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object
+    ) -> object:
+        from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
+
+        try:
+            with fail_after(optional_callback_budget(LOGGING_WORKER_MAX_TIME_PER_COROUTINE, fraction=0.5)):
+                return await self._run_async_success_hooks(callbacks, result)
+        except TimeoutError:
+            self.shadow_eval_request_snapshot = None
+            self.model_call_details = redact_streaming_responses_for_custom_logger(
+                self.model_call_details, CustomLogger(message_logging=False)
+            )
+            verbose_logger.warning("Logging hooks timed out; recording spend with request and response bodies redacted")
+            return cast(  # cast-ok: legacy redaction returns heterogeneous response objects
+                object, perform_redaction(self.model_call_details, result)
+            )
+
     async def async_success_handler(
         self,
         result: object = None,  # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
@@ -3735,7 +3800,7 @@ class Logging(LiteLLMLoggingBaseClass):
             self._get_completed_success_callbacks(asynchronous=True) if not self.stream else callbacks
         )
 
-        result = redact_message_input_output_from_logging(
+        redacted_result: Final = redact_message_input_output_from_logging(
             model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}),
             result=result,
         )
@@ -3743,45 +3808,30 @@ class Logging(LiteLLMLoggingBaseClass):
         ## LOGGING HOOK ##
 
         hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks)))
-        for callback in hook_callbacks:
-            try:
-                if isinstance(callback, CustomGuardrail):
-                    from litellm.types.guardrails import GuardrailEventHooks
-
-                    if (
-                        callback.should_run_guardrail(
-                            data=self.model_call_details,
-                            event_type=GuardrailEventHooks.logging_only,
-                        )
-                        is not True
-                    ):
-                        continue
-
-                    self.shadow_eval_request_snapshot = None
-                    self.model_call_details, result = await callback.async_logging_hook(
-                        kwargs=self.model_call_details,
-                        result=result,
-                        call_type=self.call_type,
-                    )
-                elif isinstance(callback, CustomLogger):
-                    result = redact_message_input_output_from_custom_logger(
-                        result=result, litellm_logging_obj=self, custom_logger=callback
-                    )
-                    self.model_call_details, result = await callback.async_logging_hook(
-                        kwargs=self.model_call_details,
-                        result=result,
-                        call_type=self.call_type,
-                    )
-            except Exception:  # noqa: BLE001  # one failing hook must not skip later callbacks (slot release)
-                verbose_logger.error(
-                    "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s",
-                    traceback.format_exc(),
-                )
-                self._handle_callback_failure(callback=callback)
+        accounting_callbacks: Final = tuple(
+            {
+                type(callback): callback
+                for callback in delivery_callbacks
+                if isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback
+            }.values()
+        )
+        result = (
+            await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result)
+            if accounting_callbacks
+            else await self._run_async_success_hooks(hook_callbacks, redacted_result)
+        )
+        ordered_callbacks: Final = (
+            *accounting_callbacks,
+            *(
+                callback
+                for callback in delivery_callbacks
+                if not (isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback)
+            ),
+        )
 
         self.mark_logging_complete(event_type="async_success")
 
-        for callback in delivery_callbacks:
+        for callback in ordered_callbacks:
             # check if callback can run for this request
             litellm_params = self.model_call_details.get("litellm_params", {})
             should_run = self.should_run_callback(

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -3,7 +3,7 @@ import traceback
 from collections.abc import Callable, Mapping, Sequence
 from datetime import datetime
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Protocol, cast
+from typing import TYPE_CHECKING, ClassVar, Final, Protocol, cast
 
 from pydantic import TypeAdapter, ValidationError
 
@@ -154,6 +154,8 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
 
 
 class _ProxyDBLogger(CustomLogger):
+    is_internal_accounting_callback: ClassVar[bool] = True
+
     def __init__(
         self,
         spend_event_producer: SpendEventProducer | None = None,

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -1,5 +1,6 @@
 import asyncio
 import contextlib
+import contextvars
 import copy
 import datetime
 import importlib.abc
@@ -48,7 +49,7 @@ from litellm.litellm_core_utils.litellm_logging import (
     get_standard_logging_object_payload,
     set_callbacks,
 )
-from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
 from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck
@@ -2125,7 +2126,7 @@ async def test_aresponses_websocket_turn_reaches_a_custom_logger_through_the_asy
 async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(
     monkeypatch: pytest.MonkeyPatch,
 ):
-    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
     from litellm.responses.main import base_llm_http_handler
 
     success_events = []
@@ -12958,3 +12959,93 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a
     assert s3_writer.log_event.call_count == expected_s3_writes
     if expected_s3_writes:
         assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("accounting_first", [False, True])
+@pytest.mark.parametrize("duplicate_accounting", [False, True])
+@pytest.mark.parametrize("blocked_phase", ["hook", "delivery", "neither", "redacted"])
+async def test_slow_optional_logging_cannot_cancel_accounting(
+    monkeypatch: pytest.MonkeyPatch,
+    stream: bool,
+    accounting_first: bool,
+    duplicate_accounting: bool,
+    blocked_phase: str,
+) -> None:
+    from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
+    from litellm.proxy.proxy_server import general_settings
+    from litellm.proxy.spend_tracking.spend_event import SpendEvent, decode_spend_event
+    from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress
+
+    recorded: Final[asyncio.Queue[tuple[bytes, str]]] = asyncio.Queue()
+    hook_context: Final = contextvars.ContextVar("accounting_hook_context", default="outside")
+    delivered: Final = asyncio.Event()
+
+    async def store_event(line: bytes) -> None:
+        recorded.put_nowait((line, hook_context.get()))
+
+    producer: Final = SpendEventProducer(
+        address=UnixAddress("/unused"),
+        on_unavailable="fallback",
+        buffer_size=1,
+        connect_timeout=1,
+        fallback=store_event,
+    )
+    await producer.close(drain_timeout=1)
+    monkeypatch.setitem(general_settings, "store_prompts_in_spend_logs", True)
+    cancelled: Final = asyncio.Event()
+    never: Final = asyncio.Event()
+
+    class OptionalLogger(CustomLogger):
+        async def async_logging_hook(
+            self, kwargs: dict[str, object], result: object, call_type: str
+        ) -> tuple[dict[str, object], object]:
+            hook_context.set("inside")
+            if blocked_phase == "hook":
+                try:
+                    await never.wait()
+                finally:
+                    cancelled.set()
+            return kwargs, result
+
+        async def async_log_success_event(
+            self,
+            kwargs: dict[str, object],
+            response_obj: object,
+            start_time: datetime.datetime,
+            end_time: datetime.datetime,
+        ) -> None:
+            delivered.set()
+            if blocked_phase == "delivery":
+                try:
+                    await never.wait()
+                finally:
+                    cancelled.set()
+
+    optional: Final = OptionalLogger(message_logging=blocked_phase != "redacted")
+    accounting: Final = _ProxyDBLogger(producer)
+    accounting_callbacks: Final = [accounting, _ProxyDBLogger(producer)] if duplicate_accounting else [accounting]
+    callbacks: Final = [*accounting_callbacks, optional] if accounting_first else [optional, *accounting_callbacks]
+    monkeypatch.setattr(litellm, "_async_success_callback", [])
+    logging_obj: Final = _async_request_logging_obj("accounting-timeout", callbacks, stream=stream)
+    response: Final = ModelResponse(
+        choices=[{"message": {"role": "assistant", "content": "private-marker"}}],
+        usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+    )
+    worker: Final = LoggingWorker(timeout=0.1)
+    worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=response))
+    try:
+        await worker.flush()
+        line, context = recorded.get_nowait()
+        assert context == "inside"
+        assert delivered.is_set()
+        event: Final = decode_spend_event(line)
+        assert isinstance(event, SpendEvent)
+        assert event.response is not None
+        assert Usage.model_validate(event.response["usage"]).total_tokens == 15
+        assert recorded.empty()
+        assert cancelled.is_set() is (blocked_phase in ("hook", "delivery"))
+        assert ("private-marker" in line.decode()) is (blocked_phase not in ("hook", "redacted"))
+        assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase in ("hook", "redacted"))
+        assert response.choices[0].message.content == "private-marker"
+    finally:
+        await worker.stop()