fix(proxy): scrub the prompt out of proxy failure hook errors when message logging is redacted

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

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:

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 {}

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"

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

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"