fix(admission): reserve Responses and Messages bodies before clone (#10814)

Merged — validated together with a batch of related RaviTharuma PRs in one combined worktree (typecheck:core clean, complexity/file-size/changelog gates green, focused tests passing). Thanks for the contribution!

src/app/api/v1/antigravity/route.ts

@@ -1,5 +1,6 @@
 import { handleChat } from "@/sse/handlers/chat";
 import { initTranslators } from "@omniroute/open-sse/translator/index.ts";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 
 let initialized = false;
 
@@ -41,7 +42,9 @@ export async function OPTIONS() {
  * already-registered bidirectional translators. The AgentBridge MITM proxy
  * (`server.cjs`) forwards the IDE's intercepted cloudcode request here.
  */
-export async function POST(request: Request): Promise<Response> {
+async function postHandler(request: Request): Promise<Response> {
   await ensureInitialized();
   return await handleChat(request);
 }
+
+export const POST = withChatAdmission(postHandler);

src/app/api/v1/api/chat/route.ts

@@ -1,6 +1,7 @@
 import { handleChat } from "@/sse/handlers/chat";
 import { initTranslators } from "@omniroute/open-sse/translator/index.ts";
 import { transformToOllama } from "@omniroute/open-sse/utils/ollamaTransform.ts";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 
 let initialized = false;
 
@@ -21,7 +22,7 @@ export async function OPTIONS() {
   });
 }
 
-export async function POST(request) {
+async function postHandler(request) {
   await ensureInitialized();
 
   const clonedReq = request.clone();
@@ -34,3 +35,5 @@ export async function POST(request) {
   const response = await handleChat(request);
   return transformToOllama(response, modelName);
 }
+
+export const POST = withChatAdmission(postHandler);

src/app/api/v1/completions/route.ts

@@ -7,6 +7,7 @@ import {
   readCompressionRequestHeader,
   withCompressionHeaderEcho,
 } from "@/shared/utils/compressionHeaderEcho";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 
 let initPromise = null;
 const injectionGuard = createInjectionGuard();
@@ -41,7 +42,7 @@ export async function OPTIONS() {
  *
  * @see https://platform.openai.com/docs/api-reference/completions
  */
-export async function POST(request: Request) {
+async function postHandler(request: Request) {
   await ensureInitialized();
 
   // #6422 — capture the compression request header once so we can echo it back
@@ -122,3 +123,5 @@ export async function POST(request: Request) {
     compressionRequestHeader
   );
 }
+
+export const POST = withChatAdmission(postHandler);

src/app/api/v1/messages/route.ts

@@ -1,6 +1,7 @@
 import { handleChat } from "@/sse/handlers/chat";
 import { initTranslators } from "@omniroute/open-sse/translator/index.ts";
 import { withInjectionGuard } from "@/middleware/promptInjectionGuard";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 import { requireJsonContentType } from "@/shared/middleware/requireJsonContentType";
 import {
   withEarlyStreamKeepalive,
@@ -78,4 +79,4 @@ async function postHandler(request: any, context: any, preParsedBody: any = null
   return await handleChat(request, null, body);
 }
 
-export const POST = withInjectionGuard(postHandler);
+export const POST = withChatAdmission(withInjectionGuard(postHandler));

src/app/api/v1/providers/[provider]/chat/completions/route.ts

@@ -4,6 +4,7 @@ import { initTranslators } from "@omniroute/open-sse/translator/index.ts";
 import { errorResponse } from "@omniroute/open-sse/utils/error.ts";
 import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts";
 import { getRegistryEntry } from "@omniroute/open-sse/config/providerRegistry.ts";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 
 let initialized = false;
 
@@ -31,7 +32,7 @@ export async function OPTIONS() {
  * Routes to the specified provider, validating model/provider match.
  * Full body format validation is delegated to handleChat.
  */
-export async function POST(request, { params }) {
+async function postHandler(request, { params }) {
   const { provider: rawProvider } = await params;
 
   const providerEntry = getRegistryEntry(rawProvider);
@@ -103,3 +104,5 @@ export async function POST(request, { params }) {
 
   return await handleChat(newRequest, () => buildClientRawRequest(request, rawBody));
 }
+
+export const POST = withChatAdmission(postHandler);

src/app/api/v1/relay/chat/completions/route.ts

@@ -8,6 +8,7 @@
 
 import { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors";
 import { handleChat } from "@/sse/handlers/chat";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 import { createInjectionGuard } from "@/middleware/promptInjectionGuard";
 import { getRelayTokenByHash, checkRateLimit, recordRelayUsage } from "@/lib/db/relayProxies";
 import {
@@ -199,7 +200,7 @@ export async function OPTIONS() {
   return handleCorsOptions();
 }
 
-export async function POST(request: Request) {
+async function postHandler(request: Request) {
   const startTime = Date.now();
   const clientIp = getClientIp(request);
   const userAgent = sanitizeForensicHeader(request.headers.get("user-agent"));
@@ -433,3 +434,5 @@ export async function POST(request: Request) {
     });
   }
 }
+
+export const POST = withChatAdmission(postHandler);

src/app/api/v1/responses/[...path]/route.ts

@@ -1,5 +1,6 @@
 import { handleChat } from "@/sse/handlers/chat";
 import { initTranslators } from "@omniroute/open-sse/translator/index.ts";
+import { withChatAdmission } from "@/shared/middleware/withChatAdmission";
 
 let initialized = false;
 
@@ -25,7 +26,9 @@ export async function OPTIONS() {
  * Reuses the shared chat handler so native Codex passthrough can keep
  * arbitrary Responses suffixes all the way to the upstream provider.
  */
-export async function POST(request) {
+async function postHandler(request) {
   await ensureInitialized();
   return await handleChat(request);
 }
+
+export const POST = withChatAdmission(postHandler);

src/shared/middleware/withChatAdmission.ts

@@ -0,0 +1,48 @@
+/**
+ * Compose process-wide chat admission in front of a route handler.
+ *
+ * Uses the shipped `admitChatRequest` budget/fairness controller — it does not
+ * introduce a second admission path. Call this *outside* `withInjectionGuard`
+ * so a large `/v1/responses` or `/v1/messages` body is reserved (or 503-shed)
+ * before `request.clone()` / `.json()`.
+ */
+import {
+  admitChatRequest,
+  CHAT_ADMISSION_QUEUE_MAX_MS,
+  releaseChatAdmissionAfterHandler,
+  resolveSessionId,
+  type ChatAdmissionController,
+} from "./chatBodyAdmission";
+
+type RouteHandler = (request: Request, ...args: any[]) => Promise<Response> | Response;
+
+export function withChatAdmission(
+  handler: RouteHandler,
+  options: {
+    controller?: ChatAdmissionController;
+    queueMs?: number;
+    largeBodyBytes?: number;
+    hardMaxBytes?: number;
+  } = {}
+): RouteHandler {
+  return async function admittedHandler(request: Request, ...args: any[]) {
+    const sessionId = resolveSessionId(request);
+    const admission = await admitChatRequest(request, {
+      sessionId,
+      queueMs: options.queueMs ?? CHAT_ADMISSION_QUEUE_MAX_MS,
+      controller: options.controller,
+      largeBodyBytes: options.largeBodyBytes,
+      hardMaxBytes: options.hardMaxBytes,
+    });
+    if (admission.admit === false) return admission.response;
+    try {
+      return await releaseChatAdmissionAfterHandler(
+        Promise.resolve(handler(admission.request, ...args)),
+        admission.lease
+      );
+    } catch (error) {
+      admission.lease?.release();
+      throw error;
+    }
+  };
+}

tests/unit/with-chat-admission-10786.test.ts

@@ -0,0 +1,94 @@
+// #10786: process-wide admitChatRequest must run before Responses/Messages clone/parse.
+import test from "node:test";
+import assert from "node:assert/strict";
+
+const { ChatAdmissionController, admitChatRequest } = await import(
+  "../../src/shared/middleware/chatBodyAdmission.ts"
+);
+const { withChatAdmission } = await import("../../src/shared/middleware/withChatAdmission.ts");
+
+function largeBody(n = 64): string {
+  return JSON.stringify({ messages: [{ role: "user", content: "x".repeat(n) }] });
+}
+
+function chatRequest(url: string, body: string): Request {
+  return new Request(url, {
+    method: "POST",
+    headers: { "content-type": "application/json", "content-length": String(body.length) },
+    body,
+  });
+}
+
+test("withChatAdmission does not invoke the handler when a second large body is shed", async () => {
+  const controller = new ChatAdmissionController(1);
+  const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 0 };
+  const body = largeBody();
+
+  const first = await admitChatRequest(chatRequest("http://x/v1/responses", body), options);
+  assert.equal(first.admit, true);
+
+  let called = false;
+  const wrapped = withChatAdmission(async () => {
+    called = true;
+    return new Response("ok");
+  }, options);
+
+  const res = await wrapped(chatRequest("http://x/v1/responses", body));
+  assert.equal(called, false);
+  assert.equal(res.status, 503);
+  assert.equal(res.headers.get("Retry-After"), "2");
+  const json = await res.json();
+  assert.equal(json.error.code, "chat_admission_busy");
+  first.lease?.release();
+});
+
+test("withChatAdmission invokes the handler and forwards the admitted request", async () => {
+  const controller = new ChatAdmissionController(1);
+  const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 0 };
+  const body = largeBody();
+  let seen: Request | null = null;
+  const wrapped = withChatAdmission(async (request: Request) => {
+    seen = request;
+    return new Response("ok", { status: 200 });
+  }, options);
+  const res = await wrapped(chatRequest("http://x/v1/messages", body));
+  assert.ok(seen);
+  assert.equal(res.status, 200);
+  assert.equal(await seen.text(), body);
+});
+
+test("responses admits inline before json; messages wrap withChatAdmission before withInjectionGuard", async () => {
+  const { readFileSync } = await import("node:fs");
+  const responses = readFileSync(new URL("../../src/app/api/v1/responses/route.ts", import.meta.url), "utf8");
+  const messages = readFileSync(new URL("../../src/app/api/v1/messages/route.ts", import.meta.url), "utf8");
+  const catchAll = readFileSync(
+    new URL("../../src/app/api/v1/responses/[...path]/route.ts", import.meta.url),
+    "utf8"
+  );
+  // Keepalive #10806 inlined admitChatRequest into /v1/responses. Wrapping
+  // withChatAdmission on top would double-admit. Body is reserved before json().
+  const admitAt = responses.indexOf("await admitChatRequest(request");
+  const jsonAt = responses.indexOf("parsedBody = await request.json()");
+  assert.ok(admitAt >= 0, "responses route must call admitChatRequest");
+  assert.ok(jsonAt > admitAt, "admitChatRequest must run before request.json()");
+  assert.doesNotMatch(responses, /withChatAdmission/);
+  assert.match(messages, /withChatAdmission\(\s*withInjectionGuard\(postHandler\)\s*\)/);
+  assert.match(catchAll, /export const POST = withChatAdmission\(postHandler\)/);
+  assert.doesNotMatch(catchAll, /export async function POST/);
+});
+
+test("remaining handleChat aliases wrap POST with withChatAdmission (#10790)", async () => {
+  const { readFileSync } = await import("node:fs");
+  const files = [
+    "src/app/api/v1/antigravity/route.ts",
+    "src/app/api/v1/api/chat/route.ts",
+    "src/app/api/v1/completions/route.ts",
+    "src/app/api/v1/providers/[provider]/chat/completions/route.ts",
+    "src/app/api/v1/relay/chat/completions/route.ts",
+  ];
+  for (const rel of files) {
+    const src = readFileSync(new URL("../../" + rel, import.meta.url), "utf8");
+    assert.match(src, /withChatAdmission/, rel);
+    assert.match(src, /export const POST = withChatAdmission\(postHandler\)/, rel);
+  }
+});