fix(coding-agent): keep MCP manager responsive (closes #10562)

Open the manager before startup connections finish and keep connection actions in the background. Serialize reconnects and replacements, wait for current readiness, and resolve the active connection when executing tools.
This commit is contained in:
Armin Ronacher 2026-10-07 08:35:24 +02:00
parent 6b58544540
commit 2db5e359bf
2 changed files with 136 additions and 46 deletions

View file

@ -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 `<console_output>` 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

View file

@ -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<void>;
/** Cleanup of a detached connection; a replacement waits for it before opening a transport. */
closing?: Promise<void>;
}
const EXPOSURE_DESCRIPTIONS: Record<Exclude<McpExposure, "hidden">, 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<Promise<void>>();
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<McpServerConnection> => {
const createConnection = async (
server: McpServer,
isCurrent: () => boolean,
): Promise<McpServerConnection | undefined> => {
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<unknown>,
): Promise<void> => {
const startConnection = (server: McpServer, after?: Promise<unknown>): Promise<void> => {
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<void>((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<void>((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<string | undefined> => {
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<string | undefined>,
) => {
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);