From 5d6fec94b74371a7314a4900f3cf025a54f8bcff Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 10 Sep 2026 11:33:21 -0700 Subject: [PATCH] fix(router): strip Codex harness envelopes before classification --- .../complexity_router/README.md | 8 +- .../complexity_router/complexity_router.py | 9 +- .../complexity_router/config.py | 4 +- .../router_strategy/test_complexity_router.py | 103 +++++++++++++++++- 4 files changed, 118 insertions(+), 6 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index d605b43e42a..3d2b33a2ecd 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -446,7 +446,11 @@ Reasoning markers in the system prompt do **not** trigger the reasoning override Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead -By default a block is anything between `` and ``. `reminder_markers` replaces that with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit: +By default the router strips complete ``, ``, ``, ``, and `` blocks. It also strips Codex repository instructions from the fixed heading prefix `# AGENTS.md instructions for ` through ``, regardless of the repository path + +The Codex `Message Type: NEW_TASK` wrapper and its delegated-task payload remain available for classification. Cleanup applies to the current ask and quoted prior turns; the routed request retains its original content + +`reminder_markers` replaces these defaults with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit: ```yaml model_list: @@ -461,7 +465,7 @@ model_list: close: "[[SUBAGENT_CONTEXT_END]]" ``` -Setting `reminder_markers` replaces the built-in `` pair rather than adding to it, so list that pair too if your harness also emits it. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten +Setting `reminder_markers` replaces all built-in pairs, including the Codex heading pair, so include every default your harness still needs. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten ### Code Detection diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index faafcea404a..923aac85ee2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -386,7 +386,14 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | N _REMINDER_OPEN: Final = "" _REMINDER_CLOSE: Final = "" -_DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),) +_DEFAULT_REMINDER_MARKERS: Final = ( + (_REMINDER_OPEN, _REMINDER_CLOSE), + ("", ""), + ("", ""), + ("", ""), + ("", ""), + ("# agents.md instructions for ", ""), +) _TRUNCATION_MARKER: Final = "..." _TRUNCATION_HEAD_FRACTION: Final = 0.3 diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d28924c69b2..6ab4aa85148 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1292,8 +1292,8 @@ class ComplexityRouterConfig(BaseModel): "Override the delimiter pairs used to recognize and strip harness-injected reminder " "blocks before classification. A harness that wraps injected context differently per " "agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than " - "adds to, the built-in default of ('', ''), so a " - "harness that also emits that pair lists it too. Matching is case-insensitive." + "adds to, the built-in system-reminder and Codex envelope pairs, so list every " + "built-in pair your harness also emits. Matching is case-insensitive." ), ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 51103297c58..792963fde23 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5,6 +5,7 @@ Tests the rule-based complexity scoring and tier assignment logic. """ import asyncio +from copy import deepcopy import logging import sys import time @@ -194,7 +195,14 @@ class TestComplexityRouterInit: complexity_router_config=basic_config, ) - assert router._reminder_markers == (("", ""),) + from litellm.router_strategy.complexity_router.complexity_router import _extract_current_ask_and_system_prompt + + assert ( + _extract_current_ask_and_system_prompt( + [{"role": "user", "content": "noisehello"}], router._reminder_markers + )[0] + == "hello" + ) def test_init_without_config(self, mock_router_instance): """Test initialization without configuration uses defaults.""" @@ -7601,11 +7609,103 @@ _ASKED = {"role": "user", "content": _ASK} _ANSWERED = {"role": "assistant", "content": "Working on it."} _TOOL_RESULT = {"type": "tool_result", "tool_use_id": "x", "content": "out"} _REMINDER = "Budget: 42 tokens remaining. Do not mention this." +_CODEX_NEW_TASK: Final = ( + "Message Type: NEW_TASK\nTask name: /root/cache_worker\nSender: /root\nPayload:\n" + "Implement and test a thread-safe bounded LRU cache." +) +_CODEX_ENVELOPES: Final = ( + "LITELLM ESCALATE cwd=/repo", + "LITELLM ESCALATE plugin list", + "LITELLM ESCALATE preferences", + "LITELLM ESCALATE environment", + "# AGENTS.md instructions for /repo with spaces/中文\nLITELLM ESCALATE instructions", +) class TestContextAwareClassifier: """Test the new classifier context window and trajectory signals.""" + @pytest.mark.parametrize("envelope", _CODEX_ENVELOPES) + def test_codex_envelopes_preserve_delegated_task_and_prior_context(self, envelope: str) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ( + _extract_current_ask_and_system_prompt, + _extract_prior_turns, + _newest_turn_ask, + _newest_turn_is_human_ask, + ) + + messages: Final = [ + {"role": "user", "content": f"{envelope}\nDesign cache invalidation"}, + { + "role": "user", + "content": [{"type": "text", "text": envelope}, {"type": "text", "text": _CODEX_NEW_TASK}], + }, + {"role": "developer", "content": "developer scope"}, + {"role": "user", "content": envelope}, + ] + + assert _extract_current_ask_and_system_prompt(messages)[0] == _CODEX_NEW_TASK + assert _extract_prior_turns(messages, _CODEX_NEW_TASK, 1, 100, None, False) == ( + ("user", "Design cache invalidation"), + ) + assert _newest_turn_ask(messages) is None + assert _newest_turn_is_human_ask(messages) is False + assert _extract_current_ask_and_system_prompt([messages[-1]])[0] is None + + @pytest.mark.parametrize("envelope", _CODEX_ENVELOPES) + def test_codex_marker_override_and_incomplete_blocks_preserve_text(self, envelope: str) -> None: + from litellm.router_strategy.complexity_router.complexity_router import _strip_reminder_blocks + + incomplete: Final = envelope.rsplit("noise{envelope}", (("", ""),)) == envelope + + @pytest.mark.asyncio + @pytest.mark.parametrize("responses_api", (False, True)) + async def test_codex_routing_preserves_original_request(self, responses_api: bool) -> None: + completion: Final = AsyncMock(return_value=_llm_response('{"tier":"COMPLEX"}')) + router: Final = ComplexityRouter( + model_name="codex-router", + litellm_router_instance=MagicMock(acompletion=completion), + complexity_router_config={ + "tiers": {"COMPLEX": "task-model", "REASONING": "escalated-model"}, + "classifier_type": "llm", + "classifier_llm_config": {"model": "classifier-model"}, + "keyword_tier_rules": [{"keywords": ["LITELLM ESCALATE"], "tier": "REASONING"}], + }, + ) + messages: Final = [ + {"role": "user", "content": _CODEX_NEW_TASK}, + {"role": "user", "content": "\n".join(_CODEX_ENVELOPES)}, + ] + original: Final = deepcopy(messages) + request_kwargs: Final = ( + {"input": messages, "litellm_metadata": {"user_api_key_request_route": "/v1/responses"}} + if responses_api + else {} + ) + + response: Final = await router.async_pre_routing_hook( + model="codex-router", + request_kwargs=request_kwargs, + messages=None if responses_api else messages, + input=messages if responses_api else None, + ) + + assert response is not None + assert response.model == "task-model" + completion.assert_awaited_once() + assert completion.call_args.kwargs["messages"][1]["content"].strip() == ( + f"Classify this message:\n{_CODEX_NEW_TASK}" + ) + assert messages == original + if responses_api: + assert response.messages is None + assert request_kwargs["input"] == original + else: + assert response.messages == original + @pytest.mark.parametrize( "messages,expected_ask", [ @@ -13662,6 +13762,7 @@ class _OutputCeilingRecorder(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): self.seen.append((model, kwargs.get("optional_params", {}).get("max_tokens"))) + NON_REASONING_TIERS: Final = { "NON_REASONING": "gpt-4o-mini", "SIMPLE": "gpt-4o-mini",