From 08a21bcf27a504c752b038d6d5c4f6d91e301bb3 Mon Sep 17 00:00:00 2001 From: Ravi Tharuma <25951435+RaviTharuma@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:59:20 +0200 Subject: [PATCH] fix(backend): add structure-aware chat admission (#8296) Co-authored-by: Ravi Tharuma Co-authored-by: Claude --- .env.example | 8 ++ docs/reference/ENVIRONMENT.md | 4 + src/app/api/v1/chat/completions/route.ts | 13 +- src/shared/middleware/chatBodyAdmission.ts | 139 +++++++++++++++++++ tests/unit/chat-body-admission.test.ts | 147 +++++++++++++++++++++ 5 files changed, 309 insertions(+), 2 deletions(-) diff --git a/.env.example b/.env.example index 25e68bddc8..e15f497d76 100644 --- a/.env.example +++ b/.env.example @@ -312,6 +312,14 @@ ALLOW_API_KEY_REVEAL=false # OMNIROUTE_CHAT_HARD_MAX_BODY_BYTES=52428800 # Maximum heavyweight requests simultaneously admitted in one process. Default 1. # OMNIROUTE_CHAT_MAX_HEAVY_IN_FLIGHT=1 +# Message count that classifies an otherwise small body as heavyweight. Default 200. +# OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT=200 +# Tool count that classifies an otherwise small body as heavyweight. Default 64. +# OMNIROUTE_CHAT_HEAVY_TOOL_COUNT=64 +# Conservative string-size token estimate that classifies a request as heavyweight. Default 32000. +# OMNIROUTE_CHAT_HEAVY_ESTIMATED_TOKENS=32000 +# Hard message-count cap; excess receives compact-required 413. Default 800. +# OMNIROUTE_CHAT_HARD_MAX_MESSAGES=800 # 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 diff --git a/docs/reference/ENVIRONMENT.md b/docs/reference/ENVIRONMENT.md index de60a4e8ec..446c7a6b9d 100644 --- a/docs/reference/ENVIRONMENT.md +++ b/docs/reference/ENVIRONMENT.md @@ -188,6 +188,10 @@ OmniRoute uses **SQLite** (via `better-sqlite3`) for all persistence. These vari | `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. When capacity is unavailable, OmniRoute returns retryable `503` with `Retry-After`. | +| `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. | +| `OMNIROUTE_CHAT_HEAVY_TOOL_COUNT` | `64` | `src/shared/middleware/chatBodyAdmission.ts` | Tool count that classifies a chat request as heavyweight even when its body is below the byte threshold. | +| `OMNIROUTE_CHAT_HEAVY_ESTIMATED_TOKENS` | `32000` | `src/shared/middleware/chatBodyAdmission.ts` | Conservative string-size token estimate that classifies a request as heavyweight; this is an admission-cost proxy, not provider billing tokenization. | +| `OMNIROUTE_CHAT_HARD_MAX_MESSAGES` | `800` | `src/shared/middleware/chatBodyAdmission.ts` | Hard chat history cap. Requests above it receive structured compact-required `413` before compression, translation, or provider dispatch. | | `OMNIROUTE_MAX_NONSTREAMING_RESPONSE_BYTES` | `67108864` (64 MB) | `open-sse/handlers/chatCore/nonStreamingResponseBody.ts` | Hard cap for a non-streaming upstream response buffered fully into memory. Past this the upstream reader is cancelled and the request fails fast instead of growing an unbounded string until the heap is exhausted. | | `CORS_ORIGIN` | _(unset)_ | `src/server/cors/origins.ts` | Legacy single-origin CORS allowlist. Prefer `CORS_ALLOWED_ORIGINS` for new deployments. CORS is only for cross-origin browser API clients; authenticated dashboard writes use same-origin requests plus session-bound CSRF protection instead. | | `CORS_ALLOWED_ORIGINS` | _(unset)_ | `src/server/cors/origins.ts` | Comma-separated CORS allowlist. No wildcard is sent unless `CORS_ALLOW_ALL=true` is explicitly configured. | diff --git a/src/app/api/v1/chat/completions/route.ts b/src/app/api/v1/chat/completions/route.ts index 59c1c22af7..7c281a6ff7 100644 --- a/src/app/api/v1/chat/completions/route.ts +++ b/src/app/api/v1/chat/completions/route.ts @@ -14,6 +14,7 @@ import { import { resolveKeepaliveThreshold } from "@omniroute/open-sse/utils/keepaliveThreshold"; import { admitChatRequest, + admitChatStructure, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, } from "@/shared/middleware/chatBodyAdmission"; @@ -73,8 +74,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 admission = await admitChatRequest(request); - if (!admission.admit) return admission.response; + const admissionResult = await admitChatRequest(request); + if (admissionResult.admit === false) return admissionResult.response; + const admission = admissionResult; request = admission.request; const finishAdmission = (response: Response) => releaseChatAdmissionWhenDone(response, admission.lease); @@ -101,6 +103,13 @@ export async function POST(request) { try { parsedBody = await request.json().catch(() => null); if (parsedBody) { + const structuralAdmission = admitChatStructure(parsedBody, admission.lease); + if (structuralAdmission.admit === false) { + admission.lease?.release(); + return structuralAdmission.response; + } + admission.lease = structuralAdmission.lease; + const { blocked, result } = injectionGuard(parsedBody); if (blocked) { return finishAdmission( diff --git a/src/shared/middleware/chatBodyAdmission.ts b/src/shared/middleware/chatBodyAdmission.ts index 3aed384749..57719d4e0b 100644 --- a/src/shared/middleware/chatBodyAdmission.ts +++ b/src/shared/middleware/chatBodyAdmission.ts @@ -30,6 +30,23 @@ const CHAT_MAX_HEAVY_IN_FLIGHT = parsePositiveInt( 1 ); +export const CHAT_HEAVY_MESSAGE_COUNT = parsePositiveInt( + process.env.OMNIROUTE_CHAT_HEAVY_MESSAGE_COUNT, + 200 +); +export const CHAT_HEAVY_TOOL_COUNT = parsePositiveInt( + process.env.OMNIROUTE_CHAT_HEAVY_TOOL_COUNT, + 64 +); +export const CHAT_HEAVY_ESTIMATED_TOKENS = parsePositiveInt( + process.env.OMNIROUTE_CHAT_HEAVY_ESTIMATED_TOKENS, + 32_000 +); +export const CHAT_HARD_MAX_MESSAGES = parsePositiveInt( + process.env.OMNIROUTE_CHAT_HARD_MAX_MESSAGES, + 800 +); + export interface ChatAdmissionLease { readonly released: boolean; release(): void; @@ -76,6 +93,10 @@ export type ChatRequestAdmission = | { admit: true; request: Request; lease: ChatAdmissionLease | null } | { admit: false; response: Response }; +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 = { @@ -99,6 +120,124 @@ function rejectionResponse(status: 413 | 503, hardMaxBytes: number): Response { ); } +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 }; +} + +export function admitChatStructure( + body: unknown, + lease: ChatAdmissionLease | null, + options: { + controller?: ChatAdmissionController; + maxMessages?: number; + heavyMessages?: number; + heavyTools?: number; + heavyTokens?: number; + } = {} +): ChatStructureAdmission { + if (!body || typeof body !== "object" || Array.isArray(body)) return { admit: true, lease }; + + const record = body as Record; + const messages = Array.isArray(record.messages) ? record.messages : []; + const tools = Array.isArray(record.tools) ? record.tools : []; + const maxMessages = options.maxMessages ?? CHAT_HARD_MAX_MESSAGES; + if (messages.length > maxMessages) { + return { admit: false, response: structuralRejectionResponse(413, maxMessages) }; + } + + const heavyMessages = options.heavyMessages ?? CHAT_HEAVY_MESSAGE_COUNT; + const heavyTools = options.heavyTools ?? CHAT_HEAVY_TOOL_COUNT; + const heavyTokens = options.heavyTokens ?? CHAT_HEAVY_ESTIMATED_TOKENS; + const countHeavy = messages.length >= heavyMessages || tools.length >= heavyTools; + if (!countHeavy && lease) return { admit: true, lease }; + + const messageEstimate = estimateStructureTokens(messages, heavyTokens); + const toolEstimate = messageEstimate.exhausted + ? { tokens: 0, exhausted: true } + : estimateStructureTokens(tools, heavyTokens - messageEstimate.tokens); + const estimatedTokens = Math.min(heavyTokens, messageEstimate.tokens + toolEstimate.tokens); + const heavy = + countHeavy || + messageEstimate.exhausted || + toolEstimate.exhausted || + estimatedTokens >= heavyTokens; + if (!heavy || lease) return { admit: true, lease }; + + const acquired = (options.controller ?? defaultAdmissionController).tryAcquireHeavy(); + return acquired + ? { admit: true, lease: acquired } + : { admit: false, response: structuralRejectionResponse(503, maxMessages) }; +} + function parseContentLength(header: string | null): number | null { if (header === null || !/^(0|[1-9]\d*)$/.test(header.trim())) return null; const parsed = Number(header); diff --git a/tests/unit/chat-body-admission.test.ts b/tests/unit/chat-body-admission.test.ts index ec5c551d3d..ee70ec1ded 100644 --- a/tests/unit/chat-body-admission.test.ts +++ b/tests/unit/chat-body-admission.test.ts @@ -4,6 +4,7 @@ import assert from "node:assert/strict"; const { admitChatRequest, + admitChatStructure, ChatAdmissionController, releaseChatAdmissionAfterHandler, releaseChatAdmissionWhenDone, @@ -32,6 +33,152 @@ test("small known body is admitted without consuming heavyweight capacity", asyn if (result.admit) assert.equal(await result.request.text(), "{}"); }); +test("a byte-light request above the message threshold acquires heavyweight capacity", async () => { + const controller = new ChatAdmissionController(1); + const result = admitChatStructure( + { messages: [{ role: "user", content: "one" }, { role: "user", content: "two" }] }, + null, + { controller, maxMessages: 10, heavyMessages: 2, heavyTools: 10, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); + assert.equal(controller.activeHeavy, 0); +}); + +test("a byte-light request above the tool threshold is rejected when heavy capacity is busy", async () => { + const controller = new ChatAdmissionController(1); + const occupied = controller.tryAcquireHeavy(); + assert.ok(occupied); + + const result = admitChatStructure( + { messages: [], tools: [{ type: "function" }, { type: "function" }] }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 2, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, false); + if (result.admit) return; + 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"); + occupied.release(); +}); + +test("a request above the hard history cap returns structured compact-required 413", async () => { + const controller = new ChatAdmissionController(1); + const result = admitChatStructure( + { messages: Array.from({ length: 3 }, () => ({ role: "user", content: "x" })) }, + null, + { controller, maxMessages: 2, heavyMessages: 1, heavyTools: 10, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, false); + if (result.admit) return; + assert.equal(result.response.status, 413); + const payload = await result.response.json(); + assert.equal(payload.error.code, "chat_history_too_large"); + assert.equal(payload.error.reason, "message_limit"); + assert.equal(controller.activeHeavy, 0); +}); + +test("a conservative token estimate classifies string messages and tool schemas as heavy", () => { + const controller = new ChatAdmissionController(1); + const result = admitChatStructure( + { + messages: [{ role: "user", content: "abcdefgh" }], + tools: [{ type: "function", function: { name: "tool", description: "abcdefgh" } }], + }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 10, heavyTokens: 4 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); +}); + +test("exhausting the bounded structural inspection is conservatively heavyweight", () => { + const controller = new ChatAdmissionController(1); + const result = admitChatStructure( + { + messages: [ + { + role: "user", + content: Array.from({ length: 10_001 }, () => ({ value: 0 })), + }, + ], + }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 10, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); +}); + +test("tool-schema property names contribute to the conservative token estimate", () => { + const controller = new ChatAdmissionController(1); + const properties = Object.fromEntries( + Array.from({ length: 5 }, (_, index) => [`${index}${"k".repeat(99)}`, {}]) + ); + const result = admitChatStructure( + { messages: [], tools: [{ function: { parameters: { properties } } }] }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 10, heavyTokens: 100 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); +}); + +test("non-ASCII strings use a conservative UTF-8 token estimate", () => { + const controller = new ChatAdmissionController(1); + const result = admitChatStructure( + { messages: [{ role: "user", content: "漢".repeat(100) }] }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 10, heavyTokens: 100 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); +}); + +test("wide objects exhaust bounded inspection without materializing all property values", () => { + const controller = new ChatAdmissionController(1); + const wide = Object.fromEntries(Array.from({ length: 10_001 }, (_, index) => [`k${index}`, 0])); + const result = admitChatStructure( + { messages: [{ role: "user", content: wide }] }, + null, + { controller, maxMessages: 10, heavyMessages: 10, heavyTools: 10, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) result.lease?.release(); +}); + +test("an existing byte-heavy lease is reused for structure-heavy admission", () => { + const controller = new ChatAdmissionController(1); + const lease = controller.tryAcquireHeavy(); + assert.ok(lease); + + const result = admitChatStructure( + { messages: [{ role: "user", content: "one" }, { role: "user", content: "two" }] }, + lease, + { controller, maxMessages: 10, heavyMessages: 2, heavyTools: 10, heavyTokens: 10_000 } + ); + + assert.equal(result.admit, true); + assert.equal(controller.activeHeavy, 1); + if (result.admit) assert.equal(result.lease, lease); + lease.release(); +}); + test("heavyweight admission is atomic and returns retryable 503 at capacity", async () => { const controller = new ChatAdmissionController(1); const body = JSON.stringify({ messages: [{ role: "user", content: "x".repeat(40) }] });