From c9050a8fbe68be0ae0be744f5ca788e40c77f87f Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:20:03 +0000 Subject: [PATCH] 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 | 1 - .../integrations/langfuse/langfuse_otel.py | 11 +++- litellm/integrations/otel/langfuse_logger.py | 5 +- .../integrations/otel/model/trace_controls.py | 6 +-- tests/integration/contracts.json | 3 ++ .../test_langfuse_otel_identity.py | 52 +++++++++++++++++++ .../integrations/test_langfuse_otel.py | 12 +++++ 7 files changed, 81 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 9594bfc268f..0f5addeb39d 100644 --- a/litellm/integrations/arize/_utils.py +++ b/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 diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 91771995d1d..342b8ef5148 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/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)) diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py index a690b33f0f0..2578777f303 100644 --- a/litellm/integrations/otel/langfuse_logger.py +++ b/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() diff --git a/litellm/integrations/otel/model/trace_controls.py b/litellm/integrations/otel/model/trace_controls.py index 2f683c7fd1b..dc64d5d5743 100644 --- a/litellm/integrations/otel/model/trace_controls.py +++ b/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") diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 22a1323f311..50ff6072ec2 100644 --- a/tests/integration/contracts.json +++ b/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" ], diff --git a/tests/integration/observability/test_langfuse_otel_identity.py b/tests/integration/observability/test_langfuse_otel_identity.py index 0d04899585f..2c9ed62fde0 100644 --- a/tests/integration/observability/test_langfuse_otel_identity.py +++ b/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 diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index e34670c8488..48473523352 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/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(