From 888bba6fc45bae6949d8111e2116b492edb2059e Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:47:15 -0700 Subject: [PATCH] Fail loudly when a DO restart kills an in-flight MCP tool call --- packages/hosts/cloudflare/package.json | 4 + .../mcp/agent-session-durable-object.test.ts | 226 +++++++++++- .../src/mcp/agent-session-durable-object.ts | 95 +++++ .../cloudflare/src/mcp/restart-reaper.test.ts | 330 ++++++++++++++++++ .../cloudflare/src/mcp/restart-reaper.ts | 258 ++++++++++++++ 5 files changed, 903 insertions(+), 10 deletions(-) create mode 100644 packages/hosts/cloudflare/src/mcp/restart-reaper.test.ts create mode 100644 packages/hosts/cloudflare/src/mcp/restart-reaper.ts diff --git a/packages/hosts/cloudflare/package.json b/packages/hosts/cloudflare/package.json index 861447689..eb398cef6 100644 --- a/packages/hosts/cloudflare/package.json +++ b/packages/hosts/cloudflare/package.json @@ -23,6 +23,10 @@ "./mcp/session-stub": { "types": "./src/mcp/session-stub.ts", "default": "./src/mcp/session-stub.ts" + }, + "./mcp/restart-reaper": { + "types": "./src/mcp/restart-reaper.ts", + "default": "./src/mcp/restart-reaper.ts" } }, "scripts": { diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts index c5fabeb2b..8e98e3e81 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts @@ -1,8 +1,14 @@ import { describe, expect, it } from "@effect/vitest"; -import { Cause, Effect } from "effect"; +import { Cause, Effect, Exit } from "effect"; +import type * as Tracer from "effect/Tracer"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; -import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; +import type { + JSONRPCMessage, + MessageExtraInfo, + RequestId, +} from "@modelcontextprotocol/sdk/types.js"; +import { DurableObjectEventStore } from "agents/mcp"; import { defaultMcpResource } from "@executor-js/host-mcp"; import type { ExecutionEngine, ExecutionResult, ResumeResponse } from "@executor-js/execution"; @@ -13,6 +19,7 @@ import { type McpSessionModelResumeResult, type SessionMeta, } from "./agent-session-durable-object"; +import { RESTART_REAP_ERROR_CODE } from "./restart-reaper"; class MemoryStorage { private readonly data = new Map(); @@ -52,16 +59,26 @@ class MemoryStorage { this.data.clear(); } + // Mirrors DurableObjectStorage's SORTED key-value surface, including `start` + // and `reverse`. The real agents `DurableObjectEventStore` depends on that + // ordering for both sequence recovery (`reverse`+`limit: 1`) and replay + // (`start`), so a list that returned insertion order would quietly make + // event-store assertions meaningless. async list( - options: { readonly prefix?: string; readonly limit?: number } = {}, + options: { + readonly prefix?: string; + readonly start?: string; + readonly limit?: number; + readonly reverse?: boolean; + } = {}, ): Promise> { - const rows = new Map(); - for (const [key, value] of this.data) { - if (options.prefix && !key.startsWith(options.prefix)) continue; - rows.set(key, value as T); - if (options.limit && rows.size >= options.limit) break; - } - return rows; + const keys = [...this.data.keys()] + .filter((key) => (options.prefix === undefined ? true : key.startsWith(options.prefix))) + .filter((key) => (options.start === undefined ? true : key >= options.start)) + .sort(); + if (options.reverse === true) keys.reverse(); + const limited = options.limit === undefined ? keys : keys.slice(0, options.limit); + return new Map(limited.map((key) => [key, this.data.get(key) as T])); } async blockConcurrencyWhile(callback: () => T | Promise): Promise { @@ -93,6 +110,11 @@ type HarnessSession = { props: Record; runMcpAgentOnStart: () => Promise; server?: McpServer; + /** Inherited from the real `McpAgent`; the reaper tests drive them directly. */ + setStreamRequestIds: (streamId: string, requestIds: RequestId[]) => Promise; + getStreamRequestIds: (streamId: string) => Promise; + /** The base's telemetry seam; the span test swaps in a recording tracer. */ + withTelemetry: (effect: Effect.Effect) => Effect.Effect; sessionMeta: SessionMeta; sessionTimeoutMs: () => number; resumeExecutionForModel: ( @@ -342,3 +364,187 @@ describe("McpAgentSessionDOBase transport restore", () => { expect(restoredEngine.calls).toEqual([{ executionId: "exec-model", response: approval }]); }); }); + +// Restart reaping, driven through the DO's real start path against the real +// agents storage layout (`McpAgent.getOpenStreamRequestIds` and +// `DurableObjectEventStore`, both inherited off the live prototype chain rather +// than stubbed). This is the integration half of restart-reaper.test.ts: it +// pins that the reaper actually FIRES on start, that it is wired to the same +// event store the transport writes through, and that a normal session start +// stays inert. +// +// The failure being covered: a deploy resets the isolate mid tool call. The +// execution dies unrecoverably and the client's POST stream closes cleanly, so +// without this the client hangs to its own timeout with no error and no span. +describe("McpAgentSessionDOBase restart reaping", () => { + /** + * A tracer that records ended spans with their final `Exit`. Effect's OTEL + * tracer maps a failed Exit to `SpanStatusCode.ERROR` plus a recorded + * exception, so asserting the Exit here is asserting what Axiom would show. + */ + const recordingTracer = () => { + const ended: Array<{ + readonly name: string; + readonly attributes: Map; + exit?: Exit.Exit; + }> = []; + const tracer: Tracer.Tracer = { + span: (options) => { + const attributes = new Map(); + const record = { name: options.name, attributes } as (typeof ended)[number]; + let status: Tracer.SpanStatus = { _tag: "Started", startTime: options.startTime }; + return { + _tag: "Span", + name: options.name, + spanId: "1234567890abcdef", + traceId: "4268a606000000000000000000000000", + parent: options.parent, + annotations: options.annotations, + get status() { + return status; + }, + attributes, + links: options.links, + sampled: options.sampled, + kind: options.kind, + end: (endTime, exit) => { + status = { _tag: "Ended", startTime: options.startTime, endTime, exit }; + record.exit = exit; + ended.push(record); + }, + attribute: (key, value) => { + attributes.set(key, value); + }, + event: () => undefined, + addLinks: () => undefined, + }; + }, + }; + return { tracer, ended }; + }; + + const replayFrom = async ( + session: HarnessSession, + lastEventId: string, + ): Promise> => { + const store = new DurableObjectEventStore(session.ctx as never); + const replayed: JSONRPCMessage[] = []; + await store.replayEventsAfter(lastEventId, { + send: async (_eventId: string, message: JSONRPCMessage) => { + replayed.push(message); + }, + }); + return replayed; + }; + + /** + * Reproduce the pre-restart on-disk state: the transport primed the POST + * stream and persisted its in-flight request ids, then the isolate died + * before any response was written. Only durable state is set up — exactly + * what actually survives a reset. + */ + const seedKilledCall = async ( + session: HarnessSession, + streamId: string, + requestIds: ReadonlyArray, + ): Promise => { + const store = new DurableObjectEventStore(session.ctx as never); + const primingId = await store.storeEvent(streamId, { + jsonrpc: "2.0", + method: "notifications/message", + params: { level: "debug", data: "mcp-stream-priming" }, + }); + await session.setStreamRequestIds(streamId, [...requestIds]); + return primingId; + }; + + it("turns a call killed by the restart into a replayable error on the next start", async () => { + const session = await makeHarnessSession(); + const primingId = await seedKilledCall(session, "killed-stream", [11]); + + await session.onStart(); + + const replayed = await replayFrom(session, primingId); + expect(replayed.length, "the reconnect replays exactly one error").toBe(1); + expect(replayed[0], "and it names the restart without inviting a blind retry").toMatchObject({ + jsonrpc: "2.0", + id: 11, + error: { code: RESTART_REAP_ERROR_CODE, data: { retrySafe: false } }, + }); + expect( + await session.getStreamRequestIds("killed-stream"), + "the reaped stream no longer looks like running work to the idle alarm", + ).toBeUndefined(); + }); + + it("reaps every concurrent in-flight request on the session", async () => { + // Several tool calls can be in flight on one session at once; a reset + // kills all of them. Missing any one leaves that client call hanging. + const session = await makeHarnessSession(); + const first = await seedKilledCall(session, "killed-a", [1]); + const second = await seedKilledCall(session, "killed-b", [2, 3]); + + await session.onStart(); + + expect( + (await replayFrom(session, first)).map((m) => (m as { readonly id?: RequestId }).id), + "the single-request stream is reaped", + ).toEqual([1]); + expect( + (await replayFrom(session, second)).map((m) => (m as { readonly id?: RequestId }).id), + "and both requests batched on the other stream are too", + ).toEqual([2, 3]); + }); + + it("ends the reap span as a failure so the loss is visible in the trace store", async () => { + // The observability half of the bug: a killed execution exports NOTHING — + // no exception, no span — so an outage of this shape is invisible. The + // reap deliberately ends its span on a failed Exit, which the OTEL tracer + // renders as status ERROR with a recorded exception; absence would look + // exactly like health. Asserted against the real span lifecycle via a + // recording tracer rather than trusting the annotation call. + const session = await makeHarnessSession(); + await seedKilledCall(session, "traced-stream", [1, 2]); + // The DO runs each method in its own `Effect.runPromise`, so a tracer + // cannot be injected from outside. `withTelemetry` is the base's own seam + // for installing one — the same hook cloud uses to provide the real OTEL + // tracer — so the recording tracer goes in exactly where production's does. + const spans = recordingTracer(); + session.withTelemetry = (effect) => Effect.withTracer(effect, spans.tracer); + + await session.onStart(); + + const reapSpan = spans.ended.find((span) => span.name === "McpSessionDO.restart_reap"); + expect(reapSpan, "the reap exports its own span").toBeDefined(); + expect(reapSpan?.exit && Exit.isFailure(reapSpan.exit), "and ends it as a FAILURE").toBe(true); + expect( + reapSpan?.attributes.get("exception.message"), + "carrying how much work the restart destroyed", + ).toContain("2 in-flight MCP request(s)"); + }); + + it("writes nothing when a start finds no interrupted work", async () => { + // The overwhelmingly common case — an ordinary cold start or a restore + // after idle disposal. A false positive here would inject a spurious error + // into a healthy session, so inertness is asserted directly. + const session = await makeHarnessSession(); + const store = new DurableObjectEventStore(session.ctx as never); + const streamId = "completed-stream"; + const primingId = await store.storeEvent(streamId, { + jsonrpc: "2.0", + method: "notifications/message", + params: { level: "debug", data: "mcp-stream-priming" }, + }); + const response: JSONRPCMessage = { jsonrpc: "2.0", id: 1, result: { ok: true } }; + await store.storeEvent(streamId, response); + // The transport deletes the request-ids key as it writes the final + // response, so a completed call leaves nothing behind to reap. + + await session.onStart(); + + expect( + await replayFrom(session, primingId), + "the completed call's result is the only thing replayed", + ).toEqual([response]); + }); +}); diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts index aa069be60..b97ec0336 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts @@ -33,6 +33,13 @@ import { pausedLeaseExtensionLog, runningLeaseExtensionLog, } from "./session-alarm-policy"; +import { + McpRestartReapedExecutions, + asRestartReapAgent, + asRestartReapEventStore, + reapOrphanedRequests, + restartReapLog, +} from "./restart-reaper"; export type IncomingTraceHeaders = IncomingPropagationHeaders; @@ -618,9 +625,97 @@ export abstract class McpAgentSessionDOBase< yield* self.closeRuntime(); return yield* Effect.failCause(started.cause); } + yield* self.reapRestartOrphanedRequests(); }); } + /** + * Fail loudly for tool calls the previous isolate was killed mid-flight. + * + * Runs on every runtime start, immediately after `closeRuntime()` has torn + * down whatever ran before and the fresh transport is up. That ordering IS + * the safety argument: at this instant no execution from any earlier runtime + * survives, so a POST stream that still carries persisted in-flight request + * ids can only be work that is already dead — a deploy that reset the + * isolate, or an idle disposal that outlived its running lease. Reaping is + * therefore never racing live work, and the same code covers both causes. + * + * The client half of the fix is the priming SSE event the patched agents + * transport writes at the head of every POST tools/call stream: without it + * the MCP SDK never reconnects a cleanly-closed POST stream, and the error + * written here would sit in storage unread while the client hangs to its own + * timeout. See `restart-reaper.ts` and `patches/agents@0.17.3.patch`. + * + * Failures here are logged and swallowed: a session must still start even if + * the event store rejects a write, and the pre-existing behavior (a hang) is + * no worse than before. + */ + private reapRestartOrphanedRequests(): Effect.Effect { + const self = this; + return Effect.gen(function* () { + yield* self.prepareErrorCaptureScope(); + // `getEventStore()` is the agents SDK's documented override point for + // resumability, so it is the supported way to reach the same store the + // transport writes through. A subclass that disables resumability gets + // `undefined` here and the reaper stands down — correctly, since with no + // event store there is no replay for a client to collect. + const agent = asRestartReapAgent(self); + const eventStore = asRestartReapEventStore(self.getEventStore()); + if (!agent || !eventStore) { + console.warn( + JSON.stringify({ + event: "mcp_session_restart_reap_unavailable", + sessionId: self.sessionId, + reason: agent ? "no_event_store" : "agent_api_missing", + }), + ); + return; + } + + const outcome = yield* reapOrphanedRequests({ agent, eventStore }); + if (outcome.kind === "unavailable" || outcome.requestCount === 0) return; + + console.error( + JSON.stringify( + restartReapLog({ + sessionId: self.sessionId, + streamIds: outcome.streamIds, + requestCount: outcome.requestCount, + }), + ), + ); + // Fail the span deliberately. Killed executions export NOTHING today — + // that invisibility is half the bug — so the reap is surfaced as a real + // span failure, which the OTEL tracer renders as status ERROR with a + // recorded exception. The failure is caught immediately below, so DO + // startup still succeeds. + return yield* new McpRestartReapedExecutions({ + streamCount: outcome.streamIds.length, + requestCount: outcome.requestCount, + }); + }).pipe( + // Inside the span, so the annotations land on `restart_reap` itself. The + // OTEL tracer additionally sets status ERROR and records the exception + // when the span ends on a failure, which is what makes the reap + // queryable in the trace store. + Effect.tapCause((cause) => + Effect.gen(function* () { + yield* self.captureCauseEffect(cause); + yield* self.recordCauseOnSpan(cause); + }), + ), + Effect.withSpan("McpSessionDO.restart_reap"), + // Restart reaping starts a fresh trace: the isolate that owned the + // original request's trace context died with it, and a start carries no + // incoming propagation. Without the telemetry layer these spans land in + // Effect's no-op tracer and never export — which is the invisibility + // this whole change exists to remove. + (effect) => self.withTelemetry(effect), + Effect.ignore, + (effect) => self.withSpanFlush(effect), + ); + } + protected runMcpAgentOnStart(props?: McpSessionProps): Promise { return super.onStart(props); } diff --git a/packages/hosts/cloudflare/src/mcp/restart-reaper.test.ts b/packages/hosts/cloudflare/src/mcp/restart-reaper.test.ts new file mode 100644 index 000000000..e4ba14986 --- /dev/null +++ b/packages/hosts/cloudflare/src/mcp/restart-reaper.test.ts @@ -0,0 +1,330 @@ +// Unit coverage for the restart reaper (see restart-reaper.ts). A Durable +// Object reset kills an in-flight tool execution outright: the result is never +// produced, so resumability has nothing to replay, and the client's POST stream +// closes CLEANLY so the MCP client neither errors nor reconnects — it hangs to +// its own timeout with no server-side signal at all. +// +// The reaper's job is to make that failure REACHABLE: at DO start, any POST +// stream still carrying persisted in-flight request ids was killed mid-flight, +// so an error is appended to its event log and the request-id key dropped. The +// reconnecting client (primed by the SSE priming event — see +// agents-priming-event.test.ts) replays from `last-event-id` and collects it. +// +// These tests drive the REAL `DurableObjectEventStore` so the assertions cover +// actual event ordering and replay, not a hand-rolled model of it. What's +// pinned: +// 1. An orphaned in-flight request yields an error that a `last-event-id` +// reconnect after the priming event actually replays. +// 2. No false positives: a stream the transport already finished (its +// request-ids key deleted on the final response) is left untouched. +// 3. Concurrent in-flight requests — several streams, and several batched +// requests on ONE stream — are ALL reaped; missing any one still hangs +// the client. +// 4. The standalone GET stream is never reaped. +// 5. Retry safety: the error does not invite an automatic retry. +import { describe, expect, it } from "@effect/vitest"; +import { Effect, Exit } from "effect"; +import { DurableObjectEventStore } from "agents/mcp"; +import type { JSONRPCMessage, RequestId } from "@modelcontextprotocol/sdk/types.js"; + +import { + RESTART_REAP_ERROR_CODE, + RESTART_REAP_ERROR_MESSAGE, + asRestartReapAgent, + asRestartReapEventStore, + reapOrphanedRequests, + type RestartReapAgent, +} from "./restart-reaper"; + +type ListOptions = { + readonly prefix?: string; + readonly start?: string; + readonly limit?: number; + readonly reverse?: boolean; +}; + +/** Minimal in-memory stand-in for DurableObjectStorage's sorted KV surface. */ +const makeFakeStorage = () => { + const entries = new Map(); + return { + entries, + put: (key: string, value: unknown) => { + entries.set(key, value); + return Promise.resolve(); + }, + delete: (keys: string | ReadonlyArray) => { + for (const key of Array.isArray(keys) ? keys : [keys]) entries.delete(key); + return Promise.resolve(); + }, + list: (options: ListOptions = {}) => { + const keys = [...entries.keys()] + .filter((key) => (options.prefix === undefined ? true : key.startsWith(options.prefix))) + .filter((key) => (options.start === undefined ? true : key >= options.start)) + .sort(); + if (options.reverse === true) keys.reverse(); + const limited = options.limit === undefined ? keys : keys.slice(0, options.limit); + return Promise.resolve(new Map(limited.map((key) => [key, entries.get(key)]))); + }, + }; +}; + +/** + * Stands in for the patched `McpAgent`'s stream-request-id bookkeeping. The + * real transport writes an entry when a POST dispatches and deletes it the + * instant the final response is written, so "an entry survived a restart" + * means "dispatched, never answered" — the exact condition the reaper acts on. + */ +const makeStreamRequestIds = ( + initial: ReadonlyArray<{ readonly streamId: string; readonly requestIds: RequestId[] }> = [], +) => { + const open = new Map(initial.map((entry) => [entry.streamId, entry.requestIds] as const)); + const agent: RestartReapAgent = { + getOpenStreamRequestIds: () => + Promise.resolve([...open].map(([streamId, requestIds]) => ({ streamId, requestIds }))), + deleteStreamRequestIds: (streamId: string) => { + open.delete(streamId); + return Promise.resolve(); + }, + }; + return { agent, open }; +}; + +/** The priming notification the patched transport writes at the head of every POST tools/call stream. */ +const PRIMING_MESSAGE: JSONRPCMessage = { + jsonrpc: "2.0", + method: "notifications/message", + params: { level: "debug", data: "mcp-stream-priming" }, +}; + +const replayAfter = async ( + store: DurableObjectEventStore, + lastEventId: string, +): Promise> => { + const replayed: JSONRPCMessage[] = []; + await store.replayEventsAfter(lastEventId, { + send: async (_eventId: string, message: JSONRPCMessage) => { + replayed.push(message); + }, + }); + return replayed; +}; + +const eventKeys = (storage: ReturnType): ReadonlyArray => + [...storage.entries.keys()].sort(); + +const runReap = (agent: RestartReapAgent, eventStore: DurableObjectEventStore) => + Effect.runPromise(Effect.exit(reapOrphanedRequests({ agent, eventStore }))); + +describe("restart reaper: orphaned in-flight requests", () => { + it("writes a replayable error for a request the restart killed mid-flight", async () => { + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + const streamId = "post-stream"; + + // The pre-restart state: the transport primed the stream and persisted the + // request id, then the isolate died before any response was written. + const primingId = await store.storeEvent(streamId, PRIMING_MESSAGE); + const { agent, open } = makeStreamRequestIds([{ streamId, requestIds: [7] }]); + + const outcome = await runReap(agent, store); + + expect(Exit.isSuccess(outcome), "the reap succeeds").toBe(true); + expect(outcome).toMatchObject({ + value: { kind: "reaped", streamIds: [streamId], requestCount: 1 }, + }); + + const replayed = await replayAfter(store, primingId); + expect(replayed, "a last-event-id reconnect after the priming event replays the error").toEqual( + [ + { + jsonrpc: "2.0", + id: 7, + error: { + code: RESTART_REAP_ERROR_CODE, + message: RESTART_REAP_ERROR_MESSAGE, + data: { reason: "server_restart", retrySafe: false }, + }, + }, + ], + ); + expect(open.has(streamId), "the reaped stream's request-id key is dropped").toBe(false); + }); + + it("reaps every request in a batched POST, not just the first", async () => { + // The transport only closes a POST stream once ALL its request ids have + // responded, so leaving one unanswered hangs the client exactly as before. + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + const streamId = "batched-stream"; + const primingId = await store.storeEvent(streamId, PRIMING_MESSAGE); + const { agent } = makeStreamRequestIds([{ streamId, requestIds: [1, 2, "three"] }]); + + const outcome = await runReap(agent, store); + + expect(outcome).toMatchObject({ value: { kind: "reaped", requestCount: 3 } }); + const replayed = await replayAfter(store, primingId); + expect( + replayed.map((message) => (message as { readonly id?: RequestId }).id), + "every batched request id gets its own error, in order", + ).toEqual([1, 2, "three"]); + }); + + it("reaps concurrent in-flight requests across several streams on one session", async () => { + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + const primingIds = new Map(); + for (const streamId of ["stream-a", "stream-b", "stream-c"]) { + primingIds.set(streamId, await store.storeEvent(streamId, PRIMING_MESSAGE)); + } + const { agent, open } = makeStreamRequestIds([ + { streamId: "stream-a", requestIds: [1] }, + { streamId: "stream-b", requestIds: [2] }, + { streamId: "stream-c", requestIds: [3] }, + ]); + + const outcome = await runReap(agent, store); + + expect(outcome).toMatchObject({ + value: { kind: "reaped", streamIds: ["stream-a", "stream-b", "stream-c"], requestCount: 3 }, + }); + for (const [streamId, primingId] of primingIds) { + const replayed = await replayAfter(store, primingId); + expect(replayed.length, `${streamId} replays exactly one error`).toBe(1); + expect(replayed[0]).toMatchObject({ error: { code: RESTART_REAP_ERROR_CODE } }); + } + expect(open.size, "no stream is left holding in-flight request ids").toBe(0); + }); +}); + +describe("restart reaper: no false positives", () => { + it("leaves a normally-completed call alone", async () => { + // A completed call: the transport wrote the response and deleted the + // stream's request-ids key in the same `shouldClose` step, so there is + // nothing for the reaper to see. + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + const streamId = "completed-stream"; + const primingId = await store.storeEvent(streamId, PRIMING_MESSAGE); + const response: JSONRPCMessage = { + jsonrpc: "2.0", + id: 1, + result: { content: [{ type: "text", text: "MARKER" }] }, + }; + await store.storeEvent(streamId, response); + const keysBefore = eventKeys(storage); + const { agent } = makeStreamRequestIds(); + + const outcome = await runReap(agent, store); + + expect(outcome).toMatchObject({ value: { kind: "reaped", streamIds: [], requestCount: 0 } }); + expect(eventKeys(storage), "no event was appended").toEqual(keysBefore); + expect( + await replayAfter(store, primingId), + "the reconnect still replays exactly the real result", + ).toEqual([response]); + }); + + it("leaves a stream whose persisted request-id list is empty", async () => { + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + await store.storeEvent("empty-stream", PRIMING_MESSAGE); + const keysBefore = eventKeys(storage); + const { agent } = makeStreamRequestIds([{ streamId: "empty-stream", requestIds: [] }]); + + const outcome = await runReap(agent, store); + + expect(outcome).toMatchObject({ value: { kind: "reaped", streamIds: [], requestCount: 0 } }); + expect(eventKeys(storage)).toEqual(keysBefore); + }); + + it("never reaps the standalone GET listen stream", async () => { + // `_GET_stream` is the session's long-lived server-to-client channel. Its + // events are deliberately never auto-cleared, and it carries no POST + // request ids — but the skip is explicit so a future SDK change that starts + // recording ids against it cannot make the reaper cancel a live listener. + const storage = makeFakeStorage(); + const store = new DurableObjectEventStore(storage as never); + await store.storeEvent("_GET_stream", PRIMING_MESSAGE); + const keysBefore = eventKeys(storage); + const { agent, open } = makeStreamRequestIds([ + { streamId: "_GET_stream", requestIds: ["listen"] }, + ]); + + const outcome = await runReap(agent, store); + + expect(outcome).toMatchObject({ value: { kind: "reaped", streamIds: [], requestCount: 0 } }); + expect(eventKeys(storage), "the listen stream's event log is untouched").toEqual(keysBefore); + expect(open.has("_GET_stream"), "its request-id entry is left in place").toBe(true); + }); +}); + +describe("restart reaper: retry safety", () => { + it("does not invite an automatic retry", async () => { + // `execute` runs arbitrary user code, which may already have charged a + // card or sent mail before the isolate died. The server cannot know how + // far it got, so the error must not pre-approve a retry — in prose or in + // its machine-readable data. + const store = new DurableObjectEventStore(makeFakeStorage() as never); + const streamId = "side-effect-stream"; + const primingId = await store.storeEvent(streamId, PRIMING_MESSAGE); + const { agent } = makeStreamRequestIds([{ streamId, requestIds: [1] }]); + + await runReap(agent, store); + const [replayed] = await replayAfter(store, primingId); + // Destructured off the JSON-RPC error OBJECT (a wire payload with a + // `message` field), not a JS Error — hence the local names. + const { message: wireMessage, data: wireData } = ( + replayed as { readonly error: { readonly message: string; readonly data: unknown } } + ).error; + + expect(wireMessage.toLowerCase(), "the message never says 'please retry'").not.toContain( + "please retry", + ); + expect(wireMessage, "it states the outcome is unknown").toContain("outcome is unknown"); + expect(wireData, "and marks the request machine-readably retry-unsafe").toMatchObject({ + retrySafe: false, + }); + }); +}); + +describe("restart reaper: SDK seam", () => { + it("stands down when the agents API the reaper depends on is missing", async () => { + // `getOpenStreamRequestIds` is `@internal` in the agents SDK. If an upgrade + // removes it the DO must still start; the reaper reports unavailable + // rather than throwing on every cold start. + expect(asRestartReapAgent({}), "an agent without the API is rejected").toBeNull(); + expect(asRestartReapAgent(null)).toBeNull(); + expect( + asRestartReapAgent(makeStreamRequestIds().agent), + "the real shape is accepted", + ).not.toBeNull(); + }); + + it("accepts the real DurableObjectEventStore as its event-store seam", () => { + const store = new DurableObjectEventStore(makeFakeStorage() as never); + expect(asRestartReapEventStore(store)).not.toBeNull(); + expect( + asRestartReapEventStore(undefined), + "a disabled event store stands the reaper down", + ).toBeNull(); + }); + + it("reports a failure when the event store cannot record the error", async () => { + const failingStore = { + storeEvent: () => + // oxlint-disable-next-line executor/no-promise-reject, executor/no-error-constructor -- boundary: models a rejecting SDK EventStore, whose contract is a rejected Promise. + Promise.reject(new Error("storage unavailable")), + }; + const { agent, open } = makeStreamRequestIds([{ streamId: "doomed", requestIds: [1] }]); + + const outcome = await Effect.runPromise( + Effect.exit(reapOrphanedRequests({ agent, eventStore: failingStore })), + ); + + expect(Exit.isFailure(outcome), "the failure is reported, not swallowed").toBe(true); + expect( + open.has("doomed"), + "the request-id key survives, so the next start can retry the reap", + ).toBe(true); + }); +}); diff --git a/packages/hosts/cloudflare/src/mcp/restart-reaper.ts b/packages/hosts/cloudflare/src/mcp/restart-reaper.ts new file mode 100644 index 000000000..b8490977f --- /dev/null +++ b/packages/hosts/cloudflare/src/mcp/restart-reaper.ts @@ -0,0 +1,258 @@ +// --------------------------------------------------------------------------- +// Restart reaper — turn a Durable Object reset that killed an in-flight tool +// call into a prompt, replayable JSON-RPC error instead of a silent hang. +// +// ## The failure this exists for +// +// A deploy (or any isolate reset) destroys the DO mid-execution. The work is +// gone and unrecoverable: the result was never produced, so MCP stream +// resumability has nothing to replay. Worse, the client's POST SSE stream +// closes CLEANLY, so the MCP client neither errors nor reconnects — it hangs +// until its own request timeout with no server-side exception and no exported +// span. Production traces show the reset itself (`Durable Object reset because +// its code was updated`) only on session-establishment spans; the killed +// execution never exports anything at all. +// +// ## The two halves of the fix +// +// 1. A PRIMING EVENT on every POST tools/call stream, so the MCP TS SDK's +// `hasPrimingEvent` is set and it will actually reconnect a broken POST +// stream (`canResume = isReconnectable || hasPrimingEvent` in the SDK's +// `_handleSseStream`). That half already ships in +// `patches/agents@0.17.3.patch` (`emitPrimingEvent`); without it the error +// this module writes would never be collected by anyone. +// 2. This reaper. On DO start, any POST stream that still has persisted +// in-flight request ids has, by definition, not been answered — its +// execution died with the previous isolate. We append a JSON-RPC error into +// that stream's event log and drop the request-id key, so the reconnecting +// client replays from `last-event-id` and receives a real error in seconds. +// +// Priming alone does NOT fix the hang (the reconnect finds nothing to replay); +// the reaper alone is never collected (the client never reconnects). Both are +// required. +// +// ## Scope +// +// This converts a silent hang into a prompt visible error. It does NOT recover +// the lost work — that is durable execution, and is deliberately out of scope. +// --------------------------------------------------------------------------- + +import { Cause, Data, Effect, Exit } from "effect"; +import type { EventStore } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import type { JSONRPCMessage, RequestId } from "@modelcontextprotocol/sdk/types.js"; + +/** + * The agents SDK's fixed stream id for the standalone GET listen stream. Its + * events are deliberately never auto-cleared and it carries no POST request + * ids, so it can never look orphaned — but it is skipped explicitly rather + * than relied on to be absent, so a future SDK change that starts recording + * request ids against it cannot make the reaper cancel a live listen stream. + */ +const STANDALONE_STREAM_ID = "_GET_stream"; + +/** + * JSON-RPC error code for a reaped request. `-32001` is the MCP SDK's + * `RequestTimeout`: the client-visible semantics are exactly right (the request + * will never be answered) and every MCP client already understands it, whereas a + * private code lands in generic error handling. + */ +export const RESTART_REAP_ERROR_CODE = -32001; + +/** + * The client-facing message for a reaped request. + * + * ## Retry safety — deliberately does NOT say "please retry" + * + * `execute` runs ARBITRARY user code. An execution killed by an isolate reset + * may have already performed real, non-idempotent side effects — a charge, an + * email, a write — before it died. The server cannot know how far it got: the + * only evidence was in the dead isolate. A message that invites a retry invites + * duplicating those effects, and an LLM client will take that invitation. + * + * So the message states what is known (the run was interrupted, the outcome is + * unknown, effects may have been applied) and hands the retry decision to the + * caller instead of pre-approving it. `data.retrySafe: false` carries the same + * fact machine-readably. + */ +export const RESTART_REAP_ERROR_MESSAGE = + "Execution was interrupted by a server restart and its outcome is unknown. " + + "Any side effects it had already performed have been applied. " + + "Do not retry automatically — verify the current state first, then re-run only if it is safe."; + +export class McpRestartReapError extends Data.TaggedError("McpRestartReapError")<{ + readonly streamId: string; + readonly requestId: RequestId; + readonly cause: unknown; +}> {} + +/** + * Raised (and immediately caught) purely so the reap ends its span as a + * FAILURE. Killed executions currently export no span at all, so their + * invisibility in o11y is half the reported bug; failing the span is what + * makes "a deploy ate N tool calls" a queryable error rather than an absence. + */ +export class McpRestartReapedExecutions extends Data.TaggedError("McpRestartReapedExecutions")<{ + readonly streamCount: number; + readonly requestCount: number; +}> { + override get message(): string { + return `Reaped ${this.requestCount} in-flight MCP request(s) across ${this.streamCount} stream(s) killed by a server restart`; + } +} + +/** + * The narrow slice of the agents `McpAgent` the reaper drives. + * + * THIS IS THE SEAM. `getOpenStreamRequestIds` is `@internal` in the agents SDK + * (added by `patches/agents@0.17.3.patch`) and the event store's `storeEvent` + * is public-ish. Nothing else in this module knows a storage key or prefix — no + * `__mcp_stream_reqs__:`, no `__mcp_event__::`, no + * zero-padding. An `agents` upgrade that moves this surface breaks HERE, at one + * structurally-typed interface with a compile error, rather than silently + * reading nothing from hand-rolled key scans. + */ +export interface RestartReapAgent { + readonly getOpenStreamRequestIds: () => Promise< + ReadonlyArray<{ readonly streamId: string; readonly requestIds: ReadonlyArray }> + >; + readonly deleteStreamRequestIds: (streamId: string) => Promise; +} + +/** The event-store slice the reaper needs: append one message to a stream. */ +export type RestartReapEventStore = Pick; + +/** + * Structurally narrow an unknown value to {@link RestartReapAgent}. + * + * The `agents` type declarations do not describe `getOpenStreamRequestIds` + * (the patch adds it to the JS dist only), so this is the single place the + * shape is asserted at runtime rather than assumed. When an upgrade drops the + * method the reaper reports `unavailable` and the DO starts normally, instead + * of throwing on every cold start. + */ +export const asRestartReapAgent = (candidate: unknown): RestartReapAgent | null => { + if (typeof candidate !== "object" || candidate === null) return null; + const agent = candidate as Partial>; + return typeof agent.getOpenStreamRequestIds === "function" && + typeof agent.deleteStreamRequestIds === "function" + ? (candidate as RestartReapAgent) + : null; +}; + +/** Structurally narrow an unknown value to {@link RestartReapEventStore}. */ +export const asRestartReapEventStore = (candidate: unknown): RestartReapEventStore | null => { + if (typeof candidate !== "object" || candidate === null) return null; + const store = candidate as { readonly storeEvent?: unknown }; + return typeof store.storeEvent === "function" ? (candidate as RestartReapEventStore) : null; +}; + +export const restartReapErrorMessage = (requestId: RequestId): JSONRPCMessage => ({ + jsonrpc: "2.0", + id: requestId, + error: { + code: RESTART_REAP_ERROR_CODE, + message: RESTART_REAP_ERROR_MESSAGE, + // Machine-readable companion to the prose above. `retrySafe: false` is the + // load-bearing field: a client that automates retries must not treat this + // like a transport blip. + data: { reason: "server_restart", retrySafe: false }, + }, +}); + +export type RestartReapOutcome = + | { readonly kind: "unavailable" } + | { + readonly kind: "reaped"; + readonly streamIds: ReadonlyArray; + readonly requestCount: number; + }; + +export const restartReapLog = (input: { + readonly sessionId: string; + readonly streamIds: ReadonlyArray; + readonly requestCount: number; +}): Record => ({ + event: "mcp_session_restart_reaped_requests", + sessionId: input.sessionId, + streamIds: [...input.streamIds], + streamCount: input.streamIds.length, + requestCount: input.requestCount, +}); + +/** + * Reap every POST stream left with in-flight request ids by a previous isolate. + * + * Correctness rests on one invariant of the patched agents transport: a POST + * stream's `__mcp_stream_reqs__` entry is written when the request is dispatched + * and deleted the moment its final response is written (`sendOnStream`'s + * `shouldClose` branch). So at DO start, a surviving entry means "dispatched, + * never answered" — which after a restart can only be work the reset destroyed. + * That is why this runs at START and nowhere else: mid-life the same entry + * legitimately means "still running", and reaping it would cancel live work. + * + * The error is appended AFTER the priming event already in the stream, so its + * event id sorts later and a `last-event-id: ` reconnect replays it. + * + * Per-request failures are collected rather than thrown: one unwritable stream + * must not strand the others, and the reaper must never block DO startup. + */ +export const reapOrphanedRequests = (input: { + readonly agent: RestartReapAgent; + readonly eventStore: RestartReapEventStore; +}): Effect.Effect => + Effect.gen(function* () { + const open = yield* Effect.promise(() => + Promise.resolve(input.agent.getOpenStreamRequestIds()), + ); + const orphaned = open.filter( + (entry) => entry.streamId !== STANDALONE_STREAM_ID && entry.requestIds.length > 0, + ); + if (orphaned.length === 0) { + return { kind: "reaped", streamIds: [], requestCount: 0 } as const; + } + + const reapedStreamIds: string[] = []; + let requestCount = 0; + const failures: McpRestartReapError[] = []; + + for (const entry of orphaned) { + const before = requestCount; + // Every request id on the stream is reaped, not just the first: a client + // may batch several JSON-RPC requests into one POST, and the transport + // only closes the stream once ALL of them have responded. Leaving one + // unanswered would hang the client exactly as before. + for (const requestId of entry.requestIds) { + const stored = yield* Effect.exit( + Effect.promise(() => + Promise.resolve( + input.eventStore.storeEvent(entry.streamId, restartReapErrorMessage(requestId)), + ), + ), + ); + if (Exit.isFailure(stored)) { + failures.push( + new McpRestartReapError({ + streamId: entry.streamId, + requestId, + cause: Cause.squash(stored.cause), + }), + ); + continue; + } + requestCount += 1; + } + if (requestCount === before) continue; + // Drop the request-id key only once its errors are durably stored, so a + // reset DURING the reap leaves the stream reapable on the next start + // rather than silently forgotten. Re-reaping is harmless: the client + // takes the first error and the stream is torn down. + yield* Effect.promise(() => + Promise.resolve(input.agent.deleteStreamRequestIds(entry.streamId)), + ); + reapedStreamIds.push(entry.streamId); + } + + const firstFailure = failures[0]; + if (firstFailure && reapedStreamIds.length === 0) return yield* firstFailure; + return { kind: "reaped", streamIds: reapedStreamIds, requestCount } as const; + });