From 4d729e4b31c5c39bb6b1f351a0c3928909db45dc Mon Sep 17 00:00:00 2001 From: Kyzcreig <9063726+Kyzcreig@users.noreply.github.com> Date: Sat, 25 Jul 2026 09:33:41 -0700 Subject: [PATCH] fix(model-metadata): local ctx probe must not read max_tokens as the context window The two local-server context probes in _query_local_context_length read data.get("max_tokens") as a context-window candidate. On an OpenAI-compatible /v1/models passthrough max_tokens is the max OUTPUT tokens, so a 1M-context model advertising a 128K output cap resolves to 128000 and auto-compaction fires ~7x early. Route both branches through the module's own key vocabulary (_CONTEXT_LENGTH_KEYS), which already classifies max_tokens as a _MAX_COMPLETION_KEYS entry. --- agent/model_metadata.py | 43 ++++++++-- tests/agent/test_model_metadata_local_ctx.py | 88 ++++++++++++++++++++ 2 files changed, 123 insertions(+), 8 deletions(-) diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 73783ae66c..b1e4724ce9 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -1152,6 +1152,25 @@ def _extract_first_int(payload: Dict[str, Any], keys: tuple[str, ...]) -> Option return None +def _extract_flat_context_length(payload: Dict[str, Any]) -> Optional[int]: + """Read a context WINDOW from the top level of a model-describe payload. + + Same key vocabulary as :func:`_extract_context_length` (the module's single + source of truth for what counts as a context window), but WITHOUT the + nested-dict walk — for callers that hold a specific model object and must + not pick up a same-named key from an unrelated nested section. + + Critically, ``max_tokens`` is NOT in ``_CONTEXT_LENGTH_KEYS``: it lives in + ``_MAX_COMPLETION_KEYS`` because on an OpenAI-compatible ``/v1/models`` + passthrough it is the max *output* tokens, not the context window. + """ + for key in _CONTEXT_LENGTH_KEYS: + coerced = _coerce_reasonable_int(payload.get(key)) + if coerced is not None: + return coerced + return None + + def _extract_context_length(payload: Dict[str, Any]) -> Optional[int]: return _extract_first_int(payload, _CONTEXT_LENGTH_KEYS) @@ -1164,7 +1183,7 @@ def _context_length_from_model_payload(payload: Dict[str, Any]) -> Optional[int] """Extract a context *window* from a ``/v1/models`` model object. Prefers input-window keys (``max_model_len``, ``max_input_tokens``, - ``context_length``, …) via :func:`_extract_context_length`. Falls back to + ``context_length``, …) via :func:`_extract_flat_context_length`. Falls back to ``max_tokens`` only when no input-window field is present. Anthropic (and Anthropic-compatible proxies such as local reverse @@ -1176,7 +1195,7 @@ def _context_length_from_model_payload(payload: Dict[str, Any]) -> Optional[int] """ if not isinstance(payload, dict): return None - ctx = _extract_context_length(payload) + ctx = _extract_flat_context_length(payload) if ctx is not None: return ctx # Last resort for OpenAI-compat servers that only report max_tokens as @@ -2341,8 +2360,15 @@ def _query_local_context_length_uncached(model: str, base_url: str, api_key: str if resp.status_code == 200: data = resp.json() if isinstance(data, dict): - # Prefer max_model_len / max_input_tokens / context_length - # over max_tokens (Anthropic max_tokens = max OUTPUT). + # Context-WINDOW keys only (canonical _CONTEXT_LENGTH_KEYS + # vocabulary). `max_tokens` is the max *output* tokens on + # OpenAI-compatible passthroughs (LiteLLM, Anthropic-compat + # shims, cloud proxies) — e.g. 393216 for a 1M-context + # model — so reading it ahead of real window keys collapses + # the window to the output cap and poisons the context + # cache. It is consulted only as an explicit last resort + # inside _context_length_from_model_payload, for servers + # that report nothing else. ctx = _context_length_from_model_payload(data) if ctx is not None: return ctx @@ -2380,10 +2406,11 @@ def _query_local_context_length_uncached(model: str, base_url: str, api_key: str if isinstance(val, (int, float)) and val: return int(val) # Canonical context-WINDOW keys (via _CONTEXT_LENGTH_KEYS) - # with max_tokens demoted to an explicit last resort — see - # _context_length_from_model_payload for why max_tokens - # must never win over a real window key (it is the max - # OUTPUT cap on Anthropic/OpenAI-compatible passthroughs). + # with max_tokens demoted to an explicit last resort — + # sibling of the /v1/models/{id} path above; see that + # comment for why max_tokens must never win over a real + # window key (it is the max OUTPUT cap on + # Anthropic/OpenAI-compatible passthroughs). for source in sources: ctx = _context_length_from_model_payload(source) if ctx is not None: diff --git a/tests/agent/test_model_metadata_local_ctx.py b/tests/agent/test_model_metadata_local_ctx.py index c7b3b7e939..ab7e39777f 100644 --- a/tests/agent/test_model_metadata_local_ctx.py +++ b/tests/agent/test_model_metadata_local_ctx.py @@ -130,6 +130,94 @@ class TestQueryLocalContextLengthVllm: assert result == 32768 + def test_detail_branch_reads_context_window_not_output_cap(self): + """A payload carrying BOTH a context window and an output cap must + resolve to the context window. + + An OpenAI-compatible ``/v1/models/{id}`` passthrough (LiteLLM, an + Anthropic-compat shim, a cloud proxy) returns ``max_input_tokens`` — + the context window — alongside ``max_tokens``, the max *output* + tokens. Reading ``max_tokens`` collapses a 1M-context model to its + 128K output cap and drives premature auto-compaction. + + Contract asserted: when a describe payload contains both classes of + key, the resolver returns the ``_CONTEXT_LENGTH_KEYS`` value, never + the ``_MAX_COMPLETION_KEYS`` one. + """ + from agent.model_metadata import _query_local_context_length + + detail_resp = self._make_resp(200, { + "type": "model", + "id": "some-model", + "max_input_tokens": 1000000, # context window + "max_tokens": 128000, # max OUTPUT tokens — not a window + }) + + client_mock = MagicMock() + client_mock.__enter__ = lambda s: client_mock + client_mock.__exit__ = MagicMock(return_value=False) + client_mock.post.return_value = self._make_resp(404, {}) + client_mock.get.return_value = detail_resp + + with patch("agent.model_metadata.detect_local_server_type", return_value="vllm"), \ + patch("httpx.Client", return_value=client_mock): + result = _query_local_context_length("some-model", "http://localhost:8000/v1") + + assert result == 1000000, ( + f"must resolve the context window, not the output cap; got {result}" + ) + + def test_list_branch_reads_context_window_not_output_cap(self): + """Same contract on the sibling ``/v1/models`` LIST branch. + + Both probe branches must share one definition of "context window"; + fixing only the detail branch would leave the identical bug reachable + whenever the per-model describe endpoint 404s. + """ + from agent.model_metadata import _query_local_context_length + + detail_miss = self._make_resp(404, {}) + list_resp = self._make_resp(200, {"data": [ + {"id": "some-model", "max_input_tokens": 1000000, "max_tokens": 128000}, + ]}) + + client_mock = MagicMock() + client_mock.__enter__ = lambda s: client_mock + client_mock.__exit__ = MagicMock(return_value=False) + client_mock.post.return_value = self._make_resp(404, {}) + # first GET is /v1/models/{model} (miss), second is /v1/models (list) + client_mock.get.side_effect = [detail_miss, list_resp] + + with patch("agent.model_metadata.detect_local_server_type", return_value="vllm"), \ + patch("httpx.Client", return_value=client_mock): + result = _query_local_context_length("some-model", "http://localhost:8000/v1") + + assert result == 1000000, ( + f"list branch must resolve the context window, not the output cap; got {result}" + ) + + def test_probe_agrees_with_the_module_key_vocabulary(self): + """Invariant: the probe's notion of a context window is the module's. + + ``_CONTEXT_LENGTH_KEYS`` / ``_MAX_COMPLETION_KEYS`` are the single + source of truth for this distinction. Asserting the relation (rather + than a frozen key list) keeps the guard correct as the vocabulary + grows, and fails if a probe branch ever re-hardcodes its own keys. + """ + from agent import model_metadata as mm + + assert "max_tokens" in mm._MAX_COMPLETION_KEYS + assert "max_tokens" not in mm._CONTEXT_LENGTH_KEYS + # No key may be classified as both a window and an output cap. + assert not (set(mm._CONTEXT_LENGTH_KEYS) & set(mm._MAX_COMPLETION_KEYS)) + + # Every context key the module recognises is honoured by the flat + # reader the probe branches use, and no completion key ever is. + for key in mm._CONTEXT_LENGTH_KEYS: + assert mm._extract_flat_context_length({key: 123456}) == 123456, key + for key in mm._MAX_COMPLETION_KEYS: + assert mm._extract_flat_context_length({key: 123456}) is None, key + class TestQueryLocalContextLengthModelsList: """_query_local_context_length: falls back to /v1/models list."""