Skip to content

Commit 1bccd29

Browse files
wangmiao0668000666claudevincentkoc
authored
fix(provider-transport-fetch): bound SSE buffer to prevent OOM (#96989)
* fix(provider-transport-fetch): bound SSE buffer to prevent OOM * fix(provider-transport-fetch): appease oxlint curly rule in test * fix(provider-transport-fetch): drain events before cap + cancel reader on overflow * fix(provider-transport-fetch): remove unused encoder from coalesced chunk test Co-Authored-By: Claude <[email protected]> * fix(transport): tighten SSE buffer guards --------- Co-authored-by: Claude <[email protected]> Co-authored-by: Vincent Koc <[email protected]>
1 parent 4985671 commit 1bccd29

2 files changed

Lines changed: 158 additions & 0 deletions

File tree

‎src/agents/provider-transport-fetch.test.ts‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,6 +1130,37 @@ describe("buildGuardedModelFetch", () => {
11301130
expect(items).toEqual([{ ok: true }]);
11311131
});
11321132

1133+
it("handles a large transport chunk containing many valid small SSE events", async () => {
1134+
// Regression: one TCP read can deliver >64 KiB of already-delimited SSE
1135+
// events; the cap must apply only to the unterminated tail, not the full chunk.
1136+
const eventCount = 5_000;
1137+
const manyEvents = `data: ${JSON.stringify({ ok: true })}\n\n`.repeat(eventCount);
1138+
fetchWithSsrFGuardMock.mockResolvedValue({
1139+
response: new Response(manyEvents, {
1140+
headers: { "content-type": "text/event-stream" },
1141+
}),
1142+
finalUrl: "https://openrouter.ai/api/v1/chat/completions",
1143+
release: vi.fn(async () => undefined),
1144+
});
1145+
const model = {
1146+
id: "gpt-5.4",
1147+
provider: "openrouter",
1148+
api: "openai-completions",
1149+
baseUrl: "https://openrouter.ai/api/v1",
1150+
} as unknown as Model<"openai-completions">;
1151+
1152+
const response = await buildGuardedModelFetch(model)(
1153+
"https://openrouter.ai/api/v1/chat/completions",
1154+
{ method: "POST" },
1155+
);
1156+
const items: unknown[] = [];
1157+
for await (const item of Stream.fromSSEResponse(response, new AbortController())) {
1158+
items.push(item);
1159+
}
1160+
expect(items.length).toBe(eventCount);
1161+
expect(items[0]).toEqual({ ok: true });
1162+
});
1163+
11331164
it("synthesizes SSE frames for JSON bodies returned to streaming OpenAI SDK requests", async () => {
11341165
fetchWithSsrFGuardMock.mockResolvedValue({
11351166
response: new Response(' {"ok": true} ', {
@@ -1338,6 +1369,102 @@ describe("buildGuardedModelFetch", () => {
13381369
expect(refreshTimeout).toHaveBeenCalledTimes(2);
13391370
});
13401371

1372+
it("errors on oversized SSE body without event boundary in sanitizer", async () => {
1373+
const oversized = "x".repeat(65 * 1024);
1374+
const encoder = new TextEncoder();
1375+
fetchWithSsrFGuardMock.mockResolvedValue({
1376+
response: new Response(
1377+
new ReadableStream({
1378+
start(controller) {
1379+
controller.enqueue(encoder.encode(oversized));
1380+
controller.close();
1381+
},
1382+
}),
1383+
{ headers: { "content-type": "text/event-stream" } },
1384+
),
1385+
finalUrl: "https://openrouter.ai/api/v1/chat/completions",
1386+
release: vi.fn(async () => undefined),
1387+
});
1388+
const model = {
1389+
id: "gpt-5.4",
1390+
provider: "openrouter",
1391+
api: "openai-completions",
1392+
baseUrl: "https://openrouter.ai/api/v1",
1393+
} as unknown as Model<"openai-completions">;
1394+
1395+
const response = await buildGuardedModelFetch(model)(
1396+
"https://openrouter.ai/api/v1/chat/completions",
1397+
{ method: "POST" },
1398+
);
1399+
1400+
const reader = response.body?.getReader();
1401+
let caught: unknown = null;
1402+
try {
1403+
while (true) {
1404+
const { done } = await reader!.read();
1405+
if (done) {
1406+
break;
1407+
}
1408+
}
1409+
} catch (e) {
1410+
caught = e;
1411+
}
1412+
expect(caught).toBeTruthy();
1413+
expect(String(caught)).toMatch(/exceeded max buffer size/i);
1414+
});
1415+
1416+
it("errors on oversized streaming JSON body without content-length in SSE synthesis", async () => {
1417+
const CHUNK = 1024 * 1024;
1418+
let sends = 0;
1419+
fetchWithSsrFGuardMock.mockResolvedValue({
1420+
response: new Response(
1421+
new ReadableStream({
1422+
pull(controller) {
1423+
if (sends < 17) {
1424+
sends++;
1425+
controller.enqueue(new Uint8Array(CHUNK));
1426+
} else {
1427+
controller.close();
1428+
}
1429+
},
1430+
}),
1431+
{ headers: { "content-type": "application/json" } },
1432+
),
1433+
finalUrl: "https://openrouter.ai/api/v1/chat/completions",
1434+
release: vi.fn(async () => undefined),
1435+
});
1436+
const model = {
1437+
id: "moonshotai/kimi-k2.6",
1438+
provider: "openrouter",
1439+
api: "openai-completions",
1440+
baseUrl: "https://openrouter.ai/api/v1",
1441+
} as unknown as Model<"openai-completions">;
1442+
1443+
const response = await buildGuardedModelFetch(model)(
1444+
"https://openrouter.ai/api/v1/chat/completions",
1445+
{
1446+
method: "POST",
1447+
headers: { "content-type": "application/json" },
1448+
body: JSON.stringify({ model: "moonshotai/kimi-k2.6", stream: true }),
1449+
},
1450+
);
1451+
1452+
const reader = response.body?.getReader();
1453+
let caught: unknown = null;
1454+
try {
1455+
while (true) {
1456+
const { done } = await reader!.read();
1457+
if (done) {
1458+
break;
1459+
}
1460+
}
1461+
} catch (e) {
1462+
caught = e;
1463+
}
1464+
expect(caught).toBeTruthy();
1465+
expect(String(caught)).toMatch(/exceeded.*bytes while synthesizing SSE/i);
1466+
});
1467+
13411468
describe("long retry-after handling", () => {
13421469
const anthropicModel = {
13431470
id: "sonnet-4.6",

‎src/agents/provider-transport-fetch.ts‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,17 @@ import {
4545
const DEFAULT_MAX_SDK_RETRY_WAIT_SECONDS = 60;
4646
const OPENAI_SDK_STREAM_CONTENT_SNIFF_BYTES = 2 * 1024;
4747
const log = createSubsystemLogger("provider-transport-fetch");
48+
49+
/** Max bytes for an entire JSON body synthesized into SSE frames. Prevents OOM
50+
* when a hostile streaming endpoint returns a never-ending JSON response
51+
* without Content-Length. */
52+
const SSE_SYNTHESIZE_JSON_MAX_BYTES = 16 * 1024 * 1024;
53+
54+
/** Max bytes for the internal SSE sanitization buffer between event boundaries.
55+
* A response that cannot find a \n\n boundary within this many characters is
56+
* almost certainly hostile or broken — cap the buffer rather than let it grow. */
57+
const SSE_SANITIZE_BUFFER_MAX_BYTES = 64 * 1024;
58+
4859
const BLOCKED_EXACT_ORIGIN_TRUST_HOSTNAME_LABELS = new Set(["instance-data"]);
4960
const PLAIN_DECIMAL_NUMBER_RE = /^\d+(?:\.\d+)?$/;
5061
const RETRY_AFTER_HTTP_DATE_RE =
@@ -102,6 +113,7 @@ function sanitizeOpenAISdkSseResponse(
102113
const encoder = new TextEncoder();
103114
let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
104115
let buffer = "";
116+
let totalBytes = 0;
105117
const sseBody = new ReadableStream<Uint8Array>({
106118
start() {
107119
reader = source.getReader();
@@ -120,9 +132,17 @@ function sanitizeOpenAISdkSseResponse(
120132
controller.close();
121133
return;
122134
}
135+
const nextTotalBytes = totalBytes + chunk.value.byteLength;
136+
if (nextTotalBytes > SSE_SYNTHESIZE_JSON_MAX_BYTES) {
137+
throw new Error(
138+
`Streaming JSON body exceeded ${SSE_SYNTHESIZE_JSON_MAX_BYTES} bytes while synthesizing SSE frames`,
139+
);
140+
}
141+
totalBytes = nextTotalBytes;
123142
buffer += decoder.decode(chunk.value, { stream: true });
124143
}
125144
} catch (error) {
145+
await reader?.cancel(error).catch(() => {});
126146
controller.error(error);
127147
}
128148
},
@@ -157,6 +177,11 @@ function sanitizeOpenAISdkSseResponse(
157177
for (;;) {
158178
const boundary = findSseEventBoundary(buffer);
159179
if (!boundary) {
180+
if (buffer.length > SSE_SANITIZE_BUFFER_MAX_BYTES) {
181+
throw new Error(
182+
`SSE response exceeded max buffer size (${SSE_SANITIZE_BUFFER_MAX_BYTES} bytes) without event boundary`,
183+
);
184+
}
160185
return enqueued;
161186
}
162187
const block = buffer.slice(0, boundary.index);
@@ -167,6 +192,7 @@ function sanitizeOpenAISdkSseResponse(
167192
if (hasReadableSseData(block)) {
168193
controller.enqueue(encoder.encode(`${block}${separator}`));
169194
enqueued += 1;
195+
return enqueued;
170196
}
171197
}
172198
};
@@ -178,6 +204,10 @@ function sanitizeOpenAISdkSseResponse(
178204
async pull(controller) {
179205
try {
180206
for (;;) {
207+
const pending = enqueueSanitized(controller, "");
208+
if (pending > 0) {
209+
return;
210+
}
181211
const chunk = await reader?.read();
182212
if (!chunk || chunk.done) {
183213
const tail = decoder.decode();
@@ -200,6 +230,7 @@ function sanitizeOpenAISdkSseResponse(
200230
}
201231
}
202232
} catch (error) {
233+
await reader?.cancel(error).catch(() => {});
203234
controller.error(error);
204235
}
205236
},

0 commit comments

Comments
 (0)