diff --git a/src/main.ts b/src/main.ts index 0743843ba0..a01bd2398e 100644 --- a/src/main.ts +++ b/src/main.ts @@ -21,6 +21,7 @@ import { } from "./composition/optionalObservationProviders.js"; import { createToolResultDelivery } from "./server/toolResult.js"; import { createShutdown } from "./main/shutdown.js"; +import { registerProcessShutdown } from "./main/processShutdown.js"; const runtimeDependencies = (): RuntimeDependencies => ({ env: process.env, @@ -29,18 +30,7 @@ const runtimeDependencies = (): RuntimeDependencies => ({ setExitCode: (code) => { process.exitCode = code; }, - registerShutdown: (handler) => { - process.once("SIGINT", handler); - 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); - }; - }, + registerShutdown: registerProcessShutdown, }); /** diff --git a/src/main/processShutdown.ts b/src/main/processShutdown.ts new file mode 100644 index 0000000000..1daa102688 --- /dev/null +++ b/src/main/processShutdown.ts @@ -0,0 +1,38 @@ +type ShutdownSignal = "SIGINT" | "SIGTERM"; + +interface ShutdownSignalSource { + on(event: ShutdownSignal, listener: () => void): unknown; + off(event: ShutdownSignal, listener: () => void): unknown; +} + +interface ShutdownInputSource { + once(event: "end" | "close", listener: () => void): unknown; + off(event: "end" | "close", listener: () => void): unknown; +} + +interface ProcessShutdownEvents { + readonly signals: ShutdownSignalSource; + readonly stdin: ShutdownInputSource; +} + +const processShutdownEvents: ProcessShutdownEvents = { + signals: process, + stdin: process.stdin, +}; + +/** Register the process events that request shutdown and return their owner. */ +export const registerProcessShutdown = ( + handler: () => void, + events: ProcessShutdownEvents = processShutdownEvents, +): (() => void) => { + events.signals.on("SIGINT", handler); + events.signals.on("SIGTERM", handler); + events.stdin.once("end", handler); + events.stdin.once("close", handler); + return () => { + events.signals.off("SIGINT", handler); + events.signals.off("SIGTERM", handler); + events.stdin.off("end", handler); + events.stdin.off("close", handler); + }; +}; diff --git a/tests/boundary/mcp/shutdown.test.ts b/tests/boundary/mcp/shutdown.test.ts index f1720bb62d..7288ec71ba 100644 --- a/tests/boundary/mcp/shutdown.test.ts +++ b/tests/boundary/mcp/shutdown.test.ts @@ -1,3 +1,4 @@ +import { EventEmitter } from "node:events"; import type { StdioServerHandle } from "@modelcontextprotocol/server/stdio"; import pino from "pino"; import { expect, it } from "vitest"; @@ -5,6 +6,7 @@ import { expect, it } from "vitest"; import { createManagedBinarySession } from "../../../src/composition/binary.js"; import { MCP_SHUTDOWN_FAILED } from "../../../src/main/messages.js"; import { createShutdown } from "../../../src/main/shutdown.js"; +import { registerProcessShutdown } from "../../../src/main/processShutdown.js"; import type { RuntimeDependencies } from "../../../src/main/types.js"; import { ProviderCleanupError } from "../../../src/domain/providerCleanupError.js"; import { err, ok } from "../../../src/domain/result.js"; @@ -209,3 +211,72 @@ it("retries only failed cleanup owners after concurrent shutdown requests", asyn 2, 2, 1, 1, ]); }); + +it("retries failed MCP cleanup owners on a repeated signal and unregisters after success", async () => { + const signals = new EventEmitter(); + const stdin = new EventEmitter(); + const exitCodes: number[] = []; + let handleCalls = 0; + let androidCalls = 0; + let sessionCalls = 0; + let unregisterCalls = 0; + const dependencies: RuntimeDependencies = { + env: {}, + serve: () => { + throw new Error("unused fixture transport"); + }, + writeStderr: () => undefined, + setExitCode: (code) => exitCodes.push(code), + registerShutdown: (handler) => { + const unregister = registerProcessShutdown(handler, { signals, stdin }); + return () => { + unregisterCalls += 1; + unregister(); + }; + }, + }; + const session = createTestBinarySession(() => ({ + execute: async () => observed(null), + close: async () => { + sessionCalls += 1; + return ok(null); + }, + })); + const [target] = await createBinarySessionTargets(); + if (target === undefined) throw new Error("Missing fixture target"); + expect((await session.open(target)).ok).toBe(true); + createShutdown({ + handle: { + close: async () => { + handleCalls += 1; + if (handleCalls === 1) throw new Error("first handle close failed"); + }, + }, + closeAndroid: async () => { + androidCalls += 1; + if (androidCalls === 1) throw new Error("first Android close failed"); + }, + session, + dependencies, + serverLogger: pino({ enabled: false }), + }); + + signals.emit("SIGTERM"); + await new Promise((resolve) => setImmediate(resolve)); + expect([handleCalls, androidCalls, sessionCalls, unregisterCalls]).toEqual([ + 1, 1, 1, 0, + ]); + expect(exitCodes).toEqual([1]); + + signals.emit("SIGTERM"); + await new Promise((resolve) => setImmediate(resolve)); + expect([handleCalls, androidCalls, sessionCalls, unregisterCalls]).toEqual([ + 2, 2, 1, 1, + ]); + + signals.emit("SIGTERM"); + await new Promise((resolve) => setImmediate(resolve)); + expect([handleCalls, androidCalls, sessionCalls, unregisterCalls]).toEqual([ + 2, 2, 1, 1, + ]); +});