From 896f35c7513c689469942326c3f25c75d0c1fca7 Mon Sep 17 00:00:00 2001
From: Tin Chi Lo <tin@berri.ai>
Date: Tue, 15 Sep 2026 11:52:20 -0700
Subject: [PATCH] fix(router): extract capability tasks with request scoped
 markers

---
 .../complexity_router/complexity_router.py    |  7 ++--
 .../router_strategy/test_complexity_router.py | 34 +++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py
index 72f22e0bab0..68dfde394e7 100644
--- a/litellm/router_strategy/complexity_router/complexity_router.py
+++ b/litellm/router_strategy/complexity_router/complexity_router.py
@@ -2210,7 +2210,8 @@ class ComplexityRouter(CustomLogger):
         if capability is None or classifier_system_prompt is None:
             raise ValueError("capability classifier is not configured")
 
-        asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), self._reminder_markers))
+        markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
+        asks_newest_first: Final = tuple(_iter_human_asks_newest_first(messages or (), markers))
         opening_task: Final = asks_newest_first[-1] if asks_newest_first else prompt
         latest_follow_up: Final = asks_newest_first[0] if len(asks_newest_first) > 1 else None
         task_messages: list[AllMessageValues] = [  # mutable-ok: the latest message gains optional image parts below
@@ -2239,9 +2240,7 @@ class ComplexityRouter(CustomLogger):
             messages_for_call,
             request_kwargs,
             max_output_tokens=capability.max_output_tokens,
-            encrypted_task=_encrypted_classifier_task(
-                request_kwargs, self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
-            ),
+            encrypted_task=_encrypted_classifier_task(request_kwargs, markers),
         )
         verdict: Final = parse_capability_classifier_verdict(content)
         threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step)
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py
index 38411ef52ea..c8916f10965 100644
--- a/tests/test_litellm/router_strategy/test_complexity_router.py
+++ b/tests/test_litellm/router_strategy/test_complexity_router.py
@@ -2588,6 +2588,40 @@ class TestCapabilityClassifier:
             complexity_router_config=_capability_router_config(**overrides),
         )
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("custom_markers", (False, True))
+    async def test_task_forecast_uses_request_scoped_codex_markers(
+        self, mock_router_instance: MagicMock, custom_markers: bool
+    ) -> None:
+        completion: Final = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
+        mock_router_instance.acompletion = completion
+        router: Final = self._router(
+            mock_router_instance,
+            escalation_keywords=[],
+            **({"reminder_markers": [{"open": "<custom>", "close": "</custom>"}]} if custom_markers else {}),
+        )
+        envelope: Final = "\n".join(_CODEX_ENVELOPES)
+        opening: Final = f"{envelope}\nFix nested behavior"
+        messages: Final = [
+            {"role": "user", "content": opening},
+            {"role": "user", "content": "Preserve empty inputs"},
+            {"role": "user", "content": envelope},
+        ]
+        original: Final = deepcopy(messages)
+        for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"):
+            result: Final = await router.async_pre_routing_hook(
+                model="capability-router", messages=messages, request_kwargs={"metadata": {"user_agent": user_agent}}
+            )
+            assert result is not None and result.model == "efficient-model"
+            sent: Final = completion.call_args.kwargs["messages"]
+            if user_agent.startswith("codex") and not custom_markers:
+                assert [message["content"] for message in sent[1:]] == ["Fix nested behavior", "Preserve empty inputs"]
+            else:
+                assert [message["content"] for message in sent[1:]] == [opening, envelope]
+            assert result.messages == original
+        assert completion.await_count == 3
+        assert messages == original
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("p_solve,expected_model", ((0.95, "capable-model"), (0.98, "efficient-model")))
     async def test_fitted_probability_controls_routing_and_preserves_raw_score(