From 7c91349363c33abe63937f2a2d75cb3f8aab1010 Mon Sep 17 00:00:00 2001
From: Jerome <jerome@anthropic.com>
Date: Fri, 3 Jan 2025 16:29:06 +0000
Subject: [PATCH] Added clean shutdown + tests

---
 src/integration-tests/process-cleanup.test.ts | 28 ++++++++++++++
 src/server/stdio.test.ts                      | 38 +++++++++++++++++++
 src/server/stdio.ts                           |  8 ++++
 3 files changed, 74 insertions(+)
 create mode 100644 src/integration-tests/process-cleanup.test.ts

diff --git a/src/integration-tests/process-cleanup.test.ts b/src/integration-tests/process-cleanup.test.ts
new file mode 100644
index 0000000..0dd7861
--- /dev/null
+++ b/src/integration-tests/process-cleanup.test.ts
@@ -0,0 +1,28 @@
+import { Server } from "../server/index.js";
+import { StdioServerTransport } from "../server/stdio.js";
+
+describe("Process cleanup", () => {
+  jest.setTimeout(5000); // 5 second timeout
+
+  it("should exit cleanly after closing transport", async () => {
+    const server = new Server(
+      {
+        name: "test-server",
+        version: "1.0.0",
+      },
+      {
+        capabilities: {},
+      }
+    );
+
+    const transport = new StdioServerTransport();
+    await server.connect(transport);
+
+    // Close the transport
+    await transport.close();
+
+    // If we reach here without hanging, the test passes
+    // The test runner will fail if the process hangs
+    expect(true).toBe(true);
+  });
+});
\ No newline at end of file
diff --git a/src/server/stdio.test.ts b/src/server/stdio.test.ts
index 5243268..d041422 100644
--- a/src/server/stdio.test.ts
+++ b/src/server/stdio.test.ts
@@ -100,3 +100,41 @@ test("should read multiple messages", async () => {
   await finished;
   expect(readMessages).toEqual(messages);
 });
+
+test("should properly clean up resources when closed", async () => {
+  // Create mock streams that track their destroyed state
+  const mockStdin = new Readable({
+    read() {}, // No-op implementation
+    destroy() {
+      this.destroyed = true;
+      return this;
+    }
+  });
+  const mockStdout = new Writable({
+    write(chunk, encoding, callback) {
+      callback();
+    },
+    destroy() {
+      this.destroyed = true;
+      return this;
+    }
+  });
+
+  const transport = new StdioServerTransport(mockStdin, mockStdout);
+  await transport.start();
+
+  // Send a message to potentially create 'drain' listeners
+  await transport.send({ jsonrpc: "2.0", method: "test", id: 1 });
+
+  // Close the transport
+  await transport.close();
+
+  // Check that all listeners were removed
+  expect(mockStdin.listenerCount('data')).toBe(0);
+  expect(mockStdin.listenerCount('error')).toBe(0);
+  expect(mockStdout.listenerCount('drain')).toBe(0);
+  
+  // Check that streams were properly ended
+  expect(mockStdin.destroyed).toBe(true);
+  expect(mockStdout.destroyed).toBe(true);
+});
diff --git a/src/server/stdio.ts b/src/server/stdio.ts
index f0eff82..4e44cc3 100644
--- a/src/server/stdio.ts
+++ b/src/server/stdio.ts
@@ -62,8 +62,16 @@ export class StdioServerTransport implements Transport {
   }
 
   async close(): Promise<void> {
+    // Remove all event listeners
     this._stdin.off("data", this._ondata);
     this._stdin.off("error", this._onerror);
+    this._stdout.removeAllListeners('drain');
+
+    // Destroy both streams
+    this._stdin.destroy();
+    this._stdout.destroy();
+
+    // Clear the buffer and notify closure
     this._readBuffer.clear();
     this.onclose?.();
   }