fix: gate body trace_id preference to litellm_metadata routes only

This commit is contained in:
Sahil 2026-04-07 03:24:45 +05:30
parent e29b55f409
commit 4a3647b986
2 changed files with 35 additions and 8 deletions

View file

@ -660,13 +660,17 @@ class LiteLLMProxyRequestSetup:
if chain_id:
metadata_from_headers["trace_id"] = chain_id
metadata_from_headers["session_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
)
# Only prefer body trace_id on litellm_metadata routes (/v1/messages, threads, etc.)
# On /chat/completions, metadata["trace_id"] belongs to observability tools
# (e.g. LangFuse) and should not override the header chain ID.
_body_trace_id: Optional[str] = None
if _metadata_variable_name == "litellm_metadata":
_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
@ -675,7 +679,8 @@ class LiteLLMProxyRequestSetup:
)
if isinstance(data[_metadata_variable_name], dict):
# Body values take priority — only inject from headers if the key isn't already set.
# 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.
for key, value in metadata_from_headers.items():
if key not in data[_metadata_variable_name]:
data[_metadata_variable_name][key] = value

View file

@ -1975,6 +1975,28 @@ class TestAddLitellmMetadataFromRequestHeaders:
assert result["litellm_session_id"] == "from-body"
assert result["litellm_trace_id"] == "from-body"
def test_should_not_bleed_metadata_trace_id_into_chain_ids_on_chat_completions_route(self):
"""On /chat/completions, metadata["trace_id"] (e.g. LangFuse) must not override the header chain ID."""
data = {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"trace_id": "langfuse-id"},
}
headers = {"x-litellm-trace-id": "chain-id", "content-type": "application/json"}
result = LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=headers,
data=data,
_metadata_variable_name="metadata",
)
# Header chain ID must win for top-level fields on this route.
assert result["litellm_session_id"] == "chain-id"
assert result["litellm_trace_id"] == "chain-id"
# The observability trace_id in the metadata dict must be untouched.
assert result["metadata"]["trace_id"] == "langfuse-id"
@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."""