From ebbfcf4ffee647d3547a79c1b720c44788dfaed7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nguy=E1=BB=85n=20Vi=E1=BA=BFt=20Tu=E1=BA=A5n?= Date: Wed, 26 Aug 2026 18:11:18 +0700 Subject: [PATCH] fix(sse): scale chat admission by ingest byte budget instead of a fixed request count (#11548) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merged via /merge-batch (lote 2026-08-26, v3.8.51). Boarded no worktree combinado junto com outras ~30 PRs; validação única: typecheck/complexity/cognitive-complexity/changelog-integrity verdes, file-size rebaseado onde necessário (crescimento legítimo), lint com os mesmos 228 achados pré-existentes confirmados via sonda contra o tip puro (não introduzidos por este lote), e ~370 testes focados (unit + vitest) passando. Obrigado pela contribuição. --- .env.example | 11 +- docs/architecture/admission-lanes.md | 47 +- docs/guides/DOCKER_GUIDE.md | 2 +- docs/guides/TROUBLESHOOTING.md | 43 +- docs/ops/VM_DEPLOYMENT_GUIDE.md | 2 +- docs/reference/ENVIRONMENT.md | 3 +- src/lib/monitoring/observability.ts | 16 + src/shared/middleware/admissionBudget.ts | 116 +++++ .../middleware/chatAdmissionIdentity.ts | 40 ++ .../middleware/chatAdmissionResponses.ts | 73 +++ .../chatAdmissionStructureEstimate.ts | 54 ++ src/shared/middleware/chatBodyAdmission.ts | 481 +++++++++--------- src/shared/middleware/ingestByteAdmission.ts | 184 +++++++ tests/unit/admission-ingest-budget.test.ts | 123 +++++ tests/unit/admission-pressure-gate.test.ts | 214 ++++++++ .../agent-fanout-admission-regression.test.ts | 182 +++++++ tests/unit/observability-payloads.test.ts | 15 +- .../per-connection-admission-9654.test.ts | 13 +- 18 files changed, 1340 insertions(+), 279 deletions(-) create mode 100644 src/shared/middleware/admissionBudget.ts create mode 100644 src/shared/middleware/chatAdmissionIdentity.ts create mode 100644 src/shared/middleware/chatAdmissionResponses.ts create mode 100644 src/shared/middleware/chatAdmissionStructureEstimate.ts create mode 100644 src/shared/middleware/ingestByteAdmission.ts create mode 100644 tests/unit/admission-ingest-budget.test.ts create mode 100644 tests/unit/admission-pressure-gate.test.ts create mode 100644 tests/unit/agent-fanout-admission-regression.test.ts diff --git a/.env.example b/.env.example index 4b10004a8d..86ea45e2e4 100644 --- a/.env.example +++ b/.env.example @@ -403,8 +403,17 @@ ALLOW_API_KEY_REVEAL=false # OMNIROUTE_CHAT_LARGE_BODY_BYTES=262144 # Actual-byte hard cap enforced during bounded ingestion. Default 52428800 (50 MB). # OMNIROUTE_CHAT_HARD_MAX_BODY_BYTES=52428800 -# Maximum heavyweight requests simultaneously admitted in one process. Default 1. +# Legacy request-COUNT cap (#503-fanout). Now binds only when explicitly set here — +# left unset, heavyweight admission is gated by OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES below +# instead (an auto-derived byte budget), fixing coding-agent fan-out (multiple +# subagents/CLIs) collapsing to an effective concurrency of ~1 and 503ing. # OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT=1 +# Override for the auto-derived ingest byte budget (#503-fanout). Default: 25% of the +# process's effective memory ceiling (V8 heap limit, or the tighter cgroup/container +# limit) divided by an 8x transient-amplification factor, clamped between 8 MiB and +# 2 GiB; explicit overrides are clamped to the same safe range. Read +# chatAdmission.maxInflightBytes/budgetSource at /api/monitoring/health before overriding. +# OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES=134217728 # Heap-pressure shed ratio (heapUsed/heap_size_limit) for the structural admission gate # (#10183, #10268): a second concurrent heavyweight request past OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT # is only shed with a retryable 503 when the heap is ALSO under this much pressure — on a diff --git a/docs/architecture/admission-lanes.md b/docs/architecture/admission-lanes.md index 5a7aa5b0f3..2b2a4cf96b 100644 --- a/docs/architecture/admission-lanes.md +++ b/docs/architecture/admission-lanes.md @@ -9,20 +9,49 @@ lastUpdated: 2026-08-10 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`) +## 1. Byte-level process-wide admission (`chatBodyAdmission.ts`) -- **Scope:** the buffered-body/heap path for `POST /v1/chat/completions`. Guards +- **Scope:** the buffered-body/heap path for `POST /v1/chat/completions`, + `/v1/messages`, `/v1/responses`, and the other chat-shaped routes. 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. +- **One process-global controller, not per-key lanes (#10110).** Every API key + (hashed) or `anonymous` session admits against the **same** shared budget — + the hashed session id is used ONLY as a fairness scheduling key (round-robin + dispatch across waiters), never as a capacity shard. A prior version of this + doc described per-key lanes with independent capacity; that model was + removed in #10110 because it let unauthenticated fake credentials multiply + the process-wide bound. +- **Gate (#503-fanout): an auto-derived ingest BYTE budget, not a fixed request + count.** The legacy `CHAT_MAX_HEAVY_IN_FLIGHT` request-count cap (default `1` + before this fix) collapsed coding-agent fan-out (multiple subagents/CLIs, + bodies routinely > 256 KB) to an effective concurrency of ~1, which 503'd + under completely normal load. It now binds only when an operator explicitly + sets `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT`. Left unset, admission is instead + gated by `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES` — a budget auto-derived from the + process's real memory ceiling (`src/shared/middleware/admissionBudget.ts`): + 25% of the tighter of the V8 heap limit and any cgroup/container limit, + divided by an 8x transient-amplification factor, clamped between 8 MiB and + 2 GiB. Explicit overrides use the same clamps. This scales itself from a + 512 MB container to a 32 GB desktop with no env tuning. A body that cannot + fit within the effective budget fails immediately with `413 body_exceeds_budget`; + only contention among individually serviceable bodies enters the bounded + fairness queue. A live multi-signal resource-pressure tracker (V8 heap ratio, + cgroup, PSI, OOM events — `open-sse/utils/resourcePressurePolicy.ts`) shortens + the bounded wait under `high` pressure and sheds immediately with + `503 resource_pressure` under `critical` pressure, before any bytes are even + ingested. - **Tuning:** - - `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` — idle-lane eviction (default 60000) - - `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` — lane count cap (default 64) + - `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES` — override for the auto-derived byte budget + - `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` — legacy request-count cap, opt-in only - `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). + - `OMNIROUTE_CHAT_VIRTUAL_TTL_MS` / `OMNIROUTE_CHAT_VIRTUAL_MAX_SESSIONS` — deprecated + no-ops since #10110 (accepted for config compatibility, ignored) +- **Reports:** `GET /api/monitoring/health` → `chatAdmission` (#11244) — including + the #503-fanout additions `inflightBytes`, `maxInflightBytes`, `budgetSource` + (`v8_heap` | `cgroup` | `override`), `pressureSeverity`, and `countCapEnabled` + (false on a default deployment — confirms the byte budget, not the legacy + count cap, is what is actually binding). ## 2. Adaptive runtime virtual lanes (`open-sse/services/admission`) diff --git a/docs/guides/DOCKER_GUIDE.md b/docs/guides/DOCKER_GUIDE.md index 5c650df5da..295d8aa800 100644 --- a/docs/guides/DOCKER_GUIDE.md +++ b/docs/guides/DOCKER_GUIDE.md @@ -561,7 +561,7 @@ External Postgres / multi-writer HA is **not** a documented stock path. If you n ## Scale-out: N independent processes -One Node process is **one V8 heap**. Two overlapping ~3 MiB / ~750k-token coding-agent `POST /v1/responses` (RTK + Caveman) abort that heap at ~12 Gi (`FATAL ERROR: Reached heap limit`) and can OOM a 16 Gi cgroup. See [#7849](https://github.com/diegosouzapw/OmniRoute/issues/7849). Raising `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` on that process reintroduces the abort. Small chats, `/healthz`, `/v1/models`, and MCP are **not** in that cap. +One Node process is **one V8 heap**. Two overlapping ~3 MiB / ~750k-token coding-agent `POST /v1/responses` (RTK + Caveman) abort that heap at ~12 Gi (`FATAL ERROR: Reached heap limit`) and can OOM a 16 Gi cgroup. See [#7849](https://github.com/diegosouzapw/OmniRoute/issues/7849). Heavyweight chat admission is gated by an auto-derived ingest byte budget (`OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES`, `src/shared/middleware/admissionBudget.ts`) sized from that same V8/cgroup ceiling -- it already scales itself to the process's real memory, so overriding it upward (or setting the legacy `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` request-count cap) on an already-sized process reintroduces the abort. Small chats, `/healthz`, `/v1/models`, and MCP are **not** in that cap. To go beyond two concurrent **large** jobs **today**: diff --git a/docs/guides/TROUBLESHOOTING.md b/docs/guides/TROUBLESHOOTING.md index 75c8e4b601..e81ef8a95c 100644 --- a/docs/guides/TROUBLESHOOTING.md +++ b/docs/guides/TROUBLESHOOTING.md @@ -58,11 +58,11 @@ export OMNIROUTE_CHAT_ADMISSION_QUEUE_MS=5000 # longer bounded wait for heavywei Set these in the OmniRoute process environment (the daemon, e.g. via the LaunchAgent plist or `systemctl edit`), then restart OmniRoute. The rotation flag is the single highest-leverage lever: it converts a hard failure into a transparent retry against a healthy provider in the pool. -**Note**: `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` (default `1`, per process) caps how many heavyweight — long-context — requests run at once; the bound is an admission gate, not a provider rate limiter. Raising it only reduces client-visible `503 chat_admission_busy` rejects for heavy requests. The per-provider rate limiting (`open-sse/services/rateLimitManager.ts`) is governed separately by `RATE_LIMIT_MAX_WAIT_MS`, `RATE_LIMIT_MAX_QUEUE_DEPTH`, and `RATE_LIMIT_AUTO_ENABLE` — see `.env.example`. +**Note**: `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` caps how many heavyweight — long-context — requests run at once; the bound is an admission gate, not a provider rate limiter. **#503-fanout update:** this var is no longer set by default (it now binds only when explicitly configured, as above) — heavyweight admission is instead gated by an auto-derived byte budget (`OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES`) that scales itself from the host's real memory ceiling, so a fresh deployment should see far fewer `503 chat_admission_busy` rejects without setting this var at all; explicitly setting it here still works exactly as documented. Explicit byte-budget overrides clamp to 8 MiB–2 GiB. A `413 body_exceeds_budget` is not transient: increase that byte budget, lower `OMNIROUTE_CHAT_HARD_MAX_BODY_BYTES`, or increase the process memory ceiling. An `inflight_bytes_budget` shed is temporary contention and remains retryable. The per-provider rate limiting (`open-sse/services/rateLimitManager.ts`) is governed separately by `RATE_LIMIT_MAX_WAIT_MS`, `RATE_LIMIT_MAX_QUEUE_DEPTH`, and `RATE_LIMIT_AUTO_ENABLE` — see `.env.example`. **How to verify it worked**: run your agent/cron twice in quick succession and confirm both succeed. Before the fix, the second run typically throws `429`/`401`. After the fix, failures (if any) are retried transparently and the call completes. You can also `curl /monitoring/health` and watch the `rateLimitedUntil` field on the provider connections and the `circuitBreakers.providerBreakers[].state` for the affected providers — the state is one of `CLOSED`, `DEGRADED`, `OPEN`, or `HALF_OPEN` (see `src/shared/utils/circuitBreaker.ts`), and a provider that keeps failing will flip `CLOSED → DEGRADED → OPEN` before the reset window lets a probe through (`HALF_OPEN`). -**If you still see 429**: the active account for that provider has genuinely exhausted its *quota* (not just rate). Add a second account for the same provider in the OmniRoute dashboard → Providers → Accounts, or mix in another free provider (e.g. `routeway`, `auggie`). Rotation only helps with transient rate/400/401; a hard quota exhaustion requires a second credential or a different provider. +**If you still see 429**: the active account for that provider has genuinely exhausted its _quota_ (not just rate). Add a second account for the same provider in the OmniRoute dashboard → Providers → Accounts, or mix in another free provider (e.g. `routeway`, `auggie`). Rotation only helps with transient rate/400/401; a hard quota exhaustion requires a second credential or a different provider. **If you see 403 on vision models (`auto/vision`, `bazaarlink/*`)**: the connected account lacks a paid plan that includes vision, or the API key has insufficient permissions. Verify in the provider dashboard that the key scope includes vision/multimodal, or connect a paid tier account and keep it as the vision target. @@ -567,34 +567,41 @@ Each process uses a process-local guard to reserve limited heavyweight capacity and parsing a large request body. A heavyweight lease remains held for the lifetime of an SSE response. +**#503-fanout:** before this fix, the guard capped concurrency at a fixed request COUNT +(`OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT`, default `1`) regardless of host memory, so coding-agent +fan-out (multiple subagents/CLIs, bodies routinely > 256 KB) collapsed to an effective +concurrency of ~1 and 503'd under completely normal load. The guard now self-tunes: it is gated +by an auto-derived ingest BYTE budget (`OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES`) sized from the +process's real memory ceiling, and it also consults a live resource-pressure signal — so it +only sheds when the host is genuinely under memory pressure, not merely because more than one +heavy request arrived at once. The old count cap (`OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT`) is +still honored, but only if you explicitly set it. + When capacity is busy, a heavyweight request first waits up to -`OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` (default `5000`, `0` disables the wait) for a slot to free up +`OMNIROUTE_CHAT_ADMISSION_QUEUE_MS` (default `2000`, `0` disables the wait) for a slot to free up before answering the retryable `503`. The bounded wait exists so agent-style clients (OpenCode, Claude Code, Cursor) that fan out heavy sub-requests concurrently serialize the burst instead of burning their whole retry budget on immediate rejections and dying mid-task. -Current heavyweight lease occupancy is not surfaced in the dashboard. +Current heavyweight lease occupancy, the resolved byte budget, and live pressure severity are +surfaced at `GET /api/monitoring/health` → `chatAdmission` (`inflightBytes`, `maxInflightBytes`, +`budgetSource`, `pressureSeverity`, `countCapEnabled`) — check these before touching any env var. Settings → Resilience → Request Queue → Concurrent Requests does not control this; that setting governs a separate provider request-queue mechanism. **Fix:** 1. Retry first. Clients should honor `Retry-After` and use backoff rather than immediately - repeating the request. Note that with the default `OMNIROUTE_CHAT_ADMISSION_QUEUE_MS=5000` - a heavy request already waited up to 5 seconds before the `503`, so a client retry loop should - back off beyond that instead of hammering. -2. If normal deployment traffic repeatedly exhausts the guard, you can cautiously raise - `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` from its default of `1`. Increase it one step at a time, - restart OmniRoute after each change, and observe memory headroom under representative load. - Every additional heavyweight request can increase concurrent V8 heap use and container or - host OOM risk. No value is safe for every deployment; validate the setting against your own - traffic and memory limits rather than assuming that `2` is universally safe. -3. Prefer widening the wait (`OMNIROUTE_CHAT_ADMISSION_QUEUE_MS`) over raising the in-flight - limit when bursts are short: waiting costs latency, while an extra concurrent heavyweight - request costs heap residency for the whole request lifetime. + repeating the request. +2. Check `/api/monitoring/health` → `chatAdmission` before tuning anything. `countCapEnabled: +false` and a generous `maxInflightBytes` mean the auto-derived budget is already doing its + job; a `pressureSeverity` of `high`/`critical` means the host is genuinely low on memory — + that is not fixable by an admission env var, it needs more RAM or a smaller workload. +3. Only if `/api/monitoring/health` shows the auto-derived budget is genuinely too small for + your host (rare — it already scales from container to bare-metal), override it directly with + `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES` rather than falling back to the legacy request-count cap. See the [environment-variable reference](../reference/ENVIRONMENT.md#4-security--authentication) -for the authoritative admission settings. Loosening the heavyweight classification thresholds -can let expensive requests bypass this guard and is riskier than a cautious in-flight increase. +for the authoritative admission settings. --- diff --git a/docs/ops/VM_DEPLOYMENT_GUIDE.md b/docs/ops/VM_DEPLOYMENT_GUIDE.md index c0e64c74e7..41a77d9e98 100644 --- a/docs/ops/VM_DEPLOYMENT_GUIDE.md +++ b/docs/ops/VM_DEPLOYMENT_GUIDE.md @@ -430,6 +430,6 @@ For deployments on small VPS instances (1 GB RAM or less): - **Disable background services** — set `OMNIROUTE_DISABLE_BACKGROUND_SERVICES=1` to skip scheduler, MCP server, and periodic maintenance tasks. See `docs/reference/ENVIRONMENT.md`. - **Use SQLite WAL mode** — enabled by default, reduces peak memory during concurrent reads. - **Cap the V8 heap** — set `OMNIROUTE_MEMORY_MB` (e.g. `512`) so the runtime does not calibrate a ceiling larger than the VM. See `docs/reference/ENVIRONMENT.md`. -- **Limit concurrent heavy requests** — lower `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` (default `1`); excess requests get a retryable `503` with `Retry-After` instead of competing for memory. +- **Heavyweight admission auto-scales with the heap cap** -- once `OMNIROUTE_MEMORY_MB` is set above, the ingest byte budget (`OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES`) derives itself from that same ceiling, so a memory-constrained VM already gets a smaller concurrent-request budget with no extra tuning; excess requests get a retryable `503` with `Retry-After` instead of competing for memory. Set the legacy `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` request-count cap only if you need a hard ceiling on top of that. - **Avoid `next build` on the VPS** — build locally and deploy the standalone output (`.next/standalone/`). - **Monitor with `top` / `free -m`** — OmniRoute typically uses 200-400 MB RSS at idle on a 1 GB VM. diff --git a/docs/reference/ENVIRONMENT.md b/docs/reference/ENVIRONMENT.md index e3459a6117..b782e53607 100644 --- a/docs/reference/ENVIRONMENT.md +++ b/docs/reference/ENVIRONMENT.md @@ -203,7 +203,8 @@ OmniRoute uses **SQLite** (via `better-sqlite3`) for all persistence. These vari | `MAX_BODY_SIZE_BYTES` | `10485760` (10 MB) | `src/shared/middleware/bodySizeGuard.ts` | Maximum allowed request body size. Rejects payloads exceeding this limit. | | `OMNIROUTE_CHAT_LARGE_BODY_BYTES` | `262144` (256 KB) | `src/shared/middleware/chatBodyAdmission.ts` | Actual request bodies at or above this threshold require an atomic process-local heavyweight admission lease before JSON parsing. | | `OMNIROUTE_CHAT_HARD_MAX_BODY_BYTES` | `52428800` (50 MB) | `src/shared/middleware/chatBodyAdmission.ts` | Chat-route hard cap enforced against bytes read during bounded ingestion, including requests with missing, invalid, or dishonest `Content-Length`; excess receives `413`. | -| `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` | `1` | `src/shared/middleware/chatBodyAdmission.ts` | Maximum heavyweight chat requests admitted concurrently in **one process** (one V8 heap). Overload is retryable `503` with `Retry-After`. Two overlapping ~750k-token `/v1/responses` already abort ~12 Gi heaps (#7849); do not raise this to “use the host.” Multiply capacity with **N independent `DATA_DIR`s** (#11024), not `replicas>1` on one SQLite file. | +| `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` | _(unset — no request-count cap)_ | `src/shared/middleware/chatBodyAdmission.ts` | **#503-fanout:** this legacy request-COUNT cap now binds only when explicitly set. Left unset (the default), heavyweight chat admission is instead gated by `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES` — an auto-derived BYTE budget sized from the process's real memory ceiling in **one process** (one V8 heap), fixing a bug where coding-agent fan-out (multiple subagents/CLIs, bodies routinely > 256 KB) collapsed to an effective concurrency of ~1 and 503'd under normal load. Setting this var restores the old fixed-count behavior on top of the byte budget for a deployment that already tuned it. Overload is retryable `503` with `Retry-After`. Two overlapping ~750k-token `/v1/responses` already abort ~12 Gi heaps (#7849) — the byte budget accounts for that ceiling automatically, so raising this manually is no longer the recommended lever. Multiply capacity with **N independent `DATA_DIR`s** (#11024), not `replicas>1` on one SQLite file. | +| `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES` | _(auto-derived)_ | `src/shared/middleware/admissionBudget.ts` | **#503-fanout:** override for the auto-derived ingest byte budget (25% of the tighter V8/cgroup memory ceiling divided by 8x transient amplification). Derived and explicit values clamp to 8 MiB–2 GiB. A body larger than the effective budget fails immediately with `413 body_exceeds_budget`; contention between individually serviceable bodies remains retryable `503`. Read `chatAdmission.maxInflightBytes` / `budgetSource` / `pressureSeverity` at `/api/monitoring/health` before tuning. | | `OMNIROUTE_CHAT_ADMISSION_HEAP_SHED_RATIO` | `0.75` | `src/shared/middleware/chatBodyAdmission.ts` | Heap-pressure shed ratio (`heapUsed / heap_size_limit`) for the structural admission gate (#10183, #10268). A second concurrent heavyweight request past `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` is only shed with the retryable `503` when the heap is ALSO at or above this ratio; on a healthy heap it is admitted instead. | | `OMNIROUTE_CHAT_ADMISSION_HEALTHY_HEADROOM` | `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT` (default `1`) | `src/shared/middleware/chatBodyAdmission.ts` | Bounded extra capacity for the healthy-heap fast path above (#10437). Without this bound, every busy-but-healthy-heap request bypassed admission with no ceiling at all — a slow leak or a burst that never quite trips the heap-shed ratio could still pile up unlimited concurrent heavyweight work. Once this many concurrent leases are active through the healthy-heap path, further busy requests fall through to the SAME bounded-wait/shed path used under real heap pressure. `0` disables the bypass entirely. | | `OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT` | `200` | `src/shared/middleware/chatBodyAdmission.ts` | Message count that classifies a chat request as heavyweight even when its body is below the byte threshold. | diff --git a/src/lib/monitoring/observability.ts b/src/lib/monitoring/observability.ts index cb4d239a65..747afa0b91 100644 --- a/src/lib/monitoring/observability.ts +++ b/src/lib/monitoring/observability.ts @@ -24,6 +24,17 @@ export type ChatAdmissionHealthSummary = { shedTotal: number; shedsByReason: Record; lanes: Array<{ key: string; waiting: number }>; + /** #503-fanout: live ingest bytes reserved through the byte-budget gate. */ + inflightBytes: number; + /** #503-fanout: the auto-derived (or overridden) budget ceiling. */ + maxInflightBytes: number; + /** #503-fanout: which signal the budget was derived from. */ + budgetSource: string; + /** #503-fanout: live multi-signal resource-pressure severity. */ + pressureSeverity: string; + /** #503-fanout: false on a default deployment — proves the byte-budget + * gate, not the legacy request-count cap, is what is actually binding. */ + countCapEnabled: boolean; }; /** @@ -42,6 +53,11 @@ export function projectChatAdmissionSummary( shedTotal: snapshot.shedTotal, shedsByReason: { ...(snapshot.shedsByReason ?? {}) }, lanes: (snapshot.lanes ?? []).map((lane) => ({ key: lane.key, waiting: lane.waiting })), + inflightBytes: snapshot.inflightBytes, + maxInflightBytes: snapshot.maxInflightBytes, + budgetSource: snapshot.budgetSource, + pressureSeverity: snapshot.pressureSeverity, + countCapEnabled: snapshot.countCapEnabled, }; } diff --git a/src/shared/middleware/admissionBudget.ts b/src/shared/middleware/admissionBudget.ts new file mode 100644 index 0000000000..bb6e7a2a27 --- /dev/null +++ b/src/shared/middleware/admissionBudget.ts @@ -0,0 +1,116 @@ +/** + * Auto-derived ingest byte budget for chat admission (#503-fanout). + * + * The legacy `chatBodyAdmission.ts` gate counted *requests* (default cap: 1) + * instead of *bytes*, so any deployment serving coding-agent traffic (bodies + * routinely > 256 KB) collapsed to an effective concurrency of 1-2 regardless + * of how much RAM the host actually has. This module derives a byte budget + * from the process's real memory ceiling — mirroring the auto-calibration + * pattern already proven by `open-sse/utils/heapPressure.ts::computeHeapPressureThresholdMb` + * (the v3.8.8 "resource pressure" outage was a fixed-number version of this + * same mistake) — so the gate scales itself on a 512 MB container and a + * 32 GB desktop alike, with no env tuning required. + */ +import v8 from "node:v8"; + +/** Fraction of the effective memory ceiling reserved for raw, not-yet-dispatched request bytes. */ +export const INGEST_HEAP_FRACTION = 0.25; +/** + * Transient heap multiplier per raw ingest byte during buffer → parse → + * translate → dispatch (UTF-8 buffer + JS string + parsed object graph + + * translated graph coexist briefly). Folded into the budget itself so a + * request is charged its exact raw byte count, never a guessed multiple. + */ +export const INGEST_AMPLIFICATION = 8; +/** + * Preserve liveness for ordinary agent requests on tiny hosts, accepting that + * the floor can reserve a large share of a sub-128 MiB container. + */ +export const MIN_INGEST_BUDGET_BYTES = 8 * 1024 * 1024; +export const MAX_INGEST_BUDGET_BYTES = 2 * 1024 * 1024 * 1024; + +export type IngestBudgetSource = "override" | "cgroup" | "v8_heap"; + +export interface IngestBudget { + bytes: number; + source: IngestBudgetSource; + effectiveCeilingBytes: number; +} + +function parsePositiveFinite(value: string | number | null | undefined): number | null { + if (value === null || value === undefined || value === "") return null; + const parsed = Number(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : null; +} + +/** + * Pure budget calculation — no I/O, no env reads. Callers supply the live + * V8 heap ceiling and (optionally) a cgroup-derived constrained-memory + * figure; the tighter of the two wins. `override` (when positive) always wins. + */ +export function computeIngestByteBudget(input: { + heapSizeLimitBytes: number; + constrainedMemoryBytes?: number | null; + override?: string | number | null; +}): IngestBudget { + const override = parsePositiveFinite(input.override); + if (override !== null) { + const bytes = Math.min( + MAX_INGEST_BUDGET_BYTES, + Math.max(MIN_INGEST_BUDGET_BYTES, Math.floor(override)) + ); + return { bytes, source: "override", effectiveCeilingBytes: bytes }; + } + + const heapLimit = parsePositiveFinite(input.heapSizeLimitBytes) ?? MIN_INGEST_BUDGET_BYTES; + const constrained = parsePositiveFinite(input.constrainedMemoryBytes ?? null); + const ceiling = constrained !== null ? Math.min(heapLimit, constrained) : heapLimit; + const source: IngestBudgetSource = + constrained !== null && constrained <= heapLimit ? "cgroup" : "v8_heap"; + + const raw = Math.floor((ceiling * INGEST_HEAP_FRACTION) / INGEST_AMPLIFICATION); + const bytes = Math.min(MAX_INGEST_BUDGET_BYTES, Math.max(MIN_INGEST_BUDGET_BYTES, raw)); + + return { bytes, source, effectiveCeilingBytes: Math.floor(ceiling) }; +} + +/** + * Best-effort cgroup/container memory ceiling. `process.constrainedMemory()` + * (Node >=19.6/20.13) returns the cgroup limit on Linux containers and is + * absent/undefined elsewhere (e.g. plain Windows/macOS hosts) — treated the + * same as "no cgroup limit" so the V8 heap ceiling is used instead. + */ +function readConstrainedMemoryBytes(): number | null { + try { + const proc = process as NodeJS.Process & { constrainedMemory?: () => number }; + if (typeof proc.constrainedMemory !== "function") return null; + const value = proc.constrainedMemory(); + return Number.isFinite(value) && value > 0 ? value : null; + } catch { + return null; + } +} + +let cachedBudget: IngestBudget | null = null; + +/** + * Resolve the process-wide ingest byte budget, cached after first call (like + * `HEAP_PRESSURE_THRESHOLD_MB`) so the gate never re-reads cgroup/V8 state on + * the hot path. `override` defaults to `OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES`. + */ +export function resolveIngestByteBudget( + override: string | number | null | undefined = process.env.OMNIROUTE_CHAT_MAX_INFLIGHT_BYTES +): IngestBudget { + if (cachedBudget) return cachedBudget; + cachedBudget = computeIngestByteBudget({ + heapSizeLimitBytes: v8.getHeapStatistics().heap_size_limit, + constrainedMemoryBytes: readConstrainedMemoryBytes(), + override, + }); + return cachedBudget; +} + +/** Test seam: force the next `resolveIngestByteBudget()` call to recompute. */ +export function reloadIngestBudgetForTests(): void { + cachedBudget = null; +} diff --git a/src/shared/middleware/chatAdmissionIdentity.ts b/src/shared/middleware/chatAdmissionIdentity.ts new file mode 100644 index 0000000000..0be286af58 --- /dev/null +++ b/src/shared/middleware/chatAdmissionIdentity.ts @@ -0,0 +1,40 @@ +import { createHmac } from "crypto"; + +const ADMISSION_BYPASS_VALUE = "internal"; +const SELF_LOOP_KEY = "sk_omniroute"; +const FINGERPRINT_KEY = "omniroute-admission-fingerprint-v1"; + +export const ADMISSION_BYPASS_HEADER = "x-omniroute-admission-bypass"; + +export function resolveSessionId(request: Request): string { + const authHeader = request.headers.get("authorization") || ""; + const bearerMatch = /^bearer\s+(\S+)$/i.exec(authHeader.trim()); + if (bearerMatch) return fingerprint(bearerMatch[1]); + + const xApiKey = request.headers.get("x-api-key")?.trim(); + if (xApiKey) return fingerprint(xApiKey); + + const xGoogApiKey = request.headers.get("x-goog-api-key")?.trim(); + return xGoogApiKey ? fingerprint(xGoogApiKey) : "anonymous"; +} + +export function resolveSelfLoopBearer(): string { + return ( + process.env.OMNIROUTE_API_KEY?.trim() || process.env.ROUTER_API_KEY?.trim() || SELF_LOOP_KEY + ); +} + +export function isInternalAdmissionBypass(request: Request): boolean { + const bypass = + request.headers.get(ADMISSION_BYPASS_HEADER)?.trim().toLowerCase() === ADMISSION_BYPASS_VALUE; + if (!bypass) return false; + + const auth = request.headers.get("authorization") || ""; + const match = /^bearer\s+(\S+)$/i.exec(auth.trim()); + return Boolean(match && match[1].trim().toLowerCase() === resolveSelfLoopBearer().toLowerCase()); +} + +function fingerprint(value: string): string { + // Deterministic admission-lane fingerprint, never password verification. + return `key_${createHmac("sha256", FINGERPRINT_KEY).update(value).digest("hex").slice(0, 16)}`; +} diff --git a/src/shared/middleware/chatAdmissionResponses.ts b/src/shared/middleware/chatAdmissionResponses.ts new file mode 100644 index 0000000000..8fca594b44 --- /dev/null +++ b/src/shared/middleware/chatAdmissionResponses.ts @@ -0,0 +1,73 @@ +import { buildErrorBody } from "@omniroute/open-sse/utils/error.ts"; + +import { CORS_HEADERS } from "../utils/cors"; + +const JSON_HEADERS = { ...CORS_HEADERS, "Content-Type": "application/json" }; + +export function chatAdmissionRejectionResponse(status: 413 | 503, hardMaxBytes: number): Response { + const isPayload = status === 413; + const headers: Record = { ...JSON_HEADERS }; + if (!isPayload) headers["Retry-After"] = "2"; + const message = isPayload + ? `Request body too large for chat completions (max ${Math.floor( + hardMaxBytes / (1024 * 1024) + )} MB).` + : "Chat admission capacity is temporarily unavailable. Retry shortly."; + return new Response( + JSON.stringify( + buildErrorBody(status, message, undefined, { + type: isPayload ? "payload_too_large" : "server_error", + code: isPayload ? "PAYLOAD_TOO_LARGE" : "chat_admission_busy", + }) + ), + { status, headers } + ); +} + +export function bodyExceedsBudgetResponse(maxInflightBytes: number): Response { + const maxMiB = Math.max(1, Math.floor(maxInflightBytes / (1024 * 1024))); + return new Response( + JSON.stringify( + buildErrorBody( + 413, + `Request body exceeds the chat ingest budget (max ${maxMiB} MB).`, + undefined, + { type: "payload_too_large", code: "body_exceeds_budget" } + ) + ), + { status: 413, headers: JSON_HEADERS } + ); +} + +export function resourcePressureRejectionResponse(): Response { + return new Response( + JSON.stringify( + buildErrorBody( + 503, + "Service temporarily unavailable due to resource pressure. Retry shortly.", + undefined, + { type: "server_error", code: "resource_pressure" } + ) + ), + { status: 503, headers: { ...JSON_HEADERS, "Retry-After": "2" } } + ); +} + +export function structuralRejectionResponse(status: 413 | 503, maxMessages: number): Response { + const historyLimit = status === 413; + const headers: Record = { ...JSON_HEADERS }; + if (!historyLimit) headers["Retry-After"] = "1"; + const body = buildErrorBody( + status, + historyLimit + ? `Chat history exceeds the ${maxMessages}-message limit; compact the conversation and retry.` + : "Structurally heavy chat request capacity is busy; retry shortly.", + undefined, + { + type: historyLimit ? "payload_too_large" : "server_error", + code: historyLimit ? "chat_history_too_large" : "chat_admission_busy", + } + ); + body.error.reason = historyLimit ? "message_limit" : "structure_limit"; + return new Response(JSON.stringify(body), { status, headers }); +} diff --git a/src/shared/middleware/chatAdmissionStructureEstimate.ts b/src/shared/middleware/chatAdmissionStructureEstimate.ts new file mode 100644 index 0000000000..602c03378b --- /dev/null +++ b/src/shared/middleware/chatAdmissionStructureEstimate.ts @@ -0,0 +1,54 @@ +export interface TokenEstimate { + tokens: number; + exhausted: boolean; +} + +export function estimateStructureTokens(value: unknown, limit: number): TokenEstimate { + let tokens = 0; + let visited = 0; + const maxNodes = 10_000; + const stack: Array<{ value: unknown; depth: number }> = [{ value, depth: 0 }]; + while (stack.length > 0 && tokens < limit && visited < maxNodes) { + const current = stack.pop(); + if (!current) break; + visited += 1; + if (typeof current.value === "string") { + tokens += conservativeStringTokens(current.value, limit - tokens); + continue; + } + if (!current.value || typeof current.value !== "object") continue; + if (current.depth >= 12) return { tokens, exhausted: true }; + + const remainingNodes = maxNodes - visited - stack.length; + if (Array.isArray(current.value)) { + if (current.value.length > remainingNodes) return { tokens, exhausted: true }; + for (let index = current.value.length - 1; index >= 0; index -= 1) { + stack.push({ value: current.value[index], depth: current.depth + 1 }); + } + continue; + } + + let children = 0; + for (const key in current.value) { + if (!Object.hasOwn(current.value, key)) continue; + children += 1; + if (children > remainingNodes) return { tokens, exhausted: true }; + tokens += conservativeStringTokens(key, limit - tokens); + if (tokens >= limit) return { tokens: limit, exhausted: false }; + stack.push({ + value: (current.value as Record)[key], + depth: current.depth + 1, + }); + } + } + return { tokens, exhausted: stack.length > 0 && tokens < limit }; +} + +function conservativeStringTokens(value: string, remaining: number): number { + let tokens = 0; + for (const character of value) { + tokens += character.codePointAt(0)! < 0x80 ? 0.25 : 1; + if (tokens >= remaining) return remaining; + } + return tokens; +} diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 3e37f3ca3b..0fc550ba6b 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -15,11 +15,32 @@ * connection's burst cannot starve others (#9654). */ -import { CORS_HEADERS } from "../utils/cors"; import { createLogger } from "../utils/logger"; -import { createHmac } from "crypto"; import v8 from "node:v8"; import { trackRequest } from "../../lib/gracefulShutdown"; +import { resolveIngestByteBudget, type IngestBudgetSource } from "./admissionBudget"; +import { + ADMISSION_BYPASS_HEADER, + isInternalAdmissionBypass, + resolveSelfLoopBearer, + resolveSessionId, +} from "./chatAdmissionIdentity"; +import { + bodyExceedsBudgetResponse, + chatAdmissionRejectionResponse, + resourcePressureRejectionResponse, + structuralRejectionResponse, +} from "./chatAdmissionResponses"; +import { estimateStructureTokens } from "./chatAdmissionStructureEstimate"; +import { + composeAdmissionLease, + IngestByteAdmissionController, + type IngestBudgetAcquireResult, +} from "./ingestByteAdmission"; +import { + getResourcePressureObservation, + type PressureSeverity, +} from "@omniroute/open-sse/utils/resourcePressure.ts"; function parsePositiveInt(value: string | undefined, fallback: number): number { const parsed = Number.parseInt(String(value), 10); @@ -180,7 +201,21 @@ interface AdmissionWaiter { * A client abort mid-wait is deliberately NOT a shed: capacity was never denied, * the caller simply left (its 503 is dropped on the dead connection). */ -export type ChatAdmissionShedReason = "queue_timeout" | "queued_bytes_budget"; +export type ChatAdmissionShedReason = + | "queue_timeout" + | "queued_bytes_budget" + | "body_exceeds_budget" + | "inflight_bytes_budget" + | "resource_pressure"; + +/** Read cached pressure severity; sampling failures must not cause false sheds. */ +export function defaultPressureSeverity(): PressureSeverity { + try { + return getResourcePressureObservation().state.severity; + } catch { + return "normal"; + } +} /** * One structural-shed observation, emitted to the shed sink at warn level. @@ -239,6 +274,8 @@ export class ChatAdmissionController { #shedsByReason = new Map(); readonly #onShed: ChatAdmissionShedSink; + readonly #ingestBudget: IngestByteAdmissionController; + constructor( readonly maxHeavyInFlight = 1, readonly maxQueuedBytes = CHAT_ADMISSION_MAX_QUEUED_BYTES, @@ -248,7 +285,13 @@ export class ChatAdmissionController { readonly healthyHeadroom = CHAT_ADMISSION_HEALTHY_HEADROOM, /** #11244: sink notified once per structural shed. Defaults to the shared pino * logger (warn); tests inject a capture/no-op sink. */ - onShed: ChatAdmissionShedSink = defaultChatAdmissionShedSink + onShed: ChatAdmissionShedSink = defaultChatAdmissionShedSink, + /** #503-fanout: see the field-level comment above `#inflightBytes`. */ + budgetOptions: { + maxInflightBytes?: number; + budgetSource?: IngestBudgetSource; + checkPressureSeverity?: () => PressureSeverity; + } = {} ) { if (!Number.isSafeInteger(maxHeavyInFlight) || maxHeavyInFlight < 1) { throw new RangeError("maxHeavyInFlight must be a positive integer"); @@ -260,6 +303,10 @@ export class ChatAdmissionController { throw new RangeError("healthyHeadroom must be a non-negative integer"); } this.#onShed = onShed; + this.#ingestBudget = new IngestByteAdmissionController({ + ...budgetOptions, + onShed: (reason, lane) => this.recordShed(reason, lane), + }); } get activeHeavy(): number { @@ -505,6 +552,39 @@ export class ChatAdmissionController { return; } } + + get inflightBytes(): number { + return this.#ingestBudget.inflightBytes; + } + + get maxInflightBytes(): number { + return this.#ingestBudget.maxInflightBytes; + } + + get budgetSource(): IngestBudgetSource { + return this.#ingestBudget.budgetSource; + } + + pressureSeverity(): PressureSeverity { + return this.#ingestBudget.pressureSeverity(); + } + + canFitBudget(bytes: number): boolean { + return this.#ingestBudget.canFit(bytes); + } + + tryAcquireBudget(bytes: number): ChatAdmissionLease | null { + return this.#ingestBudget.tryAcquire(bytes); + } + + acquireBudgetWithin( + bytes: number, + timeoutMs: number, + signal?: AbortSignal, + sessionKey = "default" + ): Promise { + return this.#ingestBudget.acquireWithin(bytes, timeoutMs, signal, sessionKey); + } } const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN_FLIGHT); @@ -526,53 +606,12 @@ const defaultAdmissionController = new ChatAdmissionController(CHAT_MAX_HEAVY_IN * per-key capacity being allocated. */ -export function resolveSessionId(request: Request): string { - // Fairness scheduling key ONLY (never a capacity shard): hashed so raw key - // material never appears in diagnostics. Reuses the internal-bypass auth - // extraction: bearer token from Authorization, x-api-key (Anthropic-style), - // or Google API key header. - // CodeQL: Intentionally HMAC-SHA256 with a fixed context key, NOT password hashing. The - // digest is a deterministic, non-reversible per-key fairness key for the shared admission - // budget — never stored or used for password-style verification. - const authHeader = request.headers.get("authorization") || ""; - const bearerMatch = /^bearer\s+(\S+)$/i.exec(authHeader.trim()); - if (bearerMatch) { - // Fingerprint for the admission-budget bucket key, not a password/credential hash — keyed - // with a fixed context label so it reads as a domain-separated digest, not a bare hash. - return ( - "key_" + - createHmac("sha256", "omniroute-admission-fingerprint-v1") - .update(bearerMatch[1]) - .digest("hex") - .slice(0, 16) - ); - } - const xApiKey = request.headers.get("x-api-key") || ""; - if (xApiKey.trim().length > 0) { - // Fingerprint for the admission-budget bucket key, not a password/credential hash — keyed - // with a fixed context label so it reads as a domain-separated digest, not a bare hash. - return ( - "key_" + - createHmac("sha256", "omniroute-admission-fingerprint-v1") - .update(xApiKey.trim()) - .digest("hex") - .slice(0, 16) - ); - } - const xGoogApiKey = request.headers.get("x-goog-api-key") || ""; - if (xGoogApiKey.trim().length > 0) { - // Fingerprint for the admission-budget bucket key, not a password/credential hash — keyed - // with a fixed context label so it reads as a domain-separated digest, not a bare hash. - return ( - "key_" + - createHmac("sha256", "omniroute-admission-fingerprint-v1") - .update(xGoogApiKey.trim()) - .digest("hex") - .slice(0, 16) - ); - } - return "anonymous"; -} +export { ADMISSION_BYPASS_HEADER, resolveSelfLoopBearer, resolveSessionId }; + +const NULL_LEASE: ChatAdmissionLease = { + released: true, + release() {}, +}; export class PerConnectionAdmissionController { readonly #controller: ChatAdmissionController; @@ -583,13 +622,26 @@ export class PerConnectionAdmissionController { // accepted for API compatibility and ignored — there are no per-session lanes // to evict. `onShed` (#11244) is live: it replaces the shed sink of the shared // controller (tests inject a capture/no-op sink; production keeps the pino warn). - _opts?: { maxSessions?: number; sessionTtlMs?: number; onShed?: ChatAdmissionShedSink } + // `budget` (#503-fanout) is live: the additive ingest byte-budget gate — see + // `ChatAdmissionController`'s constructor comment. Absent for every caller + // except the production singleton below. + _opts?: { + maxSessions?: number; + sessionTtlMs?: number; + onShed?: ChatAdmissionShedSink; + budget?: { + maxInflightBytes?: number; + budgetSource?: IngestBudgetSource; + checkPressureSeverity?: () => PressureSeverity; + }; + } ) { this.#controller = new ChatAdmissionController( maxHeavyInFlight, undefined, undefined, - _opts?.onShed + _opts?.onShed, + _opts?.budget ); } @@ -611,6 +663,17 @@ export class PerConnectionAdmissionController { lanes: ReadonlyArray<{ key: string; waiting: number }>; shedTotal: number; shedsByReason: Record; + /** #503-fanout: live ingest bytes reserved through the byte-budget gate. */ + inflightBytes: number; + /** #503-fanout: the auto-derived (or overridden) budget ceiling. */ + maxInflightBytes: number; + /** #503-fanout: which signal the budget was derived from. */ + budgetSource: IngestBudgetSource; + /** #503-fanout: live multi-signal resource-pressure severity. */ + pressureSeverity: PressureSeverity; + /** #503-fanout: false on a default deployment — the legacy count cap only + * binds when the operator explicitly set OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT. */ + countCapEnabled: boolean; } { return { activeHeavy: this.#controller.activeHeavy, @@ -620,6 +683,11 @@ export class PerConnectionAdmissionController { lanes: this.#controller.waitersByKey, shedTotal: this.#controller.shedTotal, shedsByReason: this.#controller.shedsByReason, + inflightBytes: this.#controller.inflightBytes, + maxInflightBytes: this.#controller.maxInflightBytes, + budgetSource: this.#controller.budgetSource, + pressureSeverity: this.#controller.pressureSeverity(), + countCapEnabled: this.#controller.maxHeavyInFlight < Number.MAX_SAFE_INTEGER, }; } @@ -641,8 +709,32 @@ export class PerConnectionAdmissionController { } } +/** + * The legacy count cap (#503-fanout) now binds ONLY when the operator has + * explicitly set `OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT`. Left unset — the + * default on every deployment that produced the multi-subagent 503 storm — + * it resolves to effectively unlimited, so the auto-derived ingest byte + * budget below (`resolveIngestByteBudget()`) is the gate that actually binds. + * A deployment that already tuned this env var (e.g. `infra/app.env.example` + * setting `=5`) keeps its exact prior behavior layered on top of the budget. + */ +function resolveLegacyCountCap(): number { + const raw = process.env.OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT; + if (raw === undefined || raw.trim() === "") return Number.MAX_SAFE_INTEGER; + return CHAT_MAX_HEAVY_IN_FLIGHT; +} + +const productionIngestBudget = resolveIngestByteBudget(); + export const perConnectionAdmissionController = new PerConnectionAdmissionController( - CHAT_MAX_HEAVY_IN_FLIGHT + resolveLegacyCountCap(), + { + budget: { + maxInflightBytes: productionIngestBudget.bytes, + budgetSource: productionIngestBudget.source, + checkPressureSeverity: defaultPressureSeverity, + }, + } ); export type ChatRequestAdmission = @@ -652,101 +744,7 @@ export type ChatRequestAdmission = export type ChatStructureAdmission = { admit: true; lease: ChatAdmissionLease | null } | { admit: false; response: Response }; -function rejectionResponse(status: 413 | 503, hardMaxBytes: number): Response { - const isPayload = status === 413; - const headers: Record = { - ...CORS_HEADERS, - "Content-Type": "application/json", - }; - if (!isPayload) headers["Retry-After"] = "2"; - return new Response( - JSON.stringify({ - error: { - message: isPayload - ? `Request body too large for chat completions (max ${Math.floor( - hardMaxBytes / (1024 * 1024) - )} MB).` - : "Chat admission capacity is temporarily unavailable. Retry shortly.", - type: isPayload ? "payload_too_large" : "server_error", - code: isPayload ? "PAYLOAD_TOO_LARGE" : "chat_admission_busy", - }, - }), - { status, headers } - ); -} - -function structuralRejectionResponse(status: 413 | 503, maxMessages: number): Response { - const historyLimit = status === 413; - const headers: Record = { - ...CORS_HEADERS, - "Content-Type": "application/json", - }; - if (!historyLimit) headers["Retry-After"] = "1"; - - return new Response( - JSON.stringify({ - error: { - message: historyLimit - ? `Chat history exceeds the ${maxMessages}-message limit; compact the conversation and retry.` - : "Structurally heavy chat request capacity is busy; retry shortly.", - type: historyLimit ? "payload_too_large" : "server_error", - code: historyLimit ? "chat_history_too_large" : "chat_admission_busy", - reason: historyLimit ? "message_limit" : "structure_limit", - }, - }), - { status, headers } - ); -} - -type TokenEstimate = { tokens: number; exhausted: boolean }; - -function conservativeStringTokens(value: string, remaining: number): number { - let tokens = 0; - for (const character of value) { - tokens += character.codePointAt(0)! < 0x80 ? 0.25 : 1; - if (tokens >= remaining) return remaining; - } - return tokens; -} - -function estimateStructureTokens(value: unknown, limit: number): TokenEstimate { - let tokens = 0; - let visited = 0; - const maxNodes = 10_000; - const stack: Array<{ value: unknown; depth: number }> = [{ value, depth: 0 }]; - while (stack.length > 0 && tokens < limit && visited < maxNodes) { - const current = stack.pop(); - if (!current) break; - visited += 1; - if (typeof current.value === "string") { - tokens += conservativeStringTokens(current.value, limit - tokens); - continue; - } - if (!current.value || typeof current.value !== "object") continue; - if (current.depth >= 12) return { tokens, exhausted: true }; - - const remainingNodes = maxNodes - visited - stack.length; - if (Array.isArray(current.value)) { - if (current.value.length > remainingNodes) return { tokens, exhausted: true }; - for (const child of current.value) stack.push({ value: child, depth: current.depth + 1 }); - continue; - } - - let children = 0; - for (const key in current.value) { - if (!Object.hasOwn(current.value, key)) continue; - children += 1; - if (children > remainingNodes) return { tokens, exhausted: true }; - tokens += conservativeStringTokens(key, limit - tokens); - if (tokens >= limit) return { tokens: limit, exhausted: false }; - stack.push({ - value: (current.value as Record)[key], - depth: current.depth + 1, - }); - } - } - return { tokens, exhausted: stack.length > 0 && tokens < limit }; -} +const INGEST_NORMAL_MAX_WAIT_MS = 250; export async function admitChatStructure( body: unknown, @@ -803,9 +801,21 @@ export async function admitChatStructure( ? perConnectionAdmissionController.getController(options.sessionId) : defaultAdmissionController); - // Uncontended fast path: capacity is free, no need to consult heap pressure at all. - const immediate = controller.tryAcquireHeavy(); - if (immediate) return { admit: true, lease: immediate }; + // Uncontended fast path: capacity is free on BOTH the legacy count gate and + // the byte-budget gate (#503-fanout) — mirrors admitChatRequest's composed + // reserve(). When the count cap is unlimited (the production default since + // this fix), the byte-budget gate is what actually decides "uncontended": + // without composing both here, a structurally-heavy-but-byte-light request + // would always take this fast path and the heap-pressure-conditional shed + // below would never be reachable in production. + const immediateCount = controller.tryAcquireHeavy(); + if (immediateCount) { + const immediateBudget = controller.tryAcquireBudget(CHAT_LARGE_BODY_BYTES); + if (immediateBudget) { + return { admit: true, lease: composeAdmissionLease(immediateCount, immediateBudget) }; + } + immediateCount.release(); + } // Heavyweight capacity is momentarily busy (a concurrent heavy request holds the // lease). #10183 / #10268: only enter the bounded-wait / shed path — with its @@ -833,15 +843,31 @@ export async function admitChatStructure( // 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( + const acquiredCount = await controller.acquireHeavyWithin( options.queueMs ?? 0, options.signal, CHAT_LARGE_BODY_BYTES, options.sessionId ); - return acquired - ? { admit: true, lease: acquired } - : { admit: false, response: structuralRejectionResponse(503, maxMessages) }; + if (!acquiredCount) { + return { admit: false, response: structuralRejectionResponse(503, maxMessages) }; + } + + // #503-fanout: same composed count+budget gate as the fast path above. + const acquiredBudget = await controller.acquireBudgetWithin( + CHAT_LARGE_BODY_BYTES, + options.queueMs ?? 0, + options.signal, + options.sessionId + ); + if (acquiredBudget.status !== "acquired") { + acquiredCount.release(); + return { admit: false, response: structuralRejectionResponse(503, maxMessages) }; + } + return { + admit: true, + lease: composeAdmissionLease(acquiredCount, acquiredBudget.lease), + }; } function parseContentLength(header: string | null): number | null { @@ -863,83 +889,6 @@ function rebuildRequest(request: Request, body: Uint8Array): Request { } as RequestInit & { duplex: "half" }); } -/** - * Internal self-loop bypass marker for the vision-bridge describe call (and any - * other trusted in-process sub-request). An external client cannot spoof it: - * it is honored ONLY when combined with a trusted self-loop credential — the - * local-mode `sk_omniroute` sentinel or the operator-configured env key - * (`OMNIROUTE_API_KEY` / `ROUTER_API_KEY`, #1350) so REQUIRE_API_KEY=true - * deployments can run the describe sub-request. - */ -export const ADMISSION_BYPASS_HEADER = "x-omniroute-admission-bypass"; -const ADMISSION_BYPASS_VALUE = "internal"; -const SELF_LOOP_KEY = "sk_omniroute"; - -/** - * Resolve the bearer credential used by trusted in-process self-loop - * sub-requests (the vision-bridge describe call). - * - * Local mode uses the `sk_omniroute` sentinel. Deployments that force API key - * auth (`REQUIRE_API_KEY=true`) reject that sentinel with 401, so they must use - * a real key — the persistent env-var key (#1350, `OMNIROUTE_API_KEY` / - * `ROUTER_API_KEY`) is the natural choice because it always validates and - * survives restarts. Falls back to the sentinel when no env key is configured - * so local-mode behavior is unchanged. - */ -export function resolveSelfLoopBearer(): string { - return ( - process.env.OMNIROUTE_API_KEY?.trim() || process.env.ROUTER_API_KEY?.trim() || SELF_LOOP_KEY - ); -} - -/** - * Sentinel lease returned by the admission byte stage for an internal self-loop - * sub-request (the vision-bridge describe call). The parent request already holds - * the single heavyweight lease, so the describe call must never reserve again — - * but a NON-NULL lease is still required so the route's later structural stage - * (`admitChatStructure`) treats the body as covered. With `lease: null` the - * structural stage classifies the base64-heavy describe body as "heavy" and tries - * to acquire the busy capacity, returning 503 `chat_admission_busy` anyway — the - * gap that kept the Zoo Code / api-key describe call failing even after the byte - * stage was bypassed. Release is a no-op; capacity was never reserved. - */ -function createNoopLease(): ChatAdmissionLease { - return { - get released() { - return true; - }, - release() { - // No-op: this sentinel never reserved heavyweight capacity. - }, - }; -} - -const NULL_LEASE: ChatAdmissionLease = createNoopLease(); - -/** - * True when the request is a trusted in-process self-loop sub-request that must - * not consume a heavyweight admission lease. The describe call runs WHILE the - * parent request already holds the single heavyweight lease (`CHAT_MAX_HEAVY_IN_FLIGHT=1`), - * so without this bypass it is rejected with 503 `chat_admission_busy` and the - * image is never described (#vision-bridge self-loop). - */ -function isInternalAdmissionBypass(request: Request): boolean { - const bypass = - request.headers.get(ADMISSION_BYPASS_HEADER)?.trim().toLowerCase() === ADMISSION_BYPASS_VALUE; - if (!bypass) return false; - - // Credential gate: the bypass only applies to trusted self-loop credentials — - // the local `sk_omniroute` sentinel OR the operator-configured env key - // (`OMNIROUTE_API_KEY` / `ROUTER_API_KEY`, #1350) so REQUIRE_API_KEY=true - // deployments can still run the vision-bridge describe sub-request. The env - // key is a secret like any other API key, so honoring it here does not widen - // the attack surface: a third-party that holds it can already call every API. - const auth = request.headers.get("authorization") || ""; - const match = /^bearer\s+(\S+)$/i.exec(auth.trim()); - if (!match) return false; - return match[1].trim().toLowerCase() === resolveSelfLoopBearer().toLowerCase(); -} - /** * Reserve heavyweight capacity and ingest the body with a hard byte bound before JSON * parsing. Missing/invalid Content-Length is sniffed only up to the heavyweight threshold; @@ -970,9 +919,8 @@ export async function admitChatRequest( // Internal self-loop: skip the heavyweight reservation entirely (the parent // request already holds the single lease) but still enforce the hard byte bound. if (internalBypass) { - const contentLengthHeader = request.headers.get("content-length"); if (contentLength !== null && contentLength > hardMaxBytes) { - return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(413, hardMaxBytes) }; } // Sniff bytes for the hard bound without reserving a lease. const reader = request.body?.getReader(); @@ -986,7 +934,7 @@ export async function admitChatRequest( totalBytes += value.byteLength; if (totalBytes > hardMaxBytes) { await reader.cancel("chat request exceeds hard body limit").catch(() => undefined); - return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(413, hardMaxBytes) }; } chunks.push(value); } @@ -1004,15 +952,58 @@ export async function admitChatRequest( return { admit: true, request: rebuildRequest(request, body), lease: NULL_LEASE }; } + // #503-fanout: shed before spending any bytes on ingestion when the process + // is under genuine critical resource pressure. No-op for every controller a + // test constructs directly (default severity is always "normal"). + if (controller.pressureSeverity() === "critical") { + controller.recordShed("resource_pressure", sessionId); + return { admit: false, response: resourcePressureRejectionResponse() }; + } + if (contentLength !== null && contentLength > hardMaxBytes) { - return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(413, hardMaxBytes) }; + } + if ( + contentLength !== null && + contentLength >= largeBodyBytes && + !controller.canFitBudget(contentLength) + ) { + controller.recordShed("body_exceeds_budget", sessionId); + return { admit: false, response: bodyExceedsBudgetResponse(controller.maxInflightBytes) }; } let lease: ChatAdmissionLease | null = null; const reserve = async (bytes = 0): Promise => { if (lease) return true; - lease = await controller.acquireHeavyWithin(queueMs, request.signal, bytes, sessionId); - return lease !== null; + const countLease = await controller.acquireHeavyWithin( + queueMs, + request.signal, + bytes, + sessionId + ); + if (!countLease) return false; + + // Additive ingest byte-budget gate (#503-fanout), layered on top of the + // legacy count gate above. `maxInflightBytes` defaults to unlimited for + // every controller a test constructs directly, so this resolves + // synchronously true there — only the production singleton (built with a + // real host-derived budget) is ever actually gated by it. + const severity = controller.pressureSeverity(); + const budgetWaitMs = + severity === "high" ? queueMs : Math.min(queueMs, INGEST_NORMAL_MAX_WAIT_MS); + const budgetResult = await controller.acquireBudgetWithin( + bytes, + budgetWaitMs, + request.signal, + sessionId + ); + if (budgetResult.status !== "acquired") { + countLease.release(); + return false; + } + + lease = composeAdmissionLease(countLease, budgetResult.lease); + return true; }; // A known-large declaration can reserve before ingestion. Unknown lengths are boundedly @@ -1022,7 +1013,7 @@ export async function admitChatRequest( contentLength >= largeBodyBytes && !(await reserve(Math.min(contentLength, hardMaxBytes))) ) { - return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(503, hardMaxBytes) }; } const reader = request.body?.getReader(); @@ -1038,11 +1029,17 @@ export async function admitChatRequest( if (totalBytes > hardMaxBytes) { await reader.cancel("chat request exceeds hard body limit").catch(() => undefined); lease?.release(); - return { admit: false, response: rejectionResponse(413, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(413, hardMaxBytes) }; + } + if (totalBytes >= largeBodyBytes && !controller.canFitBudget(totalBytes)) { + controller.recordShed("body_exceeds_budget", sessionId); + await reader.cancel("chat request exceeds ingest budget").catch(() => undefined); + lease?.release(); + return { admit: false, response: bodyExceedsBudgetResponse(controller.maxInflightBytes) }; } if (totalBytes >= largeBodyBytes && !(await reserve(totalBytes))) { await reader.cancel("chat admission capacity unavailable").catch(() => undefined); - return { admit: false, response: rejectionResponse(503, hardMaxBytes) }; + return { admit: false, response: chatAdmissionRejectionResponse(503, hardMaxBytes) }; } chunks.push(value); } diff --git a/src/shared/middleware/ingestByteAdmission.ts b/src/shared/middleware/ingestByteAdmission.ts new file mode 100644 index 0000000000..2420074855 --- /dev/null +++ b/src/shared/middleware/ingestByteAdmission.ts @@ -0,0 +1,184 @@ +import type { PressureSeverity } from "@omniroute/open-sse/utils/resourcePressure.ts"; + +import type { IngestBudgetSource } from "./admissionBudget"; + +export interface IngestByteLease { + readonly released: boolean; + release(): void; +} + +export function composeAdmissionLease(...leases: IngestByteLease[]): IngestByteLease { + let released = false; + return { + get released() { + return released; + }, + release: () => { + if (released) return; + released = true; + for (const lease of leases) lease.release(); + }, + }; +} + +export type IngestBudgetAcquireResult = + | { status: "acquired"; lease: IngestByteLease } + | { status: "body_exceeds_budget" } + | { status: "unavailable" }; + +export interface IngestByteAdmissionOptions { + maxInflightBytes?: number; + budgetSource?: IngestBudgetSource; + checkPressureSeverity?: () => PressureSeverity; + onShed: (reason: "body_exceeds_budget" | "inflight_bytes_budget", lane: string) => void; +} + +interface BudgetWaiter { + resolve: () => void; +} + +export class IngestByteAdmissionController { + #inflightBytes = 0; + readonly maxInflightBytes: number; + readonly budgetSource: IngestBudgetSource; + readonly #checkPressureSeverity: () => PressureSeverity; + readonly #onShed: IngestByteAdmissionOptions["onShed"]; + #queues = new Map(); + #fairKeys: string[] = []; + #fairCursor = 0; + + constructor(options: IngestByteAdmissionOptions) { + this.maxInflightBytes = options.maxInflightBytes ?? Number.MAX_SAFE_INTEGER; + this.budgetSource = options.budgetSource ?? "v8_heap"; + this.#checkPressureSeverity = options.checkPressureSeverity ?? (() => "normal"); + this.#onShed = options.onShed; + } + + get inflightBytes(): number { + return this.#inflightBytes; + } + + pressureSeverity(): PressureSeverity { + return this.#checkPressureSeverity(); + } + + canFit(bytes: number): boolean { + return normalizeCharge(bytes) <= this.maxInflightBytes; + } + + tryAcquire(bytes: number): IngestByteLease | null { + const charge = normalizeCharge(bytes); + if (this.#inflightBytes + charge > this.maxInflightBytes) return null; + this.#inflightBytes += charge; + let released = false; + return { + get released() { + return released; + }, + release: () => { + if (released) return; + released = true; + this.#inflightBytes = Math.max(0, this.#inflightBytes - charge); + this.#dispatchFair(); + }, + }; + } + + async acquireWithin( + bytes: number, + timeoutMs: number, + signal?: AbortSignal, + sessionKey = "default" + ): Promise { + if (!this.canFit(bytes)) { + this.#onShed("body_exceeds_budget", sessionKey); + return { status: "body_exceeds_budget" }; + } + const deadline = Date.now() + Math.max(0, Math.floor(timeoutMs)); + for (;;) { + if (signal?.aborted) return { status: "unavailable" }; + const lease = this.tryAcquire(bytes); + if (lease) return { status: "acquired", lease }; + const remaining = deadline - Date.now(); + if (remaining <= 0) return this.#timeout(sessionKey); + + let queue = this.#queues.get(sessionKey); + if (!queue) { + queue = []; + this.#queues.set(sessionKey, queue); + this.#fairKeys.push(sessionKey); + } + let resolveParked: (() => void) | null = null; + const waiter = { resolve: () => resolveParked?.() }; + const parked = new Promise((resolve) => { + resolveParked = resolve; + queue.push(waiter); + }); + let timer: ReturnType | null = null; + const races: Array> = [ + parked.then(() => false), + new Promise((resolve) => { + timer = setTimeout(() => resolve(true), remaining); + }), + ]; + let onAbort: (() => void) | null = null; + if (signal) { + races.push( + new Promise((resolve) => { + onAbort = () => resolve(true); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) resolve(true); + }) + ); + } + const timedOut = await Promise.race(races); + this.#removeWaiter(sessionKey, waiter); + if (timer) clearTimeout(timer); + if (onAbort) signal?.removeEventListener("abort", onAbort); + if (timedOut) { + if (signal?.aborted) return { status: "unavailable" }; + return this.#timeout(sessionKey); + } + } + } + + #timeout(sessionKey: string): IngestBudgetAcquireResult { + this.#onShed("inflight_bytes_budget", sessionKey); + return { status: "unavailable" }; + } + + #removeWaiter(key: string, waiter: BudgetWaiter): void { + const queue = this.#queues.get(key); + if (!queue) return; + const index = queue.indexOf(waiter); + if (index >= 0) queue.splice(index, 1); + if (queue.length === 0) this.#removeFairKey(key); + } + + #removeFairKey(key: string): void { + this.#queues.delete(key); + const index = this.#fairKeys.indexOf(key); + if (index < 0) return; + this.#fairKeys.splice(index, 1); + if (index < this.#fairCursor) this.#fairCursor -= 1; + if (this.#fairKeys.length === 0) this.#fairCursor = 0; + } + + #dispatchFair(): void { + if (this.#fairKeys.length === 0) return; + for (let i = 0; i < this.#fairKeys.length; i++) { + const key = this.#fairKeys[this.#fairCursor % this.#fairKeys.length]; + this.#fairCursor += 1; + const queue = this.#queues.get(key); + if (!queue || queue.length === 0) continue; + const waiter = queue.shift() as BudgetWaiter; + if (queue.length === 0) this.#removeFairKey(key); + waiter.resolve(); + return; + } + } +} + +function normalizeCharge(bytes: number): number { + return Math.max(0, Math.floor(bytes)); +} diff --git a/tests/unit/admission-ingest-budget.test.ts b/tests/unit/admission-ingest-budget.test.ts new file mode 100644 index 0000000000..cda4b55064 --- /dev/null +++ b/tests/unit/admission-ingest-budget.test.ts @@ -0,0 +1,123 @@ +// Auto-derived ingest byte budget (#503-fanout): the byte budget that replaces +// the legacy request-count admission cap must scale itself from the host's +// real memory ceiling instead of a fixed number, on both bare-metal (V8 heap +// limit only) and containerized (cgroup-constrained) hosts. +import test from "node:test"; +import assert from "node:assert/strict"; +import { + computeIngestByteBudget, + MIN_INGEST_BUDGET_BYTES, + MAX_INGEST_BUDGET_BYTES, + INGEST_HEAP_FRACTION, + INGEST_AMPLIFICATION, +} from "../../src/shared/middleware/admissionBudget.ts"; + +const GiB = 1024 * 1024 * 1024; +const MiB = 1024 * 1024; + +test("desktop host (no cgroup limit): budget derives from the V8 heap ceiling", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + constrainedMemoryBytes: null, + }); + assert.equal(budget.source, "v8_heap"); + assert.equal(budget.effectiveCeilingBytes, 4 * GiB); + const expected = Math.floor((4 * GiB * INGEST_HEAP_FRACTION) / INGEST_AMPLIFICATION); + assert.equal(budget.bytes, expected); + assert.ok(budget.bytes > 32 * MiB, "a 4 GiB heap must yield a generous budget, not the floor"); +}); + +test("container host: a tighter cgroup limit wins over a larger V8 heap ceiling", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, // Node's default heap sizing can exceed the container's real limit + constrainedMemoryBytes: 512 * MiB, + }); + assert.equal(budget.source, "cgroup"); + assert.equal(budget.effectiveCeilingBytes, 512 * MiB); + assert.equal(budget.bytes, Math.floor((512 * MiB * INGEST_HEAP_FRACTION) / INGEST_AMPLIFICATION)); +}); + +test("a looser cgroup limit than the V8 heap ceiling does not widen the budget", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 1 * GiB, + constrainedMemoryBytes: 16 * GiB, + }); + assert.equal(budget.source, "v8_heap"); + assert.equal(budget.effectiveCeilingBytes, 1 * GiB); +}); + +test("tiny container: the budget clamps to the floor instead of rejecting at idle", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 256 * MiB, + constrainedMemoryBytes: 256 * MiB, + }); + assert.equal(budget.bytes, MIN_INGEST_BUDGET_BYTES); +}); + +test("an enormous ceiling clamps to the max, never grows unbounded", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 1024 * GiB, + constrainedMemoryBytes: null, + }); + assert.equal(budget.bytes, MAX_INGEST_BUDGET_BYTES); +}); + +test("a positive override always wins over both heap and cgroup signals", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + constrainedMemoryBytes: 512 * MiB, + override: 64 * MiB, + }); + assert.equal(budget.source, "override"); + assert.equal(budget.bytes, 64 * MiB); + assert.equal(budget.effectiveCeilingBytes, 64 * MiB); +}); + +test("a string override (env var shape) parses the same as a number", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + override: "134217728", // 128 MiB + }); + assert.equal(budget.source, "override"); + assert.equal(budget.bytes, 128 * MiB); +}); + +test("an override below the safe range clamps to the minimum", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + override: 1024, + }); + assert.equal(budget.source, "override"); + assert.equal(budget.bytes, MIN_INGEST_BUDGET_BYTES); + assert.equal(budget.effectiveCeilingBytes, MIN_INGEST_BUDGET_BYTES); +}); + +test("an override above the safe range clamps to the maximum", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + override: 4 * GiB, + }); + assert.equal(budget.source, "override"); + assert.equal(budget.bytes, MAX_INGEST_BUDGET_BYTES); + assert.equal(budget.effectiveCeilingBytes, MAX_INGEST_BUDGET_BYTES); +}); + +test("invalid overrides (NaN, zero, negative, empty) fall through to auto-derivation", () => { + for (const bad of ["not-a-number", "0", "-5", "", null, undefined, NaN, 0, -1]) { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: 4 * GiB, + constrainedMemoryBytes: null, + override: bad as never, + }); + assert.equal(budget.source, "v8_heap", `override ${JSON.stringify(bad)} must not be honored`); + } +}); + +test("a non-finite heap limit input falls back to the floor input rather than throwing", () => { + const budget = computeIngestByteBudget({ + heapSizeLimitBytes: NaN, + constrainedMemoryBytes: null, + }); + assert.ok(Number.isFinite(budget.bytes)); + assert.equal(budget.bytes, MIN_INGEST_BUDGET_BYTES); +}); diff --git a/tests/unit/admission-pressure-gate.test.ts b/tests/unit/admission-pressure-gate.test.ts new file mode 100644 index 0000000000..2db9d1d9af --- /dev/null +++ b/tests/unit/admission-pressure-gate.test.ts @@ -0,0 +1,214 @@ +// #503-fanout: the ingest byte-budget gate must be pressure-driven, not +// unconditional. `normal` admits within budget (bounded wait capped short); +// `high` uses the caller's full bounded wait; `critical` sheds before any +// bytes are even ingested. This is the counterpart to +// agent-fanout-admission-regression.test.ts, focused on the pressure +// dimension rather than the fan-out/concurrency dimension. +import test from "node:test"; +import assert from "node:assert/strict"; + +const { ChatAdmissionController, admitChatRequest } = + await import("../../src/shared/middleware/chatBodyAdmission.ts"); + +const silentSink = () => {}; + +function bodyOf(bytes: number): string { + return "x".repeat(bytes); +} + +function requestFor(body: string): Request { + return new Request("http://x/v1/messages", { + method: "POST", + headers: { "content-type": "application/json", "content-length": String(body.length) }, + body, + }); +} + +test("normal pressure: a request within the byte budget is admitted", async () => { + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + silentSink, + { + maxInflightBytes: 1024 * 1024, + checkPressureSeverity: () => "normal", + } + ); + + const result = await admitChatRequest(requestFor(bodyOf(4096)), { + controller, + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 5000, + }); + + assert.equal(result.admit, true); + if (result.admit) result.lease?.release(); +}); + +test("normal pressure: contention sheds within the short wait instead of the full queueMs", async () => { + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + silentSink, + { + maxInflightBytes: 4096, + checkPressureSeverity: () => "normal", + } + ); + const occupied = controller.tryAcquireBudget(4096); + assert.ok(occupied); + + const start = Date.now(); + const result = await admitChatRequest(requestFor(bodyOf(2048)), { + controller, + sessionId: "budget-exhausted", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 5000, + }); + const elapsedMs = Date.now() - start; + occupied.release(); + + assert.equal(result.admit, false); + if (!result.admit) assert.equal(result.response.status, 503); + assert.ok( + elapsedMs < 2000, + `normal pressure must cap the ingest wait well under the full queueMs (took ${elapsedMs}ms)` + ); +}); + +test("a body larger than the whole budget fails immediately with a distinct diagnosis", async () => { + const sheds: string[] = []; + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + (event) => sheds.push(event.reason), + { maxInflightBytes: 1024, checkPressureSeverity: () => "high" } + ); + + const start = Date.now(); + const result = await admitChatRequest(requestFor(bodyOf(4096)), { + controller, + sessionId: "unservable-body", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 5000, + }); + const elapsedMs = Date.now() - start; + + assert.equal(result.admit, false); + if (result.admit) return; + const payload = (await result.response.json()) as { error: { code: string } }; + assert.equal(result.response.status, 413); + assert.equal(result.response.headers.get("retry-after"), null); + assert.equal(payload.error.code, "body_exceeds_budget"); + assert.deepEqual(sheds, ["body_exceeds_budget"]); + assert.equal(controller.activeHeavy, 0); + assert.equal(controller.inflightBytes, 0); + assert.ok( + elapsedMs < 1000, + `an impossible charge must not enter the wait queue (took ${elapsedMs}ms)` + ); +}); + +test("high pressure: contention waits up to the full queueMs before shedding", async () => { + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + silentSink, + { + maxInflightBytes: 4096, + checkPressureSeverity: () => "high", + } + ); + const occupied = controller.tryAcquireBudget(4096); + assert.ok(occupied); + + const start = Date.now(); + const result = await admitChatRequest(requestFor(bodyOf(2048)), { + controller, + sessionId: "high-pressure-wait", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 300, + }); + const elapsedMs = Date.now() - start; + occupied.release(); + + assert.equal(result.admit, false); + assert.ok( + elapsedMs >= 280, + `high pressure must honor the full bounded wait (took ${elapsedMs}ms)` + ); +}); + +test("high pressure: budget freed mid-wait is claimed instead of shedding", async () => { + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + silentSink, + { + maxInflightBytes: 4096, + checkPressureSeverity: () => "high", + } + ); + + // Occupy the entire budget first. + const occupied = controller.tryAcquireBudget(4096); + assert.ok(occupied); + + const pending = admitChatRequest(requestFor(bodyOf(2048)), { + controller, + sessionId: "high-pressure-freed", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 2000, + }); + + setTimeout(() => occupied.release(), 30); + const result = await pending; + assert.equal(result.admit, true, "freeing budget mid-wait must let the waiter through"); + if (result.admit) result.lease?.release(); +}); + +test("critical pressure: the whole request is shed before ingestion, with a distinct code", async () => { + const controller = new ChatAdmissionController( + Number.MAX_SAFE_INTEGER, + undefined, + 0, + silentSink, + { + maxInflightBytes: 1024 * 1024 * 1024, // budget is not the limiting factor here + checkPressureSeverity: () => "critical", + } + ); + + const result = await admitChatRequest(requestFor(bodyOf(64)), { + controller, + sessionId: "critical-shed", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 5000, + }); + + assert.equal(result.admit, false); + if (!result.admit) { + assert.equal(result.response.status, 503); + assert.equal(result.response.headers.get("Retry-After"), "2"); + const payload = await result.response.json(); + assert.equal(payload.error.code, "resource_pressure"); + } + assert.deepEqual(controller.shedsByReason, { resource_pressure: 1 }); +}); + +test("pressureSeverity() defaults to normal for a controller with no injected probe", () => { + const controller = new ChatAdmissionController(1); + assert.equal(controller.pressureSeverity(), "normal"); + assert.equal(controller.maxInflightBytes, Number.MAX_SAFE_INTEGER); +}); diff --git a/tests/unit/agent-fanout-admission-regression.test.ts b/tests/unit/agent-fanout-admission-regression.test.ts new file mode 100644 index 0000000000..ec1775edcc --- /dev/null +++ b/tests/unit/agent-fanout-admission-regression.test.ts @@ -0,0 +1,182 @@ +// #503-fanout: coding-agent fan-out (multiple subagents / CLIs hitting the +// proxy concurrently, often sharing one API key) landed on the legacy +// count-based chat admission gate (default CHAT_MAX_HEAVY_IN_FLIGHT=1) and +// serialized to an effective concurrency of ~1, so a burst of concurrent +// requests mostly hung, retried, and then received 503 chat_admission_busy. +// +// This is the Hard-Rule-#18 TDD regression guard: the first test below +// reproduces the OLD behavior exactly as it shipped (a plain +// `new ChatAdmissionController(1)`, matching the pre-fix production default) +// to document why the bug happened. The remaining tests exercise the SAME +// `admitChatRequest` code path with the controller shape the FIXED +// production singleton now builds (an effectively unlimited legacy count cap +// plus a real, host-derived ingest byte budget) and prove the fan-out no +// longer serializes or 503s. +import test from "node:test"; +import assert from "node:assert/strict"; + +const { ChatAdmissionController, admitChatRequest } = await import( + "../../src/shared/middleware/chatBodyAdmission.ts" +); + +const silentSink = () => {}; + +function agentBody(bytes = 400_000): string { + // Shape of a real coding-agent turn: a big system prompt + many tool + // definitions serialized as one JSON string. The exact content does not + // matter here — only the byte size, which is what the ingest gate charges. + const filler = "x".repeat(Math.max(0, bytes - 64)); + return JSON.stringify({ model: "test-model", messages: [{ role: "user", content: filler }] }); +} + +function agentRequest(body: string): Request { + return new Request("http://x/v1/messages", { + method: "POST", + headers: { "content-type": "application/json", "content-length": String(body.length) }, + body, + }); +} + +test("OLD behavior: 8 concurrent agent-fanout requests mostly 503 under the count=1 gate", async () => { + // Matches the pre-fix production default exactly (CHAT_MAX_HEAVY_IN_FLIGHT + // defaulted to 1 unconditionally). This documents the bug; it is not + // expected to change as part of this fix — the legacy count cap still + // behaves identically when an operator (or a test) constructs it directly. + const controller = new ChatAdmissionController(1, undefined, 0, silentSink); + const body = agentBody(); + const requests = Array.from({ length: 8 }, () => agentRequest(body)); + + const results = await Promise.all( + requests.map((request) => + admitChatRequest(request, { + controller, + sessionId: "same-api-key", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 0, + }) + ) + ); + + const admitted = results.filter((r) => r.admit === true); + const shed = results.filter((r) => r.admit === false); + assert.equal(admitted.length, 1, "the count=1 gate admits exactly one of the eight"); + assert.equal(shed.length, 7, "the remaining seven are shed under the old gate"); + for (const result of shed) { + if (result.admit) continue; + assert.equal(result.response.status, 503); + const payload = await result.response.json(); + assert.equal(payload.error.code, "chat_admission_busy"); + } + + for (const result of admitted) { + if (result.admit) result.lease?.release(); + } +}); + +test("FIX: the same 8-request agent fan-out all admit under the byte-budget gate", async () => { + // Shape of the fixed production singleton: legacy count cap disabled + // (Number.MAX_SAFE_INTEGER — matches resolveLegacyCountCap()'s unset-env + // default) and a real ingest byte budget generously larger than the burst. + const controller = new ChatAdmissionController(Number.MAX_SAFE_INTEGER, undefined, 0, silentSink, { + maxInflightBytes: 8 * 1024 * 1024, // 8 MiB — 8 x ~400 KB bodies fit comfortably + checkPressureSeverity: () => "normal", + }); + const body = agentBody(); + const requests = Array.from({ length: 8 }, () => agentRequest(body)); + + const start = Date.now(); + const results = await Promise.all( + requests.map((request) => + admitChatRequest(request, { + controller, + sessionId: "same-api-key", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 2000, + }) + ) + ); + const elapsedMs = Date.now() - start; + + const shed = results.filter((r) => r.admit === false); + assert.equal(shed.length, 0, "no request should be shed once concurrency is byte-budgeted"); + assert.equal(results.length, 8); + assert.ok( + elapsedMs < 1000, + `admission must resolve promptly, not serialize behind the old count=1 gate (took ${elapsedMs}ms)` + ); + + assert.equal( + controller.inflightBytes, + body.length * 8, + "sanity: every one of the eight requests charged its real body size before any release" + ); + + for (const result of results) { + if (result.admit) result.lease?.release(); + } + assert.equal(controller.inflightBytes, 0, "every charge is released"); +}); + +test("FIX: distinct sessions in the same fan-out are also all admitted (not just a shared API key)", async () => { + const controller = new ChatAdmissionController(Number.MAX_SAFE_INTEGER, undefined, 0, silentSink, { + maxInflightBytes: 8 * 1024 * 1024, + checkPressureSeverity: () => "normal", + }); + const body = agentBody(); + + const results = await Promise.all( + Array.from({ length: 8 }, (_, index) => + admitChatRequest(agentRequest(body), { + controller, + sessionId: `session-${index}`, + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 2000, + }) + ) + ); + + assert.equal( + results.filter((r) => r.admit === false).length, + 0, + "distinct sessions must not starve each other or the shared budget" + ); + + for (const result of results) { + if (result.admit) result.lease?.release(); + } +}); + +test("critical resource pressure still sheds the entire fan-out (the gate is pressure-driven, not removed)", async () => { + const controller = new ChatAdmissionController(Number.MAX_SAFE_INTEGER, undefined, 0, silentSink, { + maxInflightBytes: 8 * 1024 * 1024, + checkPressureSeverity: () => "critical", + }); + const body = agentBody(); + + const results = await Promise.all( + Array.from({ length: 8 }, () => + admitChatRequest(agentRequest(body), { + controller, + sessionId: "same-api-key", + largeBodyBytes: 1024, + hardMaxBytes: 10 * 1024 * 1024, + queueMs: 2000, + }) + ) + ); + + assert.equal( + results.filter((r) => r.admit === true).length, + 0, + "critical pressure must shed the whole burst, not just serialize it" + ); + for (const result of results) { + if (result.admit) continue; + assert.equal(result.response.status, 503); + const payload = await result.response.json(); + assert.equal(payload.error.code, "resource_pressure"); + } +}); diff --git a/tests/unit/observability-payloads.test.ts b/tests/unit/observability-payloads.test.ts index 5cb9241b28..05bc8f6b76 100644 --- a/tests/unit/observability-payloads.test.ts +++ b/tests/unit/observability-payloads.test.ts @@ -348,11 +348,17 @@ test("buildHealthPayload projects allowlisted structural chatAdmission fields on waiting: 2, queuedBytes: 524_288, shedTotal: 3, - shedsByReason: { queue_timeout: 2, queued_bytes_budget: 1 }, + shedsByReason: { queue_timeout: 2, body_exceeds_budget: 1 }, lanes: [ { key: "key_c49d1c242feda590", waiting: 1 }, { key: "anonymous", waiting: 1 }, ], + // #503-fanout additions. + inflightBytes: 131_072, + maxInflightBytes: 134_217_728, + budgetSource: "v8_heap", + pressureSeverity: "normal", + countCapEnabled: false, // Extra keys that must never leak into the public payload. internalController: { secret: "controller-state" }, rawAuthorization: "Bearer raw-SHOULD-NOT-LEAK", @@ -387,11 +393,16 @@ test("buildHealthPayload projects allowlisted structural chatAdmission fields on waiting: 2, queuedBytes: 524_288, shedTotal: 3, - shedsByReason: { queue_timeout: 2, queued_bytes_budget: 1 }, + shedsByReason: { queue_timeout: 2, body_exceeds_budget: 1 }, lanes: [ { key: "key_c49d1c242feda590", waiting: 1 }, { key: "anonymous", waiting: 1 }, ], + inflightBytes: 131_072, + maxInflightBytes: 134_217_728, + budgetSource: "v8_heap", + pressureSeverity: "normal", + countCapEnabled: false, }); // The adaptive projection is untouched by the new key. assert.equal(payload.adaptiveAdmission, null); diff --git a/tests/unit/per-connection-admission-9654.test.ts b/tests/unit/per-connection-admission-9654.test.ts index 05af2deda6..9085163b41 100644 --- a/tests/unit/per-connection-admission-9654.test.ts +++ b/tests/unit/per-connection-admission-9654.test.ts @@ -170,9 +170,13 @@ test("admitChatRequest with explicit controller overrides per-connection lookup" }); test("admitChatStructure routes structural rejection to per-connection controller when heap pressure is genuinely high (#10183/#10268)", async () => { - // occupy sess-a's per-connection controller via the module-level instance + // occupy sess-a's per-connection controller via the module-level instance. + // #503-fanout: the production singleton's legacy count cap is unlimited by + // default (OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT unset in tests) — the real + // capacity dimension is now the auto-derived ingest byte budget, so + // "occupied" must exhaust that budget, not the (now unlimited) count. const controller = perConnectionAdmissionController.getController("sess-a"); - const occupied = controller.tryAcquireHeavy(); + const occupied = controller.tryAcquireBudget(controller.maxInflightBytes); assert.ok(occupied); const result = await admitChatStructure( @@ -199,9 +203,10 @@ test("admitChatStructure routes structural rejection to per-connection controlle }); test("admitChatStructure with different sessionId shares the global budget", async () => { - // occupy the shared process-global budget via sess-a + // occupy the shared process-global budget via sess-a (byte budget — see + // the #503-fanout comment in the previous test). const ctrlA = perConnectionAdmissionController.getController("sess-a"); - const occupied = ctrlA.tryAcquireHeavy(); + const occupied = ctrlA.tryAcquireBudget(ctrlA.maxInflightBytes); assert.ok(occupied); // Session B must NOT get independent capacity (pre-#10110 it did — that was