ocr: resumption must respect the context window, and never extend a loop (v0.3.50)

Two defects that a throughput collapse on a live install made visible. With
surya_max_tokens raised to 12288, 262 resumption segments spent 1.6M generated
tokens in 80 minutes, 122 of them running the full cap dry, while 253 of those
262 ended in "came back repeating itself".

Resumption re-sends the text so far as INPUT, so the model's context window --
not the continuation budget -- is the real ceiling. The budget did not know
that: at a 12288-token cap, the third segment of a page asked for 12288 more on
top of ~28,500 already in a 32,768-token window and was rejected outright. The
page was then reported as one that could not be resumed AT ALL, discarding text
already paid for.

  - surya_context_tokens, or 0 to read the server's own max_model_len from
    /v1/models. Configuration wins: an operator may serve a shorter window than
    the checkpoint advertises.
  - Each segment asks for the smaller of the cap and the room left, measured
    from what the server said the prompt cost (usage.prompt_tokens), not
    guessed; the guess covers the first resumption only.
  - Out of room stops, KEEPS what it has, and says so -- naming the fix that
    helps, which is a SMALLER max_tokens, not a bigger budget.
  - A refusal for length is not retried: the fallback resumption adds a user
    turn, so it is longer than the request just refused.
  - Not knowing the window changes nothing and guesses nothing.

The second defect: a reply that trips surya's repeat detector hit the cap
because it is LOOPING, not because the page is long. Resuming it can only
extend the loop, and surya re-runs the whole page anyway -- so one page on that
path cost four full generations. It is no longer resumed, and it is not marked
truncated either: that is surya's retry to make, and "truncated" would send the
caller looking for a bigger budget it cannot use.

A fixed token cap is also an accidental repetition circuit-breaker, which is
why raising it slowed the LOOPING pages by the ratio of the caps rather than
every page. Documented in docs/ocr.md with the window arithmetic (2048 -> ~13
segments, 12288 -> 2).

Tests: 8 new (window clamp, stop-one-segment-early, length refusal not
retried, unknown window, configuration over server, a cap below the 256-token
floor is never mistaken for a full window, a looping reply is not resumed, a
genuinely long page still is). 62 pass in the surya/blank-page suites; the
floor-versus-cap bug was caught by the existing 44.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PxYHWbCAjjjD7A1rhFKtaW

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.3.49"
+__version__ = "0.3.50"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/admin/routes.py

@@ -4142,6 +4142,7 @@ def build_settings_dict(c, gpu_cards):
             "surya_max_retries": getattr(c.ocr, "surya_max_retries", 1),
             "surya_max_tokens": getattr(c.ocr, "surya_max_tokens", 0),
             "surya_max_continuations": getattr(c.ocr, "surya_max_continuations", 2),
+            "surya_context_tokens": getattr(c.ocr, "surya_context_tokens", 0),
             "surya_langs": c.ocr.surya_langs,
             "surya_token_trace": c.ocr.surya_token_trace,
             "surya_repair_accents": c.ocr.surya_repair_accents,
@@ -4971,6 +4972,10 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
             o.surya_max_continuations = _i(
                 "surya_max_continuations",
                 getattr(o, "surya_max_continuations", 2), 0)
+        if "surya_context_tokens" in d:
+            o.surya_context_tokens = _i(
+                "surya_context_tokens",
+                getattr(o, "surya_context_tokens", 0), 0)
         if "surya_max_tokens" in d:
             o.surya_max_tokens = _i("surya_max_tokens",
                                     getattr(o, "surya_max_tokens", 0), 0)

codai/admin/templates/settings.html

@@ -1208,6 +1208,7 @@
       <div class="form-group" style="flex:1;min-width:130px"><label class="form-label" title="How many times Surya may RE-run a page whose generation repeated itself. Surya's own number is 3 and each retry is another FULL generation at a higher temperature: a batch measured on 2026-10-13 spent 2,644 generated tokens per page on text worth ~400, with a quarter of all requests re-sending an image the pod had already read. 1 keeps the rescue a retry does buy and drops the long tail; 0 disables it; -1 means whatever Surya decides. A retried page is reported as model_attempts in the result.">Retries / page</label><input type="number" id="s-ocr-surya-max-retries" class="form-input" min="-1" placeholder="1"></div>
       <div class="form-group" style="flex:1;min-width:140px"><label class="form-label" title="Cap on the tokens Surya may generate for ONE page. 0 uses Surya's own 2048, which its vLLM backend never overrides. Surya-ocr-2 emits block markup (a category and a bbox per block) BEFORE the text, so a reply stopped at the cap is incomplete markup and parses to a few blocks or to none — the page then arrives as a short fragment or as no text at all, which looks exactly like a page the model could not read. A dense page can legitimately need several thousand tokens. model_attempts.truncated in a page's result tells you when this is what went wrong.">Max tokens / page</label><input type="number" id="s-ocr-surya-max-tokens" class="form-input" min="0" placeholder="0 = 2048"></div>
       <div class="form-group" style="flex:1;min-width:140px"><label class="form-label" title="How many times a page whose reply hit the token cap may be RESUMED from where it stopped. Raising Max tokens only moves the cliff — a dense enough page overflows any fixed number, and incomplete block markup parses to a fragment or to nothing. Resumption sends the text so far back as a prefilled assistant turn so the model carries on the exact string. 2 by default, because returning a page known to be incomplete is a correctness failure, not a speed trade-off; 0 restores the old behaviour. A page still unfinished afterwards is marked model_attempts.complete = false so a caller can refuse it.">Resume cut-off pages</label><input type="number" id="s-ocr-surya-max-continuations" class="form-input" min="0" placeholder="2"></div>
+      <div class="form-group" style="flex:1;min-width:150px"><label class="form-label" title="The serving model's context window, in tokens. 0 asks the server for its own max_model_len (vLLM reports it on /v1/models) — leave it at 0 unless the server will not say, or is served with a shorter --max-model-len than the checkpoint advertises. This is the ceiling resumption really runs into: every resumption re-sends the text so far as INPUT, so once the window is nearly full, asking for another whole Max tokens is not a bigger segment, it is an HTTP 400 — and a page that was nearly whole gets reported as a page that could not be resumed at all. Knowing the number lets such a page stop one segment early, keep what it already has and say so. It also means a SMALLER Max tokens fits more text in the same window, because the last segment wastes less of it.">Context window</label><input type="number" id="s-ocr-surya-context-tokens" class="form-input" min="0" placeholder="0 = ask the server"></div>
       <label style="display:flex;align-items:center;gap:.4rem;cursor:pointer" title="Local mode only. On, Surya runs on this machine's GPU or refuses to load — it will not quietly fall back to the CPU, where a page takes minutes. Off, the CPU is used deliberately."><input type="checkbox" id="s-ocr-surya-use-gpu" checked><span>Require GPU (Local)</span></label>
       <div class="form-group" style="flex:1;min-width:120px"><label class="form-label">Langs</label><input type="text" id="s-ocr-surya-langs" class="form-input" placeholder="it"></div>
       <label style="display:flex;align-items:center;gap:.4rem;cursor:pointer" title="Return the model's own per-token bytes wherever a page came back unreadable. A diagnostic — leave off unless you are chasing lost characters."><input type="checkbox" id="s-ocr-surya-token-trace"><span>Trace unreadable tokens</span></label>
