mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-10 17:22:17 +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 * fix(admission): restore bounded queue-wait on per-connection lanes (#9654) The per-connection lane refactor dropped the bounded queue-wait (acquireHeavyWithin / #waiters / queueMs). #9654's acceptance criteria and #9608 section C prefer server-side wait/pacing up to defaultMaxWaitMs over an instant retryable 503. - ChatAdmissionController: re-add #waiters FIFO + acquireHeavyWithin(timeoutMs); queueMs: 0 preserves the instant-503 path - admitChatStructure and admitChatRequest.reserve are async again and take queueMs - route: pass CHAT_ADMISSION_QUEUE_MAX_MS and await the admission calls - per-connection lane tests await the async admitChatStructure Admission suite: 114/114 pass (bun test, 7 files). * chore: re-trigger CI after dast-smoke infra cancellation (#9654) * feat(admission): cancel queue-wait on client abort (#9654) U2 from KC plan 2026-08-09-001. Thread the request AbortSignal through acquireHeavyWithin so a disconnected client stops parking in the FIFO for the full queueMs. - acquireHeavyWithin(timeoutMs, signal?): on abort the waiter is removed from the FIFO immediately and the promise resolves null early; pre-aborted signals never park; the deadline timer is cleared when abort/release wins the race - admitChatRequest reserve() passes request.signal; admitChatStructure gains options.signal; the route threads request.signal - 5 exact-assertion tests (settle-early, pre-aborted, byte-heavy, structural, FIFO-preservation): 119/119 across the 7-file suite * fix(admission): bound queued bytes for the queue-wait heap valve (#9654) U3 from KC plan 2026-08-09-001. The restored queue-wait parks fully-buffered bodies; without a cap, several large coding-agent bodies (~750 KB) waiting at once recreates the #4380 heap amplification this module was built to stop. - acquireHeavyWithin(timeoutMs, signal?, queuedBytes): each parked waiter is charged its buffered size against CHAT_ADMISSION_MAX_QUEUED_BYTES (default 4 MB); over-budget waits reject immediately with a retryable 503 and never park. The charge is released on wake, abort, or timeout. - Real sizes threaded from admitChatRequest (declared length / sniffed bytes); structural waits charge the conservative 256 KB weight. - Lower default OMNIROUTE_CHAT_ADMISSION_QUEUE_MS to 2000ms (was 5000ms). - Env vars documented in .env.example; 6 exact-assertion tests: 125/125 across the 7-file admission suite (was 119). * docs: map the two admission-lane systems for operators (#9654) U5 from KC plan 2026-08-09-001. Verifies lane metrics are exposed by the health payload (GET /api/monitoring/health -> adaptiveAdmission -> lane* fields) and records which lane system reports where: byte-level per-connection lanes (always on, memory scope) vs adaptive virtual lanes (opt-in via OMNIROUTE_CHAT_VIRTUAL_LANES, dispatch scope) plus the explicit opt-in ops note. * docs: add required frontmatter to admission-lanes doc (dast-smoke build fix) * docs: sync env vars with .env.example and ENVIRONMENT.md (docs gate fix) * fix(admission): complete REJECT_MAP, literal lane env read, split oversized test file Three CI-gate fixes surfaced by the post-merge check run (head3de77166e): 1. open-sse-typecheck (TS2741): REJECT_MAP was missing the ADMISSION_LANE_EVICTED entry that controller.ts:662 emits on lane eviction. Add the 503 mapping so the Record<AdmissionRejectCode, RejectHttpMapping> is total. 2. Docs Gates fabricated-claim: OMNIROUTE_CHAT_VIRTUAL_LANES was read dynamically via ENV_KEYS.virtualLanes (env[key]), invisible to the literal env.X scanner. Read it literally — behavior-identical, doc claim now verifiable. 3. check:file-size: chat-body-admission.test.ts (1307 lines) exceeded the 1000-line new-file cap. Split the queue-wait/abort/heap-valve section into chat-body-admission-queue.test.ts (818 + 513 lines, both under cap). Suite: 125/125 across 8 files. All three checkers pass locally. * refactor(admission): drop dead ENV_KEYS.virtualLanes entry + lock lane-evicted mapping test Code-review follow-up on50c93d266: 1. ENV_KEYS.virtualLanes is now unreferenced since the literal env read landed; remove it so the config map only lists keys actually read through the map. 2. Add an exact-assertion runtime test for the ADMISSION_LANE_EVICTED mapping: a queued lane waiter evicted by the 60s idle TTL rejects with 503 / admission_lane_evicted / Retry-After 1 / sanitized body (no raw tenant key). Proves the REJECT_MAP entry end-to-end through buildAdmissionRejectResponse. Suite: 126/126 (17 in runtime file, 125 in the 8-file admission suite). --------- Co-authored-by: Brandon Bennett <brandonbennett@macbookair.myfiosgateway.com>
This commit is contained in:
Brandon Bennett
12
.env.example
12
.env.example
@@ -350,6 +350,18 @@ ALLOW_API_KEY_REVEAL=false
|
||||
# by OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT and the heap-pressure shed instead. Set a positive
|
||||
# value only on memory-constrained deployments that need a hard ceiling.
|
||||
# OMNIROUTE_CHAT_HARD_MAX_MESSAGES=0
|
||||
# How long a heavy request waits for heavyweight capacity before a retryable 503.
|
||||
# A short bounded wait serializes agent bursts instead of an instant 503; 0 = instant.
|
||||
# Default 2000 (2s).
|
||||
# OMNIROUTE_CHAT_ADMISSION_QUEUE_MS=2000
|
||||
# Queued-bytes budget for the admission wait: bounds total buffered body bytes parked
|
||||
# per lane so the wait cannot amplify the heap (#4380). Over-budget waits 503 immediately.
|
||||
# Default 4194304 (4 MB).
|
||||
# OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES=4194304
|
||||
# Per-connection virtual admission lanes (#9654): idle-lane eviction TTL. Default 60000 (60s).
|
||||
# OMNIROUTE_CHAT_VIRTUAL_TTL_MS=60000
|
||||
# Per-connection virtual admission lanes (#9654): max concurrent sessions (lanes). Default 64.
|
||||
# OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS=64
|
||||
|
||||
# Hard cap (bytes) for a non-streaming upstream response buffered fully into memory
|
||||
# (#5152). Past this the upstream reader is cancelled and the request fails fast
|
||||
|
||||
@@ -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)
|
||||
50
docs/architecture/admission-lanes.md
Normal file
50
docs/architecture/admission-lanes.md
Normal file
@@ -0,0 +1,50 @@
|
||||
---
|
||||
title: "Admission lanes — two lane systems, what gates each, where each reports"
|
||||
status: active
|
||||
lastUpdated: 2026-08-09
|
||||
---
|
||||
|
||||
# Admission lanes (#9654) — two lane systems, what gates each, where each reports
|
||||
|
||||
OmniRoute has **two** process-local lane systems with different scopes. They are
|
||||
complementary; operators should know which one they are looking at.
|
||||
|
||||
## 1. Byte-level per-connection lanes (`chatBodyAdmission.ts`)
|
||||
|
||||
- **Scope:** the buffered-body/heap path for `POST /v1/chat/completions`. Guards
|
||||
against heap amplification from large coding-agent bodies (#4380).
|
||||
- **Gate:** **always on.** Each distinct API key (hashed) — or `anonymous` — gets its
|
||||
own lane with `CHAT_MAX_HEAVY_IN_FLIGHT` capacity, so one session's burst cannot
|
||||
starve another session's heavyweight slot.
|
||||
- **Tuning:**
|
||||
- `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` — idle-lane eviction (default 60000)
|
||||
- `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` — lane count cap (default 64)
|
||||
- `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` — queue-wait before 503 (default 2000)
|
||||
- `OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES` — queued-bytes heap valve (default 4 MB)
|
||||
- **Reports:** not in `GET /api/monitoring/health` today; observable via
|
||||
`PerConnectionAdmissionController.snapshot()` (sessionId hash, activeHeavy, idleMs).
|
||||
|
||||
## 2. Adaptive runtime virtual lanes (`open-sse/services/admission`)
|
||||
|
||||
- **Scope:** tenant-key admission for provider dispatch — queue cost, latency-guided
|
||||
limit adaptation, lane queueing, and lane metrics.
|
||||
- **Gate:** **opt-in.** Disabled unless `OMNIROUTE_CHAT_VIRTUAL_LANES=true`. Without it,
|
||||
the adaptive controller keeps the shared queue behavior (criterion 1 of #9654 only
|
||||
holds once an operator enables lanes).
|
||||
- **Tuning:** `OMNIROUTE_CHAT_VIRTUAL_LANES` + adaptive config (`maxQueueCount`,
|
||||
`maxQueueCost`, `defaultMaxWaitMs`, …).
|
||||
- **Reports:** `GET /api/monitoring/health` → `adaptiveAdmission` → `laneCount`,
|
||||
`laneQueuedCount`, `laneQueuedCost`, `laneTenants` (opaque lane IDs, never raw keys).
|
||||
|
||||
## Which one is showing in a dashboard
|
||||
|
||||
- `adaptiveAdmission.laneCount` / `laneTenants` → **adaptive virtual lanes** (system 2).
|
||||
- A health payload with **no** `adaptiveAdmission.lane*` fields usually means
|
||||
`OMNIROUTE_CHAT_VIRTUAL_LANES` is unset — the byte-level lanes (system 1) are still
|
||||
active, but nothing under `adaptiveAdmission` will report lane data until it is enabled.
|
||||
|
||||
## Why both exist
|
||||
|
||||
The byte-level lanes bound the memory-heavy parse/compress path; the adaptive lanes
|
||||
bound dispatch cost per tenant. #9654's criterion 1 ("one session's burst does not 503
|
||||
another") is enforced by system 1 unconditionally and by system 2 once opt-in is enabled.
|
||||
@@ -1398,7 +1398,10 @@ These settings were introduced after the previous environment-contract snapshot.
|
||||
|
||||
| Variable | Default | Source File | Description |
|
||||
| --- | --- | --- | --- |
|
||||
| `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` | `5000` | `src/shared/middleware/chatBodyAdmission.ts` | Maximum wait for a heavyweight chat admission slot before a retryable `503`; `0` restores immediate rejection. |
|
||||
| `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` | `2000` | `src/shared/middleware/chatBodyAdmission.ts` | Maximum wait for a heavyweight chat admission slot before a retryable `503`; a short bounded wait serializes agent bursts instead of an instant `503`. `0` restores immediate rejection. |
|
||||
| `OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES` | `4194304` (4 MB) | `src/shared/middleware/chatBodyAdmission.ts` | Queued-bytes budget for the admission wait (#9654): bounds total buffered body bytes parked per lane so the wait cannot amplify the heap (#4380). Over-budget waits receive a retryable `503` immediately. |
|
||||
| `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` | `60000` (60 s) | `src/shared/middleware/chatBodyAdmission.ts` | Per-connection virtual admission lanes (#9654): idle-lane eviction TTL. |
|
||||
| `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` | `64` | `src/shared/middleware/chatBodyAdmission.ts` | Per-connection virtual admission lanes (#9654): max concurrent sessions (lanes). |
|
||||
| `OMNIROUTE_RUNNOW_TIMEOUT_MS` | `30000` | `src/app/api/jobs/[id]/run-now/route.ts` | Bounds how long a run-now call waits for an in-flight job before starting the queued run. |
|
||||
| `CHAT_LOG_MAX_BODY_KB` | `1024` | `src/lib/logEnv.ts` | Maximum request or response body size before log summarization, in KiB. |
|
||||
| `ADOBE_FIREFLY_BROWSER_REFRESH` | enabled | `open-sse/services/adobeFireflySession.ts` | Keeps IMS and browser-risk state fresh through account-scoped Chrome CDP sessions; set `0` to disable. |
|
||||
|
||||
@@ -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");
|
||||
@@ -117,6 +118,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.OMNIROUTE_CHAT_VIRTUAL_LANES;
|
||||
cfg.virtualLanes = vlRaw === "1" || vlRaw === "true";
|
||||
|
||||
return cfg;
|
||||
}
|
||||
|
||||
@@ -236,6 +242,12 @@ const REJECT_MAP: Record<AdmissionRejectCode, RejectHttpMapping> = {
|
||||
message: "Service temporarily unavailable",
|
||||
retryAfter: "1",
|
||||
},
|
||||
ADMISSION_LANE_EVICTED: {
|
||||
status: 503,
|
||||
code: "admission_lane_evicted",
|
||||
message: "Connection lane evicted",
|
||||
retryAfter: "1",
|
||||
},
|
||||
};
|
||||
|
||||
function isAdmissionRejectError(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -20,6 +20,7 @@ import {
|
||||
CHAT_ADMISSION_QUEUE_MAX_MS,
|
||||
releaseChatAdmissionAfterHandler,
|
||||
releaseChatAdmissionWhenDone,
|
||||
resolveSessionId,
|
||||
} from "@/shared/middleware/chatBodyAdmission";
|
||||
import {
|
||||
readCompressionRequestHeader,
|
||||
@@ -101,7 +102,9 @@ export async function POST(request) {
|
||||
// Reserve heavyweight capacity atomically and ingest the body with a hard byte bound
|
||||
// BEFORE JSON parsing. Missing or dishonest Content-Length values cannot bypass
|
||||
// the actual-byte limit. Capacity exhaustion is retryable rather than process-fatal.
|
||||
const sessionId = resolveSessionId(request);
|
||||
const admissionResult = await admitChatRequest(request, {
|
||||
sessionId,
|
||||
queueMs: CHAT_ADMISSION_QUEUE_MAX_MS,
|
||||
});
|
||||
if (admissionResult.admit === false) return admissionResult.response;
|
||||
@@ -147,7 +150,9 @@ export async function POST(request) {
|
||||
}
|
||||
|
||||
const structuralAdmission = await admitChatStructure(parsedBody, admission.lease, {
|
||||
sessionId,
|
||||
queueMs: CHAT_ADMISSION_QUEUE_MAX_MS,
|
||||
signal: request.signal,
|
||||
});
|
||||
if (structuralAdmission.admit === false) {
|
||||
admission.lease?.release();
|
||||
|
||||
@@ -1,14 +1,25 @@
|
||||
/**
|
||||
* Bounded admission for POST /v1/chat/completions.
|
||||
* Process-local bounded admission for POST /v1/chat/completions.
|
||||
*
|
||||
* Large chat bodies amplify into multiple transient representations while they are parsed,
|
||||
* translated, compressed, and dispatched. A heap snapshot alone cannot prevent two healthy
|
||||
* requests from entering that allocation-heavy path together. This module reserves process-
|
||||
* local heavyweight capacity before parsing and enforces the hard limit against bytes read,
|
||||
* not an untrusted Content-Length header.
|
||||
*
|
||||
* Per-connection virtual admission lanes (#9654): each distinct API-key (or anonymous)
|
||||
* bucket gets its own FairCostQueue so one connection cannot exhaust heavyweight capacity
|
||||
* and starve others. Idle sessions are auto-evicted after a TTL.
|
||||
*/
|
||||
|
||||
import { CORS_HEADERS } from "../utils/cors";
|
||||
import { createHash } from "crypto";
|
||||
|
||||
|
||||
const OMNIROUTE_CHAT_VIRTUAL_TTL_MS = parsePositiveInt(
|
||||
process.env.OMNIROUTE_CHAT_VIRTUAL_TTL_MS,
|
||||
60_000
|
||||
);
|
||||
|
||||
function parsePositiveInt(value: string | undefined, fallback: number): number {
|
||||
const parsed = Number.parseInt(String(value), 10);
|
||||
@@ -30,7 +41,7 @@ export const CHAT_HARD_MAX_BODY_BYTES = parsePositiveInt(
|
||||
50 * 1024 * 1024
|
||||
);
|
||||
|
||||
const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt(
|
||||
export const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt(
|
||||
process.env.OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT,
|
||||
1
|
||||
);
|
||||
@@ -44,7 +55,20 @@ const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt(
|
||||
*/
|
||||
export const CHAT_ADMISSION_QUEUE_MAX_MS = parseNonNegativeInt(
|
||||
process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS,
|
||||
5000
|
||||
2000
|
||||
);
|
||||
|
||||
/**
|
||||
* Queued-bytes budget for the admission wait (#9654 / U3). A parked waiter holds a
|
||||
* fully-buffered request body; several large coding-agent bodies (~750 KB) waiting at
|
||||
* once is exactly the heap-amplification scenario chatBodyAdmission was built to stop
|
||||
* (#4380). Each lane's controller charges every parked waiter's buffered size against
|
||||
* this budget and rejects over-budget waits immediately (retryable 503) instead of
|
||||
* parking. Bytes are released when a waiter wakes, aborts, or times out.
|
||||
*/
|
||||
export const CHAT_ADMISSION_MAX_QUEUED_BYTES = parsePositiveInt(
|
||||
process.env.OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES,
|
||||
4 * 1024 * 1024
|
||||
);
|
||||
|
||||
export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt(
|
||||
@@ -94,18 +118,30 @@ export interface ChatAdmissionLease {
|
||||
*/
|
||||
export class ChatAdmissionController {
|
||||
#activeHeavy = 0;
|
||||
#queuedBytes = 0;
|
||||
#waiters: Array<() => void> = [];
|
||||
|
||||
constructor(readonly maxHeavyInFlight = 1) {
|
||||
constructor(
|
||||
readonly maxHeavyInFlight = 1,
|
||||
readonly maxQueuedBytes = CHAT_ADMISSION_MAX_QUEUED_BYTES
|
||||
) {
|
||||
if (!Number.isSafeInteger(maxHeavyInFlight) || maxHeavyInFlight < 1) {
|
||||
throw new RangeError("maxHeavyInFlight must be a positive integer");
|
||||
}
|
||||
if (!Number.isSafeInteger(maxQueuedBytes) || maxQueuedBytes < 0) {
|
||||
throw new RangeError("maxQueuedBytes must be a non-negative integer");
|
||||
}
|
||||
}
|
||||
|
||||
get activeHeavy(): number {
|
||||
return this.#activeHeavy;
|
||||
}
|
||||
|
||||
/** Total buffered bytes currently parked in the FIFO (heap valve accounting). */
|
||||
get queuedBytes(): number {
|
||||
return this.#queuedBytes;
|
||||
}
|
||||
|
||||
tryAcquireHeavy(): ChatAdmissionLease | null {
|
||||
if (this.#activeHeavy >= this.maxHeavyInFlight) return null;
|
||||
this.#activeHeavy += 1;
|
||||
@@ -128,27 +164,71 @@ export class ChatAdmissionController {
|
||||
* release. Resolves `null` when the deadline expires with no capacity freed, in
|
||||
* which case the caller answers the retryable 503. `timeoutMs <= 0` is the
|
||||
* legacy immediate-reject path. Waiters are served FIFO.
|
||||
*
|
||||
* When `signal` aborts while parked (client disconnect), the waiter is removed
|
||||
* from the FIFO immediately and the promise resolves `null` early instead of
|
||||
* parking for the full `timeoutMs` — the caller's 503 is dropped on the dead
|
||||
* connection, so no capacity is consumed and the freed slot never wakes a
|
||||
* waiter the client no longer needs. A signal that is already aborted never
|
||||
* parks at all.
|
||||
*
|
||||
* `queuedBytes` is the buffered body size this waiter will hold while parked;
|
||||
* it is charged against `maxQueuedBytes` so a burst of large bodies cannot
|
||||
* amplify the heap (#4380). An over-budget wait is rejected immediately with
|
||||
* `null` (retryable 503) and never parks; the charge is released on wake,
|
||||
* abort, or timeout.
|
||||
*/
|
||||
async acquireHeavyWithin(timeoutMs: number): Promise<ChatAdmissionLease | null> {
|
||||
async acquireHeavyWithin(
|
||||
timeoutMs: number,
|
||||
signal?: AbortSignal,
|
||||
queuedBytes = 0
|
||||
): Promise<ChatAdmissionLease | null> {
|
||||
const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs));
|
||||
for (;;) {
|
||||
if (signal?.aborted) return null;
|
||||
const lease = this.tryAcquireHeavy();
|
||||
if (lease) return lease;
|
||||
const remaining = deadline - Date.now();
|
||||
if (remaining <= 0) return null;
|
||||
// Heap valve: refuse to park when the queued-bytes budget is exhausted.
|
||||
if (queuedBytes > 0 && this.#queuedBytes + queuedBytes > this.maxQueuedBytes) {
|
||||
return null;
|
||||
}
|
||||
this.#queuedBytes += queuedBytes;
|
||||
let resolver: (() => void) | null = null;
|
||||
const released = new Promise<void>((resolve) => {
|
||||
resolver = () => resolve();
|
||||
this.#waiters.push(resolver);
|
||||
});
|
||||
const timedOut = await Promise.race([
|
||||
let deadlineTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
const races: Array<Promise<boolean>> = [
|
||||
released.then(() => false),
|
||||
new Promise<boolean>((resolve) => setTimeout(() => resolve(true), remaining)),
|
||||
]);
|
||||
new Promise<boolean>((resolve) => {
|
||||
deadlineTimer = setTimeout(() => resolve(true), remaining);
|
||||
}),
|
||||
];
|
||||
let onAbort: (() => void) | null = null;
|
||||
if (signal) {
|
||||
races.push(
|
||||
new Promise<boolean>((resolve) => {
|
||||
const listener = () => resolve(true);
|
||||
onAbort = listener;
|
||||
signal.addEventListener("abort", listener, { once: true });
|
||||
// Already-aborted signals must settle without parking.
|
||||
if (signal.aborted) resolve(true);
|
||||
})
|
||||
);
|
||||
}
|
||||
const timedOut = await Promise.race(races);
|
||||
// The waiter has left the FIFO (wake, abort, or timeout) — release its charge.
|
||||
this.#queuedBytes = Math.max(0, this.#queuedBytes - queuedBytes);
|
||||
if (resolver) {
|
||||
const index = this.#waiters.indexOf(resolver);
|
||||
if (index >= 0) this.#waiters.splice(index, 1);
|
||||
}
|
||||
// Cancel the deadline timer when abort/release wins; a fired timer is a no-op.
|
||||
if (deadlineTimer) clearTimeout(deadlineTimer);
|
||||
if (onAbort) signal?.removeEventListener("abort", onAbort);
|
||||
if (timedOut) return null;
|
||||
}
|
||||
}
|
||||
@@ -156,6 +236,142 @@ export class ChatAdmissionController {
|
||||
|
||||
const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT);
|
||||
|
||||
/**
|
||||
* Per-connection virtual admission lanes (#9654).
|
||||
*
|
||||
* Maps a sessionId (API-key hash or "anonymous") → ChatAdmissionController.
|
||||
Each connection gets its own bounded heavyweight capacity so one connection
|
||||
* cannot exhaust `CHAT_MAX_HEAVY_IN_FLIGHT` and starve others at the byte-level
|
||||
* admission stage.
|
||||
*
|
||||
* Idle sessions are auto-evicted after OMNIROUTE_CHAT_VIRTUAL_TTL_MS
|
||||
* (default 60s) to prevent unbounded Map growth.
|
||||
*/
|
||||
const OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS = parsePositiveInt(
|
||||
process.env.OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS,
|
||||
64
|
||||
);
|
||||
|
||||
export function resolveSessionId(request: Request): string {
|
||||
// Reuse the existing internal-bypass auth extraction: bearer token from
|
||||
// Authorization, x-api-key (Anthropic-style), or Google API key header.
|
||||
const authHeader = request.headers.get("authorization") || "";
|
||||
const bearerMatch = /^bearer\s+(\S+)$/i.exec(authHeader.trim());
|
||||
if (bearerMatch) {
|
||||
return "key_" + createHash("sha256").update(bearerMatch[1]).digest("hex").slice(0, 16);
|
||||
}
|
||||
const xApiKey = request.headers.get("x-api-key") || "";
|
||||
if (xApiKey.trim().length > 0) {
|
||||
return "key_" + createHash("sha256").update(xApiKey.trim()).digest("hex").slice(0, 16);
|
||||
}
|
||||
const xGoogApiKey = request.headers.get("x-goog-api-key") || "";
|
||||
if (xGoogApiKey.trim().length > 0) {
|
||||
return "key_" + createHash("sha256").update(xGoogApiKey.trim()).digest("hex").slice(0, 16);
|
||||
}
|
||||
return "anonymous";
|
||||
}
|
||||
|
||||
interface SessionRecord {
|
||||
controller: ChatAdmissionController;
|
||||
lastUsedMs: number;
|
||||
}
|
||||
|
||||
export class PerConnectionAdmissionController {
|
||||
#sessions = new Map<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 };
|
||||
@@ -264,11 +480,13 @@ export async function admitChatStructure(
|
||||
lease: ChatAdmissionLease | null,
|
||||
options: {
|
||||
controller?: ChatAdmissionController;
|
||||
sessionId?: string;
|
||||
maxMessages?: number;
|
||||
heavyMessages?: number;
|
||||
heavyTools?: number;
|
||||
heavyTokens?: number;
|
||||
queueMs?: number;
|
||||
signal?: AbortSignal;
|
||||
} = {}
|
||||
): Promise<ChatStructureAdmission> {
|
||||
if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease };
|
||||
@@ -301,8 +519,18 @@ export async function admitChatStructure(
|
||||
estimatedTokens >= heavyTokens;
|
||||
if (!heavy || lease) return { admit: true, lease };
|
||||
|
||||
const acquired = await (options.controller ?? defaultAdmissionController).acquireHeavyWithin(
|
||||
options.queueMs ?? 0
|
||||
const controller =
|
||||
options.controller ??
|
||||
(options.sessionId
|
||||
? perConnectionAdmissionController.getController(options.sessionId)
|
||||
: defaultAdmissionController);
|
||||
// Structural-only waits happen on byte-light bodies (a byte-heavy body already
|
||||
// holds the byte-stage lease), so the conservative 256KB weight bounds the
|
||||
// parsed JSON the waiter keeps resident while parked.
|
||||
const acquired = await controller.acquireHeavyWithin(
|
||||
options.queueMs ?? 0,
|
||||
options.signal,
|
||||
CHAT_LARGE_BODY_BYTES
|
||||
);
|
||||
return acquired
|
||||
? { admit: true, lease: acquired }
|
||||
@@ -413,12 +641,15 @@ export async function admitChatRequest(
|
||||
request: Request,
|
||||
options: {
|
||||
controller?: ChatAdmissionController;
|
||||
sessionId?: string;
|
||||
largeBodyBytes?: number;
|
||||
hardMaxBytes?: number;
|
||||
queueMs?: number;
|
||||
} = {}
|
||||
): Promise<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;
|
||||
@@ -467,15 +698,19 @@ export async function admitChatRequest(
|
||||
}
|
||||
|
||||
let lease: ChatAdmissionLease | null = null;
|
||||
const reserve = async (): Promise<boolean> => {
|
||||
const reserve = async (bytes = 0): Promise<boolean> => {
|
||||
if (lease) return true;
|
||||
lease = await controller.acquireHeavyWithin(queueMs);
|
||||
lease = await controller.acquireHeavyWithin(queueMs, request.signal, bytes);
|
||||
return lease !== null;
|
||||
};
|
||||
|
||||
// A known-large declaration can reserve before ingestion. Unknown lengths are boundedly
|
||||
// sniffed below; this avoids consuming scarce heavyweight capacity for small chunked bodies.
|
||||
if (contentLength !== null && contentLength >= largeBodyBytes && !(await reserve())) {
|
||||
if (
|
||||
contentLength !== null &&
|
||||
contentLength >= largeBodyBytes &&
|
||||
!(await reserve(Math.min(contentLength, hardMaxBytes)))
|
||||
) {
|
||||
return { admit: false, response: rejectionResponse(503, hardMaxBytes) };
|
||||
}
|
||||
|
||||
@@ -494,7 +729,7 @@ export async function admitChatRequest(
|
||||
lease?.release();
|
||||
return { admit: false, response: rejectionResponse(413, hardMaxBytes) };
|
||||
}
|
||||
if (totalBytes >= largeBodyBytes && !(await reserve())) {
|
||||
if (totalBytes >= largeBodyBytes && !(await reserve(totalBytes))) {
|
||||
await reader.cancel("chat admission capacity unavailable").catch(() => undefined);
|
||||
return { admit: false, response: rejectionResponse(503, hardMaxBytes) };
|
||||
}
|
||||
|
||||
@@ -600,6 +600,51 @@ describe("rejection mapping", () => {
|
||||
runtime.dispose();
|
||||
oversizedRuntime.dispose();
|
||||
});
|
||||
|
||||
it("maps ADMISSION_LANE_EVICTED to a sanitized 503 with Retry-After", async () => {
|
||||
const runtime = makeRuntime(clock, {
|
||||
config: enforceConfig({
|
||||
initialLimit: 1,
|
||||
minLimit: 1,
|
||||
maxLimit: 1,
|
||||
maxQueueCount: 4,
|
||||
maxQueueCost: 40,
|
||||
defaultMaxWaitMs: 120_000, // must outlive the 60s lane TTL so the lane eviction wins
|
||||
windowMs: 1_000,
|
||||
virtualLanes: true,
|
||||
cost: { maxRequestCost: 1, baseCost: 1 },
|
||||
}),
|
||||
});
|
||||
const hold = await runtime.acquire({
|
||||
tenantKey: "hold",
|
||||
body: { stream: true },
|
||||
});
|
||||
assert.equal(hold.status, "admitted");
|
||||
|
||||
// Park a waiter in a virtual lane; its own deadline is far beyond the TTL.
|
||||
const pending = runtime.acquire({
|
||||
tenantKey: "lane-waiter",
|
||||
body: { stream: true },
|
||||
maxWaitMs: 120_000,
|
||||
});
|
||||
|
||||
// Advance past the 60s lane TTL: the window tick evicts idle lanes, which
|
||||
// drains and rejects the queued waiter with ADMISSION_LANE_EVICTED.
|
||||
clock.advance(60_001);
|
||||
|
||||
const rejected = await pending;
|
||||
assert.equal(rejected.status, "rejected");
|
||||
if (rejected.status === "rejected") {
|
||||
assert.equal(rejected.code, "admission_lane_evicted");
|
||||
assert.equal(rejected.response.status, 503);
|
||||
assert.equal(rejected.response.headers.get("Retry-After"), "1");
|
||||
const body = await parseJson(rejected.response);
|
||||
assert.equal(body.error.code, "admission_lane_evicted");
|
||||
assert.ok(!JSON.stringify(body).includes("lane-waiter"));
|
||||
}
|
||||
if (hold.status === "admitted") hold.lease.release();
|
||||
runtime.dispose();
|
||||
});
|
||||
});
|
||||
|
||||
describe("resource pressure integration", () => {
|
||||
|
||||
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();
|
||||
});
|
||||
});
|
||||
513
tests/unit/chat-body-admission-queue.test.ts
Normal file
513
tests/unit/chat-body-admission-queue.test.ts
Normal file
@@ -0,0 +1,513 @@
|
||||
// #9654: queue-wait, AbortSignal cancellation, and the queued-bytes heap valve.
|
||||
// Split from chat-body-admission.test.ts to stay under the 1000-line new-file cap.
|
||||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
|
||||
const admissionModule = await import("../../src/shared/middleware/chatBodyAdmission.ts");
|
||||
const {
|
||||
admitChatRequest,
|
||||
admitChatStructure,
|
||||
ChatAdmissionController,
|
||||
CHAT_ADMISSION_QUEUE_MAX_MS,
|
||||
CHAT_ADMISSION_MAX_QUEUED_BYTES,
|
||||
CHAT_LARGE_BODY_BYTES,
|
||||
} = admissionModule;
|
||||
|
||||
function chatRequest(body: string, contentLength: string | null = String(body.length)): Request {
|
||||
const headers: Record<string, string> = { "content-type": "application/json" };
|
||||
if (contentLength !== null) headers["content-length"] = contentLength;
|
||||
return new Request("http://x/v1/chat/completions", {
|
||||
method: "POST",
|
||||
headers,
|
||||
body,
|
||||
});
|
||||
}
|
||||
|
||||
test("a heavy structural request waits for capacity instead of failing immediately", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const pending = admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 500,
|
||||
}
|
||||
);
|
||||
|
||||
// Capacity is still busy: the request must not have resolved (admit/reject) yet.
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(settled, false, "must wait while capacity is busy");
|
||||
|
||||
held.release();
|
||||
const result = await pending;
|
||||
assert.equal(result.admit, true);
|
||||
if (result.admit) {
|
||||
assert.equal(controller.activeHeavy, 1, "waiting request acquires the freed lease");
|
||||
result.lease?.release();
|
||||
}
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("waiting for admission times out into a retryable 503", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const started = Date.now();
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 50,
|
||||
}
|
||||
);
|
||||
|
||||
assert.equal(result.admit, false);
|
||||
if (!result.admit) {
|
||||
assert.equal(result.response.status, 503);
|
||||
assert.equal(result.response.headers.get("retry-after"), "1");
|
||||
assert.equal((await result.response.json()).error.code, "chat_admission_busy");
|
||||
}
|
||||
assert.ok(Date.now() - started >= 40, "must wait for the queue deadline before rejecting");
|
||||
assert.equal(controller.activeHeavy, 1, "the holder keeps its lease");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("byte-heavy admission waits for capacity when queueMs is set", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] });
|
||||
const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 500 };
|
||||
|
||||
const first = await admitChatRequest(chatRequest(body), options);
|
||||
assert.equal(first.admit, true);
|
||||
if (!first.admit) return;
|
||||
|
||||
const second = admitChatRequest(chatRequest(body), options);
|
||||
let secondSettled = false;
|
||||
void second.then(() => {
|
||||
secondSettled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(
|
||||
secondSettled,
|
||||
false,
|
||||
"second heavy request must queue while the first holds capacity"
|
||||
);
|
||||
|
||||
first.lease?.release();
|
||||
const secondResult = await second;
|
||||
assert.equal(secondResult.admit, true, "second request acquires capacity after release");
|
||||
if (secondResult.admit) secondResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("expired admission queue keeps the legacy immediate 503 behaviour", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 0,
|
||||
}
|
||||
);
|
||||
|
||||
assert.equal(result.admit, false);
|
||||
if (!result.admit) assert.equal(result.response.status, 503);
|
||||
held.release();
|
||||
});
|
||||
|
||||
test("admission waiters are served FIFO as capacity frees", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const body = {
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
};
|
||||
const options = {
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 500,
|
||||
};
|
||||
const first = admitChatStructure(body, null, options);
|
||||
const second = admitChatStructure(body, null, options);
|
||||
|
||||
held.release();
|
||||
const firstResult = await first;
|
||||
assert.equal(firstResult.admit, true);
|
||||
if (firstResult.admit) firstResult.lease?.release();
|
||||
const secondResult = await second;
|
||||
assert.equal(secondResult.admit, true);
|
||||
if (secondResult.admit) secondResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
// ── AbortSignal support in acquireHeavyWithin (#9654 / U2) ────────────────
|
||||
// A disconnected client must not keep parking in the admission queue for the
|
||||
// full queueMs. On abort the waiter is removed from the FIFO immediately and
|
||||
// the acquire resolves `null` early (the caller's 503 is dropped on the dead
|
||||
// connection); no capacity is consumed and the freed slot does not wake it.
|
||||
|
||||
test("aborting the admission wait settles early, grants no lease, and removes the waiter", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const abortController = new AbortController();
|
||||
const pending = controller.acquireHeavyWithin(2_000, abortController.signal);
|
||||
|
||||
// Parked while capacity is busy.
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(settled, false, "must be parked while capacity is busy");
|
||||
|
||||
abortController.abort();
|
||||
|
||||
// Must settle well before the 2s deadline.
|
||||
let settledAfterAbort = false;
|
||||
void pending.then(() => {
|
||||
settledAfterAbort = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
assert.equal(settledAfterAbort, true, "abort must settle the wait promptly, not park for queueMs");
|
||||
|
||||
const lease = await pending;
|
||||
assert.equal(lease, null, "abort must not grant a lease");
|
||||
assert.equal(controller.activeHeavy, 1, "the holder keeps its lease; the aborted wait consumed nothing");
|
||||
|
||||
// Releasing must NOT wake the removed waiter: capacity stays free.
|
||||
held.release();
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
assert.equal(
|
||||
controller.activeHeavy,
|
||||
0,
|
||||
"releasing after abort must not wake the removed waiter"
|
||||
);
|
||||
});
|
||||
|
||||
test("aborting the head waiter preserves FIFO order for remaining waiters", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const firstAbort = new AbortController();
|
||||
const first = controller.acquireHeavyWithin(2_000, firstAbort.signal);
|
||||
const second = controller.acquireHeavyWithin(2_000);
|
||||
|
||||
// Both are parked, head-first.
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
|
||||
// Abort the HEAD waiter: it must leave the queue without disturbing the rest.
|
||||
firstAbort.abort();
|
||||
assert.equal(await first, null, "head waiter returns null on abort");
|
||||
|
||||
// The remaining waiter is now first in line and must get the freed capacity.
|
||||
held.release();
|
||||
const secondLease = await second;
|
||||
assert.ok(secondLease, "remaining waiter must acquire the freed capacity");
|
||||
secondLease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
// ── Heap-pressure safety valve (#9654 / U3) ───────────────────────────────
|
||||
// The queue-wait parks fully-buffered bodies; the queued-bytes cap bounds the
|
||||
// total buffered memory parked per lane so the wait cannot recreate the #4380
|
||||
// heap amplification. Over-budget waits are rejected immediately (503).
|
||||
|
||||
test("queued-bytes cap rejects an over-budget wait without parking", async () => {
|
||||
const controller = new ChatAdmissionController(1, 200);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
// First waiter parks within budget.
|
||||
const first = controller.acquireHeavyWithin(2_000, undefined, 150);
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(controller.queuedBytes, 150);
|
||||
|
||||
// Second waiter would push the total over the 200-byte budget → must NOT park.
|
||||
const started = Date.now();
|
||||
const second = await controller.acquireHeavyWithin(2_000, undefined, 100);
|
||||
assert.equal(second, null, "over-budget wait must be rejected");
|
||||
assert.ok(Date.now() - started < 500, "rejection must be immediate, not park for queueMs");
|
||||
assert.equal(controller.queuedBytes, 150, "rejected waiter must not be charged");
|
||||
assert.equal(controller.activeHeavy, 1, "holder keeps its lease");
|
||||
|
||||
// Free the slot: the parked waiter acquires and its bytes leave the queue.
|
||||
held.release();
|
||||
const firstLease = await first;
|
||||
assert.ok(firstLease, "in-budget waiter acquires the freed slot");
|
||||
assert.equal(controller.queuedBytes, 0, "acquired waiter's bytes must leave the queue");
|
||||
firstLease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("aborting a parked wait releases its queued bytes", async () => {
|
||||
const controller = new ChatAdmissionController(1, 1_000);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const abortController = new AbortController();
|
||||
const pending = controller.acquireHeavyWithin(2_000, abortController.signal, 400);
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(controller.queuedBytes, 400);
|
||||
|
||||
abortController.abort();
|
||||
assert.equal(await pending, null);
|
||||
assert.equal(controller.queuedBytes, 0, "abort must release the charged bytes");
|
||||
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("a timed-out wait releases its queued bytes", async () => {
|
||||
const controller = new ChatAdmissionController(1, 1_000);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const pending = controller.acquireHeavyWithin(50, undefined, 400);
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(controller.queuedBytes, 400);
|
||||
|
||||
assert.equal(await pending, null);
|
||||
assert.equal(controller.queuedBytes, 0, "timeout must release the charged bytes");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("byte-heavy admission enforces the queued-bytes cap end-to-end", async () => {
|
||||
const controller = new ChatAdmissionController(1, 100);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] });
|
||||
const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 2_000 };
|
||||
|
||||
// First request parks: declared length (~70B) fits the budget.
|
||||
const first = admitChatRequest(chatRequest(body), options);
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
|
||||
// Second request would exceed the 100-byte budget → rejected immediately.
|
||||
const started = Date.now();
|
||||
const second = await admitChatRequest(chatRequest(body), options);
|
||||
assert.equal(second.admit, false, "over-budget byte-heavy wait must not admit");
|
||||
if (!second.admit) assert.equal(second.response.status, 503);
|
||||
assert.ok(Date.now() - started < 500, "over-budget wait must reject immediately");
|
||||
|
||||
held.release();
|
||||
const firstResult = await first;
|
||||
assert.equal(firstResult.admit, true);
|
||||
if (firstResult.admit) firstResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("structural admission enforces the queued-bytes cap end-to-end", async () => {
|
||||
const controller = new ChatAdmissionController(1, CHAT_LARGE_BODY_BYTES);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const structural = {
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
};
|
||||
const options = {
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 2_000,
|
||||
};
|
||||
|
||||
// First structural wait parks, charging the conservative 256KB weight.
|
||||
const first = admitChatStructure(structural, null, options);
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
|
||||
// Second would double the charge → rejected immediately.
|
||||
const started = Date.now();
|
||||
const second = await admitChatStructure(structural, null, options);
|
||||
assert.equal(second.admit, false, "over-budget structural wait must not admit");
|
||||
if (!second.admit) assert.equal(second.response.status, 503);
|
||||
assert.ok(Date.now() - started < 500, "over-budget structural wait must reject immediately");
|
||||
|
||||
held.release();
|
||||
const firstResult = await first;
|
||||
assert.equal(firstResult.admit, true);
|
||||
if (firstResult.admit) firstResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("queue-wait defaults are bounded (2s wait, 4MB queued-bytes budget)", () => {
|
||||
if (process.env.OMNIROUTE_CHAT_ADMISSION_QUEUE_MS === undefined) {
|
||||
assert.equal(CHAT_ADMISSION_QUEUE_MAX_MS, 2_000);
|
||||
}
|
||||
if (process.env.OMNIROUTE_CHAT_ADMISSION_MAX_QUEUED_BYTES === undefined) {
|
||||
assert.equal(CHAT_ADMISSION_MAX_QUEUED_BYTES, 4 * 1024 * 1024);
|
||||
}
|
||||
});
|
||||
|
||||
test("a pre-aborted signal never parks in the admission queue", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const abortController = new AbortController();
|
||||
abortController.abort("client already disconnected");
|
||||
|
||||
const pending = controller.acquireHeavyWithin(2_000, abortController.signal);
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
assert.equal(settled, true, "a pre-aborted signal must settle immediately, not park");
|
||||
|
||||
const lease = await pending;
|
||||
assert.equal(lease, null, "no lease is granted after abort");
|
||||
assert.equal(controller.activeHeavy, 1, "holder keeps capacity; aborted wait consumed nothing");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("aborting the request signal cancels a queued byte-heavy wait", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const abortController = new AbortController();
|
||||
const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] });
|
||||
const request = new Request("http://x/v1/chat/completions", {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body,
|
||||
signal: abortController.signal,
|
||||
});
|
||||
const pending = admitChatRequest(request, {
|
||||
controller,
|
||||
largeBodyBytes: 32,
|
||||
hardMaxBytes: 1024,
|
||||
queueMs: 2_000,
|
||||
});
|
||||
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(settled, false, "must queue while capacity is busy");
|
||||
|
||||
abortController.abort();
|
||||
const started = Date.now();
|
||||
const result = await pending;
|
||||
assert.ok(
|
||||
Date.now() - started < 500,
|
||||
"abort must cancel the queue-wait early, not park the full queueMs"
|
||||
);
|
||||
assert.equal(result.admit, false, "abort must not admit");
|
||||
if (!result.admit) {
|
||||
assert.equal(result.response.status, 503);
|
||||
assert.equal((await result.response.json()).error.code, "chat_admission_busy");
|
||||
}
|
||||
assert.equal(controller.activeHeavy, 1, "holder keeps capacity; aborted wait consumed nothing");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("aborting the signal cancels a structural queue-wait", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const abortController = new AbortController();
|
||||
const pending = admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 2_000,
|
||||
signal: abortController.signal,
|
||||
}
|
||||
);
|
||||
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(settled, false, "must queue while capacity is busy");
|
||||
|
||||
abortController.abort();
|
||||
const started = Date.now();
|
||||
const result = await pending;
|
||||
assert.ok(
|
||||
Date.now() - started < 500,
|
||||
"abort must cancel the queue-wait early, not park the full queueMs"
|
||||
);
|
||||
assert.equal(result.admit, false, "abort must not admit");
|
||||
if (!result.admit) {
|
||||
assert.equal(result.response.status, 503);
|
||||
assert.equal((await result.response.json()).error.code, "chat_admission_busy");
|
||||
}
|
||||
assert.equal(controller.activeHeavy, 1, "holder keeps its lease");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
@@ -8,6 +8,9 @@ const {
|
||||
admitChatStructure,
|
||||
ChatAdmissionController,
|
||||
CHAT_HARD_MAX_MESSAGES,
|
||||
CHAT_ADMISSION_QUEUE_MAX_MS,
|
||||
CHAT_ADMISSION_MAX_QUEUED_BYTES,
|
||||
CHAT_LARGE_BODY_BYTES,
|
||||
releaseChatAdmissionAfterHandler,
|
||||
releaseChatAdmissionWhenDone,
|
||||
resolveSelfLoopBearer,
|
||||
@@ -813,168 +816,3 @@ test("sk_omniroute sentinel is rejected once an env key is configured (REQUIRE_A
|
||||
restore();
|
||||
}
|
||||
});
|
||||
|
||||
test("a heavy structural request waits for capacity instead of failing immediately", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const pending = admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 500,
|
||||
}
|
||||
);
|
||||
|
||||
// Capacity is still busy: the request must not have resolved (admit/reject) yet.
|
||||
let settled = false;
|
||||
void pending.then(() => {
|
||||
settled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(settled, false, "must wait while capacity is busy");
|
||||
|
||||
held.release();
|
||||
const result = await pending;
|
||||
assert.equal(result.admit, true);
|
||||
if (result.admit) {
|
||||
assert.equal(controller.activeHeavy, 1, "waiting request acquires the freed lease");
|
||||
result.lease?.release();
|
||||
}
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("waiting for admission times out into a retryable 503", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const started = Date.now();
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 50,
|
||||
}
|
||||
);
|
||||
|
||||
assert.equal(result.admit, false);
|
||||
if (!result.admit) {
|
||||
assert.equal(result.response.status, 503);
|
||||
assert.equal(result.response.headers.get("retry-after"), "1");
|
||||
assert.equal((await result.response.json()).error.code, "chat_admission_busy");
|
||||
}
|
||||
assert.ok(Date.now() - started >= 40, "must wait for the queue deadline before rejecting");
|
||||
assert.equal(controller.activeHeavy, 1, "the holder keeps its lease");
|
||||
held.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("byte-heavy admission waits for capacity when queueMs is set", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] });
|
||||
const options = { controller, largeBodyBytes: 32, hardMaxBytes: 1024, queueMs: 500 };
|
||||
|
||||
const first = await admitChatRequest(chatRequest(body), options);
|
||||
assert.equal(first.admit, true);
|
||||
if (!first.admit) return;
|
||||
|
||||
const second = admitChatRequest(chatRequest(body), options);
|
||||
let secondSettled = false;
|
||||
void second.then(() => {
|
||||
secondSettled = true;
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, 30));
|
||||
assert.equal(
|
||||
secondSettled,
|
||||
false,
|
||||
"second heavy request must queue while the first holds capacity"
|
||||
);
|
||||
|
||||
first.lease?.release();
|
||||
const secondResult = await second;
|
||||
assert.equal(secondResult.admit, true, "second request acquires capacity after release");
|
||||
if (secondResult.admit) secondResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
test("expired admission queue keeps the legacy immediate 503 behaviour", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
},
|
||||
null,
|
||||
{
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 0,
|
||||
}
|
||||
);
|
||||
|
||||
assert.equal(result.admit, false);
|
||||
if (!result.admit) assert.equal(result.response.status, 503);
|
||||
held.release();
|
||||
});
|
||||
|
||||
test("admission waiters are served FIFO as capacity frees", async () => {
|
||||
const controller = new ChatAdmissionController(1);
|
||||
const held = controller.tryAcquireHeavy();
|
||||
assert.ok(held);
|
||||
|
||||
const body = {
|
||||
messages: [
|
||||
{ role: "user", content: "one" },
|
||||
{ role: "user", content: "two" },
|
||||
],
|
||||
};
|
||||
const options = {
|
||||
controller,
|
||||
maxMessages: 10,
|
||||
heavyMessages: 2,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
queueMs: 500,
|
||||
};
|
||||
const first = admitChatStructure(body, null, options);
|
||||
const second = admitChatStructure(body, null, options);
|
||||
|
||||
held.release();
|
||||
const firstResult = await first;
|
||||
assert.equal(firstResult.admit, true);
|
||||
if (firstResult.admit) firstResult.lease?.release();
|
||||
const secondResult = await second;
|
||||
assert.equal(secondResult.admit, true);
|
||||
if (secondResult.admit) secondResult.lease?.release();
|
||||
assert.equal(controller.activeHeavy, 0);
|
||||
});
|
||||
|
||||
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", async () => {
|
||||
// occupy sess-a's per-connection controller via the module-level instance
|
||||
const controller = perConnectionAdmissionController.getController("sess-a");
|
||||
const occupied = controller.tryAcquireHeavy();
|
||||
assert.ok(occupied);
|
||||
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: Array.from({ length: 3 }, () => ({ role: "user", content: "x" })),
|
||||
},
|
||||
null,
|
||||
{
|
||||
sessionId: "sess-a",
|
||||
maxMessages: 10,
|
||||
heavyMessages: 1,
|
||||
heavyTools: 10,
|
||||
heavyTokens: 10_000,
|
||||
}
|
||||
);
|
||||
// Session A is busy → 503
|
||||
assert.equal(result.admit, false);
|
||||
if (result.admit) return;
|
||||
assert.equal(result.response.status, 503);
|
||||
assert.equal(result.response.headers.get("Retry-After"), "1");
|
||||
occupied.release();
|
||||
});
|
||||
|
||||
test("admitChatStructure with different sessionId gets independent capacity", async () => {
|
||||
// occupy sess-a's per-connection controller
|
||||
const ctrlA = perConnectionAdmissionController.getController("sess-a");
|
||||
const occupied = ctrlA.tryAcquireHeavy();
|
||||
assert.ok(occupied);
|
||||
|
||||
// Session B should get its own controller → admitted
|
||||
const result = await admitChatStructure(
|
||||
{
|
||||
messages: Array.from({ length: 500 }, () => ({ role: "user", content: "x" })),
|
||||
},
|
||||
null,
|
||||
{
|
||||
sessionId: "sess-b",
|
||||
maxMessages: 0,
|
||||
heavyMessages: 200,
|
||||
heavyTools: 64,
|
||||
heavyTokens: 32_000,
|
||||
}
|
||||
);
|
||||
assert.equal(result.admit, true);
|
||||
if (result.admit) {
|
||||
assert.notEqual(result.lease, null);
|
||||
result.lease?.release();
|
||||
}
|
||||
occupied.release();
|
||||
});
|
||||
Reference in New Issue
Block a user