Browse Source

fix: handle MCP transport disconnect and concurrent tool re-sync

- Add client.onclose handler to unregister tools when the MCP server
  disconnects (crash or intentional close)
- Replace bare `let disposers` with a shared mutable state object so the
  effect disposer and notification handler always reference the current
  generation
- Serialize tools/list_changed resyncs with latest-wins coalescing
  (syncing + pendingResync flags) to prevent concurrent races
lintianle 3 months ago
parent
commit
f38111e5ca
1 changed files with 50 additions and 14 deletions
  1. 50 14
      packages/mcp/mcp-client/src/index.ts

+ 50 - 14
packages/mcp/mcp-client/src/index.ts

@@ -86,6 +86,16 @@ export const Config = z.union([
 
 // ---- Plugin apply ----
 
+/** Mutable state shared between the async connect path, notification handler, and disposers. */
+interface PluginState {
+  /** Current generation of tool disposers (keyed by registered name). */
+  disposers: Map<string, () => void>
+  /** Whether a syncTools call is currently in-flight. */
+  syncing: boolean
+  /** Whether another tools/list_changed arrived while syncing (coalesce flag). */
+  pendingResync: boolean
+}
+
 export function apply(ctx: Context, config: Config): void {
   const transport = createTransport(config)
   const client = new Client(
@@ -93,36 +103,62 @@ export function apply(ctx: Context, config: Config): void {
     { capabilities: {} },
   )
 
+  const state: PluginState = { disposers: new Map(), syncing: false, pendingResync: false }
+
+  const opts = { toolPrefix: config.toolPrefix, toolCallTimeoutMs: config.toolCallTimeoutMs }
+
+  /** Dispose all currently registered tools. */
+  function disposeTools(): void {
+    for (const dispose of state.disposers.values()) dispose()
+    state.disposers = new Map()
+  }
+
+  /** Run syncTools with latest-wins coalescing. */
+  async function resync(): Promise<void> {
+    if (state.syncing) {
+      state.pendingResync = true
+      return
+    }
+    state.syncing = true
+    try {
+      state.disposers = await syncTools(client, ctx, opts, state.disposers)
+    } finally {
+      state.syncing = false
+    }
+    // If another notification arrived while we were syncing, run once more.
+    if (state.pendingResync) {
+      state.pendingResync = false
+      await resync()
+    }
+  }
+
+  // When the connection closes (server crash or intentional close), unregister
+  // all tools so the model no longer sees them in the system prompt.
+  client.onclose = () => {
+    disposeTools()
+    ctx.logger.info('mcp-client: connection closed, tools unregistered')
+  }
+
   // Connect and set up tools. Errors during connect are logged, not thrown
   // (the plugin simply has no tools registered).
   const ready = (async () => {
     await client.connect(transport)
-
-    let disposers = await syncTools(client, ctx, {
-      toolPrefix: config.toolPrefix,
-      toolCallTimeoutMs: config.toolCallTimeoutMs,
-    }, new Map())
+    await resync()
 
     client.setNotificationHandler(
       ToolListChangedNotificationSchema,
       async () => {
         ctx.logger.info('mcp-client: tool list changed, re-syncing')
-        disposers = await syncTools(client, ctx, {
-          toolPrefix: config.toolPrefix,
-          toolCallTimeoutMs: config.toolCallTimeoutMs,
-        }, disposers)
+        await resync()
       },
     )
-
-    return disposers
   })().catch((error: unknown) => {
     ctx.logger.error(`mcp-client: failed to connect: ${String(error)}`)
-    return new Map<string, () => void>()
   })
 
+  // Fiber disposal: close the client (triggers onclose → tools unregistered).
   ctx.effect(() => async () => {
-    const disposers = await ready
-    for (const dispose of disposers.values()) dispose()
+    await ready
     try { await client.close() } catch { /* transport already gone */ }
   }, 'mcp-client.connection')
 }