Merge 5274719d66 into 0c98afa780

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

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -978,6 +978,7 @@ async def pass_through_request(
             headers=headers,
             forward_headers=forward_headers,
         )
+        headers = _with_trace_context(headers, inbound_headers=_safe_get_request_headers(request))
 
         requested_query_params: dict | None = query_params or dict(request.query_params)
 
@@ -2132,6 +2133,17 @@ 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"))
+
+
+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,
@@ -2174,20 +2186,16 @@ 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()
-
-    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: 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 _WEBSOCKET_FORWARDED_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(

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
 

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
@@ -45,6 +47,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
@@ -4242,6 +4253,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():
     """
@@ -4838,6 +4883,65 @@ 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
+@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
+    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 = {"authorization": "Bearer client"}
+    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=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"].get("authorization") == ("Bearer client" if forward_headers else None)
+
+
 class ClosingUpstreamWebSocket:
     def __init__(self, close_exc: Exception):
         self._close_exc = close_exc