diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3b22397559f..eccba99d749 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -39,6 +39,13 @@ _EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session- # (covers UUIDs and most common session-id formats). _SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +# W3C Trace Context traceparent format: +# {version}-{trace-id}-{parent-id}-{trace-flags} +# e.g. 00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01 +_TRACEPARENT_RE = re.compile( + r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE +) + def _sanitize_for_log(value: Any) -> str: """ @@ -309,6 +316,27 @@ def _extract_generic_session_id_from_headers( return None +def _extract_trace_id_from_traceparent( + normalized: Dict[str, str], +) -> Optional[str]: + """ + Extract the ``trace-id`` component from a W3C ``traceparent`` header. + + The traceparent format is ``{version}-{trace-id}-{parent-id}-{trace-flags}``, + e.g. ``00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01``. + + Returns the 32-hex-char trace-id if the header is present and well-formed, + otherwise ``None``. + """ + traceparent = normalized.get("traceparent") + if not traceparent or not isinstance(traceparent, str): + return None + match = _TRACEPARENT_RE.match(traceparent.strip()) + if match: + return match.group(1) + return None + + def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]: """ Extract chain id for call chaining from request headers. @@ -318,6 +346,9 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str 2. ``x-litellm-session-id`` (explicit) 3. Any ``x--session-id`` header whose value looks like a session id (alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``. + 4. W3C ``traceparent`` header — the 32-hex-char trace-id is extracted and + used as the session id. This allows chaining LLM calls that are part + of one agent / distributed-trace interaction. Header keys are matched case-insensitively so this works with raw header dicts from any transport. @@ -332,6 +363,7 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str normalized.get("x-litellm-trace-id") or normalized.get("x-litellm-session-id") or _extract_generic_session_id_from_headers(normalized) + or _extract_trace_id_from_traceparent(normalized) ) 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 d2a1468be2a..790ae688b03 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2276,6 +2276,104 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): ) +def test_get_chain_id_from_headers_traceparent(): + """get_chain_id_from_headers extracts trace-id from a W3C traceparent header.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + assert ( + get_chain_id_from_headers({"traceparent": traceparent}) + == "0af7651916cd43dd8448eb211c80319c" + ) + + +def test_get_chain_id_from_headers_traceparent_case_insensitive(): + """traceparent header key matching is case-insensitive.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + traceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-00" + assert ( + get_chain_id_from_headers({"Traceparent": traceparent}) + == "abcdef1234567890abcdef1234567890" + ) + + +def test_get_chain_id_from_headers_traceparent_malformed_ignored(): + """Malformed traceparent values are ignored.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert get_chain_id_from_headers({"traceparent": "not-a-valid-traceparent"}) is None + assert get_chain_id_from_headers({"traceparent": ""}) is None + assert get_chain_id_from_headers({"traceparent": "00-short-abc-01"}) is None + + +def test_get_chain_id_from_headers_traceparent_lower_priority_than_explicit(): + """Explicit litellm headers take precedence over traceparent.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + { + "x-litellm-trace-id": "explicit-id", + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + } + ) + == "explicit-id" + ) + + +def test_get_chain_id_from_headers_traceparent_lower_priority_than_generic_session(): + """Generic x--session-id headers take precedence over traceparent.""" + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + assert ( + get_chain_id_from_headers( + { + "x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01", + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + } + ) + == "e96634a3-fa28-4083-b354-55542e2dca01" + ) + + +def test_add_litellm_metadata_from_request_headers_traceparent_sets_session_id(): + """traceparent header sets session_id and trace_id in metadata.""" + traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + headers = {"traceparent": traceparent} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + expected_trace_id = "0af7651916cd43dd8448eb211c80319c" + assert data["metadata"]["session_id"] == expected_trace_id + assert data["metadata"]["trace_id"] == expected_trace_id + assert data["litellm_session_id"] == expected_trace_id + assert data["litellm_trace_id"] == expected_trace_id + + +def test_extract_trace_id_from_traceparent_directly(): + """Unit test for _extract_trace_id_from_traceparent helper.""" + from litellm.proxy.litellm_pre_call_utils import _extract_trace_id_from_traceparent + + assert ( + _extract_trace_id_from_traceparent( + {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"} + ) + == "0af7651916cd43dd8448eb211c80319c" + ) + assert _extract_trace_id_from_traceparent({}) is None + assert _extract_trace_id_from_traceparent({"traceparent": ""}) is None + assert _extract_trace_id_from_traceparent({"traceparent": "invalid-format"}) is None + # Uppercase hex should also work + assert ( + _extract_trace_id_from_traceparent( + {"traceparent": "00-0AF7651916CD43DD8448EB211C80319C-B7AD6B7169203331-01"} + ) + == "0AF7651916CD43DD8448EB211C80319C" + ) + + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},