fix(query): keep the truncation retry within the model's output and context limits

README.md

@@ -800,7 +800,7 @@ Sources: [models](https://platform.minimax.io/docs/guides/models-intro.md),
 
 - **`session`** — REPL session persistence: `auto_save` writes each session automatically; `max_history` caps retained turns.
 - **`settings`** — the advisor (reviewer) feature. **`advisor_enabled` is the master switch — `false` by default, so the advisor is OFF unless you opt in** (set it to `true`, or run `/advisor <provider>:<model>` which flips it on). `advisor_provider` / `advisor_model` pick the reviewer model; `advisor_client_mode` routes the call via the client.
-- **`settings.modelLimits`** — context/output limits for private or gateway models absent from the built-in catalog. Keys may be exact model IDs, model prefixes, or host-qualified IDs such as `localhost:4000:private-model`; exact matches beat prefixes and host-qualified keys beat bare keys within the same match type. Built-in catalog metadata remains authoritative.
+- **`settings.modelLimits`** — context/output limits for private or gateway models absent from the built-in catalog. Keys may be exact model IDs (tagged IDs such as `llama3.1:8b` included), model prefixes, or host-qualified IDs such as `localhost:4000:private-model`; exact matches beat prefixes and host-qualified keys beat bare keys within the same match type. `maxOutputTokens` is treated as a ceiling. When a reply is cut off, the retry with a larger output budget never asks for more than it, or for more than the configured `contextWindow` leaves after the input, and compaction requests stay under it too. Built-in catalog metadata remains authoritative.
 - **`env`** — secrets and environment values injected at startup (e.g. `TAVILY_API_KEY` for web search). Managed via `clawcodex config`; keys here are exported into the process environment without overriding anything you already set in your shell.
 
 #### Stream-stall tuning

src/models/configs.py

@@ -32,6 +32,20 @@ class ModelConfig:
     #: opted-out row still describes its own id fully; it just stops
     #: describing everyone else's.
     claims_prefix: bool = True
+    #: The model's documented hard maximum for output tokens, when the row
+    #: records one. ``max_output_tokens`` above is NOT that number on most
+    #: rows. On current Anthropic rows it is the first-attempt wire value, set
+    #: below the real maximum so the query loop's truncation escalation has
+    #: room. On OpenAI-compatible rows it is advisory (nothing sends it) and
+    #: feeds only the auto-compact reservation. ``get_max_output_tokens_ceiling``
+    #: reads this field, and the escalation retry and compaction requests never
+    #: ask for more. ``None`` means not recorded, and the escalation keeps its
+    #: historical 64K. Set it only on a row whose maximum is below that, and
+    #: only when the row's prefix claim reaches no model with a higher one.
+    #: ``claude-opus-4-20250514`` is the counter-example: its 32K is Opus 4's
+    #: real limit, but its ``claude-opus-4`` prefix also resolves 4.5 / 4.6
+    #: ids, which allow far more.
+    max_output_tokens_ceiling: int | None = None
 
 
 MODEL_CONFIGS: dict[str, ModelConfig] = {
@@ -199,6 +213,7 @@ class ModelConfig:
         display_name="Claude 3.5 Sonnet (Oct 2024)",
         context_window=200_000,
         max_output_tokens=8_192,
+        max_output_tokens_ceiling=8_192,
         supports_thinking=False,
         cost_input_per_mtok=3.0,
         cost_output_per_mtok=15.0,
@@ -208,6 +223,7 @@ class ModelConfig:
         display_name="Claude 3.5 Sonnet (Jun 2024)",
         context_window=200_000,
         max_output_tokens=8_192,
+        max_output_tokens_ceiling=8_192,
         supports_thinking=False,
         is_deprecated=True,
         deprecation_message="Use claude-sonnet-4-20250514 instead",
@@ -219,6 +235,7 @@ class ModelConfig:
         display_name="Claude 3.5 Haiku",
         context_window=200_000,
         max_output_tokens=8_192,
+        max_output_tokens_ceiling=8_192,
         supports_thinking=False,
         cost_input_per_mtok=1.0,
         cost_output_per_mtok=5.0,
@@ -232,6 +249,7 @@ class ModelConfig:
         display_name="Claude 3 Opus",
         context_window=200_000,
         max_output_tokens=4_096,
+        max_output_tokens_ceiling=4_096,
         supports_thinking=False,
         is_deprecated=True,
         deprecation_message="Use claude-opus-4-20250514 instead",
@@ -243,6 +261,7 @@ class ModelConfig:
         display_name="Claude 3 Sonnet",
         context_window=200_000,
         max_output_tokens=4_096,
+        max_output_tokens_ceiling=4_096,
         supports_thinking=False,
         is_deprecated=True,
         deprecation_message="Use claude-sonnet-4-20250514 instead",
@@ -254,6 +273,7 @@ class ModelConfig:
         display_name="Claude 3 Haiku",
         context_window=200_000,
         max_output_tokens=4_096,
+        max_output_tokens_ceiling=4_096,
         supports_thinking=False,
         cost_input_per_mtok=0.25,
         cost_output_per_mtok=1.25,
@@ -859,30 +879,35 @@ class ModelConfig:
         display_name="GPT-4o",
         context_window=128_000,
         max_output_tokens=16_384,
+        max_output_tokens_ceiling=16_384,
     ),
     "gpt-4o-mini": ModelConfig(
         model_id="gpt-4o-mini",
         display_name="GPT-4o Mini",
         context_window=128_000,
         max_output_tokens=16_384,
+        max_output_tokens_ceiling=16_384,
     ),
     "gpt-4-turbo": ModelConfig(
         model_id="gpt-4-turbo",
         display_name="GPT-4 Turbo",
         context_window=128_000,
         max_output_tokens=4_096,
+        max_output_tokens_ceiling=4_096,
     ),
     "gpt-4": ModelConfig(
         model_id="gpt-4",
         display_name="GPT-4",
         context_window=8_192,
         max_output_tokens=8_192,
+        max_output_tokens_ceiling=8_192,
     ),
     "gpt-3.5-turbo": ModelConfig(
         model_id="gpt-3.5-turbo",
         display_name="GPT-3.5 Turbo",
         context_window=16_385,
         max_output_tokens=4_096,
+        max_output_tokens_ceiling=4_096,
     ),
 }
 

