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 ea4ede7e513..2b545ae83fc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/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( 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 d57bed430c1..6d941bca485 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 @@ -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