Fail loudly when a DO restart kills an in-flight MCP tool call

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": {

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<string, unknown>();
@@ -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<T>(
-    options: { readonly prefix?: string; readonly limit?: number } = {},
+    options: {
+      readonly prefix?: string;
+      readonly start?: string;
+      readonly limit?: number;
+      readonly reverse?: boolean;
+    } = {},
   ): Promise<Map<string, T>> {
-    const rows = new Map<string, T>();
-    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<T>(callback: () => T | Promise<T>): Promise<T> {
@@ -93,6 +110,11 @@ type HarnessSession = {
   props: Record<string, unknown>;
   runMcpAgentOnStart: () => Promise<void>;
   server?: McpServer;
+  /** Inherited from the real `McpAgent`; the reaper tests drive them directly. */
+  setStreamRequestIds: (streamId: string, requestIds: RequestId[]) => Promise<void>;
+  getStreamRequestIds: (streamId: string) => Promise<RequestId[] | undefined>;
+  /** The base's telemetry seam; the span test swaps in a recording tracer. */
+  withTelemetry: <A, E>(effect: Effect.Effect<A, E>) => Effect.Effect<A, E>;
   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<string, unknown>;
+      exit?: Exit.Exit<unknown, unknown>;
+    }> = [];
+    const tracer: Tracer.Tracer = {
+      span: (options) => {
+        const attributes = new Map<string, unknown>();
+        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<ReadonlyArray<JSONRPCMessage>> => {
+    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<RequestId>,
+  ): Promise<string> => {
+    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]);
+  });
+});

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<void> {
+    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<void> {
     return super.onStart(props);
   }