mcp: exit the daemon process after shutdown (#98)

Author: benvinegar

src/main.tsx

@@ -21,8 +21,9 @@ async function main() {
   }
 
   if (startupPlan.kind === "mcp-serve") {
-    serveHunkMcpServer();
-    await new Promise<never>(() => {});
+    const server = serveHunkMcpServer();
+    await server.stopped;
+    return;
   }
 
   if (startupPlan.kind === "session-command") {

src/mcp/server.ts

@@ -33,6 +33,10 @@ export interface ServeHunkMcpServerOptions {
   staleSessionSweepIntervalMs?: number;
 }
 
+export type RunningHunkMcpServer = ReturnType<typeof Bun.serve<{}>> & {
+  stopped: Promise<void>;
+};
+
 function formatDaemonServeError(error: unknown, host: string, port: number) {
   const message = error instanceof Error ? error.message : String(error);
   const normalized = message.toLowerCase();
@@ -161,14 +165,18 @@ async function handleSessionApiRequest(state: HunkDaemonState, request: Request)
 }
 
 /** Serve the local Hunk session daemon and websocket session broker. */
-export function serveHunkMcpServer(options: ServeHunkMcpServerOptions = {}) {
+export function serveHunkMcpServer(options: ServeHunkMcpServerOptions = {}): RunningHunkMcpServer {
   const config = resolveHunkMcpConfig();
   const idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS;
   const staleSessionTtlMs = options.staleSessionTtlMs ?? DEFAULT_STALE_SESSION_TTL_MS;
   const staleSessionSweepIntervalMs =
     options.staleSessionSweepIntervalMs ?? DEFAULT_STALE_SESSION_SWEEP_INTERVAL_MS;
   const state = new HunkDaemonState();
   const startedAt = Date.now();
+  let resolveStopped: (() => void) | null = null;
+  const stopped = new Promise<void>((resolve) => {
+    resolveStopped = resolve;
+  });
   let lastActivityAt = startedAt;
   let shuttingDown = false;
   let sweepTimer: Timer | null = null;
@@ -203,6 +211,8 @@ export function serveHunkMcpServer(options: ServeHunkMcpServerOptions = {}) {
 
     state.shutdown();
     server?.stop(true);
+    resolveStopped?.();
+    resolveStopped = null;
   };
 
   const refreshIdleShutdownTimer = () => {
@@ -354,5 +364,5 @@ export function serveHunkMcpServer(options: ServeHunkMcpServerOptions = {}) {
   console.log(`Hunk session daemon listening on ${config.httpOrigin}${HUNK_SESSION_API_PATH}`);
   console.log(`Hunk session websocket listening on ${config.wsOrigin}${HUNK_SESSION_SOCKET_PATH}`);
 
-  return server;
+  return Object.assign(server, { stopped }) as RunningHunkMcpServer;
 }

test/mcp-serve-process.test.ts

@@ -0,0 +1,100 @@
+import { afterEach, describe, expect, test } from "bun:test";
+import type { Subprocess } from "bun";
+import { createServer } from "node:net";
+
+const repoRoot = process.cwd();
+const spawned: Subprocess[] = [];
+
+async function reserveLoopbackPort() {
+  const listener = createServer(() => undefined);
+  await new Promise<void>((resolve, reject) => {
+    listener.once("error", reject);
+    listener.listen(0, "127.0.0.1", () => resolve());
+  });
+
+  const address = listener.address();
+  const port = typeof address === "object" && address ? address.port : 0;
+  await new Promise<void>((resolve) => listener.close(() => resolve()));
+  return port;
+}
+
+async function waitUntil<T>(
+  label: string,
+  fn: () => Promise<T | null> | T | null,
+  timeoutMs = 1_500,
+  intervalMs = 20,
+) {
+  const deadline = Date.now() + timeoutMs;
+
+  for (;;) {
+    const value = await fn();
+    if (value !== null) {
+      return value;
+    }
+
+    if (Date.now() >= deadline) {
+      throw new Error(`Timed out waiting for ${label}.`);
+    }
+
+    await Bun.sleep(intervalMs);
+  }
+}
+
+async function readHealth(port: number) {
+  try {
+    const response = await fetch(`http://127.0.0.1:${port}/health`);
+    if (!response.ok) {
+      return null;
+    }
+
+    return (await response.json()) as { ok: boolean; pid: number };
+  } catch {
+    return null;
+  }
+}
+
+afterEach(async () => {
+  await Promise.allSettled(
+    spawned.splice(0).map(async (proc) => {
+      try {
+        proc.kill();
+      } catch {
+        // Ignore processes that already exited.
+      }
+
+      await proc.exited.catch(() => undefined);
+    }),
+  );
+});
+
+describe("mcp serve process lifecycle", () => {
+  test("exits cleanly after SIGTERM instead of hot-looping after server shutdown", async () => {
+    const port = await reserveLoopbackPort();
+    const proc = Bun.spawn(["bun", "run", "src/main.tsx", "mcp", "serve"], {
+      cwd: repoRoot,
+      stdin: "ignore",
+      stdout: "pipe",
+      stderr: "pipe",
+      env: {
+        ...process.env,
+        HUNK_MCP_PORT: String(port),
+      },
+    });
+    spawned.push(proc);
+
+    const health = await waitUntil("daemon health", () => readHealth(port), 3_000, 50);
+    expect(health).toMatchObject({ ok: true, pid: proc.pid });
+
+    let exited = false;
+    void proc.exited.then(() => {
+      exited = true;
+    });
+
+    process.kill(proc.pid, "SIGTERM");
+
+    await waitUntil("mcp serve process exit", () => (exited ? true : null), 1_500, 25);
+    await waitUntil("daemon port close", async () =>
+      (await readHealth(port)) === null ? true : null,
+    );
+  }, 10_000);
+});