diff --git a/src/main.ts b/src/main.ts index 19669f2e5..5908d565f 100644 --- a/src/main.ts +++ b/src/main.ts @@ -34,7 +34,7 @@ interface RuntimeDependencies { readonly serve: typeof serveStdio; readonly writeStderr: (text: string) => void; readonly setExitCode: (code: number) => void; - readonly registerShutdown: (handler: () => void) => void; + readonly registerShutdown: (handler: () => void) => () => void; readonly registerReload?: (handler: () => void) => () => void; readonly createServer?: typeof createServer; readonly readProjectPermissionStore?: typeof readProjectPermissionStore; @@ -52,6 +52,12 @@ const runtimeDependencies = (): RuntimeDependencies => ({ process.once("SIGTERM", handler); process.stdin.once("end", handler); process.stdin.once("close", handler); + return () => { + process.off("SIGINT", handler); + process.off("SIGTERM", handler); + process.stdin.off("end", handler); + process.stdin.off("close", handler); + }; }, registerReload: (handler) => { process.on("SIGHUP", handler); @@ -232,9 +238,11 @@ export const run = async ( }) ?? (() => undefined); let shutdownPromise: Promise | undefined; + let unregisterShutdown = (): void => undefined; const shutdown = (): Promise => { shutdownPromise ??= (async () => { unregisterReload(); + unregisterShutdown(); await handle.close(); await session.close(); permissionAuthority.value.clearSessionGrants(); @@ -249,7 +257,7 @@ export const run = async ( }); }; - dependencies.registerShutdown(requestShutdown); + unregisterShutdown = dependencies.registerShutdown(requestShutdown); return 0; }; diff --git a/tests/mainErrors.test.ts b/tests/mainErrors.test.ts index 4b24f61f4..b0ee88447 100644 --- a/tests/mainErrors.test.ts +++ b/tests/mainErrors.test.ts @@ -13,18 +13,23 @@ const dependencies = ( output: string[], shutdown: Array<() => void>, exitCodes: number[], + shutdownUnregistrations: string[] = [], ): RuntimeDependencies => ({ env: {}, serve, writeStderr: (text) => output.push(text), setExitCode: (code) => exitCodes.push(code), - registerShutdown: (handler) => shutdown.push(handler), + registerShutdown: (handler) => { + shutdown.push(handler); + return () => shutdownUnregistrations.push("shutdown"); + }, }); describe("MCP runtime errors", () => { it("reports connection loss without the transport cause", async () => { const output: string[] = []; const shutdown: Array<() => void> = []; + const shutdownUnregistrations: string[] = []; let options: ServeStdioOptions | undefined; let closeCalls = 0; const close = (): Promise => { @@ -42,6 +47,7 @@ describe("MCP runtime errors", () => { output, shutdown, exitCodes, + shutdownUnregistrations, ), ), ).toBe(0); @@ -52,6 +58,42 @@ describe("MCP runtime errors", () => { shutdown[0]?.(); await nextTurn(); expect(closeCalls).toBe(1); + expect(shutdownUnregistrations).toEqual(["shutdown"]); + }); + + it("unregisters every process-lifetime handler during idempotent shutdown", async () => { + const shutdown: Array<() => void> = []; + const reload: Array<() => void> = []; + const unregistrations: string[] = []; + let closeCalls = 0; + const runtime: RuntimeDependencies = { + ...dependencies( + () => ({ + close: async () => { + closeCalls += 1; + }, + }), + [], + shutdown, + [], + unregistrations, + ), + registerReload: (handler) => { + reload.push(handler); + return () => unregistrations.push("reload"); + }, + }; + + expect(await run(runtime)).toBe(0); + expect(shutdown).toHaveLength(1); + expect(reload).toHaveLength(1); + + shutdown[0]?.(); + shutdown[0]?.(); + await nextTurn(); + + expect(closeCalls).toBe(1); + expect(unregistrations).toEqual(["reload", "shutdown"]); }); it("reports transport startup failure without its cause", async () => {