Merge fix/exit-on-stdio-close: exit on MCP stdio disconnect (upstream PR ktnyt#53)

Author: sudhirj

index.ts

@@ -4,6 +4,7 @@ import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
 import { logger } from './src/logger.js';
 import { LSPClient } from './src/lsp-client.js';
+import { installShutdownHandlers } from './src/shutdown.js';
 import { diagnosticsTools } from './src/tools/diagnostics.js';
 import { hoverTools } from './src/tools/hover.js';
 import { navigationTools } from './src/tools/navigation.js';
@@ -57,15 +58,8 @@ const allTools = [
 
 registerTools(server, allTools, lspClient);
 
-process.on('SIGINT', () => {
-  lspClient.dispose();
-  process.exit(0);
-});
-
-process.on('SIGTERM', () => {
-  lspClient.dispose();
-  process.exit(0);
-});
+// Covers SIGINT/SIGTERM/SIGHUP and the MCP client closing its end of stdio.
+const shutdown = installShutdownHandlers(lspClient);
 
 async function main() {
   const transport = new StdioServerTransport();
@@ -82,6 +76,5 @@ async function main() {
 
 main().catch((error) => {
   logger.error(`Server error: ${error}\n`);
-  lspClient.dispose();
-  process.exit(1);
+  shutdown(1);
 });

src/shutdown.test.ts

@@ -0,0 +1,115 @@
+import { describe, expect, it, jest } from 'bun:test';
+import { installShutdownHandlers } from './shutdown.js';
+import type { ShutdownProcess } from './shutdown.js';
+
+/** A stand-in for `process` that records handlers so tests can fire them. */
+function createFakeProcess() {
+  const signalHandlers = new Map<string, Array<() => void>>();
+  const stdinHandlers = new Map<string, Array<() => void>>();
+  const exit = jest.fn();
+
+  const record = (map: Map<string, Array<() => void>>, event: string, listener: () => void) => {
+    const existing = map.get(event) ?? [];
+    existing.push(listener);
+    map.set(event, existing);
+  };
+
+  const proc: ShutdownProcess = {
+    on(event: string, listener: () => void) {
+      record(signalHandlers, event, listener);
+      return proc;
+    },
+    stdin: {
+      on(event: string, listener: () => void) {
+        record(stdinHandlers, event, listener);
+        return proc.stdin;
+      },
+    },
+    exit,
+  };
+
+  return {
+    proc,
+    exit,
+    fireSignal: (event: string) => {
+      for (const handler of signalHandlers.get(event) ?? []) handler();
+    },
+    fireStdin: (event: string) => {
+      for (const handler of stdinHandlers.get(event) ?? []) handler();
+    },
+    hasSignal: (event: string) => signalHandlers.has(event),
+    hasStdin: (event: string) => stdinHandlers.has(event),
+  };
+}
+
+describe('installShutdownHandlers', () => {
+  it('subscribes to the termination signals', () => {
+    const fake = createFakeProcess();
+    installShutdownHandlers({ dispose: jest.fn() }, fake.proc);
+
+    expect(fake.hasSignal('SIGINT')).toBe(true);
+    expect(fake.hasSignal('SIGTERM')).toBe(true);
+    expect(fake.hasSignal('SIGHUP')).toBe(true);
+  });
+
+  // The regression: an MCP client that exits closes the pipe without
+  // signalling, and StdioServerTransport never surfaces EOF.
+  it('disposes and exits when stdin ends', () => {
+    const fake = createFakeProcess();
+    const dispose = jest.fn();
+    installShutdownHandlers({ dispose }, fake.proc);
+
+    expect(fake.hasStdin('end')).toBe(true);
+    fake.fireStdin('end');
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+    expect(fake.exit).toHaveBeenCalledWith(0);
+  });
+
+  it('disposes and exits when stdin closes', () => {
+    const fake = createFakeProcess();
+    const dispose = jest.fn();
+    installShutdownHandlers({ dispose }, fake.proc);
+
+    fake.fireStdin('close');
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+    expect(fake.exit).toHaveBeenCalledWith(0);
+  });
+
+  it.each(['SIGINT', 'SIGTERM', 'SIGHUP'])('disposes and exits on %s', (signal) => {
+    const fake = createFakeProcess();
+    const dispose = jest.fn();
+    installShutdownHandlers({ dispose }, fake.proc);
+
+    fake.fireSignal(signal);
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+    expect(fake.exit).toHaveBeenCalledWith(0);
+  });
+
+  // stdin emits both 'end' and 'close', and signals can race each other.
+  it('disposes once even when several shutdown triggers fire', () => {
+    const fake = createFakeProcess();
+    const dispose = jest.fn();
+    installShutdownHandlers({ dispose }, fake.proc);
+
+    fake.fireStdin('end');
+    fake.fireStdin('close');
+    fake.fireSignal('SIGTERM');
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+    expect(fake.exit).toHaveBeenCalledTimes(1);
+  });
+
+  it('returns a shutdown function that follows the same path', () => {
+    const fake = createFakeProcess();
+    const dispose = jest.fn();
+    const shutdown = installShutdownHandlers({ dispose }, fake.proc);
+
+    shutdown(1);
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+    expect(fake.exit).toHaveBeenCalledWith(1);
+  });
+});

src/shutdown.ts

@@ -0,0 +1,62 @@
+/** Anything holding resources that must be released before exit. */
+export interface Disposable {
+  dispose(): void;
+}
+
+/**
+ * The parts of `process` this module touches, named structurally so tests can
+ * pass a fake.
+ */
+export interface ShutdownProcess {
+  on(event: string, listener: () => void): unknown;
+  stdin: { on(event: string, listener: () => void): unknown };
+  exit(code?: number): unknown;
+}
+
+/** Signals that mean "stop now". */
+const SIGNALS = ['SIGINT', 'SIGTERM', 'SIGHUP'] as const;
+
+/**
+ * stdin reaching EOF means the MCP client has gone away and nothing further can
+ * arrive.
+ */
+const STDIN_CLOSED = ['end', 'close'] as const;
+
+/**
+ * Exit cleanly when the MCP client disconnects or we are signalled, releasing
+ * the LSP servers we spawned.
+ *
+ * The stdio half matters as much as the signals. An MCP client that exits
+ * simply closes its end of the pipe, often without signalling us, and
+ * `StdioServerTransport` subscribes only to stdin's `data` and `error` events —
+ * never `end` or `close` — so `Server.onclose` does not fire on EOF. Without
+ * this, cclsp keeps running after its client is gone, and every LSP server it
+ * spawned stays resident with it. Long-lived language servers are not cheap
+ * (pylsp holds a few hundred MB), so across repeated client invocations the
+ * orphans accumulate until the machine is under real memory pressure.
+ *
+ * Returns the shutdown function, so a caller can trigger the same path itself.
+ */
+export function installShutdownHandlers(
+  client: Disposable,
+  proc: ShutdownProcess = process
+): (code?: number) => void {
+  let shuttingDown = false;
+
+  const shutdown = (code = 0): void => {
+    // Signals can arrive together, and stdin emits both `end` and `close`.
+    if (shuttingDown) return;
+    shuttingDown = true;
+    client.dispose();
+    proc.exit(code);
+  };
+
+  for (const signal of SIGNALS) {
+    proc.on(signal, () => shutdown(0));
+  }
+  for (const event of STDIN_CLOSED) {
+    proc.stdin.on(event, () => shutdown(0));
+  }
+
+  return shutdown;
+}