1  """Tests that switch_model preserves config_context_length."""
 2  
 3  from unittest.mock import MagicMock, patch
 4  
 5  from run_agent import AIAgent
 6  from agent.context_compressor import ContextCompressor
 7  
 8  
 9  def _make_agent_with_compressor(config_context_length=None) -> AIAgent:
10      """Build a minimal AIAgent with a context_compressor, skipping __init__."""
11      agent = AIAgent.__new__(AIAgent)
12  
13      # Primary model settings
14      agent.model = "primary-model"
15      agent.provider = "openrouter"
16      agent.base_url = "https://openrouter.ai/api/v1"
17      agent.api_key = "sk-primary"
18      agent.api_mode = "chat_completions"
19      agent.client = MagicMock()
20      agent.quiet_mode = True
21  
22      # Store config_context_length for later use in switch_model
23      agent._config_context_length = config_context_length
24  
25      # Context compressor with primary model values
26      compressor = ContextCompressor(
27          model="primary-model",
28          threshold_percent=0.50,
29          base_url="https://openrouter.ai/api/v1",
30          api_key="sk-primary",
31          provider="openrouter",
32          quiet_mode=True,
33          config_context_length=config_context_length,
34      )
35      agent.context_compressor = compressor
36  
37      # For switch_model
38      agent._primary_runtime = {}
39  
40      return agent
41  
42  
43  @patch("agent.model_metadata.get_model_context_length", return_value=131_072)
44  def test_switch_model_preserves_config_context_length(mock_ctx_len):
45      """When switching models, config_context_length should be passed to get_model_context_length."""
46      agent = _make_agent_with_compressor(config_context_length=32_768)
47  
48      assert agent.context_compressor.model == "primary-model"
49      assert agent.context_compressor.context_length == 32_768  # From config override
50  
51      # Switch model
52      agent.switch_model("new-model", "openrouter", api_key="sk-new", base_url="https://openrouter.ai/api/v1")
53  
54      # Verify get_model_context_length was called with config_context_length
55      mock_ctx_len.assert_called_once()
56      call_kwargs = mock_ctx_len.call_args.kwargs
57      assert call_kwargs.get("config_context_length") == 32_768
58  
59      # Verify compressor was updated
60      assert agent.context_compressor.model == "new-model"
61  
62  
63  def test_switch_model_without_config_context_length():
64      """When switching models without config override, config_context_length should be None."""
65      agent = _make_agent_with_compressor(config_context_length=None)
66  
67      with patch("agent.model_metadata.get_model_context_length", return_value=128_000) as mock_ctx_len:
68          # Switch model
69          agent.switch_model("new-model", "openrouter", api_key="sk-new", base_url="https://openrouter.ai/api/v1")
70  
71          # Verify get_model_context_length was called with None
72          mock_ctx_len.assert_called_once()
73          call_kwargs = mock_ctx_len.call_args.kwargs
74          assert call_kwargs.get("config_context_length") is None