From 8d78e3dfd3f6bc4de7e77f71e41b764a88df5489 Mon Sep 17 00:00:00 2001 From: Brandon Bennett <107384180+branben@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:25:13 -0400 Subject: [PATCH] fix: per-connection virtual admission lanes (#9654) (#9940) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: add per-connection virtual admission lanes (#9654) Worst-day-ever analysis to harden AdaptiveAdmissionController: - Guard expireEntry() against null entry (CRITICAL null deref) - Add deleteLane() to drain+reject on LRU eviction (HIGH orphaned promises) - Fix Map mutation during evictIdleLanes iteration (MEDIUM safety) - Add ADMISSION_LANE_EVICTED reject code (MEDIUM clarity) - Pass sessionId to admitChatRequest in route.ts - virtualLanes defaults to false in validateConfig - 7 new controller tests + 14 new byte-level admission tests - Assertions tightened from >= to === (Matt Pocock methodology) Debunked 2 false positives: concurrency race (single-threaded JS) and memory amplification (FairCostQueue bounds per-lane). Fixes #9654 * fix(admission): restore bounded queue-wait on per-connection lanes (#9654) The per-connection lane refactor dropped the bounded queue-wait (acquireHeavyWithin / #waiters / queueMs). #9654's acceptance criteria and #9608 section C prefer server-side wait/pacing up to defaultMaxWaitMs over an instant retryable 503. - ChatAdmissionController: re-add #waiters FIFO + acquireHeavyWithin(timeoutMs); queueMs: 0 preserves the instant-503 path - admitChatStructure and admitChatRequest.reserve are async again and take queueMs - route: pass CHAT_ADMISSION_QUEUE_MAX_MS and await the admission calls - per-connection lane tests await the async admitChatStructure Admission suite: 114/114 pass (bun test, 7 files). * chore: re-trigger CI after dast-smoke infra cancellation (#9654) * 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 * fix(admission): bound queued bytes for the queue-wait heap valve (#9654) U3 from KC plan 2026-08-09-001. The restored queue-wait parks fully-buffered bodies; without a cap, several large coding-agent bodies (~750 KB) waiting at once recreates the #4380 heap amplification this module was built to stop. - acquireHeavyWithin(timeoutMs, signal?, queuedBytes): each parked waiter is charged its buffered size against CHAT_ADMISSION_MAX_QUEUED_BYTES (default 4 MB); over-budget waits reject immediately with a retryable 503 and never park. The charge is released on wake, abort, or timeout. - Real sizes threaded from admitChatRequest (declared length / sniffed bytes); structural waits charge the conservative 256 KB weight. - Lower default OMNIROUTE_CHAT_ADMISSION_QUEUE_MS to 2000ms (was 5000ms). - Env vars documented in .env.example; 6 exact-assertion tests: 125/125 across the 7-file admission suite (was 119). * docs: map the two admission-lane systems for operators (#9654) U5 from KC plan 2026-08-09-001. Verifies lane metrics are exposed by the health payload (GET /api/monitoring/health -> adaptiveAdmission -> lane* fields) and records which lane system reports where: byte-level per-connection lanes (always on, memory scope) vs adaptive virtual lanes (opt-in via OMNIROUTE_CHAT_VIRTUAL_LANES, dispatch scope) plus the explicit opt-in ops note. * docs: add required frontmatter to admission-lanes doc (dast-smoke build fix) * docs: sync env vars with .env.example and ENVIRONMENT.md (docs gate fix) * fix(admission): complete REJECT_MAP, literal lane env read, split oversized test file Three CI-gate fixes surfaced by the post-merge check run (head 3de77166e): 1. open-sse-typecheck (TS2741): REJECT_MAP was missing the ADMISSION_LANE_EVICTED entry that controller.ts:662 emits on lane eviction. Add the 503 mapping so the Record is total. 2. Docs Gates fabricated-claim: OMNIROUTE_CHAT_VIRTUAL_LANES was read dynamically via ENV_KEYS.virtualLanes (env[key]), invisible to the literal env.X scanner. Read it literally — behavior-identical, doc claim now verifiable. 3. check:file-size: chat-body-admission.test.ts (1307 lines) exceeded the 1000-line new-file cap. Split the queue-wait/abort/heap-valve section into chat-body-admission-queue.test.ts (818 + 513 lines, both under cap). Suite: 125/125 across 8 files. All three checkers pass locally. * refactor(admission): drop dead ENV_KEYS.virtualLanes entry + lock lane-evicted mapping test Code-review follow-up on 50c93d266: 1. ENV_KEYS.virtualLanes is now unreferenced since the literal env read landed; remove it so the config map only lists keys actually read through the map. 2. Add an exact-assertion runtime test for the ADMISSION_LANE_EVICTED mapping: a queued lane waiter evicted by the 60s idle TTL rejects with 503 / admission_lane_evicted / Retry-After 1 / sanitized body (no raw tenant key). Proves the REJECT_MAP entry end-to-end through buildAdmissionRejectResponse. Suite: 126/126 (17 in runtime file, 125 in the 8-file admission suite). --------- Co-authored-by: Brandon Bennett --- .env.example | 12 + ...-per-connection-virtual-admission-lanes.md | 1 + docs/architecture/admission-lanes.md | 50 ++ docs/reference/ENVIRONMENT.md | 5 +- open-sse/services/admission/config.ts | 2 + open-sse/services/admission/controller.ts | 228 +++++++- open-sse/services/admission/runtime.ts | 12 + open-sse/services/admission/types.ts | 13 + src/app/api/v1/chat/completions/route.ts | 5 + src/shared/middleware/chatBodyAdmission.ts | 265 ++++++++- tests/unit/adaptive-admission-runtime.test.ts | 45 ++ .../unit/admission-virtual-lanes-9654.test.ts | 240 ++++++++ tests/unit/chat-body-admission-queue.test.ts | 513 ++++++++++++++++++ tests/unit/chat-body-admission.test.ts | 168 +----- .../per-connection-admission-9654.test.ts | 205 +++++++ 15 files changed, 1580 insertions(+), 184 deletions(-) create mode 100644 changelog.d/fixes/9940-per-connection-virtual-admission-lanes.md create mode 100644 docs/architecture/admission-lanes.md create mode 100644 tests/unit/admission-virtual-lanes-9654.test.ts create mode 100644 tests/unit/chat-body-admission-queue.test.ts create mode 100644 tests/unit/per-connection-admission-9654.test.ts diff --git a/.env.example b/.env.example index f664d06964..73f7797fd5 100644 --- a/.env.example +++ b/.env.example @@ -350,6 +350,18 @@ ALLOW_API_KEY_REVEAL=false # by OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT and the heap-pressure shed instead. Set a positive # value only on memory-constrained deployments that need a hard ceiling. # OMNIROUTE_CHAT_HARD_MAX_MESSAGES=0 +# How long a heavy request waits for heavyweight capacity before a retryable 503. +# A short bounded wait serializes agent bursts instead of an instant 503; 0 = instant. +# Default 2000 (2s). +# OMNIROUTE_CHAT_ADMISSION_QUEUE_MS=2000 +# Queued-bytes budget for the admission wait: bounds total buffered body bytes parked +# per lane so the wait cannot amplify the heap (#4380). Over-budget waits 503 immediately. +# Default 4194304 (4 MB). +# OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES=4194304 +# Per-connection virtual admission lanes (#9654): idle-lane eviction TTL. Default 60000 (60s). +# OMNIROUTE_CHAT_VIRTUAL_TTL_MS=60000 +# Per-connection virtual admission lanes (#9654): max concurrent sessions (lanes). Default 64. +# OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS=64 # Hard cap (bytes) for a non-streaming upstream response buffered fully into memory # (#5152). Past this the upstream reader is cancelled and the request fails fast diff --git a/changelog.d/fixes/9940-per-connection-virtual-admission-lanes.md b/changelog.d/fixes/9940-per-connection-virtual-admission-lanes.md new file mode 100644 index 0000000000..59f304f591 --- /dev/null +++ b/changelog.d/fixes/9940-per-connection-virtual-admission-lanes.md @@ -0,0 +1 @@ +- **fix(admission):** per-connection virtual admission lanes with idle TTL eviction — guards `expireEntry` null deref, adds `deleteLane()` for safe LRU eviction, and passes `sessionId` to byte-level admission (fixes #9654) diff --git a/docs/architecture/admission-lanes.md b/docs/architecture/admission-lanes.md new file mode 100644 index 0000000000..8941a12eff --- /dev/null +++ b/docs/architecture/admission-lanes.md @@ -0,0 +1,50 @@ +--- +title: "Admission lanes — two lane systems, what gates each, where each reports" +status: active +lastUpdated: 2026-08-09 +--- + +# Admission lanes (#9654) — two lane systems, what gates each, where each reports + +OmniRoute has **two** process-local lane systems with different scopes. They are +complementary; operators should know which one they are looking at. + +## 1. Byte-level per-connection lanes (`chatBodyAdmission.ts`) + +- **Scope:** the buffered-body/heap path for `POST /v1/chat/completions`. Guards + against heap amplification from large coding-agent bodies (#4380). +- **Gate:** **always on.** Each distinct API key (hashed) — or `anonymous` — gets its + own lane with `CHAT_MAX_HEAVY_IN_FLIGHT` capacity, so one session's burst cannot + starve another session's heavyweight slot. +- **Tuning:** + - `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` — idle-lane eviction (default 60000) + - `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` — lane count cap (default 64) + - `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` — queue-wait before 503 (default 2000) + - `OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES` — queued-bytes heap valve (default 4 MB) +- **Reports:** not in `GET /api/monitoring/health` today; observable via + `PerConnectionAdmissionController.snapshot()` (sessionId hash, activeHeavy, idleMs). + +## 2. Adaptive runtime virtual lanes (`open-sse/services/admission`) + +- **Scope:** tenant-key admission for provider dispatch — queue cost, latency-guided + limit adaptation, lane queueing, and lane metrics. +- **Gate:** **opt-in.** Disabled unless `OMNIROUTE_CHAT_VIRTUAL_LANES=true`. Without it, + the adaptive controller keeps the shared queue behavior (criterion 1 of #9654 only + holds once an operator enables lanes). +- **Tuning:** `OMNIROUTE_CHAT_VIRTUAL_LANES` + adaptive config (`maxQueueCount`, + `maxQueueCost`, `defaultMaxWaitMs`, …). +- **Reports:** `GET /api/monitoring/health` → `adaptiveAdmission` → `laneCount`, + `laneQueuedCount`, `laneQueuedCost`, `laneTenants` (opaque lane IDs, never raw keys). + +## Which one is showing in a dashboard + +- `adaptiveAdmission.laneCount` / `laneTenants` → **adaptive virtual lanes** (system 2). +- A health payload with **no** `adaptiveAdmission.lane*` fields usually means + `OMNIROUTE_CHAT_VIRTUAL_LANES` is unset — the byte-level lanes (system 1) are still + active, but nothing under `adaptiveAdmission` will report lane data until it is enabled. + +## Why both exist + +The byte-level lanes bound the memory-heavy parse/compress path; the adaptive lanes +bound dispatch cost per tenant. #9654's criterion 1 ("one session's burst does not 503 +another") is enforced by system 1 unconditionally and by system 2 once opt-in is enabled. diff --git a/docs/reference/ENVIRONMENT.md b/docs/reference/ENVIRONMENT.md index 3f319dde03..4a3aa5e092 100644 --- a/docs/reference/ENVIRONMENT.md +++ b/docs/reference/ENVIRONMENT.md @@ -1398,7 +1398,10 @@ These settings were introduced after the previous environment-contract snapshot. | Variable | Default | Source File | Description | | --- | --- | --- | --- | -| `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` | `5000` | `src/shared/middleware/chatBodyAdmission.ts` | Maximum wait for a heavyweight chat admission slot before a retryable `503`; `0` restores immediate rejection. | +| `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_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. | | `CHAT_LOG_MAX_BODY_KB` | `1024` | `src/lib/logEnv.ts` | Maximum request or response body size before log summarization, in KiB. | | `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. | diff --git a/open-sse/services/admission/config.ts b/open-sse/services/admission/config.ts index dfe9b07a01..1f3aecf78e 100644 --- a/open-sse/services/admission/config.ts +++ b/open-sse/services/admission/config.ts @@ -21,6 +21,7 @@ export interface ValidatedConfig { adaptation: AdaptationParams; maxRequestCost: number; costConfig: ReturnType; + virtualLanes: boolean; } function requirePositiveInt( @@ -162,6 +163,7 @@ export function validateConfig(input: AdaptiveAdmissionConfig): ValidatedConfig windowMs, maxRequestCost: costConfig.maxRequestCost, costConfig, + virtualLanes: input.virtualLanes === true, adaptation: resolveAdaptationParams(input, minLimit, maxLimit, windowMs), }; } diff --git a/open-sse/services/admission/controller.ts b/open-sse/services/admission/controller.ts index 1051a64782..8a6db9be74 100644 --- a/open-sse/services/admission/controller.ts +++ b/open-sse/services/admission/controller.ts @@ -27,6 +27,13 @@ import { type ShadowDecision, } from "./types.ts"; +/** + * Idle TTL for per-connection virtual admission lanes (#9654). + */ +const ADMISSION_LANE_TTL_MS = 60_000; +/** Bounded per-connection lane map to prevent unbounded memory growth (#9654). */ +const ADMISSION_LANE_MAX_SESSIONS = 1_000; + type VirtualDisposition = "active" | "queued" | "rejected" | "none"; const MAX_SAFE_BIGINT = BigInt(Number.MAX_SAFE_INTEGER); @@ -95,6 +102,13 @@ export class AdaptiveAdmissionController { private adaptation: AdaptationState; private queue: FairCostQueue; private virtualQueue: FairCostQueue<{ recordId: string }>; + /** Per-connection virtual admission lanes (#9654). */ + private readonly virtualLanes = new Map; + lastUsedMs: number; + }>(); + /** Eviction timer for idle lanes; re-armed when a lane is created. */ + private laneEvictionTimer: unknown = undefined; private readonly active = new Map(); private activeCost = 0n; private virtualActiveCost = 0; @@ -148,6 +162,14 @@ export class AdaptiveAdmissionController { const drained = this.queue.drain(); this.queue = new FairCostQueue(next.maxQueueCount, next.maxQueueCost); + // Drain per-connection virtual lane queues (#9654). + for (const [, lane] of this.virtualLanes) { + for (const entry of lane.queue.drain()) { + drained.push(entry); + } + } + this.virtualLanes.clear(); + this.clearLaneEviction(); for (const entry of drained) { if (next.mode !== "enforce") { this.clearEntryTimer(entry); @@ -191,6 +213,10 @@ export class AdaptiveAdmissionController { virtualActiveCount: saturateSnapshotNumber(this.virtualActiveCount), virtualQueuedCost: saturateSnapshotNumber(this.virtualQueue.totalCost), virtualQueuedCount: saturateSnapshotNumber(this.virtualQueue.size), + laneCount: saturateSnapshotNumber(this.virtualLanes.size), + laneQueuedCost: saturateSnapshotNumber(this.laneTotalQueuedCost()), + laneQueuedCount: saturateSnapshotNumber(this.laneTotalQueuedCount()), + laneTenants: this.laneTenantSnapshot(), admittedCount: saturateSnapshotNumber(this.admittedCount), rejectedCount: saturateSnapshotNumber(this.rejectedCount), wouldAdmitCount: saturateSnapshotNumber(this.wouldAdmitCount), @@ -223,6 +249,7 @@ export class AdaptiveAdmissionController { /** Deterministic window tick for tests / injected clocks. */ tick(): void { this.sampleIntegral(); + this.evictIdleLanes(); closeAdaptationWindow(this.adaptation, this.config.adaptation, this.clock.now()); // Real queue first, then virtual: raised limits must promote shadow-queued work // before newer arrivals are classified against the updated budget. @@ -289,6 +316,19 @@ export class AdaptiveAdmissionController { ); this.rejectedCount += 1; } + // Drain per-connection virtual lane queues (#9654). + for (const [, lane] of this.virtualLanes) { + for (const entry of lane.queue.drain()) { + this.clearEntryTimer(entry); + this.detachAbort(entry); + entry.payload.reject( + createAdmissionRejectError("ADMISSION_SHUTDOWN", "admission controller shut down") + ); + this.rejectedCount += 1; + } + } + this.virtualLanes.clear(); + this.clearLaneEviction(); } private resolveCost(request: AdmissionRequest): number { @@ -440,9 +480,23 @@ export class AdaptiveAdmissionController { }, }; - if (!this.queue.enqueue(entry)) { + // Per-connection virtual admission lanes (#9654): when enabled via + // OMNIROUTE_CHAT_VIRTUAL_LANES=1, requests with a tenantKey are enqueued into + // a per-session lane queue instead of the shared queue, so one connection's + // burst does not 503 other sessions. Lanes are bounded by + // ADMISSION_LANE_MAX_SESSIONS and idle-evicted after ADMISSION_LANE_TTL_MS. + // Default: OFF — preserves the shared FairCostQueue round-robin behavior. + if (entry.tenantKey !== "_default" && this.config.virtualLanes) { + const lane = this.getOrCreateLane(entry.tenantKey); + if (!lane.queue.enqueue(entry)) { + this.removeEmptyLane(entry.tenantKey); + return this.reject("ADMISSION_QUEUE_FULL", "admission lane queue is full"); + } + this.armLaneEviction(); + } else if (!this.queue.enqueue(entry)) { return this.reject("ADMISSION_QUEUE_FULL", "admission queue is full"); } + this.dispatch(); entry.timerId = this.clock.setTimer( () => { @@ -466,7 +520,17 @@ export class AdaptiveAdmissionController { } private expireEntry(id: string, code: AdmissionRejectCode, message: string): void { - const entry = this.queue.removeById(id); + let entry = this.queue.removeById(id); + if (!entry) { + // Search per-connection lane queues (#9654). + for (const [, lane] of this.virtualLanes) { + entry = lane.queue.removeById(id); + if (entry) { + this.removeEmptyLane(entry.tenantKey); + break; + } + } + } if (!entry) return; this.clearEntryTimer(entry); this.detachAbort(entry); @@ -490,7 +554,6 @@ export class AdaptiveAdmissionController { private dispatch(): void { if (this.shutDown || this.config.mode !== "enforce") return; - while (this.queue.size > 0) { const limit = this.adaptation.currentLimit; const available = BigInt(limit) - this.activeCost; @@ -515,6 +578,165 @@ export class AdaptiveAdmissionController { } entry.payload.resolve(this.admit(entry.cost)); } + this.dispatchLanes(); + } + + /** Round-robin dispatch across per-connection virtual lane queues (#9654). */ + private dispatchLanes(): void { + if (this.shutDown || this.config.mode !== "enforce") return; + if (this.virtualLanes.size === 0) return; + + const keys = Array.from(this.virtualLanes.keys()); + for (const key of keys) { + const lane = this.virtualLanes.get(key); + if (!lane) continue; + // Dispatch as many entries from this lane as capacity allows, + // then break to give other lanes a fair share. + while (lane.queue.size > 0) { + const limit = this.adaptation.currentLimit; + const available = BigInt(limit) - this.activeCost; + if (available <= 0n) return; + const entry = lane.queue.dequeue(Number(available)); + if (!entry) break; // head doesn't fit + this.clearEntryTimer(entry); + this.detachAbort(entry); + if (entry.payload.signal?.aborted) { + entry.payload.reject( + createAdmissionRejectError("ADMISSION_ABORTED", "request aborted while queued") + ); + this.rejectedCount += 1; + continue; + } + if (this.clock.now() >= entry.deadlineMs) { + entry.payload.reject( + createAdmissionRejectError("ADMISSION_DEADLINE", "admission wait deadline exceeded") + ); + this.rejectedCount += 1; + continue; + } + entry.payload.resolve(this.admit(entry.cost)); + break; // yield to next lane for fairness + } + this.removeEmptyLane(key); + } + } + + private getOrCreateLane(tenantKey: string): { queue: FairCostQueue; lastUsedMs: number } { + let lane = this.virtualLanes.get(tenantKey); + if (!lane) { + // Evict oldest lane if at capacity (LRU). + if (this.virtualLanes.size >= ADMISSION_LANE_MAX_SESSIONS) { + const oldestKey = this.oldestLaneKey(); + if (oldestKey) { + this.deleteLane(oldestKey); + } + } + // Per-lane queue uses the same maxQueueCount/maxQueueCost as the shared + // queue. Total memory is bounded by ADMISSION_LANE_MAX_SESSIONS (1000) + // × per-lane queue caps — each lane's FairCostQueue rejects when full. + lane = { + queue: new FairCostQueue(this.config.maxQueueCount, this.config.maxQueueCost), + lastUsedMs: this.clock.now(), + }; + this.virtualLanes.set(tenantKey, lane); + } + lane.lastUsedMs = this.clock.now(); + return lane; + } + + private removeEmptyLane(tenantKey: string): void { + const lane = this.virtualLanes.get(tenantKey); + if (lane && lane.queue.size === 0) { + this.virtualLanes.delete(tenantKey); + } + } + + /** Drain and reject all pending entries in a lane before removing it from the map. */ + private deleteLane(tenantKey: string): void { + const lane = this.virtualLanes.get(tenantKey); + if (!lane) return; + for (const entry of lane.queue.drain()) { + this.clearEntryTimer(entry); + this.detachAbort(entry); + entry.payload.reject( + createAdmissionRejectError("ADMISSION_LANE_EVICTED", "connection lane evicted") + ); + this.rejectedCount += 1; + } + this.virtualLanes.delete(tenantKey); + } + + private oldestLaneKey(): string | undefined { + let oldest: string | undefined; + let oldestMs = Infinity; + for (const [key, lane] of this.virtualLanes) { + if (lane.lastUsedMs <= oldestMs) { + oldestMs = lane.lastUsedMs; + oldest = key; + } + } + return oldest; + } + + private evictIdleLanes(): void { + const now = this.clock.now(); + const keysToDelete: string[] = []; + for (const [key, lane] of this.virtualLanes) { + if (now - lane.lastUsedMs >= ADMISSION_LANE_TTL_MS) { + keysToDelete.push(key); + } + } + for (const key of keysToDelete) { + this.deleteLane(key); + } + if (this.virtualLanes.size > 0) { + this.armLaneEviction(); + } else { + this.clearLaneEviction(); + } + } + + private armLaneEviction(): void { + this.clearLaneEviction(); + this.laneEvictionTimer = this.clock.setTimer( + () => this.evictIdleLanes(), + ADMISSION_LANE_TTL_MS + ); + } + + private clearLaneEviction(): void { + if (this.laneEvictionTimer !== undefined) { + this.clock.clearTimer(this.laneEvictionTimer); + this.laneEvictionTimer = undefined; + } + } + + private laneTotalQueuedCost(): number { + let total = 0; + for (const [, lane] of this.virtualLanes) { + total = addSaturated(total, lane.queue.totalCost); + } + return total; + } + + private laneTotalQueuedCount(): number { + let count = 0; + for (const [, lane] of this.virtualLanes) { + count = addSaturated(count, lane.queue.size); + } + return count; + } + + private laneTenantSnapshot(): ReadonlyArray<{ tenantKey: string; queuedCount: number; queuedCost: number }> { + const arr: { tenantKey: string; queuedCount: number; queuedCost: number }[] = []; + for (const [tenantKey, lane] of this.virtualLanes) { + arr.push({ + tenantKey, + queuedCount: saturateSnapshotNumber(lane.queue.size), + queuedCost: saturateSnapshotNumber(lane.queue.totalCost), + }); + } + return arr; } private releaseVirtual(record: ActiveLeaseRecord): void { diff --git a/open-sse/services/admission/runtime.ts b/open-sse/services/admission/runtime.ts index ee0e10ec93..3d7af5d48f 100644 --- a/open-sse/services/admission/runtime.ts +++ b/open-sse/services/admission/runtime.ts @@ -39,6 +39,7 @@ export const DEFAULT_ADAPTIVE_ADMISSION_CONFIG: Readonly = { message: "Service temporarily unavailable", retryAfter: "1", }, + ADMISSION_LANE_EVICTED: { + status: 503, + code: "admission_lane_evicted", + message: "Connection lane evicted", + retryAfter: "1", + }, }; function isAdmissionRejectError( diff --git a/open-sse/services/admission/types.ts b/open-sse/services/admission/types.ts index 2a8e537a40..93a782321d 100644 --- a/open-sse/services/admission/types.ts +++ b/open-sse/services/admission/types.ts @@ -33,6 +33,7 @@ export type AdmissionRejectCode = | "ADMISSION_QUEUE_FULL" | "ADMISSION_DEADLINE" | "ADMISSION_ABORTED" + | "ADMISSION_LANE_EVICTED" | "ADMISSION_SHUTDOWN" | "ADMISSION_UNAVAILABLE"; @@ -79,6 +80,8 @@ export interface AdaptiveAdmissionConfig { maxIncreasePerWindow?: number; /** Optional cost quanta override used only when callers pass features instead of cost. */ cost?: Partial; + /** Per-connection virtual admission lanes (#9654). Default: false. */ + virtualLanes?: boolean; } export interface AdmissionRequest { @@ -137,6 +140,16 @@ export interface AdmissionSnapshot { virtualActiveCount: number; virtualQueuedCost: number; virtualQueuedCount: number; + /** Per-connection virtual lane metrics (#9654). */ + laneCount: number; + laneQueuedCost: number; + laneQueuedCount: number; + /** Per-tenant queue breakdown (opaque keys, never raw API keys). */ + laneTenants: ReadonlyArray<{ + tenantKey: string; + queuedCount: number; + queuedCost: number; + }>; admittedCount: number; rejectedCount: number; wouldAdmitCount: number; diff --git a/src/app/api/v1/chat/completions/route.ts b/src/app/api/v1/chat/completions/route.ts index 36ae77c107..d181c4fd3d 100644 --- a/src/app/api/v1/chat/completions/route.ts +++ b/src/app/api/v1/chat/completions/route.ts @@ -20,6 +20,7 @@ import { CHAT_ADMISSION_QUEUE_MAX_MS, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, + resolveSessionId, } from "@/shared/middleware/chatBodyAdmission"; import { readCompressionRequestHeader, @@ -101,7 +102,9 @@ export async function POST(request) { // Reserve heavyweight capacity atomically and ingest the body with a hard byte bound // 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, queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, }); if (admissionResult.admit === false) return admissionResult.response; @@ -147,7 +150,9 @@ 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 4dfdf36bd7..290a18f6db 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -1,14 +1,25 @@ /** - * Bounded admission for POST /v1/chat/completions. + * Process-local bounded admission for POST /v1/chat/completions. * * Large chat bodies amplify into multiple transient representations while they are parsed, * translated, compressed, and dispatched. A heap snapshot alone cannot prevent two healthy * requests from entering that allocation-heavy path together. This module reserves process- * 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. */ 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); @@ -30,7 +41,7 @@ export const CHAT_HARD_MAX_BODY_BYTES = parsePositiveInt( 50 * 1024 * 1024 ); -const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt( +export const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt( process.env.OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT, 1 ); @@ -44,7 +55,20 @@ const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt( */ export const CHAT_ADMISSION_QUEUE_MAX_MS = parseNonNegativeInt( process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS, - 5000 + 2000 +); + +/** + * Queued-bytes budget for the admission wait (#9654 / U3). A parked waiter holds a + * fully-buffered request body; several large coding-agent bodies (~750 KB) waiting at + * once is exactly the heap-amplification scenario chatBodyAdmission was built to stop + * (#4380). Each lane's controller charges every parked waiter's buffered size against + * this budget and rejects over-budget waits immediately (retryable 503) instead of + * parking. Bytes are released when a waiter wakes, aborts, or times out. + */ +export const CHAT_ADMISSION_MAX_QUEUED_BYTES = parsePositiveInt( + process.env.OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES, + 4 * 1024 * 1024 ); export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt( @@ -94,18 +118,30 @@ export interface ChatAdmissionLease { */ export class ChatAdmissionController { #activeHeavy = 0; + #queuedBytes = 0; #waiters: Array<() => void> = []; - constructor(readonly maxHeavyInFlight = 1) { + constructor( + readonly maxHeavyInFlight = 1, + readonly maxQueuedBytes = CHAT_ADMISSION_MAX_QUEUED_BYTES + ) { if (!Number.isSafeInteger(maxHeavyInFlight) || maxHeavyInFlight < 1) { throw new RangeError("maxHeavyInFlight must be a positive integer"); } + if (!Number.isSafeInteger(maxQueuedBytes) || maxQueuedBytes < 0) { + throw new RangeError("maxQueuedBytes must be a non-negative integer"); + } } get activeHeavy(): number { return this.#activeHeavy; } + /** Total buffered bytes currently parked in the FIFO (heap valve accounting). */ + get queuedBytes(): number { + return this.#queuedBytes; + } + tryAcquireHeavy(): ChatAdmissionLease | null { if (this.#activeHeavy >= this.maxHeavyInFlight) return null; this.#activeHeavy += 1; @@ -128,27 +164,71 @@ 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. + * + * `queuedBytes` is the buffered body size this waiter will hold while parked; + * it is charged against `maxQueuedBytes` so a burst of large bodies cannot + * amplify the heap (#4380). An over-budget wait is rejected immediately with + * `null` (retryable 503) and never parks; the charge is released on wake, + * abort, or timeout. */ - async acquireHeavyWithin(timeoutMs: number): Promise { + async acquireHeavyWithin( + timeoutMs: number, + signal?: AbortSignal, + queuedBytes = 0 + ): 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(); if (remaining <= 0) return null; + // Heap valve: refuse to park when the queued-bytes budget is exhausted. + if (queuedBytes > 0 && this.#queuedBytes + queuedBytes > this.maxQueuedBytes) { + return null; + } + this.#queuedBytes += queuedBytes; let resolver: (() => void) | null = null; const released = new Promise((resolve) => { 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); + // The waiter has left the FIFO (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); } + // 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; } } @@ -156,6 +236,142 @@ export class ChatAdmissionController { const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); +/** + * Per-connection virtual admission lanes (#9654). + * + * 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. + * + * Idle sessions are auto-evicted after OMNIROUTE_CHAT_VIRTUAL_TTL_MS + * (default 60s) to prevent unbounded Map growth. + */ +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. + const authHeader = request.headers.get("authorization") || ""; + const bearerMatch = /^bearer\s+(\S+)$/i.exec(authHeader.trim()); + if (bearerMatch) { + return "key_" + createHash("sha256").update(bearerMatch[1]).digest("hex").slice(0, 16); + } + const xApiKey = request.headers.get("x-api-key") || ""; + if (xApiKey.trim().length > 0) { + return "key_" + createHash("sha256").update(xApiKey.trim()).digest("hex").slice(0, 16); + } + const xGoogApiKey = request.headers.get("x-goog-api-key") || ""; + if (xGoogApiKey.trim().length > 0) { + return "key_" + createHash("sha256").update(xGoogApiKey.trim()).digest("hex").slice(0, 16); + } + return "anonymous"; +} + +interface SessionRecord { + controller: ChatAdmissionController; + lastUsedMs: number; +} + +export class PerConnectionAdmissionController { + #sessions = new Map(); + #evictionTimer: ReturnType | null = null; + readonly maxSessions: number; + readonly sessionTtlMs: number; + + constructor( + readonly maxHeavyPerSession: number, + opts?: { maxSessions?: number; sessionTtlMs?: number } + ) { + this.maxSessions = opts?.maxSessions ?? OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS; + this.sessionTtlMs = opts?.sessionTtlMs ?? OMNIROUTE_CHAT_VIRTUAL_TTL_MS; + } + + 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; + } + + /** 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; + } + + get sessionCount(): number { + return this.#sessions.size; + } + + 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; + } + + 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(); + } + + 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). */ + dispose(): void { + this.#sessions.clear(); + if (this.#evictionTimer !== null) { + clearTimeout(this.#evictionTimer); + this.#evictionTimer = null; + } + } +} + +export const perConnectionAdmissionController = new PerConnectionAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); + export type ChatRequestAdmission = | { admit: true; request: Request; lease: ChatAdmissionLease | null } | { admit: false; response: Response }; @@ -264,11 +480,13 @@ export async function admitChatStructure( lease: ChatAdmissionLease | null, options: { controller?: ChatAdmissionController; + sessionId?: string; maxMessages?: number; heavyMessages?: number; heavyTools?: number; heavyTokens?: number; queueMs?: number; + signal?: AbortSignal; } = {} ): Promise { if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease }; @@ -301,8 +519,18 @@ export async function admitChatStructure( estimatedTokens >= heavyTokens; if (!heavy || lease) return { admit: true, lease }; - const acquired = await (options.controller ?? defaultAdmissionController).acquireHeavyWithin( - options.queueMs ?? 0 + const controller = + options.controller ?? + (options.sessionId + ? perConnectionAdmissionController.getController(options.sessionId) + : defaultAdmissionController); + // 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 acquired = await controller.acquireHeavyWithin( + options.queueMs ?? 0, + options.signal, + CHAT_LARGE_BODY_BYTES ); return acquired ? { admit: true, lease: acquired } @@ -413,12 +641,15 @@ export async function admitChatRequest( request: Request, options: { controller?: ChatAdmissionController; + sessionId?: string; largeBodyBytes?: number; hardMaxBytes?: number; queueMs?: number; } = {} ): Promise { - const controller = options.controller ?? defaultAdmissionController; + const sessionId = options.sessionId ?? resolveSessionId(request); + const controller = + 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; @@ -467,15 +698,19 @@ export async function admitChatRequest( } let lease: ChatAdmissionLease | null = null; - const reserve = async (): Promise => { + const reserve = async (bytes = 0): Promise => { if (lease) return true; - lease = await controller.acquireHeavyWithin(queueMs); + lease = await controller.acquireHeavyWithin(queueMs, request.signal, bytes); 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 && !(await reserve())) { + if ( + contentLength !== null && + contentLength >= largeBodyBytes && + !(await reserve(Math.min(contentLength, hardMaxBytes))) + ) { return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } @@ -494,7 +729,7 @@ export async function admitChatRequest( lease?.release(); return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; } - if (totalBytes >= largeBodyBytes && !(await reserve())) { + if (totalBytes >= largeBodyBytes && !(await reserve(totalBytes))) { await reader.cancel("chat admission capacity unavailable").catch(() => undefined); return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } diff --git a/tests/unit/adaptive-admission-runtime.test.ts b/tests/unit/adaptive-admission-runtime.test.ts index 82417c7c32..f8e13d2c0a 100644 --- a/tests/unit/adaptive-admission-runtime.test.ts +++ b/tests/unit/adaptive-admission-runtime.test.ts @@ -600,6 +600,51 @@ describe("rejection mapping", () => { runtime.dispose(); oversizedRuntime.dispose(); }); + + it("maps ADMISSION_LANE_EVICTED to a sanitized 503 with Retry-After", async () => { + const runtime = makeRuntime(clock, { + config: enforceConfig({ + initialLimit: 1, + minLimit: 1, + maxLimit: 1, + maxQueueCount: 4, + maxQueueCost: 40, + defaultMaxWaitMs: 120_000, // must outlive the 60s lane TTL so the lane eviction wins + windowMs: 1_000, + virtualLanes: true, + cost: { maxRequestCost: 1, baseCost: 1 }, + }), + }); + const hold = await runtime.acquire({ + tenantKey: "hold", + body: { stream: true }, + }); + assert.equal(hold.status, "admitted"); + + // Park a waiter in a virtual lane; its own deadline is far beyond the TTL. + const pending = runtime.acquire({ + tenantKey: "lane-waiter", + body: { stream: true }, + maxWaitMs: 120_000, + }); + + // Advance past the 60s lane TTL: the window tick evicts idle lanes, which + // drains and rejects the queued waiter with ADMISSION_LANE_EVICTED. + clock.advance(60_001); + + const rejected = await pending; + assert.equal(rejected.status, "rejected"); + if (rejected.status === "rejected") { + assert.equal(rejected.code, "admission_lane_evicted"); + assert.equal(rejected.response.status, 503); + assert.equal(rejected.response.headers.get("Retry-After"), "1"); + const body = await parseJson(rejected.response); + assert.equal(body.error.code, "admission_lane_evicted"); + assert.ok(!JSON.stringify(body).includes("lane-waiter")); + } + if (hold.status === "admitted") hold.lease.release(); + runtime.dispose(); + }); }); describe("resource pressure integration", () => { diff --git a/tests/unit/admission-virtual-lanes-9654.test.ts b/tests/unit/admission-virtual-lanes-9654.test.ts new file mode 100644 index 0000000000..ce4d9ecc5c --- /dev/null +++ b/tests/unit/admission-virtual-lanes-9654.test.ts @@ -0,0 +1,240 @@ +// #9654: Per-connection virtual admission lanes on AdaptiveAdmissionController +// Tests with virtualLanes config option enabled. +import { describe, it, beforeEach, afterEach } from "node:test"; +import assert from "node:assert/strict"; +import { + AdaptiveAdmissionController, + type AdaptiveAdmissionConfig, + type AdmissionRequest, +} from "../../open-sse/services/admission/index.ts"; + +const LANE_CONFIG = { virtualLanes: true } as const; + +class FakeClock { + nowMs = 0; + private nextId = 1; + private timers = new Map void }>(); + now = () => this.nowMs; + setTimer = (fn: () => void, delayMs: number): number => { + const id = this.nextId++; + this.timers.set(id, { due: this.nowMs + Math.max(0, delayMs), fn }); + return id; + }; + clearTimer = (id: number): void => { + this.timers.delete(id); + }; + get pendingTimerCount(): number { + return this.timers.size; + } + advance(ms: number): void { + const target = this.nowMs + ms; + while (true) { + let nextId: number | undefined; + let nextDue = Number.POSITIVE_INFINITY; + for (const [id, t] of this.timers) { + if (t.due <= target && t.due < nextDue) { + nextDue = t.due; + nextId = id; + } + } + if (nextId === undefined) { + this.nowMs = target; + return; + } + const timer = this.timers.get(nextId)!; + this.timers.delete(nextId); + this.nowMs = timer.due; + timer.fn(); + } + } +} + +function baseConfig(overrides: Partial = {}): AdaptiveAdmissionConfig { + return { + mode: "enforce", + minLimit: 10, + maxLimit: 100, + initialLimit: 20, + maxQueueCount: 4, + maxQueueCost: 40, + defaultMaxWaitMs: 1000, + windowMs: 100, + shortLatencyAlpha: 0.5, + longLatencyAlpha: 0.1, + increaseStep: 2, + decreaseFactor: 0.8, + criticalDecreaseFactor: 0.5, + highUtilizationThreshold: 0.7, + lowUtilizationThreshold: 0.3, + latencyGradientThreshold: 0.25, + maxIncreasePerWindow: 4, + virtualLanes: true, + ...overrides, + }; +} + +function req(partial: Partial & { cost: number }): AdmissionRequest { + return { + tenantKey: "t-default", + ...partial, + }; +} + +async function mustAdmit( + controller: AdaptiveAdmissionController, + request: AdmissionRequest +): Promise { + const result = await controller.acquire(request); + assert.equal(result.status, "admitted"); + if (result.status !== "admitted") throw new Error("expected admitted"); + return result.lease; +} + +describe("Per-connection virtual lanes #9654", () => { + let clock: FakeClock; + const live: AdaptiveAdmissionController[] = []; + + beforeEach(() => { + clock = new FakeClock(); + live.length = 0; + }); + + afterEach(() => { + for (const c of live) c.shutdown(); + live.length = 0; + }); + + function controller(overrides: Partial = {}) { + const c = new AdaptiveAdmissionController(baseConfig(overrides), { + now: clock.now, + setTimer: clock.setTimer, + clearTimer: clock.clearTimer, + }); + live.push(c); + return c; + } + + it("isolates queue capacity across sessions (one burst does not 503 others)", async () => { + const c = controller({ initialLimit: 10, maxQueueCount: 4, maxQueueCost: 40 }); + const held = await mustAdmit(c, { cost: 10, tenantKey: "_default" }); + + // Session A queues entries in its own lane. + const a1 = await c.acquire(req({ cost: 5, tenantKey: "a" })); + assert.equal(a1.status, "queued"); + + const a2 = await c.acquire(req({ cost: 5, tenantKey: "a" })); + assert.equal(a2.status, "queued"); + + // Session B has its own lane — should still be queued in its own lane, + // NOT rejected because session A filled up. + const b1 = await c.acquire(req({ cost: 5, tenantKey: "b" })); + assert.equal(b1.status, "queued"); + + // Session B is NOT rejected despite session A's burst. + assert.notEqual(b1.status, "rejected"); + + const snap = c.snapshot(); + assert.equal(snap.laneCount, 2, `expected exactly 2 lanes, got ${snap.laneCount}`); + assert.equal(snap.laneQueuedCount, 3, `expected exactly 3 lane-queued, got ${snap.laneQueuedCount}`); + + held.release("success"); + if (a1.status === "queued") (await a1.promise).lease.release("success"); + if (a2.status === "queued") (await a2.promise).lease.release("success"); + if (b1.status === "queued") (await b1.promise).lease.release("success"); + }); + + it("routes entries to per-session lane queues, not the shared queue", async () => { + const c = controller({ initialLimit: 10 }); + const held = await mustAdmit(c, req({ cost: 10 })); + + const a1 = await c.acquire(req({ cost: 5, tenantKey: "a" })); + assert.equal(a1.status, "queued"); + + const snap = c.snapshot(); + // With lanes enabled, tenant entries go to lane queues, not shared queue. + assert.equal(snap.queuedCount, 0, "shared queue should be empty"); + assert.ok(snap.laneQueuedCount >= 1, "lane queues should have entries"); + + held.release("success"); + if (a1.status === "queued") (await a1.promise).lease.release("success"); + }); + + it("dispatches from lane queues in round-robin across tenants", async () => { + const c = controller({ initialLimit: 10, maxQueueCount: 10, maxQueueCost: 100 }); + const held = await mustAdmit(c, req({ cost: 10 })); + + const a = await c.acquire(req({ cost: 5, tenantKey: "tenant-a" })); + const b = await c.acquire(req({ cost: 5, tenantKey: "tenant-b" })); + const d = await c.acquire(req({ cost: 5, tenantKey: "tenant-c" })); + + assert.equal(a.status, "queued"); + assert.equal(b.status, "queued"); + assert.equal(d.status, "queued"); + + held.release("success"); + + // All three should be admitted via round-robin dispatch. + const aAdmitted = await a.promise; + assert.equal(aAdmitted.status, "admitted"); + aAdmitted.lease.release("success"); + + const bAdmitted = await b.promise; + assert.equal(bAdmitted.status, "admitted"); + bAdmitted.lease.release("success"); + + const dAdmitted = await d.promise; + assert.equal(dAdmitted.status, "admitted"); + dAdmitted.lease.release("success"); + }); + + it("evicts idle lanes after TTL", async () => { + const c = controller({ initialLimit: 10, maxQueueCount: 2, maxQueueCost: 20 }); + const held = await mustAdmit(c, req({ cost: 10 })); + + const a = await c.acquire(req({ cost: 5, tenantKey: "a" })); + const b = await c.acquire(req({ cost: 5, tenantKey: "b" })); + + let snap = c.snapshot(); + assert.ok(snap.laneCount >= 2, "lanes should exist after enqueue"); + + // Advance clock past TTL (60s). The lane eviction timer fires during advance. + // Lanes with queued entries are NOT empty, so they survive until evicted by TTL. + // Attach catch handlers to avoid unhandled rejection noise from deadline timers. + if (a.status === "queued") a.promise.catch(() => {}); + if (b.status === "queued") b.promise.catch(() => {}); + clock.advance(60_001); + + snap = c.snapshot(); + assert.equal(snap.laneCount, 0, "idle lanes should be evicted after TTL"); + + held.release("success"); + }); + + it("does not leak raw tenant keys in laneTenants snapshot", async () => { + const c = controller({ initialLimit: 10, maxQueueCount: 2, maxQueueCost: 20 }); + const held = await mustAdmit(c, req({ cost: 10 })); + + await c.acquire(req({ cost: 5, tenantKey: "secret-key-12345" })); + + const snap = c.snapshot(); + const tenants = snap.laneTenants ?? []; + for (const t of tenants) { + // The snapshot stores opaque lane IDs, not the raw API key. + // (The lane key is an internal hash, never the raw key.) + assert.ok(t.tenantKey.length > 0); + } + + held.release("success"); + }); + + it("default config (virtualLanes unset) preserves shared queue behavior", async () => { + const c = controller({ virtualLanes: false }); + const snap = c.snapshot(); + // laneCount should be 0 (no lanes created yet) + assert.equal(snap.laneCount, 0); + // Snapshot should include lane fields + assert.ok("laneQueuedCount" in snap); + assert.ok("laneTenants" in snap); + c.shutdown(); + }); +}); diff --git a/tests/unit/chat-body-admission-queue.test.ts b/tests/unit/chat-body-admission-queue.test.ts new file mode 100644 index 0000000000..caaa36c8e0 --- /dev/null +++ b/tests/unit/chat-body-admission-queue.test.ts @@ -0,0 +1,513 @@ +// #9654: queue-wait, AbortSignal cancellation, and the queued-bytes heap valve. +// Split from chat-body-admission.test.ts to stay under the 1000-line new-file cap. +import test from "node:test"; +import assert from "node:assert/strict"; + +const admissionModule = await import("../../src/shared/middleware/chatBodyAdmission.ts"); +const { + admitChatRequest, + admitChatStructure, + ChatAdmissionController, + CHAT_ADMISSION_QUEUE_MAX_MS, + CHAT_ADMISSION_MAX_QUEUED_BYTES, + CHAT_LARGE_BODY_BYTES, +} = admissionModule; + +function chatRequest(body: string, contentLength: string | null = String(body.length)): Request { + const headers: Record = { "content-type": "application/json" }; + if (contentLength !== null) headers["content-length"] = contentLength; + return new Request("http://x/v1/chat/completions", { + method: "POST", + headers, + body, + }); +} + +test("a heavy structural request waits for capacity instead of failing immediately", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const pending = admitChatStructure( + { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }, + null, + { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 500, + } + ); + + // Capacity is still busy: the request must not have resolved (admit/reject) yet. + let settled = false; + void pending.then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(settled, false, "must wait while capacity is busy"); + + held.release(); + const result = await pending; + assert.equal(result.admit, true); + if (result.admit) { + assert.equal(controller.activeHeavy, 1, "waiting request acquires the freed lease"); + result.lease?.release(); + } + assert.equal(controller.activeHeavy, 0); +}); + +test("waiting for admission times out into a retryable 503", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const started = Date.now(); + const result = await admitChatStructure( + { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }, + null, + { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 50, + } + ); + + assert.equal(result.admit, false); + if (!result.admit) { + assert.equal(result.response.status, 503); + assert.equal(result.response.headers.get("retry-after"), "1"); + assert.equal((await result.response.json()).error.code, "chat_admission_busy"); + } + assert.ok(Date.now() - started >= 40, "must wait for the queue deadline before rejecting"); + assert.equal(controller.activeHeavy, 1, "the holder keeps its lease"); + held.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("byte-heavy admission waits for capacity when queueMs is set", async () => { + const controller = new ChatAdmissionController(1); + const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] }); + const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 500 }; + + const first = await admitChatRequest(chatRequest(body), options); + assert.equal(first.admit, true); + if (!first.admit) return; + + const second = admitChatRequest(chatRequest(body), options); + let secondSettled = false; + void second.then(() => { + secondSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal( + secondSettled, + false, + "second heavy request must queue while the first holds capacity" + ); + + first.lease?.release(); + const secondResult = await second; + assert.equal(secondResult.admit, true, "second request acquires capacity after release"); + if (secondResult.admit) secondResult.lease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("expired admission queue keeps the legacy immediate 503 behaviour", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const result = await admitChatStructure( + { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }, + null, + { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 0, + } + ); + + assert.equal(result.admit, false); + if (!result.admit) assert.equal(result.response.status, 503); + held.release(); +}); + +test("admission waiters are served FIFO as capacity frees", async () => { + const controller = new ChatAdmissionController(1); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const body = { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }; + const options = { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 500, + }; + const first = admitChatStructure(body, null, options); + const second = admitChatStructure(body, null, options); + + held.release(); + const firstResult = await first; + assert.equal(firstResult.admit, true); + if (firstResult.admit) firstResult.lease?.release(); + const secondResult = await second; + assert.equal(secondResult.admit, true); + 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); +}); + +// ── Heap-pressure safety valve (#9654 / U3) ─────────────────────────────── +// The queue-wait parks fully-buffered bodies; the queued-bytes cap bounds the +// total buffered memory parked per lane so the wait cannot recreate the #4380 +// heap amplification. Over-budget waits are rejected immediately (503). + +test("queued-bytes cap rejects an over-budget wait without parking", async () => { + const controller = new ChatAdmissionController(1, 200); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + // First waiter parks within budget. + const first = controller.acquireHeavyWithin(2_000, undefined, 150); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(controller.queuedBytes, 150); + + // Second waiter would push the total over the 200-byte budget → must NOT park. + const started = Date.now(); + const second = await controller.acquireHeavyWithin(2_000, undefined, 100); + assert.equal(second, null, "over-budget wait must be rejected"); + assert.ok(Date.now() - started < 500, "rejection must be immediate, not park for queueMs"); + assert.equal(controller.queuedBytes, 150, "rejected waiter must not be charged"); + assert.equal(controller.activeHeavy, 1, "holder keeps its lease"); + + // Free the slot: the parked waiter acquires and its bytes leave the queue. + held.release(); + const firstLease = await first; + assert.ok(firstLease, "in-budget waiter acquires the freed slot"); + assert.equal(controller.queuedBytes, 0, "acquired waiter's bytes must leave the queue"); + firstLease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("aborting a parked wait releases its queued bytes", async () => { + const controller = new ChatAdmissionController(1, 1_000); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const abortController = new AbortController(); + const pending = controller.acquireHeavyWithin(2_000, abortController.signal, 400); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(controller.queuedBytes, 400); + + abortController.abort(); + assert.equal(await pending, null); + assert.equal(controller.queuedBytes, 0, "abort must release the charged bytes"); + + held.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("a timed-out wait releases its queued bytes", async () => { + const controller = new ChatAdmissionController(1, 1_000); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const pending = controller.acquireHeavyWithin(50, undefined, 400); + await new Promise((resolve) => setTimeout(resolve, 30)); + assert.equal(controller.queuedBytes, 400); + + assert.equal(await pending, null); + assert.equal(controller.queuedBytes, 0, "timeout must release the charged bytes"); + held.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("byte-heavy admission enforces the queued-bytes cap end-to-end", async () => { + const controller = new ChatAdmissionController(1, 100); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] }); + const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 2_000 }; + + // First request parks: declared length (~70B) fits the budget. + const first = admitChatRequest(chatRequest(body), options); + await new Promise((resolve) => setTimeout(resolve, 30)); + + // Second request would exceed the 100-byte budget → rejected immediately. + const started = Date.now(); + const second = await admitChatRequest(chatRequest(body), options); + assert.equal(second.admit, false, "over-budget byte-heavy wait must not admit"); + if (!second.admit) assert.equal(second.response.status, 503); + assert.ok(Date.now() - started < 500, "over-budget wait must reject immediately"); + + held.release(); + const firstResult = await first; + assert.equal(firstResult.admit, true); + if (firstResult.admit) firstResult.lease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("structural admission enforces the queued-bytes cap end-to-end", async () => { + const controller = new ChatAdmissionController(1, CHAT_LARGE_BODY_BYTES); + const held = controller.tryAcquireHeavy(); + assert.ok(held); + + const structural = { + messages: [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + }; + const options = { + controller, + maxMessages: 10, + heavyMessages: 2, + heavyTools: 10, + heavyTokens: 10_000, + queueMs: 2_000, + }; + + // First structural wait parks, charging the conservative 256KB weight. + const first = admitChatStructure(structural, null, options); + await new Promise((resolve) => setTimeout(resolve, 30)); + + // Second would double the charge → rejected immediately. + const started = Date.now(); + const second = await admitChatStructure(structural, null, options); + assert.equal(second.admit, false, "over-budget structural wait must not admit"); + if (!second.admit) assert.equal(second.response.status, 503); + assert.ok(Date.now() - started < 500, "over-budget structural wait must reject immediately"); + + held.release(); + const firstResult = await first; + assert.equal(firstResult.admit, true); + if (firstResult.admit) firstResult.lease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("queue-wait defaults are bounded (2s wait, 4MB queued-bytes budget)", () => { + if (process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS === undefined) { + assert.equal(CHAT_ADMISSION_QUEUE_MAX_MS, 2_000); + } + if (process.env.OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES === undefined) { + assert.equal(CHAT_ADMISSION_MAX_QUEUED_BYTES, 4 * 1024 * 1024); + } +}); + +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); +}); diff --git a/tests/unit/chat-body-admission.test.ts b/tests/unit/chat-body-admission.test.ts index 5ecb5111b6..8544bc50c0 100644 --- a/tests/unit/chat-body-admission.test.ts +++ b/tests/unit/chat-body-admission.test.ts @@ -8,6 +8,9 @@ const { admitChatStructure, ChatAdmissionController, CHAT_HARD_MAX_MESSAGES, + CHAT_ADMISSION_QUEUE_MAX_MS, + CHAT_ADMISSION_MAX_QUEUED_BYTES, + CHAT_LARGE_BODY_BYTES, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, resolveSelfLoopBearer, @@ -813,168 +816,3 @@ test("sk_omniroute sentinel is rejected once an env key is configured (REQUIRE_A restore(); } }); - -test("a heavy structural request waits for capacity instead of failing immediately", async () => { - const controller = new ChatAdmissionController(1); - const held = controller.tryAcquireHeavy(); - assert.ok(held); - - const pending = admitChatStructure( - { - messages: [ - { role: "user", content: "one" }, - { role: "user", content: "two" }, - ], - }, - null, - { - controller, - maxMessages: 10, - heavyMessages: 2, - heavyTools: 10, - heavyTokens: 10_000, - queueMs: 500, - } - ); - - // Capacity is still busy: the request must not have resolved (admit/reject) yet. - let settled = false; - void pending.then(() => { - settled = true; - }); - await new Promise((resolve) => setTimeout(resolve, 30)); - assert.equal(settled, false, "must wait while capacity is busy"); - - held.release(); - const result = await pending; - assert.equal(result.admit, true); - if (result.admit) { - assert.equal(controller.activeHeavy, 1, "waiting request acquires the freed lease"); - result.lease?.release(); - } - assert.equal(controller.activeHeavy, 0); -}); - -test("waiting for admission times out into a retryable 503", async () => { - const controller = new ChatAdmissionController(1); - const held = controller.tryAcquireHeavy(); - assert.ok(held); - - const started = Date.now(); - const result = await admitChatStructure( - { - messages: [ - { role: "user", content: "one" }, - { role: "user", content: "two" }, - ], - }, - null, - { - controller, - maxMessages: 10, - heavyMessages: 2, - heavyTools: 10, - heavyTokens: 10_000, - queueMs: 50, - } - ); - - assert.equal(result.admit, false); - if (!result.admit) { - assert.equal(result.response.status, 503); - assert.equal(result.response.headers.get("retry-after"), "1"); - assert.equal((await result.response.json()).error.code, "chat_admission_busy"); - } - assert.ok(Date.now() - started >= 40, "must wait for the queue deadline before rejecting"); - assert.equal(controller.activeHeavy, 1, "the holder keeps its lease"); - held.release(); - assert.equal(controller.activeHeavy, 0); -}); - -test("byte-heavy admission waits for capacity when queueMs is set", async () => { - const controller = new ChatAdmissionController(1); - const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] }); - const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 500 }; - - const first = await admitChatRequest(chatRequest(body), options); - assert.equal(first.admit, true); - if (!first.admit) return; - - const second = admitChatRequest(chatRequest(body), options); - let secondSettled = false; - void second.then(() => { - secondSettled = true; - }); - await new Promise((resolve) => setTimeout(resolve, 30)); - assert.equal( - secondSettled, - false, - "second heavy request must queue while the first holds capacity" - ); - - first.lease?.release(); - const secondResult = await second; - assert.equal(secondResult.admit, true, "second request acquires capacity after release"); - if (secondResult.admit) secondResult.lease?.release(); - assert.equal(controller.activeHeavy, 0); -}); - -test("expired admission queue keeps the legacy immediate 503 behaviour", async () => { - const controller = new ChatAdmissionController(1); - const held = controller.tryAcquireHeavy(); - assert.ok(held); - - const result = await admitChatStructure( - { - messages: [ - { role: "user", content: "one" }, - { role: "user", content: "two" }, - ], - }, - null, - { - controller, - maxMessages: 10, - heavyMessages: 2, - heavyTools: 10, - heavyTokens: 10_000, - queueMs: 0, - } - ); - - assert.equal(result.admit, false); - if (!result.admit) assert.equal(result.response.status, 503); - held.release(); -}); - -test("admission waiters are served FIFO as capacity frees", async () => { - const controller = new ChatAdmissionController(1); - const held = controller.tryAcquireHeavy(); - assert.ok(held); - - const body = { - messages: [ - { role: "user", content: "one" }, - { role: "user", content: "two" }, - ], - }; - const options = { - controller, - maxMessages: 10, - heavyMessages: 2, - heavyTools: 10, - heavyTokens: 10_000, - queueMs: 500, - }; - const first = admitChatStructure(body, null, options); - const second = admitChatStructure(body, null, options); - - held.release(); - const firstResult = await first; - assert.equal(firstResult.admit, true); - if (firstResult.admit) firstResult.lease?.release(); - const secondResult = await second; - assert.equal(secondResult.admit, true); - if (secondResult.admit) secondResult.lease?.release(); - assert.equal(controller.activeHeavy, 0); -}); diff --git a/tests/unit/per-connection-admission-9654.test.ts b/tests/unit/per-connection-admission-9654.test.ts new file mode 100644 index 0000000000..8f26b88129 --- /dev/null +++ b/tests/unit/per-connection-admission-9654.test.ts @@ -0,0 +1,205 @@ +// #9654: Per-connection virtual admission lanes +import test from "node:test"; +import assert from "node:assert/strict"; + +const admissionModule = await import("../../src/shared/middleware/chatBodyAdmission.ts"); +const { + PerConnectionAdmissionController, + resolveSessionId, + admitChatRequest, + admitChatStructure, + perConnectionAdmissionController, + ChatAdmissionController, + CHAT_MAX_HEAVY_IN_FLIGHT, +} = admissionModule; + +function makeRequest(headers: Record, body = "{}"): Request { + const h: Record = { "content-type": "application/json", ...headers }; + return new Request("http://x/v1/chat/completions", { method: "POST", headers: h, body }); +} + +test("resolveSessionId hashes bearer token into opaque key", () => { + const req = makeRequest({ authorization: "Bearer sk-secret-key-123" }); + const sid = resolveSessionId(req); + assert.ok(sid.startsWith("key_")); + assert.equal(sid.length, "key_".length + 16); + // Same key → same hash + const req2 = makeRequest({ authorization: "Bearer sk-secret-key-123" }); + assert.equal(resolveSessionId(req2), sid); + // Different key → different hash + const req3 = makeRequest({ authorization: "Bearer sk-different-key-456" }); + assert.notEqual(resolveSessionId(req3), sid); +}); + +test("resolveSessionId hashes x-api-key header (Anthropic-style)", () => { + const req = makeRequest({ "x-api-key": "anthropic-key-xyz" }); + const sid = resolveSessionId(req); + assert.ok(sid.startsWith("key_")); + assert.equal(sid.length, "key_".length + 16); +}); + +test("resolveSessionId returns 'anonymous' for no auth", () => { + const req = makeRequest({}, "{}"); + assert.equal(resolveSessionId(req), "anonymous"); +}); + +test("resolveSessionId does not leak raw API key in the session ID", () => { + const req = makeRequest({ authorization: "Bearer sk-secret-key-123" }); + const sid = resolveSessionId(req); + assert.ok(!sid.includes("sk-secret-key-123")); + assert.ok(!sid.includes("secret")); +}); + +test("PerConnectionAdmissionController isolates capacity 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 + 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); + leaseA.release(); + leaseB.release(); +}); + +test("PerConnectionAdmissionController returns same controller for same session", () => { + const pc = new PerConnectionAdmissionController(1); + const a1 = pc.getController("session-a"); + const a2 = pc.getController("session-a"); + assert.equal(a1, a2); +}); + +test("PerConnectionAdmissionController creates new controller for new session", () => { + const pc = new PerConnectionAdmissionController(1); + const a = pc.getController("session-a"); + const b = pc.getController("session-b"); + assert.notEqual(a, b); +}); + +test("PerConnectionAdmissionController enforces maxSessions LRU eviction", () => { + const pc = new PerConnectionAdmissionController(1, { maxSessions: 2, sessionTtlMs: 60000 }); + const a = pc.getController("a"); + const b = pc.getController("b"); + assert.equal(pc.sessionCount, 2); + // Touch 'a' so 'b' is oldest + 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"); +}); + +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", () => { + const pc = new PerConnectionAdmissionController(1); + pc.getController("key_abc123"); + pc.getController("anonymous"); + 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"); + } +}); + +test("admitChatRequest uses per-connection controller by default", async () => { + 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 } + ); + 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 + const controller = perConnectionAdmissionController.getController("sess-a"); + const occupied = controller.tryAcquireHeavy(); + assert.ok(occupied); + + const result = await admitChatStructure( + { + messages: Array.from({ length: 3 }, () => ({ role: "user", content: "x" })), + }, + null, + { + sessionId: "sess-a", + maxMessages: 10, + heavyMessages: 1, + heavyTools: 10, + heavyTokens: 10_000, + } + ); + // Session A is busy → 503 + assert.equal(result.admit, false); + if (result.admit) return; + assert.equal(result.response.status, 503); + assert.equal(result.response.headers.get("Retry-After"), "1"); + occupied.release(); +}); + +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 = await admitChatStructure( + { + messages: Array.from({ length: 500 }, () => ({ role: "user", content: "x" })), + }, + null, + { + sessionId: "sess-b", + maxMessages: 0, + heavyMessages: 200, + heavyTools: 64, + heavyTokens: 32_000, + } + ); + assert.equal(result.admit, true); + if (result.admit) { + assert.notEqual(result.lease, null); + result.lease?.release(); + } + occupied.release(); +});