Skip to content

fix(stdio): implement graceful shutdown as per spec #821

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions src/client/stdio.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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" },
]);
});
44 changes: 35 additions & 9 deletions src/client/stdio.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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) {
Expand Down