Shut down on stdin EOF/close so host death can't orphan LSP children

src/__tests__/__shutdown_fixtures__/dummy-lsp.mjs

@@ -0,0 +1,32 @@
+import { writeFileSync } from "node:fs";
+
+writeFileSync(process.argv[2], String(process.pid));
+
+function send(message) {
+  const json = JSON.stringify(message);
+  process.stdout.write(`Content-Length: ${Buffer.byteLength(json)}\r\n\r\n${json}`);
+}
+
+let buffer = Buffer.alloc(0);
+
+process.stdin.on("data", (chunk) => {
+  buffer = Buffer.concat([buffer, chunk]);
+
+  for (;;) {
+    const separator = buffer.indexOf("\r\n\r\n");
+    if (separator === -1) return;
+
+    const length = Number(/Content-Length: (\d+)/i.exec(buffer.subarray(0, separator).toString())?.[1]);
+    const bodyStart = separator + 4;
+    if (buffer.length < bodyStart + length) return;
+
+    const message = JSON.parse(buffer.subarray(bodyStart, bodyStart + length).toString());
+    buffer = buffer.subarray(bodyStart + length);
+
+    if (message.method === "initialize") {
+      send({ jsonrpc: "2.0", id: message.id, result: { capabilities: { textDocumentSync: 1 } } });
+    }
+  }
+});
+
+setInterval(() => {}, 1000);

src/__tests__/index.shutdown.test.ts

