From 4f17349c799188ea8f0856d0f9bfca2604b28352 Mon Sep 17 00:00:00 2001 From: Karavil <46430499+Karavil@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:12:17 +0000 Subject: [PATCH] Keep a request's messages on the stream that carried it --- .../src/mcp/agents-request-routing.test.ts | 207 ++++++++++++++++++ patches/agents@0.17.3.patch | 46 +++- 2 files changed, 242 insertions(+), 11 deletions(-) create mode 100644 packages/hosts/cloudflare/src/mcp/agents-request-routing.test.ts diff --git a/packages/hosts/cloudflare/src/mcp/agents-request-routing.test.ts b/packages/hosts/cloudflare/src/mcp/agents-request-routing.test.ts new file mode 100644 index 0000000000..196f8b9e0e --- /dev/null +++ b/packages/hosts/cloudflare/src/mcp/agents-request-routing.test.ts @@ -0,0 +1,207 @@ +// Unit coverage for request routing in the Durable Object transport (see +// patches/agents@0.17.3.patch). +// +// Clients that pool one MCP session for many callers (each caller numbering +// its ids from 0) put the same JSON-RPC id on several POST streams at once. The +// transport used to pick the target stream by that id: it preferred the stream +// that carried the request only while that stream was still live. Once that +// caller had hung up, a late result went to the one other live stream holding +// the same id (a valid-looking answer to someone else's call), or, with two or +// more such streams, every one of them got `-32603 Internal error` and was +// closed. +// +// Pinned here against the real patched transport: +// 1. A late result never reaches another caller's stream that holds the same +// id, and never errors that stream; it stays on its own stream. +// 2. A GET that resumed the originating stream (Last-Event-ID) still receives +// the result, and the stream keeps its own id when the send runs under +// that GET. +// 3. Related notifications and a batch's results stay on their own stream. +// 4. The async context is trusted only while its stream owns the id: a send +// that runs under another stream's context still goes by id. +import { describe, expect, it } from "@effect/vitest"; +import { __DO_NOT_USE_WILL_BREAK__agentContext as agentContext } from "agents"; +import { McpAgent } from "agents/mcp"; + +type RequestId = number | string; + +type FakeConnection = { + readonly id: string; + readonly state: { readonly streamId: string; readonly requestIds: RequestId[] }; + readonly sent: string[]; + readonly send: (message: string) => void; +}; + +const makeConnection = ( + id: string, + requestIds: RequestId[], + streamId: string = id, +): FakeConnection => { + const sent: string[] = []; + return { id, state: { streamId, requestIds }, sent, send: (message) => void sent.push(message) }; +}; + +/** The `McpAgent` surface the transport touches, backed by plain maps. */ +const makeAgent = (live: ReadonlyArray, rows: Map) => ({ + getSessionId: () => "session-1", + getTransportType: () => "streamable-http", + getEventStore: () => undefined, + getConnections: () => live, + getStreamRequestIds: async (streamId: string) => rows.get(streamId), + getStreamForRequestId: async (requestId: RequestId) => { + for (const [streamId, requestIds] of rows) { + if (requestIds.includes(requestId)) return { streamId, requestIds }; + } + return undefined; + }, + deleteStreamRequestIds: async (streamId: string) => void rows.delete(streamId), + markStreamUndelivered: async () => {}, + _handleElicitationResponse: () => false, +}); + +type FakeAgent = ReturnType; +type Transport = { + readonly send: (message: unknown, options?: { relatedRequestId?: RequestId }) => Promise; +}; + +const makeTransport = (agent: FakeAgent): Transport => + agentContext.run( + { agent, connection: undefined, request: undefined, email: undefined } as never, + () => + // oxlint-disable-next-line executor/no-double-cast -- test double: the transport class is not exported, so it is built through McpAgent's own initTransport on a fake agent + (McpAgent.prototype as unknown as { initTransport: () => Transport }).initTransport.call( + agent, + ), + ); + +/** Send `message` the way the MCP server does: inside the originating request's context. */ +const sendFrom = ( + agent: FakeAgent, + origin: FakeConnection, + transport: Transport, + message: unknown, + relatedRequestId?: RequestId, +) => + agentContext.run( + { agent, connection: origin, request: undefined, email: undefined } as never, + () => transport.send(message, { relatedRequestId }), + ); + +const result = (id: RequestId, text: string) => ({ + jsonrpc: "2.0" as const, + id, + result: { content: [{ type: "text", text }] }, +}); + +const progress = (progressToken: RequestId) => ({ + jsonrpc: "2.0" as const, + method: "notifications/progress", + params: { progressToken, progress: 1 }, +}); + +describe("DO transport: a response goes to the stream that carried its request", () => { + it("keeps a late result off the one other live stream that reuses its id", async () => { + const gone = makeConnection("gone", [0]); + const waiting = makeConnection("waiting", [0]); + const rows = new Map([ + ["gone", [0]], + ["waiting", [0]], + ]); + const agent = makeAgent([waiting], rows); + + await sendFrom(agent, gone, makeTransport(agent), result(0, "gone's result")); + + expect(waiting.sent, "the waiting caller never sees another caller's result").toEqual([]); + expect(rows.has("gone"), "the result closes out its own stream").toBe(false); + expect(rows.has("waiting"), "the waiting caller is still owed its own answer").toBe(true); + }); + + it("does not error other live streams that reuse the id", async () => { + const gone = makeConnection("gone", [0]); + const waitingA = makeConnection("waiting-a", [0]); + const waitingB = makeConnection("waiting-b", [0]); + const rows = new Map([ + ["gone", [0]], + ["waiting-a", [0]], + ["waiting-b", [0]], + ]); + const agent = makeAgent([waitingA, waitingB], rows); + + await sendFrom(agent, gone, makeTransport(agent), result(0, "gone's result")); + + expect(waitingA.sent, "no -32603 for a request this stream did not make").toEqual([]); + expect(waitingB.sent, "no -32603 for a request this stream did not make").toEqual([]); + expect([...rows.keys()].sort(), "both waiting callers are still owed an answer").toEqual([ + "waiting-a", + "waiting-b", + ]); + }); + + it("delivers to a GET that resumed the originating stream", async () => { + const post = makeConnection("post", [0]); + const resumed = makeConnection("resumed-get", [0], "post"); + const other = makeConnection("other", [0]); + const rows = new Map([ + ["post", [0]], + ["other", [0]], + ]); + const agent = makeAgent([resumed, other], rows); + + await sendFrom(agent, post, makeTransport(agent), result(0, "post's result")); + + expect(resumed.sent, "the resumed stream receives its own result").toHaveLength(1); + expect(resumed.sent[0]).toContain("post's result"); + expect(other.sent, "the other caller receives nothing").toEqual([]); + }); + + it("keeps the stream's own id when the send runs under the GET that resumed it", async () => { + const resumed = makeConnection("resumed-get", [0], "post"); + const other = makeConnection("other", [0]); + const rows = new Map([ + ["post", [0]], + ["other", [0]], + ]); + const agent = makeAgent([resumed, other], rows); + + await sendFrom(agent, resumed, makeTransport(agent), result(0, "post's result")); + + expect(resumed.sent, "the resumed stream receives its own result").toHaveLength(1); + expect(other.sent, "the other caller receives nothing").toEqual([]); + expect(rows.has("post"), "the result closes out the original stream").toBe(false); + expect(rows.has("other"), "the other caller is still owed its own answer").toBe(true); + }); + + it("keeps a batch's related notifications and results on its own stream", async () => { + const batch = makeConnection("batch", [0, 1]); + const other = makeConnection("other", [0]); + const rows = new Map([ + ["batch", [0, 1]], + ["other", [0]], + ]); + const agent = makeAgent([other], rows); + const transport = makeTransport(agent); + + await sendFrom(agent, batch, transport, progress(0), 0); + await sendFrom(agent, batch, transport, result(0, "batch's first result")); + + expect(other.sent, "the other caller sees neither message").toEqual([]); + expect(rows.get("batch"), "the batch still owes its second answer").toEqual([0, 1]); + expect(rows.has("other"), "the other caller is still owed its own answer").toBe(true); + }); + + it("routes by id when the send runs under a stream that does not own the id", async () => { + const owner = makeConnection("owner", [7]); + const bystander = makeConnection("bystander", [5]); + const rows = new Map([ + ["owner", [7]], + ["bystander", [5]], + ]); + const agent = makeAgent([owner, bystander], rows); + + await sendFrom(agent, bystander, makeTransport(agent), result(7, "owner's result")); + + expect(owner.sent, "the stream that holds id 7 receives it").toHaveLength(1); + expect(owner.sent[0]).toContain("owner's result"); + expect(bystander.sent, "the stream in context does not own id 7").toEqual([]); + }); +}); diff --git a/patches/agents@0.17.3.patch b/patches/agents@0.17.3.patch index 7eec44acf0..5a4e84f1c6 100644 --- a/patches/agents@0.17.3.patch +++ b/patches/agents@0.17.3.patch @@ -95,7 +95,7 @@ index c8fad448e8797b89690a99d93490d1363851b225..d80f66f1532c29dda0a0477dc1ec9ef5 McpAgent, type McpAuthContext, diff --git a/dist/mcp/index.js b/dist/mcp/index.js -index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c28786cca392 100644 +index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..012ca5d0ad52384804f06feabf0654e2ceca8237 100644 --- a/dist/mcp/index.js +++ b/dist/mcp/index.js @@ -28,13 +28,60 @@ import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/ @@ -1059,7 +1059,31 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 if (standalone) this.writeSSEEvent(standalone, message, eventId); } /** -@@ -861,12 +1466,10 @@ var StreamableHTTPServerTransport = class { +@@ -819,6 +1424,23 @@ var StreamableHTTPServerTransport = class { + const agent = this._agent; + const context = getCurrentAgent(); + const originatingConnection = context.agent === agent ? context.connection : void 0; ++ // Callers that share one session can have the same JSON-RPC id in ++ // flight on several streams at once. So when the stream that carried ++ // this request is known and still owns the id, the message goes to ++ // that stream only: to the GET that resumed it, to its POST connection, ++ // or, with neither left, into its own event log for its own ++ // Last-Event-ID recovery. Never to another stream that merely holds ++ // the same id. ++ if (originatingConnection) { ++ const streamId = originatingConnection.state?.streamId ?? originatingConnection.id; ++ const connections = Array.from(agent.getConnections()); ++ const liveConnection = connections.find((conn) => conn.id !== streamId && conn.state?.streamId === streamId) ?? connections.find((conn) => conn.id === streamId) ?? null; ++ const relatedIds = liveConnection?.state?.requestIds ?? await agent.getStreamRequestIds(streamId); ++ if (relatedIds?.includes(requestId)) { ++ await this.sendOnStream(agent, streamId, relatedIds, liveConnection, message, requestId); ++ return; ++ } ++ } + const matchingConnections = Array.from(agent.getConnections()).filter((conn) => conn.state?.requestIds?.includes(requestId)); + const liveConnection = matchingConnections.find((conn) => conn.id === originatingConnection?.id) ?? (matchingConnections.length === 1 ? matchingConnections[0] : null); + if (!liveConnection && matchingConnections.length > 1) { +@@ -861,12 +1483,10 @@ var StreamableHTTPServerTransport = class { * * ## Lifecycle * @@ -1076,7 +1100,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 * * Standalone GET stream events (`_GET_stream`) are *not* cleared * automatically; they accumulate for the lifetime of the DO. Bounded -@@ -893,12 +1496,34 @@ var DurableObjectEventStore = class DurableObjectEventStore { +@@ -893,12 +1513,34 @@ var DurableObjectEventStore = class DurableObjectEventStore { } async storeEvent(streamId, message) { if (streamId.includes(":")) throw new Error(`DurableObjectEventStore: streamId must not contain ':' (got ${JSON.stringify(streamId)})`); @@ -1111,7 +1135,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 return eventId; } async getStreamIdForEventId(eventId) { -@@ -915,9 +1540,59 @@ var DurableObjectEventStore = class DurableObjectEventStore { +@@ -915,9 +1557,59 @@ var DurableObjectEventStore = class DurableObjectEventStore { start: startKey, limit: DurableObjectEventStore.REPLAY_LIMIT }); @@ -1172,7 +1196,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 /** * 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 +1648,13 @@ DurableObjectEventStore.EVENT_KEY_PREFIX = "__mcp_event__:"; +@@ -973,6 +1665,13 @@ DurableObjectEventStore.EVENT_KEY_PREFIX = "__mcp_event__:"; DurableObjectEventStore.SEQ_PAD = 16; DurableObjectEventStore.DELETE_CHUNK = 128; DurableObjectEventStore.REPLAY_LIMIT = 1e3; @@ -1186,7 +1210,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 //#endregion //#region src/mcp/client-transports.ts /** -@@ -1355,6 +2037,26 @@ function experimental_createMcpHandler(server, options = {}) { +@@ -1355,6 +2054,26 @@ function experimental_createMcpHandler(server, options = {}) { } //#endregion //#region src/mcp/index.ts @@ -1213,7 +1237,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 var McpAgent = class McpAgent extends Agent { constructor(..._args) { super(..._args); -@@ -1369,18 +2071,121 @@ var McpAgent = class McpAgent extends Agent { +@@ -1369,18 +2088,121 @@ var McpAgent = class McpAgent extends Agent { async getInitializeRequest() { return this.ctx.storage.get("initializeRequest"); } @@ -1338,7 +1362,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 /** * Reverse lookup: find which POST stream a given `requestId` belongs * to, and return the stream's full `requestIds` list in the same -@@ -1407,10 +2212,14 @@ var McpAgent = class McpAgent extends Agent { +@@ -1407,10 +2229,14 @@ var McpAgent = class McpAgent extends Agent { limit: STREAM_REQS_SCAN_LIMIT }); if (rows.size === STREAM_REQS_SCAN_LIMIT) console.warn(`McpAgent: getStreamForRequestId hit the ${STREAM_REQS_SCAN_LIMIT}-key scan cap; stale __mcp_stream_reqs__ entries may be accumulating from abandoned POSTs`); @@ -1357,7 +1381,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 } /** Read the transport type for this agent. * This relies on the naming scheme being `sse:${sessionId}`, -@@ -1498,6 +2307,12 @@ var McpAgent = class McpAgent extends Agent { +@@ -1498,6 +2324,12 @@ var McpAgent = class McpAgent extends Agent { } /** Sets up the MCP transport and server every time the Agent is started.*/ async onStart(props) { @@ -1370,7 +1394,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 if (props) await this.updateProps(props); else this.props = await this.ctx.storage.get("props"); await this.init(); -@@ -1516,23 +2331,36 @@ var McpAgent = class McpAgent extends Agent { +@@ -1516,23 +2348,36 @@ var McpAgent = class McpAgent extends Agent { return; } break; @@ -1422,7 +1446,7 @@ index 1edcf0c8c9e67aa211ae515e7672cdf79912101e..231220c26b995a48e40dba993773c287 } } } -@@ -1697,7 +2525,9 @@ var McpAgent = class McpAgent extends Agent { +@@ -1697,7 +2542,9 @@ var McpAgent = class McpAgent extends Agent { } }; McpAgent.STREAM_REQS_KEY_PREFIX = "__mcp_stream_reqs__:";