Merge pull request #2045 from modelcontextprotocol/v2/fix/1999-sse-onabort-before-priming

fix: register SSE abort cleanup before the priming write

clients/web/src/test/integration/mcp/remote/sse-abort-cleanup.test.ts

@@ -0,0 +1,178 @@
+/**
+ * Regression tests for #1999 — an SSE client that disconnects *while the
+ * priming write is in flight* used to leak the stream's subscriber and park
+ * the handler forever.
+ *
+ * Both SSE handlers in `core/mcp/remote/node/server.ts` install their
+ * subscriber (session event consumer / server-list subscriber + file watcher)
+ * synchronously, then prime the stream so Firefox's `fetch()` resolves
+ * (#1858), and only then registered their `onAbort` cleanup. Priming is an
+ * `await`, and Hono's `onAbort` has no already-aborted replay — so a
+ * disconnect inside that window fired zero listeners: the subscriber stayed
+ * installed, the `mcp.json` watcher was never stopped, and the promise the
+ * handler was holding open on never resolved.
+ *
+ * `primeAndHoldSseStream` is the fix: one listener, registered before the
+ * first `await`, that both cleans up and releases the hold. These tests drive
+ * it with a stream whose priming write is controllable, which is the only way
+ * to hit that window deterministically — against a real Hono stream the write
+ * settles on its own microtask.
+ */
+
+import { describe, it, expect } from "vitest";
+import { StreamingApi } from "hono/utils/stream";
+import {
+  primeAndHoldSseStream,
+  type PrimableSseStream,
+} from "@inspector/core/mcp/remote/node/server.js";
+
+/** A stream whose priming write stays pending until the test releases it. */
+function controllableStream(): {
+  stream: PrimableSseStream;
+  /** Deliver the disconnect Hono would deliver. */
+  abort: () => void;
+  /** Settle the pending priming write. */
+  finishWrite: () => void;
+  writes: string[];
+  closes: number;
+} {
+  const listeners: (() => void)[] = [];
+  const writes: string[] = [];
+  let releaseWrite = (): void => {};
+  const state = { closes: 0 };
+  const stream: PrimableSseStream = {
+    write: (input: string) => {
+      writes.push(input);
+      return new Promise<void>((resolve) => {
+        releaseWrite = resolve;
+      });
+    },
+    onAbort: (listener: () => void) => {
+      listeners.push(listener);
+    },
+    close: async () => {
+      state.closes += 1;
+    },
+  };
+  return {
+    stream,
+    abort: () => {
+      // Hono fires only the subscribers registered at this moment, with a
+      // bare `subscriber()` call — see `hono/utils/stream`.
+      for (const l of [...listeners]) l();
+    },
+    finishWrite: () => releaseWrite(),
+    writes,
+    get closes() {
+      return state.closes;
+    },
+  };
+}
+
+/** Reject rather than hang if the handler is never released. */
+async function withinTick(promise: Promise<void>): Promise<void> {
+  const timeout = new Promise<never>((_, reject) => {
+    setTimeout(() => reject(new Error("handler was never released")), 500);
+  });
+  await Promise.race([promise, timeout]);
+}
+
+describe("primeAndHoldSseStream (#1999)", () => {
+  it("runs cleanup when the client aborts during the priming write", async () => {
+    const s = controllableStream();
+    let cleanups = 0;
+
+    const held = primeAndHoldSseStream(s.stream, () => {
+      cleanups += 1;
+    });
+
+    // The write is still pending — exactly the window the pre-fix code
+    // registered no listener in.
+    await Promise.resolve();
+    expect(s.writes).toEqual([":\n\n"]);
+    expect(cleanups).toBe(0);
+
+    s.abort();
+    expect(cleanups).toBe(1);
+
+    // The handler must also come back, or Hono never closes the stream.
+    s.finishWrite();
+    await withinTick(held);
+    expect(s.closes).toBe(1);
+  });
+
+  it("primes the stream before holding it open", async () => {
+    const s = controllableStream();
+    const held = primeAndHoldSseStream(s.stream, () => {});
+    await Promise.resolve();
+
+    expect(s.writes).toEqual([":\n\n"]);
+
+    s.finishWrite();
+    s.abort();
+    await withinTick(held);
+  });
+
+  it("holds the handler open until the abort, not merely until the write", async () => {
+    const s = controllableStream();
+    let released = false;
+    const held = primeAndHoldSseStream(s.stream, () => {}).then(() => {
+      released = true;
+    });
+
+    s.finishWrite();
+    await new Promise((resolve) => setTimeout(resolve, 10));
+    expect(released).toBe(false);
+
+    s.abort();
+    await withinTick(held);
+    expect(released).toBe(true);
+  });
+
+  it("runs cleanup once, on the first abort", async () => {
+    const s = controllableStream();
+    let cleanups = 0;
+    const held = primeAndHoldSseStream(s.stream, () => {
+      cleanups += 1;
+    });
+
+    s.finishWrite();
+    s.abort();
+    // Hono guards re-entry itself, but the helper must not depend on that —
+    // so drive the listener a second time directly.
+    s.abort();
+    await withinTick(held);
+    expect(cleanups).toBe(1);
+    expect(s.closes).toBe(1);
+  });
+
+  it("does not reject when cleanup's own close loses the race", async () => {
+    const s = controllableStream();
+    const stream: PrimableSseStream = {
+      ...s.stream,
+      close: () => Promise.reject(new Error("peer already gone")),
+    };
+    const held = primeAndHoldSseStream(stream, () => {});
+
+    s.finishWrite();
+    s.abort();
+    await expect(withinTick(held)).resolves.toBeUndefined();
+  });
+
+  it("documents the Hono behavior it exists for: onAbort has no replay", async () => {
+    // The premise of the whole fix. If Hono ever replays the abort to a
+    // late subscriber, this fails and the ordering guard can be revisited.
+    const { readable, writable } = new TransformStream();
+    const stream = new StreamingApi(writable, readable);
+
+    await stream.responseReadable.cancel();
+
+    let late = 0;
+    stream.onAbort(() => {
+      late += 1;
+    });
+    await new Promise((resolve) => setTimeout(resolve, 10));
+
+    expect(late).toBe(0);
+  });
+});

