fix: close server on process disconnect

workspace-server/src/__tests__/utils/process-lifecycle.test.ts

@@ -0,0 +1,205 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { EventEmitter } from 'node:events';
+import { describe, expect, it, jest } from '@jest/globals';
+import { installProcessLifecycle } from '../../utils/process-lifecycle';
+
+class FakeStdin extends EventEmitter {
+  destroyed = false;
+  readableEnded = false;
+}
+
+describe('installProcessLifecycle', () => {
+  it('closes once when stdin ends and then closes', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    const close = jest.fn(async () => {});
+    const exit = jest.fn<(code: number) => void>();
+    const controller = installProcessLifecycle({
+      close,
+      stdin,
+      signals,
+      exit,
+    });
+
+    stdin.emit('end');
+    stdin.emit('close');
+
+    await controller.getShutdownPromise();
+
+    expect(close).toHaveBeenCalledTimes(1);
+    expect(exit).not.toHaveBeenCalled();
+  });
+
+  it('closes once when the process disconnects and then stdin closes', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    const close = jest.fn(async () => {});
+    const exit = jest.fn<(code: number) => void>();
+    const controller = installProcessLifecycle({
+      close,
+      stdin,
+      signals,
+      exit,
+    });
+
+    signals.emit('disconnect');
+    stdin.emit('close');
+
+    await controller.getShutdownPromise();
+
+    expect(close).toHaveBeenCalledTimes(1);
+    expect(exit).not.toHaveBeenCalled();
+  });
+
+  it('upgrades an in-flight stdin shutdown when SIGTERM arrives', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    let finishClose: (() => void) | undefined;
+    const close = jest.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          finishClose = resolve;
+        }),
+    );
+    const exit = jest.fn<(code: number) => void>();
+    const controller = installProcessLifecycle({
+      close,
+      stdin,
+      signals,
+      exit,
+    });
+
+    stdin.emit('end');
+    signals.emit('SIGTERM');
+
+    expect(close).toHaveBeenCalledTimes(1);
+    expect(exit).not.toHaveBeenCalled();
+
+    finishClose?.();
+    await controller.getShutdownPromise();
+
+    expect(close).toHaveBeenCalledTimes(1);
+    expect(exit).toHaveBeenCalledWith(143);
+  });
+
+  it('waits for close before exiting on SIGINT', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    let finishClose: (() => void) | undefined;
+    const close = jest.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          finishClose = resolve;
+        }),
+    );
+    const exit = jest.fn<(code: number) => void>();
+    const controller = installProcessLifecycle({
+      close,
+      stdin,
+      signals,
+      exit,
+    });
+
+    signals.emit('SIGINT');
+
+    expect(close).toHaveBeenCalledTimes(1);
+    expect(exit).not.toHaveBeenCalled();
+
+    finishClose?.();
+    await controller.getShutdownPromise();
+
+    expect(exit).toHaveBeenCalledWith(130);
+  });
+
+  it('reports close errors and exits with failure for a signal', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    const error = new Error('close failed');
+    const onError = jest.fn<(reason: string, error: unknown) => void>();
+    const exit = jest.fn<(code: number) => void>();
+    const controller = installProcessLifecycle({
+      close: jest.fn(async () => {
+        throw error;
+      }),
+      stdin,
+      signals,
+      exit,
+      onError,
+    });
+
+    signals.emit('SIGTERM');
+    await controller.getShutdownPromise();
+
+    expect(onError).toHaveBeenCalledWith('SIGTERM', error);
+    expect(exit).toHaveBeenCalledWith(1);
+  });
+
+  it('reports close errors and exits with failure after disconnect', async () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    const error = new Error('close failed');
+    const callOrder: string[] = [];
+    const onError = jest.fn<(reason: string, error: unknown) => void>(() => {
+      callOrder.push('error');
+    });
+    const exit = jest.fn<(code: number) => void>(() => {
+      callOrder.push('exit');
+    });
+    const controller = installProcessLifecycle({
+      close: jest.fn(async () => {
+        throw error;
+      }),
+      stdin,
+      signals,
+      exit,
+      onError,
+    });
+
+    signals.emit('disconnect');
+    await controller.getShutdownPromise();
+
+    expect(onError).toHaveBeenCalledWith('disconnect', error);
+    expect(exit).toHaveBeenCalledWith(1);
+    expect(callOrder).toEqual(['error', 'exit']);
+  });
+
+  it('closes immediately when stdin had already ended', async () => {
+    const stdin = new FakeStdin();
+    stdin.readableEnded = true;
+    const close = jest.fn(async () => {});
+    const controller = installProcessLifecycle({
+      close,
+      stdin,
+      signals: new EventEmitter(),
+      exit: jest.fn<(code: number) => void>(),
+    });
+
+    await controller.getShutdownPromise();
+
+    expect(close).toHaveBeenCalledTimes(1);
+  });
+
+  it('removes listeners without closing when disposed', () => {
+    const stdin = new FakeStdin();
+    const signals = new EventEmitter();
+    const close = jest.fn(async () => {});
+    const controller = installProcessLifecycle({ close, stdin, signals });
+
+    controller.dispose();
+    stdin.emit('end');
+    signals.emit('SIGTERM');
+    signals.emit('disconnect');
+
+    expect(close).not.toHaveBeenCalled();
+    expect(stdin.listenerCount('end')).toBe(0);
+    expect(stdin.listenerCount('close')).toBe(0);
+    expect(signals.listenerCount('SIGINT')).toBe(0);
+    expect(signals.listenerCount('SIGTERM')).toBe(0);
+    expect(signals.listenerCount('disconnect')).toBe(0);
+  });
+});

