ocr: per-request token and continuation budgets, and a safe ceiling (v0.3.49)

/v1/ocr took engine, dpi, lang, structured, schema and detect. FastAPI drops
unknown form fields SILENTLY, so a caller sending max_tokens=8000 on every
request was discarded without a word and every page ran at the configured
value. Both budgets are now real request parameters.

One global cap cannot serve a public deployment: a four-page contract and a
fifty-page judgment need different budgets and only the caller knows which it
is sending. The values travel WITH the page rather than with the engine load,
because a pool is built once and then serves every caller, and they are
applied thread-locally in the worker, because several pages are in flight in
one process and the hooks they feed are module-level — a per-request value in
a module cell would be overwritten by the next page to arrive. Carried on a
ContextVar rather than widening `recognize_image`, which is the one thing
every engine has, duck-typed ones included.

Omitting a parameter means "use the configured value", not zero: 0 is
meaningful for both (no cap override, and resumption off).

surya_max_continuations 2 -> 8. It is PER PAGE, not per document: a fifty-page
document is fifty independent pages each with its own budget, so what the
number covers is the densest SINGLE page. The real ceiling is not the number
but the model's context — each resumption re-sends the text so far as INPUT,
so with surya-ocr-2's 32,768-token window and ~4,000 of it on the page image,
about thirteen 2,048-token segments exhaust it whatever is configured, and
raising surya_max_tokens lowers that count in proportion.

A generous budget is only safe if a page that stops making progress stops
costing, so resumption now ends when a segment comes back repeating itself or
returns the text already in hand. The second case is also how a server that
IGNORES continue_final_message is caught: a model that starts the page over
returns what we already have, and without the check a budget of 8 would buy
eight full generations of the same page.
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.48"
+__version__ = "0.3.49"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/ocr.py

@@ -61,6 +61,22 @@ def _ocr_config():
 
 
 @router.post("/v1/ocr", summary="OCR an image or PDF with a dedicated OCR engine")
+def _page_opts(max_tokens, max_continuations) -> dict:
+    """Per-request engine overrides, or {} when the caller said nothing.
+
+    One global cap cannot serve a public deployment: a four-page contract and a
+    fifty-page judgment need different budgets, and only the caller knows which
+    it is sending. Omitted means "use the configured value" -- NOT zero, which
+    is a meaningful setting for both of these.
+    """
+    out = {}
+    if max_tokens is not None:
+        out["max_tokens"] = max(0, int(max_tokens))
+    if max_continuations is not None:
+        out["max_continuations"] = max(0, int(max_continuations))
+    return out
+
+
 async def create_ocr(
     file: UploadFile = File(...),
     engine: Optional[str] = Form(None),        # paddle|doctr|surya|olmocr|chandra (default from config)
@@ -69,6 +85,8 @@ async def create_ocr(
     structured: Optional[bool] = Form(False),  # structured JSON extraction via text model
     schema_: Optional[str] = Form(None, alias="schema"),  # override extraction schema (JSON string)
     detect: Optional[str] = Form(None),        # off|layout|detector|both
+    max_tokens: Optional[int] = Form(None),    # cap on the model's reply for ONE page (engine default if unset)
+    max_continuations: Optional[int] = Form(None),  # resumptions allowed for a page that hits the cap
 ):
     cfg = _ocr_config()
     if cfg is None:
@@ -92,6 +110,7 @@ async def create_ocr(
             structured=bool(structured),
             schema=schema_,
             lang=lang,
+            page_opts=_page_opts(max_tokens, max_continuations),
         )
     except OcrError as e:
         raise HTTPException(status_code=e.status, detail=str(e))

codai/config.py

@@ -1076,11 +1076,20 @@ class OcrConfig:
     #: prefilled assistant turn (`continue_final_message`) so the model carries
     #: on the exact string rather than answering afresh about the same image.
     #:
-    #: 2 by default, because silently returning a page known to be incomplete is
-    #: a correctness failure rather than a speed/cost trade-off. 0 restores the
-    #: old behaviour. Each resumption is one more request, so the ceiling is
-    #: what stops a pathological page costing without bound.
-    surya_max_continuations: int = 2
+    #: PER PAGE, not per document: a fifty-page document is fifty independent
+    #: pages, each with its own budget, so what this number has to cover is the
+    #: densest SINGLE page, not the longest document.
+    #:
+    #: 8 by default, because returning a page known to be incomplete is a
+    #: correctness failure rather than a speed/cost trade-off, and because the
+    #: real ceiling is the model's context rather than this number: each
+    #: resumption re-sends the text so far as input, so with surya-ocr-2's
+    #: 32,768-token context and ~4,000 of it spent on the page image, roughly
+    #: thirteen 2,048-token segments exhaust the window whatever is set here.
+    #: Raising `surya_max_tokens` lowers that count in proportion. A page that
+    #: 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
     #: 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/base.py

@@ -23,6 +23,7 @@ extraction) speaks only these types, so engines stay interchangeable.
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
+import contextvars
 from typing import Any, List, Optional
 
 
@@ -231,6 +232,29 @@ class OcrEngine(ABC):
 _IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp", ".gif", ".ppm"}
 
 
+#: Per-REQUEST engine options, for the one page being read right now.
+#:
+#: A ContextVar rather than an argument because an engine is anything with
+#: `recognize_image(image)` -- the fallback chain and several tests drive
+#: duck-typed ones -- so widening that signature would break every engine to
+#: serve the two that can use it. A ContextVar also survives the hop into
+#: `asyncio.to_thread`, which copies the context, so it reaches a pooled
+#: subprocess engine on a worker thread and a lane inside it.
+#:
+#: Why per request at all: one global cap cannot serve a public deployment.
+#: A four-page contract and a fifty-page judgment need different budgets, and
+#: the caller is the only one who knows which it is sending.
+PAGE_OPTIONS = contextvars.ContextVar("ocr_page_options", default=None)
+
+
+def page_options() -> dict:
+    """Options the caller set for this request, or {}."""
+    try:
+        return dict(PAGE_OPTIONS.get() or {})
+    except Exception:
+        return {}
+
+
 def ink_coverage(image, threshold: int = 200) -> float:
     """Fraction of a page's pixels that are not background, 0.0-1.0.
 

codai/ocr/manager.py

@@ -34,6 +34,7 @@ import time
 from typing import Dict, List, Optional
 
 from codai.ocr.base import (OcrEngine, OcrPage, OcrError, load_pages, ink_coverage,
+                            PAGE_OPTIONS,
                             unreadable_report)
 
 
@@ -1049,7 +1050,8 @@ class OcrManager:
                            engine: Optional[str] = None, dpi: Optional[int] = None,
                            detect: Optional[str] = None, structured: bool = False,
                            schema: Optional[str] = None,
-                           lang: Optional[str] = None) -> dict:
+                           lang: Optional[str] = None,
+                           page_opts: Optional[dict] = None) -> dict:
         """OCR raw bytes (image/PDF) end-to-end → response dict.
 
         ``detect`` overrides the configured stamp/signature detection mode
@@ -1068,7 +1070,14 @@ class OcrManager:
         # Share the GPU thermal governor with every other workload: block here until
         # temps are safe (wait_until_safe blocks, so run it off the event loop).
         await asyncio.to_thread(_wait_thermal_safe)
-        name, pages = await self.ocr_pages(images, engine=engine, lang=lang)
+        token = None
+        if page_opts:
+            token = PAGE_OPTIONS.set(dict(page_opts))
+        try:
+            name, pages = await self.ocr_pages(images, engine=engine, lang=lang)
+        finally:
+            if token is not None:
+                PAGE_OPTIONS.reset(token)
         blank = _mark_blank_pages(cfg, images, pages)
         full_text = "\n\n".join(p.text for p in pages)
 

codai/ocr/subprocess_engine.py

@@ -20,7 +20,7 @@ import subprocess
 import sys
 import threading
 
-from codai.ocr.base import OcrEngine, OcrPage, OcrLine, OcrRegion, OcrError
+from codai.ocr.base import page_options, OcrEngine, OcrPage, OcrLine, OcrRegion, OcrError
 
 _WORKER = os.path.join(os.path.dirname(__file__), "workers", "ocr_worker.py")
 
@@ -292,11 +292,18 @@ class SubprocessOcrEngine(OcrEngine):
         buf = io.BytesIO()
         image.convert("RGB").save(buf, "PNG")
         b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+        # Per-request overrides travel WITH the page, not with the load: a pool
+        # is built once and then serves every caller, so anything set at load
+        # time is the same for all of them.
+        msg = {"cmd": "ocr", "image": b64}
+        opts = page_options()
+        if opts:
+            msg["opts"] = opts
         if self._mux:
-            resp = self._rpc_mux({"cmd": "ocr", "image": b64})
+            resp = self._rpc_mux(msg)
         else:
             with self._lock:                   # one request at a time per subprocess
-                resp = self._rpc({"cmd": "ocr", "image": b64})
+                resp = self._rpc(msg)
         if not resp.get("ok"):
             raise OcrError(resp.get("error") or f"{self.name}: ocr failed",
                            status=int(resp.get("status", 500)))

codai/ocr/workers/ocr_worker.py

@@ -11,7 +11,10 @@
 # Protocol: newline-delimited JSON on stdin/stdout.
 #   startup            → worker prints {"ready": true}
 #   {"cmd":"load","opts":{...}}   → {"ok":true} | {"ok":false,"error":..,"status":..}
-#   {"cmd":"ocr","image":"<b64 png>"} → {"ok":true,"page":{...}} | {"ok":false,...}
+#   {"cmd":"ocr","image":"<b64 png>","opts":{...}} → {"ok":true,"page":{...}} | {...}
+#       `opts` are PER-REQUEST overrides for this page only (max_tokens,
+#       max_continuations). Applied thread-locally, so lanes cannot overwrite
+#       each other's budgets. Absent = the values sent at load.
 #   {"cmd":"ping"}     → {"ok":true}
 # argv[1] = engine name ("paddle" | "surya").
 
@@ -75,6 +78,13 @@ def _trace_begin():
     return _TRACE.calls, _TRACE.repairs
 
 
+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")
+                   and str(v).strip() not in ("", "None")}
+
+
 def _trace_calls():
     return getattr(_TRACE, "calls", None)
 
@@ -114,6 +124,23 @@ def _attempt_tokens(n: int) -> None:
         pass
 
 
+def _page_override(name: str, default):
+    """A per-request value for the page this THREAD is reading, else the default.
+
+    The hooks these feed are module-level -- there is one `_generate_one` per
+    process -- so a per-request value cannot live in a module cell: with lanes,
+    several pages are in flight in one process and they would overwrite each
+    other's budgets. Same reasoning as the token trace.
+    """
+    try:
+        ov = getattr(_TRACE, "opts", None) or {}
+        if name in ov:
+            return ov[name]
+    except Exception:
+        pass
+    return default
+
+
 def _attempts():
     return getattr(_TRACE, "attempts", 0)
 
@@ -810,6 +837,27 @@ def _continue_reply(oc, result, item, client, model_name, cap, temperature,
             break
         if not more:
             break
+        # A resumption that adds nothing new is the end of what this page is
+        # going to give. Without this, a budget raised for the one genuinely
+        # enormous page also buys a degenerate page that many more full
+        # generations -- which is the failure the retry ceiling exists to stop,
+        # arriving through the other door.
+        if more == segments[-1]:
+            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 "
+                      f"rather than spending the rest of the budget on it",
+                      file=sys.stderr, flush=True)
+                segments.append(more)
+                rounds += 1
+                total += int(last or 0)
+                break
+        except Exception:
+            pass
         segments.append(more)
         # Counted only now. `rounds` is reported to the caller as "the page was
         # resumed n times, so it is whole" — an ATTEMPT that failed must not
@@ -935,7 +983,7 @@ class SuryaWorker:
             # 2048 was the only value available -- the same reason `max_retries`
             # needed hooking. It arrives here as a keyword, so it can be raised
             # without touching surya.
-            cap = _MAX_TOKENS[0]
+            cap = _page_override("max_tokens", _MAX_TOKENS[0])
             if cap and "max_tokens_default" in kwargs:
                 kwargs["max_tokens_default"] = int(cap)
             result = inner(*args, **kwargs)
@@ -943,7 +991,8 @@ class SuryaWorker:
             # it: raising the cap only moves the cliff, and incomplete block
             # markup parses to a fragment or to nothing at all.
             try:
-                budget = int(_CONTINUE_BUDGET[0] or 0)
+                budget = int(_page_override("max_continuations",
+                                            _CONTINUE_BUDGET[0]) or 0)
                 item = args[0] if args else kwargs.get("item")
                 eff_cap = int(getattr(item, "max_tokens", 0) or 0) or int(
                     kwargs.get("max_tokens_default") or 0)
@@ -1114,7 +1163,8 @@ class SuryaWorker:
         # fewer completions than the GPU's throughput says it should.
         tries = _attempts()
         toks = _attempt_token_counts()
-        cap = int(_MAX_TOKENS[0] or _SURYA_DEFAULT_MAX_TOKENS)
+        cap = int(_page_override("max_tokens", _MAX_TOKENS[0])
+                  or _SURYA_DEFAULT_MAX_TOKENS)
         # Reported whenever there is anything to report -- a retry, or a reply
         # that stopped at its cap. A page cut off mid-markup parses to a few
         # blocks or to none, which is indistinguishable in the finished text
@@ -1241,6 +1291,10 @@ def _serve_page(lanes: "_Lanes", msg: dict, engine: str) -> None:
         from PIL import Image
         data = base64.b64decode(msg["image"])
         img = Image.open(io.BytesIO(data)); img.load()
+        # On THIS thread, before the page is read. _trace_begin() inside ocr()
+        # resets the page's trace but not the options, so they survive it -- and
+        # being thread-local they cannot leak into another lane's page.
+        _trace_opts(msg.get("opts"))
         _emit({"ok": True, "id": rid, "page": worker.ocr(img)})
     except Exception as e:
         _emit({"ok": False, "id": rid, "error": f"{engine} ocr failed: {e}",
@@ -1291,6 +1345,7 @@ def main():
                     from PIL import Image
                     data = base64.b64decode(msg["image"])
                     img = Image.open(io.BytesIO(data)); img.load()
+                    _trace_opts(msg.get("opts"))
                     _emit({"ok": True, "page": worker.ocr(img)})
                 except Exception as e:
                     _emit({"ok": False, "error": f"{engine} ocr failed: {e}",

docs/ocr.md

@@ -1074,6 +1074,44 @@ deliberately *raises* temperature to break a repetition. A page still unfinished
 after its budget is marked `"complete": false`, so a caller can refuse a partial
 rather than store it as read.
 
+**The budget is per PAGE, not per document.** A fifty-page document is fifty
+independent pages, each with its own budget, so what the number has to cover is
+the densest *single* page — not the longest document.
+
+The real ceiling is not this number but the model's context. Each resumption
+re-sends the text so far as *input*, so with surya-ocr-2's 32,768-token window
+and roughly 4,000 of it spent on the page image, about thirteen 2,048-token
+segments exhaust it whatever is configured; raising `surya_max_tokens` lowers
+that count in proportion. The default is **8**, which covers any plausible
+single page while staying inside the window.
+
+A generous budget is safe because a page that stops making progress stops
+costing: a resumption that comes back repeating itself, or that returns the text
+already in hand, ends the loop. That second case is also how a server that
+**ignores** `continue_final_message` is caught — a model that starts the page
+over returns what we already have, and without the check a budget of 8 would buy
+eight full generations of the same page.
+
+### Per-request budgets
+
+One global cap cannot serve a public deployment: a four-page contract and a
+fifty-page judgment need different budgets, and only the caller knows which it is
+sending. `/v1/ocr` and `/v1/ocr/batch` therefore accept `max_tokens` and
+`max_continuations` per request, overriding the configured values for that call
+only. Omitting them means "use the configured value" — **not** zero, which is a
+meaningful setting for both (no override, and resumption off).
+
+They travel *with the page* rather than with the engine load, because a pool is
+built once and then serves every caller, and they are applied **thread-locally**
+in the worker, because several pages are in flight in one process and the hooks
+they feed are module-level — a per-request value in a module cell would be
+overwritten by the next page to arrive.
+
+One caveat worth knowing, since it cost a day of wrong conclusions: FastAPI
+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.
+
 ---
 
 ## An empty page and a 71-character page are the same bug

tests/test_surya_retry_budget.py

@@ -569,7 +569,8 @@ def test_the_engine_sends_the_resumption_budget():
     from codai.ocr.surya import SuryaEngine
     from codai.config import OcrConfig
     cfg = OcrConfig()
-    assert SuryaEngine(cfg)._worker_opts()["max_continuations"] == 2
+    assert (SuryaEngine(cfg)._worker_opts()["max_continuations"]
+            == cfg.surya_max_continuations)
     cfg.surya_max_continuations = 0
     assert SuryaEngine(cfg)._worker_opts()["max_continuations"] == 0
 
@@ -579,3 +580,107 @@ def test_resumption_is_on_by_default():
     cost trade-off, so this is not opt-in."""
     from codai.config import OcrConfig
     assert OcrConfig().surya_max_continuations >= 1
+
+
+# ------------------------------------------- per-request budgets, and safety #
+#
+# One global cap cannot serve a public deployment: a four-page contract and a
+# fifty-page judgment need different budgets, and only the caller knows which it
+# is sending. /v1/ocr took engine, dpi, lang, structured, schema and detect —
+# anything else a caller sent was silently dropped by FastAPI, so a `max_tokens`
+# in the request was discarded and the page ran at the configured value.
+
+def test_the_request_can_set_the_budgets():
+    from codai.api.ocr import _page_opts
+    assert _page_opts(8000, 20) == {"max_tokens": 8000, "max_continuations": 20}
+
+
+def test_omitting_them_is_not_the_same_as_zero():
+    """0 is meaningful for both — no cap override, and resumption off. A caller
+    that said nothing must get the configured value, not zero."""
+    from codai.api.ocr import _page_opts
+    assert _page_opts(None, None) == {}
+    assert _page_opts(0, 0) == {"max_tokens": 0, "max_continuations": 0}
+
+
+def test_the_endpoint_actually_accepts_them():
+    """The bug this fixes: FastAPI drops unknown form fields silently, so a
+    caller can send max_tokens for a week and never be told it did nothing."""
+    import inspect
+    from codai.api.ocr import create_ocr
+    params = inspect.signature(create_ocr).parameters
+    assert "max_tokens" in params and "max_continuations" in params
+
+
+def test_the_page_carries_the_options_to_the_worker():
+    src = (ROOT / "codai/ocr/subprocess_engine.py").read_text()
+    blk = src.split("def recognize_image")[1].split("\n    @staticmethod")[0]
+    assert '"cmd": "ocr"' in blk and 'msg["opts"] = opts' in blk
+
+
+def test_overrides_are_per_thread_so_lanes_cannot_cross():
+    """Several pages are in flight in one process, and the hooks they feed are
+    module-level. A per-request value in a module cell would be overwritten by
+    the next page to arrive."""
+    mod = _worker_module()
+    import threading as _th
+    seen = {}
+
+    def lane(name, opts, expect):
+        mod._trace_opts(opts)
+        mod._trace_begin()
+        seen[name] = mod._page_override("max_tokens", 2048)
+
+    a = _th.Thread(target=lane, args=("a", {"max_tokens": 9000}, 9000))
+    b = _th.Thread(target=lane, args=("b", {}, 2048))
+    a.start(); b.start(); a.join(); b.join()
+    assert seen == {"a": 9000, "b": 2048}
+
+
+def test_beginning_a_page_does_not_clear_its_options():
+    """`_trace_begin()` resets the page's trace; the options are set just before
+    it and have to survive."""
+    mod = _worker_module()
+    mod._trace_opts({"max_continuations": 5})
+    mod._trace_begin()
+    assert mod._page_override("max_continuations", 2) == 5
+
+
+def test_only_known_options_are_accepted():
+    mod = _worker_module()
+    mod._trace_opts({"max_tokens": 10, "rm": "-rf", "lanes": 99})
+    assert mod._page_override("lanes", 1) == 1
+    assert mod._page_override("max_tokens", 0) == 10
+
+
+def test_a_resumption_that_repeats_itself_stops_early():
+    """A budget raised for the one enormous page must not also buy a degenerate
+    page that many more full generations.
+
+    It also catches the server ignoring `continue_final_message`: a model that
+    starts the page over returns the text we already have, which is a verbatim
+    repeat of the previous segment. Resuming eight times against a server that
+    does not support resumption would otherwise cost eight full generations and
+    produce the page eight times over.
+    """
+    mod = _worker_module()
+    import _pytest.monkeypatch as _mp
+    m = _mp.MonkeyPatch()
+    try:
+        out, rounds, still, cl = _resume(mod, m, ("aaa", 100),
+                                         [("aaa", 100), ("aaa", 100)], budget=8)
+    finally:
+        m.undo()
+    # Nothing was added, so nothing was resumed: reporting a successful
+    # resumption here would be the claim the counter was fixed to stop making.
+    assert rounds == 0
+    assert len(cl.sent) == 1, "it kept paying after the repeat"
+    assert out.raw == "aaa", "the repeat must not be appended"
+    assert still is True, "an unfinished page must still say so"
+
+
+def test_the_default_budget_covers_a_dense_page():
+    """Per PAGE, not per document: a fifty-page document is fifty pages each
+    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