From 455cd8bfcd01019df373c0a80bffa64efaab2f8d Mon Sep 17 00:00:00 2001 From: Falicitas Date: Sat, 11 Apr 2026 23:32:43 +0800 Subject: [PATCH] fix(agent): keep smart routing from compressing session history MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Smart model routing rebuilds AIAgent for temporary per-turn model switches. The rebuild creates a new ContextCompressor bound to the cheap model's context_length, so preflight compression compares against the cheap threshold and compresses session history that was sized for the primary model — even for trivial messages like "hi". The fix establishes this contract: smart routing either uses the cheap model when it can safely handle the current request, or falls back to the primary model. It never triggers destructive session compression. Two coordinated changes in run_conversation / smart_model_routing: 1. skip_preflight_compression param on run_conversation() Set True from the CLI chat, single-query, and /btw paths when the turn carries a smart-routing label. Prevents preflight from firing at 50% of the cheap model's (possibly smaller) context. 2. Refuse-to-route on history overflow (smart_model_routing) resolve_turn_route() accepts current_request_tokens and max_history_ratio kwargs. When the estimated request exceeds cheap_ctx * max_history_ratio, smart routing returns the primary route instead of the cheap route. CLI computes the estimate via estimate_request_tokens_rough and passes the active compressor's threshold_percent as the ratio — tuning compression.threshold in config.yaml tunes smart routing's refuse threshold at the same time. The default max_history_ratio of 0.50 mirrors ContextCompressor.threshold_percent. It leaves the upper half of the cheap context free for tool outputs and responses within the turn, so typical read_file / web_extract / search_files results can land on the cheap model without pushing the next API call past its real context limit. Behavior summary (default primary=Sonnet 1M, ratio=0.50): - history < 50% of cheap context: route to cheap, preflight skipped - history 50-100% of cheap context: refuse, stay on primary - history > cheap context: refuse, stay on primary - short message fails the simple-turn filter: stay on primary (existing) The default cheap model (gemini-2.5-flash, 1M) matches Sonnet 1M exactly, so the refusal path rarely triggers for default configs. Users who set cheap to a smaller model (gpt-4o-mini, haiku, etc.) are now protected from silent session-history loss. Also extracts an internal _primary_route() helper in smart_model_routing.py to DRY up three identical primary-route dict constructions that previously lived inline in resolve_turn_route(). Known limitations (out of scope, left for follow-up): - Mid-turn tool output explosions >50% of cheap context still hit in-loop compression on API error. Would require per-call override architecture (reusing _try_activate_fallback pattern). - /model command (switch_model) has the same compressor-rebind pattern but its semantics are "permanent switch" so compression on downsize may be intentional. - Session resume with a different current model is a separate initialization-time variant. Gateway and batch_runner use the same smart-routing code path but are not modified here. CLI-scoped for this PR. Tests: - test_413_compression.py::TestPreflightCompression adds test_skip_preflight_compression_flag and test_skip_preflight_false_still_compresses - test_smart_model_routing.py adds 5 unit tests covering refuse / allow / backward-compat / custom ratio / unknown cheap context Closes #7798 --- agent/smart_model_routing.py | 128 +++++++++++------ cli.py | 41 ++++++ run_agent.py | 10 +- tests/agent/test_credential_pool_routing.py | 8 +- tests/agent/test_smart_model_routing.py | 149 ++++++++++++++++++++ tests/run_agent/test_413_compression.py | 73 ++++++++++ 6 files changed, 361 insertions(+), 48 deletions(-) diff --git a/agent/smart_model_routing.py b/agent/smart_model_routing.py index 6d482be27051..53925dfc6b21 100644 --- a/agent/smart_model_routing.py +++ b/agent/smart_model_routing.py @@ -2,12 +2,15 @@ from __future__ import annotations +import logging import os import re from typing import Any, Dict, Optional from utils import is_truthy_value +logger = logging.getLogger(__name__) + _COMPLEX_KEYWORDS = { "debug", "debugging", @@ -59,6 +62,37 @@ def _coerce_int(value: Any, default: int) -> int: return default +def _primary_route(primary: Dict[str, Any]) -> Dict[str, Any]: + """Build the canonical primary-model route dict. + + Used whenever smart routing declines to pick a cheap model — because the + message doesn't look simple, the cheap runtime can't be resolved, or the + estimated request won't fit the cheap model's context. ``label`` is None + so callers can distinguish primary from smart-routed turns. + """ + return { + "model": primary.get("model"), + "runtime": { + "api_key": primary.get("api_key"), + "base_url": primary.get("base_url"), + "provider": primary.get("provider"), + "api_mode": primary.get("api_mode"), + "command": primary.get("command"), + "args": list(primary.get("args") or []), + "credential_pool": primary.get("credential_pool"), + }, + "label": None, + "signature": ( + primary.get("model"), + primary.get("provider"), + primary.get("base_url"), + primary.get("api_mode"), + primary.get("command"), + tuple(primary.get("args") or ()), + ), + } + + def choose_cheap_model_route(user_message: str, routing_config: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Return the configured cheap-model route when a message looks simple. @@ -107,34 +141,36 @@ def choose_cheap_model_route(user_message: str, routing_config: Optional[Dict[st return route -def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any]], primary: Dict[str, Any]) -> Dict[str, Any]: +def resolve_turn_route( + user_message: str, + routing_config: Optional[Dict[str, Any]], + primary: Dict[str, Any], + *, + current_request_tokens: int = 0, + max_history_ratio: float = 0.50, +) -> Dict[str, Any]: """Resolve the effective model/runtime for one turn. Returns a dict with model/runtime/signature/label fields. + + Args: + current_request_tokens: Best-effort token estimate of the current + request (messages + system prompt + tools). When > 0, enables the + refuse-to-route check: if the estimate exceeds the cheap model's + context window times ``max_history_ratio``, smart routing falls + back to the primary model instead. When 0 (default), the check is + skipped — preserves backward compat for callers without an estimate. + max_history_ratio: Fraction of the cheap model's context length that + the current request may occupy before smart routing refuses. The + default of 0.50 mirrors ``ContextCompressor.threshold_percent`` so + smart routing only uses the cheap model in the zone where the cheap + model wouldn't even want to compress. Leaving the upper half of the + cheap context free gives room for tool outputs and responses within + the turn without triggering destructive in-loop compression. """ route = choose_cheap_model_route(user_message, routing_config) if not route: - return { - "model": primary.get("model"), - "runtime": { - "api_key": primary.get("api_key"), - "base_url": primary.get("base_url"), - "provider": primary.get("provider"), - "api_mode": primary.get("api_mode"), - "command": primary.get("command"), - "args": list(primary.get("args") or []), - "credential_pool": primary.get("credential_pool"), - }, - "label": None, - "signature": ( - primary.get("model"), - primary.get("provider"), - primary.get("base_url"), - primary.get("api_mode"), - primary.get("command"), - tuple(primary.get("args") or ()), - ), - } + return _primary_route(primary) from hermes_cli.runtime_provider import resolve_runtime_provider @@ -150,27 +186,33 @@ def resolve_turn_route(user_message: str, routing_config: Optional[Dict[str, Any explicit_base_url=route.get("base_url"), ) except Exception: - return { - "model": primary.get("model"), - "runtime": { - "api_key": primary.get("api_key"), - "base_url": primary.get("base_url"), - "provider": primary.get("provider"), - "api_mode": primary.get("api_mode"), - "command": primary.get("command"), - "args": list(primary.get("args") or []), - "credential_pool": primary.get("credential_pool"), - }, - "label": None, - "signature": ( - primary.get("model"), - primary.get("provider"), - primary.get("base_url"), - primary.get("api_mode"), - primary.get("command"), - tuple(primary.get("args") or ()), - ), - } + return _primary_route(primary) + + # Refuse-to-route: if the request won't fit comfortably inside the cheap + # model's context, fall back to primary rather than either (a) triggering + # preflight compression against the cheap threshold or (b) letting the + # cheap API call fail and hit in-loop compression. Both outcomes would + # permanently compress a session sized for the primary model. + if current_request_tokens > 0: + try: + from agent.model_metadata import get_model_context_length + cheap_ctx = get_model_context_length( + route.get("model") or "", + base_url=runtime.get("base_url") or "", + api_key=runtime.get("api_key") or "", + provider=runtime.get("provider"), + ) + except Exception: + cheap_ctx = 0 + if cheap_ctx and current_request_tokens > int(cheap_ctx * max_history_ratio): + logger.info( + "Smart route refused: est %d tokens > %.0f%% of %s context %d (staying on primary)", + current_request_tokens, + max_history_ratio * 100, + route.get("model") or "", + cheap_ctx, + ) + return _primary_route(primary) return { "model": route.get("model"), diff --git a/cli.py b/cli.py index 4102dcf021b4..d384862cfde8 100644 --- a/cli.py +++ b/cli.py @@ -2794,6 +2794,33 @@ def _resolve_turn_agent_config(self, user_message: str) -> dict: from agent.smart_model_routing import resolve_turn_route from hermes_cli.models import resolve_fast_mode_overrides + # Best-effort token estimate for the refuse-to-route check in + # smart_model_routing. Short-circuit when smart routing is disabled + # so we don't pay the estimation cost on every turn unnecessarily. + _est_tokens = 0 + _max_history_ratio = 0.50 + if self._smart_model_routing.get("enabled"): + from agent.model_metadata import estimate_request_tokens_rough + _agent = getattr(self, "agent", None) + _sys_prompt = "" + _tools = None + if _agent is not None: + _sys_prompt = getattr(_agent, "_cached_system_prompt", "") or "" + _tools = getattr(_agent, "tools", None) + _cc = getattr(_agent, "context_compressor", None) + if _cc is not None: + _max_history_ratio = ( + getattr(_cc, "threshold_percent", 0.50) or 0.50 + ) + try: + _est_tokens = estimate_request_tokens_rough( + self.conversation_history, + system_prompt=_sys_prompt, + tools=_tools, + ) + except Exception: + _est_tokens = 0 + route = resolve_turn_route( user_message, self._smart_model_routing, @@ -2807,6 +2834,8 @@ def _resolve_turn_agent_config(self, user_message: str) -> dict: "args": list(self.acp_args or []), "credential_pool": getattr(self, "_credential_pool", None), }, + current_request_tokens=_est_tokens, + max_history_ratio=_max_history_ratio, ) service_tier = getattr(self, "service_tier", None) @@ -5911,6 +5940,7 @@ def _handle_btw_command(self, cmd: str): turn_route = self._resolve_turn_agent_config(question) history_snapshot = list(self.conversation_history) + _is_smart_routed_turn = bool(turn_route.get("label")) preview = question[:60] + ("..." if len(question) > 60 else "") _cprint(f' 💬 /btw: "{preview}"') @@ -5956,6 +5986,7 @@ def run_btw(): user_message=btw_prompt, conversation_history=history_snapshot, task_id=task_id, + skip_preflight_compression=_is_smart_routed_turn, ) response = (result.get("final_response") or "") if result else "" @@ -7719,6 +7750,13 @@ def chat(self, message, images: list = None) -> Optional[str]: if turn_route["signature"] != self._active_agent_route_signature: self.agent = None + # Smart routing picks a temporary per-turn model (e.g. cheap fallback), + # which rebuilds the agent with a ContextCompressor bound to that + # model's context_length. Preflight compression would then fire + # against the temporary threshold and compress history sized for the + # primary model. Mark the turn so run_conversation skips preflight. + _is_smart_routed_turn = bool(turn_route.get("label")) + # Initialize agent if needed if self.agent is None: _cprint(f"{_DIM}Initializing agent...{_RST}") @@ -7867,6 +7905,7 @@ def run_agent(): stream_callback=stream_callback, task_id=self.session_id, persist_user_message=message if _voice_prefix else None, + skip_preflight_compression=_is_smart_routed_turn, ) except Exception as exc: logging.error("run_conversation raised: %s", exc, exc_info=True) @@ -10282,6 +10321,7 @@ def main( turn_route = cli._resolve_turn_agent_config(effective_query) if turn_route["signature"] != cli._active_agent_route_signature: cli.agent = None + _is_smart_routed_turn = bool(turn_route.get("label")) if cli._init_agent( model_override=turn_route["model"], runtime_override=turn_route["runtime"], @@ -10298,6 +10338,7 @@ def main( result = cli.agent.run_conversation( user_message=effective_query, conversation_history=cli.conversation_history, + skip_preflight_compression=_is_smart_routed_turn, ) response = result.get("final_response", "") if isinstance(result, dict) else str(result) if response: diff --git a/run_agent.py b/run_agent.py index 8db57a703e51..2fd32d55af72 100644 --- a/run_agent.py +++ b/run_agent.py @@ -8289,6 +8289,7 @@ def run_conversation( task_id: str = None, stream_callback: Optional[callable] = None, persist_user_message: Optional[str] = None, + skip_preflight_compression: bool = False, ) -> Dict[str, Any]: """ Run a complete conversation with tool calling until completion. @@ -8305,6 +8306,10 @@ def run_conversation( transcripts/history when user_message contains API-only synthetic prefixes. or queuing follow-up prefetch work. + skip_preflight_compression: Skip the preflight compression check for + this turn. Used when the agent was temporarily swapped to a + smaller-context model (e.g. smart routing to a cheap model). + In-loop compression on API error still handles the fallback. Returns: Dict: Complete conversation result with final response and message history @@ -8496,8 +8501,11 @@ def run_conversation( # while having a large existing session — compress proactively rather # than waiting for an API error (which might be caught as a non-retryable # 4xx and abort the request entirely). + # Skipped when the caller temporarily swapped to a smaller-context model + # for one turn (e.g. smart routing) — the session is sized for primary. if ( - self.compression_enabled + not skip_preflight_compression + and self.compression_enabled and len(messages) > self.context_compressor.protect_first_n + self.context_compressor.protect_last_n + 1 ): diff --git a/tests/agent/test_credential_pool_routing.py b/tests/agent/test_credential_pool_routing.py index 38f5c6dfd058..0269a2c972b6 100644 --- a/tests/agent/test_credential_pool_routing.py +++ b/tests/agent/test_credential_pool_routing.py @@ -115,9 +115,9 @@ def test_resolve_turn_includes_pool(self, monkeypatch, tmp_path): from agent.smart_model_routing import resolve_turn_route captured = {} - def spy_resolve(user_message, routing_config, primary): + def spy_resolve(user_message, routing_config, primary, **kwargs): captured["primary"] = primary - return resolve_turn_route(user_message, routing_config, primary) + return resolve_turn_route(user_message, routing_config, primary, **kwargs) monkeypatch.setattr( "agent.smart_model_routing.resolve_turn_route", spy_resolve @@ -151,9 +151,9 @@ def test_resolve_turn_includes_pool(self, monkeypatch): from agent.smart_model_routing import resolve_turn_route captured = {} - def spy_resolve(user_message, routing_config, primary): + def spy_resolve(user_message, routing_config, primary, **kwargs): captured["primary"] = primary - return resolve_turn_route(user_message, routing_config, primary) + return resolve_turn_route(user_message, routing_config, primary, **kwargs) monkeypatch.setattr( "agent.smart_model_routing.resolve_turn_route", spy_resolve diff --git a/tests/agent/test_smart_model_routing.py b/tests/agent/test_smart_model_routing.py index 7e902560953f..005804f1c58c 100644 --- a/tests/agent/test_smart_model_routing.py +++ b/tests/agent/test_smart_model_routing.py @@ -59,3 +59,152 @@ def test_resolve_turn_route_falls_back_to_primary_when_route_runtime_cannot_be_r assert result["model"] == "anthropic/claude-sonnet-4" assert result["runtime"]["provider"] == "openrouter" assert result["label"] is None + + +# --------------------------------------------------------------------------- +# Refuse-to-route: when estimated request tokens exceed cheap_ctx * ratio, +# smart routing must fall back to primary instead of letting the request hit +# the cheap model (which would trigger destructive preflight/fallback +# compression on a session sized for the primary model). +# --------------------------------------------------------------------------- + +_REFUSE_PRIMARY = { + "model": "anthropic/claude-sonnet-4.6", + "provider": "openrouter", + "base_url": "https://openrouter.ai/api/v1", + "api_mode": "chat_completions", + "api_key": "sk-primary", +} + + +def _mock_runtime_ok(monkeypatch): + """Make resolve_runtime_provider return a fixed cheap-side runtime.""" + monkeypatch.setattr( + "hermes_cli.runtime_provider.resolve_runtime_provider", + lambda **kwargs: { + "api_key": "sk-cheap", + "base_url": "https://openrouter.ai/api/v1", + "provider": "openrouter", + "api_mode": "chat_completions", + "command": None, + "args": (), + "credential_pool": None, + }, + ) + + +def _mock_cheap_context_length(monkeypatch, ctx_length: int): + """Make get_model_context_length always return ``ctx_length``.""" + monkeypatch.setattr( + "agent.model_metadata.get_model_context_length", + lambda *args, **kwargs: ctx_length, + ) + + +def test_resolve_turn_route_refuses_when_history_exceeds_cheap_ratio(monkeypatch): + """80K > 128K * 0.50 = 64K → refuse and fall back to primary.""" + from agent.smart_model_routing import resolve_turn_route + + _mock_runtime_ok(monkeypatch) + _mock_cheap_context_length(monkeypatch, 128_000) + + result = resolve_turn_route( + "hi", + _BASE_CONFIG, + _REFUSE_PRIMARY, + current_request_tokens=80_000, + max_history_ratio=0.50, + ) + + assert result["model"] == _REFUSE_PRIMARY["model"] + assert result["label"] is None + + +def test_resolve_turn_route_allows_when_history_fits_cheap_ratio(monkeypatch): + """50K < 128K * 0.50 = 64K → use cheap route.""" + from agent.smart_model_routing import resolve_turn_route + + _mock_runtime_ok(monkeypatch) + _mock_cheap_context_length(monkeypatch, 128_000) + + result = resolve_turn_route( + "hi", + _BASE_CONFIG, + _REFUSE_PRIMARY, + current_request_tokens=50_000, + max_history_ratio=0.50, + ) + + assert result["model"] == "google/gemini-2.5-flash" + assert result["label"] is not None + + +def test_resolve_turn_route_zero_token_count_skips_check(monkeypatch): + """Backward compat: current_request_tokens=0 (caller didn't estimate) + must skip the refuse check entirely, even when cheap_ctx is small.""" + from agent.smart_model_routing import resolve_turn_route + + _mock_runtime_ok(monkeypatch) + + # get_model_context_length should never be called in this path. + def _should_not_be_called(*args, **kwargs): + raise AssertionError("get_model_context_length should not be called when current_request_tokens=0") + + monkeypatch.setattr( + "agent.model_metadata.get_model_context_length", + _should_not_be_called, + ) + + result = resolve_turn_route( + "hi", + _BASE_CONFIG, + _REFUSE_PRIMARY, + # current_request_tokens defaults to 0 + ) + + assert result["model"] == "google/gemini-2.5-flash" + assert result["label"] is not None + + +def test_resolve_turn_route_custom_max_history_ratio(monkeypatch): + """A caller that tunes threshold_percent (e.g. to 0.80) propagates to + smart routing refusal — 90K fits under 128K * 0.80 = 102K → cheap route.""" + from agent.smart_model_routing import resolve_turn_route + + _mock_runtime_ok(monkeypatch) + _mock_cheap_context_length(monkeypatch, 128_000) + + # At default 0.50: 90K > 64K would refuse. + # At 0.80: 90K < 102K → allow. + result = resolve_turn_route( + "hi", + _BASE_CONFIG, + _REFUSE_PRIMARY, + current_request_tokens=90_000, + max_history_ratio=0.80, + ) + + assert result["model"] == "google/gemini-2.5-flash" + assert result["label"] is not None + + +def test_resolve_turn_route_refuses_when_cheap_context_unknown_and_estimate_positive(monkeypatch): + """Defensive: if get_model_context_length returns 0 (unknown), we can't + compare against the threshold, so the refusal check is a no-op and the + cheap route is used. This preserves existing behavior for providers + whose context length can't be resolved at all.""" + from agent.smart_model_routing import resolve_turn_route + + _mock_runtime_ok(monkeypatch) + _mock_cheap_context_length(monkeypatch, 0) # unknown + + result = resolve_turn_route( + "hi", + _BASE_CONFIG, + _REFUSE_PRIMARY, + current_request_tokens=1_000_000, # huge, but cheap_ctx=0 means skip check + max_history_ratio=0.50, + ) + + assert result["model"] == "google/gemini-2.5-flash" + assert result["label"] is not None diff --git a/tests/run_agent/test_413_compression.py b/tests/run_agent/test_413_compression.py index 1d6f6cebb822..e76c118b7e70 100644 --- a/tests/run_agent/test_413_compression.py +++ b/tests/run_agent/test_413_compression.py @@ -492,6 +492,79 @@ def test_no_preflight_when_compression_disabled(self, agent): mock_compress.assert_not_called() + def test_skip_preflight_compression_flag(self, agent): + """skip_preflight_compression=True must bypass preflight even when over threshold. + + Regression for smart_model_routing: when the agent is temporarily swapped to a + smaller-context model for one turn, preflight compression fires against the + temporary model's threshold and permanently compresses history sized for the + primary model. Callers can pass skip_preflight_compression=True to opt out. + """ + agent.compression_enabled = True + # Small context so history easily exceeds threshold + agent.context_compressor.context_length = 2000 + agent.context_compressor.threshold_tokens = 200 + + big_history = [] + for i in range(20): + big_history.append({"role": "user", "content": f"Message {i} with padding"}) + big_history.append({"role": "assistant", "content": f"Response {i} with padding"}) + + ok_resp = _mock_response(content="Not compressed", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [ok_resp] + + with ( + patch.object(agent, "_compress_context") as mock_compress, + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation( + "hello", + conversation_history=big_history, + skip_preflight_compression=True, + ) + + mock_compress.assert_not_called() + assert result["completed"] is True + assert result["final_response"] == "Not compressed" + + def test_skip_preflight_false_still_compresses(self, agent): + """skip_preflight_compression=False (default) must preserve existing behavior.""" + agent.compression_enabled = True + agent.context_compressor.context_length = 2000 + agent.context_compressor.threshold_tokens = 200 + + big_history = [] + for i in range(20): + big_history.append({"role": "user", "content": f"Message {i} with padding"}) + big_history.append({"role": "assistant", "content": f"Response {i} with padding"}) + + ok_resp = _mock_response(content="After compression", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [ok_resp] + + with ( + patch.object(agent, "_compress_context") as mock_compress, + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + mock_compress.return_value = ( + [ + {"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}, + {"role": "user", "content": "hello"}, + ], + "new system prompt", + ) + result = agent.run_conversation( + "hello", + conversation_history=big_history, + skip_preflight_compression=False, + ) + + mock_compress.assert_called_once() + assert result["completed"] is True + class TestToolResultPreflightCompression: """Compression should trigger when tool results push context past the threshold."""