From e71a1bab703d5c2059868315372dc6367189c457 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 18:39:25 -0700 Subject: [PATCH] fix(proxy): scrub the prompt out of proxy failure hook errors when message logging is redacted --- litellm/integrations/datadog/datadog.py | 14 ++-- litellm/integrations/otel/logger.py | 27 ++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 16 +++-- .../datadog/test_datadog_logger_batching.py | 25 +++++++ .../integrations/otel/test_otel_v2_logger.py | 69 +++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 27 ++++++++ 6 files changed, 162 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 6518cda5ae3..3c569ff4fe9 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -18,7 +18,7 @@ import datetime import os import time import traceback -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime as datetimeObj from typing import TYPE_CHECKING, Any, Final @@ -287,7 +287,7 @@ class DataDogLogger( async def async_post_call_failure_hook( self, - request_data: dict, + request_data: Mapping[str, object], original_exception: Exception, user_api_key_dict: "UserAPIKeyAuth", traceback_str: str | None = None, @@ -303,11 +303,15 @@ class DataDogLogger( from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, ) + from litellm.litellm_core_utils.redact_messages import scrubbed_error_information_for_request from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - error_information: Final = StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - traceback_str=traceback_str, + error_information: Final = scrubbed_error_information_for_request( + StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + traceback_str=traceback_str, + ), + request_data, ) _code: Final = error_information.get("error_code") or "" status_code: int | None = None diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index b5b65d299e0..6496987c1bd 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -1,6 +1,7 @@ """``CustomLogger`` adapter on the OpenTelemetry span engine.""" import sys +import weakref from collections import OrderedDict from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext @@ -106,18 +107,22 @@ def _span_error_from_exception( *, status_code: int | None = None, traceback_str: str | None = None, + request_data: Mapping[str, object] | None = None, ) -> SpanError: """A ``SpanError`` for a proxy-level failure that never produced a ``StandardLoggingPayload`` (auth / validation / malformed-body rejections), mirroring ``_parse_error``'s field mapping so it stamps the same v2 keys a failed LLM call does. ``status_code`` pins ``error.code`` to the real response - status, matching v1's SERVER-span behavior.""" + status, matching v1's SERVER-span behavior. ``request_data`` scrubs the + request's own text out of the error when its message logging is redacted.""" from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.litellm_core_utils.redact_messages import scrubbed_error_information_for_request - info: Final = StandardLoggingPayloadSetup.get_error_information( + raw_info: Final = StandardLoggingPayloadSetup.get_error_information( original_exception=exception, traceback_str=traceback_str, ) + info: Final = raw_info if request_data is None else scrubbed_error_information_for_request(raw_info, request_data) return SpanError( error_type=info.get("error_class") or info.get("error_code") or None, message=info.get("error_message") or None, @@ -127,6 +132,9 @@ def _span_error_from_exception( ) +_failed_request_by_span: Final[weakref.WeakKeyDictionary[Span, Mapping[str, object]]] = weakref.WeakKeyDictionary() + + # Any callback whose class belongs to one of these modules is "the OTel # callback" for proxy-global-registration purposes. _OTEL_MODULES: Final = ( @@ -843,19 +851,22 @@ class OpenTelemetryV2(CustomLogger): this span already — ``async_post_call_failure_hook`` and the ``auth`` phase span record their own, and a second event would duplicate it — while the attributes are always restamped so ``error.code`` stays pinned to the real - response status.""" + response status. The restamp scrubs the text of the request the failure hook + stamped this span for, as the hook did.""" if span is None or not is_recordable_span(span): return already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ()) stamp_error( span, - _span_error_from_exception(exception, status_code=status_code), + _span_error_from_exception( + exception, status_code=status_code, request_data=_failed_request_by_span.get(span) + ), record_event=not already_stamped, ) async def async_post_call_failure_hook( self, - request_data: dict, + request_data: Mapping[str, object], original_exception: Exception, user_api_key_dict: "UserAPIKeyAuth", traceback_str: "str | None" = None, @@ -875,7 +886,11 @@ class OpenTelemetryV2(CustomLogger): span: Final = mcp_message_transport_span() or request_root_span() or user_api_key_dict.parent_otel_span if span is None or not is_recordable_span(span): return - stamp_error(span, _span_error_from_exception(original_exception, traceback_str=traceback_str)) + _failed_request_by_span[span] = request_data + stamp_error( + span, + _span_error_from_exception(original_exception, traceback_str=traceback_str, request_data=request_data), + ) return def emit_guardrail_span(self, entry: "StandardLoggingGuardrailInformation") -> None: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 90887422567..7a02ec86283 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -28,7 +28,10 @@ from litellm.litellm_core_utils.litellm_logging import ( UsageObservation, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost -from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS +from litellm.litellm_core_utils.redact_messages import ( + MESSAGE_REDACTION_ENABLE_HEADERS, + scrubbed_error_information_for_request, +) from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( @@ -488,9 +491,13 @@ class _ProxyDBLogger(CustomLogger): ) _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) _metadata["status"] = "failure" - _error_information = StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - traceback_str=traceback_str, + request_fields: Final = _str_keyed_fields(request_data) + _error_information = scrubbed_error_information_for_request( + StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + traceback_str=traceback_str, + ), + request_fields, ) if should_suppress_spend_log_tracebacks(): # Drop the traceback key entirely so the per-row Metadata pane in @@ -507,7 +514,6 @@ class _ProxyDBLogger(CustomLogger): ) ) - request_fields: Final = _str_keyed_fields(request_data) _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} diff --git a/tests/unit/integrations/datadog/test_datadog_logger_batching.py b/tests/unit/integrations/datadog/test_datadog_logger_batching.py index e2707d321bf..512d30d7fa6 100644 --- a/tests/unit/integrations/datadog/test_datadog_logger_batching.py +++ b/tests/unit/integrations/datadog/test_datadog_logger_batching.py @@ -1,4 +1,5 @@ import asyncio +import json from unittest.mock import AsyncMock, Mock, patch import httpx @@ -6,6 +7,7 @@ import pytest from httpx import Request, Response from pydantic import BaseModel, computed_field +import litellm from litellm.integrations.datadog.datadog import DataDogLogger from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError from litellm.types.integrations.datadog import ( @@ -558,3 +560,26 @@ async def test_raised_intake_error_preserves_datadog_requeue_behavior(datadog_en await logger.async_send_batch() assert [event["message"] for event in logger.log_queue] == ['{"event": 0}', '{"event": 1}'] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"]) +async def test_failure_hook_scrubs_the_prompt_out_of_the_logged_exception(datadog_env, monkeypatch, redact): + secret = "my ssn is 123-45-6789" + monkeypatch.setattr(litellm, "turn_off_message_logging", redact) + with patch("asyncio.create_task"): + logger = DataDogLogger() + + await logger.async_post_call_failure_hook( + request_data={"messages": [{"role": "user", "content": secret}], "metadata": {}}, + original_exception=litellm.BadRequestError( + message=f"Unsupported content: {secret}", model="gpt-4o", llm_provider="openai" + ), + user_api_key_dict=type("UserKey", (), {})(), + traceback_str=f"Traceback: {secret}", + ) + + message = json.loads(logger.log_queue[0]["message"]) + assert (secret in json.dumps(message)) is not redact + assert "Unsupported content: " in message["exception"] + assert message["error_class"] == "BadRequestError" diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index ac1f5557bf6..5667cb4c6a5 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -1553,6 +1553,75 @@ def test_record_error_attributes_on_span_decorates_without_ending(): assert [e.name for e in span.events] == ["exception"] +@pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"]) +def test_failure_hook_scrubs_the_prompt_out_of_the_server_span_and_the_restamp_keeps_it(monkeypatch, redact): + import json + + import litellm + from litellm.proxy._types import UserAPIKeyAuth + + secret = "my ssn is 123-45-6789" + monkeypatch.setattr(litellm, "turn_off_message_logging", redact) + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc(f"blocked: {secret}", 400) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={"messages": [{"role": "user", "content": secret}]}, + original_exception=exc, + user_api_key_dict=UserAPIKeyAuth(), + ) + ) + logger.record_error_attributes_on_span(server, exc, 400) + server.end() + (span,) = exporter.get_finished_spans() + logged = json.dumps( + [dict(span.attributes), [dict(event.attributes or {}) for event in span.events], span.status.description] + ) + assert (secret in logged) is not redact + assert span.attributes["error.message"].startswith("blocked: ") + assert span.attributes["error.type"] == "ProxyException" + + +def _raised(exc): + try: + raise exc + except Exception as raised: + return raised + + +@pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"]) +def test_restamp_shows_the_handlers_own_error_with_the_prompt_scrubbed(monkeypatch, redact): + """The proxy's failure hook and its exception handler stamp the same SERVER span from different + exceptions. The handler's restamp keeps its own wording and stack trace, scrubbed like the hook's.""" + from fastapi import HTTPException + + import litellm + from litellm.proxy._types import UserAPIKeyAuth + + secret = "my ssn is 123-45-6789" + monkeypatch.setattr(litellm, "turn_off_message_logging", redact) + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={"messages": [{"role": "user", "content": secret}]}, + original_exception=HTTPException(status_code=400, detail={"error": f"blocked: {secret}", "mode": "pre"}), + user_api_key_dict=UserAPIKeyAuth(), + traceback_str=f' File "/srv/litellm/router.py", line 7, in route\nHTTPException: blocked: {secret}\n', + ) + ) + logger.record_error_attributes_on_span(server, _raised(_proxy_exc(f"blocked: {secret}", 500)), 500) + server.end() + (span,) = exporter.get_finished_spans() + + assert span.attributes["error.message"] == f"blocked: {'redacted-by-litellm' if redact else secret}" + assert "_raised" in span.attributes["litellm.provider.error.stack_trace"] + assert "router.py" not in span.attributes["litellm.provider.error.stack_trace"] + + def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error(): """A failure that already went through ``async_post_call_failure_hook`` reaches the exception handler too; the second stamp must keep one exception event while diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c6366a40504..31e8bb10f67 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4056,3 +4056,30 @@ async def test_a_poll_of_an_already_charged_job_still_charges_its_own_guardrails assert await _charge_request_cost(charge_spend, dict(poll), 0.43, reservation) == charged assert charge_spend.calls == ((0.43, 0.43, reservation), *(() if second_call is None else (second_call,))) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"]) +async def test_async_post_call_failure_hook_scrubs_the_prompt_out_of_the_persisted_error( + monkeypatch: pytest.MonkeyPatch, redact: bool +): + secret: Final = "my ssn is 123-45-6789" + monkeypatch.setattr(litellm, "turn_off_message_logging", redact) + spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) + + await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook( + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": secret}], "metadata": {}}, + original_exception=litellm.BadRequestError( + message=f"Unsupported content: {secret}", model="gpt-4o", llm_provider="openai" + ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + traceback_str=f"Traceback: {secret}", + ) + + persisted: Final = spend_writer.update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"][ + "error_information" + ] + assert (secret in json.dumps(persisted)) is not redact + assert "Unsupported content: " in persisted["error_message"] + assert persisted["error_class"] == "BadRequestError" + assert persisted["error_code"] == "400"