mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(router): isolate Claude Code subagent tier pins in user turn mode
Claude Code subagents send the parent's x-claude-code-session-id plus x-claude-code-agent-id, so in classification_mode user_turn a subagent's new ask overwrote the parent's tier pin and the parent's next tool-result continuation replayed the subagent's tier. Derive the child tier-pin identity from the session and agent id pair, reusing the child tier session metadata key from #48701. Session id, tracing, deployment affinity and the Claude Code session router binding keep the parent session. session_affinity keeps its shared pin. Fixes #50099
This commit is contained in:
parent
784754ca48
commit
110b9ad9ee
4 changed files with 172 additions and 10 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue