fix(otel): keep only the quoted request text for the restamp, drop it once read, and stamp spans the map cannot hold

litellm/integrations/otel/logger.py

@@ -132,7 +132,43 @@ def _span_error_from_exception(
     )
 
 
-_scrub_texts_by_span: Final[weakref.WeakKeyDictionary[Span, frozenset[str] | None]] = weakref.WeakKeyDictionary()
+_scrub_texts_by_span: Final[weakref.WeakKeyDictionary[Span, frozenset[str]]] = weakref.WeakKeyDictionary()
+
+
+def _keep_scrub_texts_for_restamp(
+    span: Span, exception: Exception, traceback_str: str | None, scrub_texts: frozenset[str] | None
+) -> None:
+    """Keep for the exception handler's restamp of ``span`` the request text that ``exception``'s error quotes, and
+    nothing for a failure that is not redacted or quotes none of it. A span the map cannot hold keeps nothing."""
+    quoted: Final[frozenset[str]] = (
+        frozenset() if scrub_texts is None else _quoted_request_texts(exception, traceback_str, scrub_texts)
+    )
+    try:
+        if quoted:
+            _scrub_texts_by_span[span] = quoted
+        else:
+            _scrub_texts_by_span.pop(span, None)
+    except TypeError:
+        return
+
+
+def _quoted_request_texts(exception: Exception, traceback_str: str | None, texts: frozenset[str]) -> frozenset[str]:
+    """The request ``texts`` that the error logged for ``exception`` quotes."""
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+    from litellm.litellm_core_utils.redact_messages import quoted_request_texts
+
+    info: Final = StandardLoggingPayloadSetup.get_error_information(
+        original_exception=exception, traceback_str=traceback_str
+    )
+    return quoted_request_texts(f"{info.get('error_message') or ''}\n{info.get('traceback') or ''}", texts)
+
+
+def _restamp_scrub_texts(span: Span) -> frozenset[str] | None:
+    """The request text the failure hook kept for ``span``'s restamp, forgotten once read."""
+    try:
+        return _scrub_texts_by_span.pop(span, None)
+    except TypeError:
+        return None
 
 
 # Any callback whose class belongs to one of these modules is "the OTel
@@ -858,7 +894,7 @@ class OpenTelemetryV2(CustomLogger):
         already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ())
         stamp_error(
             span,
-            _span_error_from_exception(exception, status_code=status_code, scrub_texts=_scrub_texts_by_span.get(span)),
+            _span_error_from_exception(exception, status_code=status_code, scrub_texts=_restamp_scrub_texts(span)),
             record_event=not already_stamped,
         )
 
@@ -887,7 +923,7 @@ class OpenTelemetryV2(CustomLogger):
         from litellm.litellm_core_utils.redact_messages import failure_request_texts, should_redact_failed_request
 
         scrub_texts: Final = failure_request_texts(request_data) if should_redact_failed_request(request_data) else None
-        _scrub_texts_by_span[span] = scrub_texts
+        _keep_scrub_texts_for_restamp(span, original_exception, traceback_str, scrub_texts)
         stamp_error(
             span,
             _span_error_from_exception(original_exception, traceback_str=traceback_str, scrub_texts=scrub_texts),

tests/unit/integrations/otel/test_otel_v2_logger.py

@@ -13,6 +13,7 @@ import contextvars
 import gc
 import os
 import re
+import sys
 import weakref
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
@@ -1650,6 +1651,97 @@ def test_failure_hook_lets_the_failed_requests_span_be_freed(monkeypatch):
     assert server_ref() is None
 
 
+def _hook_failure(logger: OpenTelemetryV2, server: object, error: str, *prompts: str) -> None:
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    asyncio.run(
+        logger.async_post_call_failure_hook(
+            request_data={"messages": [{"role": "user", "content": prompt} for prompt in prompts]},
+            original_exception=_proxy_exc(error, 400),
+            user_api_key_dict=UserAPIKeyAuth(parent_otel_span=server),
+        )
+    )
+
+
+def test_failure_hook_keeps_no_request_text_its_error_does_not_quote(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A keep-alive timer can hold a failed request's span past the request. The span must not keep alive request
+    text that the failure's error never quoted, whether or not it quoted another message of the request."""
+    import litellm
+
+    monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+    logger, _ = _logger()
+    server: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
+    prompt: Final = " ".join(("my ssn is", "123-45-6789"))
+    held: Final = sys.getrefcount(prompt)
+
+    _hook_failure(logger, server, "blocked by the upstream provider", prompt)
+    held_after_an_unquoting_error: Final = sys.getrefcount(prompt)
+    _hook_failure(logger, server, "blocked: say hi", prompt, "say hi")
+
+    assert (held_after_an_unquoting_error, sys.getrefcount(prompt)) == (held, held)
+
+
+def test_restamp_lets_go_of_the_request_text_it_scrubbed(monkeypatch: pytest.MonkeyPatch) -> None:
+    import litellm
+
+    monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+    logger, exporter = _logger()
+    server: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
+    prompt: Final = " ".join(("my ssn is", "123-45-6789"))
+    held: Final = sys.getrefcount(prompt)
+
+    _hook_failure(logger, server, f"blocked: {prompt}", prompt)
+    logger.record_error_attributes_on_span(server, _proxy_exc(f"blocked: {prompt}", 400), 400)
+    server.end()
+
+    assert sys.getrefcount(prompt) == held
+    assert exporter.get_finished_spans()[0].attributes["error.message"] == "blocked: redacted-by-litellm"
+
+
+def test_restamp_of_a_failure_that_is_not_redacted_keeps_its_own_text(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A span the proxy reuses across requests holds nothing from an earlier redacted failure once a later failure
+    on it is not redacted."""
+    import litellm
+
+    logger, exporter = _logger()
+    server: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
+    monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+    _hook_failure(logger, server, "blocked: my ssn is 123-45-6789", "my ssn is 123-45-6789")
+    monkeypatch.setattr(litellm, "turn_off_message_logging", False)
+    _hook_failure(logger, server, "blocked: my ssn is 123-45-6789", "my ssn is 123-45-6789")
+
+    logger.record_error_attributes_on_span(server, _proxy_exc("blocked: my ssn is 123-45-6789", 400), 400)
+    server.end()
+
+    assert exporter.get_finished_spans()[0].attributes["error.message"] == "blocked: my ssn is 123-45-6789"
+
+
+@pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"])
+def test_failure_hook_and_restamp_stamp_a_span_that_cannot_be_hashed(
+    monkeypatch: pytest.MonkeyPatch, redact: bool
+) -> None:
+    import litellm
+
+    logger, exporter = _logger()
+    concurrent: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
+    monkeypatch.setattr(litellm, "turn_off_message_logging", True)
+    _hook_failure(logger, concurrent, "blocked: another request", "another request")
+    monkeypatch.setattr(litellm, "turn_off_message_logging", redact)
+    server: Final = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
+    server.__class__ = type("UnhashableSpan", (type(server),), {"__slots__": (), "__hash__": None})
+
+    _hook_failure(logger, server, "blocked: my ssn is 123-45-6789", "my ssn is 123-45-6789")
+    logger.record_error_attributes_on_span(server, _proxy_exc("blocked", 403), 403)
+    server.end()
+    (span,) = exporter.get_finished_spans()
+
+    assert concurrent.is_recording()
+    assert span.attributes["error.type"] == "ProxyException"
+    assert span.attributes["litellm.provider.error.code"] == "403"
+    assert span.status.status_code is StatusCode.ERROR
+    assert [event.name for event in span.events] == ["exception"]
+
+
 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