Merge 9e09a5ecc5 into f5a1c9f1f1

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -882,9 +882,8 @@ 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."""
+    """Resolve the session id for the spend log row: `omit` honors only metadata.session_id,
+    batch sessions outrank everything, then the caller's litellm_session_id, then trace ids."""
     if omit_when_missing:
         session_id: Final = metadata.get("session_id") if metadata else None
         return str(session_id) if session_id else None
@@ -893,6 +892,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:
+        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:

tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py

@@ -237,6 +237,27 @@ 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():
+    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():
+    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={},