@@ -2210,6 +2211,7 @@ async function loadSettings(){
     document.getElementById('s-ocr-surya-max-retries').value = ocr.surya_max_retries ?? 1;
     document.getElementById('s-ocr-surya-max-tokens').value = ocr.surya_max_tokens ?? 0;
     document.getElementById('s-ocr-surya-max-continuations').value = ocr.surya_max_continuations ?? 2;
+    document.getElementById('s-ocr-surya-context-tokens').value = ocr.surya_context_tokens ?? 0;
     document.getElementById('s-ocr-surya-use-gpu').checked = ocr.surya_use_gpu !== false;
     document.getElementById('s-ocr-surya-langs').value = ocr.surya_langs ?? 'it';
     document.getElementById('s-ocr-surya-token-trace').checked = !!ocr.surya_token_trace;
@@ -2503,6 +2505,7 @@ async function saveSettings(silent){
       // 0 is meaningful (resumption off), so an explicit NaN check.
       surya_max_continuations: (function(v){ v = parseInt(v); return isNaN(v) ? 2 : v; })(
         document.getElementById('s-ocr-surya-max-continuations').value),
+      surya_context_tokens: parseInt(document.getElementById('s-ocr-surya-context-tokens').value) || 0,
       surya_langs: document.getElementById('s-ocr-surya-langs').value.trim() || 'it',
       surya_token_trace: document.getElementById('s-ocr-surya-token-trace').checked,
       surya_repair_accents: document.getElementById('s-ocr-surya-repair-accents').checked,

codai/config.py

@@ -1090,6 +1090,19 @@ class OcrConfig:
     #: stops making progress stops early regardless, so a generous budget costs
     #: nothing on a page that cannot use it. 0 disables resumption.
     surya_max_continuations: int = 8
+    #: The serving model's context window, or 0 to ask the server for its own
+    #: `max_model_len` (vLLM reports it on `/v1/models`).
+    #:
+    #: This is the ceiling resumption actually runs into, and the one it used to
+    #: hit the hard way: each round re-sends the text so far as INPUT, so asking
+    #: for another full `surya_max_tokens` once the window is nearly full is not
+    #: a bigger segment, it is an HTTP 400 -- and a page that was nearly whole
+    #: gets reported as a page that could not be resumed at all. Knowing the
+    #: number lets a page stop one segment early, keep what it has, and say so.
+    #:
+    #: Set it when the server does not report a window, or when it is served
+    #: with a shorter `--max-model-len` than the checkpoint advertises.
+    surya_context_tokens: int = 0
     #: A page whose text is this short gets its ink measured, so a short result
     #: can be explained rather than guessed at. 0 disables the measurement.
     blank_page_max_chars: int = 200

codai/ocr/surya.py

@@ -80,7 +80,10 @@ class SuryaEngine(SubprocessOcrEngine):
                 "max_tokens": max(0, int(getattr(self.cfg, "surya_max_tokens", 0) or 0)),
                 # Resumptions allowed for a page whose reply hit the cap.
                 "max_continuations": max(0, int(getattr(
-                    self.cfg, "surya_max_continuations", 2) or 0))}
+                    self.cfg, "surya_max_continuations", 2) or 0)),
+                # 0 = ask the server for its own max_model_len.
+                "context_tokens": max(0, int(getattr(
+                    self.cfg, "surya_context_tokens", 0) or 0))}
 
     #: Modes where the recognition VLM lives in a model server we only talk to.
     SERVED_MODES = ("model", "vllm", "llamacpp", "server")

codai/ocr/workers/ocr_worker.py

@@ -66,6 +66,96 @@ _CONTINUE_BUDGET = [0]
 #: (`chat_completions_batch(max_tokens_default=2048)`).
 _SURYA_DEFAULT_MAX_TOKENS = 2048
 
