feat: add graceful stdio shutdown (#570)

* feat: add graceful stdio shutdown

* fix: harden graceful stdio shutdown

---------

Co-authored-by: CharlieHelps <charlie@charlielabs.ai>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

.changeset/graceful-stdio-shutdown.md

@@ -0,0 +1,6 @@
+---
+"hevy-mcp": patch
+---
+
+Gracefully close and flush the stdio transport on SIGINT or SIGTERM, with a
+bounded forced-exit fallback when shutdown stalls or other handles remain open.

src/index.test.ts

@@ -29,6 +29,11 @@ const testDoubles = vi.hoisted(() => ({
 	tool: vi.fn(),
 	registerTool: vi.fn(),
 	directRegisterToolCalls: 0,
+	close: vi.fn().mockResolvedValue(undefined),
+	installGracefulShutdown: vi.fn(() => ({
+		cleanup: vi.fn(),
+		getShutdownPromise: vi.fn(),
+	})),
 	sentry: {
 		init: vi.fn(() => ({})),
 		setUser: vi.fn(),
@@ -50,6 +55,10 @@ vi.mock("./utils/hevyClient.js", () => ({
 	createClient: vi.fn().mockReturnValue({ mockedClient: true }),
 }));
 
+vi.mock("./utils/graceful-shutdown.js", () => ({
+	installGracefulShutdown: testDoubles.installGracefulShutdown,
+}));
+
 vi.mock("./utils/telemetry.js", () => ({
 	Sentry: testDoubles.sentry,
 	tracer: {
@@ -100,6 +109,7 @@ vi.mock("@modelcontextprotocol/sdk/server/mcp.js", () => {
 		}
 
 		connect = testDoubles.connect;
+		close = testDoubles.close;
 		isConnected = vi.fn(() => true);
 		sendLoggingMessage = testDoubles.sendLoggingMessage;
 		registerPrompt = testDoubles.registerPrompt;
@@ -132,6 +142,8 @@ describe("Server entry", () => {
 		process.env = { ...originalEnv };
 		process.argv = [...originalArgv];
 		vi.clearAllMocks();
+		testDoubles.connect.mockResolvedValue(undefined);
+		testDoubles.close.mockResolvedValue(undefined);
 		testDoubles.directRegisterToolCalls = 0;
 		testDoubles.tool.mockImplementation(
 			function (this: { registerTool: () => void }) {
@@ -290,6 +302,7 @@ describe("Server entry", () => {
 				expect(exitSpy).not.toHaveBeenCalled();
 				expect(createClient).not.toHaveBeenCalled();
 				expect(testDoubles.startActiveSpan).not.toHaveBeenCalled();
+				expect(testDoubles.installGracefulShutdown).not.toHaveBeenCalled();
 
 				const anyStdioModule = stdioModule as { __transports?: unknown[] };
 				expect(anyStdioModule.__transports).toHaveLength(0);
@@ -323,6 +336,7 @@ describe("Server entry", () => {
 				expect(helpText).toContain("Examples:");
 				expect(createClient).not.toHaveBeenCalled();
 				expect(testDoubles.startActiveSpan).not.toHaveBeenCalled();
+				expect(testDoubles.installGracefulShutdown).not.toHaveBeenCalled();
 
 				const anyStdioModule = stdioModule as { __transports?: unknown[] };
 				expect(anyStdioModule.__transports).toHaveLength(0);
@@ -357,6 +371,12 @@ describe("Server entry", () => {
 			);
 			expect(spanNames).toContain("mcp.server.run");
 			expect(spanNames).toContain("mcp.server.connect");
+			expect(testDoubles.installGracefulShutdown).toHaveBeenCalledWith({
+				target: expect.objectContaining({ close: testDoubles.close }),
+			});
+			expect(testDoubles.connect.mock.invocationCallOrder[0]).toBeLessThan(
+				testDoubles.installGracefulShutdown.mock.invocationCallOrder[0] ?? 0,
+			);
 		});
 
 		it("prefers CLI --hevy-api-key argument over environment variable", async () => {
@@ -393,6 +413,7 @@ describe("Server entry", () => {
 
 			await expect(runServer()).rejects.toThrow("connect failed");
 			expect(testDoubles.connect).toHaveBeenCalled();
+			expect(testDoubles.installGracefulShutdown).not.toHaveBeenCalled();
 			expect(testDoubles.span.setStatus).toHaveBeenCalledWith({ code: 2 });
 		});
 

src/index.ts

@@ -23,6 +23,7 @@ import { registerTemplateTools } from "./tools/templates.js";
 import { registerUserTools } from "./tools/user.js";
 import { registerWorkoutTools } from "./tools/workouts.js";
 import { assertApiKey, parseConfig } from "./utils/config.js";
+import { installGracefulShutdown } from "./utils/graceful-shutdown.js";
 import { createClient } from "./utils/hevyClient.js";
 import { createMcpClientLogger } from "./utils/mcp-client-logger.js";
 import { createInstrumentedStdioTransport } from "./utils/stdio-observability.js";
@@ -307,6 +308,7 @@ export async function runServer() {
 						}
 					},
 				);
+				installGracefulShutdown({ target: server });
 
 				span.setStatus({ code: SpanStatusCode.OK });
 			} catch (e) {

src/utils/graceful-shutdown.child-process.test.ts

@@ -0,0 +1,100 @@
+import { spawn } from "node:child_process";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import { describe, expect, it } from "vitest";
+import type { ShutdownSignal } from "./graceful-shutdown.js";
+
+const repositoryRoot = path.resolve(
+	path.dirname(fileURLToPath(import.meta.url)),
+	"../..",
+);
+const tsxCli = path.join(repositoryRoot, "node_modules/tsx/dist/cli.mjs");
+const childFixture = path.join(
+	repositoryRoot,
+	"tests/fixtures/graceful-shutdown-child.ts",
+);
+
+describe("graceful stdio shutdown regression", () => {
+	it.each(["SIGTERM", "SIGINT"] satisfies ShutdownSignal[])(
+		"flushes backpressured JSON-RPC frames before exiting on %s",
+		async (signal) => {
+			const child = spawn(process.execPath, [tsxCli, childFixture], {
+				cwd: repositoryRoot,
+				stdio: ["pipe", "pipe", "pipe"],
+			});
+			const stdoutChunks: Buffer[] = [];
+			const stderrChunks: Buffer[] = [];
+			let exited = false;
+			const exitResult = new Promise<{
+				code: number | null;
+				signal: NodeJS.Signals | null;
+			}>((resolve) => {
+				child.once("exit", (code, exitSignal) => {
+					exited = true;
+					resolve({ code, signal: exitSignal });
+				});
+			});
+
+			try {
+				let expectedFrameCount = 0;
+				await new Promise<void>((resolve, reject) => {
+					const timeout = setTimeout(() => {
+						reject(new Error("Timed out waiting for backpressure marker"));
+					}, 5_000);
+
+					child.stderr.on("data", (chunk: Buffer) => {
+						stderrChunks.push(chunk);
+						const stderr = Buffer.concat(stderrChunks).toString("utf8");
+						const marker = stderr.match(/BACKPRESSURED:(\d+)/);
+						if (marker) {
+							expectedFrameCount = Number(marker[1]);
+							clearTimeout(timeout);
+							resolve();
+						}
+					});
+					child.once("error", reject);
+					child.once("exit", (code, exitSignal) => {
+						clearTimeout(timeout);
+						reject(
+							new Error(
+								`Child exited before marker: code=${code}, signal=${exitSignal}`,
+							),
+						);
+					});
+				});
+
+				child.kill(signal);
+				await new Promise((resolve) => setTimeout(resolve, 50));
+				expect(exited).toBe(false);
+
+				child.stdout.on("data", (chunk: Buffer) => stdoutChunks.push(chunk));
+				const timeout = setTimeout(() => child.kill("SIGKILL"), 10_000);
+				const result = await exitResult;
+				clearTimeout(timeout);
+
+				expect(result).toEqual({ code: 0, signal: null });
+				const stdout = Buffer.concat(stdoutChunks).toString("utf8");
+				const stderr = Buffer.concat(stderrChunks).toString("utf8");
+				const frames = stdout.trimEnd().split("\n");
+
+				expect(expectedFrameCount).toBeGreaterThan(0);
+				expect(frames).toHaveLength(expectedFrameCount);
+				for (const [index, frame] of frames.entries()) {
+					expect(JSON.parse(frame)).toEqual({
+						jsonrpc: "2.0",
+						id: index + 1,
+						result: { payload: "x".repeat(128 * 1024) },
+					});
+				}
+				expect(stdout).not.toContain("Shutting down gracefully");
+				expect(stderr).toContain(`Shutting down gracefully after ${signal}`);
+			} finally {
+				if (!exited) {
+					child.kill("SIGKILL");
+					await exitResult;
+				}
+			}
+		},
+		10_000,
+	);
+});