diff --git a/changelog.d/fixes/12395-heavy-admission-retry-after.md b/changelog.d/fixes/12395-heavy-admission-retry-after.md new file mode 100644 index 0000000000..4cc5badb6c --- /dev/null +++ b/changelog.d/fixes/12395-heavy-admission-retry-after.md @@ -0,0 +1 @@ +- **fix(chat-admission):** derive the `chat_admission_busy` 503 `Retry-After` from observed heavyweight-lease occupancy — the larger of the exhausted `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` window and the time since capacity last turned over, capped at 60 s — instead of a fixed 1 s (structural) / 2 s (byte-stage) hint that invited Codex/agent fan-out clients to re-send ~1 MiB `/v1/responses` bodies every second into a gate held for the whole SSE lifetime; an idle gate keeps the historical floors ([#12135](https://github.com/diegosouzapw/OmniRoute/issues/12135)) (#12395 — thanks @pacocartones) diff --git a/docs/guides/TROUBLESHOOTING.md b/docs/guides/TROUBLESHOOTING.md index 4e0cb8883b..e512afbfee 100644 --- a/docs/guides/TROUBLESHOOTING.md +++ b/docs/guides/TROUBLESHOOTING.md @@ -538,8 +538,12 @@ When many concurrent requests hit a rate-limited provider, OmniRoute uses mutex - The chat completions endpoint returns a retryable `503` response whose error code is `chat_admission_busy`. -- The response includes `Retry-After`; the byte-based path uses 2 seconds, while the - structure-based path uses 1 second and includes `reason: "structure_limit"`. +- The response includes `Retry-After`. Since #12135 the value is derived from observed + occupancy — the larger of the `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` window the request already + waited and the time the current heavyweight leases have been held — rounded up to whole + seconds and capped at 60. On an idle gate it keeps the historical floors: 2 seconds on the + byte-based path, 1 second on the structure-based path (which also includes + `reason: "structure_limit"`). - This can happen while another heavyweight chat or long-running streaming response is still in flight. diff --git a/src/shared/middleware/chatAdmissionResponses.ts b/src/shared/middleware/chatAdmissionResponses.ts index ac1dfdce66..0d848b1fba 100644 --- a/src/shared/middleware/chatAdmissionResponses.ts +++ b/src/shared/middleware/chatAdmissionResponses.ts @@ -4,10 +4,34 @@ import { CORS_HEADERS } from "../utils/cors"; const JSON_HEADERS = { ...CORS_HEADERS, "Content-Type": "application/json" }; -export function chatAdmissionRejectionResponse(status: 413 | 503, hardMaxBytes: number): Response { +/** + * `Retry-After` floors for the retryable 503s — the pre-#12135 fixed values. A caller + * passes an occupancy-derived hint (`ChatAdmissionController#retryAfterSeconds`) and the + * header carries whichever is larger, so an idle gate still answers exactly as before + * while a gate whose leases have been busy for a whole SSE stream stops inviting a + * 1-second retry storm. + */ +const BYTE_STAGE_RETRY_AFTER_FLOOR_SECONDS = 2; +const STRUCTURAL_RETRY_AFTER_FLOOR_SECONDS = 1; + +function retryAfterHeader(floorSeconds: number, hintSeconds: number | undefined): string { + const hint = Number.isFinite(hintSeconds) ? Math.ceil(hintSeconds as number) : 0; + return String(Math.max(floorSeconds, hint)); +} + +export function chatAdmissionRejectionResponse( + status: 413 | 503, + hardMaxBytes: number, + retryAfterSeconds?: number +): Response { const isPayload = status === 413; const headers: Record = { ...JSON_HEADERS }; - if (!isPayload) headers["Retry-After"] = "2"; + if (!isPayload) { + headers["Retry-After"] = retryAfterHeader( + BYTE_STAGE_RETRY_AFTER_FLOOR_SECONDS, + retryAfterSeconds + ); + } const message = isPayload ? `Request body too large for chat completions (max ${Math.floor( hardMaxBytes / (1024 * 1024) @@ -53,10 +77,19 @@ export function resourcePressureRejectionResponse(): Response { ); } -export function structuralRejectionResponse(status: 413 | 503, maxMessages: number): Response { +export function structuralRejectionResponse( + status: 413 | 503, + maxMessages: number, + retryAfterSeconds?: number +): Response { const historyLimit = status === 413; const headers: Record = { ...JSON_HEADERS }; - if (!historyLimit) headers["Retry-After"] = "1"; + if (!historyLimit) { + headers["Retry-After"] = retryAfterHeader( + STRUCTURAL_RETRY_AFTER_FLOOR_SECONDS, + retryAfterSeconds + ); + } const body = buildErrorBody( status, historyLimit diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 0fc550ba6b..b30ad4cad7 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -92,6 +92,15 @@ export const CHAT_ADMISSION_MAX_QUEUED_BYTES = parsePositiveInt( 4 * 1024 * 1024 ); +/** + * Ceiling for the occupancy-derived `Retry-After` on a capacity 503 (#12135). A + * heavyweight lease is held for the whole SSE lifetime, so the hint is derived from how + * long capacity has demonstrably been busy (`ChatAdmissionController#retryAfterSeconds`); + * this cap keeps a multi-minute stream from telling a client to sleep for minutes when + * another slot may free far sooner. + */ +export const CHAT_ADMISSION_RETRY_AFTER_MAX_SECONDS = 60; + export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt( process.env.OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT, 200 @@ -260,6 +269,9 @@ export class ChatAdmissionController { * `CHAT_MAX_HEAVY_IN_FLIGHT` bound, but still a real, finite ceiling instead of * the unconditional bypass this replaces. */ #activeHealthy = 0; + /** #12135: acquisition time of every live heavy lease, keyed by an opaque token, so the + * capacity 503 can advertise a `Retry-After` derived from observed occupancy. */ + #heavyLeaseStartedAt = new Map(); /** Per-key FIFOs. A key groups one client's waiters so they are served * round-robin against the shared budget instead of monopolizing a strict * FIFO (see #dispatchFair). */ @@ -397,6 +409,8 @@ export class ChatAdmissionController { tryAcquireHeavy(): ChatAdmissionLease | null { if (this.#activeHeavy >= this.maxHeavyInFlight) return null; this.#activeHeavy += 1; + const token = Symbol("heavy-lease"); + this.#heavyLeaseStartedAt.set(token, Date.now()); const done = trackRequest(); let released = false; return { @@ -407,12 +421,39 @@ export class ChatAdmissionController { if (released) return; released = true; this.#activeHeavy = Math.max(0, this.#activeHeavy - 1); + this.#heavyLeaseStartedAt.delete(token); done(); this.#dispatchFair(); }, }; } + /** + * `Retry-After` (whole seconds) for a capacity 503, derived from live occupancy instead + * of a fixed constant (#12135). A heavyweight lease is held for the ENTIRE SSE lifetime + * (tens of seconds to minutes), so a fixed 1–2 s hint invited clients to re-send the + * same ~1 MiB body every second into a gate that could not possibly have cleared. The + * hint is the larger of: + * - `queueMs`, the bounded wait the caller already exhausted — the server itself needed + * longer than that, so advertising less is dishonest; and + * - the age of the YOUNGEST live heavy lease: the time since heavyweight capacity last + * turned over. Every slot has been continuously held at least that long, so it is the + * observed floor on how long "busy" has lasted (the oldest lease would be a pessimist + * with N slots in flight). + * Rounded up and capped at `CHAT_ADMISSION_RETRY_AFTER_MAX_SECONDS`. The response + * builders floor the result at their historical value (1 s structural, 2 s byte-stage), + * so an idle gate answers exactly as before. + */ + retryAfterSeconds(queueMs: number, now = Date.now()): number { + let youngestAgeMs = Number.POSITIVE_INFINITY; + for (const startedAt of this.#heavyLeaseStartedAt.values()) { + youngestAgeMs = Math.min(youngestAgeMs, now - startedAt); + } + const occupancyMs = Number.isFinite(youngestAgeMs) ? youngestAgeMs : 0; + const hintSeconds = Math.ceil(Math.max(0, queueMs, occupancyMs) / 1000); + return Math.min(CHAT_ADMISSION_RETRY_AFTER_MAX_SECONDS, Math.max(1, hintSeconds)); + } + /** * Wait up to `timeoutMs` for heavyweight capacity, retrying atomically on each * release. Resolves `null` when the deadline expires with no capacity freed, in @@ -843,26 +884,41 @@ export async function admitChatStructure( // Structural-only waits happen on byte-light bodies (a byte-heavy body already // holds the byte-stage lease), so the conservative 256KB weight bounds the // parsed JSON the waiter keeps resident while parked. + const queueMs = options.queueMs ?? 0; const acquiredCount = await controller.acquireHeavyWithin( - options.queueMs ?? 0, + queueMs, options.signal, CHAT_LARGE_BODY_BYTES, options.sessionId ); if (!acquiredCount) { - return { admit: false, response: structuralRejectionResponse(503, maxMessages) }; + return { + admit: false, + response: structuralRejectionResponse( + 503, + maxMessages, + controller.retryAfterSeconds(queueMs) + ), + }; } // #503-fanout: same composed count+budget gate as the fast path above. const acquiredBudget = await controller.acquireBudgetWithin( CHAT_LARGE_BODY_BYTES, - options.queueMs ?? 0, + queueMs, options.signal, options.sessionId ); if (acquiredBudget.status !== "acquired") { acquiredCount.release(); - return { admit: false, response: structuralRejectionResponse(503, maxMessages) }; + return { + admit: false, + response: structuralRejectionResponse( + 503, + maxMessages, + controller.retryAfterSeconds(queueMs) + ), + }; } return { admit: true, @@ -1006,6 +1062,10 @@ export async function admitChatRequest( return true; }; + // #12135: the capacity 503 advertises an occupancy-derived Retry-After. + const busyResponse = () => + chatAdmissionRejectionResponse(503, hardMaxBytes, controller.retryAfterSeconds(queueMs)); + // 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 ( @@ -1013,7 +1073,7 @@ export async function admitChatRequest( contentLength >= largeBodyBytes && !(await reserve(Math.min(contentLength, hardMaxBytes))) ) { - return { admit: false, response: chatAdmissionRejectionResponse(503, hardMaxBytes) }; + return { admit: false, response: busyResponse() }; } const reader = request.body?.getReader(); @@ -1039,7 +1099,7 @@ export async function admitChatRequest( } if (totalBytes >= largeBodyBytes && !(await reserve(totalBytes))) { await reader.cancel("chat admission capacity unavailable").catch(() => undefined); - return { admit: false, response: chatAdmissionRejectionResponse(503, hardMaxBytes) }; + return { admit: false, response: busyResponse() }; } chunks.push(value); } diff --git a/tests/unit/heavy-admission-retry-after-12135.test.ts b/tests/unit/heavy-admission-retry-after-12135.test.ts new file mode 100644 index 0000000000..fda3b5f259 --- /dev/null +++ b/tests/unit/heavy-admission-retry-after-12135.test.ts @@ -0,0 +1,229 @@ +// #12135: "[BUG] Heavy /v1/responses still 503 chat_admission_busy after MAX_HEAVY is +// raised: QUEUE_MS and Retry-After: 1 are far shorter than SSE occupancy". +// +// A heavyweight admission lease is held for the ENTIRE SSE lifetime (tens of seconds to +// minutes), but the retryable 503 advertised a fixed `Retry-After: 1` (structural path) +// or `Retry-After: 2` (byte-stage path) regardless of how long capacity had actually been +// busy or how long the waiter had already spent in the bounded queue. Clients that honor +// the header (Codex CLI, agent fan-out) re-sent the same ~1 MiB body every second into a +// gate that could not possibly have cleared, producing a `queue_timeout` retry storm. +// +// The maintainer scoped the fix on the issue: "A `Retry-After` derived from observed +// lease age/occupancy would be honest." These tests pin that contract WITHOUT touching +// the queue posture (`OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` default), which the maintainer +// explicitly left as a separate decision: +// (a) after a `queue_timeout`, `Retry-After` is at least the queue window the waiter +// already exhausted — never less than what the server itself needed; +// (b) `Retry-After` reflects the observed age of the in-flight heavy lease (time since +// heavyweight capacity last turned over), on BOTH the structural and byte-stage 503s; +// (c) the hint is capped so a multi-minute stream never tells a client to sleep for +// minutes when another slot may free sooner; +// (d) an idle gate (fresh lease, no queue) still answers exactly as before (1 s / 2 s), +// so no existing client behavior changes on a quiet host; +// (e) with the count cap raised, a third heavy `/v1/responses` request is queued and +// admitted when a lease frees inside `queueMs` instead of being rejected. +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { + admitChatRequest, + admitChatStructure, + ChatAdmissionController, + type ChatAdmissionLease, +} from "../../src/shared/middleware/chatBodyAdmission.ts"; + +/** The reporter's shape: a Codex `/v1/responses` session with ~70 function tools. */ +function responsesHeavyBody() { + const tools = Array.from({ length: 70 }, (_, i) => ({ + type: "function", + name: `tool_${i}`, + description: "a".repeat(64), + parameters: { type: "object", properties: {} }, + })); + return { + model: "gpt-5.6-sol", + input: [{ role: "user", content: "run the plan" }], + tools, + stream: true, + }; +} + +const NOOP_SHED_SINK = () => {}; + +function heavyController(maxHeavyInFlight: number): ChatAdmissionController { + // healthyHeadroom=0 forces the bounded-wait/shed path; the sink keeps pino quiet. + return new ChatAdmissionController(maxHeavyInFlight, undefined, 0, NOOP_SHED_SINK); +} + +type Rejected = { admit: false; response: Response }; + +function admitStructure( + controller: ChatAdmissionController, + queueMs: number +): ReturnType { + return admitChatStructure(responsesHeavyBody(), null, { + controller, + queueMs, + heapPressureCheck: () => true, + }); +} + +async function holdStructural(controller: ChatAdmissionController): Promise { + const holder = await admitStructure(controller, 0); + assert.equal(holder.admit, true); + const lease = (holder as { admit: true; lease: ChatAdmissionLease | null }).lease; + assert.ok(lease, "the first heavy request must hold the heavyweight lease"); + return lease; +} + +/** Let a pending admission park in the queue before the mocked clock advances. */ +async function settleMicrotasks(): Promise { + for (let i = 0; i < 8; i++) await Promise.resolve(); +} + +test("#12135 (a): structural queue_timeout 503 advertises at least the exhausted queue window", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(1); + const lease = await holdStructural(controller); + try { + const pending = admitStructure(controller, 5_000); + await settleMicrotasks(); + assert.equal(controller.waitingCount, 1, "the second heavy request must park, not fail fast"); + t.mock.timers.tick(5_000); + const result = (await pending) as Rejected; + assert.equal(result.admit, false); + assert.equal(result.response.status, 503); + assert.equal( + result.response.headers.get("Retry-After"), + "5", + "Retry-After must not be shorter than the queue window the waiter already burned" + ); + const body = await result.response.json(); + assert.equal(body.error?.code, "chat_admission_busy"); + assert.equal(body.error?.reason, "structure_limit"); + } finally { + lease.release(); + } +}); + +test("#12135 (b): structural 503 Retry-After reflects the observed age of the in-flight lease", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(1); + const lease = await holdStructural(controller); + try { + // The holder streams for 45 s; a fast-fail (queueMs=0) arrival must be told to wait + // on the order of what capacity has demonstrably been busy for, not 1 s. + t.mock.timers.tick(45_000); + const result = (await admitStructure(controller, 0)) as Rejected; + assert.equal(result.admit, false); + assert.equal(result.response.status, 503); + assert.equal(result.response.headers.get("Retry-After"), "45"); + } finally { + lease.release(); + } +}); + +test("#12135 (b): Retry-After is the age of the YOUNGEST live lease — time since capacity last turned over", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(2); + const first = await holdStructural(controller); + t.mock.timers.tick(40_000); + const second = await holdStructural(controller); + try { + t.mock.timers.tick(7_000); + const result = (await admitStructure(controller, 0)) as Rejected; + assert.equal(result.admit, false); + // first is 47 s old, second is 7 s old: every slot has been continuously held for + // at least 7 s, so that is the honest occupancy floor — not the 47 s pessimist. + assert.equal(result.response.headers.get("Retry-After"), "7"); + } finally { + second.release(); + first.release(); + } +}); + +test("#12135 (c): the occupancy-derived hint is capped", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(1); + const lease = await holdStructural(controller); + try { + t.mock.timers.tick(10 * 60_000); + const result = (await admitStructure(controller, 0)) as Rejected; + assert.equal(result.admit, false); + // CHAT_ADMISSION_RETRY_AFTER_MAX_SECONDS: a 10-minute-old stream must not advertise + // 600 s — another slot may free long before that. + assert.equal(result.response.headers.get("Retry-After"), "60"); + } finally { + lease.release(); + } +}); + +function largeRequest(): Request { + const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] }); + return new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json", "content-length": String(body.length) }, + body, + }); +} + +test("#12135 (b): byte-stage 503 (admitChatRequest) carries the same occupancy-derived Retry-After", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(1); + const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 0 }; + const first = await admitChatRequest(largeRequest(), options); + assert.equal(first.admit, true); + if (!first.admit) return; + try { + t.mock.timers.tick(30_000); + const second = (await admitChatRequest(largeRequest(), options)) as Rejected; + assert.equal(second.admit, false); + assert.equal(second.response.status, 503); + assert.equal(second.response.headers.get("Retry-After"), "30"); + assert.equal((await second.response.json()).error.code, "chat_admission_busy"); + } finally { + first.lease?.release(); + } +}); + +test("#12135 (d): an idle gate keeps the historical 1 s (structural) and 2 s (byte-stage) floors", async () => { + const structural = heavyController(1); + const structuralLease = await holdStructural(structural); + try { + const result = (await admitStructure(structural, 0)) as Rejected; + assert.equal(result.admit, false); + assert.equal(result.response.headers.get("Retry-After"), "1"); + } finally { + structuralLease.release(); + } + + const byteStage = heavyController(1); + const options = { controller: byteStage, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 0 }; + const first = await admitChatRequest(largeRequest(), options); + assert.equal(first.admit, true); + if (!first.admit) return; + try { + const second = (await admitChatRequest(largeRequest(), options)) as Rejected; + assert.equal(second.admit, false); + assert.equal(second.response.headers.get("Retry-After"), "2"); + } finally { + first.lease?.release(); + } +}); + +test("#12135 (e): with the count cap raised, a third heavy /v1/responses request queues and is admitted when a lease frees inside queueMs", async (t) => { + t.mock.timers.enable({ apis: ["Date", "setTimeout"] }); + const controller = heavyController(2); + const first = await holdStructural(controller); + const second = await holdStructural(controller); + assert.equal(controller.activeHeavy, 2); + const pending = admitStructure(controller, 10_000); + await settleMicrotasks(); + assert.equal(controller.waitingCount, 1, "the third request must wait, not 503"); + t.mock.timers.tick(1_000); + first.release(); + const third = await pending; + assert.equal(third.admit, true, "a freed lease inside the queue window must admit the waiter"); + if (third.admit) third.lease?.release(); + second.release(); + assert.equal(controller.activeHeavy, 0); +});