Skip to content

Commit ba5cb68

Browse files
RhysSullivanRhysSullivan
committed
Close max-age SSE rotations gracefully instead of aborting
1 parent 9bbf5d4 commit ba5cb68

2 files changed

Lines changed: 82 additions & 28 deletions

File tree

‎packages/hosts/cloudflare/src/mcp/agents-sse-max-age.test.ts‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -55,20 +55,30 @@ const waitFor = async (predicate: () => boolean): Promise<void> => {
5555
expect(predicate()).toBe(true);
5656
};
5757

58-
const drainResponse = (response: Response): Promise<void> => {
59-
return Effect.runPromise(
58+
const drainResponse = async (response: Response): Promise<string> => {
59+
const decoder = new TextDecoder();
60+
let body = "";
61+
62+
await Effect.runPromise(
6063
Effect.ignore(
6164
Effect.tryPromise({
6265
try: () =>
6366
response.body?.pipeTo(
6467
new WritableStream<Uint8Array>({
65-
write: () => {},
68+
close: () => {
69+
body += decoder.decode();
70+
},
71+
write: (chunk) => {
72+
body += decoder.decode(chunk, { stream: true });
73+
},
6674
}),
6775
) ?? Promise.resolve(),
6876
catch: () => undefined,
6977
}),
7078
),
7179
);
80+
81+
return body;
7282
};
7383

7484
const installStallingTransformStream = () => {
@@ -233,14 +243,19 @@ const rotationLogs = (logs: ReadonlyArray<string>): ReadonlyArray<RotationLog> =
233243

234244
describe("agents SSE max-age rotation", () => {
235245
let errorLogs: string[] = [];
246+
let infoLogs: string[] = [];
236247

237248
beforeEach(() => {
238249
errorLogs = [];
250+
infoLogs = [];
239251
vi.useFakeTimers();
240252
vi.setSystemTime(0);
241253
vi.spyOn(console, "error").mockImplementation((line) => {
242254
errorLogs.push(String(line));
243255
});
256+
vi.spyOn(console, "log").mockImplementation((line) => {
257+
infoLogs.push(String(line));
258+
});
244259
});
245260

246261
afterEach(() => {
@@ -254,17 +269,18 @@ describe("agents SSE max-age rotation", () => {
254269
const drained = drainResponse(response);
255270

256271
await vi.advanceTimersByTimeAsync(MAX_SSE_AGE_MS + KEEPALIVE_INTERVAL_MS);
257-
await waitFor(() => ws.closeCode === 1013);
272+
await waitFor(() => ws.closeCode === 1000);
258273

259-
expect(ws.closeReason).toBe("SSE client not draining");
260-
const [rotationLog] = rotationLogs(errorLogs);
274+
expect(ws.closeReason).toBe("sse_max_age_rotation");
275+
const [rotationLog] = rotationLogs(infoLogs);
261276
expect(rotationLog?.event).toBe("sse_max_age_close");
262277
expect(rotationLog?.ageMs).toBeGreaterThan(MAX_SSE_AGE_MS);
263278
expect(rotationLog?.ageMs).toBeLessThanOrEqual(MAX_SSE_AGE_MS + KEEPALIVE_INTERVAL_MS);
264279
expect(rotationLog?.pendingBytes).toBeGreaterThanOrEqual(0);
280+
expect(errorLogs).toEqual([]);
265281
expect(vi.getTimerCount()).toBe(0);
266282

267-
await drained;
283+
await expect(drained).resolves.toContain(": max-age rotation, reconnect\n\n");
268284
});
269285

270286
it("does not rotate an in-flight POST response past max age", async () => {
@@ -277,7 +293,7 @@ describe("agents SSE max-age rotation", () => {
277293

278294
expect(ws.closeCode).toBeUndefined();
279295
expect(ws.closeReason).toBeUndefined();
280-
expect(rotationLogs(errorLogs)).toEqual([]);
296+
expect(rotationLogs(infoLogs)).toEqual([]);
281297

282298
emitAgentEvent(
283299
ws,
@@ -299,7 +315,7 @@ describe("agents SSE max-age rotation", () => {
299315
await flushMicrotasks();
300316

301317
expect(ws.closeCode).toBeUndefined();
302-
expect(rotationLogs(errorLogs)).toEqual([]);
318+
expect(rotationLogs(infoLogs)).toEqual([]);
303319

304320
emitClose(ws);
305321
await drained;
@@ -325,7 +341,7 @@ describe("agents SSE max-age rotation", () => {
325341
expect(ws.closeCode).toBe(1013);
326342
expect(ws.closeReason).toBe("SSE client not draining");
327343
expect(transform.abortReason()).toBeInstanceOf(Error);
328-
expect(rotationLogs(errorLogs)).toEqual([]);
344+
expect(rotationLogs(infoLogs)).toEqual([]);
329345
expect(vi.getTimerCount()).toBe(0);
330346
});
331347

@@ -338,7 +354,7 @@ describe("agents SSE max-age rotation", () => {
338354
await drained;
339355

340356
expect(ws.closeCode).toBeUndefined();
341-
expect(rotationLogs(errorLogs)).toEqual([]);
357+
expect(rotationLogs(infoLogs)).toEqual([]);
342358
expect(vi.getTimerCount()).toBe(0);
343359
});
344360
});

‎patches/agents@0.17.3.patch‎

Lines changed: 55 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
diff --git a/dist/mcp/index.d.ts b/dist/mcp/index.d.ts
2-
index c8fad44..439da01 100644
2+
index c8fad448e8797b89690a99d93490d1363851b225..439da01e789713d217778e8bdebfa1a9d52a71e2 100644
33
--- a/dist/mcp/index.d.ts
44
+++ b/dist/mcp/index.d.ts
55
@@ -29,6 +29,7 @@ import {
@@ -19,7 +19,7 @@ index c8fad44..439da01 100644
1919
McpAgent,
2020
type McpAuthContext,
2121
diff --git a/dist/mcp/index.js b/dist/mcp/index.js
22-
index 1edcf0c..707f1cf 100644
22+
index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..0ba4f5bff8e6a968aa18c609952e576cedf13d8f 100644
2323
--- a/dist/mcp/index.js
2424
+++ b/dist/mcp/index.js
2525
@@ -28,13 +28,14 @@ import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/
@@ -178,7 +178,7 @@ index 1edcf0c..707f1cf 100644
178178
props: ctx.props,
179179
jurisdiction: options.jurisdiction
180180
});
181-
@@ -306,27 +344,66 @@ const createStreamingHttpHandler = (basePath, namespace, options = {}) => {
181+
@@ -306,27 +344,86 @@ const createStreamingHttpHandler = (basePath, namespace, options = {}) => {
182182
if (!ws) {
183183
await writer.close();
184184
return new Response("Failed to establish WS to DO", { status: 500 });
@@ -187,7 +187,7 @@ index 1edcf0c..707f1cf 100644
187187
- ws.addEventListener("message", (event) => {
188188
+ }
189189
+ ws.accept();
190-
+ const __closeSse = () => {
190+
+ const __abortSse = () => {
191191
+ __sseClosed = true;
192192
+ try {
193193
+ clearInterval(keepAlive);
@@ -197,22 +197,42 @@ index 1edcf0c..707f1cf 100644
197197
+ } catch {}
198198
+ writer.abort(new Error("SSE client not draining")).catch(() => {});
199199
+ };
200+
+ const __closeSseGracefully = () => {
201+
+ __sseClosed = true;
202+
+ try {
203+
+ clearInterval(keepAlive);
204+
+ } catch {}
205+
+ try {
206+
+ const finalFrame = encoder.encode(": max-age rotation, reconnect\n\n");
207+
+ __writeChain = __writeChain.then(() => writer.write(finalFrame)).catch(() => {}).then(() => writer.close()).catch(() => {
208+
+ writer.abort().catch(() => {});
209+
+ });
210+
+ } catch {
211+
+ __writeChain = writer.close().catch(() => {
212+
+ writer.abort().catch(() => {});
213+
+ });
214+
+ }
215+
+ try {
216+
+ ws.close(1000, "sse_max_age_rotation");
217+
+ } catch {}
218+
+ return __writeChain;
219+
+ };
200220
+ const __forwardSse = (frame) => {
201221
+ if (__sseClosed) return __writeChain;
202222
+ const ageMs = Date.now() - __openedAt;
203223
+ if (ageMs > MAX_SSE_AGE_MS) {
204-
+ console.error(JSON.stringify({
224+
+ console.log(JSON.stringify({
205225
+ event: "sse_max_age_close",
206226
+ sessionId,
207227
+ variant: "streamable-get",
208228
+ ageMs,
209229
+ pendingBytes: __pendingBytes
210230
+ }));
211-
+ __closeSse();
231+
+ __closeSseGracefully();
212232
+ return Promise.resolve();
213233
+ }
214234
+ if (__pendingBytes + frame.byteLength > MAX_PENDING_SSE_BYTES) {
215-
+ __closeSse();
235+
+ __abortSse();
216236
+ return Promise.resolve();
217237
+ }
218238
+ __pendingBytes += frame.byteLength;
@@ -260,7 +280,7 @@ index 1edcf0c..707f1cf 100644
260280
return new Response(readable, {
261281
headers: {
262282
"Cache-Control": "no-cache",
263-
@@ -389,10 +466,16 @@ const createLegacySseHandler = (basePath, namespace, options = {}) => {
283+
@@ -389,10 +486,16 @@ const createLegacySseHandler = (basePath, namespace, options = {}) => {
264284
const url = new URL(request.url);
265285
if (request.method === "GET" && basePattern.test(url)) {
266286
const sessionId = url.searchParams.get("sessionId") || namespace.newUniqueId().toString();
@@ -281,7 +301,7 @@ index 1edcf0c..707f1cf 100644
281301
endpointUrl.pathname = encodeURI(`${basePath}/message`);
282302
endpointUrl.searchParams.set("sessionId", sessionId);
283303
const endpointMessage = `event: endpoint\ndata: ${endpointUrl.pathname + endpointUrl.search + endpointUrl.hash}\n\n`;
284-
@@ -414,35 +497,74 @@ const createLegacySseHandler = (basePath, namespace, options = {}) => {
304+
@@ -414,35 +517,94 @@ const createLegacySseHandler = (basePath, namespace, options = {}) => {
285305
console.error("Failed to establish WebSocket connection");
286306
await writer.close();
287307
return new Response("Failed to establish WebSocket connection", { status: 500 });
@@ -290,7 +310,7 @@ index 1edcf0c..707f1cf 100644
290310
- ws.addEventListener("message", (event) => {
291311
+ }
292312
+ ws.accept();
293-
+ const __closeSse = () => {
313+
+ const __abortSse = () => {
294314
+ __sseClosed = true;
295315
+ try {
296316
+ clearInterval(keepAlive);
@@ -300,22 +320,42 @@ index 1edcf0c..707f1cf 100644
300320
+ } catch {}
301321
+ writer.abort(new Error("SSE client not draining")).catch(() => {});
302322
+ };
323+
+ const __closeSseGracefully = () => {
324+
+ __sseClosed = true;
325+
+ try {
326+
+ clearInterval(keepAlive);
327+
+ } catch {}
328+
+ try {
329+
+ const finalFrame = encoder.encode(": max-age rotation, reconnect\n\n");
330+
+ __writeChain = __writeChain.then(() => writer.write(finalFrame)).catch(() => {}).then(() => writer.close()).catch(() => {
331+
+ writer.abort().catch(() => {});
332+
+ });
333+
+ } catch {
334+
+ __writeChain = writer.close().catch(() => {
335+
+ writer.abort().catch(() => {});
336+
+ });
337+
+ }
338+
+ try {
339+
+ ws.close(1000, "sse_max_age_rotation");
340+
+ } catch {}
341+
+ return __writeChain;
342+
+ };
303343
+ const __forwardSse = (frame) => {
304344
+ if (__sseClosed) return __writeChain;
305345
+ const ageMs = Date.now() - __openedAt;
306346
+ if (ageMs > MAX_SSE_AGE_MS) {
307-
+ console.error(JSON.stringify({
347+
+ console.log(JSON.stringify({
308348
+ event: "sse_max_age_close",
309349
+ sessionId,
310350
+ variant: "legacy-sse",
311351
+ ageMs,
312352
+ pendingBytes: __pendingBytes
313353
+ }));
314-
+ __closeSse();
354+
+ __closeSseGracefully();
315355
+ return Promise.resolve();
316356
+ }
317357
+ if (__pendingBytes + frame.byteLength > MAX_PENDING_SSE_BYTES) {
318-
+ __closeSse();
358+
+ __abortSse();
319359
+ return Promise.resolve();
320360
+ }
321361
+ __pendingBytes += frame.byteLength;
@@ -374,13 +414,11 @@ index 1edcf0c..707f1cf 100644
374414
console.error("Error closing SSE connection:", error);
375415
}
376416
}
377-
@@ -1698,6 +1820,6 @@ var McpAgent = class McpAgent extends Agent {
417+
@@ -1698,6 +1860,6 @@ var McpAgent = class McpAgent extends Agent {
378418
};
379419
McpAgent.STREAM_REQS_KEY_PREFIX = "__mcp_stream_reqs__:";
380420
//#endregion
381421
-export { DurableObjectEventStore, ElicitRequestSchema, MCP_SERVER_ID_MAX_LENGTH, McpAgent, RPCClientTransport, RPCServerTransport, RPC_DO_PREFIX, SSEEdgeClientTransport, StreamableHTTPEdgeClientTransport, WorkerTransport, createMcpHandler, experimental_createMcpHandler, getMcpAuthContext, normalizeServerId };
382422
+export { DurableObjectEventStore, ElicitRequestSchema, MAX_SSE_AGE_MS, MCP_SERVER_ID_MAX_LENGTH, McpAgent, RPCClientTransport, RPCServerTransport, RPC_DO_PREFIX, SSEEdgeClientTransport, StreamableHTTPEdgeClientTransport, WorkerTransport, createMcpHandler, experimental_createMcpHandler, getMcpAuthContext, normalizeServerId };
383423

384-
-//# sourceMappingURL=index.js.map
385-
\ No newline at end of file
386-
+//# sourceMappingURL=index.js.map
424+
//# sourceMappingURL=index.js.map

0 commit comments

Comments
 (0)