Merge pull request #41786 from BerriAI/litellm_passthrough_xpass_trace

Pass-through requests inject the proxy span into upstream headers since #40669, which
replaced an explicit x-pass-traceparent with an unrelated trace and dropped its
x-pass-tracestate. Keep the caller's context when the carrier already names a
different trace, and keep the proxy child span for same-trace or missing headers.

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/otel/plumbing/context.py

@@ -325,12 +325,27 @@ def _outgoing_trace_context(parent_span: object) -> Context | None:
     return None
 
 
+def _propagated_context(headers: Mapping[str, str], request_context: Context) -> Context:
+    """``request_context`` when it continues the trace ``headers`` already name, else the
+    caller's own context, so an explicit upstream ``traceparent`` (``x-pass-traceparent``)
+    is never swapped for an unrelated trace and its ``tracestate`` survives."""
+    caller: Final = extract_traceparent(headers)
+    if caller is None:
+        return request_context
+    caller_span: Final = get_current_span(caller).get_span_context()
+    request_span: Final = get_current_span(request_context).get_span_context()
+    if not caller_span.is_valid or caller_span.trace_id == request_span.trace_id:
+        return request_context
+    return caller
+
+
 def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
     """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
 
     Parent preference: ``parent_span`` (the request span auth stashed on the key), then
     the anchored request root span, then the ambient active span. Only trace context is
-    injected, never Baggage. Unchanged when no valid span exists anywhere.
+    injected, never Baggage. Unchanged when no valid span exists anywhere. A ``traceparent``
+    already in ``headers`` from a different trace is forwarded as-is instead of replaced.
     """
     context: Final = _outgoing_trace_context(parent_span)
     if context is None:
@@ -338,7 +353,7 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None)
     carrier: Final = {  # mutable-ok: OpenTelemetry propagator requires a mutable carrier
         key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
     }
-    _PROPAGATOR.inject(carrier, context=context)
+    _PROPAGATOR.inject(carrier, context=_propagated_context(headers, context))
     return carrier
 
 

tests/test_litellm/integrations/otel/test_otel_v2_components.py

@@ -477,19 +477,21 @@ def _test_tracer():
     return provider.get_tracer("test")
 
 
+_CALLER_TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01"
+
+
 def test_inject_trace_context_prefers_request_root_span():
     def run():
         tracer = _test_tracer()
-        with tracer.start_as_current_span("root") as root:
+        inbound = TraceContextTextMapPropagator().extract({"traceparent": _CALLER_TRACEPARENT})
+        with tracer.start_as_current_span("root", context=inbound) as root:
             ctx_mod.set_request_root_span(root)
