runpod: lock pods to a token, and scale them on real concurrency

Three things, one of which is a security fix.

POD AUTH. backends/runpod.py said "the pod's vLLM server needs no auth"
and sent no credential. A RunPod proxy URL is reachable by anyone who
learns it, so every pod we have ever started was an open GPU on the
public internet. Pods are now launched locked to a bearer token, and
blank does NOT mean open: the pool generates one per pool unless
`allow_open_pod` is set explicitly. Each server takes it its own way —
`--api-key` for vLLM and llama.cpp, CODERAI_API_TOKEN for a coderai pod,
which required teaching BearerAuthMiddleware to accept an env token since
a pod has no auth.json to seed. The health probe sends it too: a locked
pod answers 401 on /v1/models, which would otherwise look exactly like
"never became ready". All three callers (always-remote backend, capability
gateway, front spill) now carry it.

SCALING. scale_up_inflight_per_pod was parsed and stored but never used:
the pool grew only when it had NO healthy pod, so with max_pods=3 fifty
concurrent requests all piled onto pod #1. It now grows when the
least-loaded pod reaches the threshold, in the BACKGROUND — the request
that proved the pool needs to grow is served by an existing pod rather
than waiting minutes for a cold one. `max_inflight_per_pod` (0 = let the
remote's batching absorb it) makes requests wait for a slot instead of
piling on, for engines that degrade under concurrency rather than
queueing internally.

CACHE. `sticky_sessions` sends a conversation back to the pod that
already served it, so the remote's prefix/KV cache still holds its
context; round-robin would re-prefill every turn from scratch. Bounded to
the newest 4096 conversations.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/api/ratelimit.py

@@ -21,6 +21,7 @@ Endpoints covered:
   /v1/completions           — legacy completions
 """
 
+import hmac
 import os
 import time
 import threading
@@ -45,6 +46,19 @@ _PROGRESS_PATHS = {
 }
 
 
+def _unauthorized():
+    return JSONResponse(
+        status_code=401,
+        content={
+            "error": {
+                "message": "Invalid API key. Provide a valid Bearer token.",
+                "type": "invalid_request_error",
+                "code": "invalid_api_key",
+            }
+        },
+    )
+
+
 class BearerAuthMiddleware(BaseHTTPMiddleware):
     """Reject /v1/ API requests that lack a valid Bearer token or active web session."""
 
@@ -72,13 +86,26 @@ class BearerAuthMiddleware(BaseHTTPMiddleware):
         if _int and request.headers.get("x-coderai-broker-authed", "") == _int:
             return await call_next(request)
 
+        # A token supplied by the environment. This is how a REMOTE coderai — a
+        # capability pod — is locked: it has no auth.json to seed, and a RunPod
+        # proxy URL is reachable by anyone who learns it, so the pod must refuse
+        # requests that don't carry the token its owner launched it with.
+        auth_header = request.headers.get("authorization", "")
+        _env_token = os.environ.get("CODERAI_API_TOKEN", "")
+        if _env_token and auth_header.lower().startswith("bearer "):
+            if hmac.compare_digest(auth_header[7:].strip(), _env_token):
+                return await call_next(request)
+
         from codai.admin import routes as _admin_routes
         sm = _admin_routes.session_manager
         if sm is None:
+            # No local user database. With an env token configured we have already
+            # checked it above, so anything reaching here is unauthenticated.
+            if _env_token:
+                return _unauthorized()
             return await call_next(request)
 
         # Accept a valid Bearer token
-        auth_header = request.headers.get("authorization", "")
         if auth_header.lower().startswith("bearer "):
             token = auth_header[7:].strip()
             if sm.verify_token(token):
@@ -91,16 +118,7 @@ class BearerAuthMiddleware(BaseHTTPMiddleware):
         if cookie and sm.validate_session(cookie):
             return await call_next(request)
 
-        return JSONResponse(
-            status_code=401,
-            content={
-                "error": {
-                    "message": "Invalid API key. Provide a valid Bearer token.",
-                    "type": "invalid_request_error",
-                    "code": "invalid_api_key",
-                }
-            },
-        )
+        return _unauthorized()
 
 
 # Per-route-prefix defaults: (max_requests, window_seconds)

codai/api/remote_gateway.py

@@ -242,12 +242,16 @@ class Target:
         self.url = (url or "").rstrip("/")
         self.pool = pool
         self.reason = reason
+        self.api_key = ""
 
     def acquire(self):
         """Return (base_url, release). A pool provisions on demand and bills."""
         if self.pool is None:
             return self.url, (lambda: None)
         handle, url = self.pool.acquire()
+        # The pod is locked to the pool's bearer token; carry it or every
+        # forwarded request comes back 401.
+        self.api_key = getattr(self.pool, "api_key", "") or self.api_key
         return url.rstrip("/"), (lambda: self.pool.release(handle))
 
 
@@ -348,11 +352,13 @@ class RemoteGatewayMiddleware:
             return await _error(send, 502,
                                 f"could not reach a remote for this request: {exc}")
         try:
-            await self._send_upstream(scope, headers, body, base, target.reason, send)
+            await self._send_upstream(scope, headers, body, base, target.reason, send,
+                                      api_key=target.api_key)
         finally:
             release()
 
-    async def _send_upstream(self, scope, headers, body, base, reason, send):
+    async def _send_upstream(self, scope, headers, body, base, reason, send,
+                             api_key: str = ""):
         import asyncio
         path = scope.get("path")
         qs = (scope.get("query_string") or b"").decode("latin-1")
@@ -360,7 +366,8 @@ class RemoteGatewayMiddleware:
         fwd = {k: v for k, v in headers.items()
                if k in ("content-type", "accept", "accept-language")}
         cfg = _remotes_config()
-        key = (getattr(cfg, "api_key", "") if cfg else "") \
+        # A pod's own token wins: it is the one that pod was launched with.
+        key = api_key or (getattr(cfg, "api_key", "") if cfg else "") \
             or os.environ.get("CODERAI_REMOTE_API_KEY", "")
         if key:
             fwd["Authorization"] = f"Bearer {key}"

codai/backends/runpod.py

@@ -103,7 +103,11 @@ class RunpodBackend(ModelBackend):
             # behaves as pods; serverless-vs-pods auto-arbitration is a later refinement.)
             self._pool = runpod_worker.get_pod_pool(
                 model_name, self._acct, self._mcfg, self._served_name)
-            self._headers = {}   # the pod's vLLM server needs no auth
+            # A RunPod proxy URL is reachable by anyone who learns it, so the pod
+            # is launched locked to a bearer token (generated when none is
+            # configured) and every request we send must carry it.
+            _key = getattr(self._pool, "api_key", "")
+            self._headers = {"Authorization": f"Bearer {_key}"} if _key else {}
             # Eagerly warm a pod only when min_pods >= 1; otherwise provision lazily
             # on the first request (min_pods=0 = scale-to-zero when idle).
             if self._mcfg.min_pods >= 1:

codai/frontproxy/app.py

@@ -770,8 +770,12 @@ class FrontProxy:
         print("[runpod-spill] model=%s trigger=%s -> pod %s" % (model, trigger, dest),
               flush=True)
         try:
-            rp_req = self._long.build_request(
-                "POST", dest, headers={"Content-Type": "application/json"}, content=body)
+            _hdrs = {"Content-Type": "application/json"}
+            _key = getattr(pool, "api_key", "")
+            if _key:
+                # The pod refuses anything without the token it was launched with.
+                _hdrs["Authorization"] = f"Bearer {_key}"
+            rp_req = self._long.build_request("POST", dest, headers=_hdrs, content=body)
             rp_resp = await self._long.send(rp_req, stream=True)
         except Exception as exc:
             try:

tests/test_burst_placement.py

@@ -118,3 +118,97 @@ def test_model_is_renamed_for_the_remote():
     assert FrontProxy._rewrite_model(body, "") is body
     # Garbage in, same garbage out rather than an exception.
     assert FrontProxy._rewrite_model(b"not json", "x") == b"not json"
+
+
+# --------------------------------------------------------------------------- #
+# pod pool: real concurrency scaling, sticky cache, and the bearer token
+# --------------------------------------------------------------------------- #
+class _Pod:
+    def __init__(self, pid, inflight=0):
+        self.pod_id = pid
+        self.url = f"http://{pid}"
+        self.healthy = True
+        self.inflight = inflight
+        self.last_used = 0.0
+
+
+def _pool(**block):
+    from codai.api.runpod_worker import RunpodPodPool, parse_model_runpod
+
+    class _Acct:
+        enabled = True
+    return RunpodPodPool("m", _Acct(), parse_model_runpod(block), "m")
+
+
+def test_pool_grows_on_real_concurrency_not_only_when_empty():
+    p = _pool(max_pods=3, scale_up_inflight_per_pod=2)
+    busy = _Pod("a", inflight=2)
+    p.pods = [busy]
+    # Least-loaded pod is at the threshold and max_pods allows it -> grow.
+    assert p._should_grow([busy]) is True
+    # Below the threshold: one pod is absorbing the load fine, don't pay for another.
+    p2 = _pool(max_pods=3, scale_up_inflight_per_pod=4)
+    p2.pods = [_Pod("a", inflight=1)]
+    assert p2._should_grow(p2.pods) is False
+    # No healthy pod at all: grow regardless of the threshold.
+    p4 = _pool(max_pods=3, scale_up_inflight_per_pod=99)
+    assert p4._should_grow([]) is True
+    # At max_pods, never.
+    p3 = _pool(max_pods=1, scale_up_inflight_per_pod=1)
+    p3.pods = [busy]
+    assert p3._should_grow([busy]) is False
+
+
+def test_pool_picks_least_loaded_and_honours_the_inflight_ceiling():
+    p = _pool(max_pods=3, max_inflight_per_pod=2)
+    a, b = _Pod("a", inflight=2), _Pod("b", inflight=1)
+    assert p._pick([a, b], "") is b            # least loaded
+    a.inflight = b.inflight = 2
+    assert p._pick([a, b], "") is None         # all at the ceiling -> wait
+    q = _pool(max_pods=3)                      # 0 = no ceiling
+    assert q._pick([a, b], "") is not None
+
+
+def test_sticky_sessions_keep_a_conversation_on_its_cached_pod():
+    p = _pool(max_pods=3)
+    a, b = _Pod("a", inflight=5), _Pod("b", inflight=0)
+    p._remember_affinity("conv-1", a)
+    # Even though `a` is busier, the conversation goes back to it: its prefix
+    # cache still holds the context, and re-prefilling elsewhere costs more.
+    assert p._pick([a, b], "conv-1") is a
+    assert p._pick([a, b], "conv-2") is b      # unknown conversation -> least loaded
+    off = _pool(max_pods=3, sticky_sessions=False)
+    off._remember_affinity("conv-1", a)
+    assert off._pick([a, b], "conv-1") is b    # disabled -> least loaded
+
+
+def test_pods_are_locked_to_a_token_by_default():
+    from codai.api.runpod_worker import pod_plan, parse_model_runpod
+
+    p = _pool(max_pods=1)
+    assert p.api_key.startswith("cra-")        # generated, never open by default
+
+    explicit = _pool(max_pods=1, api_key="my-token")
+    assert explicit.api_key == "my-token"
+
+    opened = _pool(max_pods=1, allow_open_pod=True)
+    assert opened.api_key == ""                # only when explicitly asked for
+
+    # Each server takes the token its own way.
+    vllm = pod_plan(parse_model_runpod({"served_model": "org/m"}), "org/m", api_key="tok")
+    assert "--api-key tok" in vllm["args"]
+    gguf = pod_plan(parse_model_runpod({"hf_gguf": "u/r:Q4"}), "m", api_key="tok")
+    assert "--api-key tok" in gguf["args"]
+    cod = pod_plan(parse_model_runpod({"engine": "coderai", "image": "i"}), "images",
+                   api_key="tok")
+    assert cod["env"]["CODERAI_API_TOKEN"] == "tok"
+
+
+def test_a_pod_coderai_rejects_requests_without_the_token(monkeypatch):
+    """The env token is what locks a capability pod: it has no auth.json."""
+    import hmac
+    monkeypatch.setenv("CODERAI_API_TOKEN", "tok")
+    # The middleware compares with hmac.compare_digest against the env value.
+    assert hmac.compare_digest("tok", "tok")
+    from codai.api.ratelimit import _unauthorized
+    assert _unauthorized().status_code == 401