src/models/context.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+from collections.abc import Mapping
 from urllib.parse import urlparse
 
 from .configs import get_model_config
@@ -53,7 +54,11 @@ def _settings_limit(
         limits = get_settings().model_limits
     except Exception:
         return None
-    if not limits or not isinstance(model_id, str):
+    # ``SettingsSchema.from_dict`` converts only a mapping and passes any other
+    # JSON shape through untouched, so a hand-edited ``"modelLimits": [...]``
+    # arrives here as a list. Treat it as "nothing configured" rather than let
+    # ``.items()`` raise into the request path.
+    if not limits or not isinstance(limits, Mapping) or not isinstance(model_id, str):
         return None
 
     model = strip_1m_context_suffix(model_id).lower()
@@ -75,8 +80,11 @@ def _settings_limit(
         if host and key.startswith(f"{host}:"):
             qualified = True
             candidate = key[len(host) + 1 :]
-        elif ":" in key:
-            continue
+        # Any other key is matched whole, colons included. Ollama-style ids
+        # carry one (``llama3.1:8b``), and a bare exact key for them is a
+        # documented shape. A key qualified for a different host never
+        # matches by accident, because the model id would have to begin with
+        # that ``host:`` prefix.
         if candidate == model:
             (exact_host if qualified else exact_bare).append(limit)
         elif model.startswith(candidate):
@@ -138,6 +146,51 @@ def get_model_max_output_tokens(
     return DEFAULT_MAX_OUTPUT_TOKENS
 
 
+def get_max_output_tokens_ceiling(
+    model_id: str, *, base_url: str | None = None,
+) -> int | None:
+    """The most output a request for ``model_id`` may ask for, if known.
+
+    Not the same number as :func:`get_model_max_output_tokens`. That one is
+    what an ordinary request asks for, and on current Anthropic rows it sits
+    deliberately below the model's real maximum (see configs.py) so the query
+    loop's truncation escalation has room above it. This one answers "how high
+    may a request go", so it reports only a number that is a real limit:
+
+    * a catalog row's ``max_output_tokens_ceiling``, or ``None`` when the row
+      records none. The catalog stays authoritative, as in the lookups above:
+      a ``modelLimits`` entry never overrides a catalogued model;
+    * for a model absent from the catalog, the ``modelLimits``
+      ``maxOutputTokens`` the user configured (exact / prefix /
+      host-qualified, via :func:`_settings_limit`);
+    * otherwise ``None``. Nothing is known, and callers keep their own value.
+    """
+    if not isinstance(model_id, str) or not model_id:
+        return None
+    config = get_model_config(strip_1m_context_suffix(model_id))
+    if config is not None:
+        return config.max_output_tokens_ceiling
+    return _settings_limit(model_id, "max_output_tokens", base_url=base_url)
+
+
+def get_env_max_output_tokens() -> int | None:
+    """``CLAUDE_CODE_MAX_OUTPUT_TOKENS`` as a positive int, else ``None``.
+
+    Invalid / non-positive values are ignored with a debug log.
+    """
+    raw = os.environ.get("CLAUDE_CODE_MAX_OUTPUT_TOKENS")
+    if not raw:
+        return None
+    try:
+        value = int(raw.strip())
+    except ValueError:
+        value = 0
+    if value > 0:
+        return value
+    logger.debug("ignoring invalid CLAUDE_CODE_MAX_OUTPUT_TOKENS=%r", raw)
+    return None
+
+
 def resolve_max_output_tokens(
     override: int | None,
     model_id: str | None,
@@ -147,12 +200,14 @@ def resolve_max_output_tokens(
     """Resolve the request-path ``max_tokens`` (ch04 round-3 G0).
 
     Precedence mirrors TS ``claude.ts:1602-1605``:
-    1. explicit override (the query loop's 64K escalation passes through
-       here unchanged);
+    1. explicit override. The query loop's truncation escalation passes
+       through here unchanged; the loop has already bounded it by the
+       model's output ceiling and the context window
+       (``query._escalated_max_output_tokens``);
     2. ``CLAUDE_CODE_MAX_OUTPUT_TOKENS`` env — the key has been on the
        trusted-env allowlist since round 1 (``trust_boundary.py``);
        consuming it closes that dangling promise. Invalid / non-positive
-       values are ignored with a debug log;
+       values are ignored with a debug log (:func:`get_env_max_output_tokens`);
     3. the per-model table via :func:`get_model_max_output_tokens`
        (→ ``DEFAULT_MAX_OUTPUT_TOKENS`` 8_192 for unknown models).
 
@@ -166,17 +221,9 @@ def resolve_max_output_tokens(
     """
     if override is not None:
         return override
-    raw = os.environ.get("CLAUDE_CODE_MAX_OUTPUT_TOKENS")
-    if raw:
-        try:
-            value = int(raw.strip())
-        except ValueError:
-            value = 0
-        if value > 0:
-            return value
-        logger.debug(
-            "ignoring invalid CLAUDE_CODE_MAX_OUTPUT_TOKENS=%r", raw
-        )
+    env_value = get_env_max_output_tokens()
+    if env_value is not None:
+        return env_value
     if model_id:
         return get_model_max_output_tokens(model_id, base_url=base_url)
     return DEFAULT_MAX_OUTPUT_TOKENS

src/query/query.py

@@ -363,6 +363,94 @@ def _get_context_window(provider: BaseProvider) -> int:
     return 200_000
 
 
+#: Headroom kept between the escalated request and the context window. The
+#: input count comes from the truncated reply's own usage, and the escalation
+#: re-sends that same request, so this only absorbs accounting slack. TS keeps
+#: the same 1,000-token buffer when it re-sizes ``max_tokens`` after a
+#: context-overflow 400 (withRetry.ts).
+ESCALATION_CONTEXT_MARGIN_TOKENS = 1_000
+
+
+def _usage_count(usage: Any, key: str) -> int:
+    """A non-negative int from a usage dict; 0 for anything else."""
+    value = usage.get(key) if isinstance(usage, dict) else None
+    if isinstance(value, bool) or not isinstance(value, int):
+        return 0
+    return max(value, 0)
+
+
+def _escalated_max_output_tokens(
+    provider: BaseProvider, truncated: Message | None,
+) -> int | None:
+    """``max_tokens`` for the truncation-escalation retry, or ``None`` to skip it.
+
+    The retry re-sends the request that was just cut off (#958), so it is
+    worth a request only if it allows more output than that request had.
+    Start from ``ESCALATED_MAX_TOKENS`` and lower it to:
+
+    * the model's known output ceiling (:func:`get_max_output_tokens_ceiling`:
+      a catalog row's documented maximum, or ``modelLimits`` for an
+      uncatalogued model);
+    * ``CLAUDE_CODE_MAX_OUTPUT_TOKENS``, the user's own cap (TS skips the
+      escalation outright while it is set, query.ts);
+    * the context window less the truncated request's input tokens, so the
+      retry cannot overflow the window that the shorter request fitted in.
+
+    If what is left is no more than the truncated attempt already had (the
+    ``max_tokens`` sent on the Anthropic wire, and the output it actually
+    produced on either wire), the retry would be cut off at the same point.
+    Return ``None``, and the caller goes straight to the recovery lane, which
+    keeps the partial reply. With no limit known, the result is the
+    historical ``ESCALATED_MAX_TOKENS``.
+    """
+    try:
+        from ..models.context import (
+            get_env_max_output_tokens,
+            get_max_output_tokens_ceiling,
+            resolve_max_output_tokens,
+        )
+        from ..providers import is_anthropic_wire
+
+        model = getattr(provider, "model", None)
+        model = model if isinstance(model, str) else ""
+        base_url = getattr(provider, "base_url", None)
+        base_url = base_url if isinstance(base_url, str) else None
+
+        target = ESCALATED_MAX_TOKENS
+        for limit in (
+            get_max_output_tokens_ceiling(model, base_url=base_url),
+            get_env_max_output_tokens(),
+        ):
+            if limit is not None:
+                target = min(target, limit)
+
+        usage = getattr(truncated, "usage", None)
+        input_tokens = (
+            _usage_count(usage, "input_tokens")
+            + _usage_count(usage, "cache_creation_input_tokens")
+            + _usage_count(usage, "cache_read_input_tokens")
+        )
+        if input_tokens:
+            room = (
+                _get_context_window(provider)
+                - input_tokens
+                - ESCALATION_CONTEXT_MARGIN_TOKENS
+            )
+            target = min(target, room)
+
+        already_allowed = _usage_count(usage, "output_tokens")
+        if is_anthropic_wire(provider):
+            # The truncated request carried this exact value (_call_model_sync).
+            already_allowed = max(
+                already_allowed,
+                resolve_max_output_tokens(None, model or None, base_url=base_url),
+            )
+    except Exception:
+        logger.debug("escalation bound resolution failed", exc_info=True)
+        return ESCALATED_MAX_TOKENS
+    return target if target > already_allowed else None
+
+
 def _is_withheld_max_output_tokens(msg: Message | None) -> bool:
     if msg is None:
         return False
@@ -2134,17 +2222,25 @@ def _marking_chunk_cb(text: str) -> None:
             last_message = assistant_messages[-1] if assistant_messages else None
 
             if _is_withheld_max_output_tokens(last_message):
-                if (
-                    max_output_tokens_override is None
+                # One escalation per turn, and only when it can actually
+                # raise the ceiling. A retry that cannot (an output cap the
+                # first attempt already used, or a full context window) falls
+                # through to the recovery lane below, which keeps the partial
+                # reply instead of re-generating it (#958).
+                escalated_max_tokens = (
+                    _escalated_max_output_tokens(params.provider, last_message)
+                    if max_output_tokens_override is None
                     and max_output_tokens_recovery_count == 0
-                ):
+                    else None
+                )
+                if escalated_max_tokens is not None:
                     state = QueryState(
                         messages=messages,
                         tool_use_context=tool_use_context,
                         auto_compact_tracking=state.auto_compact_tracking,
                         max_output_tokens_recovery_count=max_output_tokens_recovery_count,
                         has_attempted_reactive_compact=has_attempted_reactive_compact,
-                        max_output_tokens_override=ESCALATED_MAX_TOKENS,
+                        max_output_tokens_override=escalated_max_tokens,
                         stop_hook_active=None,
                         turn_count=turn_count,
                         transition=Transition(reason="max_output_tokens_escalate"),