From 7bce15f7d8493096ced51023f7a239091bfcc1f2 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:45:43 +0000 Subject: [PATCH 01/11] fix(otel): propagate W3C trace context on passthrough upstream requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 36 ++++++++ .../pass_through_endpoints.py | 30 +++--- .../otel/test_otel_v2_components.py | 81 ++++++++++++++++- .../test_pass_through_endpoints.py | 91 +++++++++++++++++++ 4 files changed, 224 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 21e61c71fb7..851097a17e7 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -310,6 +310,42 @@ 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: + root: Final = request_root_span() + if root is not None: + return context_from_span(root) + + 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 + + +def inject_trace_context( + headers: Mapping[str, str], + inbound_headers: Mapping[str, str] | None = 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. + """ + context: Final = _outgoing_trace_context(inbound_headers) + 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 + _PROPAGATOR.inject(carrier, context=context) + return carrier + + # The OTLP destinations this request's key or team pointed its traces at, resolved # once during auth. A ``ContextVar`` for the same reason the root span above is one: # it rides the request task's context into the ``asyncio.create_task`` children that diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..1c4f8d59c9c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -955,6 +955,7 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. """ from litellm.exceptions import ModifyResponseException + from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( PassthroughGuardrailHandler, @@ -985,6 +986,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) + headers = inject_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2199,20 +2201,22 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - # Prepare headers for the upstream connection - upstream_headers: Final = custom_headers.copy() + from litellm.integrations.otel.plumbing.context import inject_trace_context - if forward_headers: - # Forward relevant headers from the incoming request - incoming_headers: Final = dict(websocket.headers) - for header_name, header_value in incoming_headers.items(): - # Only forward certain headers to avoid conflicts - if header_name.lower() in [ - "authorization", - "x-api-key", - "x-goog-user-project", - ]: - upstream_headers[header_name] = header_value + incoming_headers: Final = dict(websocket.headers) # mutable-ok: websocket headers are copied for context extraction + forwarded_headers: Final = { # mutable-ok: assembled as the upstream header carrier + **custom_headers, + **{ + header_name: header_value + for header_name, header_value in incoming_headers.items() + if forward_headers + and header_name.lower() in frozenset(("authorization", "x-api-key", "x-goog-user-project")) + }, + } + upstream_headers: Final = inject_trace_context( + forwarded_headers, + inbound_headers=incoming_headers, + ) # 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 ae41c74944d..b1e65da1661 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -5,6 +5,7 @@ builders, and the registry validator's failure paths. Needs the OTel SDK.""" import json import threading from collections.abc import Iterator +from contextvars import Context as ContextVarContext from dataclasses import replace from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer @@ -15,6 +16,8 @@ pytest.importorskip("opentelemetry") from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: E402 ExportTraceServiceRequest, ) +from opentelemetry import baggage # noqa: E402 +from opentelemetry.context import attach, detach # noqa: E402 from opentelemetry.sdk.metrics import MeterProvider # noqa: E402 from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402 from opentelemetry.sdk.trace import TracerProvider # noqa: E402 @@ -26,7 +29,10 @@ from opentelemetry.sdk.trace.export import ( # noqa: E402 from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402 InMemorySpanExporter, ) -from opentelemetry.trace import SpanKind # noqa: E402 +from opentelemetry.trace import SpanKind, get_current_span # noqa: E402 +from opentelemetry.trace.propagation.tracecontext import ( # noqa: E402 + TraceContextTextMapPropagator, +) from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 @@ -464,6 +470,79 @@ def test_extract_traceparent(): assert ctx_mod.extract_traceparent({"x": "y"}) is None +def _test_tracer(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider.get_tracer("test") + + +def test_inject_trace_context_prefers_request_root_span(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("root") as root: + ctx_mod.set_request_root_span(root) + result = ctx_mod.inject_trace_context( + {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} + ) + 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 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 + + +def test_inject_trace_context_uses_ambient_span_without_request_root(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient") as ambient: + result = ctx_mod.inject_trace_context({}) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return ambient, propagated + + ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + 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 run(): + result = ctx_mod.inject_trace_context({}, inbound_headers=inbound) + return get_current_span(TraceContextTextMapPropagator().extract(result)) + + propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == int("0af7651916cd43dd8448eb211c80319c", 16) + assert propagated.get_span_context().span_id == int("b7ad6b7169203331", 16) + + +def test_inject_trace_context_returns_headers_unchanged_without_context(): + headers = {"x-custom": "value"} + + result = ContextVarContext().run(lambda: ctx_mod.inject_trace_context(headers)) + + assert result == headers + assert "traceparent" not in result + assert result is not headers + + +def test_inject_trace_context_does_not_forward_baggage(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient"): + token = attach(baggage.set_baggage("litellm.team.id", "team")) + try: + return ctx_mod.inject_trace_context({}) + finally: + detach(token) + + result = ContextVarContext().run(run) + assert "baggage" not in result + + def test_set_request_baggage_empty_returns_context(): assert ctx_mod.set_request_baggage({}) is not None 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 0fc961cf8c9..c3ad900cc9e 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 @@ -4270,6 +4270,40 @@ def _relay_client_request(method="GET"): return mock_request +@pytest.mark.asyncio +async def test_pass_through_request_propagates_active_trace_context(): + 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] = {} + + 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) + 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"), + ) + finally: + cleanup() + await fake_client.aclose() + + 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 + + @pytest.mark.asyncio async def test_pass_through_request_relays_non_json_body_without_buffering(): """ @@ -4866,6 +4900,63 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): assert all(call.kwargs.get("code") != 1011 for call in websocket.close.await_args_list) +@pytest.mark.asyncio +async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch): + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + from starlette.websockets import WebSocketState + + captured: dict[str, dict[str, str]] = {} + upstream_ws = FakeUpstreamWebSocket(b"{}") + + def fake_connect(target, additional_headers): + captured["headers"] = additional_headers + return FakeUpstreamConnect(upstream_ws) + + websocket = MagicMock() + websocket.accept = AsyncMock() + websocket.send_text = AsyncMock() + websocket.send_bytes = AsyncMock() + websocket.receive = AsyncMock(return_value={"type": "websocket.disconnect"}) + websocket.close = AsyncMock() + websocket.headers = {} + websocket.client_state = WebSocketState.CONNECTED + websocket.application_state = WebSocketState.CONNECTED + tracer = TracerProvider().get_tracer("test") + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_worker = MagicMock() + mock_worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: async_coroutine.close() + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + fake_connect, + ) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER", + mock_worker, + ) + with tracer.start_as_current_span("websocket_passthrough") as span: + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example.test/v1/realtime", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/realtime", + accept_websocket=True, + ) + + propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) + assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + + class ClosingUpstreamWebSocket: def __init__(self, close_exc: Exception): self._close_exc = close_exc From 7bcf848b39883c651514d8df0c366a8c89ce54e2 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:47:01 +0000 Subject: [PATCH 02/11] refactor(passthrough): hoist websocket forwarded header set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1c4f8d59c9c..186318dccc1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2159,6 +2159,9 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: return upstream_close +_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project")) + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2183,6 +2186,7 @@ async def websocket_passthrough_request( cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream """ + from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2201,22 +2205,16 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - from litellm.integrations.otel.plumbing.context import inject_trace_context - - incoming_headers: Final = dict(websocket.headers) # mutable-ok: websocket headers are copied for context extraction - forwarded_headers: Final = { # mutable-ok: assembled as the upstream header carrier + incoming_headers: Final = dict(websocket.headers) # mutable-ok: propagator carrier + forwarded_headers: Final = { # mutable-ok: propagator carrier **custom_headers, **{ header_name: header_value for header_name, header_value in incoming_headers.items() - if forward_headers - and header_name.lower() in frozenset(("authorization", "x-api-key", "x-goog-user-project")) + if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = inject_trace_context( - forwarded_headers, - inbound_headers=incoming_headers, - ) + upstream_headers: Final = inject_trace_context(forwarded_headers, inbound_headers=incoming_headers) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( From d2e8b9c6565478c9ce1f5ba1cca834537a1759af Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:05:05 +0000 Subject: [PATCH 03/11] fix(otel): keep passthrough working when opentelemetry is not installed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 14 ++++++++++---- .../test_pass_through_endpoints.py | 11 +++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 186318dccc1..6d44b4fa6dc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -955,7 +955,6 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. """ from litellm.exceptions import ModifyResponseException - from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( PassthroughGuardrailHandler, @@ -986,7 +985,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) - headers = inject_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) + headers = _with_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2162,6 +2161,14 @@ 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]: + 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) + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2186,7 +2193,6 @@ async def websocket_passthrough_request( cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream """ - from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2214,7 +2220,7 @@ async def websocket_passthrough_request( if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = inject_trace_context(forwarded_headers, inbound_headers=incoming_headers) + upstream_headers: Final = _with_trace_context(forwarded_headers, inbound_headers=incoming_headers) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( 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 c3ad900cc9e..05371dccb23 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 @@ -2,6 +2,7 @@ import asyncio import json import logging import os +import sys from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO @@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( resolve_pass_through_request_timeout, resolve_llm_passthrough_timeout, websocket_passthrough_request, + _with_trace_context, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -46,6 +48,15 @@ import litellm MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n' +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"}, {}) + + assert headers == {"authorization": "x"} + assert "traceparent" not in headers + + # Test is_multipart def test_is_multipart(): # Test with multipart content type From e4a12510f647fb9ef5a790f93f1b829b78f7d905 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:22:10 +0000 Subject: [PATCH 04/11] test(otel): cover websocket forwarded headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints/test_pass_through_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 05371dccb23..d81e3ae5af8 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 @@ -4931,7 +4931,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch websocket.send_bytes = AsyncMock() websocket.receive = AsyncMock(return_value={"type": "websocket.disconnect"}) websocket.close = AsyncMock() - websocket.headers = {} + websocket.headers = {"authorization": "Bearer client"} websocket.client_state = WebSocketState.CONNECTED websocket.application_state = WebSocketState.CONNECTED tracer = TracerProvider().get_tracer("test") @@ -4959,7 +4959,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch target="wss://upstream.example.test/v1/realtime", custom_headers={}, user_api_key_dict=UserAPIKeyAuth(), - forward_headers=False, + forward_headers=True, endpoint="/realtime", accept_websocket=True, ) From 68cc12e848fb445b42aa27e6c005e91b17595e41 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:36:35 +0000 Subject: [PATCH 05/11] test(otel): assert websocket forwarded header reaches upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/pass_through_endpoints/test_pass_through_endpoints.py | 1 + 1 file changed, 1 insertion(+) 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 d81e3ae5af8..a21f695493b 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 @@ -4966,6 +4966,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 captured["headers"]["authorization"] == "Bearer client" class ClosingUpstreamWebSocket: From 7ba073aa2668820a35c69dede3ef46ed7b837cc6 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:47:09 +0000 Subject: [PATCH 06/11] test(otel): cover websocket trace propagation with forwarding on and off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints/test_pass_through_endpoints.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 a21f695493b..0603b5cb2ff 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 @@ -4912,7 +4912,8 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): @pytest.mark.asyncio -async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch): +@pytest.mark.parametrize("forward_headers", [True, False]) +async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch, forward_headers: bool): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator @@ -4959,14 +4960,14 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch target="wss://upstream.example.test/v1/realtime", custom_headers={}, user_api_key_dict=UserAPIKeyAuth(), - forward_headers=True, + forward_headers=forward_headers, endpoint="/realtime", accept_websocket=True, ) propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id - assert captured["headers"]["authorization"] == "Bearer client" + assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None) class ClosingUpstreamWebSocket: From 05ededf8a05c3b290c38e73a747bd602a49033d3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:24:52 +0000 Subject: [PATCH 07/11] 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) From 8899583d06ea4f570136cdcdecceb41f2de652f6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:33:39 +0000 Subject: [PATCH 08/11] docs(otel): tighten inject_trace_context docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 3dabd34c81c..57e220815b5 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -325,13 +325,11 @@ def _outgoing_trace_context(parent_span: object) -> Context | None: 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. + """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span. - 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. + 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. """ context: Final = _outgoing_trace_context(parent_span) if context is None: From b82de95625cb807e892a223f90004d46c43589c0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 08:04:27 +0000 Subject: [PATCH 09/11] refactor(passthrough): bind trace-enriched headers to a Final local Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index c8c2db0dbd6..7d2fe8c45a9 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, parent_span=user_api_key_dict.parent_otel_span) + upstream_headers: Final = _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) @@ -1019,7 +1019,7 @@ async def pass_through_request( verbose_proxy_logger.debug( "Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n", url, - headers, + upstream_headers, _parsed_body, ) @@ -1257,7 +1257,7 @@ async def pass_through_request( additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": headers, + "headers": upstream_headers, }, ) stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( @@ -1274,7 +1274,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=headers, + headers=upstream_headers, requested_query_params=requested_query_params, stream=True, ) @@ -1286,7 +1286,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=headers, + headers=upstream_headers, content=state_raw_body, ) if state_raw_body is not None @@ -1294,7 +1294,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=headers, + headers=upstream_headers, json=_parsed_body, ) ) @@ -1371,7 +1371,7 @@ async def pass_through_request( raw_body_request: Final = async_client.build_request( request.method, url, - headers=headers, + headers=upstream_headers, params=requested_query_params, content=state_raw_body, ) @@ -1381,7 +1381,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=headers, + headers=upstream_headers, requested_query_params=requested_query_params, _parsed_body=_parsed_body, forward_multipart=is_multipart, From cd594f104af2ab72242f6512e3b6aa240df658a4 Mon Sep 17 00:00:00 2001 From: shivam Date: Wed, 16 Sep 2026 08:34:53 +0000 Subject: [PATCH 10/11] fix(otel): drop stale trace headers before injecting passthrough trace context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 5 ++++- .../otel/test_otel_v2_components.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 57e220815b5..1282e654365 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -26,6 +26,7 @@ if TYPE_CHECKING: from litellm.integrations.otel.model.destination import OtelDestination _PROPAGATOR: Final = TraceContextTextMapPropagator() +_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate")) # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the # proxy first resolves it, so request-level spans (the LLM call, guardrails) can @@ -334,7 +335,9 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) 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 + 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) 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 43546241d06..72f6213e880 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -507,6 +507,26 @@ 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 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", + } + 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 result["x-keep"] == "1" + assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + + def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): def run(): tracer = _test_tracer() From df9a87f44a09d77b10b3868700ae6569bdedb57d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 20:28:03 +0000 Subject: [PATCH 11/11] fix(otel): keep caller tracestate on the legacy request span Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/opentelemetry.py | 11 +--- .../integrations/test_opentelemetry.py | 59 +++++++++++++++++++ 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index d4e7fcb577e..2623a9b5a56 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -2932,16 +2932,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) propagator: Final = TraceContextTextMapPropagator() - carrier: Final = {"traceparent": _traceparent} + carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None} _parent_context: Final = propagator.extract(carrier=carrier) return _parent_context def _get_span_context(self, kwargs, default_span: Span | None = None): from opentelemetry import context, trace - from opentelemetry.trace.propagation.tracecontext import ( - TraceContextTextMapPropagator, - ) litellm_params: Final = kwargs.get("litellm_params", {}) or {} proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {} @@ -2965,11 +2962,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Priority 2: HTTP traceparent header if traceparent is not None: verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation") - carrier: Final = {"traceparent": traceparent} - return ( - TraceContextTextMapPropagator().extract(carrier=carrier), - None, - ) + return self.get_traceparent_from_header(headers=headers), None # Priority 3: Active span from global context (auto-detection) try: diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 9ec8489f784..443b9012235 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -5423,6 +5423,65 @@ class TestGetSpanContextLitellmMetadataFallback(unittest.TestCase): self.assertIsNone(detected_span) +class TestInboundTraceContextKeepsCallerTracestate(unittest.TestCase): + """The request span built from inbound W3C headers must carry the caller's + tracestate so outbound propagation (passthrough) re-emits it instead of + dropping it alongside the stripped stale header.""" + + CALLER_TRACEPARENT = "00-" + "a" * 32 + "-" + "b" * 16 + "-01" + CALLER_TRACESTATE = "vendor=abc,other=xyz" + + def _otel(self): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + otel = OpenTelemetry() + otel.tracer = provider.get_tracer(__name__) + return otel + + def test_request_span_propagates_caller_tracestate_downstream(self): + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + from litellm.integrations.otel.plumbing.context import inject_trace_context + + inbound = {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE} + span = self._otel().create_litellm_proxy_request_started_span( + start_time=datetime.now(timezone.utc), headers=inbound + ) + outbound = inject_trace_context(inbound, parent_span=span) + span.end() + + propagated = trace.get_current_span(TraceContextTextMapPropagator().extract(outbound)).get_span_context() + self.assertEqual(outbound["tracestate"], self.CALLER_TRACESTATE) + self.assertEqual(propagated.trace_id, span.get_span_context().trace_id) + self.assertEqual(propagated.span_id, span.get_span_context().span_id) + self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT) + + def test_request_span_without_caller_tracestate_emits_none(self): + from litellm.integrations.otel.plumbing.context import inject_trace_context + + inbound = {"traceparent": self.CALLER_TRACEPARENT} + span = self._otel().create_litellm_proxy_request_started_span( + start_time=datetime.now(timezone.utc), headers=inbound + ) + outbound = inject_trace_context(inbound, parent_span=span) + span.end() + + self.assertNotIn("tracestate", outbound) + self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT) + + def test_span_context_from_header_keeps_caller_tracestate(self): + kwargs = { + "litellm_params": { + "proxy_server_request": { + "headers": {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE} + } + } + } + ctx, detected_span = self._otel()._get_span_context(kwargs) + self.assertIsNone(detected_span) + self.assertEqual(trace.get_current_span(ctx).get_span_context().trace_state.to_header(), self.CALLER_TRACESTATE) + + class TestEndProxySpanLitellmMetadataFallback(unittest.TestCase): """ Tests for _end_proxy_span_from_kwargs() falling back to litellm_metadata.