diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..7a0393a4f7d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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() diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index d06248a555f..f85fada18b1 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/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 diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..933b2974c63 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/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"} diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index d0e2b381253..3e5e3b65f21 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7823,7 +7823,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, @@ -7834,12 +7835,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, ) @@ -7848,11 +7853,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"}]} @@ -7871,7 +7877,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, ) @@ -7945,6 +7951,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()