fix(mitm): bound SSE transcript retention and cancel abandoned upstream reads (#13702)

Handler-side collected strings grew without bound before the inspector
clamp; abandoned streams kept the reader alive for the full upstream
lifetime. createBoundedCollector caps retention at 1 MiB while keeping
true responseSize; pipeSSE and server.cjs cancel on downstream close.
Fixes #13395.

Co-authored-by: oyi77 <oyi77@users.noreply.github.com>
This commit is contained in:
Paijo
2026-09-18 22:25:45 +07:00
committed by GitHub
parent 994245a476
commit 2374bbf1ee
13 changed files with 334 additions and 83 deletions

View File

@@ -0,0 +1 @@
- **fix(mitm):** bound per-request SSE transcript retention to 1 MiB and stop the upstream read when the downstream disconnects — handler-side `collected` strings grew without bound before the inspector clamp, and abandoned streams kept the reader alive for the full upstream lifetime ([#13395](https://github.com/diegosouzapw/OmniRoute/issues/13395))

View File

@@ -21,7 +21,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
import { TOOL_RENAME_MAP } from "@omniroute/open-sse/services/claudeCodeToolRemapper";
interface GeminiPart {
@@ -171,7 +171,7 @@ export class AntigravityHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
let chunkStr = chunk.toString();
for (const [lower, capitalized] of Object.entries(TOOL_RENAME_MAP)) {
@@ -180,15 +180,15 @@ export class AntigravityHandler extends MitmHandlerBase {
`"name":"${capitalized}"`
);
}
collected += chunkStr;
sink.push(chunkStr);
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -76,6 +76,51 @@ async function loadAgentBridgeHook(): Promise<{
}
}
/**
* Ceiling for the per-request SSE text a MITM handler retains for the Traffic
* Inspector (#13395). The inspector buffer re-clamps per body
* (`INSPECTOR_MAX_BODY_KB`, default 1 MiB), but the handler-side `collected`
* string grew without bound BEFORE reaching that clamp — one long-lived stream
* kept the whole transcript in the handler closure for the request lifetime.
* Aligned with the inspector default so the bound never hides data the UI shows.
*/
export const MITM_PIPE_MAX_COLLECT_BYTES = 1 * 1024 * 1024;
/**
* Bounded string accumulator for piped SSE transcripts. Stops retaining past
* `maxBytes` but keeps counting true bytes, so the inspector still reports
* the exact `responseSize` it reported before (`Buffer.byteLength` of the
* full transcript) and the pipe itself is unaffected — every chunk is still
* written downstream regardless of the cap.
*/
export function createBoundedCollector(maxBytes: number = MITM_PIPE_MAX_COLLECT_BYTES): {
push: (chunk: string) => void;
text: string;
totalBytes: number;
truncated: boolean;
} {
let collected = "";
let totalBytes = 0;
const acc = {
push(chunk: string): void {
totalBytes += Buffer.byteLength(chunk);
if (collected.length < maxBytes) {
collected += chunk.slice(0, maxBytes - collected.length);
}
},
get text(): string {
return collected;
},
get totalBytes(): number {
return totalBytes;
},
get truncated(): boolean {
return totalBytes > Buffer.byteLength(collected);
},
};
return acc;
}
export abstract class MitmHandlerBase {
abstract readonly agentId: AgentId;
@@ -151,13 +196,6 @@ export abstract class MitmHandlerBase {
});
}
/**
* Pipe an SSE (or any chunked) upstream Response straight to the downstream
* ServerResponse, optionally invoking `onChunk` for each received Buffer.
*
* Writes SSE-friendly headers before the first chunk (only if `res.headersSent`
* is still false — handlers MAY have set custom headers first).
*/
protected async pipeSSE(
upstream: Response,
res: ServerResponse,
@@ -179,8 +217,18 @@ export abstract class MitmHandlerBase {
}
const reader = upstream.body.getReader();
// #13395: a downstream disconnect must stop the upstream read — otherwise an
// abandoned stream keeps the reader (and its buffers) alive for the full
// upstream lifetime and the handler closure retains the transcript.
let downstreamClosed = false;
const onClose = () => {
downstreamClosed = true;
reader.cancel().catch(() => {});
};
res.once("close", onClose);
try {
while (true) {
if (downstreamClosed) break;
const { done, value } = await reader.read();
if (done) break;
const buf = Buffer.from(value);
@@ -191,9 +239,16 @@ export abstract class MitmHandlerBase {
// Inspector hook must never break the upstream pipe.
}
}
if (downstreamClosed || res.closed || res.destroyed) break;
res.write(buf);
}
} finally {
res.off("close", onClose);
try {
reader.releaseLock();
} catch {
// Reader already cancelled or released.
}
try {
res.end();
} catch {

View File

@@ -8,7 +8,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class ClaudeCodeHandler extends MitmHandlerBase {
readonly agentId: AgentId = "claude-code";
@@ -51,17 +51,17 @@ export class ClaudeCodeHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -7,7 +7,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class CodexHandler extends MitmHandlerBase {
readonly agentId: AgentId = "codex";
@@ -33,17 +33,17 @@ export class CodexHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -7,7 +7,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class CopilotHandler extends MitmHandlerBase {
readonly agentId: AgentId = "copilot";
@@ -33,17 +33,17 @@ export class CopilotHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -7,7 +7,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class CursorHandler extends MitmHandlerBase {
readonly agentId: AgentId = "cursor";
@@ -33,17 +33,17 @@ export class CursorHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -10,7 +10,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class KiroHandler extends MitmHandlerBase {
readonly agentId: AgentId = "kiro";
@@ -36,17 +36,17 @@ export class KiroHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -7,7 +7,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class OpenCodeHandler extends MitmHandlerBase {
readonly agentId: AgentId = "open-code";
@@ -33,17 +33,17 @@ export class OpenCodeHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -7,7 +7,7 @@
*/
import type { IncomingMessage, ServerResponse } from "node:http";
import type { AgentId } from "../types";
import { MitmHandlerBase } from "./base";
import { MitmHandlerBase, createBoundedCollector } from "./base";
export class ZedHandler extends MitmHandlerBase {
readonly agentId: AgentId = "zed";
@@ -33,17 +33,17 @@ export class ZedHandler extends MitmHandlerBase {
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
}
let collected = "";
const sink = createBoundedCollector();
await this.pipeSSE(upstream, res, (chunk) => {
collected += chunk.toString();
sink.push(chunk.toString());
});
const total = this.now() - startedAt;
this.hookBufferUpdate(intercepted, {
status: upstream.status,
responseHeaders: Object.fromEntries(upstream.headers.entries()),
responseBody: collected,
responseSize: Buffer.byteLength(collected),
responseBody: sink.text,
responseSize: sink.totalBytes,
proxyLatencyMs: upstreamStart - startedAt,
upstreamLatencyMs: total - (upstreamStart - startedAt),
});

View File

@@ -542,47 +542,80 @@ async function intercept(req, res, bodyBuffer, override, sourceModel) {
vlog(1, `[MITM] → forward ${forward.format} ${forward.url}`);
upstreamStartedAt = Date.now();
const response = await fetch(forward.url, {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${API_KEY}`,
"x-omniroute-source": "agent-bridge",
"x-omniroute-agent": agentId,
},
body: JSON.stringify(body),
});
// #13395: an abandoned downstream must not keep the upstream fetch + reader
// alive. Abort the router fetch and cancel the reader on client close.
const upstreamAbort = new AbortController();
let downstreamClosed = false;
const onDownstreamClose = () => {
downstreamClosed = true;
try {
upstreamAbort.abort();
} catch {
// Abort is best-effort.
}
};
res.once("close", onDownstreamClose);
let reader = null;
try {
const response = await fetch(forward.url, {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${API_KEY}`,
"x-omniroute-source": "agent-bridge",
"x-omniroute-agent": agentId,
},
body: JSON.stringify(body),
signal: upstreamAbort.signal,
});
captureStatus = response.status;
respHeaders = headersToObject(response.headers);
captureStatus = response.status;
respHeaders = headersToObject(response.headers);
if (!response.ok) {
const errText = await response.text().catch(() => "");
respBody = errText.slice(0, INGEST_MAX_BODY);
respSize = Buffer.byteLength(errText);
throw new Error(`OmniRoute ${response.status}: ${errText}`);
}
if (!response.ok) {
const errText = await response.text().catch(() => "");
respBody = errText.slice(0, INGEST_MAX_BODY);
respSize = Buffer.byteLength(errText);
throw new Error(`OmniRoute ${response.status}: ${errText}`);
}
res.writeHead(200, {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
"X-Accel-Buffering": "no",
});
res.writeHead(200, {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
"X-Accel-Buffering": "no",
});
const reader = response.body.getReader();
const decoder = new TextDecoder();
reader = response.body.getReader();
const decoder = new TextDecoder();
while (true) {
const { done, value } = await reader.read();
if (done) {
res.end();
break;
while (true) {
if (downstreamClosed) break;
const { done, value } = await reader.read();
if (done) {
res.end();
break;
}
const text = decoder.decode(value, { stream: true });
if (respBody.length < INGEST_MAX_BODY) respBody += text;
respSize += value ? value.length : 0;
if (downstreamClosed || res.closed || res.destroyed) break;
res.write(text);
}
} finally {
res.off("close", onDownstreamClose);
if (reader) {
try {
await reader.cancel();
} catch {
// Reader already done or released.
}
try {
reader.releaseLock();
} catch {
// Already released.
}
}
const text = decoder.decode(value, { stream: true });
if (respBody.length < INGEST_MAX_BODY) respBody += text;
respSize += value ? value.length : 0;
res.write(text);
}
} catch (error) {
// Log the raw message locally (server console only) but never expose it

View File

@@ -44,10 +44,29 @@ function fakeRes(): { res: ServerResponse; out: HarnessResult } {
responseChunks: [],
};
let headersSent = false;
const listeners = new Map<string, Set<(...args: unknown[]) => void>>();
const res = {
get headersSent() {
return headersSent;
},
closed: false,
destroyed: false,
once(event: string, fn: (...args: unknown[]) => void) {
let set = listeners.get(event);
if (!set) {
set = new Set();
listeners.set(event, set);
}
set.add(fn);
return res;
},
off(event: string, fn: (...args: unknown[]) => void) {
listeners.get(event)?.delete(fn);
return res;
},
emitClose() {
for (const fn of [...(listeners.get("close") ?? [])]) fn();
},
writeHead(s: number) {
out.status = s;
headersSent = true;
@@ -59,7 +78,7 @@ function fakeRes(): { res: ServerResponse; out: HarnessResult } {
end(c?: Buffer | string) {
if (c) out.responseChunks.push(typeof c === "string" ? c : c.toString());
},
} as unknown as ServerResponse;
} as unknown as ServerResponse & { emitClose: () => void };
return { res, out };
}

View File

@@ -0,0 +1,143 @@
// #13395: MITM pipe paths must not retain unbounded transcripts, and an
// abandoned downstream must stop the upstream read.
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import path from "node:path";
import url from "node:url";
import type { ServerResponse } from "node:http";
import { Readable } from "node:stream";
import {
MitmHandlerBase,
createBoundedCollector,
MITM_PIPE_MAX_COLLECT_BYTES,
} from "../../src/mitm/handlers/base.ts";
import type { AgentId } from "../../src/mitm/types.ts";
class ExposedHandler extends MitmHandlerBase {
readonly agentId: AgentId = "antigravity";
async intercept(): Promise<void> {
throw new Error("not used");
}
pipe(
upstream: Response,
res: ServerResponse,
onChunk?: (c: Buffer) => void
): Promise<void> {
return this.pipeSSE(upstream, res, onChunk);
}
}
function trackingRes() {
const listeners = new Map<string, Set<(...a: unknown[]) => void>>();
const written: string[] = [];
let offCalls = 0;
const res = {
headersSent: false,
closed: false,
destroyed: false,
once(event: string, fn: (...a: unknown[]) => void) {
let s = listeners.get(event);
if (!s) {
s = new Set();
listeners.set(event, s);
}
s.add(fn);
return res;
},
off(event: string, fn: (...a: unknown[]) => void) {
offCalls += 1;
listeners.get(event)?.delete(fn);
return res;
},
emitClose() {
for (const fn of [...(listeners.get("close") ?? [])]) fn();
},
writeHead() {
(res as { headersSent: boolean }).headersSent = true;
},
write(c: Buffer | string) {
written.push(typeof c === "string" ? c : c.toString());
return true;
},
end() {},
} as unknown as ServerResponse;
return { res, written, listeners, offCalls: () => offCalls };
}
function chunkedUpstream(chunks: string[], delayMs = 0): Response {
const iterable = (async function* () {
for (const c of chunks) {
if (delayMs > 0) await new Promise((r) => setTimeout(r, delayMs));
yield Buffer.from(c);
}
})();
const stream = Readable.toWeb(Readable.from(iterable)) as unknown as ReadableStream<Uint8Array>;
return new Response(stream, { status: 200 });
}
test("bounded collector retains at most the cap but counts the true total", () => {
const sink = createBoundedCollector(100);
sink.push("a".repeat(60));
sink.push("b".repeat(60));
assert.equal(sink.text.length, 100);
assert.equal(sink.totalBytes, 120);
assert.equal(sink.truncated, true);
});
test("bounded collector passes small transcripts through untouched", () => {
const sink = createBoundedCollector(100);
sink.push("hello");
assert.equal(sink.text, "hello");
assert.equal(sink.totalBytes, 5);
assert.equal(sink.truncated, false);
});
test("default collect ceiling is 1 MiB", () => {
assert.equal(MITM_PIPE_MAX_COLLECT_BYTES, 1 * 1024 * 1024);
});
test("pipeSSE delivers every chunk when the downstream stays open", async () => {
const h = new ExposedHandler();
const { res, written } = trackingRes();
await h.pipe(chunkedUpstream(["x".repeat(10), "y".repeat(10), "z".repeat(10)]), res);
assert.equal(written.join(""), "x".repeat(10) + "y".repeat(10) + "z".repeat(10));
});
test("pipeSSE stops the upstream read after downstream close", async () => {
const h = new ExposedHandler();
const { res, written } = trackingRes();
const seen: string[] = [];
const pipe = h.pipe(
chunkedUpstream(Array.from({ length: 50 }, (_, i) => `c${i};`), 5),
res,
(c) => seen.push(c.toString())
);
// Let a few chunks flow, then abandon the downstream.
await new Promise((r) => setTimeout(r, 25));
(res as unknown as { emitClose: () => void }).emitClose();
await pipe;
const totalUpstream = Array.from({ length: 50 }, (_, i) => `c${i};`).join("");
assert.ok(
seen.join("").length < totalUpstream.length,
`abandoned pipe must stop early (read ${seen.join("").length} of ${totalUpstream.length})`
);
assert.ok(written.join("").length <= seen.join("").length, "no writes after close");
});
test("pipeSSE detaches its close listener when the stream completes", async () => {
const h = new ExposedHandler();
const t = trackingRes();
await h.pipe(chunkedUpstream(["done"]), t.res);
assert.equal(t.listeners.get("close")?.size ?? 0, 0, "close listener must be removed");
assert.ok(t.offCalls() >= 1, "off(close) must run");
});
test("server.cjs aborts the router fetch and cancels the reader on downstream close (#13395)", async () => {
const here = path.dirname(url.fileURLToPath(import.meta.url));
const src = fs.readFileSync(path.resolve(here, "../../src/mitm/server.cjs"), "utf8");
assert.match(src, /new AbortController\(\)/, "must create an abort controller per intercept");
assert.match(src, /signal:\s*upstreamAbort\.signal/, "router fetch must take the abort signal");
assert.match(src, /res\.once\(\s*"close"/, "must listen for downstream close");
assert.match(src, /reader\.cancel\(\)/, "must cancel the upstream reader on close");
});