From b4c47f5cbf0f63137b5df03f777b980d14e04ff0 Mon Sep 17 00:00:00 2001
From: Paijo <14921983+oyi77@users.noreply.github.com>
Date: Sun, 12 Jul 2026 20:50:27 +0700
Subject: [PATCH] feat(sse): add connection backpressure for chat handler
 (#6590)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add checkConnectionCapacity guard with 429 + Retry-After in handleChat().
Introduce OMNI_MAX_CONCURRENT_CONNECTIONS env-bound cap, disabled (0) by
default so existing deployments are unaffected until an operator opts in.

Reconstructed from PR #6590, isolating only the backpressure change —
the original branch also carried unrelated headroom/docker/perf work
from the author's separate #6572 branch.

Co-authored-by: oyi77 <oyi77@users.noreply.github.com>
Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>
---
 src/sse/handlers/chat.ts        |   7 ++
 src/sse/utils/backpressure.ts   |  67 ++++++++++++++++++++
 tests/unit/backpressure.test.ts | 109 ++++++++++++++++++++++++++++++++
 3 files changed, 183 insertions(+)
 create mode 100644 src/sse/utils/backpressure.ts
 create mode 100644 tests/unit/backpressure.test.ts

diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts
index 9609173633..0bee6caa42 100644
--- a/src/sse/handlers/chat.ts
+++ b/src/sse/handlers/chat.ts
@@ -132,6 +132,7 @@ import {
   waitForCooldownAwareRetry,
 } from "../services/cooldownAwareRetry";
 import { constrainConnectionsToQuota, resolveQuotaKeyScope } from "../../lib/quota/quotaKey";
+import { checkConnectionCapacity } from "../utils/backpressure";
 
 registerCodexQuotaFetcher();
 
@@ -222,6 +223,12 @@ export async function handleChat(
   const reqId = correlationId || generateRequestId();
   const telemetry = new RequestTelemetry(reqId);
 
+  const backpressure = checkConnectionCapacity();
+  if (backpressure.shouldReject) {
+    log.warn("BACKPRESSURE", "Rejecting request: at connection limit");
+    return backpressure.response;
+  }
+
   let body;
   try {
     telemetry.startPhase("parse");
diff --git a/src/sse/utils/backpressure.ts b/src/sse/utils/backpressure.ts
new file mode 100644
index 0000000000..f5b56f19d5
--- /dev/null
+++ b/src/sse/utils/backpressure.ts
@@ -0,0 +1,67 @@
+import { getActiveSessionCount } from "@omniroute/open-sse/services/sessionManager.ts";
+
+/**
+ * Connection back-pressure for SSE / streaming endpoints.
+ *
+ * Caps in-flight requests to prevent memory exhaustion.  Reads the live
+ * session count from `sessionManager` — the same source the health endpoint
+ * uses to report `activeConnections`.
+ *
+ * Set `OMNI_MAX_CONCURRENT_CONNECTIONS` to a positive integer to enable.
+ * Default is 0 (disabled) so existing deployments are unaffected until
+ * an operator explicitly opts in.
+ */
+
+type CapacityOk = { shouldReject: false };
+type CapacityExceeded = { shouldReject: true; response: Response };
+export type CapacityResult = CapacityOk | CapacityExceeded;
+
+/**
+ * Pure capacity evaluator — given an active count and a cap, returns the
+ * appropriate result.  Cap ≤ 0 means disabled (never reject).
+ *
+ * Exported for unit testing without mocking any module boundaries.
+ */
+export function evalCapacity(active: number, cap: number): CapacityResult {
+  if (cap <= 0 || active < cap) {
+    return { shouldReject: false };
+  }
+  const retryAfter = Math.max(1, Math.ceil((active / cap) * 30));
+  return {
+    shouldReject: true,
+    response: new Response(
+      JSON.stringify({
+        error: {
+          message: `Server busy — ${active} active connections (limit ${cap}). Retry after ${retryAfter}s.`,
+          type: "rate_limit",
+          retry_after: retryAfter,
+        },
+      }),
+      {
+        status: 429,
+        headers: {
+          "Content-Type": "application/json",
+          "Retry-After": String(retryAfter),
+          "X-RateLimit-Limit": String(cap),
+          "X-RateLimit-Remaining": "0",
+        },
+      }
+    ),
+  };
+}
+
+/** Read cap from env on every call so tests can mutate process.env freely. */
+function readCap(): number {
+  const raw = parseInt(process.env.OMNI_MAX_CONCURRENT_CONNECTIONS ?? "0", 10);
+  return Number.isFinite(raw) && raw > 0 ? raw : 0;
+}
+
+/**
+ * Check whether the server is over its configured connection cap.
+ *
+ * Returns `{ shouldReject: true, response }` when at limit, or
+ * `{ shouldReject: false }` when a slot is free (or the cap is disabled).
+ */
+export function checkConnectionCapacity(): CapacityResult {
+  return evalCapacity(getActiveSessionCount(), readCap());
+}
diff --git a/tests/unit/backpressure.test.ts b/tests/unit/backpressure.test.ts
new file mode 100644
index 0000000000..14ed4fe098
--- /dev/null
+++ b/tests/unit/backpressure.test.ts
@@ -0,0 +1,109 @@
+import { test, beforeEach } from "node:test";
+import assert from "node:assert/strict";
+import { evalCapacity } from "../../src/sse/utils/backpressure.ts";
+import type { CapacityResult } from "../../src/sse/utils/backpressure.ts";
+
+// ---------------------------------------------------------------------------
+// evalCapacity — pure function, no mocking required
+// ---------------------------------------------------------------------------
+
+beforeEach(() => {
+  delete process.env.OMNI_MAX_CONCURRENT_CONNECTIONS;
+});
+
+test("backpressure: cap=0 never rejects regardless of active count", () => {
+  const r: CapacityResult = evalCapacity(9999, 0);
+  assert.equal(r.shouldReject, false);
+});
+
+test("backpressure: negative cap never rejects", () => {
+  assert.equal(evalCapacity(100, -1).shouldReject, false);
+});
+
+test("backpressure: passes when active is zero", () => {
+  assert.equal(evalCapacity(0, 10).shouldReject, false);
+});
+
+test("backpressure: passes when active is below cap", () => {
+  assert.equal(evalCapacity(5, 10).shouldReject, false);
+});
+
+test("backpressure: passes when active is exactly cap minus one", () => {
+  assert.equal(evalCapacity(9, 10).shouldReject, false);
+});
+
+test("backpressure: rejects when active equals cap", () => {
+  const r = evalCapacity(10, 10);
+  assert.equal(r.shouldReject, true);
+});
+
+test("backpressure: rejects when active exceeds cap", () => {
+  assert.equal(evalCapacity(100, 10).shouldReject, true);
+});
+
+test("backpressure: rejection returns HTTP 429", async () => {
+  const r = evalCapacity(10, 10);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  assert.equal(r.response.status, 429);
+});
+
+test("backpressure: rejection includes Retry-After header >= 1", async () => {
+  const r = evalCapacity(10, 10);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  const val = Number(r.response.headers.get("Retry-After"));
+  assert.ok(val >= 1, `Retry-After must be >= 1, got ${val}`);
+});
+
+test("backpressure: X-RateLimit-Limit matches the cap", async () => {
+  const r = evalCapacity(7, 5);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  assert.equal(r.response.headers.get("X-RateLimit-Limit"), "5");
+  assert.equal(r.response.headers.get("X-RateLimit-Remaining"), "0");
+});
+
+test("backpressure: rejection body is valid JSON with error.type=rate_limit", async () => {
+  const r = evalCapacity(3, 3);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  const body = await r.response.json();
+  assert.equal(body?.error?.type, "rate_limit");
+  assert.ok(typeof body.error.message === "string");
+  assert.ok(typeof body.error.retry_after === "number");
+});
+
+test("backpressure: rejection message contains active and limit counts", async () => {
+  const r = evalCapacity(8, 5);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  const body = await r.response.json();
+  assert.ok(body.error.message.includes("8"), "message should reference active count");
+  assert.ok(body.error.message.includes("5"), "message should reference cap");
+});
+
+test("backpressure: retry-after formula — at cap gives ceil(1*30)=30", async () => {
+  const r = evalCapacity(10, 10);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  assert.equal(Number(r.response.headers.get("Retry-After")), 30);
+});
+
+test("backpressure: retry-after increases proportionally under higher load", async () => {
+  const r1 = evalCapacity(10, 10); // factor 1.0 → 30s
+  const r2 = evalCapacity(20, 10); // factor 2.0 → 60s
+  assert.equal(r1.shouldReject, true);
+  assert.equal(r2.shouldReject, true);
+  if (!r1.shouldReject || !r2.shouldReject) throw new Error("unreachable");
+  const ra1 = Number(r1.response.headers.get("Retry-After"));
+  const ra2 = Number(r2.response.headers.get("Retry-After"));
+  assert.ok(ra2 > ra1, `retry-after should grow with load factor (${ra1} → ${ra2})`);
+});
+
+test("backpressure: Content-Type is application/json", async () => {
+  const r = evalCapacity(5, 5);
+  assert.equal(r.shouldReject, true);
+  if (!r.shouldReject) throw new Error("unreachable");
+  assert.ok(r.response.headers.get("Content-Type")?.includes("application/json"));
+});