diff --git a/changelog.d/fixes/10322-process-wide-admission-budget.md b/changelog.d/fixes/10322-process-wide-admission-budget.md new file mode 100644 index 0000000000..defab2a7fe --- /dev/null +++ b/changelog.d/fixes/10322-process-wide-admission-budget.md @@ -0,0 +1 @@ +- **fix(chat-body-admission):** restore a single process-wide admission budget — heavyweight leases and queued bytes are now bounded once for the whole process instead of per session, so one session can no longer mint extra capacity or starve others; per-session fairness is preserved via round-robin dispatch ([#10110](https://github.com/diegosouzapw/OmniRoute/issues/10110)) diff --git a/docs/reference/ENVIRONMENT.md b/docs/reference/ENVIRONMENT.md index 1592a760fb..5c89b3acc6 100644 --- a/docs/reference/ENVIRONMENT.md +++ b/docs/reference/ENVIRONMENT.md @@ -1501,9 +1501,9 @@ These settings were introduced after the previous environment-contract snapshot. | Variable | Default | Source File | Description | | --- | --- | --- | --- | | `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` | `2000` | `src/shared/middleware/chatBodyAdmission.ts` | Maximum wait for a heavyweight chat admission slot before a retryable `503`; a short bounded wait serializes agent bursts instead of an instant `503`. `0` restores immediate rejection. | -| `OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES` | `4194304` (4 MB) | `src/shared/middleware/chatBodyAdmission.ts` | Queued-bytes budget for the admission wait (#9654): bounds total buffered body bytes parked per lane so the wait cannot amplify the heap (#4380). Over-budget waits receive a retryable `503` immediately. | -| `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` | `60000` (60 s) | `src/shared/middleware/chatBodyAdmission.ts` | Per-connection virtual admission lanes (#9654): idle-lane eviction TTL. | -| `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` | `64` | `src/shared/middleware/chatBodyAdmission.ts` | Per-connection virtual admission lanes (#9654): max concurrent sessions (lanes). | +| `OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES` | `4194304` (4 MB) | `src/shared/middleware/chatBodyAdmission.ts` | Queued-bytes budget for the admission wait: bounds total buffered body bytes parked process-wide so the wait cannot amplify the heap (#4380). Over-budget waits receive a retryable `503` immediately. | +| `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` | `60000` (60 s) | `src/shared/middleware/chatBodyAdmission.ts` | Deprecated no-op since #10110: per-session admission lanes were removed in favor of one process-wide budget. Accepted for configuration compatibility; ignored. | +| `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` | `64` | `src/shared/middleware/chatBodyAdmission.ts` | Deprecated no-op since #10110: per-session admission lanes were removed in favor of one process-wide budget. Accepted for configuration compatibility; ignored. | | `OMNIROUTE_RUNNOW_TIMEOUT_MS` | `30000` | `src/app/api/jobs/[id]/run-now/route.ts` | Bounds how long a run-now call waits for an in-flight job before starting the queued run. | | `ADOBE_FIREFLY_BROWSER_REFRESH` | enabled | `open-sse/services/adobeFireflySession.ts` | Keeps IMS and browser-risk state fresh through account-scoped Chrome CDP sessions; set `0` to disable. | | `ADOBE_FIREFLY_SESSION_DISK` | enabled | `open-sse/services/adobeFireflySession.ts` | Persists repaired Adobe sessions under `DATA_DIR`; set `0` for memory-only state. | diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 290a18f6db..e85fce30ec 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -7,20 +7,17 @@ * local heavyweight capacity before parsing and enforces the hard limit against bytes read, * not an untrusted Content-Length header. * - * Per-connection virtual admission lanes (#9654): each distinct API-key (or anonymous) - * bucket gets its own FairCostQueue so one connection cannot exhaust heavyweight capacity - * and starve others. Idle sessions are auto-evicted after a TTL. + * Process-wide admission budget (#10110): ALL requests — every API key, every + * session — contend for ONE global heavyweight budget, so the documented + * "in one process" bound holds against fake-credential sharding. Per-request + * session identity is used only as a fairness scheduling key: waiters are + * grouped per session and served round-robin against the shared budget, so one + * connection's burst cannot starve others (#9654). */ import { CORS_HEADERS } from "../utils/cors"; import { createHash } from "crypto"; - -const OMNIROUTE_CHAT_VIRTUAL_TTL_MS = parsePositiveInt( - process.env.OMNIROUTE_CHAT_VIRTUAL_TTL_MS, - 60_000 -); - function parsePositiveInt(value: string | undefined, fallback: number): number { const parsed = Number.parseInt(String(value), 10); return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : fallback; @@ -109,6 +106,12 @@ export interface ChatAdmissionLease { release(): void; } +/** A parked waiter, grouped by fairness key for round-robin dispatch. */ +interface AdmissionWaiter { + readonly key: string; + readonly resolve: () => void; +} + /** * Process-local heavyweight reservation. The capacity check and increment execute in one * synchronous JavaScript turn, making acquisition atomic within an OmniRoute process. @@ -119,7 +122,13 @@ export interface ChatAdmissionLease { export class ChatAdmissionController { #activeHeavy = 0; #queuedBytes = 0; - #waiters: Array<() => void> = []; + /** 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). */ + #queues = new Map(); + /** Keys in creation order; #fairCursor scans them round-robin. */ + #fairKeys: string[] = []; + #fairCursor = 0; constructor( readonly maxHeavyInFlight = 1, @@ -137,11 +146,25 @@ export class ChatAdmissionController { return this.#activeHeavy; } - /** Total buffered bytes currently parked in the FIFO (heap valve accounting). */ + /** Total buffered bytes currently parked across all queues (heap valve accounting). */ get queuedBytes(): number { return this.#queuedBytes; } + /** Total waiters parked across all keys (diagnostics). */ + get waitingCount(): number { + let total = 0; + for (const queue of this.#queues.values()) total += queue.length; + return total; + } + + /** Per-key waiter depths (diagnostics) — opaque scheduler keys, never raw credentials. */ + get waitersByKey(): ReadonlyArray<{ key: string; waiting: number }> { + const out: Array<{ key: string; waiting: number }> = []; + for (const [key, queue] of this.#queues) out.push({ key, waiting: queue.length }); + return out; + } + tryAcquireHeavy(): ChatAdmissionLease | null { if (this.#activeHeavy >= this.maxHeavyInFlight) return null; this.#activeHeavy += 1; @@ -154,7 +177,7 @@ export class ChatAdmissionController { if (released) return; released = true; this.#activeHeavy = Math.max(0, this.#activeHeavy - 1); - this.#waiters.shift()?.(); + this.#dispatchFair(); }, }; } @@ -163,10 +186,14 @@ export class ChatAdmissionController { * 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. + * legacy immediate-reject path. + * + * Waiters are grouped by `sessionKey` and served round-robin across keys + * (#dispatchFair), so one client's burst cannot starve another's bounded wait + * while every key contends for the SAME process-wide budget. * * When `signal` aborts while parked (client disconnect), the waiter is removed - * from the FIFO immediately and the promise resolves `null` early instead of + * from its queue 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 @@ -181,7 +208,8 @@ export class ChatAdmissionController { async acquireHeavyWithin( timeoutMs: number, signal?: AbortSignal, - queuedBytes = 0 + queuedBytes = 0, + sessionKey = "default" ): Promise { const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs)); for (;;) { @@ -195,14 +223,26 @@ export class ChatAdmissionController { return null; } this.#queuedBytes += queuedBytes; - let resolver: (() => void) | null = null; - const released = new Promise((resolve) => { - resolver = () => resolve(); - this.#waiters.push(resolver); + // Park into this key's FIFO (creating the key on first use). + let queue = this.#queues.get(sessionKey); + if (!queue) { + queue = []; + this.#queues.set(sessionKey, queue); + this.#fairKeys.push(sessionKey); + } + const lane = queue; + let resolveParked: (() => void) | null = null; + const waiter: AdmissionWaiter = { + key: sessionKey, + resolve: () => resolveParked?.(), + }; + const parked = new Promise((resolve) => { + resolveParked = () => resolve(); + lane.push(waiter); }); let deadlineTimer: ReturnType | null = null; const races: Array> = [ - released.then(() => false), + parked.then(() => false), new Promise((resolve) => { deadlineTimer = setTimeout(() => resolve(true), remaining); }), @@ -220,41 +260,80 @@ export class ChatAdmissionController { ); } const timedOut = await Promise.race(races); - // The waiter has left the FIFO (wake, abort, or timeout) — release its charge. + // The waiter has left its queue (wake, abort, or timeout) — release its charge. this.#queuedBytes = Math.max(0, this.#queuedBytes - queuedBytes); - if (resolver) { - const index = this.#waiters.indexOf(resolver); - if (index >= 0) this.#waiters.splice(index, 1); - } + this.#removeWaiter(waiter); // 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; } } + + /** Remove a parked waiter from its key's queue, dropping empty keys. Idempotent. */ + #removeWaiter(waiter: AdmissionWaiter): void { + const queue = this.#queues.get(waiter.key); + if (!queue) return; + const index = queue.indexOf(waiter); + if (index >= 0) queue.splice(index, 1); + if (queue.length === 0) this.#removeFairKey(waiter.key); + } + + #removeFairKey(key: string): void { + this.#queues.delete(key); + const index = this.#fairKeys.indexOf(key); + if (index < 0) return; + this.#fairKeys.splice(index, 1); + if (index < this.#fairCursor) this.#fairCursor -= 1; + if (this.#fairKeys.length === 0) this.#fairCursor = 0; + } + + /** + * Round-robin dispatch across per-key queues (#9654 fairness, #10110 global + * budget). Called on every release; wakes exactly ONE waiter — the head of + * the next key in rotation — so the freed slot is claimed atomically by the + * woken waiter's re-loop. A strict FIFO would let one client's burst consume + * every freed slot; rotating the cursor gives each contending key a turn. + */ + #dispatchFair(): void { + if (this.#fairKeys.length === 0) return; + for (let i = 0; i < this.#fairKeys.length; i++) { + const key = this.#fairKeys[this.#fairCursor % this.#fairKeys.length]; + this.#fairCursor += 1; + const queue = this.#queues.get(key); + if (!queue || queue.length === 0) continue; + const waiter = queue.shift() as AdmissionWaiter; + if (queue.length === 0) this.#removeFairKey(key); + waiter.resolve(); + return; + } + } } const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); /** - * Per-connection virtual admission lanes (#9654). + * Process-wide byte-level admission budget (#10110). * - * Maps a sessionId (API-key hash or "anonymous") → ChatAdmissionController. - Each connection gets its own bounded heavyweight capacity so one connection - * cannot exhaust `CHAT_MAX_HEAVY_IN_FLIGHT` and starve others at the byte-level - * admission stage. + * Every request — every session, every API key — admits against ONE global + * ChatAdmissionController, so `CHAT_MAX_HEAVY_IN_FLIGHT` and + * `CHAT_ADMISSION_MAX_QUEUED_BYTES` are enforced process-wide, exactly as + * documented in docs/reference/ENVIRONMENT.md. The pre-#10110 design minted a + * per-session controller per request, multiplying the process bound by up to + * 64 lanes and letting unauthenticated fake credentials shard capacity. * - * Idle sessions are auto-evicted after OMNIROUTE_CHAT_VIRTUAL_TTL_MS - * (default 60s) to prevent unbounded Map growth. + * Per-request session identity survives ONLY as a fairness scheduling key: + * waiters are grouped per key and served round-robin against the shared + * budget (ChatAdmissionController#dispatchFair), preserving the #9654 + * guarantee that one connection's burst cannot starve others — without any + * per-key capacity being allocated. */ -const OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS = parsePositiveInt( - process.env.OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS, - 64 -); export function resolveSessionId(request: Request): string { - // Reuse the existing internal-bypass auth extraction: bearer token from - // Authorization, x-api-key (Anthropic-style), or Google API key header. + // Fairness scheduling key ONLY (never a capacity shard): hashed so raw key + // material never appears in diagnostics. Reuses the internal-bypass auth + // extraction: bearer token from Authorization, x-api-key (Anthropic-style), + // or Google API key header. const authHeader = request.headers.get("authorization") || ""; const bearerMatch = /^bearer\s+(\S+)$/i.exec(authHeader.trim()); if (bearerMatch) { @@ -271,106 +350,63 @@ export function resolveSessionId(request: Request): string { return "anonymous"; } -interface SessionRecord { - controller: ChatAdmissionController; - lastUsedMs: number; -} - export class PerConnectionAdmissionController { - #sessions = new Map(); - #evictionTimer: ReturnType | null = null; - readonly maxSessions: number; - readonly sessionTtlMs: number; + readonly #controller: ChatAdmissionController; constructor( - readonly maxHeavyPerSession: number, - opts?: { maxSessions?: number; sessionTtlMs?: number } + readonly maxHeavyInFlight = 1, + // Deprecated pre-#10110 lane-eviction knobs: accepted for API + // compatibility and ignored — there are no per-session lanes to evict. + _opts?: { maxSessions?: number; sessionTtlMs?: number } ) { - this.maxSessions = opts?.maxSessions ?? OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS; - this.sessionTtlMs = opts?.sessionTtlMs ?? OMNIROUTE_CHAT_VIRTUAL_TTL_MS; + this.#controller = new ChatAdmissionController(maxHeavyInFlight); } - getController(sessionId: string): ChatAdmissionController { - this.evictIfDue(); - const existing = this.#sessions.get(sessionId); - if (existing) { - existing.lastUsedMs = Date.now(); - return existing.controller; - } - // Evict oldest if at capacity (LRU fallback when TTL hasn't fired). - if (this.#sessions.size >= this.maxSessions) { - const oldestKey = this.oldestKey(); - if (oldestKey) this.#sessions.delete(oldestKey); - } - const controller = new ChatAdmissionController(this.maxHeavyPerSession); - this.#sessions.set(sessionId, { controller, lastUsedMs: Date.now() }); - this.armEviction(); - return controller; + /** Returns the process-global budget — the same instance for every session. */ + getController(_sessionId: string): ChatAdmissionController { + return this.#controller; } - /** Snapshot for observability — never exposes raw API keys. */ - snapshot(): ReadonlyArray<{ sessionId: string; activeHeavy: number; idleMs: number }> { - const now = Date.now(); - const arr: Array<{ sessionId: string; activeHeavy: number; idleMs: number }> = []; - for (const [sessionId, record] of this.#sessions) { - arr.push({ - sessionId, - activeHeavy: record.controller.activeHeavy, - idleMs: now - record.lastUsedMs, - }); - } - return arr; + /** + * Process-wide aggregate snapshot for observability: global totals plus + * per-key waiter depths. Keys are opaque scheduler keys, never raw + * credentials. + */ + snapshot(): { + activeHeavy: number; + queuedBytes: number; + waiting: number; + lanes: ReadonlyArray<{ key: string; waiting: number }>; + } { + return { + activeHeavy: this.#controller.activeHeavy, + queuedBytes: this.#controller.queuedBytes, + waiting: this.#controller.waitingCount, + lanes: this.#controller.waitersByKey, + }; } - get sessionCount(): number { - return this.#sessions.size; + get activeHeavy(): number { + return this.#controller.activeHeavy; } - private oldestKey(): string | undefined { - let oldest: string | undefined; - let oldestMs = Infinity; - for (const [key, record] of this.#sessions) { - // Use <= so that for equal timestamps, later-inserted entries win, - // preserving LRU semantics when Date.now() returns the same value. - if (record.lastUsedMs <= oldestMs) { - oldestMs = record.lastUsedMs; - oldest = key; - } - } - return oldest; + get queuedBytes(): number { + return this.#controller.queuedBytes; } - private evictIfDue(): void { - const now = Date.now(); - let evicted = false; - for (const [sessionId, record] of this.#sessions) { - if (now - record.lastUsedMs >= this.sessionTtlMs) { - this.#sessions.delete(sessionId); - evicted = true; - } - } - if (evicted) this.armEviction(); + get waitingCount(): number { + return this.#controller.waitingCount; } - private armEviction(): void { - if (this.#evictionTimer !== null) return; - this.#evictionTimer = setTimeout(() => { - this.#evictionTimer = null; - this.evictIfDue(); - }, this.sessionTtlMs).unref(); - } - - /** Force cleanup of all sessions (used by shutdown / tests). */ + /** No per-session state to clean; kept for API compatibility. */ dispose(): void { - this.#sessions.clear(); - if (this.#evictionTimer !== null) { - clearTimeout(this.#evictionTimer); - this.#evictionTimer = null; - } + // Intentionally empty: the process-global controller owns no session state. } } -export const perConnectionAdmissionController = new PerConnectionAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); +export const perConnectionAdmissionController = new PerConnectionAdmissionController( + CHAT_MAX_HEAVY_IN_FLIGHT +); export type ChatRequestAdmission = | { admit: true; request: Request; lease: ChatAdmissionLease | null } @@ -530,7 +566,8 @@ export async function admitChatStructure( const acquired = await controller.acquireHeavyWithin( options.queueMs ?? 0, options.signal, - CHAT_LARGE_BODY_BYTES + CHAT_LARGE_BODY_BYTES, + options.sessionId ); return acquired ? { admit: true, lease: acquired } @@ -700,7 +737,7 @@ export async function admitChatRequest( let lease: ChatAdmissionLease | null = null; const reserve = async (bytes = 0): Promise => { if (lease) return true; - lease = await controller.acquireHeavyWithin(queueMs, request.signal, bytes); + lease = await controller.acquireHeavyWithin(queueMs, request.signal, bytes, sessionId); return lease !== null; }; diff --git a/tests/unit/chat-body-admission-aggregate-10110.test.ts b/tests/unit/chat-body-admission-aggregate-10110.test.ts new file mode 100644 index 0000000000..22e67fd763 --- /dev/null +++ b/tests/unit/chat-body-admission-aggregate-10110.test.ts @@ -0,0 +1,271 @@ +// #10110: Aggregate process-wide bounds for byte-level chat admission. +// +// The always-on per-connection admission layer (#9940 / #9654) enforces +// CHAT_MAX_HEAVY_IN_FLIGHT and CHAT_ADMISSION_MAX_QUEUED_BYTES PER LANE, so the +// documented "in one process" contract (docs/reference/ENVIRONMENT.md:193) is +// multiplied by OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS (default 64): up to 64 +// concurrent heavy requests and 256 MiB of parked bodies process-wide. +// +// These tests assert the AGGREGATE contract that the issue's acceptance +// criteria demand. They are intentionally deterministic (exact === assertions, +// no <= fudge) and are RED on release/v3.8.50 — they only pass once the byte +// level admits against one process-global budget with per-key fairness. +import test from "node:test"; +import assert from "node:assert/strict"; + +const admissionModule = await import("../../src/shared/middleware/chatBodyAdmission.ts"); +const { PerConnectionAdmissionController, CHAT_ADMISSION_MAX_QUEUED_BYTES } = admissionModule; + +const GLOBAL_QUEUED_BUDGET = CHAT_ADMISSION_MAX_QUEUED_BYTES; // 4 MiB default + +// Aggregate across distinct controllers. With the fix every key resolves to +// ONE process-global controller (shared budget), so dedupe-by-identity yields +// the true process-wide totals — never double-counted, never multiplied by +// the number of keys. +function aggregateActiveHeavy( + pc: InstanceType, + keys: string[] +): number { + const seen = new Set(); + let total = 0; + for (const key of keys) { + const controller = pc.getController(key); + if (seen.has(controller)) continue; + seen.add(controller); + total += controller.activeHeavy; + } + return total; +} + +function aggregateQueuedBytes( + pc: InstanceType, + keys: string[] +): number { + const seen = new Set(); + let total = 0; + for (const key of keys) { + const controller = pc.getController(key); + if (seen.has(controller)) continue; + seen.add(controller); + total += controller.queuedBytes; + } + return total; +} + +// ── Family 1: active-LRU eviction must not mint replacement capacity ────── +// Issue repro: maxSessions=1, key A acquires; key B admission LRU-evicts A; +// A re-admits and gets a FRESH controller with fresh capacity while the old +// lease still holds → effectiveActiveForA = 2. The aggregate must stay 1. + +test("LRU eviction of a live lane does not mint a second capacity slot", () => { + // Red on release/v3.8.50: B's admission LRU-evicts A's lane, and A's re-admit + // gets a FRESH controller with fresh capacity while the old lease still holds. + // With the fix there are no per-session lanes at all: getController returns the + // one shared process-global controller, so no capacity can ever be minted. + const pc = new PerConnectionAdmissionController(1, { maxSessions: 1, sessionTtlMs: 60_000 }); + + const ctrlA1 = pc.getController("A"); + const leaseA = ctrlA1.tryAcquireHeavy(); + assert.ok(leaseA, "A acquires the only slot"); + + // B's admission must NOT evict the lane holding a live lease; and even if + // the lane is retired/recreated, it must not mint fresh capacity. + pc.getController("B"); + + // A re-admits: no fresh capacity may appear while the old lease is live. + const ctrlA2 = pc.getController("A"); + assert.equal( + ctrlA2.tryAcquireHeavy(), + null, + "a live lease must keep its slot; no second capacity may be minted" + ); + + // Aggregate active heavy across every lane stays at the process-wide bound. + assert.equal( + aggregateActiveHeavy(pc, ["A", "B"]), + 1, + "process-wide active heavy must be 1, not 2 (orphaned lease + minted slot)" + ); + + leaseA.release(); +}); + +// ── Family 2: active-TTL eviction must not mint replacement capacity ────── +// Same invariant via the idle-TTL path: a lane that still holds a live lease +// must not be evicted (or, if retired, must not hand out fresh capacity). + +test("TTL eviction of a live lane does not mint a second capacity slot", async () => { + // Red on release/v3.8.50: the idle-TTL evicts A's lane mid-lease; a re-admit + // then mints a fresh controller with fresh capacity (orphaned lease + new slot). + // With the fix the shared controller outlives any session and never mints. + const pc = new PerConnectionAdmissionController(1, { maxSessions: 64, sessionTtlMs: 50 }); + + const ctrlA1 = pc.getController("A"); + const leaseA = ctrlA1.tryAcquireHeavy(); + assert.ok(leaseA, "A acquires the only slot"); + + // Wait past the idle TTL so evictIfDue() would mark A's lane stale. + await new Promise((resolve) => setTimeout(resolve, 120)); + + // Re-admitting A must not produce a controller with fresh capacity. + const ctrlA2 = pc.getController("A"); + assert.equal( + ctrlA2.tryAcquireHeavy(), + null, + "a live lease must survive TTL; no fresh capacity may be minted" + ); + assert.equal( + aggregateActiveHeavy(pc, ["A"]), + 1, + "process-wide active heavy must stay 1 after TTL with a live lease" + ); + + leaseA.release(); +}); + +// ── Family 3: aggregate parked bytes stay within the process-wide budget ── +// Regression guard for the byte side of the multiplication: waiters parked +// from DIFFERENT lanes must share ONE process-wide queued-bytes budget. +// +// Note on shape: on the buggy code an idle lane never parks (its waiter +// acquires on its own free capacity instantly), so cross-lane bytes are only +// observable while multiple lanes are simultaneously busy — an arrangement +// the global-budget fix makes impossible by construction. The active-heavy +// families (1/2/4) are the RED probes; this family locks in the byte budget +// once the shared budget exists: one busy slot + waiters parked from two +// lanes, aggregate must never exceed the single process-wide budget. + +test("parked bytes across lanes share one process-wide budget", async () => { + const pc = new PerConnectionAdmissionController(1, { maxSessions: 64, sessionTtlMs: 60_000 }); + + // One lane holds the single busy slot. + const ctrlA = pc.getController("A"); + const heldA = ctrlA.tryAcquireHeavy(); + assert.ok(heldA); + + // Two waiters park — one keyed A, one keyed B — against the SAME busy slot. + // Together they must respect the single process-wide budget. + const wA = ctrlA.acquireHeavyWithin(2_000, undefined, GLOBAL_QUEUED_BUDGET, "A"); + const wB = pc.getController("B").acquireHeavyWithin(2_000, undefined, GLOBAL_QUEUED_BUDGET, "B"); + await new Promise((resolve) => setTimeout(resolve, 30)); + + const aggregate = aggregateQueuedBytes(pc, ["A", "B"]); + assert.ok( + aggregate <= GLOBAL_QUEUED_BUDGET, + `aggregate queued bytes (${aggregate}) must stay within the single process-wide budget (${GLOBAL_QUEUED_BUDGET})` + ); + + heldA.release(); + const leases = await Promise.all([wA, wB]); + for (const lease of leases) lease?.release(); + assert.equal(aggregateQueuedBytes(pc, ["A", "B"]), 0, "all parked bytes released"); +}); + +test("a 16 MiB per-lane config still respects the process-wide byte budget", async () => { + // The issue's 1 GiB scenario shape: per-lane budgets that would multiply + // into 1 GiB must instead be capped by the single process-wide budget. + const MiB = 1024 * 1024; + const pc = new PerConnectionAdmissionController(1, { maxSessions: 64, sessionTtlMs: 60_000 }); + + const ctrlA = pc.getController("A"); + const heldA = ctrlA.tryAcquireHeavy(); + assert.ok(heldA); + + const wA = ctrlA.acquireHeavyWithin(2_000, undefined, 16 * MiB, "A"); + const wB = pc.getController("B").acquireHeavyWithin(2_000, undefined, 16 * MiB, "B"); + await new Promise((resolve) => setTimeout(resolve, 30)); + + const aggregate = aggregateQueuedBytes(pc, ["A", "B"]); + assert.ok( + aggregate <= GLOBAL_QUEUED_BUDGET, + `16 MiB per-lane config must still respect the process-wide budget; aggregate was ${aggregate}` + ); + + heldA.release(); + const leases = await Promise.all([wA, wB]); + for (const lease of leases) lease?.release(); +}); + +// ── Family 4: same-session recreation waits on the global slot ──────────── +// A session that released and re-admits while ANOTHER session holds the +// process-wide slot must queue, not bypass. + +test("same-session recreation waits while another session holds the global slot", () => { + const pc = new PerConnectionAdmissionController(1, { maxSessions: 64, sessionTtlMs: 60_000 }); + + const ctrlA = pc.getController("A"); + const leaseA = ctrlA.tryAcquireHeavy(); + assert.ok(leaseA); + leaseA.release(); // A releases; the global slot is free. + + // B takes the single process-wide slot. + const ctrlB = pc.getController("B"); + const leaseB = ctrlB.tryAcquireHeavy(); + assert.ok(leaseB); + + // A re-admits while B holds the slot → must wait, not bypass. + assert.equal( + ctrlA.tryAcquireHeavy(), + null, + "recreated A must not bypass the process-wide slot held by B" + ); + assert.equal( + aggregateActiveHeavy(pc, ["A", "B"]), + 1, + "aggregate active heavy is 1 with B holding the slot" + ); + + leaseB.release(); +}); + +// ── Fairness guard (B2): lanes are served round-robin over the shared budget ─ +// A session that queues a burst must not consume every dispatch turn: when A +// queues two waiters and B queues one behind the same busy slot, B's waiter +// must be served BEFORE A's second follow-up (round-robin across lanes, the +// adaptive dispatchLanes precedent). A strict single FIFO would serve A-A-B +// and starve B under sustained load. + +test("one session's burst does not starve another session's bounded wait", async () => { + const pc = new PerConnectionAdmissionController(1, { maxSessions: 64, sessionTtlMs: 60_000 }); + + // A holds the single slot and queues two follow-ups. + const ctrlA = pc.getController("A"); + const leaseA = ctrlA.tryAcquireHeavy(); + assert.ok(leaseA); + + const aWaiters = [ + ctrlA.acquireHeavyWithin(2_000, undefined, 0, "A"), + ctrlA.acquireHeavyWithin(2_000, undefined, 0, "A"), + ]; + + // B queues one bounded wait behind the shared slot. + const ctrlB = pc.getController("B"); + const bWaiter = ctrlB.acquireHeavyWithin(2_000, undefined, 0, "B"); + + // Free the slot, then release each lease the moment it arrives so the next + // waiter can proceed. Record acquisition order. + leaseA.release(); + const order: string[] = []; + const track = (label: string) => (lease: unknown) => { + if (lease) { + order.push(label); + (lease as { release: () => void }).release(); + } + }; + void aWaiters[0].then(track("a1")); + void aWaiters[1].then(track("a2")); + void bWaiter.then(track("b1")); + await Promise.all([...aWaiters, bWaiter]); + + assert.equal( + order.length, + 3, + "all three queued sessions must acquire within the bounded wait; none starve" + ); + assert.equal( + order.indexOf("b1"), + 1, + `round-robin must serve B before A's second follow-up (strict FIFO would starve B); got order ${order.join(" -> ")}` + ); + assert.equal(aggregateActiveHeavy(pc, ["A", "B"]), 0); +}); diff --git a/tests/unit/per-connection-admission-9654.test.ts b/tests/unit/per-connection-admission-9654.test.ts index 8f26b88129..ccfd494c0a 100644 --- a/tests/unit/per-connection-admission-9654.test.ts +++ b/tests/unit/per-connection-admission-9654.test.ts @@ -1,4 +1,10 @@ -// #9654: Per-connection virtual admission lanes +// #9654/#10110: process-wide admission budget with per-key fair scheduling. +// +// The pre-#10110 per-connection lanes minted a controller per session, so the +// documented "in one process" heavy/bytes bound multiplied by up to 64 lanes +// (#10110). The fix removes the lanes entirely: every session resolves to ONE +// process-global controller, and per-request session identity is used only as +// a fairness scheduling key (round-robin dispatch, not capacity allocation). import test from "node:test"; import assert from "node:assert/strict"; @@ -10,7 +16,6 @@ const { admitChatStructure, perConnectionAdmissionController, ChatAdmissionController, - CHAT_MAX_HEAVY_IN_FLIGHT, } = admissionModule; function makeRequest(headers: Record, body = "{}"): Request { @@ -50,21 +55,26 @@ test("resolveSessionId does not leak raw API key in the session ID", () => { assert.ok(!sid.includes("secret")); }); -test("PerConnectionAdmissionController isolates capacity across sessions", () => { +// ── #10110: one process-global budget, shared by every session ──────────── +// The pre-fix lanes isolated capacity per session (up to 64× the process +// bound, and fake credentials could shard capacity). The fix returns the SAME +// controller for every session so the process-wide bound is real. + +test("PerConnectionAdmissionController shares ONE global budget across sessions", () => { const pc = new PerConnectionAdmissionController(1); const ctrlA = pc.getController("session-a"); const ctrlB = pc.getController("session-b"); - // Session A acquires the only slot + // A acquires the only process-wide slot. const leaseA = ctrlA.tryAcquireHeavy(); assert.ok(leaseA); - // Session A is now full - assert.equal(ctrlA.tryAcquireHeavy(), null); - // Session B still has capacity — isolation works - const leaseB = ctrlB.tryAcquireHeavy(); - assert.ok(leaseB); + // B shares the SAME budget: no per-session capacity is available while A + // holds the single process-wide slot (pre-#10110 this minted B its own slot). + assert.equal(ctrlB.tryAcquireHeavy(), null); leaseA.release(); - leaseB.release(); + // Slot freed → either session may now acquire. + assert.ok(ctrlB.tryAcquireHeavy()); + ctrlB.tryAcquireHeavy()?.release(); }); test("PerConnectionAdmissionController returns same controller for same session", () => { @@ -74,83 +84,93 @@ test("PerConnectionAdmissionController returns same controller for same session" assert.equal(a1, a2); }); -test("PerConnectionAdmissionController creates new controller for new session", () => { +test("PerConnectionAdmissionController returns the same controller for ALL sessions", () => { const pc = new PerConnectionAdmissionController(1); const a = pc.getController("session-a"); const b = pc.getController("session-b"); - assert.notEqual(a, b); + const c = pc.getController("session-c"); + // The shared process-global budget is a single instance (#10110): no per-key + // controllers exist to multiply the bound. + assert.equal(a, b); + assert.equal(b, c); }); -test("PerConnectionAdmissionController enforces maxSessions LRU eviction", () => { - const pc = new PerConnectionAdmissionController(1, { maxSessions: 2, sessionTtlMs: 60000 }); - const a = pc.getController("a"); +test("PerConnectionAdmissionController never evicts a live lease (no lane lifecycle)", () => { + // The pre-#10110 LRU/TTL lane eviction could drop a controller mid-lease and + // mint a fresh one on re-admit — silently doubling capacity (#10110). With a + // single shared controller there is nothing to evict and nothing to mint. + const pc = new PerConnectionAdmissionController(1, { maxSessions: 1, sessionTtlMs: 50 }); + const ctrlA1 = pc.getController("a"); + const leaseA = ctrlA1.tryAcquireHeavy(); + assert.ok(leaseA); + + // Touching another session (LRU pressure) and waiting past the TTL must not + // replace the controller holding the live lease. const b = pc.getController("b"); - assert.equal(pc.sessionCount, 2); - // Touch 'a' so 'b' is oldest + assert.equal(b, ctrlA1, "same global controller, no per-session lane to evict"); const aAgain = pc.getController("a"); - assert.equal(aAgain, a, "same a reference"); - // Creating 'c' should evict 'b' (oldest) - const c = pc.getController("c"); - assert.equal(pc.sessionCount, 2); - // 'a' survives, 'b' is evicted - const aAfter = pc.getController("a"); - assert.equal(aAfter, a, "a should still exist after c added"); - // 'b' gets a fresh controller (old one was evicted) - const newB = pc.getController("b"); - assert.notEqual(newB, b, "b should be evicted and recreated"); + assert.equal(aAgain, ctrlA1, "controller identity is stable across TTL"); + assert.equal(aAgain.tryAcquireHeavy(), null, "live lease keeps the only slot"); + leaseA.release(); }); -test("PerConnectionAdmissionController evicts idle sessions after TTL", async () => { - const pc = new PerConnectionAdmissionController(1, { - sessionTtlMs: 50, - maxSessions: 64, - }); - const ctrl = pc.getController("idle-session"); - assert.ok(ctrl); - assert.equal(pc.sessionCount, 1); - - // Wait past TTL + eviction tick - await new Promise((resolve) => setTimeout(resolve, 120)); - // Accessing again should trigger eviction → fresh controller - const fresh = pc.getController("idle-session"); - assert.notEqual(fresh, ctrl); -}); - -test("PerConnectionAdmissionController snapshot does not leak raw keys", () => { +test("PerConnectionAdmissionController snapshot reports process-wide aggregates", async () => { const pc = new PerConnectionAdmissionController(1); pc.getController("key_abc123"); pc.getController("anonymous"); + + const empty = pc.snapshot(); + assert.equal(empty.activeHeavy, 0); + assert.equal(empty.queuedBytes, 0); + assert.equal(empty.waiting, 0); + assert.deepEqual(empty.lanes, []); + + // Occupy the global slot and park a waiter from a second key. + const ctrlA = pc.getController("key_abc123"); + const leaseA = ctrlA.tryAcquireHeavy(); + assert.ok(leaseA); + const wB = pc.getController("anonymous").acquireHeavyWithin(500, undefined, 100, "anonymous"); + await new Promise((resolve) => setTimeout(resolve, 30)); + const snap = pc.snapshot(); - assert.equal(snap.length, 2); - for (const entry of snap) { - assert.ok(typeof entry.sessionId === "string"); - assert.ok(entry.sessionId.includes("key_abc123") || entry.sessionId === "anonymous"); - assert.ok(typeof entry.activeHeavy === "number"); - assert.ok(typeof entry.idleMs === "number"); + assert.equal(snap.activeHeavy, 1); + assert.equal(snap.queuedBytes, 100); + assert.equal(snap.waiting, 1); + assert.ok(Array.isArray(snap.lanes)); + for (const lane of snap.lanes) { + assert.ok(typeof lane.key === "string"); + assert.ok(lane.waiting === 0 || lane.waiting === 1); + // Keys are opaque hashed scheduler keys — raw credentials never appear. + assert.ok(!lane.key.includes("secret")); } + + const waiterLease = await wB; + waiterLease?.release(); + leaseA.release(); }); test("admitChatRequest uses per-connection controller by default", async () => { - const result = await admitChatRequest( - makeRequest({ authorization: "Bearer sk-test-key" }), - { largeBodyBytes: 32, hardMaxBytes: 1024 } - ); + const result = await admitChatRequest(makeRequest({ authorization: "Bearer sk-test-key" }), { + largeBodyBytes: 32, + hardMaxBytes: 1024, + }); assert.equal(result.admit, true); if (result.admit) result.lease?.release(); }); test("admitChatRequest with explicit controller overrides per-connection lookup", async () => { const explicitController = new ChatAdmissionController(1); - const result = await admitChatRequest( - makeRequest({ authorization: "Bearer sk-test-key" }), - { controller: explicitController, largeBodyBytes: 32, hardMaxBytes: 1024 } - ); + const result = await admitChatRequest(makeRequest({ authorization: "Bearer sk-test-key" }), { + controller: explicitController, + largeBodyBytes: 32, + hardMaxBytes: 1024, + }); assert.equal(result.admit, true); if (result.admit) result.lease?.release(); }); test("admitChatStructure routes structural rejection to per-connection controller", async () => { - // occupy sess-a's per-connection controller via the module-level instance + // occupy sess-a's controller — which is the shared process-global budget const controller = perConnectionAdmissionController.getController("sess-a"); const occupied = controller.tryAcquireHeavy(); assert.ok(occupied); @@ -168,7 +188,7 @@ test("admitChatStructure routes structural rejection to per-connection controlle heavyTokens: 10_000, } ); - // Session A is busy → 503 + // The process-wide slot is busy → 503 assert.equal(result.admit, false); if (result.admit) return; assert.equal(result.response.status, 503); @@ -176,13 +196,14 @@ test("admitChatStructure routes structural rejection to per-connection controlle occupied.release(); }); -test("admitChatStructure with different sessionId gets independent capacity", async () => { - // occupy sess-a's per-connection controller +test("admitChatStructure with different sessionId shares the global budget", async () => { + // occupy the shared process-global budget via sess-a const ctrlA = perConnectionAdmissionController.getController("sess-a"); const occupied = ctrlA.tryAcquireHeavy(); assert.ok(occupied); - // Session B should get its own controller → admitted + // Session B must NOT get independent capacity (pre-#10110 it did — that was + // the defect): it shares the one process-wide slot and must be rejected. const result = await admitChatStructure( { messages: Array.from({ length: 500 }, () => ({ role: "user", content: "x" })), @@ -196,10 +217,7 @@ test("admitChatStructure with different sessionId gets independent capacity", as heavyTokens: 32_000, } ); - assert.equal(result.admit, true); - if (result.admit) { - assert.notEqual(result.lease, null); - result.lease?.release(); - } + assert.equal(result.admit, false); + assert.equal(result.response.status, 503); occupied.release(); });