From 318b027782b0b6ab1aec7103a619d308661894e6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 18 Sep 2026 08:26:39 +0000 Subject: [PATCH] fix(otel): keep caller traceparent and tracestate on pass-through relays 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: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 19 +++++- .../otel/test_otel_v2_components.py | 63 ++++++++++++++----- .../test_pass_through_endpoints.py | 59 +++++++++++++++++ 3 files changed, 125 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 1282e654365..19243d64c64 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/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 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 72f6213e880..0c95049ce05 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/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(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e3e7ad618e0..df4842a7887 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4328,6 +4328,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(): """