From 70511b728257d65c63e307173bf424b6dec7e211 Mon Sep 17 00:00:00 2001 From: Brandon Bennett Date: Sun, 9 Aug 2026 16:26:52 -0700 Subject: [PATCH] feat(admission): cancel queue-wait on client abort (#9654) U2 from KC plan 2026-08-09-001. Thread the request AbortSignal through acquireHeavyWithin so a disconnected client stops parking in the FIFO for the full queueMs. - acquireHeavyWithin(timeoutMs, signal?): on abort the waiter is removed from the FIFO immediately and the promise resolves null early; pre-aborted signals never park; the deadline timer is cleared when abort/release wins the race - admitChatRequest reserve() passes request.signal; admitChatStructure gains options.signal; the route threads request.signal - 5 exact-assertion tests (settle-early, pre-aborted, byte-heavy, structural, FIFO-preservation): 119/119 across the 7-file suite --- src/app/api/v1/chat/completions/route.ts | 1 + src/shared/middleware/chatBodyAdmission.ts | 43 ++++- tests/unit/chat-body-admission.test.ts | 186 +++++++++++++++++++++ 3 files changed, 224 insertions(+), 6 deletions(-) diff --git a/src/app/api/v1/chat/completions/route.ts b/src/app/api/v1/chat/completions/route.ts index defd2fc38f..7160fdf5c2 100644 --- a/src/app/api/v1/chat/completions/route.ts +++ b/src/app/api/v1/chat/completions/route.ts @@ -151,6 +151,7 @@ export async function POST(request) { const structuralAdmission = await admitChatStructure(parsedBody, admission.lease, { sessionId, queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, + signal: request.signal, }); if (structuralAdmission.admit === false) { admission.lease?.release(); diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 46f7b32f22..ea14097038 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -139,10 +139,21 @@ export class ChatAdmissionController { * 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. + * + * When `signal` aborts while parked (client disconnect), the waiter is removed + * from the FIFO immediately and the promise resolves `null` early instead of + * parking for the full `timeoutMs` — the caller's 503 is dropped on the dead + * connection, so no capacity is consumed and the freed slot never wakes a + * waiter the client no longer needs. A signal that is already aborted never + * parks at all. */ - async acquireHeavyWithin(timeoutMs: number): Promise { + async acquireHeavyWithin( + timeoutMs: number, + signal?: AbortSignal + ): Promise { const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs)); for (;;) { + if (signal?.aborted) return null; const lease = this.tryAcquireHeavy(); if (lease) return lease; const remaining = deadline - Date.now(); @@ -152,14 +163,33 @@ export class ChatAdmissionController { resolver = () => resolve(); this.#waiters.push(resolver); }); - const timedOut = await Promise.race([ + let deadlineTimer: ReturnType | null = null; + const races: Array> = [ released.then(() => false), - new Promise((resolve) => setTimeout(() => resolve(true), remaining)), - ]); + new Promise((resolve) => { + deadlineTimer = setTimeout(() => resolve(true), remaining); + }), + ]; + let onAbort: (() => void) | null = null; + if (signal) { + races.push( + new Promise((resolve) => { + const listener = () => resolve(true); + onAbort = listener; + signal.addEventListener("abort", listener, { once: true }); + // Already-aborted signals must settle without parking. + if (signal.aborted) resolve(true); + }) + ); + } + const timedOut = await Promise.race(races); if (resolver) { const index = this.#waiters.indexOf(resolver); if (index >= 0) this.#waiters.splice(index, 1); } + // Cancel the deadline timer when abort/release wins; a fired timer is a no-op. + if (deadlineTimer) clearTimeout(deadlineTimer); + if (onAbort) signal?.removeEventListener("abort", onAbort); if (timedOut) return null; } } @@ -417,6 +447,7 @@ export async function admitChatStructure( heavyTools?: number; heavyTokens?: number; queueMs?: number; + signal?: AbortSignal; } = {} ): Promise { if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease }; @@ -454,7 +485,7 @@ export async function admitChatStructure( (options.sessionId ? perConnectionAdmissionController.getController(options.sessionId) : defaultAdmissionController); - const acquired = await controller.acquireHeavyWithin(options.queueMs ?? 0); + const acquired = await controller.acquireHeavyWithin(options.queueMs ?? 0, options.signal); return acquired ? { admit: true, lease: acquired } : { admit: false, response: structuralRejectionResponse(503, maxMessages) }; @@ -623,7 +654,7 @@ export async function admitChatRequest( let lease: ChatAdmissionLease | null = null; const reserve = async (): Promise => { if (lease) return true; - lease = await controller.acquireHeavyWithin(queueMs); + lease = await controller.acquireHeavyWithin(queueMs, request.signal); return lease !== null; }; diff --git a/tests/unit/chat-body-admission.test.ts b/tests/unit/chat-body-admission.test.ts index 5ecb5111b6..0011b5cecc 100644 --- a/tests/unit/chat-body-admission.test.ts +++ b/tests/unit/chat-body-admission.test.ts @@ -978,3 +978,189 @@ test("admission waiters are served FIFO as capacity frees", async () => { if (secondResult.admit) secondResult.lease?.release(); assert.equal(controller.activeHeavy, 0); }); + +// ── AbortSignal support in acquireHeavyWithin (#9654 / U2) ──────────────── +// A disconnected client must not keep parking in the admission queue for the +// full queueMs. On abort the waiter is removed from the FIFO immediately and +// the acquire resolves `null` early (the caller's 503 is dropped on the dead +// connection); no capacity is consumed and the freed slot does not wake it. + +test("aborting the admission wait settles early, grants no lease, and removes the waiter", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const abortController = new AbortController(); + const pending = controller.acquireHeavyWithin(2_000, abortController.signal); + + // Parked while capacity is busy. + let settled = false; + void pending.then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(settled, false, "must be parked while capacity is busy"); + + abortController.abort(); + + // Must settle well before the 2s deadline. + let settledAfterAbort = false; + void pending.then(() => { + settledAfterAbort = true; + }); + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.equal(settledAfterAbort, true, "abort must settle the wait promptly, not park for queueMs"); + + const lease = await pending; + assert.equal(lease, null, "abort must not grant a lease"); + assert.equal(controller.activeHeavy, 1, "the holder keeps its lease; the aborted wait consumed nothing"); + + // Releasing must NOT wake the removed waiter: capacity stays free. + held.release(); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal( + controller.activeHeavy, + 0, + "releasing after abort must not wake the removed waiter" + ); +}); + +test("aborting the head waiter preserves FIFO order for remaining waiters", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const firstAbort = new AbortController(); + const first = controller.acquireHeavyWithin(2_000, firstAbort.signal); + const second = controller.acquireHeavyWithin(2_000); + + // Both are parked, head-first. + await new Promise((resolve) => setTimeout(resolve, 30)); + + // Abort the HEAD waiter: it must leave the queue without disturbing the rest. + firstAbort.abort(); + assert.equal(await first, null, "head waiter returns null on abort"); + + // The remaining waiter is now first in line and must get the freed capacity. + held.release(); + const secondLease = await second; + assert.ok(secondLease, "remaining waiter must acquire the freed capacity"); + secondLease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("a pre-aborted signal never parks in the admission queue", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const abortController = new AbortController(); + abortController.abort("client already disconnected"); + + const pending = controller.acquireHeavyWithin(2_000, abortController.signal); + let settled = false; + void pending.then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.equal(settled, true, "a pre-aborted signal must settle immediately, not park"); + + const lease = await pending; + assert.equal(lease, null, "no lease is granted after abort"); + assert.equal(controller.activeHeavy, 1, "holder keeps capacity; aborted wait consumed nothing"); + held.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("aborting the request signal cancels a queued byte-heavy wait", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const abortController = new AbortController(); + const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] }); + const request = new Request("http://x/v1/chat/completions", { + method: "POST", + headers: { "content-type": "application/json" }, + body, + signal: abortController.signal, + }); + const pending = admitChatRequest(request, { + controller, + largeBodyBytes: 32, + hardMaxBytes: 1024, + queueMs: 2_000, + }); + + let settled = false; + void pending.then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(settled, false, "must queue while capacity is busy"); + + abortController.abort(); + const started = Date.now(); + const result = await pending; + assert.ok( + Date.now() - started < 500, + "abort must cancel the queue-wait early, not park the full queueMs" + ); + assert.equal(result.admit, false, "abort must not admit"); + if (!result.admit) { + assert.equal(result.response.status, 503); + assert.equal((await result.response.json()).error.code, "chat_admission_busy"); + } + assert.equal(controller.activeHeavy, 1, "holder keeps capacity; aborted wait consumed nothing"); + held.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("aborting the signal cancels a structural queue-wait", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const abortController = new AbortController(); + const pending = admitChatStructure( + { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }, + null, + { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 2_000, + signal: abortController.signal, + } + ); + + let settled = false; + void pending.then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(settled, false, "must queue while capacity is busy"); + + abortController.abort(); + const started = Date.now(); + const result = await pending; + assert.ok( + Date.now() - started < 500, + "abort must cancel the queue-wait early, not park the full queueMs" + ); + assert.equal(result.admit, false, "abort must not admit"); + if (!result.admit) { + assert.equal(result.response.status, 503); + assert.equal((await result.response.json()).error.code, "chat_admission_busy"); + } + assert.equal(controller.activeHeavy, 1, "holder keeps its lease"); + held.release(); + assert.equal(controller.activeHeavy, 0); +});