core/mcp/remote/node/server.ts

@@ -115,6 +115,57 @@ export function closeAbortedStream(stream: { close(): Promise<void> }): void {
   stream.close().catch(() => {});
 }
 
+/** The slice of Hono's `StreamingApi` the SSE prime-and-hold helper needs. */
+export interface PrimableSseStream {
+  write(input: string): Promise<unknown>;
+  onAbort(listener: () => void): void;
+  close(): Promise<void>;
+}
+
+/**
+ * Prime an SSE stream, then hold the handler open until the client aborts,
+ * running `cleanup` exactly once when it does.
+ *
+ * The registration order is the whole point (#1999). Hono's `onAbort` is a
+ * plain push onto a subscriber array with **no already-aborted replay**, and
+ * `abort()` fires only the subscribers registered at that moment
+ * (`hono/utils/stream`). Priming is an `await`, so a listener registered
+ * *after* it never runs if the client disconnects while that write is in
+ * flight — leaving the caller's subscriber/consumer installed forever and the
+ * handler parked on a promise nothing will resolve.
+ *
+ * So the listener goes in before the first `await`, and the same one both
+ * cleans up and releases the hold. Note this must not be inlined back into a
+ * caller in a way that reintroduces an `await` ahead of the registration —
+ * priming itself cannot move earlier, because callers deliberately install
+ * their subscriber before the first bytes go out (see GET /api/servers/events).
+ *
+ * `cleanup` runs inside Hono's bare `subscriber()` call, so it must be
+ * synchronous and must not throw — a rejection there has no owner.
+ */
+export async function primeAndHoldSseStream(
+  stream: PrimableSseStream,
+  cleanup: () => void,
+): Promise<void> {
+  let done = false;
+  const aborted = new Promise<void>((resolve) => {
+    stream.onAbort(() => {
+      // Hono's `abort()` guards its own re-entry, but the guard belongs here
+      // too: this helper's contract is exactly-once teardown, and running a
+      // caller's cleanup twice would double-delete a subscriber that has
+      // since been re-added by a reconnect.
+      if (done) return;
+      done = true;
+      cleanup();
+      closeAbortedStream(stream);
+      resolve();
+    });
+  });
+
+  await stream.write(SSE_PRIMING_COMMENT);
+  await aborted;
+}
+
 /**
  * Shape of the initial config returned by GET /api/config (defaults for client).
  */
