fix(router): preserve native encrypted capability tasks

This commit is contained in:
Tin Chi Lo 2026-09-15 12:05:30 -07:00
parent 896f35c751
commit cadb7ee44d
2 changed files with 32 additions and 3 deletions

View file

@ -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)

View file

@ -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(