From a4e40f20772fc5f9eb845d045d49fad64172d097 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Tue, 29 Jul 2025 15:54:16 +0200
Subject: [PATCH] fix(stdio): implement graceful shutdown as per spec

---
 src/client/stdio.test.ts | 31 ++++++++++++++++++++++++++++
 src/client/stdio.ts      | 44 ++++++++++++++++++++++++++++++++--------
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/src/client/stdio.test.ts b/src/client/stdio.test.ts
index 2e4d92c25..ecf083992 100644
--- a/src/client/stdio.test.ts
+++ b/src/client/stdio.test.ts
@@ -75,3 +75,34 @@ test("should return child process pid", async () => {
   await client.close();
   expect(client.pid).toBeNull();
 });
+
+test("should handle process exit gracefully", async () => {
+  const server = `
+    setTimeout(() => {}, 100000); // stubborn process doesnt exit gracefully
+    process.stdin.resume();
+    process.stdin.on('close', () => {
+      process.stdout.write(JSON.stringify({ jsonrpc: "2.0", method: "stdin closed" }) + "\\n");
+    });
+    process.on('SIGTERM', () => {
+      process.stdout.write(JSON.stringify({ jsonrpc: "2.0", method: "received sigterm" }) + "\\n");
+    });
+  `;
+
+  const client = new StdioClientTransport({
+    command: process.execPath,
+    args: ["-e", server],
+  });
+  client.onerror = (error) => {
+    throw error;
+  };
+  const messages: JSONRPCMessage[] = []
+  client.onmessage = (message) => messages.push(message);
+
+  await client.start();
+  await client.close();
+
+  expect(messages).toEqual([
+    { jsonrpc: "2.0", method: "stdin closed" },
+    { jsonrpc: "2.0", method: "received sigterm" },
+  ]);
+});
diff --git a/src/client/stdio.ts b/src/client/stdio.ts
index 62292ce10..3ab56bce4 100644
--- a/src/client/stdio.ts
+++ b/src/client/stdio.ts
@@ -2,6 +2,7 @@ import { ChildProcess, IOType } from "node:child_process";
 import spawn from "cross-spawn";
 import process from "node:process";
 import { Stream, PassThrough } from "node:stream";
+import timers from "node:timers/promises";
 import { ReadBuffer, serializeMessage } from "../shared/stdio.js";
 import { Transport } from "../shared/transport.js";
 import { JSONRPCMessage } from "../types.js";
@@ -91,7 +92,6 @@ export function getDefaultEnvironment(): Record<string, string> {
  */
 export class StdioClientTransport implements Transport {
   private _process?: ChildProcess;
-  private _abortController: AbortController = new AbortController();
   private _readBuffer: ReadBuffer = new ReadBuffer();
   private _serverParams: StdioServerParameters;
   private _stderrStream: PassThrough | null = null;
@@ -129,19 +129,12 @@ export class StdioClientTransport implements Transport {
           },
           stdio: ["pipe", "pipe", this._serverParams.stderr ?? "inherit"],
           shell: false,
-          signal: this._abortController.signal,
           windowsHide: process.platform === "win32" && isElectron(),
           cwd: this._serverParams.cwd,
         }
       );
 
       this._process.on("error", (error) => {
-        if (error.name === "AbortError") {
-          // Expected when close() is called.
-          this.onclose?.();
-          return;
-        }
-
         reject(error);
         this.onerror?.(error);
       });
@@ -214,11 +207,44 @@ export class StdioClientTransport implements Transport {
   }
 
   async close(): Promise<void> {
-    this._abortController.abort();
+    await this._shutdownGracefully();
     this._process = undefined;
     this._readBuffer.clear();
   }
 
+  // https://modelcontextprotocol.io/specification/2025-06-18/basic/lifecycle#stdio
+  private async _shutdownGracefully() {
+    if (!this._process) {
+      return;
+    }
+
+    const closed = new Promise<'closed'>(resolve => this._process!.once('close', () => resolve('closed')));
+
+    this._process.stdin?.end();
+    let success = await Promise.race([
+      closed,
+      timers.setTimeout(100, undefined, { ref: false })
+    ]);
+    if (success)
+      return;
+
+    // Attempt to gracefully terminate the process.
+    process.kill(this._process.pid!, "SIGTERM");
+    success = await Promise.race([
+      closed,
+      timers.setTimeout(1000, undefined, { ref: false })
+    ]);
+    if (success)
+      return;
+
+    // Kill the process group forcefully if it didn't shut down in time.
+    process.kill(this._process.pid!, "SIGKILL");
+    await Promise.race([
+      closed,
+      timers.setTimeout(5000, undefined, { ref: false }).then(() => { throw new Error("Process did not exit after SIGKILL, something is horribly wrong.") })
+    ]);
+  }
+
   send(message: JSONRPCMessage): Promise<void> {
     return new Promise((resolve) => {
       if (!this._process?.stdin) {