-            result = ctx_mod.inject_trace_context(
-                {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"}
-            )
+            result = ctx_mod.inject_trace_context({"traceparent": _CALLER_TRACEPARENT})
             propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
             return result, root, propagated
 
     result, root, propagated = ContextVarContext().run(run)
-    assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01"
+    assert result["traceparent"] != _CALLER_TRACEPARENT
     assert propagated.get_span_context().trace_id == root.get_span_context().trace_id
     assert propagated.get_span_context().span_id == root.get_span_context().span_id
 
@@ -507,24 +509,57 @@ def test_inject_trace_context_uses_ambient_span_without_request_root():
     assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
 
 
-def test_inject_trace_context_replaces_stale_trace_headers():
+def test_inject_trace_context_replaces_same_trace_headers_with_request_span():
     def run():
         tracer = _test_tracer()
-        with tracer.start_as_current_span("ambient") as ambient:
-            headers = {
-                "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01",
-                "Tracestate": "vendor=old",
-                "x-keep": "1",
-            }
+        headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"}
+        inbound = TraceContextTextMapPropagator().extract({key.lower(): value for key, value in headers.items()})
+        with tracer.start_as_current_span("ambient", context=inbound) as ambient:
             result = ctx_mod.inject_trace_context(headers)
             propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
             return result, ambient, propagated
 
     result, ambient, propagated = ContextVarContext().run(run)
     assert sum(key.lower() == "traceparent" for key in result) == 1
-    assert not any(key.lower() == "tracestate" for key in result)
+    assert sum(key.lower() == "tracestate" for key in result) == 1
     assert result["x-keep"] == "1"
-    assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id
+    assert result["tracestate"] == "vendor=caller"
+    assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
+
+
+def test_inject_trace_context_keeps_caller_traceparent_from_another_trace():
+    def run():
+        tracer = _test_tracer()
+        parent = tracer.start_span("litellm_request")
+        with tracer.start_as_current_span("ambient") as ambient:
+            ctx_mod.set_request_root_span(ambient)
+            headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"}
+            result = ctx_mod.inject_trace_context(headers, parent_span=parent)
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return result, parent, propagated
+
+    result, parent, propagated = ContextVarContext().run(run)
+    assert result["traceparent"] == _CALLER_TRACEPARENT
+    assert result["tracestate"] == "vendor=caller"
+    assert result["x-keep"] == "1"
+    assert sum(key.lower() == "traceparent" for key in result) == 1
+    assert sum(key.lower() == "tracestate" for key in result) == 1
+    assert propagated.get_span_context().trace_id != parent.get_span_context().trace_id
+
+
+def test_inject_trace_context_replaces_malformed_caller_traceparent():
+    def run():
+        tracer = _test_tracer()
+        parent = tracer.start_span("litellm_request")
+        with tracer.start_as_current_span("ambient"):
+            headers = {"traceparent": "not-a-traceparent", "tracestate": "vendor=caller"}
+            result = ctx_mod.inject_trace_context(headers, parent_span=parent)
+            propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
+            return result, parent, propagated
+
+    result, parent, propagated = ContextVarContext().run(run)
+    assert propagated.get_span_context().span_id == parent.get_span_context().span_id
+    assert "tracestate" not in result
 
 
 def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient():

tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -4418,6 +4418,65 @@ async def test_pass_through_request_propagates_active_trace_context(span_source:
     assert propagated.get_span_context().span_id == span.get_span_context().span_id
 
 
+async def _relay_with_trace_headers(inbound_headers: dict[str, str], forward_headers: bool):
+    from opentelemetry.sdk.trace import TracerProvider
+
+    captured: dict[str, httpx.Headers] = {}
+
+    def transport_handler(upstream_request: httpx.Request) -> httpx.Response:
+        captured["headers"] = upstream_request.headers
+        return httpx.Response(200, json={"ok": True}, request=upstream_request)
+
+    fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None)
+    tracer = TracerProvider().get_tracer("test")
+    try:
+        with ExitStack() as stack:
+            _enter_relay_logging_mocks(stack, {})
+            span = tracer.start_span("litellm_request")
+            stack.callback(span.end)
+            request = _relay_client_request(method="POST")
+            request.headers = Headers(inbound_headers)
+            response = await pass_through_request(
+                request=request,
+                target="http://internal-api.test/v1/generate",
+                custom_headers={},
+                user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span),
+                forward_headers=forward_headers,
+            )
+    finally:
+        cleanup()
+        await fake_client.aclose()
+    assert response.status_code == 200
+    return captured["headers"], span
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("forward_headers", [False, True])
+async def test_pass_through_request_keeps_x_pass_trace_headers_when_otel_span_is_active(forward_headers: bool):
+    caller_traceparent = "00-11111111111111111111111111111111-2222222222222222-01"
+
+    upstream_headers, span = await _relay_with_trace_headers(
+        {"x-pass-traceparent": caller_traceparent, "x-pass-tracestate": "vendor=caller"},
+        forward_headers=forward_headers,
+    )
+
+    assert upstream_headers["traceparent"] == caller_traceparent
+    assert upstream_headers["tracestate"] == "vendor=caller"
+    assert format(span.get_span_context().trace_id, "032x") not in upstream_headers["traceparent"]
+
+
+@pytest.mark.asyncio
+async def test_pass_through_request_without_caller_trace_headers_still_propagates_proxy_span():
+    from opentelemetry.trace import get_current_span
+    from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+    upstream_headers, span = await _relay_with_trace_headers({"x-pass-anthropic-beta": "beta-1"}, forward_headers=False)
+
+    propagated = get_current_span(TraceContextTextMapPropagator().extract(upstream_headers))
+    assert propagated.get_span_context().span_id == span.get_span_context().span_id
+    assert upstream_headers["anthropic-beta"] == "beta-1"
+
+
 @pytest.mark.asyncio
 async def test_pass_through_request_relays_non_json_body_without_buffering():
     """