From 6d55ca851ecd1d5d3e9d3f82140e7d79b152178c Mon Sep 17 00:00:00 2001 From: Jad Date: Tue, 4 Aug 2026 15:26:44 +0200 Subject: [PATCH] fix: scope MCP replay to the requested stream --- .changeset/scope-mcp-stream-replay.md | 5 ++ e2e/cloud/mcp-sse-replay.test.ts | 120 ++++++++++++++++++++++++-- patches/agents@0.17.3.patch | 77 ++++++++--------- 3 files changed, 158 insertions(+), 44 deletions(-) create mode 100644 .changeset/scope-mcp-stream-replay.md diff --git a/.changeset/scope-mcp-stream-replay.md b/.changeset/scope-mcp-stream-replay.md new file mode 100644 index 000000000..0374d7460 --- /dev/null +++ b/.changeset/scope-mcp-stream-replay.md @@ -0,0 +1,5 @@ +--- +"executor": patch +--- + +Keep Last-Event-ID recovery scoped to its originating MCP stream. diff --git a/e2e/cloud/mcp-sse-replay.test.ts b/e2e/cloud/mcp-sse-replay.test.ts index 3ab672296..a0df02bf6 100644 --- a/e2e/cloud/mcp-sse-replay.test.ts +++ b/e2e/cloud/mcp-sse-replay.test.ts @@ -31,13 +31,20 @@ const initializedNotification = { method: "notifications/initialized", }; -const executeBody = (id: string, code: string) => ({ +const executeBody = (id: string | number, code: string) => ({ jsonrpc: "2.0" as const, id, method: "tools/call", params: { name: "execute", arguments: { code } }, }); +const toolsListBody = (id: number) => ({ + jsonrpc: "2.0" as const, + id, + method: "tools/list", + params: {}, +}); + const mcpHeaders = (bearer: string, sessionId?: string) => ({ accept: JSON_AND_SSE, authorization: `Bearer ${bearer}`, @@ -67,6 +74,8 @@ const openSession = async (mcpUrl: string, bearer: string): Promise => { return sessionId; }; +type JsonRpcId = string | number; + type JsonRpcMessage = { readonly id?: unknown; readonly result?: unknown; @@ -78,9 +87,10 @@ class SseCapture { readonly eventIds: string[] = []; private reader: ReadableStreamDefaultReader | null = null; private readonly waiters = new Map< - string, + JsonRpcId, Array<{ readonly resolve: (message: JsonRpcMessage) => void }> >(); + private responseWaiters: Array<{ readonly resolve: (message: JsonRpcMessage) => void }> = []; readonly finished: Promise; constructor( @@ -90,7 +100,7 @@ class SseCapture { this.finished = this.consume(); } - waitForId(id: string, timeoutMs: number): Promise { + waitForId(id: JsonRpcId, timeoutMs: number): Promise { const existing = this.messages.find((message) => message.id === id); if (existing) return Promise.resolve(existing); return new Promise((resolve, reject) => { @@ -108,6 +118,25 @@ class SseCapture { }); } + waitForFirstResponse(timeoutMs: number): Promise { + const existing = this.messages.find( + (message) => typeof message.id === "string" || typeof message.id === "number", + ); + if (existing) return Promise.resolve(existing); + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + // oxlint-disable-next-line executor/no-promise-reject, executor/no-error-constructor -- boundary: Promise timeout adapter for e2e polling. + reject(new Error("timed out waiting for the first JSON-RPC response")); + }, timeoutMs); + this.responseWaiters.push({ + resolve: (message) => { + clearTimeout(timeout); + resolve(message); + }, + }); + }); + } + abort(reason: string): void { this.abortController?.abort(reason); this.reader?.cancel(reason).catch(() => undefined); @@ -159,14 +188,22 @@ class SseCapture { if (!trimmed) return; const parsed = JSON.parse(trimmed) as JsonRpcMessage; this.messages.push(parsed); - if (typeof parsed.id !== "string") return; + if (typeof parsed.id !== "string" && typeof parsed.id !== "number") return; + const responseWaiters = this.responseWaiters; + this.responseWaiters = []; + for (const waiter of responseWaiters) waiter.resolve(parsed); const waiters = this.waiters.get(parsed.id) ?? []; this.waiters.delete(parsed.id); for (const waiter of waiters) waiter.resolve(parsed); } } -const openGet = async (mcpUrl: string, bearer: string, sessionId: string): Promise => { +const openGet = async ( + mcpUrl: string, + bearer: string, + sessionId: string, + lastEventId?: string, +): Promise => { const abortController = new AbortController(); const response = await fetch(mcpUrl, { method: "GET", @@ -175,6 +212,7 @@ const openGet = async (mcpUrl: string, bearer: string, sessionId: string): Promi authorization: `Bearer ${bearer}`, "mcp-protocol-version": PROTOCOL_VERSION, "mcp-session-id": sessionId, + ...(lastEventId ? { "last-event-id": lastEventId } : {}), }, signal: abortController.signal, }); @@ -246,6 +284,78 @@ scenario( }), ); +scenario( + "MCP streamable HTTP · Last-Event-ID recovery stays scoped to its originating request", + { timeout: 90_000 }, + Effect.gen(function* () { + const target = yield* Target; + const mcp = yield* Mcp; + const identity = yield* target.newIdentity(); + const bearer = yield* mcp.mintBearer(emailOf(identity)); + const sessionId = yield* Effect.promise(() => openSession(target.mcpUrl, bearer)); + + // Initialize's POST response is intentionally retained as undelivered. + // Drain it through the cursorless compatibility path first so id=20 below + // is the only unrelated response available to expose cross-stream replay. + const initializeReplay = yield* Effect.promise(() => openGet(target.mcpUrl, bearer, sessionId)); + yield* Effect.promise(() => initializeReplay.waitForId("initialize", 15_000)); + yield* Effect.promise(() => initializeReplay.finished); + yield* Effect.promise(() => delay(200)); + + const expectedId = 21; + const unrelatedId = 20; + const marker = `MARKER_STREAM_SCOPED_${randomUUID()}`; + const interrupted = new AbortController(); + const expectedPost = yield* Effect.promise(() => + startPostCapture( + target.mcpUrl, + bearer, + sessionId, + executeBody(expectedId, delayedCode(marker, 5_000)), + interrupted, + ), + ); + + // The priming event is always the first event on this POST stream. Its id + // is the cursor the real SDK carries when that stream is interrupted. + yield* Effect.promise(async () => { + const deadline = Date.now() + 10_000; + while (expectedPost.eventIds.length === 0 && Date.now() < deadline) await delay(25); + if (expectedPost.eventIds.length === 0) { + throw new Error("timed out waiting for the expected POST priming cursor"); + } + }); + const expectedCursor = expectedPost.eventIds[0]; + if (!expectedCursor) return yield* Effect.die("expected POST priming cursor is missing"); + + // Fully consume another POST response. workerd cannot prove that delivery, + // so Executor deliberately leaves id=20 persisted and marked undelivered. + const unrelated = yield* Effect.promise(() => + postJson(target.mcpUrl, bearer, toolsListBody(unrelatedId), sessionId), + ); + yield* Effect.promise(() => unrelated.text()); + expect(unrelated.status, "the unrelated tools/list completed normally").toBe(200); + + interrupted.abort("simulate the id=21 POST stream disconnecting after priming"); + yield* Effect.promise(() => expectedPost.finished.catch(() => undefined)); + + const recovery = yield* Effect.promise(() => + openGet(target.mcpUrl, bearer, sessionId, expectedCursor), + ); + const firstResponse = yield* Effect.promise(() => recovery.waitForFirstResponse(15_000)); + recovery.abort("scenario complete"); + + expect( + firstResponse.id, + "a targeted recovery never receives the unrelated persisted response", + ).toBe(expectedId); + expect( + JSON.stringify(firstResponse), + "the targeted request receives its own completed result", + ).toContain(marker); + }), +); + scenario( "MCP streamable HTTP · in-flight call survives the session idle timeout", { timeout: 90_000 }, diff --git a/patches/agents@0.17.3.patch b/patches/agents@0.17.3.patch index df7396f95..e78418ff8 100644 --- a/patches/agents@0.17.3.patch +++ b/patches/agents@0.17.3.patch @@ -38,7 +38,7 @@ index c8fad448e8797b89690a99d93490d1363851b225..77f9fe3f6f2375eadc9f7a2974f0d202 McpAgent, type McpAuthContext, diff --git a/dist/mcp/index.js b/dist/mcp/index.js -index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756cd72faef 100644 +index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..f168ef8efadc3b3a1700ee90801ecc2f8916c9ec 100644 --- a/dist/mcp/index.js +++ b/dist/mcp/index.js @@ -28,13 +28,17 @@ import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/ @@ -167,13 +167,13 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 + // canceled the POST response body, and request.signal does + // not reliably fire for that cancellation, so a successful + // close is NOT proof of delivery. Never ack POST-stream -+ // deliveries: the DO keeps the response persisted and the -+ // client's own reconnect GET replays and acks it. A client -+ // that DID receive the result closes the POST body reader -+ // without a Last-Event-ID reconnect, and the SDK drops -+ // responses for request ids it no longer tracks, so the -+ // worst case of this at-least-once choice is a benign -+ // replay to a fresh GET, not a wedged tool call. ++ // deliveries: the DO keeps each response persisted for its ++ // own Last-Event-ID recovery or an explicitly cursorless ++ // fresh GET. A client that DID receive the result closes the ++ // POST body reader without recovery, so the worst case of ++ // this at-least-once choice is a duplicate on that fresh-GET ++ // fallback, never a response injected into another cursor's ++ // recovery stream. + ws?.close(1000, "SSE response delivered"); + } + } @@ -494,7 +494,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 console.error("Error closing SSE connection:", error); } } -@@ -634,7 +857,23 @@ var StreamableHTTPServerTransport = class { +@@ -634,7 +857,21 @@ var StreamableHTTPServerTransport = class { } this.supersedePriorStreamConnections(agent, connection.id, resumedStreamId); connection.setState(resumeState); @@ -502,12 +502,10 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 + const ackStreamIds = []; + const replayedResponse = await this.replayEvents(lastEventId); + if (resumedStreamId !== STANDALONE_STREAM_ID && replayedResponse) ackStreamIds.push(resumedStreamId); -+ // A reconnect can carry a Last-Event-ID for an already-delivered -+ // stream (e.g. the initialize response) while a tool result -+ // completed on a since-abandoned POST stream. Replay those other -+ // undelivered responses on this connection too, otherwise they are -+ // stranded until the session is torn down. -+ ackStreamIds.push(...await this.replayUndeliveredResponses(agent, connection, resumedStreamId)); ++ // Last-Event-ID identifies one disconnected stream. Keep replay ++ // scoped to that stream: mixing another POST's response into this ++ // recovery stream can make a client stop reconnecting before the ++ // response it is actually waiting for arrives. + // Storage is NOT cleared here: replayed events are only enqueued + // on the WS bridge, and workerd cannot tell a dead client from a + // live one at write time. The close frame below makes the bridge @@ -519,14 +517,14 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 return; } } -@@ -644,6 +883,26 @@ var StreamableHTTPServerTransport = class { +@@ -644,6 +881,26 @@ var StreamableHTTPServerTransport = class { _standaloneSse: true }; connection.setState(standaloneState); -+ const replayedStreamIds = await this.replayUndeliveredResponses(agent, connection); -+ // Same delivery-confirmed clearing as the resume branch above. When -+ // nothing was replayed no close frame is sent and this connection stays -+ // open as the session's long-lived standalone listener. ++ const replayedStreamIds = await this.replayUndeliveredResponsesOnFreshGet(agent, connection); ++ // A GET without Last-Event-ID has no identified stream to preserve, so ++ // retain the existing fallback that drains completed POST responses. ++ // When nothing was replayed this stays the long-lived listener. + if (replayedStreamIds.length > 0) this.sendReplayComplete(connection, replayedStreamIds); + } + /** @@ -546,7 +544,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 } /** * Close any connection (other than `selfId`) currently bound to -@@ -664,12 +923,14 @@ var StreamableHTTPServerTransport = class { +@@ -664,12 +921,14 @@ var StreamableHTTPServerTransport = class { * Only used when resumability is enabled */ async replayEvents(lastEventId) { @@ -562,25 +560,26 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 this.writeSSEEvent(connection, message, eventId); } catch (error) { this.onerror?.(error); -@@ -678,6 +939,33 @@ var StreamableHTTPServerTransport = class { +@@ -678,6 +937,34 @@ var StreamableHTTPServerTransport = class { } catch (error) { this.onerror?.(error); } + return replayedResponse; + } + /** -+ * Enqueue every undelivered stream's events on `connection` and return -+ * the stream ids whose replay included a response. Deliberately does -+ * NOT clear storage: the caller sends a replay-complete close frame and -+ * the bridge acks each stream only after the client-facing writer -+ * drained and closed with the client still attached. ++ * For a fresh GET without Last-Event-ID, enqueue every undelivered ++ * stream's events on `connection` and return the stream ids whose ++ * replay included a response. A cursor-bearing GET never calls this: ++ * MCP requires that replay to stay on the stream the cursor identifies. ++ * Deliberately does NOT clear storage: the caller sends a replay-complete ++ * close frame and the bridge acks each stream only after the client-facing ++ * writer drained and closed with the client still attached. + */ -+ async replayUndeliveredResponses(agent, connection, skipStreamId) { ++ async replayUndeliveredResponsesOnFreshGet(agent, connection) { + const replayedStreamIds = []; + if (!this._eventStore?.replayEventsForStream) return replayedStreamIds; + const streamIds = await agent.getUndeliveredStreamIds(); + for (const streamId of streamIds) { -+ if (skipStreamId !== void 0 && streamId === skipStreamId) continue; + let replayedResponse = false; + await this._eventStore.replayEventsForStream(streamId, { send: async (eventId, message) => { + try { @@ -596,7 +595,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 } /** * Writes an event to the SSE stream with proper formatting -@@ -689,10 +977,65 @@ var StreamableHTTPServerTransport = class { +@@ -689,10 +976,65 @@ var StreamableHTTPServerTransport = class { return connection.send(JSON.stringify({ type: "cf_mcp_agent_event", event: eventData, @@ -662,7 +661,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 * Handles POST requests containing JSON-RPC messages */ async handlePostRequest(req, parsedBody) { -@@ -733,6 +1076,22 @@ var StreamableHTTPServerTransport = class { +@@ -733,6 +1075,22 @@ var StreamableHTTPServerTransport = class { }; connection.setState(postState); if (this._eventStore) await agent.setStreamRequestIds(streamId, requestIds); @@ -685,7 +684,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 for (const message of messages) { if (this.messageInterceptor) { if (await this.messageInterceptor(message, { -@@ -760,7 +1119,22 @@ var StreamableHTTPServerTransport = class { +@@ -760,7 +1118,22 @@ var StreamableHTTPServerTransport = class { * when the originating WS has dropped. */ async sendOnStream(agent, streamId, relatedIds, liveConnection, message, requestId) { @@ -709,7 +708,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 let shouldClose = false; if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { let responseIds = this._streamResponseIds.get(streamId); -@@ -777,9 +1151,11 @@ var StreamableHTTPServerTransport = class { +@@ -777,9 +1150,11 @@ var StreamableHTTPServerTransport = class { } catch (error) { this.onerror?.(error); } @@ -723,7 +722,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 } } async send(message, options) { -@@ -861,12 +1237,10 @@ var StreamableHTTPServerTransport = class { +@@ -861,12 +1236,10 @@ var StreamableHTTPServerTransport = class { * * ## Lifecycle * @@ -740,7 +739,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 * * Standalone GET stream events (`_GET_stream`) are *not* cleared * automatically; they accumulate for the lifetime of the DO. Bounded -@@ -893,12 +1267,34 @@ var DurableObjectEventStore = class DurableObjectEventStore { +@@ -893,12 +1266,34 @@ var DurableObjectEventStore = class DurableObjectEventStore { } async storeEvent(streamId, message) { if (streamId.includes(":")) throw new Error(`DurableObjectEventStore: streamId must not contain ':' (got ${JSON.stringify(streamId)})`); @@ -775,7 +774,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 return eventId; } async getStreamIdForEventId(eventId) { -@@ -915,9 +1311,59 @@ var DurableObjectEventStore = class DurableObjectEventStore { +@@ -915,9 +1310,59 @@ var DurableObjectEventStore = class DurableObjectEventStore { start: startKey, limit: DurableObjectEventStore.REPLAY_LIMIT }); @@ -836,7 +835,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 /** * Drop the event log for a single stream. Called by the transport * immediately after a POST's final response has been written to the -@@ -973,6 +1419,13 @@ DurableObjectEventStore.EVENT_KEY_PREFIX = "__mcp_event__:"; +@@ -973,6 +1418,13 @@ DurableObjectEventStore.EVENT_KEY_PREFIX = "__mcp_event__:"; DurableObjectEventStore.SEQ_PAD = 16; DurableObjectEventStore.DELETE_CHUNK = 128; DurableObjectEventStore.REPLAY_LIMIT = 1e3; @@ -850,7 +849,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 //#endregion //#region src/mcp/client-transports.ts /** -@@ -1381,6 +1834,47 @@ var McpAgent = class McpAgent extends Agent { +@@ -1381,6 +1833,47 @@ var McpAgent = class McpAgent extends Agent { async deleteStreamRequestIds(streamId) { await this.ctx.storage.delete(`${McpAgent.STREAM_REQS_KEY_PREFIX}${streamId}`); } @@ -898,7 +897,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..444312ff41fbcff60e2ea681a9d5c756 /** * Reverse lookup: find which POST stream a given `requestId` belongs * to, and return the stream's full `requestIds` list in the same -@@ -1697,7 +2191,8 @@ var McpAgent = class McpAgent extends Agent { +@@ -1697,7 +2190,8 @@ var McpAgent = class McpAgent extends Agent { } }; McpAgent.STREAM_REQS_KEY_PREFIX = "__mcp_stream_reqs__:";