From 9eddafff60c555cecc25f907b2f1accbddb7fe90 Mon Sep 17 00:00:00 2001 From: Ravi Tharuma <25951435+RaviTharuma@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:48:14 +0200 Subject: [PATCH] fix(admission): reserve Responses and Messages bodies before clone (#10814) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 5 +- src/app/api/v1/api/chat/route.ts | 5 +- src/app/api/v1/completions/route.ts | 5 +- src/app/api/v1/messages/route.ts | 3 +- .../[provider]/chat/completions/route.ts | 5 +- .../api/v1/relay/chat/completions/route.ts | 5 +- src/app/api/v1/responses/[...path]/route.ts | 5 +- src/shared/middleware/withChatAdmission.ts | 48 ++++++++++ tests/unit/with-chat-admission-10786.test.ts | 94 +++++++++++++++++++ 9 files changed, 168 insertions(+), 7 deletions(-) create mode 100644 src/shared/middleware/withChatAdmission.ts create mode 100644 tests/unit/with-chat-admission-10786.test.ts diff --git a/src/app/api/v1/antigravity/route.ts b/src/app/api/v1/antigravity/route.ts index 89f12036ae..aba16415ee 100644 --- a/src/app/api/v1/antigravity/route.ts +++ b/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 { +async function postHandler(request: Request): Promise { await ensureInitialized(); return await handleChat(request); } + +export const POST = withChatAdmission(postHandler); diff --git a/src/app/api/v1/api/chat/route.ts b/src/app/api/v1/api/chat/route.ts index c531911121..d07a7509a0 100644 --- a/src/app/api/v1/api/chat/route.ts +++ b/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); diff --git a/src/app/api/v1/completions/route.ts b/src/app/api/v1/completions/route.ts index f11f7cb904..c2106271f1 100644 --- a/src/app/api/v1/completions/route.ts +++ b/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); diff --git a/src/app/api/v1/messages/route.ts b/src/app/api/v1/messages/route.ts index cbbc87f202..97f6af1fb7 100644 --- a/src/app/api/v1/messages/route.ts +++ b/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)); diff --git a/src/app/api/v1/providers/[provider]/chat/completions/route.ts b/src/app/api/v1/providers/[provider]/chat/completions/route.ts index f064163b9e..1149ef8933 100644 --- a/src/app/api/v1/providers/[provider]/chat/completions/route.ts +++ b/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); diff --git a/src/app/api/v1/relay/chat/completions/route.ts b/src/app/api/v1/relay/chat/completions/route.ts index 12ff30bea4..28cc5160db 100644 --- a/src/app/api/v1/relay/chat/completions/route.ts +++ b/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); diff --git a/src/app/api/v1/responses/[...path]/route.ts b/src/app/api/v1/responses/[...path]/route.ts index e2f7062f5a..12eb4e05b3 100644 --- a/src/app/api/v1/responses/[...path]/route.ts +++ b/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); diff --git a/src/shared/middleware/withChatAdmission.ts b/src/shared/middleware/withChatAdmission.ts new file mode 100644 index 0000000000..7a762b6f9e --- /dev/null +++ b/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; + +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; + } + }; +} diff --git a/tests/unit/with-chat-admission-10786.test.ts b/tests/unit/with-chat-admission-10786.test.ts new file mode 100644 index 0000000000..409d1c8a8e --- /dev/null +++ b/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); + } +});