diff --git a/apps/cli/tests/mcp-subcommand.test.ts b/apps/cli/tests/mcp-subcommand.test.ts index 1408dc28e..b44192ac5 100644 --- a/apps/cli/tests/mcp-subcommand.test.ts +++ b/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, 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((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"); + }); }); diff --git a/packages/browser/src/mcp/register-process-cleanup.ts b/packages/browser/src/mcp/register-process-cleanup.ts new file mode 100644 index 000000000..51443b082 --- /dev/null +++ b/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; + readonly afterCleanup?: (reason: ShutdownReason) => Promise | 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 | 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); + }); + } +}; diff --git a/packages/browser/src/mcp/start-http.ts b/packages/browser/src/mcp/start-http.ts index d4ec77b8c..4c6b0185a 100644 --- a/packages/browser/src/mcp/start-http.ts +++ b/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((resolve, reject) => { + httpServer.close((error) => { + if (error) { + reject(error); + return; + } + resolve(); + }); + }); + }, }); httpServer.listen(0, "127.0.0.1", () => { diff --git a/packages/browser/src/mcp/start.ts b/packages/browser/src/mcp/start.ts index c0052ba1a..3a4e53747 100644 --- a/packages/browser/src/mcp/start.ts +++ b/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);