@@ -0,0 +1,138 @@
+import { describe, it, expect, beforeAll, beforeEach, afterEach } from "vitest";
+import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
+import { once } from "node:events";
+import * as fs from "node:fs";
+import * as path from "node:path";
+import { createInterface } from "node:readline";
+import { setTimeout as sleep } from "node:timers/promises";
+
+const ENTRY = path.resolve(import.meta.dirname, "../../dist/index.js");
+const FIXTURES = path.join(import.meta.dirname, "__shutdown_fixtures__");
+const DUMMY_LSP = path.join(FIXTURES, "dummy-lsp.mjs");
+const PROJECT = path.join(FIXTURES, "project");
+const EMPTY_PROJECT = path.join(FIXTURES, "empty-project");
+const PID_FILE = path.join(PROJECT, "dummy-lsp.pid");
+
+let running: ChildProcessWithoutNullStreams | undefined;
+
+function startSquiggles(projectDir: string) {
+  const child = spawn(process.execPath, [ENTRY, projectDir]);
+  const answered = new Set<number>();
+  let stderr = "";
+
+  running = child;
+  child.stderr.setEncoding("utf8").on("data", (chunk: string) => {
+    stderr += chunk;
+  });
+  createInterface({ input: child.stdout }).on("line", (line) => {
+    const message = JSON.parse(line);
+    if (typeof message.id === "number") answered.add(message.id);
+  });
+
+  return {
+    child,
+    answered,
+    get stderr() {
+      return stderr;
+    },
+    send(message: object) {
+      child.stdin.write(`${JSON.stringify(message)}\n`);
+    },
+  };
+}
+
+async function exitCode(child: ChildProcessWithoutNullStreams, ms: number): Promise<number | null> {
+  const [code] = await once(child, "close", { signal: AbortSignal.timeout(ms) });
+  return code;
+}
+
+async function waitFor(condition: () => boolean, ms: number, what: string): Promise<void> {
+  const deadline = Date.now() + ms;
+  while (!condition()) {
+    if (Date.now() > deadline) throw new Error(`Timed out waiting for ${what}`);
+    await sleep(50);
+  }
+}
+
+function lspPid(): number | undefined {
+  try {
+    return Number(fs.readFileSync(PID_FILE, "utf8")) || undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function alive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+beforeAll(() => {
+  if (!fs.existsSync(ENTRY)) throw new Error(`${ENTRY} not found, run npm run build first`);
+});
+
+beforeEach(() => {
+  fs.mkdirSync(PROJECT, { recursive: true });
+  fs.mkdirSync(EMPTY_PROJECT, { recursive: true });
+  fs.writeFileSync(
+    path.join(PROJECT, "squiggles.yaml"),
+    [
+      "servers:",
+      "  dummy:",
+      `    command: ${JSON.stringify([process.execPath, DUMMY_LSP, PID_FILE])}`,
+      `    filePatterns: ["*.txt"]`,
+      "",
+    ].join("\n"),
+  );
+  fs.writeFileSync(path.join(PROJECT, "sample.txt"), "hello\n");
+});
+
+afterEach(() => {
+  running?.kill("SIGKILL");
+  const pid = lspPid();
+  if (pid && alive(pid)) process.kill(pid, "SIGKILL");
+  fs.rmSync(PROJECT, { recursive: true, force: true });
+  fs.rmSync(EMPTY_PROJECT, { recursive: true, force: true });
+});
+
+describe("when stdin closes", () => {
+  it("exits and kills its LSP servers", async () => {
+    const squiggles = startSquiggles(PROJECT);
+
+    squiggles.send({
+      jsonrpc: "2.0",
+      id: 1,
+      method: "initialize",
+      params: { protocolVersion: "2024-11-05", capabilities: {}, clientInfo: { name: "test", version: "0" } },
+    });
+    await waitFor(() => squiggles.answered.has(1), 5000, "initialize response");
+    squiggles.send({ jsonrpc: "2.0", method: "notifications/initialized" });
+    squiggles.send({
+      jsonrpc: "2.0",
+      id: 2,
+      method: "tools/call",
+      params: { name: "document_symbols", arguments: { file: "sample.txt" } },
+    });
+    await waitFor(() => lspPid() !== undefined, 10_000, "dummy LSP to start");
+    const pid = lspPid()!;
+    expect(alive(pid)).toBe(true);
+
+    squiggles.child.stdin.end();
+
+    expect(await exitCode(squiggles.child, 15_000)).toBe(0);
+    await waitFor(() => !alive(pid), 3000, "dummy LSP to exit");
+  }, 30_000);
+
+  it("exits through shutdown before any requests", async () => {
+    const squiggles = startSquiggles(EMPTY_PROJECT);
+
+    squiggles.child.stdin.end();
+
+    expect(await exitCode(squiggles.child, 5000)).toBe(0);
+    expect(squiggles.stderr).toContain("Shutting down...");
+  }, 15_000);
+});

src/index.ts

@@ -20,6 +20,10 @@ const require = createRequire(import.meta.url);
 const { version } = require("../package.json");
 
 async function main(): Promise<void> {
+  const ignoreWriteErrors = () => {};
+  process.stdout.on("error", ignoreWriteErrors);
+  process.stderr.on("error", ignoreWriteErrors);
+
   if (process.argv[2] === "init") {
     const result = runInit(process.argv.slice(3), process.cwd());
     process.stdout.write(result.message + "\n");
@@ -80,18 +84,31 @@ async function main(): Promise<void> {
   });
 
   const transport = new StdioServerTransport();
-  await server.connect(transport);
-  log("MCP server connected via stdio");
 
+  let shuttingDown = false;
   const shutdown = async () => {
+    if (shuttingDown) return;
+    shuttingDown = true;
     log("Shutting down...");
     if (manager) await manager.shutdownAll();
     await server.close();
     process.exit(0);
   };
+  const requestShutdown = () => {
+    shutdown().catch((err) => {
+      logError("Shutdown failed", err);
+      process.exit(1);
+    });
+  };
 
-  process.on("SIGINT", shutdown);
-  process.on("SIGTERM", shutdown);
+  process.on("SIGINT", requestShutdown);
+  process.on("SIGTERM", requestShutdown);
+  process.stdin.on("end", requestShutdown);
+  process.stdin.on("close", requestShutdown);
+  server.onclose = requestShutdown;
+
+  await server.connect(transport);
+  log("MCP server connected via stdio");
 }
 
 main().catch((err) => {