mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(proxy): traceparent/baggage fallback must not override caller metadata (#43688)
* fix(proxy): traceparent/baggage fallback must not override caller metadata
The W3C traceparent/baggage fallback in
add_litellm_metadata_from_request_headers documents itself as last-resort:
Lower priority than everything above - only fires when neither the
explicit litellm headers nor the Anthropic-metadata path found anything
But it guards on the top-level `litellm_trace_id` / `litellm_session_id` body
keys and never checks `metadata`, which is the documented way callers set
trace_id / session_id (`metadata: {"trace_id": ...}` on /chat/completions).
A caller that explicitly sets metadata.trace_id has it silently replaced by the
header value, so the implementation contradicts its own stated precedence.
This is not a corner case on managed platforms: GCP's front end injects a
traceparent into every inbound request, so the fallback fires on traffic whose
caller never sent the header. The request still returns 200 and the trace still
reaches the logging backend, just under an id the caller never chose, so any
caller correlating by its own id silently fails to find its trace.
Guard both fallbacks on the caller's request-body metadata as well.
Deliberately narrow:
- x-litellm-trace-id still outranks the body (documented priority #1)
- a request that steers neither field still adopts traceparent/baggage exactly
as before
- steering is per-field: setting only trace_id still lets session_id come from
baggage
- litellm_metadata is checked too, for LITELLM_METADATA_ROUTES (/v1/responses,
/v1/messages, batches, files)
Also corrects the comment, which understated the guard.
4 new tests; each fails without the source change. The existing traceparent and
baggage tests are unchanged and still pass.
* fix(proxy): check only the active metadata container and ignore empty values
Addresses review. The first version treated any value in either `metadata` or
`litellm_metadata` as caller steering. On LITELLM_METADATA_ROUTES the body
`metadata` is provider-facing and is only promoted later, so a session_id there
suppressed baggage before apply_missing_session_id_policy ran, and a request
with a usable baggage session id got a 400 under `missing_session_id: reject`.
An empty session_id did the same, since the policy treats "" as absent
Now only the active metadata container counts, and only a truthy value, which
matches how apply_missing_session_id_policy decides a session id is present.
The helper takes a typed `object` rather than a bare dict
Adds regression tests for both cases, including the end to end reject path,
and drops the test docstrings
* fix(proxy): count promoted caller trace ids on litellm_metadata routes
Addresses review. On LITELLM_METADATA_ROUTES the proxy promotes the caller's
trace control fields (trace_id, session_id, ...) from `metadata` into
`litellm_metadata`, but only after the header fallback runs. Checking only
`litellm_metadata` let traceparent and baggage fill those fields first, and the
promotion then skipped them because they were already set, so the trace was
recorded under the header ids
The check now also reads `metadata` for fields in
LITELLM_TRACE_CONTROL_METADATA_FIELDS on those routes, and
apply_missing_session_id_policy uses the same check, so `reject` no longer
refuses a request whose session id is about to be promoted
The earlier test asserting that a session_id in `metadata` must not block
baggage on /v1/responses had the premise wrong, since that value is promoted.
It is replaced by tests that assert the promoted caller ids win on
/v1/responses and /v1/messages, with no policy, reject, and generate
* fix(proxy): an empty trace field in litellm_metadata shadows the promoted one
Addresses review. Promotion copies a trace control field from `metadata` into
`litellm_metadata` only when the key is absent, so a key that is present but
empty in `litellm_metadata` wins and the `metadata` value is never used. The
check still looked at `metadata` in that case, which let
`missing_session_id: reject` accept a request that ends up with an empty
session id, and let an empty trace_id or session_id block the headers
The check now follows the same rule as promotion. If the active container has
the key, its value decides. Only when the key is absent does the promoted
`metadata` value count. This makes the conflicting-bucket cases behave exactly
as on main
* fix(proxy): stay within the LIT006 cast budget
The type-discipline gate failed: this branch added 4 unsuppressed `cast()`
calls and LIT006 was already at its limit. Each cast now carries a
`# cast-ok` reason. The two in the helper follow an `isinstance` check that
proves the Mapping. The two at the call sites stay because the method's
`data` parameter is a bare `dict`, and dropping those casts adds
basedpyright unknown-type errors instead. No behavior change
* test(proxy): cover caller metadata precedence over W3C trace headers
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): generated session uses promoted caller trace id
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(proxy): move test context into assertion messages
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(integration): audit cells for w3c fallback precedence
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(integration): tighten audit chaos cells
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style(test): drop stray blank lines
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(integration): ignore body-less model-info probes in langfuse precedence upstreams
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): keep root trace/session fields on equal ids and usable caller values
The _caller_trace_field gating added for caller-metadata precedence was
presence-aware, which regressed three pre-call behaviors versus main:
- Equal ids in W3C headers and body metadata suppressed the
traceparent/baggage branch entirely, leaving the root
litellm_trace_id/litellm_session_id unset so router fallbacks minted a
fresh uuid4 per request. The header branch now also fires when the
caller value equals the header-derived id, so equal ids stamp the root
fields exactly like main.
- A truthy but non-string metadata value (e.g. session_id 4815162342)
counted as caller-supplied and suppressed the baggage fallback, so
downstream str-only consumers (code interpreter sandbox reuse) got a
new sandbox per turn. _caller_trace_field now counts only non-empty
string values, and the generate policy falls through to generation
when the caller value is not usable.
- On litellm_metadata routes a usable caller session id suppressed the
missing_session_id policies while never landing on the root field, so
the root session stayed unset. The policy now promotes the caller's
usable session id to litellm_session_id instead of leaving it stranded.
* test(integration): non-string caller ids fall back to W3C like empty ids
The audit cell compared a headers leg against a headerless leg for
equality, which only held while a truthy non-string id suppressed the
W3C fallback on both legs. With unusable values ignored again, the
headers leg resolves to the header ids while the headerless leg cannot,
so assert the headers leg's concrete outcome instead, mirroring the
empty/null cell.
* style(proxy): trim the session promotion comment to the non-obvious why
---------
Co-authored-by: Filipe Andujar <filipeandujar@gmail.com>
Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
dc2a14a656
commit
02b8b5cc80
3 changed files with 1857 additions and 15 deletions
|
|
@ -154,6 +154,27 @@ def _session_id_from_baggage(baggage: str) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str, field: str) -> str | None:
|
||||
"""The caller's value for a trace-control field, counted only when it is a
|
||||
usable id: a non-empty string. An explicitly empty/unusable value on the
|
||||
active metadata container still shadows the promoted requester value, but
|
||||
neither ever counts as "the caller supplied this field" on its own, so a
|
||||
numeric session id or an empty string cannot suppress the W3C header
|
||||
fallback or satisfy a missing-session-id policy."""
|
||||
active: Final = data.get(metadata_variable_name)
|
||||
if isinstance(active, Mapping) and field in active:
|
||||
active_map: Final = cast(Mapping[str, object], active) # cast-ok: isinstance above, free-form JSON values
|
||||
active_value: Final = active_map[field]
|
||||
return active_value if isinstance(active_value, str) and active_value else None
|
||||
promoted: Final = metadata_variable_name == "litellm_metadata" and field in LITELLM_TRACE_CONTROL_METADATA_FIELDS
|
||||
requester: Final = data.get("metadata")
|
||||
if not promoted or not isinstance(requester, Mapping):
|
||||
return None
|
||||
requester_map: Final = cast(Mapping[str, object], requester) # cast-ok: isinstance above, free-form JSON values
|
||||
requester_value: Final = requester_map.get(field)
|
||||
return requester_value if isinstance(requester_value, str) and requester_value else None
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -839,11 +860,22 @@ def apply_missing_session_id_policy(
|
|||
):
|
||||
metadata["session_id"] = body_session_id
|
||||
return
|
||||
if data.get("litellm_session_id") or metadata.get("session_id"):
|
||||
caller_session_id: Final = _caller_trace_field(data, _metadata_variable_name, "session_id")
|
||||
if caller_session_id is not None:
|
||||
# Consumers that read the root field (router fallbacks, spend logs,
|
||||
# sandbox reuse) otherwise see no session and mint a uuid4 per request.
|
||||
if not data.get("litellm_session_id"):
|
||||
data["litellm_session_id"] = caller_session_id # rebind-ok: data is an out-param
|
||||
return
|
||||
if data.get("litellm_session_id"):
|
||||
return
|
||||
match policy:
|
||||
case "generate":
|
||||
session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4())
|
||||
session_id: Final = str(
|
||||
data.get("litellm_trace_id")
|
||||
or _caller_trace_field(data, _metadata_variable_name, "trace_id")
|
||||
or uuid.uuid4()
|
||||
)
|
||||
data["litellm_session_id"] = session_id # rebind-ok: data is an out-param
|
||||
data.setdefault("litellm_trace_id", session_id)
|
||||
metadata["session_id"] = session_id
|
||||
|
|
@ -1596,16 +1628,28 @@ 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
|
||||
# (from real OTel instrumentation) correlate with litellm's own logs
|
||||
# instead of generating an unrelated trace_id.
|
||||
# explicit litellm headers, the Anthropic-metadata path, nor the
|
||||
# caller's own request metadata set the field to a DIFFERENT usable id
|
||||
# - 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:
|
||||
traceparent: Final = normalized_headers.get("traceparent")
|
||||
if isinstance(traceparent, str):
|
||||
trace_id_from_traceparent: Final = _trace_id_from_traceparent(traceparent)
|
||||
if trace_id_from_traceparent:
|
||||
# The caller's metadata wins over the header fallback unless
|
||||
# both carry the same id: stamping the root field then claims
|
||||
# nothing the caller did not already ask for, and keeps the
|
||||
# W3C-correlated root trace id instead of a generated uuid4.
|
||||
caller_trace_id: Final = _caller_trace_field(
|
||||
cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
|
||||
_metadata_variable_name,
|
||||
"trace_id",
|
||||
)
|
||||
if trace_id_from_traceparent and (
|
||||
caller_trace_id is None or caller_trace_id == trace_id_from_traceparent
|
||||
):
|
||||
metadata_from_headers["trace_id"] = trace_id_from_traceparent
|
||||
data["litellm_trace_id"] = trace_id_from_traceparent # rebind-ok: data is an out-param
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -1615,7 +1659,14 @@ class LiteLLMProxyRequestSetup:
|
|||
baggage: Final = normalized_headers.get("baggage")
|
||||
if isinstance(baggage, str):
|
||||
session_id_from_baggage: Final = _session_id_from_baggage(baggage)
|
||||
if session_id_from_baggage:
|
||||
caller_session_id: Final = _caller_trace_field(
|
||||
cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
|
||||
_metadata_variable_name,
|
||||
"session_id",
|
||||
)
|
||||
if session_id_from_baggage and (
|
||||
caller_session_id is None or caller_session_id == session_id_from_baggage
|
||||
):
|
||||
metadata_from_headers["session_id"] = session_id_from_baggage
|
||||
data["litellm_session_id"] = session_id_from_baggage # rebind-ok: data is an out-param
|
||||
verbose_proxy_logger.debug("Extracted session_id from W3C baggage header")
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -3674,6 +3674,106 @@ 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():
|
||||
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():
|
||||
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():
|
||||
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():
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("empty_session_id", ["", None])
|
||||
def test_add_litellm_metadata_from_request_headers_empty_body_session_id_falls_back_to_baggage(
|
||||
empty_session_id: str | None,
|
||||
):
|
||||
headers = {"baggage": "session.id=baggage-session-42"}
|
||||
data = {"metadata": {"session_id": empty_session_id}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
assert data["metadata"]["session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["trace_id", "session_id"])
|
||||
def test_add_litellm_metadata_from_request_headers_promoted_metadata_beats_headers(field: str):
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
data = {"metadata": {field: "caller-chosen"}, "litellm_metadata": {}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
assert field not in data["litellm_metadata"]
|
||||
assert f"litellm_{field}" not in data
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_equal_ids_still_stamp_root_fields():
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=matching-session-42",
|
||||
}
|
||||
data = {
|
||||
"metadata": {"trace_id": "4bf92f3577b34da6a3ce929d0e0e4736", "session_id": "matching-session-42"},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["litellm_session_id"] == "matching-session-42"
|
||||
assert data["metadata"]["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["metadata"]["session_id"] == "matching-session-42"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("non_string_session_id", [4815162342, True, {"session": "nested"}])
|
||||
def test_add_litellm_metadata_from_request_headers_non_string_body_session_id_falls_back_to_baggage(
|
||||
non_string_session_id: object,
|
||||
):
|
||||
headers = {"baggage": "session.id=header-session-42"}
|
||||
data = {"metadata": {"session_id": non_string_session_id}}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="metadata"
|
||||
)
|
||||
assert data["litellm_session_id"] == "header-session-42"
|
||||
assert data["metadata"]["session_id"] == "header-session-42"
|
||||
|
||||
|
||||
def _otel_span_with_trace_id(trace_id: int) -> NonRecordingSpan:
|
||||
return NonRecordingSpan(SpanContext(trace_id=trace_id, span_id=0x00F067AA0BA902B7, is_remote=False))
|
||||
|
||||
|
|
@ -8222,6 +8322,85 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id():
|
|||
assert _spend_log_session_id(updated) == "client-session-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "client_body"),
|
||||
[
|
||||
("/v1/chat/completions", {"model": "gpt-4o", "messages": [], "metadata": {"session_id": ""}}),
|
||||
("/v1/responses", {"model": "gpt-4o", "input": "hi"}),
|
||||
],
|
||||
)
|
||||
async def test_missing_session_id_reject_accepts_baggage_session_id(path: str, client_body: dict[str, object]):
|
||||
request = _request_for(path)
|
||||
request.headers = {"baggage": "session.id=baggage-session-42"}
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data=client_body,
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "reject"},
|
||||
)
|
||||
|
||||
assert updated["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
|
||||
@pytest.mark.parametrize("policy", [None, "reject", "generate"])
|
||||
async def test_promoted_caller_trace_ids_beat_traceparent_and_baggage(path: str, policy: str | None):
|
||||
request = _request_for(path)
|
||||
request.headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-4o", "metadata": {"trace_id": "caller-trace", "session_id": "caller-session"}},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": policy} if policy else {},
|
||||
)
|
||||
|
||||
assert updated["litellm_metadata"]["trace_id"] == "caller-trace"
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-session"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_reject_ignores_requester_session_id_shadowed_by_empty_litellm_metadata():
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"metadata": {"session_id": "caller-session"},
|
||||
"litellm_metadata": {"session_id": ""},
|
||||
},
|
||||
request=_request_for("/v1/responses"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "reject"},
|
||||
)
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
|
||||
def test_add_litellm_metadata_from_request_headers_empty_litellm_metadata_field_falls_back_to_headers():
|
||||
headers = {
|
||||
"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
|
||||
"baggage": "session.id=baggage-session-42",
|
||||
}
|
||||
data = {
|
||||
"metadata": {"trace_id": "caller-trace", "session_id": "caller-session"},
|
||||
"litellm_metadata": {"trace_id": "", "session_id": ""},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
assert data["litellm_session_id"] == "baggage-session-42"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"client_body",
|
||||
|
|
@ -8341,6 +8520,91 @@ async def test_missing_session_id_generate_reuses_traceparent_trace_id():
|
|||
assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
|
||||
{},
|
||||
],
|
||||
ids=["with_traceparent", "no_traceparent"],
|
||||
)
|
||||
async def test_missing_session_id_generate_reuses_promoted_caller_trace_id(path: str, headers: dict[str, str]):
|
||||
request = _request_for(path)
|
||||
request.headers = headers
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-4o", "input": "hi", "metadata": {"trace_id": "caller-trace"}},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "generate"},
|
||||
)
|
||||
|
||||
caller_trace_msg: Final = "generate must derive the session from the caller's un-promoted metadata.trace_id"
|
||||
assert updated["litellm_session_id"] == "caller-trace", caller_trace_msg
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-trace", caller_trace_msg
|
||||
assert updated["litellm_metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True, (
|
||||
"the derived session id must still be marked generated"
|
||||
)
|
||||
assert _spend_log_session_id(updated, "litellm_metadata") == "caller-trace", (
|
||||
"spend log and callback session ids must agree on the caller trace id"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("policy", ["generate", "reject"])
|
||||
async def test_missing_session_id_policy_promotes_caller_session_to_root_field(policy: str):
|
||||
"""A caller-supplied usable session id satisfies the missing_session_id policies on
|
||||
litellm_metadata routes (where it is not yet the managed metadata field) and must also
|
||||
land on the root ``litellm_session_id`` field: consumers that read the root field
|
||||
(router fallbacks, spend logs, sandbox reuse) otherwise mint a fresh uuid4 per request."""
|
||||
request = _request_for("/v1/responses")
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"litellm_trace_id": "root-trace-42",
|
||||
"metadata": {"session_id": "caller-session-42"},
|
||||
},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": policy},
|
||||
)
|
||||
|
||||
assert updated["litellm_trace_id"] == "root-trace-42"
|
||||
assert updated["litellm_session_id"] == "caller-session-42"
|
||||
assert updated["litellm_metadata"]["session_id"] == "caller-session-42"
|
||||
assert SESSION_ID_GENERATED_METADATA_KEY not in updated["litellm_metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_session_id_generate_ignores_non_string_caller_session_id():
|
||||
"""A non-string session id is not a usable session: the generate policy must fall through
|
||||
to generation instead of letting an unusable value strand the root session field."""
|
||||
request = _request_for("/v1/responses")
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-4o",
|
||||
"input": "hi",
|
||||
"litellm_trace_id": "root-trace-42",
|
||||
"metadata": {"session_id": 4815162342},
|
||||
},
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={"missing_session_id": "generate"},
|
||||
)
|
||||
|
||||
assert updated["litellm_session_id"] == "root-trace-42"
|
||||
assert updated["litellm_metadata"]["session_id"] == "root-trace-42"
|
||||
assert updated["litellm_metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("policy", ["generate", "reject"])
|
||||
async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue