Merge 110b9ad9ee into 6fa0477d84

litellm/proxy/litellm_pre_call_utils.py

@@ -116,6 +116,8 @@ _BEDROCK_SESSION_ID_METADATA_KEYS: Final = ("sessionId", "threadId", "conversati
 # Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
 # (covers UUIDs and most common session-id formats).
 _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
+# Harnesses whose subagents reuse the parent's session header and add an agent id header.
+_SUBAGENT_ID_HEADERS: Final = MappingProxyType({"x-claude-code-agent-id": "x-claude-code-session-id"})
 
 _SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$")
 
@@ -215,12 +217,28 @@ def _caller_session_beside_trace_header(
     return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name)
 
 
+def _subagent_tier_session_metadata(normalized_headers: Mapping[str, str]) -> Mapping[str, str]:
+    harness_ids: Final = (
+        (normalized_headers.get(session_header, ""), normalized_headers.get(agent_header, ""))
+        for agent_header, session_header in _SUBAGENT_ID_HEADERS.items()
+    )
+    child_session: Final = next(
+        (
+            f"{session_id}:{agent_id}"
+            for session_id, agent_id in harness_ids
+            if _SESSION_ID_VALUE_RE.fullmatch(session_id) and _SESSION_ID_VALUE_RE.fullmatch(agent_id)
+        ),
+        None,
+    )
+    return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: child_session} if child_session else {}
+
+
 def _child_tier_session_metadata(
     normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str
 ) -> Mapping[str, str | None]:
     parent: Final = normalized_headers.get("x-parent-session-id")
     if not parent or not _SESSION_ID_VALUE_RE.fullmatch(parent):
