Merge b43135a445 into aeede53112

litellm/proxy/litellm_pre_call_utils.py

@@ -991,14 +991,17 @@ def _extract_bare_session_id_from_headers(
     normalized: Mapping[str, str],
 ) -> str | None:
     """
-    Read a vendor-less ``x-session-id`` header (opencode sends ``X-Session-Id``
-    alongside ``x-session-affinity`` on every turn of a session). Checked after
-    the ``x-<vendor>-session-id`` scan so a more specific header such as
+    Read a vendor-less ``x-session-id`` or ``x-session-affinity`` header (opencode sends
+    ``x-session-affinity``, occasionally alongside ``X-Session-Id``, on every turn of a session).
+    Checked after the ``x-<vendor>-session-id`` scan so a more specific header such as
     opencode's ``x-parent-session-id`` on subagent calls keeps winning.
     """
     value: Final = normalized.get("x-session-id")
     if isinstance(value, str) and _SESSION_ID_VALUE_RE.match(value):
         return value
+    fallback: Final = normalized.get("x-session-affinity")
+    if isinstance(fallback, str) and _SESSION_ID_VALUE_RE.match(fallback):
+        return fallback
     return None
 
 
@@ -1012,7 +1015,7 @@ def get_chain_id_from_headers(headers: dict[str, str] | None) -> str | None:
     3. Any ``x-<vendor>-session-id`` header whose value looks like a session id
        (alphanumeric / UUID, at least 8 chars).  E.g. ``x-claude-code-session-id``.
     4. Codex's unprefixed ``session-id`` / ``thread-id``, for Codex callers only.
-    5. A vendor-less ``x-session-id`` header (e.g. opencode), same value rules.
+    5. A vendor-less ``x-session-id`` or ``x-session-affinity`` header (e.g. opencode), same value rules.
 
     Header keys are matched case-insensitively so this works with raw header
     dicts from any transport.

tests/unit/proxy/test_litellm_pre_call_utils.py

@@ -4181,6 +4181,125 @@ def test_add_litellm_metadata_groups_opencode_turns_into_one_session():
         assert turn["litellm_trace_id"] == OPENCODE_SESSION_ID
 
 
+def test_add_litellm_metadata_groups_opencode_turns_with_session_affinity_only():
+    turns: Final = [{"metadata": {}}, {"metadata": {}}]
+    headers: Final = {"x-session-affinity": OPENCODE_SESSION_ID, "User-Agent": "opencode/1.18.28"}
+    for turn in turns:
+        LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+            headers=headers, data=turn, _metadata_variable_name="metadata"
+        )
+
+    for turn in turns:
+        assert turn["metadata"]["session_id"] == OPENCODE_SESSION_ID
+        assert turn["metadata"]["trace_id"] == OPENCODE_SESSION_ID
+        assert turn["litellm_session_id"] == OPENCODE_SESSION_ID
+        assert turn["litellm_trace_id"] == OPENCODE_SESSION_ID
+
+
+def test_get_chain_id_from_headers_bare_session_id_extracts_valid_value():
+    from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
+
+    assert get_chain_id_from_headers({"x-session-id": OPENCODE_SESSION_ID}) == OPENCODE_SESSION_ID
+    assert get_chain_id_from_headers({"X-Session-Id": OPENCODE_SESSION_ID}) == OPENCODE_SESSION_ID
+
+
+def test_get_chain_id_from_headers_session_affinity_extracts_when_session_id_absent():
+    from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
+
+    assert get_chain_id_from_headers({"x-session-affinity": OPENCODE_SESSION_ID}) == OPENCODE_SESSION_ID
+    assert get_chain_id_from_headers({"X-Session-Affinity": OPENCODE_SESSION_ID}) == OPENCODE_SESSION_ID
+
+
+def test_get_chain_id_from_headers_prefers_session_id_over_session_affinity():
+    from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
+
+    assert (
+        get_chain_id_from_headers(
+            {
+                "x-session-id": "ses_primary_12345678",
+                "x-session-affinity": "ses_fallback_87654321",
+            }
+        )
+        == "ses_primary_12345678"
+    )
+    assert (
+        get_chain_id_from_headers(
+            {
+                "x-session-id": "short",
+                "x-session-affinity": OPENCODE_SESSION_ID,
+            }
+        )
+        == OPENCODE_SESSION_ID
+    )
+
+
+@pytest.mark.parametrize("value", ["short", "has spaces!!", "", "invalid/slash"])
+def test_get_chain_id_from_headers_session_affinity_ignores_implausible_value(value: str):
+    from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
+
+    assert get_chain_id_from_headers({"x-session-affinity": value}) is None
+
+
+def test_get_chain_id_from_headers_missing_session_headers():
+    from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
+
+    assert get_chain_id_from_headers(None) is None
+    assert get_chain_id_from_headers({}) is None
+    assert get_chain_id_from_headers({"authorization": "Bearer token", "user-agent": "curl/7.68.0"}) is None
+
+
+@pytest.mark.asyncio
+async def test_session_affinity_header_pins_consecutive_requests_with_least_busy():
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "affinity-model",
+                "litellm_params": {"model": "openai/test-1", "api_key": "fake-1", "mock_response": "resp-1"},
+                "model_info": {"id": "deploy-1"},
+            },
+            {
+                "model_name": "affinity-model",
+                "litellm_params": {"model": "openai/test-2", "api_key": "fake-2", "mock_response": "resp-2"},
+                "model_info": {"id": "deploy-2"},
+            },
+            {
+                "model_name": "affinity-model",
+                "litellm_params": {"model": "openai/test-3", "api_key": "fake-3", "mock_response": "resp-3"},
+                "model_info": {"id": "deploy-3"},
+            },
+        ],
+        routing_strategy="least-busy",
+        optional_pre_call_checks=["session_affinity", "deployment_affinity"],
+    )
+    headers: Final = {"x-session-affinity": OPENCODE_SESSION_ID, "User-Agent": "opencode/1.18.28"}
+    try:
+        turn_data: Final = {"metadata": {}}
+        LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+            headers=headers, data=turn_data, _metadata_variable_name="metadata"
+        )
+        first_response: Final = await router.acompletion(
+            model="affinity-model",
+            messages=[{"role": "user", "content": "turn 0"}],
+            metadata=turn_data["metadata"],
+        )
+        pinned_id: Final = first_response._hidden_params["model_id"]
+
+        consecutive_responses: Final = tuple(
+            [
+                await router.acompletion(
+                    model="affinity-model",
+                    messages=[{"role": "user", "content": f"turn {i}"}],
+                    metadata=turn_data["metadata"],
+                )
+                for i in range(1, 6)
+            ]
+        )
+        assert len(consecutive_responses) == 5
+        assert all(resp._hidden_params["model_id"] == pinned_id for resp in consecutive_responses)
+    finally:
+        router.retire()
+
+
 @pytest.mark.parametrize("value", ["short", "has spaces!!", ""])
 def test_get_chain_id_from_headers_bare_session_id_ignores_implausible_value(value: str):
     from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers