From 2182e8590ce9857f3caecdfb4697b62a251b6ce0 Mon Sep 17 00:00:00 2001 From: adc77 Date: Tue, 20 Oct 2026 01:02:16 +0530 Subject: [PATCH 1/3] fix(logging): preserve spend accounting when callbacks time out --- litellm/integrations/custom_logger.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 124 ++++++++++++------ .../proxy/hooks/proxy_track_cost_callback.py | 4 +- .../test_litellm_logging.py | 95 +++++++++++++- 4 files changed, 184 insertions(+), 40 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..ca7877134db 100644 --- a/litellm/integrations/custom_logger.py +++ b/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 """ diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0ac216c2029..3292da43247 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3195394b9aa..967ce1ea328 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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, diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d1258e43154..4e057b7ea61 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/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() From dce9f01e7d56719d3279802f444b8a01aef4bd66 Mon Sep 17 00:00:00 2001 From: adc77 Date: Tue, 20 Oct 2026 01:27:25 +0530 Subject: [PATCH 2/3] fix(logging): drop bodies after masking hook timeouts --- litellm/litellm_core_utils/litellm_logging.py | 40 ++++-- litellm/proxy/spend_tracking/spend_event.py | 16 ++- .../test_litellm_logging.py | 135 ++++++++++++++++-- 3 files changed, 169 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3292da43247..a9ae6da58da 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -116,7 +116,6 @@ 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, @@ -3574,7 +3573,11 @@ class Logging(LiteLLMLoggingBaseClass): return result async def _prepare_accounting_success_hooks( - self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object + self, + callbacks: Sequence[str | Callable[..., object] | CustomLogger], + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, ) -> object: from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE @@ -3582,14 +3585,30 @@ class Logging(LiteLLMLoggingBaseClass): 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: + from litellm.proxy.spend_tracking.spend_event import ( + SpendEventBuildError, + build_spend_event_payload, + spend_event_callback_args, + ) + self.shadow_eval_request_snapshot = None - self.model_call_details = redact_streaming_responses_for_custom_logger( - self.model_call_details, CustomLogger(message_logging=False) + assembled_stream: Final = "async_complete_streaming_response" in self.model_call_details + accounting_response: Final[object] = ( + self.model_call_details["async_complete_streaming_response"] if assembled_stream else result ) - 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) + event: Final = build_spend_event_payload( + self.model_call_details, accounting_response, start_time, end_time, store_bodies=False ) + if isinstance(event, SpendEventBuildError): + verbose_logger.error("Cannot prepare body-free spend accounting after a logging hook timeout") + return event + args: Final = spend_event_callback_args(event) + self.model_call_details = { + **args.kwargs, + **({"async_complete_streaming_response": args.response_obj} if assembled_stream else {}), + } + verbose_logger.warning("Logging hooks timed out, recording spend without request or response bodies") + return args.response_obj async def async_success_handler( self, @@ -3816,10 +3835,15 @@ class Logging(LiteLLMLoggingBaseClass): }.values() ) result = ( - await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result) + await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result, start_time, end_time) if accounting_callbacks else await self._run_async_success_hooks(hook_callbacks, redacted_result) ) + if accounting_callbacks: + from litellm.proxy.spend_tracking.spend_event import SpendEventBuildError + + if isinstance(result, SpendEventBuildError): + return ordered_callbacks: Final = ( *accounting_callbacks, *( diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 031e84896b2..e28954e4f1d 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -342,10 +342,10 @@ def _function_tool(name: str) -> _FunctionTool: return tool -def build_spend_event( +def build_spend_event_payload( raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool -) -> bytes | SpendEventBuildError: - """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" +) -> SpendEvent | SpendEventBuildError: + """Validate and project the identity, usage, cost, and optional bodies of a completed call.""" try: kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs) except ValidationError as error: @@ -379,6 +379,16 @@ def build_spend_event( realtime_tool_calls=kwargs.get("realtime_tool_calls"), request_tool_names=_request_tool_names(kwargs), ) + return event + + +def build_spend_event( + raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool +) -> bytes | SpendEventBuildError: + """Validate the success callback's kwargs and serialize the event once, as a single JSON line.""" + event: Final = build_spend_event_payload(raw_kwargs, response_obj, start_time, end_time, store_bodies) + if isinstance(event, SpendEventBuildError): + return event return event.model_dump_json(fallback=_json_fallback).encode() + b"\n" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 4e057b7ea61..25d5394c426 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8009,12 +8009,14 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() -def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: +def _async_request_logging_obj( + call_id: str, callbacks: list[CustomLogger], stream: bool = False, call_type: str = "acompletion" +) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=stream, - call_type="acompletion", + call_type=call_type, start_time=datetime.datetime.now(), litellm_call_id=call_id, function_id=call_id, @@ -12959,14 +12961,18 @@ 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("response_kind", ["chat", "responses"]) @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, + response_kind: str, accounting_first: bool, duplicate_accounting: bool, blocked_phase: str, @@ -12979,6 +12985,7 @@ async def test_slow_optional_logging_cannot_cancel_accounting( recorded: Final[asyncio.Queue[tuple[bytes, str]]] = asyncio.Queue() hook_context: Final = contextvars.ContextVar("accounting_hook_context", default="outside") delivered: Final = asyncio.Event() + deliveries: Final[asyncio.Queue[str]] = asyncio.Queue() async def store_event(line: bytes) -> None: recorded.put_nowait((line, hook_context.get())) @@ -13015,6 +13022,7 @@ async def test_slow_optional_logging_cannot_cancel_accounting( end_time: datetime.datetime, ) -> None: delivered.set() + deliveries.put_nowait(json.dumps(kwargs, default=str) + str(response_obj)) if blocked_phase == "delivery": try: await never.wait() @@ -13026,13 +13034,60 @@ async def test_slow_optional_logging_cannot_cancel_accounting( 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), + logging_obj: Final = _async_request_logging_obj( + "accounting-timeout", + callbacks, + stream=stream, + call_type="aresponses" if response_kind == "responses" else "acompletion", ) - worker: Final = LoggingWorker(timeout=0.1) - worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=response)) + response: Final = ( + ModelResponse( + choices=[{"message": {"role": "assistant", "content": "private-marker"}}], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + if response_kind == "chat" + else ResponsesAPIResponse( + id="resp-private-instructions", + created_at=1, + model="gpt-4o-mini", + status="completed", + instructions="private-instructions-canary" if blocked_phase != "redacted" else None, + output=[ + { + "type": "message", + "id": "msg-private", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "private-marker", "annotations": []}], + } + ], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), + ) + ) + original_response: Final = response.model_dump(exclude={"usage"}) + logging_obj.update_environment_variables( + litellm_params={ + "acompletion": True, + "metadata": {}, + "proxy_server_request": { + "body": { + "instructions": "private-request-instructions", + "system": "private-system-canary", + "tools": [{"description": "private-tool-canary"}], + } + }, + } + if blocked_phase == "hook" + else {"acompletion": True, "metadata": {}}, + optional_params={}, + ) + logging_result: Final = ( + ResponseCompletedEvent(type="response.completed", sequence_number=1, response=response) + if stream and isinstance(response, ResponsesAPIResponse) + else response + ) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=logging_result)) try: await worker.flush() line, context = recorded.get_nowait() @@ -13041,11 +13096,69 @@ async def test_slow_optional_logging_cannot_cancel_accounting( 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 isinstance(event.response["usage"], dict) + assert 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" + assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase == "redacted") + assert ("private-instructions-canary" in line.decode()) is ( + response_kind == "responses" and blocked_phase in ("neither", "delivery") + ) + delivery: Final = deliveries.get_nowait() + assert not any( + value in delivery + for value in ( + "private-request-instructions", + "private-system-canary", + "private-tool-canary", + ) + ) + assert ("private-instructions-canary" in delivery) is ( + response_kind == "responses" and blocked_phase in ("neither", "delivery") + ) + assert response.model_dump(exclude={"usage"}) == original_response + finally: + await worker.stop() + + +@pytest.mark.asyncio +async def test_invalid_accounting_after_a_hook_timeout_never_exports_unmasked_data( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + + exported: Final = asyncio.Event() + never: Final = asyncio.Event() + + class InvalidMaskingHook(CustomLogger): + async def async_logging_hook( + self, + kwargs: dict[str, object], + result: object, + call_type: str, + ) -> tuple[dict[str, object], object]: + kwargs["stream"] = "invalid" + await never.wait() + 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: + exported.set() + + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = _async_request_logging_obj("invalid-accounting", [InvalidMaskingHook(), _ProxyDBLogger()]) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse())) + try: + await worker.flush() + assert not exported.is_set() + assert "Cannot prepare body-free spend accounting after a logging hook timeout" in caplog.text finally: await worker.stop() From 2ed25a2f6575327916ed996246f9611948d4d855 Mon Sep 17 00:00:00 2001 From: adc77 Date: Tue, 20 Oct 2026 13:24:40 +0530 Subject: [PATCH 3/3] fix(logging): preserve privacy controls after hook timeouts --- litellm/litellm_core_utils/litellm_logging.py | 22 +++++ .../test_litellm_logging.py | 91 +++++++++++++++++++ 2 files changed, 113 insertions(+) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a9ae6da58da..5a42ee6c049 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -116,6 +116,7 @@ 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, @@ -3585,6 +3586,7 @@ class Logging(LiteLLMLoggingBaseClass): 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: + from litellm.constants import X_LITELLM_DISABLE_CALLBACKS from litellm.proxy.spend_tracking.spend_event import ( SpendEventBuildError, build_spend_event_payload, @@ -3592,6 +3594,10 @@ class Logging(LiteLLMLoggingBaseClass): ) self.shadow_eval_request_snapshot = None + self.model_call_details = redact_streaming_responses_for_custom_logger( + self.model_call_details, CustomLogger(message_logging=False) + ) + perform_redaction(self.model_call_details, None) assembled_stream: Final = "async_complete_streaming_response" in self.model_call_details accounting_response: Final[object] = ( self.model_call_details["async_complete_streaming_response"] if assembled_stream else result @@ -3603,8 +3609,24 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.error("Cannot prepare body-free spend accounting after a logging hook timeout") return event args: Final = spend_event_callback_args(event) + mapping_adapter: Final = TypeAdapter(Mapping[str, object]) + empty_headers: Final = MappingProxyType[str, object]({}) + params: Final = mapping_adapter.validate_python(self.model_call_details["litellm_params"]) + request_data: Final = params.get("proxy_server_request") + request: Final = ( + mapping_adapter.validate_python(request_data) if request_data is not None else empty_headers + ) + header_data: Final = request.get("headers") + headers: Final = mapping_adapter.validate_python(header_data) if header_data is not None else empty_headers self.model_call_details = { **args.kwargs, + "litellm_params": { + **args.kwargs["litellm_params"], + "no-log": params.get("no-log", False), + "proxy_server_request": { + "headers": {X_LITELLM_DISABLE_CALLBACKS: headers.get(X_LITELLM_DISABLE_CALLBACKS)} + }, + }, **({"async_complete_streaming_response": args.response_obj} if assembled_stream else {}), } verbose_logger.warning("Logging hooks timed out, recording spend without request or response bodies") diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 25d5394c426..0a41585129b 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -13162,3 +13162,94 @@ async def test_invalid_accounting_after_a_hook_timeout_never_exports_unmasked_da assert "Cannot prepare body-free spend accounting after a logging hook timeout" in caplog.text finally: await worker.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("no_log", [False, True]) +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_masking_timeout_preserves_opt_outs_and_redacts_hidden_provider_content( + monkeypatch: pytest.MonkeyPatch, + no_log: bool, + metadata_field: str, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress + + recorded: Final[asyncio.Queue[str]] = asyncio.Queue() + exported: Final[asyncio.Queue[str]] = asyncio.Queue() + never: Final = asyncio.Event() + + async def store_event(line: bytes) -> None: + recorded.put_nowait(line.decode()) + + producer: Final = SpendEventProducer( + address=UnixAddress("/unused"), + on_unavailable="fallback", + buffer_size=1, + connect_timeout=1, + fallback=store_event, + ) + await producer.close(drain_timeout=1) + + class MaskingLogger(CustomLogger): + async def async_logging_hook( + self, + kwargs: dict[str, object], + result: object, + call_type: str, + ) -> tuple[dict[str, object], object]: + await never.wait() + 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: + exported.put_nowait(json.dumps(kwargs, default=str) + str(response_obj)) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "global_disable_no_log_param", False) + monkeypatch.setattr(litellm, "allow_dynamic_callback_disabling", True) + monkeypatch.setattr(proxy_server, "premium_user", True) + logging_obj: Final = _async_request_logging_obj( + "private-provider-content", [MaskingLogger(), _ProxyDBLogger(producer)] + ) + logging_obj.update_environment_variables( + litellm_params={ + "acompletion": True, + "no-log": no_log, + metadata_field: { + "hidden_params": { + "provider_specific_fields": { + "reasoning_content": "private-thinking-canary", + "thinking_blocks": [{"type": "thinking", "thinking": "private-thinking-canary"}], + "citations": [{"cited_text": "private-citation-canary"}], + } + } + }, + "proxy_server_request": {"headers": {"x-litellm-disable-callbacks": "dynamodb"}}, + }, + optional_params={}, + ) + worker: Final = LoggingWorker(timeout=0.5) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse())) + try: + await worker.flush() + assert recorded.qsize() == int(not no_log) + assert exported.qsize() == int(not no_log) + payloads: Final = tuple(recorded.get_nowait() for _ in range(recorded.qsize())) + tuple( + exported.get_nowait() for _ in range(exported.qsize()) + ) + assert all( + "private-thinking-canary" not in payload and "private-citation-canary" not in payload + for payload in payloads + ) + assert not logging_obj.should_run_callback( + "dynamodb", logging_obj.model_call_details["litellm_params"], "async_success" + ) + finally: + await worker.stop()