mirror of
https://github.com/deepseek-ai/deepseek-harness
synced 2026-08-15 21:04:50 +00:00
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 ----
|
||||
|
||||
/** 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')
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user