Added clean shutdown + tests

Author: jerome3o-anthropic

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);
+  });
+});

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);
+});

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?.();
   }