@@ -855,6 +906,10 @@ export function createRemoteApp(
       // drained the queued stderr + the transport_error event. Yield once
       // so the writeSSE writes flush, then return — closing the stream and
       // surfacing the real error instead of a bare 404 / silent hang.
+      //
+      // This branch needs no abort listener despite its `await`: it performs
+      // the same teardown unconditionally on the way out, so a disconnect
+      // during the yield changes nothing.
       if (session.isTransportDead()) {
         await new Promise((resolve) => setTimeout(resolve, 0));
         session.clearEventConsumer();
@@ -864,29 +919,18 @@ export function createRemoteApp(
 
       // Prime only after the consumer is registered, so nothing observable
       // to the client happens before this stream can actually report
-      // events. See SSE_PRIMING_COMMENT.
-      await stream.write(SSE_PRIMING_COMMENT);
-
-      stream.onAbort(() => {
+      // events. See SSE_PRIMING_COMMENT. The disconnect cleanup is
+      // registered ahead of that priming write and the stream is then held
+      // open until the client aborts — see primeAndHoldSseStream (#1999).
+      await primeAndHoldSseStream(stream, () => {
         // Client disconnected - clear event consumer
         const shouldCleanup = session.clearEventConsumer();
-        closeAbortedStream(stream);
 
         // If transport is dead and no client connected, cleanup session
         if (shouldCleanup || session.isTransportDead()) {
           sessions.delete(sessionId);
         }
       });
-
-      // Keep the stream open until the client disconnects. Hono's streamSSE
-      // closes the stream when this callback returns, so we must not return
-      // until the connection is aborted.
-      await new Promise<void>((resolve) => {
-        stream.onAbort(() => {
-          // Cleanup happens in onAbort handler above
-          resolve();
-        });
-      });
     });
   });
 
@@ -2497,20 +2541,15 @@ export function createRemoteApp(
       // registered and the watcher started: callers treat the arrival of
       // this stream's first bytes as proof they are subscribed, so priming
       // first would hand out that proof across an `await`, before the
-      // watcher exists, and drop an edit made in the gap.
-      await stream.write(SSE_PRIMING_COMMENT);
-
-      stream.onAbort(() => {
+      // watcher exists, and drop an edit made in the gap. The unsubscribe is
+      // nonetheless registered *before* that write, and the stream is held
+      // open until the client aborts — see primeAndHoldSseStream (#1999).
+      await primeAndHoldSseStream(stream, () => {
         serverEventSubscribers.delete(send);
+        // Voided rather than awaited: this runs inside Hono's synchronous
+        // abort subscriber, which cannot await. maybeStopWatcher owns its
+        // own failures (it swallows a close() that throws).
         void maybeStopWatcher();
-        closeAbortedStream(stream);
-      });
-
-      // Hono closes the stream the moment this callback returns, so hold the
-      // promise open until the client aborts. Cleanup happens in the
-      // onAbort handler registered above.
-      await new Promise<void>((resolve) => {
-        stream.onAbort(() => resolve());
       });
     });
   });