From 05ededf8a05c3b290c38e73a747bd602a49033d3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:24:52 +0000 Subject: [PATCH] fix(otel): parent passthrough trace propagation on the legacy request span Pass user_api_key_dict.parent_otel_span into the outgoing W3C injection so the legacy otel callback propagates its litellm_request span, falling back to the otel_v2 request root span and then the ambient span. Extend the mapped unit tests to assert the propagated trace and span ids over real captured headers for HTTP and WebSocket passthrough with forwarding on and off. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 29 +++++------- .../pass_through_endpoints.py | 13 +++--- .../otel/test_otel_v2_components.py | 32 +++++++++---- .../test_pass_through_endpoints.py | 46 +++++++++++++------ 4 files changed, 74 insertions(+), 46 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 851097a17e7..3dabd34c81c 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -310,7 +310,10 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None: return _PROPAGATOR.extract(carrier) -def _outgoing_trace_context(inbound_headers: Mapping[str, str] | None = None) -> Context | None: +def _outgoing_trace_context(parent_span: object) -> Context | None: + if isinstance(parent_span, Span) and is_recordable_span(parent_span): + return context_from_span(parent_span) + root: Final = request_root_span() if root is not None: return context_from_span(root) @@ -318,27 +321,19 @@ def _outgoing_trace_context(inbound_headers: Mapping[str, str] | None = None) -> current: Final = get_current() if is_recordable_span(get_current_span(current)): return current - - if inbound_headers is None: - return None - inbound_context: Final = extract_traceparent(inbound_headers) - if inbound_context is None or not is_recordable_span(get_current_span(inbound_context)): - return None - return inbound_context + return None -def inject_trace_context( - headers: Mapping[str, str], - inbound_headers: Mapping[str, str] | None = None, -) -> dict[str, str]: +def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: """``headers`` plus W3C ``traceparent``/``tracestate`` for the current request's span. - Parent preference: the anchored request root span, then the ambient active span, - then the trace context the caller sent inbound. Only trace context is injected, - never Baggage, so per-request identity baggage cannot leak upstream. Unchanged - when no valid span context exists anywhere. + Parent preference: the request span auth stashed on the key (the legacy + ``litellm_request`` SERVER span, or the FastAPI server span under otel_v2), then + the anchored request root span, then the ambient active span. Only trace context + is injected, never Baggage, so per-request identity baggage cannot leak upstream. + Unchanged when no valid span context exists anywhere. """ - context: Final = _outgoing_trace_context(inbound_headers) + context: Final = _outgoing_trace_context(parent_span) if context is None: return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier carrier: Final = dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6d44b4fa6dc..c8c2db0dbd6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -985,7 +985,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) - headers = _with_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) + headers = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2161,12 +2161,12 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: _WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project")) -def _with_trace_context(headers: Mapping[str, str], inbound_headers: Mapping[str, str]) -> dict[str, str]: +def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]: try: from litellm.integrations.otel.plumbing.context import inject_trace_context except ImportError: return dict(headers) # mutable-ok: matches inject_trace_context's carrier return type - return inject_trace_context(headers, inbound_headers=inbound_headers) + return inject_trace_context(headers, parent_span=parent_span) async def websocket_passthrough_request( @@ -2211,16 +2211,15 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - incoming_headers: Final = dict(websocket.headers) # mutable-ok: propagator carrier - forwarded_headers: Final = { # mutable-ok: propagator carrier + forwarded_headers: Final = { # mutable-ok: one-shot upstream header dict, read as a Mapping **custom_headers, **{ header_name: header_value - for header_name, header_value in incoming_headers.items() + for header_name, header_value in websocket.headers.items() if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = _with_trace_context(forwarded_headers, inbound_headers=incoming_headers) + upstream_headers: Final = _with_trace_context(forwarded_headers, parent_span=user_api_key_dict.parent_otel_span) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( 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 b1e65da1661..43546241d06 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -507,16 +507,32 @@ 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_forwards_valid_inbound_context_without_span(): - inbound = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"} - +def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): def run(): - result = ctx_mod.inject_trace_context({}, inbound_headers=inbound) - return get_current_span(TraceContextTextMapPropagator().extract(result)) + 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) + result = ctx_mod.inject_trace_context({}, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return parent, ambient, propagated - propagated = ContextVarContext().run(run) - assert propagated.get_span_context().trace_id == int("0af7651916cd43dd8448eb211c80319c", 16) - assert propagated.get_span_context().span_id == int("b7ad6b7169203331", 16) + parent, ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == parent.get_span_context().trace_id + assert propagated.get_span_context().span_id == parent.get_span_context().span_id + assert propagated.get_span_context().span_id != ambient.get_span_context().span_id + + +def test_inject_trace_context_skips_unusable_parent_span(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient") as ambient: + result = ctx_mod.inject_trace_context({}, parent_span=object()) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return ambient, propagated + + ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id def test_inject_trace_context_returns_headers_unchanged_without_context(): 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 0603b5cb2ff..d6b9a782f99 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 @@ -51,7 +51,7 @@ MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start" def test_with_trace_context_without_opentelemetry(monkeypatch: pytest.MonkeyPatch): monkeypatch.setitem(sys.modules, "litellm.integrations.otel.plumbing.context", None) - headers = _with_trace_context({"authorization": "x"}, {}) + headers = _with_trace_context({"authorization": "x"}, parent_span=None) assert headers == {"authorization": "x"} assert "traceparent" not in headers @@ -4282,11 +4282,11 @@ def _relay_client_request(method="GET"): @pytest.mark.asyncio -async def test_pass_through_request_propagates_active_trace_context(): +@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"]) +async def test_pass_through_request_propagates_active_trace_context(span_source: str): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator - from litellm.proxy._types import UserAPIKeyAuth captured: dict[str, httpx.Headers] = {} @@ -4295,17 +4295,23 @@ async def test_pass_through_request_propagates_active_trace_context(): 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, {}) - tracer = TracerProvider().get_tracer("test") - with tracer.start_as_current_span("passthrough") as span: - response = await pass_through_request( - request=_relay_client_request(method="POST"), - target="http://internal-api.test/v1/generate", - custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), - ) + if span_source == "auth_parent_span": + span = tracer.start_span("litellm_request") + stack.callback(span.end) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span) + else: + span = stack.enter_context(tracer.start_as_current_span("passthrough")) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + response = await pass_through_request( + request=_relay_client_request(method="POST"), + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) finally: cleanup() await fake_client.aclose() @@ -4313,6 +4319,7 @@ async def test_pass_through_request_propagates_active_trace_context(): assert response.status_code == 200 propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + assert propagated.get_span_context().span_id == span.get_span_context().span_id @pytest.mark.asyncio @@ -4913,7 +4920,10 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): @pytest.mark.asyncio @pytest.mark.parametrize("forward_headers", [True, False]) -async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch, forward_headers: bool): +@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"]) +async def test_websocket_passthrough_propagates_active_trace_context( + monkeypatch, forward_headers: bool, span_source: str +): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator @@ -4954,12 +4964,19 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER", mock_worker, ) - with tracer.start_as_current_span("websocket_passthrough") as span: + with ExitStack() as stack: + if span_source == "auth_parent_span": + span = tracer.start_span("litellm_request") + stack.callback(span.end) + user_api_key_dict = UserAPIKeyAuth(parent_otel_span=span) + else: + span = stack.enter_context(tracer.start_as_current_span("websocket_passthrough")) + user_api_key_dict = UserAPIKeyAuth() await websocket_passthrough_request( websocket=websocket, target="wss://upstream.example.test/v1/realtime", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=user_api_key_dict, forward_headers=forward_headers, endpoint="/realtime", accept_websocket=True, @@ -4967,6 +4984,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + assert propagated.get_span_context().span_id == span.get_span_context().span_id assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None)