fix(guardrails): redact failure logs when masking fails and scan Presidio messages in parallel

A mask that raises now redacts the failure log instead of logging the raw
prompt, and the sync SDK path gets the same 1 s bound as the async one.
A caller cancelled during the mask still gets its async failure logs.
Presidio holds its session lock only while it creates the session, so the
scans of one request and of concurrent requests no longer wait in line.

litellm/integrations/custom_logger.py

@@ -504,7 +504,7 @@ class CustomLogger:  # https://docs.litellm.ai/docs/observability/custom_callbac
         """For masking logged request/response. Return a modified version of the request/result."""
         return kwargs, result
 
-    def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
+    def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, Any], object]:
         """For masking logged request/response. Return a modified version of the request/result."""
         return kwargs, result
 

litellm/litellm_core_utils/litellm_logging.py

@@ -2,6 +2,7 @@
 ## Common Utility file for Logging handler
 # Logging function -> log the exact model details + what's being sent | Non-Blocking
 import asyncio
+import concurrent.futures
 import copy
 import datetime
 import functools
@@ -569,7 +570,16 @@ def _resolve_mantle_region_for_cost(
 _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"})
 _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
 _REDACTED_RESULT: Final = TypeAdapter(object)
-_ORPHANED_FAILURE_LOG_MASKS: Final[set["asyncio.Task[None]"]] = set()  # mutable-ok: keeps orphaned masks alive
+_ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set()  # mutable-ok: keeps orphaned tasks alive
+
+
+def _failure_log_hook_kwargs(details: Mapping[str, Any]) -> dict[str, object]:  # mutable-ok: hooks mask in place
+    return {**details, "messages": copy.deepcopy(details.get("messages"))}
+
+
+def _keep_alive_until_done(task: "asyncio.Future[None]") -> None:
+    _ORPHANED_FAILURE_LOG_TASKS.add(task)
+    task.add_done_callback(_ORPHANED_FAILURE_LOG_TASKS.discard)
 
 
 def _provider_response_id(source: object) -> str | None:
@@ -4188,9 +4198,6 @@ class Logging(LiteLLMLoggingBaseClass):
             if isinstance(callback, CustomGuardrail) and callback.masks_logged_request
         )
 
-    def _failure_log_hook_kwargs(self) -> dict[str, object]:  # mutable-ok: logging hooks mask a plain dict in place
-        return {**self.model_call_details, "messages": copy.deepcopy(self.model_call_details.get("messages"))}
-
     @property
     def messages_sent_to_model(self) -> object:
         return self._unmasked_messages if self._failure_log_masked else self.model_call_details.get("messages")
@@ -4214,21 +4221,25 @@ class Logging(LiteLLMLoggingBaseClass):
 
     def _report_failure_log_mask_error(self, guardrail: CustomGuardrail) -> None:
         verbose_logger.error(
-            "LiteLLM.LoggingError: [Non-Blocking] Exception occurred masking the failure log with %s %s",
+            "LiteLLM.LoggingError: [Non-Blocking] Exception occurred masking the failure log with %s, so its request "
+            "is redacted %s",
             guardrail.guardrail_name,
             traceback.format_exc(),
         )
         self._handle_callback_failure(callback=guardrail)
 
+    def _redact_failure_log_request_after_timeout(self) -> None:
+        verbose_logger.error(
+            "LiteLLM.LoggingError: [Non-Blocking] Masking the failure log took over %ss, so its request is redacted",
+            FAILURE_LOG_MASK_TIMEOUT_SECONDS,
+        )
+        self._redact_failure_log_request()
+
     async def _mask_failure_log_request_with(self, maskers: tuple[CustomGuardrail, ...]) -> None:
         try:
             await asyncio.wait_for(self._run_failure_log_maskers(maskers), timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS)
         except asyncio.TimeoutError:
-            verbose_logger.error(
-                "LiteLLM.LoggingError: [Non-Blocking] Masking the failure log took over %ss, so its request is redacted",
-                FAILURE_LOG_MASK_TIMEOUT_SECONDS,
-            )
-            self._redact_failure_log_request()
+            self._redact_failure_log_request_after_timeout()
 
     async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None:
         for guardrail in maskers:
@@ -4236,12 +4247,20 @@ class Logging(LiteLLMLoggingBaseClass):
                 if not _runs_on_logging_only(guardrail, self.model_call_details):
                     continue
                 self.model_call_details, _ = await guardrail.async_logging_hook(
-                    kwargs=self._failure_log_hook_kwargs(), result=None, call_type=self._logging_hook_call_type
+                    kwargs=_failure_log_hook_kwargs(self.model_call_details),
+                    result=None,
+                    call_type=self._logging_hook_call_type,
                 )
-            except Exception:  # noqa: BLE001  # a failing mask must not stop failure logging
+            except Exception:  # noqa: BLE001  # a failing mask redacts the log instead of stopping failure logging
                 self._report_failure_log_mask_error(guardrail)
+                self._redact_failure_log_request()
+                return
 
