From f38111e5ca02c8723a68f4a7199f49a2945a1ac2 Mon Sep 17 00:00:00 2001 From: lintianle Date: Wed, 8 Jul 2026 11:58:12 +0800 Subject: [PATCH] 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 --- packages/mcp/mcp-client/src/index.ts | 64 ++++++++++++++++++++++------ 1 file changed, 50 insertions(+), 14 deletions(-) diff --git a/packages/mcp/mcp-client/src/index.ts b/packages/mcp/mcp-client/src/index.ts index da3aefc3de..2e797ec199 100644 --- a/packages/mcp/mcp-client/src/index.ts +++ b/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 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 { + 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 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') }