+#: The serving model's context window, or 0 to ask the server for it once.
+#: Resumption re-sends the text so far as INPUT, so the window -- not the
+#: continuation budget -- is the real ceiling on how long a page may be. A
+#: budget that ignores it does not fail gracefully: the server rejects the
+#: request outright (HTTP 400) and a page that was nearly whole is reported as
+#: a page that could not be resumed at all.
+_CONTEXT_TOKENS = [0]
+
+#: Probed `max_model_len` per (server, model). A served model cannot change
+#: shape under us, so one question per server is enough.
+_CONTEXT_CACHE = {}
+
+#: What to assume the image and prompt cost when the server will not say and
+#: nothing is configured. surya-2 at the 6.29 Mpix cap is in this region; the
+#: number only has to be an over-estimate, because its whole job is to keep the
+#: LAST segment inside the window.
+_PROMPT_TOKENS_GUESS = 4096
+
+#: Below this there is no point asking for more: a segment this short cannot
+#: close an open block, so it would spend a round to add nothing. It is a floor
+#: on the room LEFT, never on the cap -- a deployment whose cap is smaller than
+#: this is simply never clamped, and must not be mistaken for a full window.
+_MIN_SEGMENT_TOKENS = 256
+
+
+def _context_limit(client, model_name):
+    """The serving model's context window, or 0 if it cannot be established.
+
+    Configuration wins, because an operator may be serving with a shorter
+    `--max-model-len` than the checkpoint advertises. Otherwise vLLM reports
+    `max_model_len` on `/v1/models`, which is the number that will actually
+    reject us.
+    """
+    fixed = int(_page_override("context_tokens", _CONTEXT_TOKENS[0]) or 0)
+    if fixed > 0:
+        return fixed
+    key = (str(getattr(client, "base_url", "")), str(model_name))
+    if key in _CONTEXT_CACHE:
+        return _CONTEXT_CACHE[key]
+    limit = 0
+    try:
+        for m in (client.models.list().data or []):
+            if getattr(m, "id", None) != model_name:
+                continue
+            for attr in ("max_model_len", "max_context_length", "context_length"):
+                v = getattr(m, attr, None)
+                if v is None and isinstance(getattr(m, "__dict__", None), dict):
+                    v = m.__dict__.get(attr)
+                if v:
+                    limit = int(v)
+                    break
+            break
+    except Exception as exc:
+        # Not knowing is survivable -- `_segment_cap` simply stops clamping.
+        # Being wrong is not, so nothing is guessed here.
+        print(f"[surya] could not read the model's context window ({exc}); "
+              f"resumption will not be able to stop before the server does",
+              file=sys.stderr, flush=True)
+    if limit:
+        print(f"[surya] {model_name} serves a {limit}-token context window",
+              file=sys.stderr, flush=True)
+    _CONTEXT_CACHE[key] = limit
+    return limit
+
+
+def _segment_cap(cap, limit, prompt_tokens):
+    """How many tokens the next segment may ask for.
+
+    `cap` is what the page is allowed per reply; the window is what the server
+    will accept. Asking for more than the room left is not a bigger segment, it
+    is an error instead of a segment.
+    """
+    if not limit:
+        return int(cap)
+    room = int(limit) - int(prompt_tokens or 0)
+    return max(0, min(int(cap), room))
+
+
+def _is_context_error(exc):
+    """A rejection for length, as opposed to any other 400.
+
+    Worth separating because the fallback resumption (a plain "continue" turn)
+    is LONGER than the one that was just refused, so retrying a length refusal
+    only spends another round trip to be told the same thing.
+    """
+    t = f"{exc}".lower()
+    return (("context" in t and ("length" in t or "window" in t))
+            or "maximum context" in t or "longer than the maximum" in t
+            or "reduce the length" in t or "max_model_len" in t)
+
 
 def _trace_begin():
     """Start a page's trace. Returns its (calls, repairs) — this thread's own."""
@@ -81,7 +171,8 @@ def _trace_begin():
 def _trace_opts(opts) -> None:
     """Per-request overrides for the page this thread is about to read."""
     _TRACE.opts = {k: int(v) for k, v in (opts or {}).items()
-                   if k in ("max_tokens", "max_continuations")
+                   if k in ("max_tokens", "max_continuations",
+                            "context_tokens")
                    and str(v).strip() not in ("", "None")}
 
 
@@ -802,13 +893,47 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
         print(f"[surya] cannot resume a cut-off page: {exc}", file=sys.stderr, flush=True)
         return result, 0, True
 
+    # A reply that is already repeating itself did not run out of room -- it ran
+    # out of page and kept going. Resuming it can only extend the loop, and
+    # surya is about to re-run the page anyway (that is what its own repeat
+    # detector does), so every token spent resuming here is spent twice over on
+    # a page that will be thrown away. This is the single largest waste the
+    # cap-plus-resumption combination can produce.
+    repeats = getattr(oc, "detect_repeat_token", None)
+    try:
+        if repeats is not None and repeats(segments[0]):
+            print(f"[surya] not resuming a reply that is already repeating "
+                  f"itself — it stopped at the {cap}-token cap because it is "
+                  f"looping, not because the page is long",
+                  file=sys.stderr, flush=True)
+            return result, 0, False
+    except Exception:
+        pass
+
+    limit = _context_limit(client, model_name)
+    # Until the server tells us what the prompt really costs, assume the image
+    # and the instructions, plus everything generated so far (it all comes back
+    # as input on the next round).
+    prompt_tokens = _PROMPT_TOKENS_GUESS + total
     attempts = 0
     while attempts < budget and _at_cap(last, cap):
         attempts += 1
