Handle MCP disconnect cleanup

apps/cli/tests/mcp-subcommand.test.ts

@@ -5,6 +5,24 @@ import { describe, expect, it } from "vite-plus/test";
 
 const CLI_BIN = path.resolve(__dirname, "../dist/index.js");
 const MCP_BIN = path.resolve(__dirname, "../dist/browser-mcp.js");
+const MCP_PROCESS_EXIT_TIMEOUT_MS = 3_000;
+
+const waitForExit = (child: ReturnType<typeof spawn>, timeoutMs = MCP_PROCESS_EXIT_TIMEOUT_MS) =>
+  new Promise<{ code: number | null; signal: NodeJS.Signals | null }>((resolve, reject) => {
+    let settled = false;
+    const timeout = setTimeout(() => {
+      if (settled) return;
+      settled = true;
+      reject(new Error(`Process ${child.pid} did not exit within ${timeoutMs}ms`));
+    }, timeoutMs);
+
+    child.once("exit", (code, signal) => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timeout);
+      resolve({ code, signal });
+    });
+  });
 
 describe("mcp subcommand", () => {
   it("browser-mcp.js exists in dist", () => {
@@ -59,4 +77,25 @@ describe("mcp subcommand", () => {
     child.kill();
     await new Promise<void>((resolve) => child.on("exit", resolve));
   });
+
+  it("direct browser-mcp.js exits when stdin closes", async () => {
+    const child = spawn(process.execPath, [MCP_BIN], {
+      stdio: ["pipe", "pipe", "pipe"],
+    });
+
+    let stderr = "";
+    child.stderr?.on("data", (chunk: Buffer) => {
+      stderr += chunk.toString();
+    });
+
+    await new Promise((resolve) => setTimeout(resolve, 500));
+
+    expect(child.exitCode).toBeNull();
+    child.stdin?.end();
+
+    const result = await waitForExit(child);
+    expect(result.code).toBe(0);
+    expect(result.signal).toBeNull();
+    expect(stderr).not.toContain("cleanup failed");
+  });
 });

packages/browser/src/mcp/register-process-cleanup.ts

@@ -0,0 +1,92 @@
+const UNIX_SIGNALS = ["SIGINT", "SIGTERM", "SIGHUP"] as const;
+const WINDOWS_SIGNALS = ["SIGINT", "SIGTERM", "SIGBREAK"] as const;
+
+type SignalShutdownReason = (typeof UNIX_SIGNALS)[number] | (typeof WINDOWS_SIGNALS)[number];
+
+export type ShutdownReason =
+  | SignalShutdownReason
+  | "beforeExit"
+  | "disconnect"
+  | "stdin-close"
+  | "stdin-end";
+
+interface RegisterProcessCleanupOptions {
+  readonly cleanup: (reason: ShutdownReason) => Promise<void>;
+  readonly afterCleanup?: (reason: ShutdownReason) => Promise<void> | void;
+  readonly watchStdin?: boolean;
+}
+
+let cleanupRegistered = false;
+
+const SIGNALS: readonly SignalShutdownReason[] =
+  process.platform === "win32" ? WINDOWS_SIGNALS : UNIX_SIGNALS;
+
+const formatError = (error: unknown) =>
+  error instanceof Error ? (error.stack ?? error.message) : String(error);
+
+const writeShutdownError = (
+  stage: "cleanup" | "afterCleanup",
+  reason: ShutdownReason,
+  error: unknown,
+) => {
+  process.stderr.write(`expect mcp ${stage} failed during ${reason}: ${formatError(error)}\n`);
+};
+
+export const registerProcessCleanup = (options: RegisterProcessCleanupOptions) => {
+  if (cleanupRegistered) return;
+  cleanupRegistered = true;
+
+  let exitAfterCleanup = false;
+  let shutdownPromise: Promise<void> | undefined;
+
+  const requestShutdown = (reason: ShutdownReason, shouldExit: boolean) => {
+    exitAfterCleanup = exitAfterCleanup || shouldExit;
+    if (shutdownPromise) return;
+
+    shutdownPromise = options
+      .cleanup(reason)
+      .catch((error) => {
+        writeShutdownError("cleanup", reason, error);
+      })
+      .then(() => options.afterCleanup?.(reason))
+      .catch((error) => {
+        writeShutdownError("afterCleanup", reason, error);
+      })
+      .then(() => {
+        if (exitAfterCleanup) {
+          process.exit(0);
+        }
+      });
+  };
+
+  for (const signal of SIGNALS) {
+    process.once(signal as NodeJS.Signals, () => {
+      requestShutdown(signal, true);
+    });
+  }
+
+  process.once("beforeExit", () => {
+    requestShutdown("beforeExit", false);
+  });
+
+  process.once("disconnect", () => {
+    requestShutdown("disconnect", true);
+  });
+
+  if (options.watchStdin) {
+    if (process.stdin.readableEnded) {
+      requestShutdown("stdin-end", true);
+      return;
+    }
+    if (process.stdin.destroyed) {
+      requestShutdown("stdin-close", true);
+      return;
+    }
+    process.stdin.once("end", () => {
+      requestShutdown("stdin-end", true);
+    });
+    process.stdin.once("close", () => {
+      requestShutdown("stdin-close", true);
+    });
+  }
+};

packages/browser/src/mcp/start-http.ts

@@ -2,6 +2,7 @@ import * as fs from "node:fs";
 import * as http from "node:http";
 import { Effect, Predicate } from "effect";
 import { McpSession } from "./mcp-session";
+import { registerProcessCleanup } from "./register-process-cleanup";
 import { McpRuntime } from "./runtime";
 import { createBrowserMcpServer } from "./server";
 import { CLI_SESSION_FILE, MAX_DAEMON_REQUEST_BODY_BYTES } from "./constants";
@@ -89,17 +90,20 @@ const closeSession = Effect.gen(function* () {
   yield* session.close();
 });
 
-const shutdown = () => {
-  void McpRuntime.runPromise(closeSession).finally(() => {
+registerProcessCleanup({
+  cleanup: () => McpRuntime.runPromise(closeSession),
+  afterCleanup: async () => {
     removeSessionFile();
-    process.exit(0);
-  });
-};
-
-process.once("SIGINT", shutdown);
-process.once("SIGTERM", shutdown);
-process.once("beforeExit", () => {
-  void McpRuntime.runPromise(closeSession).finally(removeSessionFile);
+    await new Promise<void>((resolve, reject) => {
+      httpServer.close((error) => {
+        if (error) {
+          reject(error);
+          return;
+        }
+        resolve();
+      });
+    });
+  },
 });
 
 httpServer.listen(0, "127.0.0.1", () => {

packages/browser/src/mcp/start.ts

@@ -1,29 +1,17 @@
 import { Effect } from "effect";
 import { McpSession } from "./mcp-session";
+import { registerProcessCleanup } from "./register-process-cleanup";
 import { McpRuntime } from "./runtime";
 import { startBrowserMcpServer } from "./server";
 
-let cleanupRegistered = false;
-
 const closeSession = Effect.gen(function* () {
   const session = yield* McpSession;
   yield* session.close();
 });
 
-const registerProcessCleanup = () => {
-  if (cleanupRegistered) return;
-  cleanupRegistered = true;
-
-  const handleShutdown = () => {
-    void McpRuntime.runPromise(closeSession).finally(() => process.exit(0));
-  };
-
-  process.once("SIGINT", handleShutdown);
-  process.once("SIGTERM", handleShutdown);
-  process.once("beforeExit", () => {
-    void McpRuntime.runPromise(closeSession);
-  });
-};
+registerProcessCleanup({
+  cleanup: () => McpRuntime.runPromise(closeSession),
+  watchStdin: true,
+});
 
-registerProcessCleanup();
 void startBrowserMcpServer(McpRuntime);