From 3b750ed29fcd8f3f1a8a5262d79d5f3cd22e4f9d Mon Sep 17 00:00:00 2001
From: inside-ziwu <22082285+inside-ziwu@users.noreply.github.com>
Date: Sun, 12 Apr 2026 16:41:46 +0800
Subject: [PATCH] fix(compression): reuse active model context for feasibility
 checks

---
 run_agent.py                                  | 28 ++++++++--
 .../run_agent/test_compression_feasibility.py | 52 +++++++++++++++++++
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/run_agent.py b/run_agent.py
index b230354542521..6080a7e5afb99 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -1733,11 +1733,29 @@ def _check_compression_model_feasibility(self) -> None:
 
             aux_base_url = str(getattr(client, "base_url", ""))
             aux_api_key = str(getattr(client, "api_key", ""))
-            aux_context = get_model_context_length(
-                aux_model,
-                base_url=aux_base_url,
-                api_key=aux_api_key,
-            )
+
+            def _normalize_model_name(model_name: str) -> str:
+                model_name = (model_name or "").strip().lower()
+                if "/" in model_name:
+                    model_name = model_name.rsplit("/", 1)[-1]
+                return model_name
+
+            same_model = _normalize_model_name(aux_model) == _normalize_model_name(self.model)
+            same_endpoint = aux_base_url.rstrip("/") == str(getattr(self, "base_url", "")).rstrip("/")
+
+            if same_model and same_endpoint:
+                # Compression defaults to the active model. Reuse the already
+                # resolved main-model context so explicit config overrides
+                # (for example model.context_length on custom endpoints) are
+                # honored instead of re-querying metadata and falling back to
+                # generic family defaults like GPT-5 = 128K.
+                aux_context = self.context_compressor.context_length
+            else:
+                aux_context = get_model_context_length(
+                    aux_model,
+                    base_url=aux_base_url,
+                    api_key=aux_api_key,
+                )
 
             threshold = self.context_compressor.threshold_tokens
             if aux_context < threshold:
diff --git a/tests/run_agent/test_compression_feasibility.py b/tests/run_agent/test_compression_feasibility.py
index 1b4423414ee6d..cd25f7b3fbb79 100644
--- a/tests/run_agent/test_compression_feasibility.py
+++ b/tests/run_agent/test_compression_feasibility.py
@@ -181,6 +181,58 @@ def test_just_below_threshold_warns(mock_get_client, mock_ctx_len):
     assert "small-model" in messages[0]
 
 
+@patch("agent.model_metadata.get_model_context_length", return_value=128_000)
+@patch("agent.auxiliary_client.get_text_auxiliary_client")
+def test_reuses_main_context_when_compression_model_matches_active_model(
+    mock_get_client, mock_ctx_len
+):
+    """If compression resolves to the active model on the same endpoint,
+    reuse the already-resolved main context instead of re-querying metadata.
+    """
+    agent = _make_agent(main_context=1_000_000, threshold_percent=0.80)
+    agent.model = "gpt-5.4"
+    agent.base_url = "https://sub2api.tizi.lifestyle/v1"
+
+    mock_client = MagicMock()
+    mock_client.base_url = "https://sub2api.tizi.lifestyle/v1/"
+    mock_client.api_key = "sk-aux"
+    mock_get_client.return_value = (mock_client, "gpt-5.4")
+
+    messages = []
+    agent._emit_status = lambda msg: messages.append(msg)
+
+    agent._check_compression_model_feasibility()
+
+    assert len(messages) == 0
+    assert agent._compression_warning is None
+    mock_ctx_len.assert_not_called()
+
+
+@patch("agent.model_metadata.get_model_context_length", return_value=32_768)
+@patch("agent.auxiliary_client.get_text_auxiliary_client")
+def test_explicit_different_compression_model_uses_its_own_context(
+    mock_get_client, mock_ctx_len
+):
+    """A distinct compression model should still be checked independently."""
+    agent = _make_agent(main_context=1_000_000, threshold_percent=0.80)
+    agent.model = "gpt-5.4"
+    agent.base_url = "https://sub2api.tizi.lifestyle/v1"
+
+    mock_client = MagicMock()
+    mock_client.base_url = "https://openrouter.ai/api/v1"
+    mock_client.api_key = "sk-aux"
+    mock_get_client.return_value = (mock_client, "google/gemini-3-flash-preview")
+
+    messages = []
+    agent._emit_status = lambda msg: messages.append(msg)
+
+    agent._check_compression_model_feasibility()
+
+    assert len(messages) == 1
+    assert "google/gemini-3-flash-preview" in messages[0]
+    mock_ctx_len.assert_called_once()
+
+
 # ── Two-phase: __init__ + run_conversation replay ───────────────────