+        seg_cap = _segment_cap(cap, limit, prompt_tokens)
+        if limit and seg_cap < min(int(cap), _MIN_SEGMENT_TOKENS):
+            # Honest and specific: this page is not short of budget, it is out
+            # of context. Raising `surya_max_continuations` cannot help it;
+            # LOWERING `surya_max_tokens` can, because a smaller segment leaves
+            # the window to be spent on text rather than on one huge reply.
+            print(f"[surya] page filled the model's {limit}-token context "
+                  f"window after {rounds} resumption(s) — returning it partial "
+                  f"and marked (a smaller ocr.surya_max_tokens fits more text "
+                  f"in the same window)", file=sys.stderr, flush=True)
+            break
         prefix = "".join(segments)
         kwargs = dict(model=model_name, messages=base_messages + [
             {"role": "assistant", "content": prefix}],
-            max_tokens=cap, temperature=temperature, top_p=top_p, timeout=timeout,
+            max_tokens=seg_cap, temperature=temperature, top_p=top_p,
+            timeout=timeout,
             extra_body={"continue_final_message": True,
                         "add_generation_prompt": False})
         if want_logprobs:
@@ -816,6 +941,11 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
         try:
             completion = client.chat.completions.create(**kwargs)
         except Exception as exc:
+            if _is_context_error(exc):
+                print(f"[surya] page filled the model's context window after "
+                      f"{rounds} resumption(s) ({exc}) — returning it partial "
+                      f"and marked", file=sys.stderr, flush=True)
+                break
             # The flag is a vLLM extension; llama.cpp and others reject it.
             try:
                 kwargs.pop("extra_body", None)
@@ -831,8 +961,16 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
                 break
         try:
             more = completion.choices[0].message.content or ""
-            last = (completion.usage.completion_tokens
-                    if getattr(completion, "usage", None) else 0)
+            usage = getattr(completion, "usage", None)
+            last = usage.completion_tokens if usage else 0
+            # The guess above is only needed for the first round. From here the
+            # server has told us exactly what the prefix and image cost, so the
+            # next round's input is that plus the segment just produced.
+            used = int(getattr(usage, "prompt_tokens", 0) or 0) if usage else 0
+            if used:
+                prompt_tokens = used + int(last or 0)
+            else:
+                prompt_tokens += int(last or 0)
         except Exception:
             break
         if not more:
@@ -846,7 +984,6 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
             print(f"[surya] resumption repeated the previous segment verbatim — "
                   f"stopping", file=sys.stderr, flush=True)
             break
-        repeats = getattr(oc, "detect_repeat_token", None)
         try:
             if repeats is not None and repeats(more):
                 print(f"[surya] resumption came back repeating itself — stopping "
@@ -875,7 +1012,9 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
                 # the repair is given up on rather than fed a mismatch.
                 logprobs = []
         print(f"[surya] page resumed after the {cap}-token cap "
-              f"(segment {rounds + 1}, {last} more tokens)",
+              f"(segment {rounds + 1}, {last} more tokens, "
+              f"{max(0, limit - prompt_tokens) if limit else '?'} of the "
+              f"context window left)",
               file=sys.stderr, flush=True)
 
     if rounds:
@@ -912,6 +1051,10 @@ class SuryaWorker:
         _MAX_TOKENS[0] = max(0, int(opts.get("max_tokens", 0) or 0))
         #: Resumptions allowed for one page. 0 = return a cut-off page partial.
         _CONTINUE_BUDGET[0] = max(0, int(opts.get("max_continuations", 2) or 0))
+        #: 0 asks the server for its own `max_model_len`. Set it when the server
+        #: does not report one, or when it serves a shorter window than the
+        #: checkpoint advertises.
+        _CONTEXT_TOKENS[0] = max(0, int(opts.get("context_tokens", 0) or 0))
         #: The page's token logprobs and accent repairs live in _TRACE, per
         #: thread — see _trace_begin(). They were per-worker-object until lanes
         #: put several pages in one process.

docs/ocr.md

@@ -1112,6 +1112,69 @@ drops unknown form fields **silently**. Before these parameters existed, a
 caller could send `max_tokens=8000` on every request for a week and never be
 told it did nothing.
 
+### The context window is the ceiling, not the budget
+
+Resumption re-sends the text so far as **input**. The continuation budget is
+therefore not what limits how long a page may be — the model's context window
+is, and the budget only decides how much of that window gets spent before
+coderai gives up.
+
+That arithmetic is the whole reason `surya_max_tokens` should be *modest* rather
+than as large as it will go. The two settings multiply into the same total
+capacity, but a large cap spends the window in coarse lumps:
+
+| `surya_max_tokens` | segments that fit a 32,768-token window (image ≈ 4,000) |
+|---|---|
+| 2,048 | ~13 |
+| 4,096 | ~7 |
+| 12,288 | **2** |
+
+A 12,288-token cap buys no more text than a 2,048-token cap — it buys the same
+text in two steps instead of thirteen, and it makes every wasted step six times
+more expensive. Measured on one install on 2026-10-13, with the cap at 12,288:
+262 resumption segments spent 1.6 million generated tokens, 122 of them running
+the full cap dry, and the third segment of any page was **rejected outright** —
+12,288 asked on top of ~28,500 already in the window.
+
+So coderai now establishes the window and respects it:
+
+- `surya_context_tokens`, or `0` to read the server's own `max_model_len` from
+  `/v1/models` (vLLM reports it). Configuration wins, because an operator may
+  serve with a shorter `--max-model-len` than the checkpoint advertises.
+- Each segment asks for **the smaller of the cap and the room left**, and the
+  room left is measured from what the server said the prompt actually cost
+  (`usage.prompt_tokens`), not guessed — the guess is used for the first
+  resumption only.
+- A page with no usable room left stops, **keeps what it has**, and is marked
+  incomplete. Previously it was refused by the server and reported as a page
+  that could not be resumed at all, discarding text already paid for.
+- A refusal *for length* is not retried. The fallback resumption adds a user
+  turn, so it is longer than the request just refused; retrying it only spends
+  another round trip to be told the same thing.
+- Not knowing the window is survivable, and nothing is guessed: with no number,
+  resumption behaves exactly as it did before and the budget rules.
+
+### A reply that repeats itself did not run out of room
+
+A generation that trips surya's repeat detector hit the cap because it is
+**looping**, not because the page is long. Resuming it can only extend the loop
+— and surya is about to re-run the whole page anyway, which is what its repeat
+detector does, so every token spent resuming is spent on text that is then
+thrown away.
+
+This was the single largest waste the cap-and-resumption combination could
+produce. On the same install, **253 of 262** resumptions ended in "came back
+repeating itself", and 187 of those pages were then re-run in full: one page on
+the bad path cost four full generations at 12,288 tokens each. Such a reply is
+now not resumed at all, and it is **not** marked truncated either — it is a
+degenerate reply, which is surya's retry to make. Calling it truncated would
+send the caller looking for a bigger budget it cannot use.
+
+Worth keeping in mind when reading a throughput regression: a fixed token cap is
+also an accidental repetition circuit-breaker. Raising it does not slow a page
+down by the ratio of the caps; it slows down the *looping* pages by that ratio,
+and they are the ones that were already costing the most.
+
 ---
 
 ## An empty page and a 71-character page are the same bug

tests/test_surya_retry_budget.py

@@ -415,16 +415,23 @@ class _Item:
 
 
 class _Completion:
-    def __init__(self, text, tokens, logprobs=None):
+    def __init__(self, text, tokens, logprobs=None, prompt_tokens=0):
         lp = type("_LP", (), {"content": logprobs or []})()
         choice = type("_C", (), {"message": type("_M", (), {"content": text})(),
                                  "logprobs": lp, "finish_reason": "length"})()
         self.choices = [choice]
-        self.usage = type("_U", (), {"completion_tokens": tokens})()
+        self.usage = type("_U", (), {"completion_tokens": tokens,
+                                     "prompt_tokens": prompt_tokens})()
 
 
-def _client(segments, reject_flag=False, reject_all=False):
-    """A chat client that hands back `segments` in order."""
+def _client(segments, reject_flag=False, reject_all=False, context=None,
+            prompt_tokens=0, reject_length=False):
+    """A chat client that hands back `segments` in order.
+
+    `context` makes it a server that reports a `max_model_len`, the way vLLM
+    does; `prompt_tokens` makes it report what the prefix cost, which is how
+    resumption learns how much of the window is left.
+    """
     sent = []
 
     class _Chat:
@@ -434,15 +441,29 @@ def _client(segments, reject_flag=False, reject_all=False):
                 sent.append(kwargs)
                 if reject_all:
                     raise RuntimeError("no")
+                if reject_length:
+                    raise RuntimeError(
+                        "This model's maximum context length is 32768 tokens. "
+                        "However, you requested 40000 tokens. Please reduce "
+                        "the length of the messages.")
                 if reject_flag and "extra_body" in kwargs:
                     raise RuntimeError("unknown field continue_final_message")
                 if not segments:
                     raise AssertionError("asked for more than it was given")
                 text, tokens = segments.pop(0)
-                return _Completion(text, tokens)
+                return _Completion(text, tokens, prompt_tokens=prompt_tokens)
+
+    class _Models:
+        @staticmethod
+        def list():
+            m = type("_M", (), {"id": "m", "max_model_len": context})()
+            return type("_L", (), {"data": [m]})()
 
     c = type("_Client", (), {"chat": _Chat()})()
     c.sent = sent
+    c.base_url = f"http://stub/{id(sent)}"
+    if context is not None:
+        c.models = _Models()
     return c
 
 
@@ -684,3 +705,190 @@ def test_the_default_budget_covers_a_dense_page():
     with its own budget. What this has to cover is the densest SINGLE page."""
     from codai.config import OcrConfig
     assert OcrConfig().surya_max_continuations >= 8
+
+
+# --- the context window is the real ceiling, not the continuation budget -----
+#
+# Each resumption re-sends the text so far as INPUT. A budget that ignores that
+# does not fail gracefully: the server refuses the request for length, and a
+# page that was nearly whole is reported as a page that could not be resumed at
+# all. Measured on Aruba 2026-10-13: at a 12288-token cap, segment 3 asked for
+# 12288 more on top of ~28.5k of prefix and image against a 32768-token window,
+# and eight pages came back as "could not resume" that were really "out of
+# room".
+
+
+def _resume_ctx(mod, monkeypatch, first, segments, budget=8, cap=12288, **kw):
+    oc = _fake_surya(monkeypatch, mod)
+    oc._build_messages = lambda img, prompt: [{"role": "user", "content": prompt}]
+    oc.scale_to_fit = lambda img: img
+    oc.PROMPT_MAPPING = {"ocr": "ocr"}
+    res = _Result(raw=first[0])
+    res.token_count = first[1]
+    cl = _client(list(segments), **kw)
+    out, rounds, still = mod._continue_reply(
+        oc, res, _Item(), cl, "m", cap, 0.0, 0.1, 60.0, False, budget)
+    return out, rounds, still, cl
+
+
+def test_a_segment_never_asks_for_more_room_than_the_window_has_left():
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 12288), [("b", 12288), ("c", 12288), ("d", 12288)],
+            context=32768, prompt_tokens=17000)
+    finally:
+        m.undo()
+    # Whatever it asked for, every request had to fit: the window minus what the
+    # server said the prompt cost.
+    assert cl.sent, "it never resumed at all"
+    for i, k in enumerate(cl.sent):
+        assert k["max_tokens"] <= 32768, f"request {i} asked past the window"
+    # The second request onwards knows the real prompt cost (17000), so it may
+    # ask for at most 32768 - 17000 - the segment just added.
+    assert cl.sent[1]["max_tokens"] <= 32768 - 17000
+    assert still is True, "out of room is still an incomplete page"
+
+
+def test_it_stops_one_segment_early_instead_of_being_refused():
+    """The page keeps what it has, and says why — a refusal would have thrown
+    the whole resumption away."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 12288), [("b", 12288), ("c", 12288), ("d", 12288)],
+            context=32768, prompt_tokens=32700)
+    finally:
+        m.undo()
+    assert out.raw == "ab", "it lost the text it had already paid for"
+    assert rounds == 1
+    assert len(cl.sent) == 1, "it asked again with no room left"
+    assert still is True
+
+
+def test_a_refusal_for_length_is_not_retried_with_a_longer_request():
+    """The fallback resumption ADDS a user turn, so it is longer than the one
+    just refused. Retrying a length refusal only buys the same answer twice."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 12288), [("b", 12288)], reject_length=True,
+            context=None)
+    finally:
+        m.undo()
+    assert len(cl.sent) == 1, "it tried a longer request after a length refusal"
+    assert rounds == 0 and still is True
+
+
+def test_a_window_the_server_will_not_name_does_not_stop_resumption():
+    """Not knowing the window is survivable; guessing it is not. Without a
+    number, resumption behaves exactly as it did before — the budget rules."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 12288), [("b", 12288), ("c", 5)], context=None)
+    finally:
+        m.undo()
+    assert out.raw == "abc" and rounds == 2 and still is False
+    assert all(k["max_tokens"] == 12288 for k in cl.sent)
+
+
+def test_a_configured_window_overrides_what_the_server_reports():
+    """An operator serving with a shorter --max-model-len than the checkpoint
+    advertises is the case the server cannot tell us about."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        mod._CONTEXT_TOKENS[0] = 32768
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 12288), [("b", 12288), ("c", 12288)], cap=12288,
+            context=100000, prompt_tokens=24000)
+        # After the first segment the server has said the prefix costs 24000,
+        # so the next round's input would be 36288 — past the configured 32768
+        # but nowhere near the 100000 the server advertises. The configured
+        # number is the one that must decide.
+        assert len(cl.sent) == 1, "it trusted the server over the operator"
+        assert rounds == 1 and still is True
+    finally:
+        mod._CONTEXT_TOKENS[0] = 0
+        m.undo()
+
+
+def test_a_cap_smaller_than_the_floor_is_never_mistaken_for_a_full_window():
+    """The 256-token floor is a floor on the room LEFT. A deployment whose cap
+    is below it must simply never be clamped."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume_ctx(
+            mod, m, ("a", 100), [("b", 100), ("c", 5)], cap=100,
+            context=100000, prompt_tokens=10)
+    finally:
+        m.undo()
+    assert out.raw == "abc" and rounds == 2, "a small cap stopped resumption"
+
+
+# --- a looping reply is not a long page -------------------------------------
+#
+# surya's own repeat detector re-runs such a page from scratch, so anything
+# resumption spends on it is spent on text that is about to be thrown away.
+# Measured on Aruba 2026-10-13: 253 of 262 resumptions ended in "came back
+# repeating itself", and 187 of those pages were then re-run in full.
+
+
+def test_a_reply_that_is_already_repeating_is_not_resumed_at_all():
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        oc = _fake_surya(m, mod)
+        oc._build_messages = lambda img, prompt: [{"role": "user", "content": prompt}]
+        oc.scale_to_fit = lambda img: img
+        oc.PROMPT_MAPPING = {"ocr": "ocr"}
+        oc.detect_repeat_token = lambda text: "loop" in (text or "")
+        res = _Result(raw="loop loop loop")
+        res.token_count = 12288
+        cl = _client([("more", 12288)])
+        out, rounds, still = mod._continue_reply(
+            oc, res, _Item(), cl, "m", 12288, 0.0, 0.1, 60.0, False, 8)
+    finally:
+        m.undo()
+    assert cl.sent == [], "it paid to extend a loop"
+    assert rounds == 0
+    # NOT truncated: it is a degenerate reply, which is surya's retry to make,
+    # not a page cut off for room. Calling it truncated would send the caller
+    # looking for a bigger budget it cannot use.
+    assert still is False
+    assert out.raw == "loop loop loop", "it altered a reply it did not resume"
+
+
+def test_a_page_that_really_is_long_is_still_resumed():
+    """The guard must key on repetition, not on having hit the cap."""
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        oc = _fake_surya(m, mod)
+        oc._build_messages = lambda img, prompt: [{"role": "user", "content": prompt}]
+        oc.scale_to_fit = lambda img: img
+        oc.PROMPT_MAPPING = {"ocr": "ocr"}
+        oc.detect_repeat_token = lambda text: "loop" in (text or "")
+        res = _Result(raw="<block>real text")
+        res.token_count = 12288
+        cl = _client([(" more real text", 40)])
+        out, rounds, still = mod._continue_reply(
+            oc, res, _Item(), cl, "m", 12288, 0.0, 0.1, 60.0, False, 8)
+    finally:
+        m.undo()
+    assert rounds == 1 and still is False
+    assert out.raw == "<block>real text more real text"