diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8fc5faee2c9..60748234873 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -150,6 +150,21 @@ def _session_id_from_baggage(baggage: str) -> str | None: return None +def _client_steered_trace_field(data: dict, field: str) -> bool: + """Did the caller explicitly steer ``field`` in the request body? + + ``metadata`` (``litellm_metadata`` on ``LITELLM_METADATA_ROUTES``) is the + documented way to set ``trace_id`` / ``session_id``, so a value there is an + explicit caller choice and must outrank the traceparent/baggage fallback + below, which only guards on the top-level ``litellm_*`` keys. + """ + for variable_name in ("metadata", "litellm_metadata"): + container = data.get(variable_name) + if isinstance(container, dict) and container.get(field) is not None: + return True + return False + + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable via_virtual_key marker AND a known non-secret shape: the sha256 hex digest @@ -1574,12 +1589,13 @@ class LiteLLMProxyRequestSetup: # Last-resort fallback: the W3C standards for trace/session propagation # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/). # Lower priority than everything above - only fires when neither the - # explicit litellm headers nor the Anthropic-metadata path found - # anything - but lets a caller's existing traceparent/baggage headers + # explicit litellm headers, the Anthropic-metadata path, nor the + # caller's own request-body metadata set the field - but lets a + # caller's existing traceparent/baggage headers # (from real OTel instrumentation) correlate with litellm's own logs # instead of generating an unrelated trace_id. normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) - if "litellm_trace_id" not in data: + if "litellm_trace_id" not in data and not _client_steered_trace_field(data, "trace_id"): traceparent: Final = normalized_headers.get("traceparent") if isinstance(traceparent, str): trace_id_from_traceparent: Final = _trace_id_from_traceparent(traceparent) @@ -1589,7 +1605,7 @@ class LiteLLMProxyRequestSetup: verbose_proxy_logger.debug( "Extracted trace_id from W3C traceparent header: %s", trace_id_from_traceparent ) - if "litellm_session_id" not in data: + if "litellm_session_id" not in data and not _client_steered_trace_field(data, "session_id"): baggage: Final = normalized_headers.get("baggage") if isinstance(baggage, str): session_id_from_baggage: Final = _session_id_from_baggage(baggage) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index b2241191ced..bd350f56410 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -3648,6 +3648,61 @@ def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_trace assert data["litellm_session_id"] == "explicit-trace-id-value" +def test_add_litellm_metadata_from_request_headers_body_trace_id_beats_traceparent(): + """A caller that set metadata.trace_id keeps it: the traceparent fallback is + documented as last-resort, so it must not overwrite an explicit choice. + + This matters in practice because some platforms (e.g. GCP) inject a + traceparent into every inbound request, so the fallback would otherwise fire + on traffic whose caller never sent the header at all.""" + headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + data = {"metadata": {"trace_id": "caller-chosen-trace-id"}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["trace_id"] == "caller-chosen-trace-id" + assert "litellm_trace_id" not in data + + +def test_add_litellm_metadata_from_request_headers_body_session_id_beats_baggage(): + """Same for session_id and the baggage header.""" + headers = {"baggage": "session.id=baggage-session-42"} + data = {"metadata": {"session_id": "caller-chosen-session-id"}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["session_id"] == "caller-chosen-session-id" + assert "litellm_session_id" not in data + + +def test_add_litellm_metadata_from_request_headers_body_steering_is_per_field(): + """Steering one field must not suppress the fallback for the other: + a caller setting only trace_id still gets session_id from baggage.""" + headers = { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "session.id=baggage-session-42", + } + data = {"metadata": {"trace_id": "caller-chosen-trace-id"}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["metadata"]["trace_id"] == "caller-chosen-trace-id" + assert data["litellm_session_id"] == "baggage-session-42" + + +def test_add_litellm_metadata_from_request_headers_litellm_metadata_steering_honoured(): + """Routes in LITELLM_METADATA_ROUTES (/v1/responses, /v1/messages, batches, + files) carry their metadata in litellm_metadata, so steering there counts + the same as steering in metadata.""" + headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + data = {"litellm_metadata": {"trace_id": "caller-chosen-trace-id"}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="litellm_metadata" + ) + assert data["litellm_metadata"]["trace_id"] == "caller-chosen-trace-id" + assert "litellm_trace_id" not in data + + def _otel_span_with_trace_id(trace_id: int) -> NonRecordingSpan: return NonRecordingSpan(SpanContext(trace_id=trace_id, span_id=0x00F067AA0BA902B7, is_remote=False))