diff --git a/src/app/api/v1/chat/completions/route.ts b/src/app/api/v1/chat/completions/route.ts index 4ab136d6d1..defd2fc38f 100644 --- a/src/app/api/v1/chat/completions/route.ts +++ b/src/app/api/v1/chat/completions/route.ts @@ -17,6 +17,7 @@ import { resolveKeepaliveThreshold } from "@omniroute/open-sse/utils/keepaliveTh import { admitChatRequest, admitChatStructure, + CHAT_ADMISSION_QUEUE_MAX_MS, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, resolveSessionId, @@ -101,7 +102,10 @@ export async function POST(request) { // BEFORE JSON parsing. Missing or dishonest Content-Length values cannot bypass // the actual-byte limit. Capacity exhaustion is retryable rather than process-fatal. const sessionId = resolveSessionId(request); - const admissionResult = await admitChatRequest(request, { sessionId }); + const admissionResult = await admitChatRequest(request, { + sessionId, + queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, + }); if (admissionResult.admit === false) return admissionResult.response; const admission = admissionResult; request = admission.request; @@ -144,8 +148,9 @@ export async function POST(request) { } } - const structuralAdmission = admitChatStructure(parsedBody, admission.lease, { + const structuralAdmission = await admitChatStructure(parsedBody, admission.lease, { sessionId, + queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, }); if (structuralAdmission.admit === false) { admission.lease?.release(); diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 37e554a13d..46f7b32f22 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -26,6 +26,11 @@ function parsePositiveInt(value: string | undefined, fallback: number): number { return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : fallback; } +function parseNonNegativeInt(value: string | undefined, fallback: number): number { + const parsed = Number.parseInt(String(value), 10); + return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : fallback; +} + export const CHAT_LARGE_BODY_BYTES = parsePositiveInt( process.env.OMNIROUTE_CHAT_LARGE_BODY_BYTES, 256 * 1024 @@ -41,6 +46,18 @@ export const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt( 1 ); +/** + * How long a heavy request waits for heavyweight capacity before giving up with a + * retryable 503. Agent loops (OpenCode, Claude Code, Cursor…) fan out sub-requests + * that routinely land on the admission gate together; an immediate 503 makes the + * client burn its retry budget in seconds and the agent dies mid-task. A short + * bounded wait serializes the burst instead. `0` (legacy) rejects immediately. + */ +export const CHAT_ADMISSION_QUEUE_MAX_MS = parseNonNegativeInt( + process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS, + 5000 +); + export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt( process.env.OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT, 200 @@ -82,10 +99,13 @@ export interface ChatAdmissionLease { /** * Process-local heavyweight reservation. The capacity check and increment execute in one * synchronous JavaScript turn, making acquisition atomic within an OmniRoute process. - * Queueing is intentionally separate: unavailable capacity is a retryable 503. + * Unavailable capacity is a bounded wait (see `acquireHeavyWithin`) and only then a + * retryable 503, so short agent bursts serialize instead of killing the client's + * retry budget. */ export class ChatAdmissionController { #activeHeavy = 0; + #waiters: Array<() => void> = []; constructor(readonly maxHeavyInFlight = 1) { if (!Number.isSafeInteger(maxHeavyInFlight) || maxHeavyInFlight < 1) { @@ -109,9 +129,40 @@ export class ChatAdmissionController { if (released) return; released = true; this.#activeHeavy = Math.max(0, this.#activeHeavy - 1); + this.#waiters.shift()?.(); }, }; } + + /** + * Wait up to `timeoutMs` for heavyweight capacity, retrying atomically on each + * release. Resolves `null` when the deadline expires with no capacity freed, in + * which case the caller answers the retryable 503. `timeoutMs <= 0` is the + * legacy immediate-reject path. Waiters are served FIFO. + */ + async acquireHeavyWithin(timeoutMs: number): Promise { + const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs)); + for (;;) { + const lease = this.tryAcquireHeavy(); + if (lease) return lease; + const remaining = deadline - Date.now(); + if (remaining <= 0) return null; + let resolver: (() => void) | null = null; + const released = new Promise((resolve) => { + resolver = () => resolve(); + this.#waiters.push(resolver); + }); + const timedOut = await Promise.race([ + released.then(() => false), + new Promise((resolve) => setTimeout(() => resolve(true), remaining)), + ]); + if (resolver) { + const index = this.#waiters.indexOf(resolver); + if (index >= 0) this.#waiters.splice(index, 1); + } + if (timedOut) return null; + } + } } const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); @@ -355,7 +406,7 @@ function estimateStructureTokens(value: unknown, limit: number): TokenEstimate { return { tokens, exhausted: stack.length > 0 && tokens < limit }; } -export function admitChatStructure( +export async function admitChatStructure( body: unknown, lease: ChatAdmissionLease | null, options: { @@ -365,8 +416,9 @@ export function admitChatStructure( heavyMessages?: number; heavyTools?: number; heavyTokens?: number; + queueMs?: number; } = {} -): ChatStructureAdmission { +): Promise { if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease }; const record = body as Record; @@ -402,7 +454,7 @@ export function admitChatStructure( (options.sessionId ? perConnectionAdmissionController.getController(options.sessionId) : defaultAdmissionController); - const acquired = controller.tryAcquireHeavy(); + const acquired = await controller.acquireHeavyWithin(options.queueMs ?? 0); return acquired ? { admit: true, lease: acquired } : { admit: false, response: structuralRejectionResponse(503, maxMessages) }; @@ -515,6 +567,7 @@ export async function admitChatRequest( sessionId?: string; largeBodyBytes?: number; hardMaxBytes?: number; + queueMs?: number; } = {} ): Promise { const sessionId = options.sessionId ?? resolveSessionId(request); @@ -522,6 +575,7 @@ export async function admitChatRequest( options.controller ?? perConnectionAdmissionController.getController(sessionId); const largeBodyBytes = options.largeBodyBytes ?? CHAT_LARGE_BODY_BYTES; const hardMaxBytes = options.hardMaxBytes ?? CHAT_HARD_MAX_BODY_BYTES; + const queueMs = options.queueMs ?? 0; const internalBypass = isInternalAdmissionBypass(request); const contentLength = parseContentLength(request.headers.get("content-length")); @@ -567,15 +621,15 @@ export async function admitChatRequest( } let lease: ChatAdmissionLease | null = null; - const reserve = (): boolean => { + const reserve = async (): Promise => { if (lease) return true; - lease = controller.tryAcquireHeavy(); + lease = await controller.acquireHeavyWithin(queueMs); return lease !== null; }; // A known-large declaration can reserve before ingestion. Unknown lengths are boundedly // sniffed below; this avoids consuming scarce heavyweight capacity for small chunked bodies. - if (contentLength !== null && contentLength >= largeBodyBytes && !reserve()) { + if (contentLength !== null && contentLength >= largeBodyBytes && !(await reserve())) { return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } @@ -594,7 +648,7 @@ export async function admitChatRequest( lease?.release(); return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; } - if (totalBytes >= largeBodyBytes && !reserve()) { + if (totalBytes >= largeBodyBytes && !(await reserve())) { await reader.cancel("chat admission capacity unavailable").catch(() => undefined); return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } diff --git a/tests/unit/per-connection-admission-9654.test.ts b/tests/unit/per-connection-admission-9654.test.ts index 2b52716240..8f26b88129 100644 --- a/tests/unit/per-connection-admission-9654.test.ts +++ b/tests/unit/per-connection-admission-9654.test.ts @@ -149,13 +149,13 @@ test("admitChatRequest with explicit controller overrides per-connection lookup" if (result.admit) result.lease?.release(); }); -test("admitChatStructure routes structural rejection to per-connection controller", () => { +test("admitChatStructure routes structural rejection to per-connection controller", async () => { // occupy sess-a's per-connection controller via the module-level instance const controller = perConnectionAdmissionController.getController("sess-a"); const occupied = controller.tryAcquireHeavy(); assert.ok(occupied); - const result = admitChatStructure( + const result = await admitChatStructure( { messages: Array.from({ length: 3 }, () => ({ role: "user", content: "x" })), }, @@ -176,14 +176,14 @@ test("admitChatStructure routes structural rejection to per-connection controlle occupied.release(); }); -test("admitChatStructure with different sessionId gets independent capacity", () => { +test("admitChatStructure with different sessionId gets independent capacity", async () => { // occupy sess-a's per-connection controller const ctrlA = perConnectionAdmissionController.getController("sess-a"); const occupied = ctrlA.tryAcquireHeavy(); assert.ok(occupied); // Session B should get its own controller → admitted - const result = admitChatStructure( + const result = await admitChatStructure( { messages: Array.from({ length: 500 }, () => ({ role: "user", content: "x" })), },