From 4839749a20cc443791243a023d65019181fd87d5 Mon Sep 17 00:00:00 2001 From: Brandon Bennett Date: Sun, 9 Aug 2026 11:18:54 -0700 Subject: [PATCH] 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 --- ...-per-connection-virtual-admission-lanes.md | 1 + open-sse/services/admission/config.ts | 2 + open-sse/services/admission/controller.ts | 228 ++++++++++++++++- open-sse/services/admission/runtime.ts | 7 + open-sse/services/admission/types.ts | 13 + src/app/api/v1/chat/completions/route.ts | 11 +- src/shared/middleware/chatBodyAdmission.ts | 234 ++++++++++++----- .../unit/admission-virtual-lanes-9654.test.ts | 240 ++++++++++++++++++ .../per-connection-admission-9654.test.ts | 205 +++++++++++++++ 9 files changed, 865 insertions(+), 76 deletions(-) create mode 100644 changelog.d/fixes/9940-per-connection-virtual-admission-lanes.md create mode 100644 tests/unit/admission-virtual-lanes-9654.test.ts create mode 100644 tests/unit/per-connection-admission-9654.test.ts 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/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..b5183bacda 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; + /** 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 e61ec10603..4ab136d6d1 100644 --- a/src/app/api/v1/chat/completions/route.ts +++ b/src/app/api/v1/chat/completions/route.ts @@ -17,9 +17,9 @@ import { resolveKeepaliveThreshold } from "@omniroute/open-sse/utils/keepaliveTh import { admitChatRequest, admitChatStructure, - CHAT_ADMISSION_QUEUE_MAX_MS, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, + resolveSessionId, } from "@/shared/middleware/chatBodyAdmission"; import { readCompressionRequestHeader, @@ -100,9 +100,8 @@ 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 admissionResult = await admitChatRequest(request, { - queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, - }); + const sessionId = resolveSessionId(request); + const admissionResult = await admitChatRequest(request, { sessionId }); if (admissionResult.admit === false) return admissionResult.response; const admission = admissionResult; request = admission.request; @@ -145,8 +144,8 @@ export async function POST(request) { } } - const structuralAdmission = await admitChatStructure(parsedBody, admission.lease, { - queueMs: CHAT_ADMISSION_QUEUE_MAX_MS, + const structuralAdmission = admitChatStructure(parsedBody, admission.lease, { + sessionId, }); if (structuralAdmission.admit === false) { admission.lease?.release(); diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 4dfdf36bd7..37e554a13d 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -1,25 +1,31 @@ /** - * 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); return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : fallback; } -function parseNonNegativeInt(value: string | undefined, fallback: number): number { - const parsed = Number.parseInt(String(value), 10); - return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : fallback; -} - export const CHAT_LARGE_BODY_BYTES = parsePositiveInt( process.env.OMNIROUTE_CHAT_LARGE_BODY_BYTES, 256 * 1024 @@ -30,23 +36,11 @@ 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 ); -/** - * How long a heavy request waits for heavyweight capacity before giving up with a - * retryable 503. Agent loops (OpenCode, Claude Code, Cursor…) fan out sub-requests - * that routinely land on the admission gate together; an immediate 503 makes the - * client burn its retry budget in seconds and the agent dies mid-task. A short - * bounded wait serializes the burst instead. `0` (legacy) rejects immediately. - */ -export const CHAT_ADMISSION_QUEUE_MAX_MS = parseNonNegativeInt( - process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS, - 5000 -); - export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt( process.env.OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT, 200 @@ -88,13 +82,10 @@ export interface ChatAdmissionLease { /** * Process-local heavyweight reservation. The capacity check and increment execute in one * synchronous JavaScript turn, making acquisition atomic within an OmniRoute process. - * Unavailable capacity is a bounded wait (see `acquireHeavyWithin`) and only then a - * retryable 503, so short agent bursts serialize instead of killing the client's - * retry budget. + * Queueing is intentionally separate: unavailable capacity is a retryable 503. */ export class ChatAdmissionController { #activeHeavy = 0; - #waiters: Array<() => void> = []; constructor(readonly maxHeavyInFlight = 1) { if (!Number.isSafeInteger(maxHeavyInFlight) || maxHeavyInFlight < 1) { @@ -118,44 +109,149 @@ export class ChatAdmissionController { if (released) return; released = true; this.#activeHeavy = Math.max(0, this.#activeHeavy - 1); - this.#waiters.shift()?.(); }, }; } - - /** - * 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. - */ - async acquireHeavyWithin(timeoutMs: number): Promise { - const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs)); - for (;;) { - const lease = this.tryAcquireHeavy(); - if (lease) return lease; - const remaining = deadline - Date.now(); - if (remaining <= 0) return null; - let resolver: (() => void) | null = null; - const released = new Promise((resolve) => { - resolver = () => resolve(); - this.#waiters.push(resolver); - }); - const timedOut = await Promise.race([ - released.then(() => false), - new Promise((resolve) => setTimeout(() => resolve(true), remaining)), - ]); - if (resolver) { - const index = this.#waiters.indexOf(resolver); - if (index >= 0) this.#waiters.splice(index, 1); - } - if (timedOut) return null; - } - } } 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 }; @@ -259,18 +355,18 @@ function estimateStructureTokens(value: unknown, limit: number): TokenEstimate { return { tokens, exhausted: stack.length > 0 && tokens < limit }; } -export async function admitChatStructure( +export function admitChatStructure( body: unknown, lease: ChatAdmissionLease | null, options: { controller?: ChatAdmissionController; + sessionId?: string; maxMessages?: number; heavyMessages?: number; heavyTools?: number; heavyTokens?: number; - queueMs?: number; } = {} -): Promise { +): ChatStructureAdmission { if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease }; const record = body as Record; @@ -301,9 +397,12 @@ 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); + const acquired = controller.tryAcquireHeavy(); return acquired ? { admit: true, lease: acquired } : { admit: false, response: structuralRejectionResponse(503, maxMessages) }; @@ -413,15 +512,16 @@ 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; const internalBypass = isInternalAdmissionBypass(request); const contentLength = parseContentLength(request.headers.get("content-length")); @@ -467,15 +567,15 @@ export async function admitChatRequest( } let lease: ChatAdmissionLease | null = null; - const reserve = async (): Promise => { + const reserve = (): boolean => { if (lease) return true; - lease = await controller.acquireHeavyWithin(queueMs); + lease = controller.tryAcquireHeavy(); 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 && !reserve()) { return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } @@ -494,7 +594,7 @@ export async function admitChatRequest( lease?.release(); return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; } - if (totalBytes >= largeBodyBytes && !(await reserve())) { + if (totalBytes >= largeBodyBytes && !reserve()) { await reader.cancel("chat admission capacity unavailable").catch(() => undefined); return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; } 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/per-connection-admission-9654.test.ts b/tests/unit/per-connection-admission-9654.test.ts new file mode 100644 index 0000000000..2b52716240 --- /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", () => { + // 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 = 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", () => { + // 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 = 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(); +});