Serve Crello corpus aafdb0e86816 from the mounted bucket

server/router.py

@@ -90,7 +90,7 @@ class Router:
     async def http_proxy(self, request: Request) -> Response:
         body = await request.body()
         worker = self.for_path(request.url.path)
-        creating = False
+        closing = None
         if worker is None and request.url.path.rstrip("/") == "/mcp" and request.method == "POST":
             try:
                 rpc = json.loads(body or b"{}")
@@ -99,9 +99,10 @@ class Router:
             sid = (rpc.get("params") or {}).get("session_id") if isinstance(rpc, dict) else None
             if sid and sid in self.sessions:
                 worker = self.sessions[sid]
+                if rpc.get("method") == "openenv/session/close":
+                    closing = sid
             elif isinstance(rpc, dict) and rpc.get("method") == "openenv/session/create":
-                worker, creating = self.least_busy(), True
-                self.load[worker] += 1
+                return await self.create_session(request, body)
         if worker is None:
             worker = next(self.rr)
         url = self.backends[worker] + request.url.path + (f"?{request.url.query}" if request.url.query else "")
@@ -110,21 +111,13 @@ class Router:
             req = self.http.build_request(request.method, url, headers=headers, content=body)
             resp = await self.http.send(req, stream=True)
         except httpx.HTTPError as e:
-            if creating:
-                self.load[worker] -= 1
             return JSONResponse({"error": f"worker {worker} unreachable: {e}"}, status_code=502)
         out_headers = {k: v for k, v in resp.headers.items() if k.lower() not in HOP}
-        if creating or request.url.path.startswith("/api/episodes"):
+        if closing:
+            self.sessions.pop(closing, None)
+        if request.url.path.startswith("/api/episodes"):
             data = await resp.aread()
             await resp.aclose()
-            if creating:
-                self.load[worker] -= 1
-                try:
-                    sid = json.loads(data).get("result", {}).get("session_id")
-                except ValueError:
-                    sid = None
-                if sid:
-                    self.sessions[sid] = worker
             self.learn(worker, data)
             return Response(data, status_code=resp.status_code, headers=out_headers)
 
@@ -137,6 +130,40 @@ class Router:
 
         return StreamingResponse(stream(), status_code=resp.status_code, headers=out_headers)
 
+    async def create_session(self, request: Request, body: bytes) -> Response:
+        """A new OpenEnv session: the least busy worker first, the next one if it is full (each worker caps its
+        sessions at MAX_CONCURRENT_ENVS), and the session's worker remembered for /mcp and /ws."""
+        headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP}
+        order = sorted(range(len(self.backends)), key=lambda i: (self.load[i], i))
+        first = self.least_busy()
+        order.remove(first)
+        data, status, out_headers = b"", 502, {}
+        for worker in [first, *order]:
+            self.load[worker] += 1
+            try:
+                resp = await self.http.post(self.backends[worker] + "/mcp", headers=headers, content=body)
+            except httpx.HTTPError as e:
+                data, status = json.dumps({"error": f"worker {worker} unreachable: {e}"}).encode(), 502
+                continue
+            finally:
+                self.load[worker] -= 1
+            data, status = resp.content, resp.status_code
+            out_headers = {k: v for k, v in resp.headers.items() if k.lower() not in HOP}
+            try:
+                reply = json.loads(data)
+            except ValueError:
+                reply = {}
+            sid = (reply.get("result") or {}).get("session_id") if isinstance(reply, dict) else None
+            if sid:
+                self.sessions[sid] = worker
+                if len(self.sessions) > 200_000:  # sessions closed without telling us, oldest first
+                    for k in list(self.sessions)[:50_000]:
+                        del self.sessions[k]
+                break
+            if "capacity" not in json.dumps(reply.get("error", "")).lower():
+                break  # a real error, not a full worker: pass it on
+        return Response(data, status_code=status, headers=out_headers)
+
     # ------------------------------------------------------------------ WebSocket
 
     async def ws_proxy(self, ws: WebSocket) -> None: