fix(logging): preserve privacy controls after hook timeouts

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")

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()