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:
@@ -86,6 +86,16 @@ export const Config = z.union([
|
|||||||
|
|
||||||
// ---- Plugin apply ----
|
// ---- 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 {
|
export function apply(ctx: Context, config: Config): void {
|
||||||
const transport = createTransport(config)
|
const transport = createTransport(config)
|
||||||
const client = new Client(
|
const client = new Client(
|
||||||
@@ -93,36 +103,62 @@ export function apply(ctx: Context, config: Config): void {
|
|||||||
{ capabilities: {} },
|
{ 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
|
// Connect and set up tools. Errors during connect are logged, not thrown
|
||||||
// (the plugin simply has no tools registered).
|
// (the plugin simply has no tools registered).
|
||||||
const ready = (async () => {
|
const ready = (async () => {
|
||||||
await client.connect(transport)
|
await client.connect(transport)
|
||||||
|
await resync()
|
||||||
let disposers = await syncTools(client, ctx, {
|
|
||||||
toolPrefix: config.toolPrefix,
|
|
||||||
toolCallTimeoutMs: config.toolCallTimeoutMs,
|
|
||||||
}, new Map())
|
|
||||||
|
|
||||||
client.setNotificationHandler(
|
client.setNotificationHandler(
|
||||||
ToolListChangedNotificationSchema,
|
ToolListChangedNotificationSchema,
|
||||||
async () => {
|
async () => {
|
||||||
ctx.logger.info('mcp-client: tool list changed, re-syncing')
|
ctx.logger.info('mcp-client: tool list changed, re-syncing')
|
||||||
disposers = await syncTools(client, ctx, {
|
await resync()
|
||||||
toolPrefix: config.toolPrefix,
|
|
||||||
toolCallTimeoutMs: config.toolCallTimeoutMs,
|
|
||||||
}, disposers)
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
return disposers
|
|
||||||
})().catch((error: unknown) => {
|
})().catch((error: unknown) => {
|
||||||
ctx.logger.error(`mcp-client: failed to connect: ${String(error)}`)
|
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 () => {
|
ctx.effect(() => async () => {
|
||||||
const disposers = await ready
|
await ready
|
||||||
for (const dispose of disposers.values()) dispose()
|
|
||||||
try { await client.close() } catch { /* transport already gone */ }
|
try { await client.close() } catch { /* transport already gone */ }
|
||||||
}, 'mcp-client.connection')
|
}, 'mcp-client.connection')
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user