diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ed8ee0b7c..0003900a6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -22,6 +22,7 @@ - Fixed `!!` command headers losing their dim color once output arrives ([#10557](https://github.com/earendil-works/pi/pull/10557) by [@rwachtler](https://github.com/rwachtler)) - Fixed the codemode description not marking `searchTools()`, `describeTool()`, and `describeNamespace()` as async, which led models to serialize the unawaited promise as `{}` ([#10555](https://github.com/earendil-works/pi/issues/10555)) - Fixed codemode output items running together, so models could not tell where one `text()` or `console.log()` output ended and the next began. With several text items, each now starts with a `==> text N/M <==` line, and `console` calls follow the other output in one `` block with one line per call +- Fixed `/mcp` waiting for all servers to connect before opening; the manager now updates live and remains usable while enabling, reconnecting, or disabling servers ([#10562](https://github.com/earendil-works/pi/issues/10562)) ## [1.0.4] - 2026-10-05 diff --git a/packages/coding-agent/src/extensions/mcp/index.ts b/packages/coding-agent/src/extensions/mcp/index.ts index b2b8cb71d..391c271b4 100644 --- a/packages/coding-agent/src/extensions/mcp/index.ts +++ b/packages/coding-agent/src/extensions/mcp/index.ts @@ -106,8 +106,12 @@ interface McpServer { registeredConfig?: string; /** Result of the last `/mcp` action that failed, shown in the manager. */ message?: string; + /** Identifies the current connection attempt; disabling or replacing it invalidates earlier work. */ + attempt?: symbol; /** Settles when the connection started for the server connected or failed. */ ready?: Promise; + /** Cleanup of a detached connection; a replacement waits for it before opening a transport. */ + closing?: Promise; } const EXPOSURE_DESCRIPTIONS: Record, string> = { @@ -308,6 +312,8 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension override: entry.override !== undefined, })); + /** Manager actions still opening or closing connections; shutdown waits for their cleanup. */ + const backgroundActions = new Set>(); const listeners = new Set<() => void>(); const emitChange = () => { for (const listener of listeners) listener(); @@ -402,8 +408,22 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension exposure: getMcpToolExposure(entry.config, tool.name), namespace, timeoutMs: connection.timeoutMs, - getClient: async () => connection, - readableResources: () => resourceServers().includes(connection), + getClient: async () => { + // A prepared call can outlive this definition's connection while the readiness + // hook waits for a disable/re-enable. Resolve the replacement at execution time. + const current = findServer(server); + if (!current || !isEnabled(current)) throw new Error(`MCP server "${server}" is disabled.`); + const client = current.connection; + if (!client) throw new Error(`MCP server "${server}" is still starting.`); + if ( + getMcpToolExposure(current.entry.config, tool.name) === "hidden" || + (client.state === "connected" && !client.tools.some((offered) => offered.name === tool.name)) + ) { + throw new Error(`MCP tool "${server}/${tool.name}" is no longer available.`); + } + return client; + }, + readableResources: () => resourceServers().some((current) => current.name === server), }); definitions.set(definition.name, definition); pi.registerTool(definition); @@ -522,8 +542,12 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension }; /** Create the server's connection, loading the MCP runtime on first use. */ - const createConnection = async (server: McpServer): Promise => { + const createConnection = async ( + server: McpServer, + isCurrent: () => boolean, + ): Promise => { const runtime = await loadMcpRuntime(); + if (!isCurrent()) return undefined; const connection = new runtime.McpServerConnection({ entry: server.entry, cwd: sessionCwd, @@ -531,8 +555,13 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension credentials: getCredentials(runtime), providerToken: async (provider) => modelRegistry?.getApiKeyForProvider(provider), log: getServerLog(runtime), - onTools: registerTools, - onChange: onConnectionChange, + onTools: (connection) => { + if (isCurrent()) registerTools(connection); + }, + onChange: (connection) => { + if (isCurrent()) onConnectionChange(connection); + else tokensAtSignIn.delete(connection); + }, }); server.connection = connection; emitChange(); @@ -541,38 +570,53 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension /** * Connect the server in the background. `server.ready` settles when it connected or failed; - * failures show in its state. It starts after `after`; `isCurrent` stops it when the session ended + * failures show in its state. It starts after `after`, unless disabled, replaced, or shut down * meanwhile. The returned promise rejects when the MCP runtime cannot be loaded. */ - const startConnection = ( - server: McpServer, - isCurrent: () => boolean, - after?: Promise, - ): Promise => { + const startConnection = (server: McpServer, after?: Promise): Promise => { + const current = generation; + const attempt = Symbol(); + server.attempt = attempt; + const isCurrent = () => + current === generation && server.attempt === attempt && isEnabled(server) && servers.includes(server); const ready = (async () => { await after; if (!isCurrent()) return; - const connection = await createConnection(server); - if (!isCurrent()) return; + const connection = await createConnection(server, isCurrent); + if (!connection || !isCurrent()) return; await connection.getClient().catch(() => undefined); })(); server.ready = ready.catch(() => undefined); return ready; }; - /** Wait for servers still connecting, until they settle or `signal` aborts. */ + /** Wait for the latest connection attempts, including ones queued while waiting, or an abort. */ const waitForServers = async (waiting: readonly McpServer[], signal: AbortSignal | undefined) => { - const ready = waiting.flatMap((server) => (server.ready ? [server.ready] : [])); - if (ready.length === 0 || signal?.aborted) return; - let onAbort: (() => void) | undefined; - await Promise.race([ - Promise.all(ready), - new Promise((resolve) => { - onAbort = () => resolve(); - signal?.addEventListener("abort", onAbort, { once: true }); - }), - ]); - if (onAbort) signal?.removeEventListener("abort", onAbort); + for (;;) { + const ready = waiting.flatMap((server) => + isEnabled(server) && servers.includes(server) && server.ready ? [{ server, ready: server.ready }] : [], + ); + if (ready.length === 0 || signal?.aborted) return; + let onAbort: (() => void) | undefined; + await Promise.race([ + Promise.all(ready.map((attempt) => attempt.ready)), + new Promise((resolve) => { + onAbort = () => resolve(); + signal?.addEventListener("abort", onAbort, { once: true }); + }), + ]); + if (onAbort) signal?.removeEventListener("abort", onAbort); + if ( + signal?.aborted || + ready.every( + (attempt) => + !isEnabled(attempt.server) || + !servers.includes(attempt.server) || + attempt.ready === attempt.server.ready, + ) + ) + return; + } }; /** @@ -650,8 +694,14 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension const reconnect = async (server: McpServer): Promise => { const connection = server.connection; if (!connection) return `MCP server "${server.entry.name}" is disabled.`; - try { + const previous = server.ready; + const ready = (async () => { + await previous; await connection.reconnect(); + })(); + server.ready = ready.catch(() => undefined); + try { + await ready; return undefined; } catch (error) { return errorMessage(error); @@ -668,13 +718,24 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension if (failed) return failed; if (!enabled) { const connection = server.connection; + server.attempt = undefined; server.connection = undefined; + if (connection) { + // A reconnect may already have detached its old client while awaiting transport + // shutdown. Closing the connection alone does not wait for that reconnect. + server.closing = Promise.all([connection.close(), server.ready]).then(() => undefined); + } + const closing = server.closing; hideTools(server.entry.name); emitChange(); - await connection?.close(); + try { + await closing; + } finally { + if (server.closing === closing) server.closing = undefined; + } return undefined; } - await startConnection(server, () => true); + await startConnection(server, server.closing); return undefined; }; @@ -848,17 +909,39 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension }); }; + /** Keep the subscribed menu usable while a connection opens or closes. */ + const runInBackground = ( + ctx: ExtensionContext, + server: McpServer, + operation: () => Promise, + ) => { + const ready = operation(); + const attempt = server.attempt; + const current = generation; + const finish = (message: string | undefined) => { + if (current !== generation || server.attempt !== attempt || !servers.includes(server)) return; + server.message = message; + ensureDiscoveryActive(ctx); + emitChange(); + }; + const task = ready + .then(finish, (error: unknown) => finish(errorMessage(error))) + .finally(() => backgroundActions.delete(task)); + backgroundActions.add(task); + }; + const runAction = async (ui: McpUi, ctx: ExtensionContext, server: McpServer, action: string) => { - const { name } = server.entry; let message: string | undefined; switch (action) { case "signin": message = await signInWithUi(ui, server); break; case "reconnect": - // A failure shows as the connection's state and error. - ui.status(`MCP server ${name}`, "Reconnecting…"); - await reconnect(server); + // Connection state and error already report failures, including required sign-ins. + runInBackground(ctx, server, async () => { + await reconnect(server); + return undefined; + }); break; case "signout": await signOut(server); @@ -874,8 +957,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension case "enable-project": case "disable-project": { const enable = action.startsWith("enable"); - ui.status(`MCP server ${name}`, enable ? "Connecting…" : "Disconnecting…"); - message = await setEnabled(server, enable, action.endsWith("-project")); + runInBackground(ctx, server, () => setEnabled(server, enable, action.endsWith("-project"))); break; } } @@ -1022,7 +1104,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension // event loop turn lets the first render happen before loading and connecting. const runtime = new Promise((resolve) => setImmediate(resolve)).then(() => loadMcpRuntime()); const isCurrent = () => current === generation; - pending = Promise.all(enabled.map((server) => startConnection(server, isCurrent, runtime))) + pending = Promise.all(enabled.map((server) => startConnection(server, runtime))) .then(() => { if (isCurrent()) reportProblems(ctx); }) @@ -1069,21 +1151,25 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension // A codemode script waits for the servers it names, or for every server when it searches or // enumerates tools, so their tools are registered before the script runs. tool_search and the - // resource tools reach every server, so they wait for all of them. + // resource tools reach every server, so they wait for all of them. Direct MCP calls wait only + // for their own server, including a reconnect still closing its old transport. pi.on("tool_call", async (event, ctx) => { const tool = pi.getAllTools().find((candidate) => candidate.name === event.toolName); if (!tool) return; - const pendingServers = servers.filter( - (server) => isEnabled(server) && server.connection?.state !== "connected" && server.ready, - ); - if (pendingServers.length === 0) return; + // Readiness can be pending while the display state is still connected: reconnect first + // closes the old transport, then opens a new one. Settled promises return immediately. + const readyServers = servers.filter((server) => isEnabled(server) && server.ready); + if (readyServers.length === 0) return; let waiting: McpServer[] = []; if (isCodemodeTool(tool)) { const { code } = event.input as { code?: unknown }; const source = typeof code === "string" ? code : ""; - waiting = pendingServers.filter((server) => scriptNeedsServer(source, server.entry.name)); + waiting = readyServers.filter((server) => scriptNeedsServer(source, server.entry.name)); } else if (isToolSearchTool(tool) || RESOURCE_TOOL_NAMES.has(tool.name)) { - waiting = pendingServers; + waiting = readyServers; + } else { + const owner = toolOwners.get(event.toolName)?.split("\0", 1)[0]; + waiting = readyServers.filter((server) => server.entry.name === owner); } await waitForServers(waiting, ctx.signal); }); @@ -1116,7 +1202,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension const connecting = added.filter(isEnabled); if (current !== generation || connecting.length === 0) return; try { - await Promise.all(connecting.map((server) => startConnection(server, () => current === generation))); + await Promise.all(connecting.map((server) => startConnection(server))); if (current !== generation) { await Promise.all(connecting.map((server) => server.connection?.close())); return; @@ -1135,7 +1221,7 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension const closing = connections(); servers = []; emitChange(); - await Promise.all(closing.map((connection) => connection.close())); + await Promise.all([...backgroundActions, ...closing.map((connection) => connection.close())]); }); pi.registerCommand("mcp", { @@ -1162,17 +1248,20 @@ export function createMcpExtension(options: McpExtensionOptions = {}): Extension return items.length > 0 ? items : null; }, handler: async (args, ctx) => { - await pending; const [action, name, ...extra] = args.trim().split(/\s+/).filter(Boolean); if (action === undefined) { if (ctx.mode === "tui") await showMcpManager(ctx, (ui) => manage(ui, ctx)); - else ctx.ui.notify(formatStatus(), "info"); + else { + await pending; + ctx.ui.notify(formatStatus(), "info"); + } return; } if (extra.length > 0) { ctx.ui.notify(MCP_USAGE, "warning"); return; } + await pending; switch (action) { case "login": { const server = await pickServer(name, ctx, oauthPick);