fix: align top-level session/trace ID fields with body trace_id

This commit is contained in:
Sahil 2026-04-07 03:11:09 +05:30
parent 4043104d43
commit e29b55f409
2 changed files with 30 additions and 4 deletions

View file

@ -660,15 +660,22 @@ class LiteLLMProxyRequestSetup:
if chain_id:
metadata_from_headers["trace_id"] = chain_id
metadata_from_headers["session_id"] = chain_id
data["litellm_session_id"] = chain_id
data["litellm_trace_id"] = chain_id
# Prefer body trace_id over header so all three fields stay consistent.
_body_metadata = data.get(_metadata_variable_name)
_body_trace_id = (
_body_metadata.get("trace_id")
if isinstance(_body_metadata, dict)
else None
)
effective_id = _body_trace_id or chain_id
data["litellm_session_id"] = effective_id
data["litellm_trace_id"] = effective_id
verbose_proxy_logger.debug(
f"Extracted chain_id from header (trace-id/session-id): {chain_id}"
)
if isinstance(data[_metadata_variable_name], dict):
# Only inject header-derived values for keys NOT already set by the user
# in the request body. This ensures body values take priority over headers.
# Body values take priority — only inject from headers if the key isn't already set.
for key, value in metadata_from_headers.items():
if key not in data[_metadata_variable_name]:
data[_metadata_variable_name][key] = value

View file

@ -1956,6 +1956,25 @@ class TestAddLitellmMetadataFromRequestHeaders:
assert result["litellm_metadata"]["trace_id"] == "from-header"
assert result["litellm_metadata"]["session_id"] == "from-header"
def test_should_keep_top_level_ids_consistent_with_body_trace_id(self):
"""litellm_session_id and litellm_trace_id should also reflect the body trace_id."""
data = {
"model": "anthropic/claude-3-5-sonnet",
"messages": [{"role": "user", "content": "hi"}],
"litellm_metadata": {"trace_id": "from-body"},
}
headers = {"x-litellm-trace-id": "from-header", "content-type": "application/json"}
result = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=headers,
data=data,
_metadata_variable_name="litellm_metadata",
)
assert result["litellm_metadata"]["trace_id"] == "from-body"
assert result["litellm_session_id"] == "from-body"
assert result["litellm_trace_id"] == "from-body"
@pytest.mark.asyncio
async def test_should_preserve_body_trace_id_in_full_pipeline_on_messages_route(self):
"""Body trace_id must survive the full add_litellm_data_to_request pipeline."""