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
This commit is contained in:
Brandon Bennett
2026-08-09 11:18:54 -07:00
parent 240b9b5bc4
commit 4839749a20
9 changed files with 865 additions and 76 deletions

View File

@@ -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)

View File

@@ -21,6 +21,7 @@ export interface ValidatedConfig {
adaptation: AdaptationParams;
maxRequestCost: number;
costConfig: ReturnType<typeof resolveCostConfig>;
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),
};
}

View File

@@ -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<QueuedPayload>;
private virtualQueue: FairCostQueue<{ recordId: string }>;
/** Per-connection virtual admission lanes (#9654). */
private readonly virtualLanes = new Map<string, {
queue: FairCostQueue<QueuedPayload>;
lastUsedMs: number;
}>();
/** Eviction timer for idle lanes; re-armed when a lane is created. */
private laneEvictionTimer: unknown = undefined;
private readonly active = new Map<string, ActiveLeaseRecord>();
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<QueuedPayload>; 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 {

View File

@@ -39,6 +39,7 @@ export const DEFAULT_ADAPTIVE_ADMISSION_CONFIG: Readonly<AdaptiveAdmissionConfig
maxQueueCost: 2000,
defaultMaxWaitMs: 5_000,
windowMs: 1_000,
virtualLanes: false,
});
const RUNTIME_STORE_KEY = Symbol.for("omniroute.adaptiveAdmission.runtime");
@@ -70,6 +71,7 @@ const ENV_KEYS = {
maxQueueCost: "ADAPTIVE_ADMISSION_MAX_QUEUE_COST",
defaultMaxWaitMs: "ADAPTIVE_ADMISSION_MAX_WAIT_MS",
windowMs: "ADAPTIVE_ADMISSION_WINDOW_MS",
virtualLanes: "OMNIROUTE_CHAT_VIRTUAL_LANES",
} as const;
function parsePositiveSafeInt(name: string, raw: string): number {
@@ -117,6 +119,11 @@ export function resolveAdaptiveAdmissionConfigFromEnv(
// Shared pure validation — accept exact documented maxima, reject core-invalid configs.
validateConfig(cfg);
// Per-connection virtual admission lanes (#9654) — opt-in via OMNIROUTE_CHAT_VIRTUAL_LANES.
const vlRaw = env[ENV_KEYS.virtualLanes];
cfg.virtualLanes = vlRaw === "1" || vlRaw === "true";
return cfg;
}

View File

@@ -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<AdmissionCostConfig>;
/** 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;

View File

@@ -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();

View File

@@ -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<ChatAdmissionLease | null> {
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<void>((resolve) => {
resolver = () => resolve();
this.#waiters.push(resolver);
});
const timedOut = await Promise.race([
released.then(() => false),
new Promise<boolean>((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<string, SessionRecord>();
#evictionTimer: ReturnType<typeof setTimeout> | 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> {
): ChatStructureAdmission {
if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease };
const record = body as Record<string, unknown>;
@@ -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<ChatRequestAdmission> {
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<boolean> => {
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) };
}

View File

@@ -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<number, { due: number; fn: () => 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> = {}): 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<AdmissionRequest> & { cost: number }): AdmissionRequest {
return {
tenantKey: "t-default",
...partial,
};
}
async function mustAdmit(
controller: AdaptiveAdmissionController,
request: AdmissionRequest
): Promise<import("../../open-sse/services/admission/types.ts").AdmissionLease> {
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<AdaptiveAdmissionConfig> = {}) {
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();
});
});

View File

@@ -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<string, string>, body = "{}"): Request {
const h: Record<string, string> = { "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();
});