fix(fireworks_ai): stop using the trace id as the session affinity key (#35754)

get_fireworks_session_id fell back to litellm_trace_id when no session id was
given. That id is generated per request (uuid4 when absent), so x-session-affinity
carried a different value every time and Fireworks prompt caching never hit;
cached_tokens stayed 0 across identical prompts.

The None path the original change described was effectively unreachable because
of it. Drop the fallback so affinity comes only from an id the caller actually
supplied: litellm_session_id, session_id, or metadata.session_id.

Callers who were relying on a trace id for affinity can pass litellm_session_id
instead, which is stable across the requests they want grouped.

Co-authored-by: mubashir1osmani <mubashir.osmani777@gmail.com>

litellm/llms/fireworks_ai/common_utils.py

@@ -13,6 +13,13 @@ class FireworksAIException(BaseLLMException):
 
 
 def get_fireworks_session_id(litellm_params: dict) -> str | None:
+    """
+    Session id to send as `x-session-affinity`, or None when the caller gave none.
+
+    Deliberately does not fall back to `litellm_trace_id`: that is generated per
+    request (`str(uuid.uuid4())` when absent), so using it pins every request to a
+    different Fireworks node and prompt caching never hits.
+    """
     params: Final = litellm_params
     for key in ("litellm_session_id", "session_id"):
         value = params.get(key)
@@ -23,9 +30,6 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
         value = metadata.get("session_id")
         if value:
             return str(value)
-    value = params.get("litellm_trace_id")
-    if value:
-        return str(value)
     return None
 
 

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -73,7 +73,13 @@ def test_validate_environment_sets_session_affinity_from_session_id():
     assert headers["x-session-affinity"] == "session-id-123"
 
 
-def test_validate_environment_sets_session_affinity_from_trace_id():
+def test_validate_environment_ignores_trace_id_for_session_affinity():
+    """A trace id must not become the session id.
+
+    litellm_trace_id defaults to a fresh uuid4 per request, so pinning
+    x-session-affinity to it sent every request to a different Fireworks node and
+    prompt caching never hit (cached_tokens stayed 0 across identical prompts).
+    """
     config = FireworksAIConfig()
 
     headers = config.validate_environment(
@@ -85,7 +91,25 @@ def test_validate_environment_sets_session_affinity_from_trace_id():
         api_key="test-key",
     )
 
-    assert headers["x-session-affinity"] == "trace-id-123"
+    assert "x-session-affinity" not in headers
+
+
+def test_validate_environment_prefers_session_id_over_trace_id():
+    config = FireworksAIConfig()
+
+    headers = config.validate_environment(
+        headers={},
+        model="accounts/fireworks/models/test-model",
+        messages=[],
+        optional_params={},
+        litellm_params={
+            "litellm_session_id": "session-123",
+            "litellm_trace_id": "trace-id-123",
+        },
+        api_key="test-key",
+    )
+
+    assert headers["x-session-affinity"] == "session-123"
 
 
 def test_validate_environment_does_not_set_session_affinity_without_session_id():