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
This commit is contained in:
lintianle
2026-07-08 11:58:12 +08:00
parent 1fbe7c39d4
commit f38111e5ca

View File

@@ -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')
}