-    async def async_mask_failure_log_request(self, on_cancelled: Callable[[], None] | None = None) -> None:
+    async def async_mask_failure_log_request(
+        self,
+        on_cancelled: Callable[[], None] | None = None,
+        on_cancelled_async: Callable[[], Awaitable[None]] | None = None,
+    ) -> None:
         if not self._claim_failure_log_mask():
             return
         maskers: Final = self._failure_log_maskers()
@@ -4251,24 +4270,45 @@ class Logging(LiteLLMLoggingBaseClass):
         try:
             await asyncio.shield(mask)
         except asyncio.CancelledError:
-            _ORPHANED_FAILURE_LOG_MASKS.add(mask)
-            mask.add_done_callback(_ORPHANED_FAILURE_LOG_MASKS.discard)
+            _keep_alive_until_done(mask)
             if on_cancelled is not None:
                 mask.add_done_callback(lambda _: on_cancelled())
+            if on_cancelled_async is not None:
+                mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async())))
             raise
 
     def mask_failure_log_request(self) -> None:
         if is_event_loop_running() or not self._claim_failure_log_mask():
             return
-        for guardrail in self._failure_log_maskers():
-            try:
-                if not _runs_on_logging_only(guardrail, self.model_call_details):
-                    continue
-                self.model_call_details, _ = guardrail.logging_hook(
-                    kwargs=self._failure_log_hook_kwargs(), result=None, call_type=self._logging_hook_call_type
-                )
-            except Exception:  # noqa: BLE001  # a failing mask must not stop failure logging
-                self._report_failure_log_mask_error(guardrail)
+        maskers: Final = self._failure_log_maskers()
+        if not maskers:
+            return
+        masking: Final = executor.submit(
+            functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details
+        )
+        try:
+            masked: Final = masking.result(timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS)
+        except concurrent.futures.TimeoutError:
+            self._redact_failure_log_request_after_timeout()
+            return
+        if masked is None:
+            self._redact_failure_log_request()
+            return
+        self.model_call_details = dict(masked)
+
+    def _sync_failure_log_mask_step(
+        self, details: Mapping[str, Any] | None, guardrail: CustomGuardrail
+    ) -> Mapping[str, Any] | None:
+        if details is None or not _runs_on_logging_only(guardrail, details):
+            return details
+        try:
+            hooked: Final = guardrail.logging_hook(
+                kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type
+            )
+        except Exception:  # noqa: BLE001  # a failing mask redacts the log instead of stopping failure logging
+            self._report_failure_log_mask_error(guardrail)
+            return None
+        return hooked[0]
 
     def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None):
         if start_time is None:

litellm/proxy/guardrails/guardrail_hooks/presidio.py

@@ -390,7 +390,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
             async with self._session_lock:
                 if self._http_session is None or self._http_session.closed:
                     self._http_session = aiohttp.ClientSession()
-                yield self._http_session
+                session: Final = self._http_session
+            yield session
         else:
             # Background thread/loop -> use loop-bound session cache
             # This avoids "attached to a different loop" or "no running event loop" errors

litellm/utils.py

@@ -2432,7 +2432,12 @@ def client(original_function):
                 run_failure_handler: Final = functools.partial(
                     logging_obj.failure_handler, e, traceback_exception, start_time, end_time, for_async_request=True
                 )
-                await logging_obj.async_mask_failure_log_request(on_cancelled=run_failure_handler)
+                await logging_obj.async_mask_failure_log_request(
+                    on_cancelled=run_failure_handler,
+                    on_cancelled_async=functools.partial(
+                        logging_obj.async_failure_handler, e, traceback_exception, start_time, end_time
+                    ),
+                )
                 run_failure_handler()  # DO NOT MAKE THREADED - router retry fallback relies on this!
                 try:
                     await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time)

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -8392,6 +8392,18 @@ class _BlockingNameMaskingGuardrail(_NameMaskingGuardrail):
         return self._mask(kwargs), result
 
 
+class _HangingSyncMaskingGuardrail(_NameMaskingGuardrail):
+    def __init__(self) -> None:
+        super().__init__()
+        self.release: Final = threading.Event()
+
+    def logging_hook(
+        self, kwargs: dict[str, object], result: object, call_type: str
+    ) -> tuple[dict[str, object], object]:
+        _ = self.release.wait(timeout=30)
+        return self._mask(kwargs), result
+
+
 class _FailureLogRecorder(CustomLogger):
     def __init__(self) -> None:
         super().__init__()
@@ -8491,15 +8503,15 @@ async def test_failure_log_keeps_the_prompt_when_no_logging_only_guardrail_masks
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("call", ["completion", "acompletion"])
-async def test_failure_is_still_logged_when_the_masking_guardrail_raises(
+async def test_failure_is_logged_redacted_when_the_masking_guardrail_raises(
     monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"]
 ) -> None:
     recorder, sync_logged = _register_failure_logging(monkeypatch, _RaisingMaskingGuardrail())
 
     await _fail_chat_call(call, copy.deepcopy(_RAW_LOG))
 
-    assert len(recorder.logged) == 1
-    assert len(sync_logged) == 1
+    assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
+    assert sync_logged == [_REDACTED_LOG]
 
 
 @pytest.mark.asyncio
@@ -8507,7 +8519,7 @@ async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mi
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     guardrail: Final = _BlockingNameMaskingGuardrail()
-    _, sync_logged = _register_failure_logging(monkeypatch, guardrail)
+    recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
     call: Final = asyncio.create_task(
         litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error())
     )
