gateway: ask a pod for the engine it can serve, in the request itself

The pod was launched with docTR as its default and still answered 'OCR engine
surya is not enabled' — because the request carried engine=surya, and the OCR
route honours the request's own field over the configured default. Setting the
pod's env changed the default; it did not change what the request asked for.

A pod target for OCR now rewrites the request's engine field on the way out,
to the engine the pod will actually serve (pod_ocr_engine mirrors the launch
env). Multipart is edited on the bytes with every other part untouched and
Content-Length corrected; verified through FastAPI's own parser with the real
probe body, file intact.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/api/remote_gateway.py

@@ -300,6 +300,35 @@ def _model_from_body(body: bytes, content_type: str, field: str = "model") -> st
 _MODEL_FIELD = {"/v1/ocr": "engine", "/v1/ocr/batch": "engine"}
 
 
+def _rewrite_request_field(body: bytes, headers: dict, field: str, value: str):
+    """Replace one field's value in a JSON or multipart body. (body, headers).
+
+    Multipart is edited in place on the bytes — the boundary and every other
+    part are untouched — and Content-Length is corrected, because the length
+    of the value changed and a stale length truncates or hangs the upload.
+    """
+    ctype = (headers.get("content-type") or "").lower()
+    if "json" in ctype or (body[:1] in (b"{", b"[")):
+        try:
+            obj = json.loads(body)
+            if isinstance(obj, dict):
+                obj[field] = value
+                body = json.dumps(obj).encode()
+        except Exception:
+            return body, headers
+    elif "multipart/form-data" in ctype:
+        pat = re.compile(rb'(name="' + field.encode() +
+                         rb'"\r?\n(?:[^\r\n]*\r?\n)*?\r?\n)([^\r\n]*)', re.IGNORECASE)
+        body, n = pat.subn(lambda m: m.group(1) + value.encode(), body, count=1)
+        if not n:
+            return body, headers
+    else:
+        return body, headers
+    headers = dict(headers)
+    headers["content-length"] = str(len(body))
+    return body, headers
+
+
 class Target:
     """Where a request is going: a fixed URL, or a pod pool to borrow one from."""
 
@@ -314,6 +343,13 @@ class Target:
         #: The model entry whose LOCAL text adapters must be pushed before the
         #: pod loads the model. Only a coderai pod can receive them.
         self.lora_entry = None
+        #: (field, value): rewrite one request field before forwarding. A pod
+        #: told to serve OCR with docTR still received `engine=surya` in the
+        #: request itself, and the OCR route honours the request's field over
+        #: its configured default — so the pod's env said docTR and the pod
+        #: still answered "surya is not enabled". The env sets the default; this
+        #: sets what the request asks for.
+        self.rewrite_field = None
 
     def acquire(self):
         """Return (base_url, release). A pool provisions on demand and bills."""
@@ -385,9 +421,13 @@ def resolve_target(path: str, method: str, query: str, body: bytes,
     if where == "pod":
         entry, block = detail
         from codai.api.runpod_worker import (get_model_pod_pool, parse_model_runpod,
-                                             resolve_model_source)
+                                             resolve_model_source, pod_ocr_engine)
         target = Target(pool=get_model_pod_pool(model, entry, block),
                         reason=f"model {model!r} on its own RunPod pod")
+        if cap == "ocr":
+            served = pod_ocr_engine(model)
+            if served and served != model.strip().lower():
+                target.rewrite_field = (_MODEL_FIELD.get(path, "model"), served)
         mcfg = parse_model_runpod(block)
         kind, _ = resolve_model_source(entry, mcfg)
         if kind == "upload":
@@ -501,6 +541,9 @@ class RemoteGatewayMiddleware:
             return await _error(send, 502,
                                 f"could not reach a remote for this request: {exc}")
         try:
+            if target.rewrite_field:
+                body, headers = _rewrite_request_field(
+                    body, headers, *target.rewrite_field)
             # Weights the remote cannot fetch itself, when asked to upload them.
             if target.upload_entry:
                 await asyncio.to_thread(ensure_model_uploaded, base, target.api_key,

codai/api/runpod_worker.py

@@ -689,6 +689,20 @@ def volume_env(mount: str) -> dict:
     }
 
 
+def pod_ocr_engine(requested: str) -> str:
+    """The OCR engine a pod will actually serve for a request naming ``requested``.
+
+    Mirrors the env a pod is launched with: surya needs a VLM server the image
+    does not carry, so a pod serves docTR for it. The gateway uses this to
+    rewrite the request's own `engine` field, because the pod honours what the
+    request asks for over its configured default.
+    """
+    r = (requested or "").strip().lower()
+    if r == "surya":
+        return "doctr"
+    return r if r in ("paddle", "doctr") else "doctr"
+
+
 def _surya_accepted() -> bool:
     """Whether this deployment has accepted Surya's licence."""
     try:

tests/test_remote_execution.py

@@ -1855,3 +1855,31 @@ def test_a_plain_service_url_carries_its_token(monkeypatch):
     t = rg.resolve_target("/v1/images/generations", "POST", "",
                           b'{"model":"m"}', "application/json")
     assert t is not None and t.url == "http://gpubox:8000" and t.api_key == "tok-123"
+
+
+def test_a_pod_that_cannot_serve_surya_is_asked_for_doctr_in_the_request_itself():
+    """The pod's env said docTR and the pod still answered 'surya is not
+    enabled', because the request carried engine=surya and the OCR route
+    honours the request's field over its configured default. The env sets the
+    default; the gateway has to change what the request ASKS for."""
+    from codai.api.remote_gateway import _rewrite_request_field, _model_from_body
+    from codai.api.model_test import _multipart
+    from codai.api.runpod_worker import pod_ocr_engine
+
+    assert pod_ocr_engine("surya") == "doctr"
+    assert pod_ocr_engine("paddle") == "paddle"
+    assert pod_ocr_engine("doctr") == "doctr"
+
+    ct, body = _multipart({"engine": "surya"}, "p.png", b"\x89PNG", content_type="image/png")
+    new, hdrs = _rewrite_request_field(body, {"content-type": ct}, "engine", "doctr")
+    assert _model_from_body(new, ct, "engine") == "doctr"
+    assert hdrs["content-length"] == str(len(new))
+    assert b"\x89PNG" in new and b'filename="p.png"' in new   # the file is untouched
+
+    j, _ = _rewrite_request_field(b'{"engine":"surya"}', {"content-type": "application/json"},
+                                  "engine", "doctr")
+    assert json.loads(j)["engine"] == "doctr"
+
+    # A field that is not there is left alone, with the body byte-identical.
+    same, h2 = _rewrite_request_field(body, {"content-type": ct}, "model", "x")
+    assert same == body and "content-length" not in h2