diff --git a/changelog.d/fixes/14456-admission-lease-pending-abort.md b/changelog.d/fixes/14456-admission-lease-pending-abort.md new file mode 100644 index 00000000000..fb550aa3271 --- /dev/null +++ b/changelog.d/fixes/14456-admission-lease-pending-abort.md @@ -0,0 +1 @@ +- fix(resilience): release the admission lease immediately when a client aborts while the chat handler is still pending, not only after it eventually settles (#14456) diff --git a/src/shared/middleware/chatAdmissionRelease.ts b/src/shared/middleware/chatAdmissionRelease.ts index 61dde6e3a8d..75ddaa275b8 100644 --- a/src/shared/middleware/chatAdmissionRelease.ts +++ b/src/shared/middleware/chatAdmissionRelease.ts @@ -87,9 +87,33 @@ export async function releaseChatAdmissionAfterHandler( lease: ChatAdmissionLease | null, options: ReleaseChatAdmissionOptions = {} ): Promise { + const { signal } = options; + let abortedWhilePending = false; + let detachAbortListener = (): void => undefined; + + if (signal && lease) { + const onAbort = (): void => { + abortedWhilePending = true; + detachAbortListener(); + if (!lease.released) lease.release(); + }; + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener("abort", onAbort, { once: true }); + detachAbortListener = () => signal.removeEventListener("abort", onAbort); + } + } + try { - return releaseChatAdmissionWhenDone(await responsePromise, lease, options); + const response = await responsePromise; + detachAbortListener(); + // The lease was already released by the pending-phase abort listener above; + // skip the redundant stream-wrapping work in releaseChatAdmissionWhenDone. + if (abortedWhilePending) return response; + return releaseChatAdmissionWhenDone(response, lease, options); } catch (error) { + detachAbortListener(); lease?.release(); throw error; } diff --git a/tests/unit/chat-admission-pending-handler-abort-14456.test.ts b/tests/unit/chat-admission-pending-handler-abort-14456.test.ts new file mode 100644 index 00000000000..17c802d15ab --- /dev/null +++ b/tests/unit/chat-admission-pending-handler-abort-14456.test.ts @@ -0,0 +1,95 @@ +/** + * Repro for #14456: an admission lease is NOT released when the client aborts + * WHILE the handler promise passed to `releaseChatAdmissionAfterHandler` is + * still pending. The abort-aware wrapper (`releaseChatAdmissionWhenDone`) is + * only installed AFTER `await responsePromise` resolves, so an abort signal + * fired during that pending phase has nothing listening to it yet. + * + * PR #14457 fixed the case where an SSE `Response` already exists and the + * client then disconnects (covered by chat-admission-abandoned-stream-lease + * .test.ts). It did not touch the pending-handler phase, which is what this + * test isolates: the handler (e.g. an in-flight upstream LLM call) has not + * resolved to a Response yet when the request aborts. + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import { + ChatAdmissionController, + releaseChatAdmissionAfterHandler, +} from "../../src/shared/middleware/chatBodyAdmission.ts"; + +test("a request that aborts while the handler is still pending must release the lease promptly, not only after the handler eventually settles", async () => { + const controller = new ChatAdmissionController(1); + const lease = controller.tryAcquireHeavy(); + assert.ok(lease, "precondition: a slot is available"); + assert.equal(controller.activeHeavy, 1); + + const abort = new AbortController(); + + let resolveHandler!: (value: Response) => void; + const handlerPromise = new Promise((resolve) => { + resolveHandler = resolve; + }); + + const released = releaseChatAdmissionAfterHandler(handlerPromise, lease, { + signal: abort.signal, + }); + + abort.abort(); + await new Promise((resolve) => setTimeout(resolve, 10)); + + assert.equal( + controller.activeHeavy, + 0, + "an aborted client must not hold a heavyweight slot while the handler is still pending" + ); + + resolveHandler( + new Response(null, { status: 504, headers: { "content-type": "application/json" } }) + ); + await released; +}); + +test("a handler that resolves normally without aborting releases exactly once and detaches its abort listener", async () => { + const controller = new ChatAdmissionController(1); + const lease = controller.tryAcquireHeavy(); + assert.ok(lease); + + const abort = new AbortController(); + const response = new Response(null, { status: 200 }); + + const result = await releaseChatAdmissionAfterHandler(Promise.resolve(response), lease, { + signal: abort.signal, + }); + + assert.equal(controller.activeHeavy, 0, "lease must be released on normal completion"); + assert.equal(result.status, 200); + + // Aborting after normal completion must not throw or double-release. + assert.doesNotThrow(() => abort.abort()); +}); + +test("an already-aborted signal before releaseChatAdmissionAfterHandler is called releases immediately", async () => { + const controller = new ChatAdmissionController(1); + const lease = controller.tryAcquireHeavy(); + assert.ok(lease); + + const abort = new AbortController(); + abort.abort(); + + let resolveHandler!: (value: Response) => void; + const handlerPromise = new Promise((resolve) => { + resolveHandler = resolve; + }); + + const released = releaseChatAdmissionAfterHandler(handlerPromise, lease, { + signal: abort.signal, + }); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + assert.equal(controller.activeHeavy, 0, "already-aborted signal must release immediately"); + + resolveHandler(new Response(null, { status: 504 })); + await released; +});