From 88e2a923f5948f11d6a94212159e819e27836dcb Mon Sep 17 00:00:00 2001 From: slicenfer <16222938+slicenfer@user.noreply.gitee.com> Date: Mon, 13 Apr 2026 11:57:25 +0800 Subject: [PATCH] fix: unify context length overrides for compression runtimes --- run_agent.py | 178 ++++++++---- tests/run_agent/test_switch_model_context.py | 270 ++++++++++++++++--- 2 files changed, 362 insertions(+), 86 deletions(-) diff --git a/run_agent.py b/run_agent.py index 37572db5e1f1..3fec7d712997 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1085,6 +1085,9 @@ def __init__( _agent_cfg = _load_agent_config() except Exception: _agent_cfg = {} + if not isinstance(_agent_cfg, dict): + _agent_cfg = {} + self._agent_config = _agent_cfg # Persistent memory (MEMORY.md + USER.md) -- loaded from disk self._memory_store = None @@ -1216,41 +1219,15 @@ def __init__( compression_target_ratio = float(_compression_cfg.get("target_ratio", 0.20)) compression_protect_last = int(_compression_cfg.get("protect_last_n", 20)) - # Read explicit context_length override from model config - _model_cfg = _agent_cfg.get("model", {}) - if isinstance(_model_cfg, dict): - _config_context_length = _model_cfg.get("context_length") - else: - _config_context_length = None - if _config_context_length is not None: - try: - _config_context_length = int(_config_context_length) - except (TypeError, ValueError): - _config_context_length = None + _model_cfg = _agent_cfg.get("model", {}) if isinstance(_agent_cfg, dict) else {} + _config_context_length = self._resolve_context_length_override( + model=self.model, + base_url=self.base_url, + config=_agent_cfg, + ) - # Store for reuse in switch_model (so config override persists across model switches) + # Store the active runtime override for reuse in switch_model. self._config_context_length = _config_context_length - - # Check custom_providers per-model context_length - if _config_context_length is None: - _custom_providers = _agent_cfg.get("custom_providers") - if isinstance(_custom_providers, list): - for _cp_entry in _custom_providers: - if not isinstance(_cp_entry, dict): - continue - _cp_url = (_cp_entry.get("base_url") or "").rstrip("/") - if _cp_url and _cp_url == self.base_url.rstrip("/"): - _cp_models = _cp_entry.get("models", {}) - if isinstance(_cp_models, dict): - _cp_model_cfg = _cp_models.get(self.model, {}) - if isinstance(_cp_model_cfg, dict): - _cp_ctx = _cp_model_cfg.get("context_length") - if _cp_ctx is not None: - try: - _config_context_length = int(_cp_ctx) - except (TypeError, ValueError): - pass - break # Select context engine: config-driven (like memory providers). # 1. Check config.yaml context.engine setting @@ -1417,6 +1394,7 @@ def __init__( "base_url": self.base_url, "api_mode": self.api_mode, "api_key": getattr(self, "api_key", ""), + "config_context_length": self._config_context_length, "client_kwargs": dict(self._client_kwargs), "use_prompt_caching": self._use_prompt_caching, # Context engine state that _try_activate_fallback() overwrites. @@ -1548,6 +1526,11 @@ def switch_model(self, new_model, new_provider, api_key='', base_url='', api_mod or is_native_anthropic ) + self._config_context_length = self._resolve_context_length_override( + model=self.model, + base_url=self.base_url, + ) + # ── Update context compressor ── if hasattr(self, "context_compressor") and self.context_compressor: from agent.model_metadata import get_model_context_length @@ -1556,7 +1539,7 @@ def switch_model(self, new_model, new_provider, api_key='', base_url='', api_mod base_url=self.base_url, api_key=self.api_key, provider=self.provider, - config_context_length=getattr(self, "_config_context_length", None), + config_context_length=self._config_context_length, ) self.context_compressor.update_model( model=self.model, @@ -1578,6 +1561,7 @@ def switch_model(self, new_model, new_provider, api_key='', base_url='', api_mod "base_url": self.base_url, "api_mode": self.api_mode, "api_key": getattr(self, "api_key", ""), + "config_context_length": self._config_context_length, "client_kwargs": dict(self._client_kwargs), "use_prompt_caching": self._use_prompt_caching, "compressor_model": getattr(_cc, "model", self.model) if _cc else self.model, @@ -1708,6 +1692,96 @@ def _current_main_runtime(self) -> Dict[str, str]: "api_mode": getattr(self, "api_mode", "") or "", } + @staticmethod + def _coerce_context_length_override(value: Any) -> Optional[int]: + """Normalize a configured context length override to a positive int.""" + if value is None: + return None + try: + value = int(value) + except (TypeError, ValueError): + return None + return value if value > 0 else None + + def _resolve_context_length_override( + self, + model: str = "", + base_url: str = "", + config: Optional[Dict[str, Any]] = None, + ) -> Optional[int]: + """Resolve a configured context_length override for a runtime. + + Resolution order mirrors the main context lookup entry points: + 0. direct context_length on the provided config block + (e.g. auxiliary.compression.context_length) + 1. top-level model.context_length for the configured default model + 2. custom_providers[].models[model].context_length for matching base_url + """ + from agent.model_metadata import _strip_provider_prefix + + cfg = config if isinstance(config, dict) else getattr(self, "_agent_config", {}) + if not isinstance(cfg, dict): + return None + + direct_override = self._coerce_context_length_override(cfg.get("context_length")) + if direct_override is not None: + return direct_override + + target_model = model or getattr(self, "model", "") or "" + target_model_stripped = _strip_provider_prefix(target_model) + target_base_url = (base_url or getattr(self, "base_url", "") or "").rstrip("/") + + model_cfg = cfg.get("model", {}) + if isinstance(model_cfg, dict): + configured_model = model_cfg.get("default") or model_cfg.get("model") or "" + configured_model_stripped = _strip_provider_prefix(configured_model) + configured_base_url = (model_cfg.get("base_url") or "").rstrip("/") + if ( + not target_model + or target_model == configured_model + or (target_model_stripped and target_model_stripped == configured_model_stripped) + ) and ( + not target_base_url + or not configured_base_url + or target_base_url == configured_base_url + ): + override = self._coerce_context_length_override( + model_cfg.get("context_length") + ) + if override is not None: + return override + + if not target_base_url: + return None + + custom_providers = cfg.get("custom_providers") + if not isinstance(custom_providers, list): + return None + + for entry in custom_providers: + if not isinstance(entry, dict): + continue + entry_base_url = (entry.get("base_url") or "").rstrip("/") + if not entry_base_url or entry_base_url != target_base_url: + continue + models_cfg = entry.get("models", {}) + if not isinstance(models_cfg, dict): + break + + model_cfg = models_cfg.get(target_model) + if not isinstance(model_cfg, dict) and target_model_stripped: + model_cfg = models_cfg.get(target_model_stripped) + + if isinstance(model_cfg, dict): + override = self._coerce_context_length_override( + model_cfg.get("context_length") + ) + if override is not None: + return override + break + + return None + def _check_compression_model_feasibility(self) -> None: """Warn at session start if the auxiliary compression model's context window is smaller than the main model's compression threshold. @@ -1748,25 +1822,24 @@ def _check_compression_model_feasibility(self) -> None: aux_base_url = str(getattr(client, "base_url", "")) aux_api_key = str(getattr(client, "api_key", "")) - - # Read user-configured context_length for the compression model. - # Custom endpoints often don't support /models API queries so - # get_model_context_length() falls through to the 128K default, - # ignoring the explicit config value. Pass it as the highest- - # priority hint so the configured value is always respected. - _aux_cfg = (self.config or {}).get("auxiliary", {}).get("compression", {}) - _aux_context_config = _aux_cfg.get("context_length") if isinstance(_aux_cfg, dict) else None - if _aux_context_config is not None: - try: - _aux_context_config = int(_aux_context_config) - except (TypeError, ValueError): - _aux_context_config = None - + _aux_cfg = (getattr(self, "config", {}) or {}).get("auxiliary", {}).get("compression", {}) + if not isinstance(_aux_cfg, dict): + _aux_cfg = {} + aux_context_override = self._resolve_context_length_override( + model=aux_model, + base_url=aux_base_url, + config=_aux_cfg, + ) + if aux_context_override is None: + aux_context_override = self._resolve_context_length_override( + model=aux_model, + base_url=aux_base_url, + ) aux_context = get_model_context_length( aux_model, base_url=aux_base_url, api_key=aux_api_key, - config_context_length=_aux_context_config, + config_context_length=aux_context_override, ) threshold = self.context_compressor.threshold_tokens @@ -5531,9 +5604,14 @@ def _try_activate_fallback(self) -> bool: # causing oversized sessions to overflow the fallback. if hasattr(self, 'context_compressor') and self.context_compressor: from agent.model_metadata import get_model_context_length + self._config_context_length = self._resolve_context_length_override( + model=self.model, + base_url=self.base_url, + ) fb_context_length = get_model_context_length( self.model, base_url=self.base_url, api_key=self.api_key, provider=self.provider, + config_context_length=self._config_context_length, ) self.context_compressor.update_model( model=self.model, @@ -5541,6 +5619,7 @@ def _try_activate_fallback(self) -> bool: base_url=self.base_url, api_key=getattr(self, "api_key", ""), provider=self.provider, + api_mode=self.api_mode, ) self._emit_status( @@ -5580,6 +5659,7 @@ def _restore_primary_runtime(self) -> bool: self.base_url = rt["base_url"] # setter updates _base_url_lower self.api_mode = rt["api_mode"] self.api_key = rt["api_key"] + self._config_context_length = rt.get("config_context_length") self._client_kwargs = dict(rt["client_kwargs"]) self._use_prompt_caching = rt["use_prompt_caching"] diff --git a/tests/run_agent/test_switch_model_context.py b/tests/run_agent/test_switch_model_context.py index 8b04a73262b5..827d45d281ec 100644 --- a/tests/run_agent/test_switch_model_context.py +++ b/tests/run_agent/test_switch_model_context.py @@ -1,74 +1,270 @@ -"""Tests that switch_model preserves config_context_length.""" +"""Tests for runtime context_length overrides in AIAgent.""" +from types import SimpleNamespace from unittest.mock import MagicMock, patch -from run_agent import AIAgent from agent.context_compressor import ContextCompressor +from run_agent import AIAgent + +CUSTOM_BASE_URL = "http://localhost:4000/v1" -def _make_agent_with_compressor(config_context_length=None) -> AIAgent: + +def _make_agent_with_compressor( + config_context_length=None, + agent_config=None, + model="primary-model", + provider="openrouter", + base_url="https://openrouter.ai/api/v1", +) -> AIAgent: """Build a minimal AIAgent with a context_compressor, skipping __init__.""" agent = AIAgent.__new__(AIAgent) - # Primary model settings - agent.model = "primary-model" - agent.provider = "openrouter" - agent.base_url = "https://openrouter.ai/api/v1" + agent.model = model + agent.provider = provider + agent.base_url = base_url agent.api_key = "sk-primary" agent.api_mode = "chat_completions" agent.client = MagicMock() agent.quiet_mode = True - - # Store config_context_length for later use in switch_model + agent._agent_config = agent_config or {} agent._config_context_length = config_context_length + agent._client_kwargs = {} + agent._use_prompt_caching = False + agent._cached_system_prompt = None + agent._create_openai_client = MagicMock(return_value=MagicMock()) - # Context compressor with primary model values compressor = ContextCompressor( - model="primary-model", + model=model, threshold_percent=0.50, - base_url="https://openrouter.ai/api/v1", + base_url=base_url, api_key="sk-primary", - provider="openrouter", + provider=provider, quiet_mode=True, config_context_length=config_context_length, ) agent.context_compressor = compressor - - # For switch_model agent._primary_runtime = {} return agent -@patch("agent.model_metadata.get_model_context_length", return_value=131_072) -def test_switch_model_preserves_config_context_length(mock_ctx_len): - """When switching models, config_context_length should be passed to get_model_context_length.""" - agent = _make_agent_with_compressor(config_context_length=32_768) +@patch("run_agent.AIAgent._check_compression_model_feasibility") +@patch("run_agent.ContextCompressor") +@patch("run_agent.get_tool_definitions", return_value=[]) +@patch("run_agent.check_toolset_requirements", return_value={}) +@patch("run_agent.AIAgent._create_openai_client", return_value=MagicMock()) +@patch("hermes_cli.config.load_config") +def test_init_stores_custom_provider_context_override( + mock_load_config, + _mock_create_client, + _mock_requirements, + _mock_get_tools, + mock_context_compressor, + _mock_feasibility, +): + """__init__ should store and pass the custom_providers per-model override.""" + mock_load_config.return_value = { + "model": { + "default": "gpt-5.4", + "provider": "custom", + "base_url": CUSTOM_BASE_URL, + }, + "compression": {"enabled": False}, + "custom_providers": [ + { + "base_url": CUSTOM_BASE_URL, + "models": { + "gpt-5.4": {"context_length": 500_000}, + }, + } + ], + } + + fake_compressor = MagicMock() + fake_compressor.context_length = 500_000 + fake_compressor.threshold_tokens = 250_000 + fake_compressor.get_tool_schemas.return_value = [] + mock_context_compressor.return_value = fake_compressor + + agent = AIAgent( + model="gpt-5.4", + provider="custom", + base_url=CUSTOM_BASE_URL, + api_key="sk-test", + quiet_mode=True, + skip_memory=True, + ) + + assert agent._config_context_length == 500_000 + assert mock_context_compressor.call_args.kwargs["config_context_length"] == 500_000 + - assert agent.context_compressor.model == "primary-model" - assert agent.context_compressor.context_length == 32_768 # From config override +@patch("agent.model_metadata.get_model_context_length", return_value=262_144) +def test_switch_model_refreshes_context_override_from_agent_config(mock_ctx_len): + """switch_model() should resolve the new runtime's config override before lookup.""" + agent = _make_agent_with_compressor( + config_context_length=128_000, + agent_config={ + "custom_providers": [ + { + "base_url": CUSTOM_BASE_URL, + "models": { + "new-model": {"context_length": 262_144}, + }, + } + ] + }, + provider="custom", + base_url=CUSTOM_BASE_URL, + ) - # Switch model - agent.switch_model("new-model", "openrouter", api_key="sk-new", base_url="https://openrouter.ai/api/v1") + agent.switch_model( + "new-model", + "custom", + api_key="sk-new", + base_url=CUSTOM_BASE_URL, + api_mode="chat_completions", + ) - # Verify get_model_context_length was called with config_context_length mock_ctx_len.assert_called_once() call_kwargs = mock_ctx_len.call_args.kwargs - assert call_kwargs.get("config_context_length") == 32_768 - - # Verify compressor was updated + assert call_kwargs["config_context_length"] == 262_144 + assert agent._config_context_length == 262_144 assert agent.context_compressor.model == "new-model" -def test_switch_model_without_config_context_length(): - """When switching models without config override, config_context_length should be None.""" - agent = _make_agent_with_compressor(config_context_length=None) +@patch("agent.model_metadata.get_model_context_length", return_value=64_000) +def test_switch_model_clears_context_override_when_new_model_has_none(mock_ctx_len): + """switch_model() should clear stale overrides when the new model has none.""" + agent = _make_agent_with_compressor( + config_context_length=128_000, + agent_config={ + "custom_providers": [ + { + "base_url": CUSTOM_BASE_URL, + "models": { + "primary-model": {"context_length": 128_000}, + }, + } + ] + }, + provider="custom", + base_url=CUSTOM_BASE_URL, + ) + + agent.switch_model( + "model-without-override", + "custom", + api_key="sk-new", + base_url=CUSTOM_BASE_URL, + api_mode="chat_completions", + ) - with patch("agent.model_metadata.get_model_context_length", return_value=128_000) as mock_ctx_len: - # Switch model - agent.switch_model("new-model", "openrouter", api_key="sk-new", base_url="https://openrouter.ai/api/v1") + mock_ctx_len.assert_called_once() + call_kwargs = mock_ctx_len.call_args.kwargs + assert call_kwargs["config_context_length"] is None + assert agent._config_context_length is None + assert agent.context_compressor.model == "model-without-override" - # Verify get_model_context_length was called with None - mock_ctx_len.assert_called_once() - call_kwargs = mock_ctx_len.call_args.kwargs - assert call_kwargs.get("config_context_length") is None + +@patch("agent.model_metadata.get_model_context_length", return_value=262_144) +@patch("agent.auxiliary_client.resolve_provider_client") +def test_fallback_refreshes_context_override_from_agent_config( + mock_resolve_provider_client, + mock_ctx_len, +): + """Fallback activation should reuse the shared override resolver.""" + agent = _make_agent_with_compressor( + config_context_length=128_000, + agent_config={ + "custom_providers": [ + { + "base_url": CUSTOM_BASE_URL, + "models": { + "fallback-model": {"context_length": 262_144}, + }, + } + ] + }, + provider="custom", + base_url=CUSTOM_BASE_URL, + ) + agent._fallback_chain = [ + { + "provider": "custom", + "model": "fallback-model", + "base_url": CUSTOM_BASE_URL, + "api_key": "sk-fallback", + } + ] + agent._fallback_index = 0 + agent._fallback_activated = False + agent._emit_status = MagicMock() + + fallback_client = MagicMock() + fallback_client.api_key = "sk-fallback" + fallback_client.base_url = CUSTOM_BASE_URL + mock_resolve_provider_client.return_value = (fallback_client, "fallback-model") + + assert agent._try_activate_fallback() is True + + mock_ctx_len.assert_called_once() + call_kwargs = mock_ctx_len.call_args.kwargs + assert call_kwargs["config_context_length"] == 262_144 + assert agent._config_context_length == 262_144 + assert agent.context_compressor.model == "fallback-model" + + +@patch("agent.model_metadata.get_model_context_length", return_value=500_000) +@patch("agent.auxiliary_client.get_text_auxiliary_client") +def test_compression_feasibility_uses_context_override_for_aux_model( + mock_get_aux_client, + mock_ctx_len, +): + """Compression feasibility should reuse the same override resolution logic.""" + agent = AIAgent.__new__(AIAgent) + agent.model = "gpt-5.4" + agent.base_url = CUSTOM_BASE_URL + agent.api_key = "sk-main" + agent.api_mode = "chat_completions" + agent.provider = "custom" + agent.compression_enabled = True + agent.context_compressor = SimpleNamespace( + threshold_tokens=250_000, + context_length=500_000, + ) + agent._agent_config = { + "custom_providers": [ + { + "base_url": CUSTOM_BASE_URL, + "models": { + "gpt-5.4": {"context_length": 500_000}, + }, + } + ] + } + agent._compression_warning = None + agent._emit_status = MagicMock() + agent._current_main_runtime = MagicMock( + return_value={ + "model": "gpt-5.4", + "provider": "custom", + "base_url": CUSTOM_BASE_URL, + "api_key": "sk-main", + "api_mode": "chat_completions", + } + ) + + mock_get_aux_client.return_value = ( + SimpleNamespace(base_url=CUSTOM_BASE_URL, api_key="sk-aux"), + "gpt-5.4", + ) + + agent._check_compression_model_feasibility() + + mock_ctx_len.assert_called_once() + call_kwargs = mock_ctx_len.call_args.kwargs + assert call_kwargs["config_context_length"] == 500_000 + agent._emit_status.assert_not_called() + assert agent._compression_warning is None