Implemented safe disconnect for quick disconnect test cases (that failed infrequently).

Author: BobDickinson

core/__tests__/mcp/state/managedPromptsState.test.ts

@@ -2,7 +2,7 @@
  * ManagedPromptsState tests use a real InspectorClient and test server (same model
  * as managedToolsState.test.ts) so we exercise actual client/server behavior.
  */
-import { describe, it, expect, afterEach, beforeEach } from "vitest";
+import { describe, it, expect, afterEach } from "vitest";
 import type { Prompt } from "@modelcontextprotocol/sdk/types.js";
 import { InspectorClient } from "../../../mcp/inspectorClient.js";
 import { createTransportNode } from "../../../mcp/node/transport.js";
@@ -20,37 +20,13 @@ describe("ManagedPromptsState", () => {
   let server: TestServerHttp | null = null;
   let state: ManagedPromptsState | null = null;
 
-  let unhandledRejectionHandler: (
-    reason: unknown,
-    promise: Promise<unknown>,
-  ) => void;
-  beforeEach(() => {
-    unhandledRejectionHandler = (
-      reason: unknown,
-      promise: Promise<unknown>,
-    ) => {
-      const err = reason as { code?: number; message?: string };
-      if (err?.code === -32000 || err?.message?.includes("Connection closed")) {
-        promise.catch(() => {});
-        return;
-      }
-      throw reason;
-    };
-    process.on("unhandledRejection", unhandledRejectionHandler);
-  });
   afterEach(async () => {
     if (state) {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -59,7 +35,6 @@ describe("ManagedPromptsState", () => {
       }
       server = null;
     }
-    process.off("unhandledRejection", unhandledRejectionHandler);
   });
 
   function waitForPromptsChange(s: ManagedPromptsState): Promise<Prompt[]> {
@@ -168,7 +143,7 @@ describe("ManagedPromptsState", () => {
     await client.connect();
     await waitForPromptsChange(state!);
     expect(state!.getPrompts().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getPrompts()).toEqual([]);
   });
 

core/__tests__/mcp/state/managedResourceTemplatesState.test.ts

@@ -2,7 +2,7 @@
  * ManagedResourceTemplatesState tests use a real InspectorClient and test server (same model
  * as managedToolsState.test.ts) so we exercise actual client/server behavior.
  */
-import { describe, it, expect, afterEach, beforeEach } from "vitest";
+import { describe, it, expect, afterEach } from "vitest";
 import type { ResourceTemplate } from "@modelcontextprotocol/sdk/types.js";
 import { InspectorClient } from "../../../mcp/inspectorClient.js";
 import { createTransportNode } from "../../../mcp/node/transport.js";
@@ -20,38 +20,13 @@ describe("ManagedResourceTemplatesState", () => {
   let server: TestServerHttp | null = null;
   let state: ManagedResourceTemplatesState | null = null;
 
-  let unhandledRejectionHandler: (
-    reason: unknown,
-    promise: Promise<unknown>,
-  ) => void;
-  beforeEach(() => {
-    unhandledRejectionHandler = (
-      reason: unknown,
-      promise: Promise<unknown>,
-    ) => {
-      const err = reason as { code?: number; message?: string };
-      if (err?.code === -32000 || err?.message?.includes("Connection closed")) {
-        promise.catch(() => {});
-        return;
-      }
-      throw reason;
-    };
-    process.on("unhandledRejection", unhandledRejectionHandler);
-  });
-
   afterEach(async () => {
     if (state) {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -60,7 +35,6 @@ describe("ManagedResourceTemplatesState", () => {
       }
       server = null;
     }
-    process.off("unhandledRejection", unhandledRejectionHandler);
   });
 
   function waitForResourceTemplatesChange(
@@ -140,7 +114,7 @@ describe("ManagedResourceTemplatesState", () => {
     await client.connect();
     await waitForResourceTemplatesChange(state!);
     expect(state!.getResourceTemplates().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getResourceTemplates()).toEqual([]);
   });
 

core/__tests__/mcp/state/managedResourcesState.test.ts

@@ -2,7 +2,7 @@
  * ManagedResourcesState tests use a real InspectorClient and test server (same model
  * as managedToolsState.test.ts) so we exercise actual client/server behavior.
  */
-import { describe, it, expect, afterEach, beforeEach } from "vitest";
+import { describe, it, expect, afterEach } from "vitest";
 import type { Resource } from "@modelcontextprotocol/sdk/types.js";
 import { InspectorClient } from "../../../mcp/inspectorClient.js";
 import { createTransportNode } from "../../../mcp/node/transport.js";
@@ -20,37 +20,13 @@ describe("ManagedResourcesState", () => {
   let server: TestServerHttp | null = null;
   let state: ManagedResourcesState | null = null;
 
-  let unhandledRejectionHandler: (
-    reason: unknown,
-    promise: Promise<unknown>,
-  ) => void;
-  beforeEach(() => {
-    unhandledRejectionHandler = (
-      reason: unknown,
-      promise: Promise<unknown>,
-    ) => {
-      const err = reason as { code?: number; message?: string };
-      if (err?.code === -32000 || err?.message?.includes("Connection closed")) {
-        promise.catch(() => {});
-        return;
-      }
-      throw reason;
-    };
-    process.on("unhandledRejection", unhandledRejectionHandler);
-  });
   afterEach(async () => {
     if (state) {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -59,7 +35,6 @@ describe("ManagedResourcesState", () => {
       }
       server = null;
     }
-    process.off("unhandledRejection", unhandledRejectionHandler);
   });
 
   function waitForResourcesChange(
@@ -112,6 +87,7 @@ describe("ManagedResourcesState", () => {
       { type: "streamable-http", url: server.url },
       {
         environment: { transport: createTransportNode },
+        clientIdentity: { name: "test", version: "1.0.0" },
       },
     );
     await client.connect();
@@ -175,7 +151,7 @@ describe("ManagedResourcesState", () => {
     await client.connect();
     await waitForResourcesChange(state!);
     expect(state!.getResources().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getResources()).toEqual([]);
   });
 

core/__tests__/mcp/state/managedToolsState.test.ts

@@ -25,14 +25,8 @@ describe("ManagedToolsState", () => {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -167,7 +161,7 @@ describe("ManagedToolsState", () => {
     await client.connect();
     await waitForToolsChange(state!);
     expect(state!.getTools().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getTools()).toEqual([]);
   });
 

core/__tests__/mcp/state/pagedPromptsState.test.ts

@@ -44,14 +44,8 @@ describe("PagedPromptsState", () => {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -161,7 +155,7 @@ describe("PagedPromptsState", () => {
     await client.connect();
     await state.loadPage();
     expect(state.getPrompts().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getPrompts()).toEqual([]);
   });
 

core/__tests__/mcp/state/pagedResourceTemplatesState.test.ts

@@ -44,14 +44,8 @@ describe("PagedResourceTemplatesState", () => {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -163,7 +157,7 @@ describe("PagedResourceTemplatesState", () => {
     await client.connect();
     await state.loadPage();
     expect(state.getResourceTemplates().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getResourceTemplates()).toEqual([]);
   });
 

core/__tests__/mcp/state/pagedResourcesState.test.ts

@@ -44,14 +44,8 @@ describe("PagedResourcesState", () => {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -161,7 +155,7 @@ describe("PagedResourcesState", () => {
     await client.connect();
     await state.loadPage();
     expect(state.getResources().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getResources()).toEqual([]);
   });
 

core/__tests__/mcp/state/pagedToolsState.test.ts

@@ -25,14 +25,8 @@ describe("PagedToolsState", () => {
       state.destroy();
       state = null;
     }
-    if (client) {
-      try {
-        await client.disconnect();
-      } catch {
-        // ignore
-      }
-      client = null;
-    }
+    if (client) await client.disconnect(100);
+    client = null;
     if (server) {
       try {
         await server.stop();
@@ -149,7 +143,7 @@ describe("PagedToolsState", () => {
     await client.connect();
     await state.loadPage();
     expect(state.getTools().length).toBeGreaterThan(0);
-    await client!.disconnect();
+    await client!.disconnect(100);
     expect(state!.getTools()).toEqual([]);
   });
 

core/mcp/inspectorClient.ts

@@ -854,10 +854,34 @@ export class InspectorClient extends InspectorClientEventTarget {
   }
 
   /**
-   * Disconnect from the MCP server
+   * Disconnect from the MCP server.
+   * @param safeDisconnectTimeout If > 0, poll every 10ms until SDK _responseHandlers is empty or this many ms have elapsed, then close. Default 0 = close immediately.
    */
-  async disconnect(): Promise<void> {
+  async disconnect(safeDisconnectTimeout = 0): Promise<void> {
     if (this.client) {
+      if (safeDisconnectTimeout > 0) {
+        // This is pretty creepy, but there are test cases where client calls return but there
+        // are still response handlers pending. Usually a single macrotask delay is enough to
+        // clear them, but not always (it's been >10ms in some cases). The pending handlers
+        // themselves get the error (and in cases where those aren't awaited, the errors fly
+        // out of the test). This workaround where we directly access the handlers (otherwise
+        // private member of the SDK client) is creepy, but the least ugly working solution.
+        // We will re-valuate this with the v2 SDK. Currenly only tests that do quick disconnects
+        // use this setting.
+        //
+        const protocol = this.client as unknown as {
+          _responseHandlers?: Map<unknown, unknown>;
+        };
+        const handlers = protocol._responseHandlers;
+        const deadline = Date.now() + safeDisconnectTimeout;
+        while (
+          handlers?.size !== undefined &&
+          handlers.size > 0 &&
+          Date.now() < deadline
+        ) {
+          await new Promise((r) => setTimeout(r, 10));
+        }
+      }
       try {
         await this.client.close();
       } catch {