mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-20 14:12:59 +03:00
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:
@@ -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)
|
||||
@@ -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),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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) };
|
||||
}
|
||||
|
||||
240
tests/unit/admission-virtual-lanes-9654.test.ts
Normal file
240
tests/unit/admission-virtual-lanes-9654.test.ts
Normal 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();
|
||||
});
|
||||
});
|
||||
205
tests/unit/per-connection-admission-9654.test.ts
Normal file
205
tests/unit/per-connection-admission-9654.test.ts
Normal 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();
|
||||
});
|
||||
Reference in New Issue
Block a user