fix(proxy): honor caller litellm_session_id in spend log session_id

The default path of _get_session_id_for_spend_log never consulted
kwargs['litellm_session_id'], so a caller-supplied session was
overwritten by the per-call trace id and LiteLLM_SpendLogs.session_id
could never group by litellm_session_id. Prefer it after the batch
lifecycle's deterministic session: when the caller sent no session the
value is only the derived trace-id copy, so sessionless requests are
unchanged.

Fixes #40851
This commit is contained in:
Deepanshu Pal 2026-09-12 18:39:14 +05:30
parent 30f33a949b
commit c07e9432ed
2 changed files with 34 additions and 3 deletions

View file

@ -682,9 +682,12 @@ def _get_session_id_for_spend_log(
omit_when_missing: bool,
batch_trace_session_id: str | None = None,
) -> str | None:
"""Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may
be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks
the per-request trace ids because those differ between the create call and the cost poller's row."""
"""Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session. Batch call types
carry a deterministic session derived from the batch id, which outranks everything else because the
per-request trace ids differ between the create call and the cost poller's row. Otherwise the caller's
`litellm_session_id` wins: it is only absent or a copy of the trace id when the caller did not send one,
so preferring it changes nothing for sessionless requests while letting a real session id reach the
spend log instead of being overwritten by the per-call trace id."""
if omit_when_missing:
session_id: Final = metadata.get("session_id") if metadata else None
return str(session_id) if session_id else None
@ -693,6 +696,9 @@ def _get_session_id_for_spend_log(
if batch_trace_session_id is not None:
return batch_trace_session_id
caller_session_id: Final = kwargs.get("litellm_session_id")
if caller_session_id is not None:
return str(caller_session_id)
if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
return str(standard_logging_payload.get("trace_id"))
if kwargs.get("litellm_trace_id") is not None:

View file

@ -233,6 +233,31 @@ def test_batch_session_outranks_the_per_request_trace_id():
assert session_id == "batch-uid-1"
def test_caller_litellm_session_id_wins_over_the_per_request_trace_id():
"""A caller-supplied litellm_session_id is a real session, not a copied trace id; the
spend log must keep it instead of overwriting it with the per-call trace id."""
session_id: Final = _get_session_id_for_spend_log(
kwargs={"litellm_session_id": "sess-1", "litellm_trace_id": "trace-abc"},
metadata={"trace_id": "trace-abc"},
standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
omit_when_missing=False,
)
assert session_id == "sess-1"
def test_batch_session_outranks_a_caller_litellm_session_id():
"""The batch lifecycle grouping must survive even when the create call carries a
caller session id, or the poller's cost row would land in a different session."""
session_id: Final = _get_session_id_for_spend_log(
kwargs={"litellm_session_id": "sess-1", "litellm_trace_id": "trace-abc"},
metadata={"trace_id": "trace-abc"},
standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
omit_when_missing=False,
batch_trace_session_id="batch-uid-1",
)
assert session_id == "batch-uid-1"
def test_omit_policy_still_suppresses_batch_sessions():
session_id: Final = _get_session_id_for_spend_log(
kwargs={},