-        return {}
+        return _subagent_tier_session_metadata(normalized_headers)
     own_headers: Final = {
         key: value
         for key, value in normalized_headers.items()

litellm/router_strategy/complexity_router/README.md

@@ -483,6 +483,11 @@ without replacing the parent's tier. Existing trace grouping, per-tier model aff
 affinity retain their session identities. A child without a distinct, usable session ID classifies
 without storing a tier pin. `session_affinity: true` retains precedence and the existing shared-pin behavior
 
+Claude Code subagents send the parent's `x-claude-code-session-id` plus `x-claude-code-agent-id`. When both
+values are valid session IDs, the subagent's tier pin is stored under that session and agent pair, so each
+subagent keeps its own tier and the parent's continuations keep the parent's tier. The session ID, trace grouping,
+deployment affinity and the Claude Code session router binding still use the parent's session
+
 `session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds.
 
 ### Mid-task stall escalation

tests/unit/proxy/test_litellm_pre_call_utils.py

@@ -3857,6 +3857,87 @@ async def test_non_child_requests_cannot_forge_a_child_tier_identity(surface: st
     assert updated[get_metadata_variable_name(request)]["session_id"] == "child-01"
 
 
+_CLAUDE_CODE_SESSION: Final = "11111111-2222-3333-4444-555555555555"
+_CLAUDE_CODE_AGENT: Final = "a01ef07d8c218e9cb"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "surface,headers,child_id",
+    [
+        (
+            "messages",
+            {"X-Claude-Code-Session-Id": _CLAUDE_CODE_SESSION, "X-Claude-Code-Agent-Id": _CLAUDE_CODE_AGENT},
+            f"{_CLAUDE_CODE_SESSION}:{_CLAUDE_CODE_AGENT}",
+        ),
+        (
+            "chat/completions",
+            {"x-claude-code-session-id": _CLAUDE_CODE_SESSION, "x-claude-code-agent-id": "agent-b-0001"},
+            f"{_CLAUDE_CODE_SESSION}:agent-b-0001",
+        ),
+        (
+            "responses",
+            {
+                "x-litellm-trace-id": "trace-01",
+                "x-claude-code-session-id": _CLAUDE_CODE_SESSION,
+                "x-claude-code-agent-id": _CLAUDE_CODE_AGENT,
+            },
+            f"{_CLAUDE_CODE_SESSION}:{_CLAUDE_CODE_AGENT}",
+        ),
+    ],
+)
+async def test_claude_code_subagent_gets_its_own_tier_identity_under_the_parent_session(
+    surface: str, headers: dict[str, str], child_id: str
+) -> None:
+    request: Final = _make_request_mock(f"/v1/{surface}", headers)
+    updated: Final = await add_litellm_data_to_request(
+        data={
+            "model": "test-router",
+            "messages": [],
+            "metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-1"},
+            "litellm_metadata": {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-2"},
+        },
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={},
+    )
+    metadata: Final = updated[get_metadata_variable_name(request)]
+    assert metadata[COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY] == child_id
+    expected_session: Final = headers.get("x-litellm-trace-id", _CLAUDE_CODE_SESSION)
+    assert metadata["session_id"] == updated["litellm_session_id"] == expected_session
+    assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "headers",
+    [
+        {"x-claude-code-session-id": _CLAUDE_CODE_SESSION},
+        {"x-claude-code-agent-id": _CLAUDE_CODE_AGENT},
+        {"x-claude-code-agent-id": _CLAUDE_CODE_AGENT, "x-session-id": _CLAUDE_CODE_SESSION},
+        {"x-claude-code-session-id": _CLAUDE_CODE_SESSION, "x-claude-code-agent-id": "short"},
+        {"x-claude-code-session-id": _CLAUDE_CODE_SESSION, "x-claude-code-agent-id": "agent id with spaces"},
+        {"x-claude-code-session-id": _CLAUDE_CODE_SESSION, "x-claude-code-agent-id": f"{_CLAUDE_CODE_AGENT}\n"},
+        {"x-claude-code-session-id": "bad", "x-claude-code-agent-id": _CLAUDE_CODE_AGENT},
+    ],
+)
+async def test_claude_code_requests_without_a_valid_subagent_identity_keep_the_session_tier_pin(
+    headers: dict[str, str],
+) -> None:
+    request: Final = _make_request_mock("/v1/messages", headers)
+    updated: Final = await add_litellm_data_to_request(
+        data={"model": "test-router", "messages": [], "litellm_metadata": {
+            COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: "forged-1"
+        }},
+        request=request,
+        user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+        proxy_config=MagicMock(),
+        general_settings={},
+    )
+    assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in updated[get_metadata_variable_name(request)]
+
+
 def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id():
     """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining."""
     headers = {"x-litellm-session-id": "bar"}

tests/unit/router_strategy/test_complexity_router.py

@@ -7636,7 +7636,8 @@ class TestChildSessionTierPins:
 
     @staticmethod
     def _normalized_request(
-        surface: str, session_id: str | None, caller: str = "caller-a", child: bool = True
+        surface: str, session_id: str | None, caller: str = "caller-a", child: bool = True,
+        headers: Mapping[str, str] | None = None,
     ) -> tuple[dict[str, object], str]:
         from litellm.proxy.litellm_pre_call_utils import (
             LiteLLMProxyRequestSetup,
@@ -7647,12 +7648,16 @@ class TestChildSessionTierPins:
             "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"
         }[surface]
         bucket: Final = metadata_variable_name_for_route(route)
-        data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}}
+        request_headers: Final = dict(headers) if headers is not None else {
+            **({"X-Session-Id": session_id} if session_id is not None else {}),
+            **({"x-parent-session-id": "parent-session"} if child else {}),
+        }
+        data: Final[dict[str, object]] = {
+            bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route},
+            "proxy_server_request": {"headers": request_headers},
+        }
         LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
-            headers={
-                **({"X-Session-Id": session_id} if session_id is not None else {}),
-                **({"x-parent-session-id": "parent-session"} if child else {}),
-            },
+            headers=request_headers,
             data=data,
             _metadata_variable_name=bucket,
         )
@@ -7661,11 +7666,12 @@ class TestChildSessionTierPins:
     async def _route(
         self, router: Router, surface: str, session_id: str | None, prompt: str | None,
         expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True,
-        expected_deployment: str | None = None,
+        expected_deployment: str | None = None, headers: Mapping[str, str] | None = None,
+        model: str = "child-tier-router",
     ) -> str:
         from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
 
-        data, bucket = self._normalized_request(surface, session_id, caller, child)
+        data, bucket = self._normalized_request(surface, session_id, caller, child, headers)
         turn: Final[dict[str, object]] = (
             {"role": "user", "content": prompt} if prompt is not None else
             {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]}
@@ -7684,7 +7690,7 @@ class TestChildSessionTierPins:
         if surface == "responses":
             data["input"] = turns
         selected: Final = await router.async_get_available_deployment(
-            model="child-tier-router", request_kwargs=data,
+            model=model, request_kwargs=data,
             messages=None if surface == "responses" else turns,
             input=turns if surface == "responses" else None,
         )
@@ -7758,6 +7764,58 @@ class TestChildSessionTierPins:
         finally:
             router.discard()
 
+    @staticmethod
+    def _claude_code_headers(agent_id: str | None = None) -> Mapping[str, str]:
+        return {
+            "x-claude-code-session-id": "parent-session",
+            "x-app": "cli",
+            **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}),
+        }
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("surface", ["chat", "responses", "messages"])
+    async def test_claude_code_subagents_keep_their_tier_without_replacing_the_parent(self, surface: str) -> None:
+        router: Final = self._router()
+        parent_headers: Final = self._claude_code_headers()
+        try:
+            parent: Final = await self._route(
+                router, surface, None, "reason carefully", "REASONING", "literal_keyword_match",
+                headers=parent_headers,
+            )
+            for agent, prompt, tier, cause, model in (
+                ("agent-a-0001", "Hello!", "SIMPLE", "heuristic_scorer", "simple"),
+                ("agent-a-0001", None, "SIMPLE", "user_turn_continuation", "child-tier-router"),
+                (None, None, "REASONING", "user_turn_continuation", "child-tier-router"),
+                ("agent-b-0001", None, "SIMPLE", "heuristic_scorer", "child-tier-router"),
+                ("agent-b-0001", "reason carefully", "REASONING", "literal_keyword_match", "simple"),
+                ("agent-a-0001", None, "SIMPLE", "user_turn_continuation", "child-tier-router"),
+            ):
+                await self._route(
+                    router, surface, None, prompt, tier, cause, headers=self._claude_code_headers(agent),
+                    model=model, expected_deployment=parent if tier == "REASONING" else None,
+                )
+            assert await self._route(
+                router, surface, None, None, "REASONING", "user_turn_continuation",
+                headers=parent_headers, expected_deployment=parent,
+            ) == parent
+        finally:
+            router.discard()
+
+    @pytest.mark.asyncio
+    async def test_session_affinity_still_shares_the_parent_tier_with_claude_code_subagents(self) -> None:
+        router: Final = self._router(session_affinity=True)
+        try:
+            parent: Final = await self._route(
+                router, "messages", None, "reason carefully", "REASONING", "literal_keyword_match",
+                headers=self._claude_code_headers(),
+            )
+            assert await self._route(
+                router, "messages", None, "Hello!", "REASONING", "session_affinity_pin",
+                headers=self._claude_code_headers("agent-a-0001"), expected_deployment=parent,
+            ) == parent
+        finally:
+            router.discard()
+
     @pytest.mark.asyncio
     async def test_child_tier_model_affinity_keeps_using_the_parent_identity(self) -> None:
         router: Final = self._router()