From cadb7ee44dd5b5b6902fb40c9e7048d494642811 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 12:05:30 -0700 Subject: [PATCH] fix(router): preserve native encrypted capability tasks --- .../complexity_router/complexity_router.py | 13 ++++++++--- .../router_strategy/test_complexity_router.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 68dfde394e7..c8de22e91fb 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2211,8 +2211,15 @@ class ComplexityRouter(CustomLogger): raise ValueError("capability classifier is not configured") 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 + encrypted_task: Final = _encrypted_classifier_task(request_kwargs, markers) + asks_newest_first: Final = ( + () if encrypted_task is not None else tuple(_iter_human_asks_newest_first(messages or (), markers)) + ) + opening_task: Final = ( + "The delegated task in the following agent_message." + if encrypted_task is not None + else 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 {"role": "user", "content": opening_task}, # mutable-ok: SDK messages are dict-shaped @@ -2240,7 +2247,7 @@ class ComplexityRouter(CustomLogger): messages_for_call, request_kwargs, max_output_tokens=capability.max_output_tokens, - encrypted_task=_encrypted_classifier_task(request_kwargs, markers), + encrypted_task=encrypted_task, ) 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 c8916f10965..0931b9d01a7 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2588,6 +2588,28 @@ class TestCapabilityClassifier: complexity_router_config=_capability_router_config(**overrides), ) + @pytest.mark.asyncio + async def test_encrypted_task_is_not_replaced_by_plaintext_envelope(self, mock_router_instance: MagicMock) -> None: + mock_router_instance.aresponses = AsyncMock( + return_value=_native_classifier_response(_capability_reply(p_solve=0.8)) + ) + router: Final = self._router(mock_router_instance) + task: Final = _encrypted_agent_task() + request: Final = {"input": [task]} + original: Final = deepcopy(request) + result: Final = await router.async_pre_routing_hook(model="capability-router", request_kwargs=request) + assert result is not None and result.model == "efficient-model" + assert result.routing_decision is not None + assert result.routing_decision["cause"] == "capability_classifier" + mock_router_instance.aresponses.assert_awaited_once() + call: Final = mock_router_instance.aresponses.call_args.kwargs + assert call["input"][-1] == task + plaintext: Final = json.dumps(call["input"][:-1]) + assert "The delegated task in the following agent_message." in plaintext + assert "Message Type: NEW_TASK" not in plaintext + assert "opaque-provider-task" not in plaintext + assert request == original + @pytest.mark.asyncio @pytest.mark.parametrize("custom_markers", (False, True)) async def test_task_forecast_uses_request_scoped_codex_markers(