mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(complexity_router): keep Claude Code subagent and side-request traffic off other agents' tier pins (#50741)
* fix(complexity_router): keep Claude Code subagent and side-request traffic off other agents' tier pins Claude Code sends one x-claude-code-session-id for the main agent, every subagent and their side requests. In user_turn mode a subagent's new ask rewrote the parent's pin (#50099), and the agent progress summaries Claude Code fires during a subagent run rewrote the subagent's pin to the cheapest tier (#50444). Subagent requests now get their own tier pin keyed on session id plus x-claude-code-agent-id. Requests marked x-claude-code-request-class: auxiliary route as housekeeping, and only main and subagent turns may write a pin. Without the hint header, the summary and prompt-suggestion prompts are matched as housekeeping sentinels. * chore(ui): regenerate schema.d.ts for the housekeeping_patterns description * fix(complexity_router): let the request-class hint only mark housekeeping, not suppress pin writes A client-sent request class must not decide whether a decision is saved. Tagging a real ask as workflow or auxiliary dropped its routing decision, so later continuations replayed an older pin with the new history. The auxiliary hint now only routes the request through the housekeeping path, which is already never pinned and defers to custom classifiers and keyword rules. * test(complexity_router): cover the auxiliary hint on a prompt no sentinel matches --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
parent
92c2e69ac5
commit
35386e3f4a
6 changed files with 240 additions and 7 deletions
|
|
@ -223,12 +223,21 @@ def _caller_session_beside_trace_header(
|
|||
return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name)
|
||||
|
||||
|
||||
def _claude_code_subagent_tier_session(normalized_headers: Mapping[str, str]) -> str | None:
|
||||
session: Final = normalized_headers.get("x-claude-code-session-id", "")
|
||||
agent: Final = normalized_headers.get("x-claude-code-agent-id", "")
|
||||
if not _SESSION_ID_VALUE_RE.fullmatch(session) or not _SESSION_ID_VALUE_RE.fullmatch(agent):
|
||||
return None
|
||||
return f"{session}:{agent}"
|
||||
|
||||
|
||||
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 {}
|
||||
claude_code_subagent: Final = _claude_code_subagent_tier_session(normalized_headers)
|
||||
return {COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY: claude_code_subagent} if claude_code_subagent else {}
|
||||
own_headers: Final = {
|
||||
key: value
|
||||
for key, value in normalized_headers.items()
|
||||
|
|
|
|||
|
|
@ -801,6 +801,27 @@ def _newest_turn_is_human_ask(
|
|||
return _last_human_ask_index(messages, marker_pairs) == newest_activity
|
||||
|
||||
|
||||
_CLAUDE_CODE_REQUEST_CLASS_HEADER: Final = "x-claude-code-request-class"
|
||||
_CLAUDE_CODE_AUXILIARY_SENTINEL: Final = f"{_CLAUDE_CODE_REQUEST_CLASS_HEADER}: auxiliary"
|
||||
|
||||
|
||||
def _claude_code_request_class(request_kwargs: Mapping[str, object]) -> str | None:
|
||||
proxy_request: Final = request_kwargs.get("proxy_server_request")
|
||||
if not isinstance(proxy_request, Mapping):
|
||||
return None
|
||||
headers: Final = cast(Mapping[str, object], proxy_request).get("headers") # cast-ok: proxy request is JSON-shaped
|
||||
if not isinstance(headers, Mapping):
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
value.strip().lower()
|
||||
for key, value in cast(Mapping[object, object], headers).items() # cast-ok: header map from the proxy
|
||||
if isinstance(key, str) and key.lower() == _CLAUDE_CODE_REQUEST_CLASS_HEADER and isinstance(value, str)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
_CODEX_TURN_METADATA_KEY: Final = "x-codex-turn-metadata"
|
||||
_CODEX_COMPACTION_SUMMARY_PREFIX: Final = (
|
||||
"Another language model started to solve this problem and produced a summary of its thinking process."
|
||||
|
|
@ -3475,8 +3496,11 @@ class ComplexityRouter(CustomLogger):
|
|||
self._reminder_markers_for_request(request_kwargs),
|
||||
)
|
||||
|
||||
def _matched_housekeeping_sentinel(self, newest_ask: str | None) -> str | None:
|
||||
"""The client housekeeping sentinel on this request's newest ask, or None.
|
||||
def _matched_housekeeping_sentinel(
|
||||
self, newest_ask: str | None, request_kwargs: Mapping[str, object]
|
||||
) -> str | None:
|
||||
"""The client housekeeping sentinel on this request's newest ask, or None. Claude Code's
|
||||
`auxiliary` request-class hint counts as one, so it keeps working when the prompts change.
|
||||
|
||||
Read from the newest ask alone, never the whole history, for the reason `_newest_turn_ask`
|
||||
exists: a title request quoted into a later turn's context would otherwise keep matching and
|
||||
|
|
@ -3490,6 +3514,8 @@ class ComplexityRouter(CustomLogger):
|
|||
"""
|
||||
if self.config.classifier_type == "custom" or not self.config.route_housekeeping_to_cheapest_tier:
|
||||
return None
|
||||
if _claude_code_request_class(request_kwargs) == "auxiliary":
|
||||
return _CLAUDE_CODE_AUXILIARY_SENTINEL
|
||||
if not newest_ask:
|
||||
return None
|
||||
return next(
|
||||
|
|
@ -4970,7 +4996,7 @@ class ComplexityRouter(CustomLogger):
|
|||
),
|
||||
)
|
||||
|
||||
housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask)
|
||||
housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask, request_kwargs)
|
||||
housekeeping_tier: Final = self._cheapest_configured_tier() if housekeeping_sentinel is not None else None
|
||||
outcome: Final = (
|
||||
ClassificationOutcome(tier=housekeeping_tier, score=None, signals=("housekeeping",), cause="housekeeping")
|
||||
|
|
|
|||
|
|
@ -414,9 +414,14 @@ PLAN_MODE_SYSTEM_SENTINELS: Final[tuple[str, ...]] = ('You are currently running
|
|||
# never a paraphrase: a sentinel that matches nothing costs a substring scan per request and reads
|
||||
# as coverage the router does not have. These are client-owned strings that drift with client
|
||||
# releases, so operators extend coverage via housekeeping_patterns rather than editing these.
|
||||
# The agent progress summary and the next-prompt suggestion were captured from Claude Code 2.1.316
|
||||
# on 2026-11-09 (#50444). Both ride the conversation they describe, so only the sentinel tells them
|
||||
# apart when the client does not send the x-claude-code-request-class hint.
|
||||
HOUSEKEEPING_ASK_SENTINELS: Final[tuple[str, ...]] = (
|
||||
"Write the title in the predominant language of the session",
|
||||
"You are coming up with a succinct title for a coding session",
|
||||
"Describe your most recent action in 3-5 words using present tense",
|
||||
"[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]",
|
||||
)
|
||||
PLAN_MODE_TOOL_NAME: Final[str] = "exit_plan_mode"
|
||||
|
||||
|
|
@ -1496,8 +1501,9 @@ class ComplexityRouterConfig(LiteLLMBaseModel):
|
|||
default=None,
|
||||
description=(
|
||||
"Additional case-sensitive literal sentinels that mark a request as client "
|
||||
"housekeeping, on top of the built-in conversation-title ones. For clients whose "
|
||||
"wording the built-ins don't cover, or after a client release changes its strings."
|
||||
"housekeeping, on top of the built-in conversation-title, agent-summary and "
|
||||
"prompt-suggestion ones. For clients whose wording the built-ins don't cover, or "
|
||||
"after a client release changes its strings."
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -3895,6 +3895,41 @@ 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"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"agent_id,session_id,child_id",
|
||||
[
|
||||
(
|
||||
"a7256c04dfd99c9b9",
|
||||
"11111111-2222-3333-4444-555555555555",
|
||||
"11111111-2222-3333-4444-555555555555:a7256c04dfd99c9b9",
|
||||
),
|
||||
(None, "11111111-2222-3333-4444-555555555555", None),
|
||||
("short", "11111111-2222-3333-4444-555555555555", None),
|
||||
("agent id with spaces", "11111111-2222-3333-4444-555555555555", None),
|
||||
("a7256c04dfd99c9b9", "bad", None),
|
||||
],
|
||||
)
|
||||
async def test_claude_code_subagents_get_their_own_tier_identity_under_the_parent_session(
|
||||
agent_id: str | None, session_id: str, child_id: str | None
|
||||
) -> None:
|
||||
headers: Final = {
|
||||
"x-claude-code-session-id": session_id,
|
||||
**({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}),
|
||||
}
|
||||
request: Final = _make_request_mock("/v1/messages", headers)
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data={"model": "test-router", "messages": []},
|
||||
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.get(COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY) == child_id
|
||||
assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {})
|
||||
|
||||
|
||||
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"}
|
||||
|
|
|
|||
|
|
@ -7772,6 +7772,163 @@ class TestChildSessionTierPins:
|
|||
router.discard()
|
||||
|
||||
|
||||
class TestClaudeCodeTierPins:
|
||||
SESSION: Final = "11111111-2222-3333-4444-555555555555"
|
||||
AGENT_A: Final = "a7256c04dfd99c9b9"
|
||||
AGENT_B: Final = "a71aeeb9c67f32cee"
|
||||
TOOL_USE: Final[dict[str, object]] = {
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}],
|
||||
}
|
||||
TOOL_RESULT: Final[dict[str, object]] = {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "t1",
|
||||
"content": "def add(a, b): return a + b",
|
||||
}
|
||||
# Captured from Claude Code 2.1.316 on 2026-11-09, trimmed to the sentinel and a little context.
|
||||
SUMMARY_ASK: Final = (
|
||||
"Describe your most recent action in 3-5 words using present tense (-ing). "
|
||||
"Name the file or function, not the branch. Do not use tools."
|
||||
)
|
||||
SUGGESTION_ASK: Final = (
|
||||
"[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]\n\n"
|
||||
"Reply with ONLY the suggestion, no quotes or explanation."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _router(route_housekeeping_to_cheapest_tier: bool = True) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "claude-code-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "simple", "REASONING": "reasoning"},
|
||||
"classification_mode": "user_turn",
|
||||
"keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}],
|
||||
"route_housekeeping_to_cheapest_tier": route_housekeeping_to_cheapest_tier,
|
||||
},
|
||||
},
|
||||
},
|
||||
*[
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {"model": "openai/test-model", "api_key": "test-only"},
|
||||
"model_info": {"id": group},
|
||||
}
|
||||
for group in ("simple", "reasoning")
|
||||
],
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _ask(prompt: str) -> list[dict[str, object]]:
|
||||
return [{"role": "user", "content": prompt}]
|
||||
|
||||
@classmethod
|
||||
def _continuation(cls) -> list[dict[str, object]]:
|
||||
return [*cls._ask("Hello!"), cls.TOOL_USE, {"role": "user", "content": [cls.TOOL_RESULT]}]
|
||||
|
||||
@classmethod
|
||||
def _side_request(cls, side_ask: str) -> list[dict[str, object]]:
|
||||
return [
|
||||
*cls._ask("Hello!"),
|
||||
cls.TOOL_USE,
|
||||
{"role": "user", "content": [cls.TOOL_RESULT, {"type": "text", "text": side_ask}]},
|
||||
]
|
||||
|
||||
async def _route(
|
||||
self,
|
||||
router: Router,
|
||||
messages: list[dict[str, object]],
|
||||
agent_id: str | None = None,
|
||||
request_class: str | None = None,
|
||||
) -> tuple[str, str]:
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
metadata_variable_name_for_route,
|
||||
)
|
||||
|
||||
headers: Final = {
|
||||
"x-claude-code-session-id": self.SESSION,
|
||||
**({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}),
|
||||
**({"x-claude-code-request-class": request_class} if request_class is not None else {}),
|
||||
}
|
||||
bucket: Final = metadata_variable_name_for_route("/v1/messages")
|
||||
data: Final[dict[str, object]] = {
|
||||
bucket: {"user_api_key_hash": "caller-a", "user_api_key_request_route": "/v1/messages"},
|
||||
"proxy_server_request": {"headers": headers},
|
||||
}
|
||||
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
|
||||
headers=headers, data=data, _metadata_variable_name=bucket
|
||||
)
|
||||
await router.async_get_available_deployment(model="claude-code-router", request_kwargs=data, messages=messages)
|
||||
metadata: Final = data[bucket]
|
||||
assert isinstance(metadata, dict)
|
||||
assert metadata["session_id"] == self.SESSION
|
||||
decision: Final = TypeAdapter(StandardLoggingRoutingDecision).validate_python(metadata["routing_decision"])
|
||||
return str(decision.get("tier")), decision["cause"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_and_each_subagent_keep_their_own_pin(self) -> None:
|
||||
router: Final = self._router()
|
||||
try:
|
||||
for messages, agent_id, expected in (
|
||||
(self._ask("reason carefully"), None, ("REASONING", "literal_keyword_match")),
|
||||
(self._ask("Hello!"), self.AGENT_A, ("SIMPLE", "heuristic_scorer")),
|
||||
(self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")),
|
||||
(self._continuation(), None, ("REASONING", "user_turn_continuation")),
|
||||
(self._ask("reason carefully"), self.AGENT_B, ("REASONING", "literal_keyword_match")),
|
||||
(self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")),
|
||||
(self._continuation(), self.AGENT_B, ("REASONING", "user_turn_continuation")),
|
||||
):
|
||||
assert await self._route(router, messages, agent_id) == expected
|
||||
finally:
|
||||
router.discard()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hint_header", [True, False], ids=["request-class-hint", "sentinel-fallback"])
|
||||
@pytest.mark.parametrize(
|
||||
"agent_id,side_ask",
|
||||
[(AGENT_A, SUMMARY_ASK), (None, SUGGESTION_ASK)],
|
||||
ids=["subagent-summary", "main-prompt-suggestion"],
|
||||
)
|
||||
async def test_side_requests_route_cheapest_and_never_move_the_pin(
|
||||
self, hint_header: bool, agent_id: str | None, side_ask: str
|
||||
) -> None:
|
||||
own_class: Final = ("subagent" if agent_id is not None else "main") if hint_header else None
|
||||
side_class: Final = "auxiliary" if hint_header else None
|
||||
router: Final = self._router()
|
||||
try:
|
||||
for messages, request_class, expected in (
|
||||
(self._ask("reason carefully"), own_class, ("REASONING", "literal_keyword_match")),
|
||||
(self._side_request(side_ask), side_class, ("SIMPLE", "housekeeping")),
|
||||
(self._continuation(), own_class, ("REASONING", "user_turn_continuation")),
|
||||
):
|
||||
assert await self._route(router, messages, agent_id, request_class) == expected
|
||||
finally:
|
||||
router.discard()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_auxiliary_hint_routes_housekeeping_when_no_sentinel_matches(self) -> None:
|
||||
router: Final = self._router()
|
||||
try:
|
||||
for messages, request_class, expected in (
|
||||
(self._ask("reason carefully"), "subagent", ("REASONING", "literal_keyword_match")),
|
||||
(
|
||||
self._side_request("Summarize what you just did in a few words."),
|
||||
"auxiliary",
|
||||
("SIMPLE", "housekeeping"),
|
||||
),
|
||||
(self._continuation(), "subagent", ("REASONING", "user_turn_continuation")),
|
||||
):
|
||||
assert await self._route(router, messages, self.AGENT_A, request_class) == expected
|
||||
finally:
|
||||
router.discard()
|
||||
|
||||
|
||||
class TestClassificationMode:
|
||||
"""Test classification_mode='user_turn': classify only requests whose newest turn is a new
|
||||
human ask; tool-loop continuation turns replay the session's held routing decision."""
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -44032,7 +44032,7 @@ export interface components {
|
|||
heuristic_v2_success_threshold?: number | null;
|
||||
/**
|
||||
* Housekeeping Patterns
|
||||
* @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings.
|
||||
* @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title, agent-summary and prompt-suggestion ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings.
|
||||
*/
|
||||
housekeeping_patterns?: string[] | null;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue