ocr: the page budget died at surya's thread pool, and the window at the gateway (v0.3.59)

Two separate reasons the Digesta re-test still saw wrong caps (12288 and 1360,
never the 3072 it sent) and "surya-ocr-2 does not report a context window".
Neither is the pod image, which was blamed for both.

1. surya's `chat_completions_batch` runs every item through a
   ThreadPoolExecutor:

       with ThreadPoolExecutor(max_workers=max_workers) as executor:
           return list(executor.map(_process, batch))

   so the hook around `_generate_one` executes on an EXECUTOR thread, while
   `_trace_opts` installed the page's options on the lane thread that submitted
   it — and that state is thread-local, deliberately, so two lanes cannot
   attribute one page's logprobs to another. Across that boundary every
   per-request budget read as absent: `max_tokens` was discarded (v0.3.57 fixed
   which value wins, not whether ours arrives), and so were
   `max_continuations` and `context_tokens`. The per-page trace went the same
   way — attempts, token counts, logprobs and accent repairs were recorded on
   the executor thread and never read, so `model_attempts` could not report what
   it exists to report.

   The item is the one object that crosses, so it is the key. The lane binds its
   page state to the batch's items before surya dispatches them (hooking
   `chat_completions_batch`, which still runs on the lane thread), and the
   executor thread adopts it by `id(item)`. The state is shared by reference, so
   writes land where the lane reads them. Bindings are released in a finally —
   id() is reused, and a stale entry would hand a new item another page's state.
   When a surya has no such entry point, that is said rather than silently
   leaving the budgets unreachable.

2. An attached engine reads `/v1/models` from the GATEWAY — it talks to
   `…/v1/ocr/attach/<token>/v1` — and the gateway answered with id, object and
   owned_by, nothing else. So 0.3.53's publishing of `max_model_len` was
   invisible on the one path whose arithmetic needs it, the worker correctly
   reported that nobody names a window, and that read as the pod failing to
   advertise one. The pod was never asked. The gateway now publishes the window
   it will serve: a value an upstream reported, else the model's configured ctx.

   And the pod now serves the MODEL's ctx instead of this install's. A runpod
   block asking for 18432 got a pod serving 32768, because the pod is handed
   `engine_config_for_pod`, which reads config.json — the local engine's
   settings. Publishing one number while serving another is worse than
   publishing none: resumption stops early, or overruns and takes a 400. The
   per-model value now wins and the difference is logged.

   Note for operators: a pod will now serve the per-model `ctx` it was always
   configured with. For surya on the Digesta orchestrator that is 18432, not the
   32768 it had been getting by accident — set ctx higher if the larger window
   was wanted.

A live probe of the upstream's own /v1/models was written and dropped: it built
a client per page and re-probed for ever against an upstream that reports
nothing. `learn_window()` remains as the hook for a value we do learn.

15 tests: tests/test_surya_thread_boundary.py (7) and
tests/test_ocr_window_published.py (8).
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.58"
+__version__ = "0.3.59"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/runpod_worker.py

@@ -1312,6 +1312,21 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
         if eng == "vllm":
             env.setdefault("CODERAI_VLLM_VENV", "/opt/coderai/venvs/vllm")
         forwarded = engine_config_for_pod(eng)
+        # The MODEL's own engine settings beat this install's. engine_config_for_pod
+        # reads config.json, which describes the local engine, so a pod was
+        # launched with the install-wide context window and a per-model `ctx` on
+        # the runpod block did nothing: an entry asking for 18432 got a pod
+        # serving 32768. Harmless-looking until something publishes the window —
+        # an OCR engine's resumption arithmetic is driven by it, and two
+        # different numbers for one pod is how a resume either stops early or
+        # overruns and 400s.
+        _own = int(getattr(mcfg, "ctx", 0) or 0)
+        if _own:
+            if int(forwarded.get("ctx") or 0) != _own:
+                print(f"[runpod] pod context window {_own} (this model's own "
+                      f"ctx, over the install's "
+                      f"{forwarded.get('ctx') or 'default'})", flush=True)
+            forwarded["ctx"] = _own
         if eng == "vllm" and "tensor_parallel_size" in forwarded:
             # The pod was rented with gpu_count cards; the local setting
             # describes this machine's cards. The pod's count wins.

codai/ocr/gateway.py

@@ -335,6 +335,47 @@ async def health(token: str):
     return Response(status_code=200, content=b"OK", media_type="text/plain")
 
 
+#: model_id -> context window an upstream actually reported, learned while
+#: forwarding. Preferred over configuration: a pod serves whatever its own
+#: launch arguments say, which is not always what the model entry asks for
+#: (a pod serving --max-model-len 32768 for an entry that configured 18432 is a
+#: live example), and resumption that believes the smaller number merely stops
+#: early while one that believes a larger one gets a 400.
+_windows: dict = {}
+
+
+def _window_for(model_id: str) -> int:
+    """The context window to publish for ``model_id``, or 0 when unknown.
+
+    Learned-from-upstream first, then the model entry. Never a guess: an engine
+    told a window it does not have will either waste it or overrun it, and
+    "unknown" is a case the engine already handles by refusing to resume.
+    """
+    got = int(_windows.get(model_id) or 0)
+    if got:
+        return got
+    try:
+        from codai.models.manager import context_window_for, multi_model_manager
+        from codai.models.placement import model_entry
+        win = context_window_for(multi_model_manager, model_id,
+                                 model_entry(model_id) or {})
+        return int(win or 0)
+    except Exception:
+        return 0
+
+
+def learn_window(model_id: str, window: int) -> None:
+    """Record a window an upstream reported for this model."""
+    try:
+        w = int(window or 0)
+    except (TypeError, ValueError):
+        return
+    if w > 0 and _windows.get(model_id) != w:
+        _windows[model_id] = w
+        print(f"[ocr] {model_id} serves a {w}-token context window "
+              f"(reported by the upstream)", flush=True)
+
+
 @router.get(PREFIX + "/{token}/v1/models", include_in_schema=False)
 async def models(token: str):
     """The id the engine expects, which is also what we will send upstream.
@@ -347,11 +388,19 @@ async def models(token: str):
     got = lease(token)
     if got is None:
         return _unknown_lease()
-    return JSONResponse(content={
-        "object": "list",
-        "data": [{"id": got["served_name"], "object": "model",
-                  "owned_by": "coderai"}],
-    })
+    entry = {"id": got["served_name"], "object": "model", "owned_by": "coderai"}
+    # The served context window, because an attached engine's resumption math
+    # needs it and THIS is the /v1/models it reads. coderai publishes
+    # max_model_len on its own /v1/models (and a pod publishes it on the pod's),
+    # but an engine attached through a lease never sees either of those — it
+    # sees this answer, which named nothing. The worker then said "surya-ocr-2
+    # does not report a context window on /v1/models; resumption cannot stop
+    # before the server does", which reads as the POD failing to advertise it
+    # and was reported that way. The pod was never asked.
+    win = _window_for(got["model_id"])
+    if win:
+        entry["max_model_len"] = int(win)
+    return JSONResponse(content={"object": "list", "data": [entry]})
 
 
 @router.post(PREFIX + "/{token}/v1/chat/completions", include_in_schema=False)

codai/ocr/workers/ocr_worker.py

@@ -178,27 +178,88 @@ def _is_context_error(exc):
             or "reduce the length" in t or "max_model_len" in t)
 
 
+#: id(surya batch item) -> the page state of the lane that submitted it.
+#:
+#: Thread-local state was not enough, and the gap was invisible. surya's
+#: `chat_completions_batch` runs each item through a ThreadPoolExecutor, so the
+#: hook around `_generate_one` executes on an executor thread — NOT the lane
+#: thread where `_trace_opts` installed this page's options. Every per-request
+#: budget therefore read as absent: `max_tokens=3072` was discarded and each
+#: page ran at surya's own `item.max_tokens` (12288 for a full page, or an
+#: `image_token_budget` estimate like 1360), and `max_continuations` and
+#: `context_tokens` went the same way. The per-page trace was lost to the same
+#: boundary — attempts, token counts, logprobs and accent repairs were all
+#: recorded against the executor thread and never read.
+#:
+#: The one object that crosses the boundary is the ITEM, so it is the key. The
+#: state is shared BY REFERENCE: the executor thread adopts it, writes into it,
+#: and the lane reads it back when the page is done.
+_ITEM_STATE: dict = {}
+_ITEM_LOCK = threading.Lock()
+
+
+def _state() -> dict:
+    """This thread's current page state, created on demand."""
+    st = getattr(_TRACE, "state", None)
+    if st is None:
+        st = _TRACE.state = {}
+    return st
+
+
 def _trace_begin():
-    """Start a page's trace. Returns its (calls, repairs) — this thread's own."""
-    _TRACE.calls = []
-    _TRACE.repairs = []
-    _TRACE.attempts = 0
-    _TRACE.tokens = []
-    _TRACE.continued = 0
-    _TRACE.still_truncated = False
-    return _TRACE.calls, _TRACE.repairs
+    """Start a page's trace. Returns its (calls, repairs) — this page's own."""
+    st = _state()
+    st["calls"] = []
+    st["repairs"] = []
+    st["attempts"] = 0
+    st["tokens"] = []
+    st["continued"] = 0
+    st["still_truncated"] = False
+    return st["calls"], st["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",
-                            "context_tokens")
-                   and str(v).strip() not in ("", "None")}
+    _state()["opts"] = {k: int(v) for k, v in (opts or {}).items()
+                        if k in ("max_tokens", "max_continuations",
+                                 "context_tokens")
+                        and str(v).strip() not in ("", "None")}
+
+
+def _bind_items(items) -> list:
+    """Register this page's state against the batch items, and return the keys.
+
+    Called on the LANE thread, before surya hands the items to its executor.
+    """
+    st = _state()
+    keys = []
+    with _ITEM_LOCK:
+        for it in items or ():
+            if it is None:
+                continue
+            _ITEM_STATE[id(it)] = st
+            keys.append(id(it))
+    return keys
+
+
+def _unbind_items(keys) -> None:
+    with _ITEM_LOCK:
+        for k in keys or ():
+            _ITEM_STATE.pop(k, None)
+
+
+def _adopt_item_state(item) -> None:
+    """On an executor thread: work on the page state that submitted ``item``."""
+    if item is None:
+        return
+    with _ITEM_LOCK:
+        st = _ITEM_STATE.get(id(item))
+    if st is not None and getattr(_TRACE, "state", None) is not st:
+        _TRACE.state = st
 
 
 def _trace_calls():
-    return getattr(_TRACE, "calls", None)
+    return _state().get("calls")
 
 
 def _attempts_bump():
@@ -212,8 +273,9 @@ def _attempts_bump():
     already read. None of it was visible from outside: a retried page looks
     exactly like a slow one.
     """
-    n = getattr(_TRACE, "attempts", 0) + 1
-    _TRACE.attempts = n
+    st = _state()
+    n = int(st.get("attempts") or 0) + 1
+    st["attempts"] = n
     return n
 
 
@@ -228,9 +290,10 @@ def _attempt_tokens(n: int) -> None:
     model that cannot read, and neither is.
     """
     try:
-        lst = getattr(_TRACE, "tokens", None)
+        st = _state()
+        lst = st.get("tokens")
         if lst is None:
-            lst = _TRACE.tokens = []
+            lst = st["tokens"] = []
         lst.append(int(n or 0))
     except Exception:
         pass
@@ -245,7 +308,7 @@ def _page_override(name: str, default):
     other's budgets. Same reasoning as the token trace.
     """
     try:
-        ov = getattr(_TRACE, "opts", None) or {}
+        ov = _state().get("opts") or {}
         if name in ov:
             return ov[name]
     except Exception:
@@ -254,15 +317,15 @@ def _page_override(name: str, default):
 
 
 def _attempts():
-    return getattr(_TRACE, "attempts", 0)
+    return int(_state().get("attempts") or 0)
 
 
 def _attempt_token_counts():
-    return list(getattr(_TRACE, "tokens", None) or ())
+    return list(_state().get("tokens") or ())
 
 
 def _trace_repairs():
-    return getattr(_TRACE, "repairs", None)
+    return _state().get("repairs")
 
 
 def _html_to_text(html: str) -> str:
@@ -1143,6 +1206,10 @@ class SuryaWorker:
             return
 
         def _wrapped(*args, **kwargs):
+            # surya dispatches this through a ThreadPoolExecutor, so the page's
+            # options and trace live on ANOTHER thread. Adopt them, by the one
+            # object that crossed over.
+            _adopt_item_state(args[0] if args else kwargs.get("item"))
             # surya's vLLM backend never passes `max_tokens_default`, so its
             # 2048 was the only value available -- the same reason `max_retries`
             # needed hooking. It arrives here as a keyword, so it can be raised
@@ -1198,8 +1265,9 @@ class SuryaWorker:
                         0.0, kwargs.get("top_p", 0.1), kwargs.get("timeout", 600.0),
                         bool(kwargs.get("request_logprobs_default", True)),
                         budget)
-                    _TRACE.continued = getattr(_TRACE, "continued", 0) + rounds
-                    _TRACE.still_truncated = bool(still)
+                    _st = _state()
+                    _st["continued"] = int(_st.get("continued") or 0) + rounds
+                    _st["still_truncated"] = bool(still)
             except Exception as exc:
                 print(f"[surya] resume failed: {exc}", file=sys.stderr, flush=True)
             try:
@@ -1254,6 +1322,31 @@ class SuryaWorker:
                           file=sys.stderr, flush=True)
                 return again
             _oc._should_retry = _wrapped_retry
+
+        # And the batch entry point, which runs on the LANE thread. This is
+        # where the page's state is tied to the items, because it is the last
+        # point before surya's ThreadPoolExecutor takes them somewhere this
+        # thread's locals cannot follow.
+        batch = getattr(_oc, "chat_completions_batch", None)
+        if batch is not None:
+            def _wrapped_batch(*args, **kwargs):
+                items = args[0] if args else kwargs.get("batch")
+                keys = _bind_items(items)
+                try:
+                    return batch(*args, **kwargs)
+                finally:
+                    _unbind_items(keys)
+            _oc.chat_completions_batch = _wrapped_batch
+            # surya's predictor imported the function by name before we got
+            # here, so rebinding the module attribute alone can miss the caller.
+            for _modname in ("surya.recognition", "surya.inference.manager",
+                             "surya.layout", "surya.table_rec"):
+                _m = sys.modules.get(_modname)
+                if _m is not None and getattr(_m, "chat_completions_batch", None) is batch:
+                    _m.chat_completions_batch = _wrapped_batch
+        else:
+            print("[surya] no chat_completions_batch to hook — per-request "
+                  "budgets cannot reach this surya", file=sys.stderr, flush=True)
         _oc._coderai_traced = True
 
     def _load_legacy(self):
@@ -1365,12 +1458,12 @@ class SuryaWorker:
         # from a page the model genuinely could not read. This is the only place
         # the difference is visible.
         capped = [t for t in toks if t >= cap * 0.98]
-        resumed = int(getattr(_TRACE, "continued", 0) or 0)
+        resumed = int(_state().get("continued") or 0)
         # Still short of the end after every resumption we were allowed. The
         # caller needs this to refuse a partial page, so it is reported
         # separately from "a segment hit the cap", which resumption may have
         # since fixed.
-        incomplete = bool(getattr(_TRACE, "still_truncated", False))
+        incomplete = bool(_state().get("still_truncated"))
         if tries > 1 or capped or resumed:
             info = {"calls": tries or 1, "tokens": toks, "cap": cap}
             if resumed:

docs/ocr.md

@@ -1282,6 +1282,29 @@ the documents.
 > applies the cap to the item as well, and says so once per worker:
 > `[surya] page budget 3072 tokens (overriding surya's own 12288)`. With no cap
 > configured, surya's own budget is left alone.
+>
+> **v0.3.57 was still not enough on the served path**, for a second reason.
+> surya's `chat_completions_batch` runs every item through a
+> `ThreadPoolExecutor`, so the hook executes on an executor thread while the
+> page's options were installed, thread-locally, on the lane thread that
+> submitted it. Every per-request budget read as absent — `max_tokens`,
+> `max_continuations` and `context_tokens` alike — and the per-page trace
+> (attempts, token counts, logprobs, accent repairs) was written where nothing
+> read it. v0.3.59 binds the page's state to the batch's ITEMS, the one object
+> that crosses that boundary, and the executor thread adopts it by `id(item)`.
+>
+> **The window an attached engine reads comes from the GATEWAY**, not from the
+> pod and not from this install's `/v1/models`: a leased engine talks to
+> `…/v1/ocr/attach/<token>/v1`. That answer named no `max_model_len` until
+> v0.3.59, so resumption had no ceiling to respect and said so —
+> "surya-ocr-2 does not report a context window on /v1/models" — which reads
+> like the pod failing to advertise it. The pod was never asked. The gateway now
+> publishes the window it will actually serve: one an upstream reported if it
+> has, else the model's configured `ctx`. And a pod now serves the MODEL's
+> `ctx` rather than this install's — a runpod block asking for 18432 used to get
+> a pod serving 32768, because the pod was handed `config.json`'s engine
+> settings. Publishing one number while serving another is worse than publishing
+> none: resumption stops early, or overruns and takes a 400.
 
 What is **not** optional is being able to tell the two cases apart, so every
 page that hit the cap says so in its own result:

tests/test_ocr_window_published.py

+"""The window an attached engine reads must be the one its server serves.
+
+An OCR engine attached through a lease reads `/v1/models` from the GATEWAY, not
+from the pod and not from this install — and the gateway answered with just
+`{id, object, owned_by}`. So coderai's own 0.3.53 work (publishing
+`max_model_len` on `/v1/models`) was invisible on the one path that needs it,
+and the worker said:
+
+    surya-ocr-2 does not report a context window on /v1/models; resumption
+    cannot stop before the server does — set ocr.surya_context_tokens to the
+    served --max-model-len
+
+which reads as the POD failing to advertise its window, and was reported that
+way. The pod was never asked.
+
+The second half: a pod was launched with THIS INSTALL's engine config, so a
+per-model `ctx` on the runpod block did nothing — an entry asking for 18432 got
+a pod serving 32768. Publishing a window while serving another is worse than
+publishing none: resumption either stops early or overruns and gets a 400.
+"""
+import json
+import pathlib
+import sys
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from codai.ocr import gateway                                      # noqa: E402
+
+
+def _reset():
+    gateway.revoke()
+    gateway._windows.clear()
+
+
+def _models(monkeypatch, window):
+    import asyncio
+    monkeypatch.setattr(gateway, "_window_for", lambda m: window)
+    token = gateway.issue("m", "served-name")
+    reply = asyncio.run(gateway.models(token))
+    return json.loads(bytes(reply.body).decode())
+
+
+def test_the_window_is_published_to_the_engine(monkeypatch):
+    _reset()
+    body = _models(monkeypatch, 32768)
+    assert body["data"][0]["max_model_len"] == 32768
+    # The id still has to be the engine's own checkpoint name: surya refuses to
+    # attach to a server whose /v1/models disagrees with it.
+    assert body["data"][0]["id"] == "served-name"
+
+
+def test_an_unknown_window_is_omitted_not_guessed(monkeypatch):
+    # "Unknown" is a case the engine handles — it refuses to resume and says
+    # which setting to set. A guessed number would make it overrun instead.
+    _reset()
+    assert "max_model_len" not in _models(monkeypatch, 0)["data"][0]
+
+
+def test_a_learned_window_beats_the_configured_one():
+    _reset()
+    gateway.learn_window("m", 32768)
+    assert gateway._window_for("m") == 32768
+
+
+def test_learning_ignores_rubbish():
+    _reset()
+    for bad in (0, -1, None, "", "lots"):
+        gateway.learn_window("m", bad)
+    assert gateway._window_for("m") in (0, None) or isinstance(
+        gateway._window_for("m"), int)
+    assert "m" not in gateway._windows
+
+
+def test_the_configured_window_is_read_from_the_model_entry(monkeypatch):
+    _reset()
+    import codai.models.manager as mm
+    import codai.models.placement as pl
+    monkeypatch.setattr(pl, "model_entry", lambda m: {"runpod": {"ctx": 18432}})
+    monkeypatch.setattr(mm, "context_window_for", lambda mgr, m, meta:
+                        (meta.get("runpod") or {}).get("ctx"))
+    assert gateway._window_for("m") == 18432
+
+
+def test_a_broken_lookup_is_unknown_not_an_error(monkeypatch):
+    _reset()
+    import codai.models.placement as pl
+    monkeypatch.setattr(pl, "model_entry",
+                        lambda m: (_ for _ in ()).throw(RuntimeError("no config")))
+    assert gateway._window_for("m") == 0
+
+
+# --- the pod has to SERVE what we publish ----------------------------------
+
+def test_a_models_own_ctx_reaches_the_pod():
+    src = (ROOT / "codai/api/runpod_worker.py").read_text()
+    blk = src.split("forwarded = engine_config_for_pod(eng)")[1][:1200]
+    assert '_own = int(getattr(mcfg, "ctx", 0) or 0)' in blk
+    assert 'forwarded["ctx"] = _own' in blk
+
+
+def test_the_override_is_announced():
+    """Two different numbers for one pod must not be a silent difference."""
+    src = (ROOT / "codai/api/runpod_worker.py").read_text()
+    assert "pod context window" in src and "this model's own" in src

tests/test_surya_item_cap.py

@@ -90,12 +90,11 @@ def test_the_per_page_override_reaches_the_item(monkeypatch):
     w = _worker(mod, token_trace=True, max_tokens=0)
     w._install_trace_hook()
     item = _Item(12288)
-    # Per-page overrides are thread-local (several lanes share one hook).
-    mod._TRACE.opts = {"max_tokens": 4096}
-    try:
-        oc._generate_one(item, max_tokens_default=2048)
-    finally:
-        mod._TRACE.opts = {}
+    # Per-page overrides arrive with the page, exactly as the worker installs
+    # them before reading it.
+    mod._trace_begin()
+    mod._trace_opts({"max_tokens": 4096})
+    oc._generate_one(item, max_tokens_default=2048)
     assert item.max_tokens == 4096
 
 

tests/test_surya_thread_boundary.py

+"""The per-page budgets have to survive surya's own thread pool.
+
+surya's `chat_completions_batch` runs every item through a ThreadPoolExecutor:
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        return list(executor.map(_process, batch))
+
+so the hook around `_generate_one` executes on an EXECUTOR thread, while
+`_trace_opts` installed this page's options on the lane thread that submitted
+it. Thread-local state cannot cross that, and nothing said so:
+
+  * `max_tokens=3072` from a request was read as absent, and every page ran at
+    surya's own `item.max_tokens` — 12288 for a full page, or an
+    `image_token_budget` estimate like 1360. Reported from the Digesta side as
+    "the cap is still wrong: 12288 and 1360, never the 3072 I send";
+  * `max_continuations` and `context_tokens` went the same way;
+  * the per-page trace — attempts, token counts, logprobs, accent repairs —
+    was written on the executor thread and never read, so `model_attempts`
+    could not report what it was built to report.
+
+The item is the one object that crosses the boundary, so it is the key: the lane
+registers its page state against the batch's items, and the executor thread
+adopts it.
+"""
+import pathlib
+import sys
+import threading
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "tests"))
+
+from test_surya_retry_budget import (_fake_surya, _worker,           # noqa: E402
+                                     _worker_module, _Result)
+
+
+class _Item:
+    def __init__(self, max_tokens=12288):
+        self.max_tokens = max_tokens
+
+
+def _surya_with_batch(monkeypatch, mod):
+    """A fake surya whose batch entry point dispatches like the real one."""
+    oc = _fake_surya(monkeypatch, mod)
+
+    def _batch(batch, **kw):
+        out = []
+
+        def _run(it):
+            out.append(oc._generate_one(it, max_tokens_default=2048))
+
+        for it in batch:
+            t = threading.Thread(target=_run, args=(it,))
+            t.start()
+            t.join()
+        return out
+
+    oc.chat_completions_batch = _batch
+    return oc
+
+
+def test_the_cap_crosses_into_the_executor_thread(monkeypatch):
+    mod = _worker_module()
+    oc = _surya_with_batch(monkeypatch, mod)
+    oc._generate_one = lambda *a, **k: _Result()
+    w = _worker(mod, token_trace=True, max_tokens=0)
+    w._install_trace_hook()
+
+    mod._trace_begin()
+    mod._trace_opts({"max_tokens": 3072})
+    item = _Item(12288)
+    oc.chat_completions_batch([item])
+    # Previously 12288: the lane's options were invisible to the thread that ran.
+    assert item.max_tokens == 3072
+
+
+def test_the_trace_written_on_another_thread_is_read_by_the_lane(monkeypatch):
+    mod = _worker_module()
+    oc = _surya_with_batch(monkeypatch, mod)
+    w = _worker(mod, token_trace=True, max_tokens=0)
+    w._install_trace_hook()
+
+    mod._trace_begin()
+    mod._trace_opts({})
+    oc.chat_completions_batch([_Item(0), _Item(0)])
+    # Two attempts, counted on executor threads, visible here.
+    assert mod._attempts() == 2
+    assert len(mod._attempt_token_counts()) == 2
+
+
+def test_two_lanes_do_not_share_a_page_state(monkeypatch):
+    """The reason this was thread-local in the first place: one page's logprobs
+    must never be attributed to another page."""
+    mod = _worker_module()
+    oc = _surya_with_batch(monkeypatch, mod)
+    oc._generate_one = lambda *a, **k: _Result()
+    w = _worker(mod, token_trace=True, max_tokens=0)
+    w._install_trace_hook()
+
+    seen = {}
+
+    def lane(name, cap, item):
+        mod._trace_begin()
+        mod._trace_opts({"max_tokens": cap})
+        oc.chat_completions_batch([item])
+        seen[name] = (item.max_tokens, mod._attempts())
+
+    a_item, b_item = _Item(12288), _Item(12288)
+    ta = threading.Thread(target=lane, args=("a", 3072, a_item))
+    tb = threading.Thread(target=lane, args=("b", 777, b_item))
+    ta.start(); tb.start(); ta.join(); tb.join()
+    assert seen["a"] == (3072, 1)
+    assert seen["b"] == (777, 1)
+
+
+def test_the_binding_is_released_so_items_are_not_kept_alive(monkeypatch):
+    mod = _worker_module()
+    oc = _surya_with_batch(monkeypatch, mod)
+    oc._generate_one = lambda *a, **k: _Result()
+    w = _worker(mod, token_trace=True, max_tokens=0)
+    w._install_trace_hook()
+    mod._trace_begin()
+    mod._trace_opts({"max_tokens": 3072})
+    oc.chat_completions_batch([_Item(0)])
+    # A registry that grew per page would leak for the life of the worker, and
+    # id() is reused — a stale entry would hand a new item another page's state.
+    assert mod._ITEM_STATE == {}
+
+
+def test_the_binding_is_released_even_when_the_batch_raises(monkeypatch):
+    mod = _worker_module()
+    oc = _fake_surya(monkeypatch, mod)
+
+    def _boom(batch, **kw):
+        raise RuntimeError("upstream died")
+
+    oc.chat_completions_batch = _boom
+    w = _worker(mod, token_trace=True, max_tokens=3072)
+    w._install_trace_hook()
+    mod._trace_begin()
+    try:
+        oc.chat_completions_batch([_Item(0)])
+    except RuntimeError:
+        pass
+    assert mod._ITEM_STATE == {}
+
+
+def test_an_unregistered_item_falls_back_to_this_threads_state(monkeypatch):
+    """Not every path goes through the batch entry point; the direct call must
+    keep working."""
+    mod = _worker_module()
+    oc = _fake_surya(monkeypatch, mod)
+    oc._generate_one = lambda *a, **k: _Result()
+    w = _worker(mod, token_trace=True, max_tokens=0)
+    w._install_trace_hook()
+    mod._trace_begin()
+    mod._trace_opts({"max_tokens": 1234})
+    item = _Item(12288)
+    oc._generate_one(item, max_tokens_default=2048)
+    assert item.max_tokens == 1234
+
+
+def test_the_batch_hook_is_reported_when_surya_has_no_such_entry_point(monkeypatch, capsys):
+    """Silence would mean the budgets are quietly unreachable on that surya."""
+    mod = _worker_module()
+    oc = _fake_surya(monkeypatch, mod)
+    if hasattr(oc, "chat_completions_batch"):
+        del oc.chat_completions_batch
+    w = _worker(mod, token_trace=True, max_tokens=3072)
+    w._install_trace_hook()
+    assert "no chat_completions_batch to hook" in capsys.readouterr().err