From 519ecb67c848fe2ae3ee5d59cf0c3df69a4f69f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:10:04 -0700 Subject: [PATCH 1/8] fix(guardrails): mask the prompt in failure logs for logging_only Presidio A Presidio guardrail in logging_only mode masked PII only in success logs. On a failed call every failure logger got the raw prompt in kwargs["messages"] and in standard_logging_object["messages"] Guardrails that mask the logged request now opt in with masks_logged_request (Presidio does). The failure path runs their logging hook once per call on a copy of the messages, before the standard logging payload is built, so the caller's messages stay untouched for retries and fallbacks --- litellm/integrations/custom_guardrail.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 105 +++++++--- .../guardrails/guardrail_hooks/presidio.py | 3 +- litellm/utils.py | 1 + .../test_litellm_logging.py | 180 +++++++++++++++++- .../unit/passthrough/test_passthrough_main.py | 1 + .../guardrail_hooks/test_presidio.py | 41 ++++ .../test_responses_prompt_management.py | 1 + 8 files changed, 304 insertions(+), 29 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..fc29742686f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -226,6 +226,7 @@ class CustomGuardrail(CustomLogger): use_native_lifecycle_hooks: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + masks_logged_request: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None logging_only_continue_on_input_failure: bool diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7f4b069ded6..95814989322 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -69,6 +69,7 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.langtrace import langtrace_trace_endpoint from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.classifier_logging import ( classifier_audit_fields, classifier_input_snapshot, @@ -140,6 +141,7 @@ from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse from litellm.types.extraction import ExtractionResponse +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.integrations.s3_v2 import S3PartitionGranularity from litellm.types.interactions import ( InteractionsAPIResponse, @@ -613,6 +615,10 @@ def _timestamp_seconds(moment: object) -> float | None: return None +def _runs_on_logging_only(guardrail: CustomGuardrail, data: Mapping[str, object]) -> bool: + return guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.logging_only) is True + + def same_attempt_marker(first: object, second: object) -> bool: return first is second or first == second @@ -641,6 +647,7 @@ class Logging(LiteLLMLoggingBaseClass): baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None _on_failure_payload_built: "Callable[[], None] | None" = None + _failure_log_masked: bool = False @property def _defer_async_logging(self) -> bool: @@ -2592,6 +2599,7 @@ class Logging(LiteLLMLoggingBaseClass): sync_sdk: Final = self.is_sync_litellm_request(litellm_params) passthrough: Final = self.call_type == CallTypes.pass_through.value if sync_sdk and not prefer_async_handlers and not passthrough: + await self.async_mask_failure_log_request() self.failure_handler(exception, traceback_exception) return @@ -3372,28 +3380,20 @@ class Logging(LiteLLMLoggingBaseClass): hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) for callback in hook_callbacks: if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - if ( - callback.should_run_guardrail( - data=call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): + if not _runs_on_logging_only(callback, call_details): continue self.shadow_eval_request_snapshot = None call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict kwargs=call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) elif isinstance(callback, CustomLogger): call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict kwargs=call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) if not for_async_request: self.model_call_details = call_details @@ -3968,22 +3968,14 @@ class Logging(LiteLLMLoggingBaseClass): 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 - ): + if not _runs_on_logging_only(callback, self.model_call_details): 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, + call_type=self._logging_hook_call_type, ) elif isinstance(callback, CustomLogger): result = redact_message_input_output_from_custom_logger( @@ -3992,7 +3984,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details, result = await callback.async_logging_hook( kwargs=self.model_call_details, result=result, - call_type=self.call_type, + call_type=self._logging_hook_call_type, ) except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) verbose_logger.error( @@ -4169,6 +4161,66 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) + @property + def _logging_hook_call_type(self) -> str: + return cast(str, self.call_type) # cast-ok: Logging.__init__ takes call_type untyped and callers pass names + + def _async_failure_callbacks(self) -> Sequence[str | Callable[..., object] | CustomLogger]: + return self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_async_failure_callbacks, + global_callbacks=litellm._async_failure_callback, + ) + + def _failure_log_maskers(self) -> tuple[CustomGuardrail, ...]: + callbacks: Final = ( + *self.get_combined_callback_list(self.dynamic_failure_callbacks, litellm.failure_callback), + *self._async_failure_callbacks(), + ) + return tuple( + callback + for callback in dict.fromkeys(callbacks) + 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"))} + + 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", + guardrail.guardrail_name, + traceback.format_exc(), + ) + self._handle_callback_failure(callback=guardrail) + + async def async_mask_failure_log_request(self) -> None: + if self._failure_log_masked: + return + self._failure_log_masked = True + for guardrail in self._failure_log_maskers(): + try: + 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 + ) + except Exception: # noqa: BLE001 # a failing mask must not stop failure logging + self._report_failure_log_mask_error(guardrail) + + def mask_failure_log_request(self) -> None: + if self._failure_log_masked or is_event_loop_running(): + return + self._failure_log_masked = True + 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) + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time @@ -4291,6 +4343,7 @@ class Logging(LiteLLMLoggingBaseClass): is_sync_request: Final = self.is_sync_litellm_request(litellm_params) try: + self.mask_failure_log_request() start_time, end_time = self._failure_handler_helper_fn( exception=exception, traceback_exception=traceback_exception, @@ -4489,6 +4542,8 @@ class Logging(LiteLLMLoggingBaseClass): await self.special_failure_handlers(exception=exception) if not self.should_run_logging(event_type="async_failure"): # prevent double logging return + self.mark_logging_complete(event_type="async_failure") + await self.async_mask_failure_log_request() start_time, end_time = self._failure_handler_helper_fn( exception=exception, traceback_exception=traceback_exception, @@ -4497,14 +4552,10 @@ class Logging(LiteLLMLoggingBaseClass): ) if self._on_failure_payload_built is not None: self._on_failure_payload_built() - callbacks: Final = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_async_failure_callbacks, - global_callbacks=litellm._async_failure_callback, - ) + callbacks: Final = self._async_failure_callbacks() result: Final = None # result sent to all loggers, init this to None incase it's not created - self.mark_logging_complete(event_type="async_failure") for callback in callbacks: try: litellm_params = self.model_call_details.get("litellm_params", {}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 7e2f6bfec33..35f92403ef6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -17,7 +17,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime from itertools import accumulate, count, pairwise -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast import aiohttp from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config @@ -238,6 +238,7 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 72c37c3760e..0366f31baf7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2428,6 +2428,7 @@ def client(original_function): end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) + await logging_obj.async_mask_failure_log_request() try: logging_obj.failure_handler( e, traceback_exception, start_time, end_time, for_async_request=True diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 299088a6c4f..06c74bf26f8 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object from importlib.machinery import ModuleSpec from types import MappingProxyType, ModuleType -from typing import Final, Literal +from typing import ClassVar, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -38,6 +38,7 @@ from litellm.constants import ( SENTRY_PII_DENYLIST, ) from litellm.cost_calculator import ocr_batch_cost +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -57,6 +58,7 @@ from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, @@ -8336,6 +8338,182 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() +_RAW_PROMPT: Final = "My name is John Smith" +_RAW_LOG: Final = [{"role": "user", "content": _RAW_PROMPT}] +_MASKED_LOG: Final = [{"role": "user", "content": "My name is "}] + + +class _NameMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + def __init__(self, event_hook: GuardrailEventHooks = GuardrailEventHooks.logging_only) -> None: + super().__init__(guardrail_name="name-masker", event_hook=event_hook) + self.scans = 0 + + def _mask(self, kwargs: dict) -> dict: + self.scans += 1 + for message in kwargs["messages"]: + message["content"] = message["content"].replace("John Smith", "") + return kwargs + + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: + return self._mask(kwargs), result + + async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: + await asyncio.sleep(0) + return self._mask(kwargs), result + + +class _UnflaggedNameMaskingGuardrail(_NameMaskingGuardrail): + masks_logged_request: ClassVar[bool] = False + + +class _RaisingMaskingGuardrail(_NameMaskingGuardrail): + def _mask(self, kwargs: dict) -> dict: + raise RuntimeError("masking failed") + + +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[list[dict], list[dict]]] = [] + + def _record(self, kwargs: dict) -> None: + self.logged.append( + (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) + ) + + def log_failure_event(self, kwargs: dict, response_obj: object, start_time: object, end_time: object) -> None: + self._record(kwargs) + + async def async_log_failure_event( + self, kwargs: dict, response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + +def _register_failure_logging( + monkeypatch: pytest.MonkeyPatch, guardrail: CustomGuardrail +) -> tuple[_FailureLogRecorder, list[list[dict]]]: + recorder: Final = _FailureLogRecorder() + sync_logged: Final[list[list[dict]]] = [] + + def record_sync_failure(kwargs: dict, completion_response: object, start_time: object, end_time: object) -> None: + sync_logged.append(copy.deepcopy(kwargs["messages"])) + + for name in ("input_callback", "success_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder]) + monkeypatch.setattr(litellm, "failure_callback", [record_sync_failure]) + return recorder, sync_logged + + +async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict]) -> None: + error: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") + request: Final = ( + litellm.acompletion(model="gpt-4o-mini", messages=messages, mock_response=error) + if call == "acompletion" + else asyncio.to_thread(litellm.completion, model="gpt-4o-mini", messages=messages, mock_response=error) + ) + with pytest.raises(litellm.AuthenticationError): + await request + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_log_carries_the_prompt_a_logging_only_guardrail_masked( + monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"] +) -> None: + guardrail: Final = _NameMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + messages: Final = copy.deepcopy(_RAW_LOG) + + await _fail_chat_call(call, messages) + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + assert sync_logged == [_MASKED_LOG] + assert messages == _RAW_LOG + assert guardrail.scans == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_guardrail", + [_UnflaggedNameMaskingGuardrail, lambda: _NameMaskingGuardrail(event_hook=GuardrailEventHooks.pre_call)], + ids=["logging_only_without_masking_opt_in", "masking_opt_in_in_pre_call_mode"], +) +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_log_keeps_the_prompt_when_no_logging_only_guardrail_masks_it( + monkeypatch: pytest.MonkeyPatch, + make_guardrail: Callable[[], _NameMaskingGuardrail], + call: Literal["completion", "acompletion"], +) -> None: + guardrail: Final = make_guardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + + await _fail_chat_call(call, copy.deepcopy(_RAW_LOG)) + + assert recorder.logged == [(_RAW_LOG, _RAW_LOG)] + assert sync_logged == [_RAW_LOG] + assert guardrail.scans == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["completion", "acompletion"]) +async def test_failure_is_still_logged_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 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_concurrent_failure_notifications_log_one_masked_failure() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="concurrent-masked-failure", + function_id="concurrent-masked-failure", + dynamic_async_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + error: Final = litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini") + + await asyncio.gather(*(logging_obj.async_failure_handler(error, "tb") for _ in range(3))) + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_sync_sdk_failure_dispatched_from_async_code_logs_the_masked_prompt() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="sync-sdk-masked-failure", + function_id="sync-sdk-masked-failure", + dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await logging_obj.dispatch_failure_handlers(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] + + def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", diff --git a/tests/unit/passthrough/test_passthrough_main.py b/tests/unit/passthrough/test_passthrough_main.py index 084a0fef9db..8c5d3fd45be 100644 --- a/tests/unit/passthrough/test_passthrough_main.py +++ b/tests/unit/passthrough/test_passthrough_main.py @@ -646,6 +646,7 @@ async def test_allm_passthrough_route_429_streaming_raises(): mock_logging_obj.update_environment_variables = MagicMock() mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock() mock_logging_obj.async_failure_handler = AsyncMock() + mock_logging_obj.async_mask_failure_log_request = AsyncMock() with ( patch( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index e312378969a..fc1daf02879 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -26,6 +26,7 @@ import litellm from litellm import mock_completion from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.presidio import ( OPTIONAL_PresidioPIIMasking, @@ -4425,6 +4426,46 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use assert later[3]["content"] == "Now compare against too." +class _FailureLogRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[list[dict], list[dict]]] = [] + + async def async_log_failure_event( + self, kwargs: dict, response_obj: object, start_time: object, end_time: object + ) -> None: + self.logged.append( + (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) + ) + + +async def test_logging_only_presidio_masks_the_prompt_in_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed call is logged with the prompt a logging_only Presidio guardrail masked, the same as a + successful one. The analyzer and anonymizer are the in-process fake handed over through api_base.""" + recorder: Final = _FailureLogRecorder() + for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + async with TestServer(_fake_presidio_app()) as server: + guardrail: Final = OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base=str(server.make_url("/")), + presidio_anonymizer_api_base=str(server.make_url("/")), + pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK}, + logging_only=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder]) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "My name is John Smith"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + await guardrail._close_http_session() + + masked: Final = [{"role": "user", "content": "My name is "}] + assert recorder.logged == [(masked, masked)] + + @pytest.mark.asyncio @pytest.mark.parametrize("surface", ["mcp_arguments", "mcp_result", "llm_output"]) @pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK]) diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 9af92e7f9ca..bc475eba095 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -52,6 +52,7 @@ def _make_logging_obj( logging_obj.get_chat_completion_prompt.return_value = prompt_return logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return) logging_obj.async_failure_handler = AsyncMock() + logging_obj.async_mask_failure_log_request = AsyncMock() logging_obj.model_call_details = {} return logging_obj From c85e9afd2a25e6f4c5ccc532924462ff081ce34a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:26:59 -0700 Subject: [PATCH 2/8] fix(guardrails): bound the failure-log mask and keep billing on the sent prompt Bound the failure-log mask at 1 s and redact the request when it runs out, so a hung Presidio cannot hold every failed request's error response. Run the sync failure handler when the caller is cancelled during the mask, so router cooldowns still happen. Redact unauthenticated rejections without calling Presidio. Estimate provider-billed failure usage on the prompt sent to the model, not the masked log copy. Mark Guardrails AI as masking the logged request, since its logging hook rewrites the prompt like Presidio's --- litellm/constants.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 69 ++++++++++-- litellm/litellm_core_utils/redact_messages.py | 4 +- .../guardrails_ai/guardrails_ai.py | 4 +- litellm/proxy/utils.py | 14 ++- litellm/utils.py | 13 ++- .../test_litellm_logging.py | 102 ++++++++++++++---- .../guardrails_ai/test_guardrails_ai.py | 40 +++++++ .../guardrail_hooks/test_presidio.py | 13 ++- ..._utils_model_creation_and_error_logging.py | 32 +++++- .../test_post_call_failure_hook.py | 62 ++++++++++- 11 files changed, 303 insertions(+), 51 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 42ad624026c..515ea843f18 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -166,6 +166,7 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 +FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 95814989322..c12a67ed7c0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1,6 +1,7 @@ # What is this? ## Common Utility file for Logging handler # Logging function -> log the exact model details + what's being sent | Non-Blocking +import asyncio import copy import datetime import functools @@ -42,6 +43,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, + FAILURE_LOG_MASK_TIMEOUT_SECONDS, INTERNAL_CALLBACK_NAMES, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -126,6 +128,7 @@ from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redacted_model_call_details_for_custom_logger, + redacted_request_entries, redacted_response_obj_for_custom_logger, should_redact_message_logging, ) @@ -566,6 +569,7 @@ 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 def _provider_response_id(source: object) -> str | None: @@ -648,6 +652,7 @@ class Logging(LiteLLMLoggingBaseClass): baseline_observation: "CapturedBaselineObservation | None" = None _on_failure_payload_built: "Callable[[], None] | None" = None _failure_log_masked: bool = False + _unmasked_messages: object = None @property def _defer_async_logging(self) -> bool: @@ -2599,8 +2604,9 @@ class Logging(LiteLLMLoggingBaseClass): sync_sdk: Final = self.is_sync_litellm_request(litellm_params) passthrough: Final = self.call_type == CallTypes.pass_through.value if sync_sdk and not prefer_async_handlers and not passthrough: - await self.async_mask_failure_log_request() - self.failure_handler(exception, traceback_exception) + run_failure_handler: Final = functools.partial(self.failure_handler, exception, traceback_exception) + await self.async_mask_failure_log_request(on_cancelled=run_failure_handler) + run_failure_handler() return try: @@ -4185,6 +4191,27 @@ class Logging(LiteLLMLoggingBaseClass): 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") + + def _claim_failure_log_mask(self) -> bool: + if self._failure_log_masked: + return False + self._failure_log_masked = True + self._unmasked_messages = self.model_call_details.get("messages") + return True + + def _redact_failure_log_request(self) -> None: + self.model_call_details = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + + def redact_failure_log_request(self) -> None: + masks_this_request: Final = any( + _runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers() + ) + if masks_this_request and self._claim_failure_log_mask(): + self._redact_failure_log_request() + 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", @@ -4193,11 +4220,18 @@ class Logging(LiteLLMLoggingBaseClass): ) self._handle_callback_failure(callback=guardrail) - async def async_mask_failure_log_request(self) -> None: - if self._failure_log_masked: - return - self._failure_log_masked = True - for guardrail in self._failure_log_maskers(): + 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() + + async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None: + for guardrail in maskers: try: if not _runs_on_logging_only(guardrail, self.model_call_details): continue @@ -4207,10 +4241,25 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: # noqa: BLE001 # a failing mask must not stop failure logging self._report_failure_log_mask_error(guardrail) - def mask_failure_log_request(self) -> None: - if self._failure_log_masked or is_event_loop_running(): + async def async_mask_failure_log_request(self, on_cancelled: Callable[[], None] | None = None) -> None: + if not self._claim_failure_log_mask(): + return + maskers: Final = self._failure_log_maskers() + if not maskers: + return + mask: Final = asyncio.create_task(self._mask_failure_log_request_with(maskers)) + try: + await asyncio.shield(mask) + except asyncio.CancelledError: + _ORPHANED_FAILURE_LOG_MASKS.add(mask) + mask.add_done_callback(_ORPHANED_FAILURE_LOG_MASKS.discard) + if on_cancelled is not None: + mask.add_done_callback(lambda _: on_cancelled()) + raise + + def mask_failure_log_request(self) -> None: + if is_event_loop_running() or not self._claim_failure_log_mask(): return - self._failure_log_masked = True for guardrail in self._failure_log_maskers(): try: if not _runs_on_logging_only(guardrail, self.model_call_details): diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..f1eae1f82c7 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -67,11 +67,11 @@ def redacted_model_call_details_for_custom_logger(model_call_details: dict, cust return model_call_details return { **{key: value for key, value in model_call_details.items() if key != SERVED_OUTPUT_TEXTS_KEY}, - **_redacted_entries_for_opted_out_logger(model_call_details), + **redacted_request_entries(model_call_details), } -def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, object]) -> Mapping[str, object]: +def redacted_request_entries(model_call_details: Mapping[str, object]) -> Mapping[str, object]: optional_params: Final = _as_mapping(model_call_details.get("optional_params")) litellm_params: Final = _as_mapping(model_call_details.get("litellm_params")) standard_logging_object: Final = _as_mapping(model_call_details.get("standard_logging_object")) diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index cf6592a3e58..fadd91f02d9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -7,7 +7,7 @@ import json import os -from typing import TYPE_CHECKING, Final, Literal, TypedDict +from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypedDict from fastapi import HTTPException @@ -48,6 +48,8 @@ class GuardrailsAIResponsePreCall(TypedDict): class GuardrailsAI(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + def __init__( self, guard_name: str, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9f58c267231..33ae99361d7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1592,6 +1592,11 @@ def _failure_fields_to_lift( if _logging_obj is None: return _EMPTY_LIFT _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) + _sent_model_call_details: Final = ( + {**_model_call_details, "messages": _logging_obj.messages_sent_to_model} + if isinstance(_logging_obj, Logging) + else _model_call_details + ) _first_handoff: Final = _model_call_details.get("first_api_call_start_time") _litellm_params: Final = getattr(_logging_obj, "litellm_params", None) _base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details) @@ -1599,7 +1604,7 @@ def _failure_fields_to_lift( _success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued _billed: Final = provider_billed and not _success_already_logged _usage_to_lift: Final = _failure_usage_to_lift( - model_call_details=_model_call_details, + model_call_details=_sent_model_call_details, request_body=request_data, dispatched=_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, provider_billed=_billed, @@ -1608,7 +1613,7 @@ def _failure_fields_to_lift( dispatched_call_type=dispatched_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) + _estimated_prompt_usage(_sent_model_call_details, request_data, dispatched_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -3947,8 +3952,9 @@ class ProxyLogging: litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None) set_up_by_request: Final = litellm_logging_obj is not None request_fields: Final = cast(Mapping[str, object], request_data) # cast-ok: request bodies are str-keyed JSON + authenticated: Final = error_type != ProxyErrorTypes.auth_error logged_metadata_keys: Final = _proxy_only_error_logged_metadata_keys( - authenticated=error_type != ProxyErrorTypes.auth_error, past_pre_call=set_up_by_request + authenticated=authenticated, past_pre_call=set_up_by_request ) user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( user_api_key_dict=user_api_key_dict @@ -4059,6 +4065,8 @@ class ProxyLogging: input=input, api_key="", ) + if not authenticated: + litellm_logging_obj.redact_failure_log_request() await self._dispatch_proxy_only_failure_handlers( litellm_logging_obj=litellm_logging_obj, diff --git a/litellm/utils.py b/litellm/utils.py index 0366f31baf7..4e792240ca1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10,6 +10,7 @@ import binascii import contextvars import copy import datetime +import functools import hashlib import inspect import io @@ -2428,13 +2429,11 @@ def client(original_function): end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) - await logging_obj.async_mask_failure_log_request() - try: - logging_obj.failure_handler( - e, traceback_exception, start_time, end_time, for_async_request=True - ) # DO NOT MAKE THREADED - router retry fallback relies on this! - except Exception as e: - raise e + 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) + 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) except Exception as e: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 06c74bf26f8..5ef152e1f76 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object from importlib.machinery import ModuleSpec from types import MappingProxyType, ModuleType -from typing import ClassVar, Final, Literal +from typing import ClassVar, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -8341,6 +8341,7 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur _RAW_PROMPT: Final = "My name is John Smith" _RAW_LOG: Final = [{"role": "user", "content": _RAW_PROMPT}] _MASKED_LOG: Final = [{"role": "user", "content": "My name is "}] +_REDACTED_LOG: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] class _NameMaskingGuardrail(CustomGuardrail): @@ -8350,16 +8351,20 @@ class _NameMaskingGuardrail(CustomGuardrail): super().__init__(guardrail_name="name-masker", event_hook=event_hook) self.scans = 0 - def _mask(self, kwargs: dict) -> dict: + def _mask(self, kwargs: dict[str, object]) -> dict[str, object]: self.scans += 1 - for message in kwargs["messages"]: + for message in cast(list[dict[str, str]], kwargs["messages"]): message["content"] = message["content"].replace("John Smith", "") return kwargs - def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: return self._mask(kwargs), result - async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]: + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: await asyncio.sleep(0) return self._mask(kwargs), result @@ -8369,36 +8374,53 @@ class _UnflaggedNameMaskingGuardrail(_NameMaskingGuardrail): class _RaisingMaskingGuardrail(_NameMaskingGuardrail): - def _mask(self, kwargs: dict) -> dict: + def _mask(self, kwargs: dict[str, object]) -> dict[str, object]: raise RuntimeError("masking failed") +class _BlockingNameMaskingGuardrail(_NameMaskingGuardrail): + def __init__(self) -> None: + super().__init__() + self.scanning: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.scanning.set() + await self.release.wait() + return self._mask(kwargs), result + + class _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.logged: list[tuple[list[dict], list[dict]]] = [] + self.logged: list[tuple[object, object]] = [] - def _record(self, kwargs: dict) -> None: - self.logged.append( - (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) - ) + def _record(self, kwargs: dict[str, object]) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"]))) - def log_failure_event(self, kwargs: dict, response_obj: object, start_time: object, end_time: object) -> None: + def log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: self._record(kwargs) async def async_log_failure_event( - self, kwargs: dict, response_obj: object, start_time: object, end_time: object + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object ) -> None: self._record(kwargs) def _register_failure_logging( monkeypatch: pytest.MonkeyPatch, guardrail: CustomGuardrail -) -> tuple[_FailureLogRecorder, list[list[dict]]]: +) -> tuple[_FailureLogRecorder, list[object]]: recorder: Final = _FailureLogRecorder() - sync_logged: Final[list[list[dict]]] = [] + sync_logged: Final[list[object]] = [] - def record_sync_failure(kwargs: dict, completion_response: object, start_time: object, end_time: object) -> None: + def record_sync_failure( + kwargs: dict[str, object], completion_response: object, start_time: object, end_time: object + ) -> None: sync_logged.append(copy.deepcopy(kwargs["messages"])) for name in ("input_callback", "success_callback", "_async_success_callback", "_async_failure_callback"): @@ -8408,8 +8430,17 @@ def _register_failure_logging( return recorder, sync_logged -async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict]) -> None: - error: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") +def _bad_key_error() -> litellm.AuthenticationError: + return litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") + + +async def _until(condition: Callable[[], bool]) -> None: + while not condition(): + await asyncio.sleep(0) + + +async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict[str, str]]) -> None: + error: Final = _bad_key_error() request: Final = ( litellm.acompletion(model="gpt-4o-mini", messages=messages, mock_response=error) if call == "acompletion" @@ -8471,6 +8502,41 @@ async def test_failure_is_still_logged_when_the_masking_guardrail_raises( assert len(sync_logged) == 1 +@pytest.mark.asyncio +async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mid_mask( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + _, 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()) + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + guardrail.release.set() + await asyncio.wait_for(_until(lambda: bool(sync_logged)), timeout=10) + + assert sync_logged == [_MASKED_LOG] + + +@pytest.mark.asyncio +async def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: + recorder, sync_logged = _register_failure_logging(monkeypatch, _BlockingNameMaskingGuardrail()) + + with pytest.raises(litellm.AuthenticationError): + await asyncio.wait_for( + litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()), + timeout=10, + ) + await asyncio.wait_for(_until(lambda: bool(recorder.logged)), timeout=10) + + assert sync_logged == [_REDACTED_LOG] + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + @pytest.mark.asyncio @pytest.mark.usefixtures("no_global_callbacks") async def test_concurrent_failure_notifications_log_one_masked_failure() -> None: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py index c6a7fbd7cab..608acc01865 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/guardrails_ai/test_guardrails_ai.py @@ -227,3 +227,43 @@ async def test_guardrails_ai_process_input(): # Should keep original content when no output fields are present assert result["messages"][0]["content"] == "Original message" + + +@pytest.mark.asyncio +async def test_guardrails_ai_logging_only_masks_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.guardrails.guardrail_hooks.guardrails_ai.guardrails_ai import GuardrailsAIResponse + from litellm.types.guardrails import GuardrailEventHooks + + logged: list[object] = [] + + class _FailureLogRecorder(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + logged.append(kwargs["messages"]) + + guardrail = GuardrailsAI( + guardrail_name="pii-guard", + guard_name="pii-guard", + api_base="http://test.example.com", + event_hook=GuardrailEventHooks.logging_only, + ) + for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", [guardrail, _FailureLogRecorder()]) + + with patch.object( + guardrail, + "make_guardrails_ai_api_request", + AsyncMock(return_value=GuardrailsAIResponse(validatedOutput="My name is ")), + ): + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "My name is John Smith"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + + assert logged == [[{"role": "user", "content": "My name is "}]] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index fc1daf02879..d6fd03f82df 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -11,7 +11,7 @@ import re from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager -from typing import Final, Literal +from typing import Final, Literal, cast from unittest.mock import MagicMock, patch import aiohttp @@ -34,7 +34,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import ( ) from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.proxy.guardrails.guardrail_hooks.presidio import PresidioAnalyzeRequest, PresidioAnalyzeResponseItem -from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices +from litellm.types.utils import Choices, Delta, Message, ModelResponse, StandardLoggingPayload, StreamingChoices def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""): @@ -4429,14 +4429,13 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use class _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() - self.logged: list[tuple[list[dict], list[dict]]] = [] + self.logged: list[tuple[object, object]] = [] async def async_log_failure_event( - self, kwargs: dict, response_obj: object, start_time: object, end_time: object + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self.logged.append( - (copy.deepcopy(kwargs["messages"]), copy.deepcopy(kwargs["standard_logging_object"]["messages"])) - ) + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"]))) async def test_logging_only_presidio_masks_the_prompt_in_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index f2151572f49..c6e77d979e3 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -4,7 +4,7 @@ import json import smtplib import traceback from collections.abc import Mapping -from typing import Final +from typing import ClassVar, Final import httpx import pytest @@ -3357,6 +3357,36 @@ def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object: return logging_obj +class _LongPlaceholderMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return {**kwargs, "messages": [{"role": "user", "content": " " * 40}]}, result + + +@pytest.mark.asyncio +async def test_a_billed_failure_is_estimated_on_the_prompt_sent_not_the_masked_failure_log(monkeypatch): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + + def billed_prompt_tokens(logging_obj: object) -> int: + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) + return lifted["combined_usage_object"].prompt_tokens + + unmasked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={})) + masked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={})) + monkeypatch.setattr( + litellm, + "_async_failure_callback", + [_LongPlaceholderMaskingGuardrail(guardrail_name="masker", event_hook=GuardrailEventHooks.logging_only)], + ) + await masked.async_mask_failure_log_request() + + assert masked.model_call_details["messages"] != ONE_USER_MESSAGE + assert billed_prompt_tokens(masked) == billed_prompt_tokens(unmasked) + + def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch): """``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat completions) next to the router-stamped ``metadata.model_info.id``; the failure row must diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index f880d3548f7..1571c29afaf 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -7,18 +7,20 @@ import asyncio from collections.abc import Callable from datetime import datetime from types import MappingProxyType -from typing import Final +from typing import ClassVar, Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException import litellm -from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY, REDACTED_BY_LITELLM from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CachingDetails @@ -499,6 +501,62 @@ async def test_proxy_only_failure_log_keeps_caller_tags_only_after_auth( assert recorded[0]["standard_logging_object"]["request_tags"] == logged_tags +class _PromptMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + def __init__(self) -> None: + super().__init__(guardrail_name="prompt-masker", event_hook=GuardrailEventHooks.logging_only) + self.scans = 0 + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.scans += 1 + return {**kwargs, "messages": [{"role": "user", "content": "My name is "}]}, result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error_type", "masker_scans", "logged_content"), + [(ProxyErrorTypes.auth_error, 0, REDACTED_BY_LITELLM), (None, 1, "My name is ")], +) +async def test_proxy_only_failure_log_masks_after_auth_and_redacts_before_it( + proxy_logging, + monkeypatch: pytest.MonkeyPatch, + error_type: ProxyErrorTypes | None, + masker_scans: int, + logged_content: str, +) -> None: + from litellm.proxy import proxy_server + + recorded: Final[list[dict[str, object]]] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + recorded.append(kwargs) + + masker: Final = _PromptMaskingGuardrail() + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [masker, _RecordingLogger()]) + proxy_logging.alert_types = [] + + await proxy_logging.post_call_failure_hook( + request_data={"model": "internal-model", "messages": [{"role": "user", "content": "My name is John Smith"}]}, + original_exception=HTTPException(status_code=401, detail="Invalid proxy server token passed"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + error_type=error_type, + route="/chat/completions", + ) + + logged_messages: Final = [{"role": "user", "content": logged_content}] + assert masker.scans == masker_scans + assert [(kwargs["messages"], kwargs["standard_logging_object"]["messages"]) for kwargs in recorded] == [ + (logged_messages, logged_messages) + ] + + @pytest.mark.asyncio async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch 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 3/8] 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}" From a4338dc0e8a60e922f5376468b34418b2b97691c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:13:04 -0700 Subject: [PATCH 4/8] fix(logging): mask every logged copy of the prompt on failure and redact when the logging queue drops the mask Failure logs now carry the masked prompt in input, the messages the proxy keeps in optional_params, and the proxy request body, so Langfuse model parameters and the spend log match. A sync mask the full logging executor drops is redacted, and the router's mid-stream failure-log wait leaves the mask its own second. --- litellm/constants.py | 2 +- litellm/integrations/custom_logger.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 65 +++++++++++++-- .../test_litellm_logging.py | 81 +++++++++++++++++++ 4 files changed, 141 insertions(+), 9 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 515ea843f18..906995a98ae 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -165,8 +165,8 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( # Streams exceeding this duration are terminated with a Timeout error. # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") -MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 +MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS + 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 66d30ab2c74..efc9be4c1bb 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[str, Any], object]: + def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, object], 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 ddcb7db06a2..0354b504790 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6,6 +6,7 @@ import concurrent.futures import copy import datetime import functools +import itertools import json import math import os @@ -573,10 +574,54 @@ _REDACTED_RESULT: Final = TypeAdapter(object) _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 +_ABSENT: Final = object() +_LOGGED_PROMPT_COPY_PATHS: Final = ( + ("input",), + ("optional_params", "messages"), + ("litellm_params", "proxy_server_request", "body", "messages"), +) + + +def _failure_log_hook_kwargs(details: Mapping[str, object]) -> dict[str, object]: # mutable-ok: hooks mask in place return {**details, "messages": copy.deepcopy(details.get("messages"))} +def _logged_field(container: object, key: str) -> object: + if not isinstance(container, Mapping): + return _ABSENT + return cast(Mapping[str, object], container).get(key, _ABSENT) # cast-ok: logged request dicts are str-keyed JSON + + +def _with_field(inner: object, step: tuple[object, str]) -> object: + container, key = step + return {**cast(Mapping[str, object], container), key: inner} # cast-ok: containers above a found copy are Mappings + + +def _with_prompt_copy_replaced( + details: Mapping[str, object], path: Sequence[str], raw: object, replacement: object +) -> Mapping[str, object]: + containers: Final = tuple(itertools.accumulate(path, _logged_field, initial=details)) + if containers[-1] is not raw: + return details + rebuilt: Final = functools.reduce(_with_field, reversed(tuple(zip(containers, path))), replacement) + return cast(Mapping[str, object], rebuilt) # cast-ok: the outermost rebuilt container is details + + +def _with_prompt_copies_replaced( + details: Mapping[str, object], raw: object +) -> dict[str, object]: # mutable-ok: becomes the Logging object's model_call_details + if raw is None: + return dict(details) + replacement: Final = details.get("messages") + return dict( + functools.reduce( + lambda fields, path: _with_prompt_copy_replaced(fields, path, raw, replacement), + _LOGGED_PROMPT_COPY_PATHS, + details, + ) + ) + + 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) @@ -4210,7 +4255,9 @@ class Logging(LiteLLMLoggingBaseClass): return True def _redact_failure_log_request(self) -> None: - self.model_call_details = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + self.model_call_details = _with_prompt_copies_replaced( + {**self.model_call_details, **redacted_request_entries(self.model_call_details)}, self._unmasked_messages + ) def redact_failure_log_request(self) -> None: masks_this_request: Final = any( @@ -4255,6 +4302,7 @@ class Logging(LiteLLMLoggingBaseClass): self._report_failure_log_mask_error(guardrail) self._redact_failure_log_request() return + self.model_call_details = _with_prompt_copies_replaced(self.model_call_details, self._unmasked_messages) async def async_mask_failure_log_request( self, @@ -4277,13 +4325,13 @@ class Logging(LiteLLMLoggingBaseClass): mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async()))) raise - def mask_failure_log_request(self) -> None: + def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = executor) -> None: if is_event_loop_running() or not self._claim_failure_log_mask(): return maskers: Final = self._failure_log_maskers() if not maskers: return - masking: Final = executor.submit( + masking: Final = mask_executor.submit( functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details ) try: @@ -4291,14 +4339,17 @@ class Logging(LiteLLMLoggingBaseClass): except concurrent.futures.TimeoutError: self._redact_failure_log_request_after_timeout() return + except concurrent.futures.CancelledError: + self._redact_failure_log_request() + return if masked is None: self._redact_failure_log_request() return - self.model_call_details = dict(masked) + self.model_call_details = _with_prompt_copies_replaced(masked, self._unmasked_messages) def _sync_failure_log_mask_step( - self, details: Mapping[str, Any] | None, guardrail: CustomGuardrail - ) -> Mapping[str, Any] | None: + self, details: Mapping[str, object] | None, guardrail: CustomGuardrail + ) -> Mapping[str, object] | None: if details is None or not _runs_on_logging_only(guardrail, details): return details try: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 7e8cfdb0a94..98087838667 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.thread_pool_executor import BoundedLoggingThreadPoolExecutor 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 @@ -8551,6 +8552,86 @@ async def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_boun assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="full-executor-failure", + function_id="full-executor-failure", + dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + logging_obj.mask_failure_log_request( + mask_executor=BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=0) + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + +class _ProxyRequestRecorder(_FailureLogRecorder): + def __init__(self) -> None: + super().__init__() + self.request_copies: list[tuple[object, object]] = [] + + def _record(self, kwargs: dict[str, object]) -> None: + super()._record(kwargs) + litellm_params: Final = cast(dict[str, object], kwargs["litellm_params"]) + self.request_copies.append( + copy.deepcopy((kwargs["optional_params"], litellm_params["proxy_server_request"])) + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +@pytest.mark.parametrize("sync", [True, False]) +@pytest.mark.parametrize( + ("guardrail_type", "logged_messages"), + [(_NameMaskingGuardrail, _MASKED_LOG), (_RaisingMaskingGuardrail, _REDACTED_LOG)], +) +async def test_failure_log_masks_the_prompt_copies_a_proxy_rejection_logs( + sync: bool, guardrail_type: type[_NameMaskingGuardrail], logged_messages: list[dict[str, str]] +) -> None: + messages: Final = copy.deepcopy(_RAW_LOG) + guardrail: Final = guardrail_type() + recorder: Final = _ProxyRequestRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=messages, + stream=False, + call_type="completion" if sync else "acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="proxy-rejection-copies", + function_id="proxy-rejection-copies", + dynamic_failure_callbacks=[guardrail, recorder], + dynamic_async_failure_callbacks=[guardrail, recorder], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {"url": "/chat/completions", "body": {"messages": messages}}}, + optional_params={"messages": messages}, + ) + logging_obj.model_call_details["input"] = messages + error: Final = litellm.RateLimitError("429", "openai", "gpt-4o-mini") + + if sync: + await asyncio.to_thread(logging_obj.failure_handler, error, "tb") + else: + await logging_obj.async_failure_handler(error, "tb") + + assert len(recorder.request_copies) == 1 + optional_params, proxy_server_request = recorder.request_copies[0] + assert cast(dict[str, object], optional_params)["messages"] == logged_messages + assert _RAW_PROMPT not in str(proxy_server_request) + assert "/chat/completions" in str(proxy_server_request) + assert messages == _RAW_LOG + + @pytest.mark.asyncio async def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: recorder, sync_logged = _register_failure_logging(monkeypatch, _BlockingNameMaskingGuardrail()) From e136bed4a70142cdf0dece7582c66a92492d1261 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:09:43 -0700 Subject: [PATCH 5/8] fix(logging): give the failure spend log the request body with the masked prompt The proxy's failure hooks and spend log read the proxy's own copy of the request body, which the failure-log mask never reached, so a failed call's spend row and the UI Logs page kept the raw prompt. The masked body is now lifted onto that request, and a logged copy counts as the prompt when its messages equal it, since the body keeps its own list. --- litellm/litellm_core_utils/litellm_logging.py | 9 +++- litellm/proxy/utils.py | 12 ++++- .../test_post_call_failure_hook.py | 50 +++++++++++++++++++ 3 files changed, 68 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d746907d05c..c9918ba18f1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -603,7 +603,7 @@ def _with_prompt_copy_replaced( details: Mapping[str, object], path: Sequence[str], raw: object, replacement: object ) -> Mapping[str, object]: containers: Final = tuple(itertools.accumulate(path, _logged_field, initial=details)) - if containers[-1] is not raw: + if containers[-1] != raw: return details rebuilt: Final = functools.reduce(_with_field, reversed(tuple(zip(containers, path))), replacement) return cast(Mapping[str, object], rebuilt) # cast-ok: the outermost rebuilt container is details @@ -4331,6 +4331,13 @@ class Logging(LiteLLMLoggingBaseClass): def messages_sent_to_model(self) -> object: return self._unmasked_messages if self._failure_log_masked else self.model_call_details.get("messages") + def failure_logged_proxy_server_request(self, proxy_server_request: Mapping[str, object]) -> Mapping[str, object]: + if not self._failure_log_masked: + return proxy_server_request + return _with_prompt_copy_replaced( + proxy_server_request, ("body", "messages"), self._unmasked_messages, self.model_call_details.get("messages") + ) + def _claim_failure_log_mask(self) -> bool: if self._failure_log_masked: return False diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ab91bb64954..0f8814bdf01 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1586,12 +1586,19 @@ def _failure_fields_to_lift( request_data (it is not serialisable), so the caller merges these fields onto request_data first: the request start and first-handoff instants for duration and preprocessing latency, the call type, recovered or estimated - usage for token counts, and the standard logging object for deployment - attribution on failed-request spend logs.""" + usage for token counts, the standard logging object for deployment + attribution on failed-request spend logs, and the request body with the + prompt the failure log masked.""" _logging_obj: Final = request_data.get("litellm_logging_obj") if _logging_obj is None: return _EMPTY_LIFT _model_call_details: Final = getattr(_logging_obj, "model_call_details", {}) + _proxy_server_request: Final = request_data.get("proxy_server_request") + _logged_proxy_server_request: Final = ( + _logging_obj.failure_logged_proxy_server_request(_as_str_mapping(_proxy_server_request)) + if isinstance(_logging_obj, Logging) and isinstance(_proxy_server_request, Mapping) + else None + ) _sent_model_call_details: Final = ( {**_model_call_details, "messages": _logging_obj.messages_sent_to_model} if isinstance(_logging_obj, Logging) @@ -1630,6 +1637,7 @@ def _failure_fields_to_lift( ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), ("billed_failed_attempts", _earlier_attempts or None), + ("proxy_server_request", _logged_proxy_server_request), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 1571c29afaf..6102ed1f0b6 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -557,6 +557,56 @@ async def test_proxy_only_failure_log_masks_after_auth_and_redacts_before_it( ] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error_type", "logged_content"), + [(ProxyErrorTypes.auth_error, REDACTED_BY_LITELLM), (None, "My name is ")], +) +async def test_failure_hooks_get_the_request_body_with_the_prompt_the_failure_log_masked( + proxy_logging, monkeypatch: pytest.MonkeyPatch, error_type: ProxyErrorTypes | None, logged_content: str +) -> None: + from litellm.proxy import proxy_server + + received_bodies: Final[list[object]] = [] + + class _SpendLogLikeHook(CustomLogger): + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + received_bodies.append(request_data["proxy_server_request"]) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [_PromptMaskingGuardrail(), _SpendLogLikeHook()]) + proxy_logging.alert_types = [] + messages: Final = [{"role": "user", "content": "My name is John Smith"}] + + await proxy_logging.post_call_failure_hook( + request_data={ + "model": "internal-model", + "messages": messages, + "proxy_server_request": { + "url": "http://localhost/chat/completions", + "body": {"model": "internal-model", "messages": list(messages)}, + }, + }, + original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + error_type=error_type, + route="/chat/completions", + ) + + assert received_bodies == [ + { + "url": "http://localhost/chat/completions", + "body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]}, + } + ] + + @pytest.mark.asyncio async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch From 997d27d27c962936a40c89643f6c2917cdb83ca0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:19:48 -0700 Subject: [PATCH 6/8] fix(logging): keep the masked failure log body across retries and bound the sync mask Keep the masked prompt in the spend log body when a router retry or fallback also fails, since every attempt resets the logging object's messages to the raw prompt. Run the sync failure mask on its own thread pool and cancel it when it is still queued at its time bound. A guardrail with only an async logging hook now redacts the sync failure log instead of passing the raw prompt through. Extend the mid-stream failure log wait by the mask time only when a masking guardrail runs, so deployments without one keep the one second wait. --- litellm/constants.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 36 +++++++--- .../litellm_core_utils/streaming_handler.py | 4 +- .../thread_pool_executor.py | 4 ++ .../test_litellm_logging.py | 63 +++++++++++++++-- .../test_streaming_handler.py | 58 +++++++++++++++- .../test_post_call_failure_hook.py | 69 +++++++++++++++++++ 7 files changed, 214 insertions(+), 22 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 43d40398853..f0aa6e83f0d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -166,7 +166,7 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0 -MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS + 1.0 +MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c9918ba18f1..e037cc8e4ef 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -139,7 +139,7 @@ from litellm.litellm_core_utils.served_output_texts import ( SERVED_OUTPUT_TEXTS_KEY, overlay_served_output_texts, ) -from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.litellm_core_utils.thread_pool_executor import executor, failure_log_mask_executor from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils @@ -676,6 +676,10 @@ def _timestamp_seconds(moment: object) -> float | None: return None +def _masks_synchronously(guardrail: CustomGuardrail) -> bool: + return next(cls for cls in type(guardrail).__mro__ if "logging_hook" in vars(cls)) is not CustomLogger + + def _runs_on_logging_only(guardrail: CustomGuardrail, data: Mapping[str, object]) -> bool: return guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.logging_only) is True @@ -710,6 +714,7 @@ class Logging(LiteLLMLoggingBaseClass): _on_failure_payload_built: "Callable[[], None] | None" = None _failure_log_masked: bool = False _unmasked_messages: object = None + _failure_logged_messages: object = None @property def _defer_async_logging(self) -> bool: @@ -4335,7 +4340,7 @@ class Logging(LiteLLMLoggingBaseClass): if not self._failure_log_masked: return proxy_server_request return _with_prompt_copy_replaced( - proxy_server_request, ("body", "messages"), self._unmasked_messages, self.model_call_details.get("messages") + proxy_server_request, ("body", "messages"), self._unmasked_messages, self._failure_logged_messages ) def _claim_failure_log_mask(self) -> bool: @@ -4343,18 +4348,24 @@ class Logging(LiteLLMLoggingBaseClass): return False self._failure_log_masked = True self._unmasked_messages = self.model_call_details.get("messages") + self._failure_logged_messages = self._unmasked_messages return True + def _set_failure_log_details(self, details: Mapping[str, object]) -> None: + self.model_call_details = _with_prompt_copies_replaced(details, self._unmasked_messages) + self._failure_logged_messages = self.model_call_details.get("messages") + def _redact_failure_log_request(self) -> None: - self.model_call_details = _with_prompt_copies_replaced( - {**self.model_call_details, **redacted_request_entries(self.model_call_details)}, self._unmasked_messages + self._set_failure_log_details({**self.model_call_details, **redacted_request_entries(self.model_call_details)}) + + @property + def masks_failure_log(self) -> bool: + return any( + _runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers() ) def redact_failure_log_request(self) -> None: - masks_this_request: Final = any( - _runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers() - ) - if masks_this_request and self._claim_failure_log_mask(): + if self.masks_failure_log and self._claim_failure_log_mask(): self._redact_failure_log_request() def _report_failure_log_mask_error(self, guardrail: CustomGuardrail) -> None: @@ -4393,7 +4404,7 @@ class Logging(LiteLLMLoggingBaseClass): self._report_failure_log_mask_error(guardrail) self._redact_failure_log_request() return - self.model_call_details = _with_prompt_copies_replaced(self.model_call_details, self._unmasked_messages) + self._set_failure_log_details(self.model_call_details) async def async_mask_failure_log_request( self, @@ -4416,7 +4427,7 @@ class Logging(LiteLLMLoggingBaseClass): mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async()))) raise - def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = executor) -> None: + def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = failure_log_mask_executor) -> None: if is_event_loop_running() or not self._claim_failure_log_mask(): return maskers: Final = self._failure_log_maskers() @@ -4428,6 +4439,7 @@ class Logging(LiteLLMLoggingBaseClass): try: masked: Final = masking.result(timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS) except concurrent.futures.TimeoutError: + _ = masking.cancel() self._redact_failure_log_request_after_timeout() return except concurrent.futures.CancelledError: @@ -4436,13 +4448,15 @@ class Logging(LiteLLMLoggingBaseClass): if masked is None: self._redact_failure_log_request() return - self.model_call_details = _with_prompt_copies_replaced(masked, self._unmasked_messages) + self._set_failure_log_details(masked) def _sync_failure_log_mask_step( self, details: Mapping[str, object] | None, guardrail: CustomGuardrail ) -> Mapping[str, object] | None: if details is None or not _runs_on_logging_only(guardrail, details): return details + if not _masks_synchronously(guardrail): + return None try: hooked: Final = guardrail.logging_hook( kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5542249d2c7..8cdd9192ac3 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( @@ -2477,7 +2478,8 @@ class CustomStreamWrapper: """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" if self._failure_logging is None: return - await asyncio.wait((self._failure_logging,), timeout=timeout) + mask_wait: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS if self.logging_obj.masks_failure_log else 0.0 + await asyncio.wait((self._failure_logging,), timeout=None if timeout is None else timeout + mask_wait) def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: self._earlier_rounds_spend = earlier diff --git a/litellm/litellm_core_utils/thread_pool_executor.py b/litellm/litellm_core_utils/thread_pool_executor.py index f989f20247f..00614769932 100644 --- a/litellm/litellm_core_utils/thread_pool_executor.py +++ b/litellm/litellm_core_utils/thread_pool_executor.py @@ -80,3 +80,7 @@ executor: Final = BoundedLoggingThreadPoolExecutor( max_workers=MAX_THREADS, max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS, ) +failure_log_mask_executor: Final = BoundedLoggingThreadPoolExecutor( + max_workers=MAX_THREADS, + max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS, +) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 423f3ad906b..e95b417ba01 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8515,6 +8515,15 @@ class _HangingSyncMaskingGuardrail(_NameMaskingGuardrail): return self._mask(kwargs), result +class _AsyncOnlyNameMaskingGuardrail(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return await _NameMaskingGuardrail().async_logging_hook(kwargs, result, call_type) + + class _FailureLogRecorder(CustomLogger): def __init__(self) -> None: super().__init__() @@ -8662,20 +8671,25 @@ async def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_boun assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] -@pytest.mark.usefixtures("no_global_callbacks") -def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None: - recorder: Final = _FailureLogRecorder() +def _sync_failure_logging(call_id: str, guardrail: CustomGuardrail, recorder: _FailureLogRecorder) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), stream=False, call_type="completion", start_time=datetime.datetime.now(), - litellm_call_id="full-executor-failure", - function_id="full-executor-failure", - dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], + litellm_call_id=call_id, + function_id=call_id, + dynamic_failure_callbacks=[guardrail, recorder], ) logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging("full-executor-failure", _NameMaskingGuardrail(), recorder) logging_obj.mask_failure_log_request( mask_executor=BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=0) @@ -8685,6 +8699,43 @@ def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mas assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_is_logged_redacted_when_the_masking_guardrail_has_only_an_async_hook() -> None: + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging( + "async-only-mask-failure", + _AsyncOnlyNameMaskingGuardrail( + guardrail_name="async-only-name-masker", event_hook=GuardrailEventHooks.logging_only + ), + recorder, + ) + + logging_obj.mask_failure_log_request() + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_mask_still_queued_at_its_time_bound_never_runs() -> None: + guardrail: Final = _NameMaskingGuardrail() + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging("queued-mask-failure", guardrail, recorder) + mask_executor: Final = BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2) + busy_worker_released: Final = threading.Event() + _ = mask_executor.submit(busy_worker_released.wait, 30) + + try: + logging_obj.mask_failure_log_request(mask_executor=mask_executor) + finally: + busy_worker_released.set() + mask_executor.shutdown(wait=True) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert guardrail.scans == 0 + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + class _ProxyRequestRecorder(_FailureLogRecorder): def __init__(self) -> None: super().__init__() diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index f4502937fce..25ee7d317eb 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -9,7 +9,7 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import CodeType, FrameType, MappingProxyType -from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast +from typing import AsyncIterator, Callable, ClassVar, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -20,7 +20,9 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var +from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging @@ -32,6 +34,7 @@ from litellm.litellm_core_utils.streaming_handler import ( from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, @@ -3124,7 +3127,9 @@ def _stuck_stream_wrapper( provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], call_type: Literal["completion", "acompletion"], recorder: _StuckStreamFailureRecorder, + masker: CustomGuardrail | None = None, ) -> CustomStreamWrapper: + failure_callbacks: Final[list[CustomLogger]] = [recorder] if masker is None else [masker, recorder] logging_obj: Final = Logging( model="test-model", messages=[{"role": "user", "content": "hi"}], @@ -3133,8 +3138,8 @@ def _stuck_stream_wrapper( start_time=datetime.datetime(2026, 1, 1), litellm_call_id=f"stuck-stream-{call_type}", function_id=f"stuck-stream-{call_type}", - dynamic_failure_callbacks=[recorder], - dynamic_async_failure_callbacks=[recorder], + dynamic_failure_callbacks=failure_callbacks, + dynamic_async_failure_callbacks=failure_callbacks, ) logging_obj.update_environment_variables( model="test-model", @@ -3149,6 +3154,53 @@ def _stuck_stream_wrapper( ) +class _BlockedStuckStreamFailureRecorder(_StuckStreamFailureRecorder): + def __init__(self, release: asyncio.Event) -> None: + super().__init__() + self.release = release + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + await self.release.wait() + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + + +class _FailureLogMasker(CustomGuardrail): + masks_logged_request: ClassVar[bool] = True + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return kwargs, result + + +async def _seconds_waited_for_a_blocked_failure_log(masker: CustomGuardrail | None) -> float: + release: Final = asyncio.Event() + wrapper: Final = _stuck_stream_wrapper( + _StuckProvider().async_stream(), "acompletion", _BlockedStuckStreamFailureRecorder(release), masker + ) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + started: Final = time.monotonic() + await wrapper.wait_for_failure_logging(timeout=0.1) + waited: Final = time.monotonic() - started + release.set() + await wrapper.wait_for_failure_logging(timeout=None) + return waited + + +@pytest.mark.asyncio +async def test_mid_stream_failure_wait_extends_by_the_mask_time_only_when_the_failure_log_is_masked() -> None: + _ = await _seconds_waited_for_a_blocked_failure_log(None) + + assert await _seconds_waited_for_a_blocked_failure_log(None) < FAILURE_LOG_MASK_TIMEOUT_SECONDS + assert await _seconds_waited_for_a_blocked_failure_log( + _FailureLogMasker(guardrail_name="failure-log-masker", event_hook=GuardrailEventHooks.logging_only) + ) >= FAILURE_LOG_MASK_TIMEOUT_SECONDS + + @pytest.mark.asyncio async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 6102ed1f0b6..1322f711f60 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -607,6 +607,75 @@ async def test_failure_hooks_get_the_request_body_with_the_prompt_the_failure_lo ] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("maskers", "logged_content"), + [((), "My name is John Smith"), ((_PromptMaskingGuardrail(),), "My name is ")], + ids=["without_a_masker", "with_a_masker"], +) +async def test_failure_hooks_get_the_request_body_the_failure_log_carried_after_the_router_retried_the_call( + proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, maskers: tuple[CustomGuardrail, ...], logged_content: str +) -> None: + from litellm.proxy import proxy_server + + received_bodies: Final[list[object]] = [] + + class _SpendLogLikeHook(CustomLogger): + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + received_bodies.append(request_data["proxy_server_request"]) + + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "callbacks", [*maskers, _SpendLogLikeHook()]) + proxy_logging.alert_types = [] + messages: Final = [{"role": "user", "content": "My name is John Smith"}] + router: Final = litellm.Router( + model_list=[ + { + "model_name": "internal-model", + "litellm_params": { + "model": "openai/gpt-4.1", + "api_key": "sk-test", + "mock_response": "litellm.RateLimitError", + }, + } + ], + num_retries=1, + retry_after=0, + ) + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="internal-model", + messages=messages, + litellm_call_id="retried-call", + ) + with pytest.raises(litellm.RateLimitError): + await router.acompletion(model="internal-model", messages=messages, litellm_logging_obj=logging_obj) + + await proxy_logging.post_call_failure_hook( + request_data={ + "model": "internal-model", + "messages": messages, + "litellm_logging_obj": logging_obj, + "proxy_server_request": {"body": {"model": "internal-model", "messages": messages}}, + }, + original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), + user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"), + route="/chat/completions", + ) + + assert received_bodies == [ + {"body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]}} + ] + + @pytest.mark.asyncio async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment( proxy_logging, make_user_api_key_auth, monkeypatch From 4042a605af577f1505773e8b20ea3d49f2f768d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:36:53 -0700 Subject: [PATCH 7/8] fix(logging): keep the fallback's deployment in success and spend logs when a rejected attempt's failure log is masked --- litellm/litellm_core_utils/litellm_logging.py | 65 ++++++++++++------- .../test_litellm_logging.py | 57 ++++++++++++++++ 2 files changed, 99 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e037cc8e4ef..ef59f0f3289 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -577,7 +577,8 @@ _ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set() # mutab _ABSENT: Final = object() -_LOGGED_PROMPT_COPY_PATHS: Final = ( +_LOGGED_PROMPT_PATHS: Final = ( + ("messages",), ("input",), ("optional_params", "messages"), ("litellm_params", "proxy_server_request", "body", "messages"), @@ -610,15 +611,14 @@ def _with_prompt_copy_replaced( def _with_prompt_copies_replaced( - details: Mapping[str, object], raw: object + details: Mapping[str, object], raw: object, replacement: object ) -> dict[str, object]: # mutable-ok: becomes the Logging object's model_call_details if raw is None: return dict(details) - replacement: Final = details.get("messages") return dict( functools.reduce( lambda fields, path: _with_prompt_copy_replaced(fields, path, raw, replacement), - _LOGGED_PROMPT_COPY_PATHS, + _LOGGED_PROMPT_PATHS, details, ) ) @@ -4351,12 +4351,13 @@ class Logging(LiteLLMLoggingBaseClass): self._failure_logged_messages = self._unmasked_messages return True - def _set_failure_log_details(self, details: Mapping[str, object]) -> None: - self.model_call_details = _with_prompt_copies_replaced(details, self._unmasked_messages) - self._failure_logged_messages = self.model_call_details.get("messages") + def _log_failure_with_messages(self, details: Mapping[str, object], logged_messages: object) -> None: + self.model_call_details = _with_prompt_copies_replaced(details, self._unmasked_messages, logged_messages) + self._failure_logged_messages = logged_messages def _redact_failure_log_request(self) -> None: - self._set_failure_log_details({**self.model_call_details, **redacted_request_entries(self.model_call_details)}) + redacted: Final = {**self.model_call_details, **redacted_request_entries(self.model_call_details)} + self._log_failure_with_messages(redacted, redacted["messages"]) @property def masks_failure_log(self) -> bool: @@ -4391,20 +4392,38 @@ class Logging(LiteLLMLoggingBaseClass): self._redact_failure_log_request_after_timeout() async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None: - for guardrail in maskers: - try: - if not _runs_on_logging_only(guardrail, self.model_call_details): - continue - self.model_call_details, _ = await guardrail.async_logging_hook( - 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 redacts the log instead of stopping failure logging - self._report_failure_log_mask_error(guardrail) - self._redact_failure_log_request() - return - self._set_failure_log_details(self.model_call_details) + masked: Final = await self._masked_failure_log_details(self.model_call_details, maskers) + if masked is None: + self._redact_failure_log_request() + return + self._log_failure_with_messages(self.model_call_details, masked.get("messages")) + + async def _masked_failure_log_details( + self, details: Mapping[str, object] | None, maskers: tuple[CustomGuardrail, ...] + ) -> Mapping[str, object] | None: + if details is None or not maskers: + return details + return await self._masked_failure_log_details( + await self._async_failure_log_mask_step(details, maskers[0]), maskers[1:] + ) + + async def _async_failure_log_mask_step( + self, details: Mapping[str, object], guardrail: CustomGuardrail + ) -> Mapping[str, object] | None: + if not _runs_on_logging_only(guardrail, details): + return details + try: + hooked_details: Final = _METADATA_MAPPING.validate_python( + ( + await guardrail.async_logging_hook( + kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type + ) + )[0] + ) + 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_details async def async_mask_failure_log_request( self, @@ -4448,7 +4467,7 @@ class Logging(LiteLLMLoggingBaseClass): if masked is None: self._redact_failure_log_request() return - self._set_failure_log_details(masked) + self._log_failure_with_messages(self.model_call_details, masked.get("messages")) def _sync_failure_log_mask_step( self, details: Mapping[str, object] | None, guardrail: CustomGuardrail diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e95b417ba01..985e4afe0ed 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -8830,6 +8830,63 @@ async def test_concurrent_failure_notifications_log_one_masked_failure() -> None assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)] +class _DeploymentSuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: list[tuple[str, str, str]] = [] + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.append((payload["model_id"] or "", payload["model_group"] or "", payload["api_base"])) + + +def _deployment_params(model_id: str, model_group: str) -> dict[str, object]: + return { + "acompletion": True, + "api_base": f"https://{model_group}.example.com", + "metadata": {"model_group": model_group, "model_info": {"id": model_id}}, + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_fallback_started_while_the_failure_log_is_masked_logs_its_own_deployment() -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + failures: Final = _FailureLogRecorder() + successes: Final = _DeploymentSuccessRecorder() + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=copy.deepcopy(_RAW_LOG), + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="fallback-during-mask", + function_id="fallback-during-mask", + dynamic_async_failure_callbacks=[guardrail, failures], + dynamic_async_success_callbacks=[successes], + ) + logging_obj.update_environment_variables(litellm_params=_deployment_params("rejected", "primary"), optional_params={}) + rejected_attempt_log: Final = asyncio.create_task( + logging_obj.async_failure_handler(litellm.RateLimitError("429", "openai", "gpt-4o-mini"), "tb") + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + + logging_obj.update_environment_variables(litellm_params=_deployment_params("served", "backup"), optional_params={}) + guardrail.release.set() + await asyncio.wait_for(rejected_attempt_log, timeout=10) + await logging_obj.async_success_handler( + result=ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + ) + ) + + assert failures.logged == [(_MASKED_LOG, _MASKED_LOG)] + assert successes.logged == [("served", "backup", "https://backup.example.com")] + + @pytest.mark.asyncio @pytest.mark.usefixtures("no_global_callbacks") async def test_sync_sdk_failure_dispatched_from_async_code_logs_the_masked_prompt() -> None: From a247d9e553b034da04c293f22824d4cb0f41bf22 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:07:38 -0700 Subject: [PATCH 8/8] test(guardrails): run the failure-log time-bound tests on controlled clocks --- .../test_litellm_logging.py | 132 +++++++++++++----- .../test_streaming_handler.py | 53 +++++-- .../test_post_call_failure_hook.py | 2 +- 3 files changed, 143 insertions(+), 44 deletions(-) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index ad22ea9f3ab..e1a8e9f7227 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 concurrent.futures import copy import datetime import importlib.abc @@ -10,11 +11,11 @@ import sys import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object from importlib.machinery import ModuleSpec from types import MappingProxyType, ModuleType -from typing import ClassVar, Final, Literal, cast +from typing import ClassVar, Final, Literal, ParamSpec, TypeVar, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -31,6 +32,7 @@ from litellm._logging import session_id_var, trace_id_var, verbose_logger from litellm._service_logger import ServiceLogging from litellm.caching.caching import DualCache from litellm.constants import ( + FAILURE_LOG_MASK_TIMEOUT_SECONDS, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, LOGGING_WORKER_MAX_TIME_PER_COROUTINE, @@ -8800,6 +8802,58 @@ async def _until(condition: Callable[[], bool]) -> None: await asyncio.sleep(0) +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock(scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, None]]) -> None: + loop: Final = _SkippableClockLoop() + try: + loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +class _OutlastedTimeBoundFuture(concurrent.futures.Future[_T]): + def __init__(self, work: concurrent.futures.Future[_T], on_wait: Callable[[float | None], None]) -> None: + super().__init__() + self._work: Final = work + self._on_wait: Final = on_wait + + def result(self, timeout: float | None = None) -> _T: + self._on_wait(timeout) + raise concurrent.futures.TimeoutError + + def cancel(self) -> bool: + return self._work.cancel() + + +class _TimeBoundOutlastingExecutor(concurrent.futures.Executor): + def __init__(self, workers: concurrent.futures.Executor) -> None: + self._workers: Final = workers + self.waits: tuple[float | None, ...] = () + + def _record_wait(self, timeout: float | None) -> None: + self.waits = (*self.waits, timeout) + + def submit(self, fn: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> concurrent.futures.Future[_T]: + return _OutlastedTimeBoundFuture(self._workers.submit(fn, *args, **kwargs), self._record_wait) + + def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None: + self._workers.shutdown(wait=wait, cancel_futures=cancel_futures) + + async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict[str, str]]) -> None: error: Final = _bad_key_error() request: Final = ( @@ -8884,29 +8938,13 @@ async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mi 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)] - - def _sync_failure_logging(call_id: str, guardrail: CustomGuardrail, recorder: _FailureLogRecorder) -> LitellmLogging: logging_obj: Final = LitellmLogging( model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), stream=False, call_type="completion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id=call_id, function_id=call_id, dynamic_failure_callbacks=[guardrail, recorder], @@ -8915,6 +8953,26 @@ def _sync_failure_logging(call_id: str, guardrail: CustomGuardrail, recorder: _F return logging_obj +@pytest.mark.usefixtures("no_global_callbacks") +def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound() -> None: + guardrail: Final = _HangingSyncMaskingGuardrail() + recorder: Final = _FailureLogRecorder() + logging_obj: Final = _sync_failure_logging("hanging-mask-failure", guardrail, recorder) + mask_executor: Final = _TimeBoundOutlastingExecutor( + BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=1) + ) + + try: + logging_obj.mask_failure_log_request(mask_executor=mask_executor) + finally: + guardrail.release.set() + mask_executor.shutdown(wait=True) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb") + + assert mask_executor.waits == (FAILURE_LOG_MASK_TIMEOUT_SECONDS,) + assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] + + @pytest.mark.usefixtures("no_global_callbacks") def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None: recorder: Final = _FailureLogRecorder() @@ -8950,7 +9008,9 @@ def test_sync_failure_mask_still_queued_at_its_time_bound_never_runs() -> None: guardrail: Final = _NameMaskingGuardrail() recorder: Final = _FailureLogRecorder() logging_obj: Final = _sync_failure_logging("queued-mask-failure", guardrail, recorder) - mask_executor: Final = BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2) + mask_executor: Final = _TimeBoundOutlastingExecutor( + BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2) + ) busy_worker_released: Final = threading.Event() _ = mask_executor.submit(busy_worker_released.wait, 30) @@ -8996,7 +9056,7 @@ async def test_failure_log_masks_the_prompt_copies_a_proxy_rejection_logs( messages=messages, stream=False, call_type="completion" if sync else "acompletion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="proxy-rejection-copies", function_id="proxy-rejection-copies", dynamic_failure_callbacks=[guardrail, recorder], @@ -9022,17 +9082,25 @@ async def test_failure_log_masks_the_prompt_copies_a_proxy_rejection_logs( assert messages == _RAW_LOG -@pytest.mark.asyncio -async def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: - recorder, sync_logged = _register_failure_logging(monkeypatch, _BlockingNameMaskingGuardrail()) - +async def _fail_once_the_mask_outlasts_its_time_bound( + clock: _SkippableClockLoop, guardrail: _BlockingNameMaskingGuardrail, recorder: _FailureLogRecorder +) -> None: + call: Final = asyncio.create_task( + litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()) + ) + await asyncio.wait_for(guardrail.scanning.wait(), timeout=10) + clock.skip(FAILURE_LOG_MASK_TIMEOUT_SECONDS) with pytest.raises(litellm.AuthenticationError): - await asyncio.wait_for( - litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error()), - timeout=10, - ) + await asyncio.wait_for(call, timeout=10) await asyncio.wait_for(_until(lambda: bool(recorder.logged)), timeout=10) + +def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _BlockingNameMaskingGuardrail() + recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail) + + _run_on_skippable_clock(lambda clock: _fail_once_the_mask_outlasts_its_time_bound(clock, guardrail, recorder)) + assert sync_logged == [_REDACTED_LOG] assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)] @@ -9046,7 +9114,7 @@ async def test_concurrent_failure_notifications_log_one_masked_failure() -> None messages=copy.deepcopy(_RAW_LOG), stream=False, call_type="acompletion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="concurrent-masked-failure", function_id="concurrent-masked-failure", dynamic_async_failure_callbacks=[_NameMaskingGuardrail(), recorder], @@ -9090,7 +9158,7 @@ async def test_fallback_started_while_the_failure_log_is_masked_logs_its_own_dep messages=copy.deepcopy(_RAW_LOG), stream=False, call_type="acompletion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="fallback-during-mask", function_id="fallback-during-mask", dynamic_async_failure_callbacks=[guardrail, failures], @@ -9125,7 +9193,7 @@ async def test_sync_sdk_failure_dispatched_from_async_code_logs_the_masked_promp messages=copy.deepcopy(_RAW_LOG), stream=True, call_type="completion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="sync-sdk-masked-failure", function_id="sync-sdk-masked-failure", dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder], diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b0f51e577ab..7a75c6f6315 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -4,11 +4,12 @@ import itertools import json import os import queue +import selectors import subprocess import sys import threading import time -from collections.abc import Mapping, Sequence +from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType @@ -3529,6 +3530,36 @@ class _FailureLogMasker(CustomGuardrail): return kwargs, result +class _VirtualClockSelector(selectors.DefaultSelector): + now: float = 0.0 + + def select(self, timeout: float | None = None) -> list[tuple[selectors.SelectorKey, int]]: + if timeout is None: + return super().select(None) + ready: Final = super().select(0) + if not ready: + self.now += timeout + return ready + + +class _VirtualClockLoop(asyncio.SelectorEventLoop): + def __init__(self) -> None: + self.clock: Final = _VirtualClockSelector() + super().__init__(self.clock) + + def time(self) -> float: + return self.clock.now + + +def _on_virtual_clock(scenario: Coroutine[object, object, float]) -> float: + loop: Final = _VirtualClockLoop() + try: + return loop.run_until_complete(scenario) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + async def _seconds_waited_for_a_blocked_failure_log(masker: CustomGuardrail | None) -> float: release: Final = asyncio.Event() wrapper: Final = _stuck_stream_wrapper( @@ -3537,22 +3568,22 @@ async def _seconds_waited_for_a_blocked_failure_log(masker: CustomGuardrail | No with pytest.raises(Exception, match="repeating the same chunk"): async for _ in wrapper: pass - started: Final = time.monotonic() + clock: Final = asyncio.get_running_loop() + started: Final = clock.time() await wrapper.wait_for_failure_logging(timeout=0.1) - waited: Final = time.monotonic() - started + waited: Final = clock.time() - started release.set() await wrapper.wait_for_failure_logging(timeout=None) return waited -@pytest.mark.asyncio -async def test_mid_stream_failure_wait_extends_by_the_mask_time_only_when_the_failure_log_is_masked() -> None: - _ = await _seconds_waited_for_a_blocked_failure_log(None) - - assert await _seconds_waited_for_a_blocked_failure_log(None) < FAILURE_LOG_MASK_TIMEOUT_SECONDS - assert await _seconds_waited_for_a_blocked_failure_log( - _FailureLogMasker(guardrail_name="failure-log-masker", event_hook=GuardrailEventHooks.logging_only) - ) >= FAILURE_LOG_MASK_TIMEOUT_SECONDS +def test_mid_stream_failure_wait_extends_by_the_mask_time_only_when_the_failure_log_is_masked() -> None: + assert _on_virtual_clock(_seconds_waited_for_a_blocked_failure_log(None)) == pytest.approx(0.1) + assert _on_virtual_clock( + _seconds_waited_for_a_blocked_failure_log( + _FailureLogMasker(guardrail_name="failure-log-masker", event_hook=GuardrailEventHooks.logging_only) + ) + ) == pytest.approx(0.1 + FAILURE_LOG_MASK_TIMEOUT_SECONDS) @pytest.mark.asyncio diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 635b4611b73..72b24189116 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -652,7 +652,7 @@ async def test_failure_hooks_get_the_request_body_the_failure_log_carried_after_ logging_obj, _ = litellm.utils.function_setup( original_function="acompletion", rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), model="internal-model", messages=messages, litellm_call_id="retried-call",