diff --git a/workspace-server/src/__tests__/utils/process-lifecycle.test.ts b/workspace-server/src/__tests__/utils/process-lifecycle.test.ts new file mode 100644 index 0000000..c6ae443 --- /dev/null +++ b/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((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((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); + }); +}); diff --git a/workspace-server/src/index.ts b/workspace-server/src/index.ts index e7a8077..bbf6c17 100644 --- a/workspace-server/src/index.ts +++ b/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...`, diff --git a/workspace-server/src/utils/process-lifecycle.ts b/workspace-server/src/utils/process-lifecycle.ts new file mode 100644 index 0000000..bbb0de2 --- /dev/null +++ b/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; + stdin?: StdinEventSource; + signals?: EventSource; + exit?: (code: number) => void; + onError?: (reason: ShutdownReason, error: unknown) => void; +} + +export interface ProcessLifecycleController { + dispose: () => void; + getShutdownPromise: () => Promise | undefined; + shutdown: (reason: ShutdownReason) => Promise; +} + +const SIGNAL_EXIT_CODES: Partial> = { + 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 | 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 { + 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, + }; +}