fix(backend): add structure-aware chat admission (#8296)

Co-authored-by: Ravi Tharuma <RaviTharuma@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

.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

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.                                                                                                                                                                                                                                                         |

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(

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<string, string> = {
@@ -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<string, string> = {
+    ...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<string, unknown>)[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<string, unknown>;
+  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);

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) }] });