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