From 16fd2cd0a46a23249e96cd34fdd26fe1a461bfb9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:21:01 -0700 Subject: [PATCH] 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 | 2 +- litellm/litellm_core_utils/litellm_logging.py | 88 ++++++++++++++----- .../guardrails/guardrail_hooks/presidio.py | 3 +- litellm/utils.py | 7 +- .../test_litellm_logging.py | 39 ++++++-- .../guardrail_hooks/test_presidio.py | 44 ++++++++-- 6 files changed, 145 insertions(+), 38 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..66d30ab2c74 100644 --- a/litellm/integrations/custom_logger.py +++ b/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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c12a67ed7c0..ddcb7db06a2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 35f92403ef6..ef8313123cb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index 4e792240ca1..ce0cd0c6914 100644 --- a/litellm/utils.py +++ b/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) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 5ef152e1f76..7e8cfdb0a94 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index d6fd03f82df..b6edd273d20 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/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}"