fix(langfuse_otel): honor caller trace_user_id sent under litellm_metadata

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/arize/_utils.py

@@ -878,7 +878,6 @@ def _set_session_and_user_attrs(span: "Span", kwargs: dict, standard_logging_pay
 
 
 def _set_request_context_attrs(span: "Span", standard_logging_payload: object) -> None:
-    """Emit `litellm.trace_id` / team / key context when source data exists."""
     payload: Final = as_str_mapping(standard_logging_payload)
     if payload is None:
         return

litellm/integrations/langfuse/langfuse_otel.py

@@ -11,6 +11,7 @@ from litellm.integrations.langfuse.langfuse_otel_attributes import (
     LangfuseLLMObsOTELAttributes,
 )
 from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from litellm.integrations.otel.model.trace_controls import metadata_bodies
 from litellm.integrations.otel.model.utils import as_str, as_str_mapping
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.types.integrations.langfuse_otel import (
@@ -267,12 +268,18 @@ class LangfuseOtelLogger(OpenTelemetry):
         metadata: Final = as_str_mapping(
             LangfuseOtelLogger._extract_langfuse_metadata(kwargs)  # pyright: ignore[reportUnknownArgumentType,reportUnknownMemberType]  # helper returns a loosely typed dict
         )
-        if metadata is None or as_str(metadata.get("trace_user_id")):
+        litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
+        bodies: Final = metadata_bodies(litellm_params)
+        caller: Final = (as_str(metadata.get("trace_user_id")) if metadata is not None else None) or next(
+            (value for body in bodies if (value := as_str(body.get("trace_user_id")))), None
+        )
+        if caller is not None:
+            safe_set_attribute(span, LangfuseSpanAttributes.TRACE_USER_ID.value, caller)
             return
         end_user: Final = (
             slp_metadata.get("user_api_key_end_user_id")
             or slp.get("end_user")
-            or as_str(metadata.get("user_api_key_end_user_id"))
+            or next((value for body in bodies if (value := as_str(body.get("user_api_key_end_user_id")))), None)
         )
         if end_user:
             safe_set_attribute(span, LangfuseSpanAttributes.TRACE_USER_ID.value, str(end_user))

litellm/integrations/otel/langfuse_logger.py

@@ -18,9 +18,8 @@ if TYPE_CHECKING:
 
 
 class LangfuseOpenTelemetryV2(OpenTelemetryV2):
-    """Stamps the caller's trace controls (name, user, session, tags) on the request, falling back to the
-    proxy's end user when the caller names no trace user. Langfuse reads them off the root observation,
-    and the proxy's root span is still recording when the LLM call starts."""
+    """Stamps the caller's trace controls (name, user, session, tags) on the request. Langfuse reads them off
+    the root observation, and the proxy's root span is still recording when the LLM call starts."""
 
     def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
         root: Final = request_root_span()

litellm/integrations/otel/model/trace_controls.py

@@ -33,7 +33,7 @@ def caller_trace_controls(kwargs: Mapping[str, object]) -> TraceControls:
         return TraceControls()
     proxy_request: Final = as_str_mapping(request.get("proxy_server_request"))
     headers: Final = as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None
-    bodies: Final = _metadata_bodies(request)
+    bodies: Final = metadata_bodies(request)
 
     def scalar(control: str) -> str | None:
         from_header: Final = as_str(headers.get(f"{LANGFUSE_HEADER_PREFIX}{control}")) if headers is not None else None
@@ -57,13 +57,13 @@ def langfuse_trace_controls(kwargs: Mapping[str, object]) -> TraceControls:
     if request is None:
         return controls
     end_user: Final = next(
-        (value for body in _metadata_bodies(request) if (value := as_str(body.get("user_api_key_end_user_id")))),
+        (value for body in metadata_bodies(request) if (value := as_str(body.get("user_api_key_end_user_id")))),
         None,
     )
     return replace(controls, user_id=end_user)
 
 
-def _metadata_bodies(request: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
+def metadata_bodies(request: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
     return tuple(
         metadata
         for key in ("metadata", "litellm_metadata")

tests/integration/contracts.json

@@ -275,6 +275,9 @@
     "tests/integration/observability/test_langfuse_otel_identity.py::test_langfuse_otel_v2_header_end_user_lands_in_user_id": [
       "other.observability.langfuse_otel.v2_header_end_user_in_user_id"
     ],
+    "tests/integration/observability/test_langfuse_otel_identity.py::test_langfuse_otel_messages_caller_trace_user_id_under_litellm_metadata_wins_over_the_end_user": [
+      "other.observability.langfuse_otel.messages_caller_trace_user_id_under_litellm_metadata"
+    ],
     "tests/integration/observability/test_guardrail_effects.py::test_guardrail_rewrites_system_and_user_in_actual_anthropic_request": [
       "other.observability.guardrails.rewrite_reaches_correct_anthropic_positions"
     ],

tests/integration/observability/test_langfuse_otel_identity.py

@@ -71,6 +71,14 @@ def _generation_span_attributes(sink: Wire, batches: list[bytes], marker: str) -
     )
 
 
+def _span_attributes_containing_marker(sink: Wire, batches: list[bytes], marker: str) -> tuple[dict[str, object], ...]:
+    return tuple(
+        attributes
+        for _trace_id, attributes in _drained_spans(sink, batches)
+        if any(isinstance(value, str) and marker in value for value in attributes.values())
+    )
+
+
 def _trace_user_span_attributes(sink: Wire, batches: list[bytes], marker: str) -> tuple[dict[str, object], ...]:
     spans: Final = _drained_spans(sink, batches)
     generation_trace: Final = next(
@@ -353,3 +361,47 @@ def test_langfuse_otel_v2_header_end_user_lands_in_user_id(gateway: Gateway, tmp
             "user.id": f"end-user-{marker}",
             "session.id": None,
         }, user_spans[0]
+
+
+@pytest.mark.covers("other.observability.langfuse_otel.messages_caller_trace_user_id_under_litellm_metadata")
+def test_langfuse_otel_messages_caller_trace_user_id_under_litellm_metadata_wins_over_the_end_user(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    marker: Final = uuid.uuid4().hex
+    upstream_bodies: Final[list[bytes]] = []
+
+    def upstream(request: Request) -> Reply:
+        upstream_bodies.append(request.body)
+        return _upstream_reply(marker)
+
+    with (
+        wire_server(upstream) as provider,
+        wire_server(_sink) as collector,
+        _langfuse_proxy(gateway, tmp_path, collector.url) as candidate,
+        candidate.scenario() as scenario,
+    ):
+        model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=provider.url + "/v1")
+        response: Final = candidate.request(
+            "POST",
+            "/v1/messages",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": marker}],
+                "max_tokens": 5,
+                "litellm_metadata": {"trace_user_id": f"caller-{marker}"},
+                "cache": {"no-cache": True},
+            },
+            headers={"x-litellm-end-user-id": f"end-user-{marker}"},
+        )
+        assert response.status_code == 200, response.text
+        assert any(marker.encode() in body for body in upstream_bodies), upstream_bodies
+        batches: Final[list[bytes]] = []
+        attributes: Final = eventually(
+            lambda: _span_attributes_containing_marker(collector, batches, marker),
+            lambda spans: len(spans) == 1,
+            seconds=30,
+        )[0]
+        assert {key: attributes.get(key) for key in ("user.id", "session.id")} == {
+            "user.id": f"caller-{marker}",
+            "session.id": None,
+        }, attributes

tests/test_litellm/integrations/test_langfuse_otel.py

@@ -1073,6 +1073,18 @@ class TestLangfuseOtelTraceIdentity:
         )
         assert identity == {"user.id": "caller-1", "session.id": None}
 
+    def test_caller_trace_user_id_under_litellm_metadata_wins_over_the_end_user(self):
+        kwargs = self._kwargs(slp_metadata={"user_api_key_end_user_id": "end-1"})
+        kwargs["litellm_params"]["litellm_metadata"] = {"trace_user_id": "caller-1"}
+        identity, _ = self._identity(kwargs)
+        assert identity == {"user.id": "caller-1", "session.id": None}
+
+    def test_end_user_only_under_litellm_metadata_lands_in_user_id(self):
+        kwargs = self._kwargs()
+        kwargs["litellm_params"]["litellm_metadata"] = {"user_api_key_end_user_id": "end-1"}
+        identity, _ = self._identity(kwargs)
+        assert identity == {"user.id": "end-1", "session.id": None}
+
     def test_caller_session_id_stays_the_session_beside_the_end_user(self):
         identity, _ = self._identity(
             self._kwargs(