@@ -8517,9 +8529,26 @@ async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mi
     with pytest.raises(asyncio.CancelledError):
         await call
     guardrail.release.set()
-    await asyncio.wait_for(_until(lambda: bool(sync_logged)), timeout=10)
+    await asyncio.wait_for(_until(lambda: bool(sync_logged) and bool(recorder.logged)), timeout=10)
 
     assert sync_logged == [_MASKED_LOG]
+    assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)]
+
+
+@pytest.mark.asyncio
+async def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    guardrail: Final = _HangingSyncMaskingGuardrail()
+    recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
+
+    try:
+        await asyncio.wait_for(_fail_chat_call("completion", copy.deepcopy(_RAW_LOG)), timeout=10)
+    finally:
+        guardrail.release.set()
+
+    assert sync_logged == [_REDACTED_LOG]
+    assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
 
 
 @pytest.mark.asyncio

tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py

@@ -4897,11 +4897,17 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]:
     )
 
 
+_PENDING_PRESIDIO_REPLIES: Final[set[asyncio.Task[None]]] = set()  # mutable-ok: keeps pending replies alive
+
+
 class _InMemoryPresidioTransport(asyncio.Transport):
-    def __init__(self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None:
+    def __init__(
+        self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int
+    ) -> None:
         super().__init__()
         self._protocol: Final = protocol
         self._analyzed: Final = analyzed
+        self._replies_wait_for: Final = replies_wait_for
         self._received = b""
         self._closing = False
 
@@ -4921,7 +4927,12 @@ class _InMemoryPresidioTransport(asyncio.Transport):
             + f"Content-Length: {len(reply)}\r\n\r\n".encode()
             + reply
         )
-        asyncio.get_running_loop().call_soon(self._protocol.data_received, response)
+        _PENDING_PRESIDIO_REPLIES.add(asyncio.get_running_loop().create_task(self._reply_once_analyzed(response)))
+
+    async def _reply_once_analyzed(self, response: bytes) -> None:
+        while self._analyzed.qsize() < self._replies_wait_for:
+            await asyncio.sleep(0.01)
+        self._protocol.data_received(response)
 
     def _reply(self, path: str, body: bytes) -> tuple[PresidioAnalyzeResponseItem, ...] | _PresidioAnonymizeReply:
         if path == "/analyze":
@@ -4962,15 +4973,16 @@ class _InMemoryPresidioTransport(asyncio.Transport):
 
 
 class _InMemoryPresidioConnector(aiohttp.BaseConnector):
-    def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None:
+    def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int) -> None:
         super().__init__()
         self._analyzed: Final = analyzed
+        self._replies_wait_for: Final = replies_wait_for
 
     async def _create_connection(  # pyright: ignore[reportImplicitOverride]  # aiohttp's connector extension point
         self, req: aiohttp.ClientRequest, traces: Sequence[object], timeout: aiohttp.ClientTimeout
     ) -> ResponseHandler:
         protocol: Final = ResponseHandler(asyncio.get_running_loop())
-        protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed))
+        protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed, self._replies_wait_for))
         return protocol
 
 
@@ -4979,17 +4991,37 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna
 
 
 def _guardrail_with_in_memory_presidio(
-    analyzed: asyncio.Queue[PresidioAnalyzeRequest],
+    analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int = 0
 ) -> OPTIONAL_PresidioPIIMasking:
     guardrail: Final = OPTIONAL_PresidioPIIMasking(
         pii_entities_config=_BLOCK_CARD_MASK_EMAIL,
         presidio_analyzer_api_base="http://presidio-analyzer.test/",
         presidio_anonymizer_api_base="http://presidio-anonymizer.test/",
     )
-    guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed))
+    guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed, replies_wait_for))
     return guardrail
 
 
+@pytest.mark.asyncio
+async def test_presidio_scans_every_message_of_a_request_at_once() -> None:
+    texts: Final = (f"My email is {_EMAIL}", f"Write to {_EMAIL} today")
+    analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
+    guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, replies_wait_for=len(texts))
+    try:
+        _ = await asyncio.wait_for(
+            guardrail.async_logging_hook(
+                kwargs={"messages": [{"role": "user", "content": text} for text in texts]},
+                result=None,
+                call_type="acompletion",
+            ),
+            timeout=5,
+        )
+    finally:
+        await guardrail._close_http_session()
+
+    assert sorted(request.get("text") or "" for request in _drain(analyzed)) == sorted(texts)
+
+
 @pytest.mark.asyncio
 async def test_check_pii_raises_blocked_entity_for_card() -> None:
     text: Final = f"My credit card number is {_CARD_NUMBER} and my email is {_EMAIL}"