History

fix(mcp): bound discovery connection teardown (#1878)

* fix(mcp): bound discovery connection teardown Discovery cleanup runs in an interruption-masked finalizer. Bound an unresponsive transport close so a completed tool listing cannot strand health checks or other callers indefinitely. * test(mcp): match current discovery client contract * test(mcp): avoid wall-clock teardown assertion * Exercise discovery teardown with real MCP connections ---------

Saatvik Arya committed 2 months ago

Commit 513426c1 parent 61f71c52 files · +82 −10Browse files

Unified · Split

2 files
  1. packages/plugins/mcp/src/sdk/discover-close.test.ts +68 −0
  2. packages/plugins/mcp/src/sdk/discover.ts +14 −10
packages/plugins/mcp/src/sdk/discover-close.test.ts+68−0
@@ -0,0 +1,68 @@
1+import { describe, expect, it } from "@effect/vitest";
2+import { Effect } from "effect";
3+
4+import { createMcpConnector, type McpConnector } from "./connection";
5+import { discoverTools } from "./discover";
6+import { makeEchoMcpServer, serveMcpServer } from "../testing";
7+
8+// Exercise the real MCP handshake and catalog. The connector owns teardown,
9+// so a wrapper can reproduce a transport that closes its sockets but never
10+// settles its close promise without replacing the protocol client.
11+const hangingCloseConnector = (connector: McpConnector, state: { closes: number }): McpConnector =>
12+ Effect.map(connector, (connection) => ({
13+ client: connection.client,
14+ close: async () => {
15+ state.closes += 1;
16+ await connection.close();
17+ return new Promise<void>(() => {});
18+ },
19+ }));
20+
21+describe("MCP discovery teardown", () => {
22+ it.live("preserves a real catalog when close never settles", () =>
23+ Effect.scoped(
24+ Effect.gen(function* () {
25+ const server = yield* serveMcpServer(() => makeEchoMcpServer({ name: "hanging-close" }));
26+ const state = { closes: 0 };
27+ const manifest = yield* discoverTools(
28+ hangingCloseConnector(
29+ createMcpConnector({
30+ transport: "remote",
31+ endpoint: server.url,
32+ remoteTransport: "streamable-http",
33+ }),
34+ state,
35+ ),
36+ );
37+ expect(state.closes).toBe(1);
38+ expect(manifest.server?.name).toBe("hanging-close");
39+ expect(manifest.tools.length).toBeGreaterThan(0);
40+ }),
41+ ),
42+ );
43+
44+ it.live("preserves listing failure when close never settles", () =>
45+ Effect.scoped(
46+ Effect.gen(function* () {
47+ const server = yield* serveMcpServer(() => makeEchoMcpServer());
48+ yield* server.rejectSessionMethod("tools/list", 403);
49+ const state = { closes: 0 };
50+ const result = yield* discoverTools(
51+ hangingCloseConnector(
52+ createMcpConnector({
53+ transport: "remote",
54+ endpoint: server.url,
55+ remoteTransport: "streamable-http",
56+ }),
57+ state,
58+ ),
59+ ).pipe(Effect.result);
60+ expect(state.closes).toBe(1);
61+ expect(result).toMatchObject({
62+ _tag: "Failure",
63+ failure: { stage: "list_tools", httpStatus: 403 },
64+ });
65+ }),
66+ ),
67+ );
68+});
packages/plugins/mcp/src/sdk/discover.ts+14−10
@@ -31,6 +31,12 @@ const MAX_LIST_TOOLS_PAGES = 100; · 30 unmodified lines
3131 // shape probe's single unauth POST.
3232 const DEFAULT_DISCOVER_TIMEOUT = Duration.seconds(15);
3333
34+// Teardown is best-effort and paid for by the request that performed discovery.
35+// A remote transport may accept close and then never settle, so use the same
36+// bound as the invocation connection pool instead of stranding the caller in an
37+// uninterruptible finalizer after discovery itself has already completed.
38+const CLOSE_TIMEOUT = Duration.seconds(2);
39+
3440 // ---------------------------------------------------------------------------
3541 // Public API
3642 // ---------------------------------------------------------------------------
@@ -243,13 +249,11 @@ export const discoverTools = ( · 206 unmodified lines
243249 const closeConnection = (connection: {
244250 readonly close: () => Promise<void>;
245251 }): Effect.Effect<void, never> =>
246− Effect.ignore(
247− Effect.tryPromise({
248− try: () => connection.close(),
249− catch: () =>
250− new McpToolDiscoveryError({
251− stage: "list_tools",
252− message: "Failed closing MCP connection",
253− }),
254− }),
255− );
252+ Effect.tryPromise({
253+ try: () => connection.close(),
254+ catch: () =>
255+ new McpToolDiscoveryError({
256+ stage: "list_tools",
257+ message: "Failed closing MCP connection",
258+ }),
259+ }).pipe(Effect.timeout(CLOSE_TIMEOUT), Effect.ignore);

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t