workspace-server/src/index.ts

@@ -24,6 +24,7 @@ import { GMAIL_SEARCH_MAX_RESULTS } from './utils/constants';
 import { gmailAttachmentSchema } from './utils/validation';
 
 import { setLoggingEnabled, logToFile } from './utils/logger';
+import { installProcessLifecycle } from './utils/process-lifecycle';
 import { applyToolNameNormalization } from './utils/tool-normalization';
 import { SCOPES } from './auth/scopes';
 import { resolveFeatures } from './features/index';
@@ -2004,6 +2005,7 @@ System labels that can be modified:
   // 4. Connect the transport layer and start listening
   const transport = new StdioServerTransport();
   await server.connect(transport);
+  installProcessLifecycle({ close: () => server.close() });
 
   console.error(
     `Google Workspace MCP Server is running (using ${separator} for tool names). Listening for requests...`,

workspace-server/src/utils/process-lifecycle.ts

@@ -0,0 +1,140 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export type ShutdownReason =
+  | 'stdin-end'
+  | 'stdin-close'
+  | 'disconnect'
+  | 'SIGINT'
+  | 'SIGTERM';
+
+interface EventSource {
+  once(event: string, listener: () => void): unknown;
+  off(event: string, listener: () => void): unknown;
+}
+
+interface StdinEventSource extends EventSource {
+  destroyed?: boolean;
+  readableEnded?: boolean;
+}
+
+export interface ProcessLifecycleOptions {
+  close: () => Promise<void>;
+  stdin?: StdinEventSource;
+  signals?: EventSource;
+  exit?: (code: number) => void;
+  onError?: (reason: ShutdownReason, error: unknown) => void;
+}
+
+export interface ProcessLifecycleController {
+  dispose: () => void;
+  getShutdownPromise: () => Promise<void> | undefined;
+  shutdown: (reason: ShutdownReason) => Promise<void>;
+}
+
+const SIGNAL_EXIT_CODES: Partial<Record<ShutdownReason, number>> = {
+  SIGINT: 130,
+  SIGTERM: 143,
+};
+
+/**
+ * Installs the process-level lifecycle events that are outside the MCP stdio
+ * transport's responsibility. All events share one shutdown promise so the
+ * server close operation runs at most once.
+ */
+export function installProcessLifecycle(
+  options: ProcessLifecycleOptions,
+): ProcessLifecycleController {
+  const stdin = options.stdin ?? process.stdin;
+  const signals = options.signals ?? process;
+  const exit = options.exit ?? ((code: number) => process.exit(code));
+  const onError =
+    options.onError ??
+    ((reason: ShutdownReason, error: unknown) => {
+      console.error(`Failed to close MCP server after ${reason}:`, error);
+    });
+
+  let disposed = false;
+  let shutdownPromise: Promise<void> | undefined;
+  let requestedExitCode: number | undefined;
+
+  const onStdinEnd = () => {
+    void shutdown('stdin-end');
+  };
+  const onStdinClose = () => {
+    void shutdown('stdin-close');
+  };
+  const onDisconnect = () => {
+    void shutdown('disconnect');
+  };
+  const onSigint = () => {
+    void shutdown('SIGINT');
+  };
+  const onSigterm = () => {
+    void shutdown('SIGTERM');
+  };
+
+  function dispose() {
+    if (disposed) {
+      return;
+    }
+
+    disposed = true;
+    stdin.off('end', onStdinEnd);
+    stdin.off('close', onStdinClose);
+    signals.off('disconnect', onDisconnect);
+    signals.off('SIGINT', onSigint);
+    signals.off('SIGTERM', onSigterm);
+  }
+
+  function shutdown(reason: ShutdownReason): Promise<void> {
+    const signalExitCode = SIGNAL_EXIT_CODES[reason];
+    if (signalExitCode !== undefined) {
+      requestedExitCode = signalExitCode;
+    }
+
+    if (shutdownPromise) {
+      return shutdownPromise;
+    }
+
+    shutdownPromise = (async () => {
+      let closeFailed = false;
+
+      try {
+        await options.close();
+      } catch (error) {
+        closeFailed = true;
+        onError(reason, error);
+      } finally {
+        dispose();
+      }
+
+      if (closeFailed) {
+        exit(1);
+      } else if (requestedExitCode !== undefined) {
+        exit(requestedExitCode);
+      }
+    })();
+
+    return shutdownPromise;
+  }
+
+  stdin.once('end', onStdinEnd);
+  stdin.once('close', onStdinClose);
+  signals.once('disconnect', onDisconnect);
+  signals.once('SIGINT', onSigint);
+  signals.once('SIGTERM', onSigterm);
+
+  if (stdin.readableEnded || stdin.destroyed) {
+    void shutdown('stdin-close');
+  }
+
+  return {
+    dispose,
+    getShutdownPromise: () => shutdownPromise,
+    shutdown,
+  };
+}