@bastani/atomic 0.9.28-alpha.1 → 0.9.28-alpha.3

package/dist/extensions/mcp/index.js

@@ -212,8 +212,6 @@ export function createMcpExtension(options = {}) {
         let overridden = [];
         let latestContext;
         let subagentMcpDirectTools;
-        /** Between session_start and session_shutdown. Registrations before that are read on session_start. */
-        let sessionActive = false;
         let autoEnableCodemode = true;
         /** Whether the "codemode tools unreachable" warning was shown since the session started. */
         let warnedUnreachable = false;
@@ -221,8 +219,6 @@ export function createMcpExtension(options = {}) {
         /** Whether a prompt already waited for the startup connections since the session started. */
         let waitedForStartup = false;
         const startupWaitMs = options.startupWaitMs ?? DEFAULT_STARTUP_WAIT_MS;
-        /** Bumped on every session start and shutdown so a runtime load that resolves late is dropped. */
-        let generation = 0;
         /** Working directory of the session, for stdio servers. */
         let sessionCwd = process.cwd();
         let credentials = options.credentials;
@@ -234,6 +230,23 @@ export function createMcpExtension(options = {}) {
             ((entry, patch) => updateMcpServerConfig(entry.override ?? entry.source, entry.name, patch, {
                 override: entry.override !== undefined,
             }));
+        /**
+         * Lifetime of the current session: aborted on session_shutdown, and before the first session_start
+         * (registrations before that are read on session_start). Work captures `session.signal` when it
+         * starts and drops late results once it aborted, for example a runtime load that resolves after
+         * shutdown. Work that outlives the command that started it also stops on the signal and is
+         * registered with `track`, so shutdown can wait for its cleanup.
+         */
+        let session = new AbortController();
+        session.abort();
+        /** Work registered with `track`: manager actions opening or closing connections, and sign-ins. */
+        const backgroundActions = new Set();
+        const track = (work) => {
+            const task = work.then(() => undefined, () => undefined);
+            backgroundActions.add(task);
+            void task.finally(() => backgroundActions.delete(task));
+            return work;
+        };
         const listeners = new Set();
         const emitChange = () => {
             for (const listener of listeners)
@@ -357,8 +370,20 @@ export function createMcpExtension(options = {}) {
                     exposure,
                     namespace,
                     timeoutMs: connection.timeoutMs,
-                    getClient: async () => connection,
-                    readableResources: () => resourceServers().includes(connection),
+                    getClient: async () => {
+                        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),
                 });
                 markNativeMcpToolDefinition(definition, { server, tool: tool.name });
                 definitions.set(definition.name, definition);
@@ -493,52 +518,59 @@ export function createMcpExtension(options = {}) {
                 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();
             return connection;
         };
-        /**
-         * 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
-         * meanwhile. The returned promise rejects when the MCP runtime cannot be loaded.
-         */
-        const startConnection = (server, isCurrent, after) => {
-            // A contribution can be replaced without changing the session generation.
-            const ownsServer = () => isCurrent() && servers.includes(server) && isEnabled(server);
+        const startConnection = (server, after) => {
+            const { signal } = session;
+            const attempt = Symbol();
+            server.attempt = attempt;
+            const isCurrent = () => !signal.aborted && server.attempt === attempt && isEnabled(server) && servers.includes(server);
             const ready = (async () => {
                 await after;
-                if (!ownsServer())
-                    return;
-                const connection = await createConnection(server, ownsServer);
-                if (!connection)
+                if (!isCurrent())
                     return;
-                if (!ownsServer()) {
-                    await connection.close();
+                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. */
         const waitForServers = async (waiting, signal) => {
-            const ready = waiting.flatMap((server) => (server.ready ? [server.ready] : []));
-            if (ready.length === 0 || signal?.aborted)
-                return;
-            let onAbort;
-            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;
+                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;
+            }
         };
         /**
          * One message for everything that needs the user after startup, or only for `only`, servers
@@ -573,20 +605,23 @@ export function createMcpExtension(options = {}) {
             server.entry = { ...entry, config: { ...entry.config, ...patch } };
             return undefined;
         };
-        const signIn = async (server, prompt) => {
+        /** Sign in, then reconnect. Returns an error message on failure. `cancel` or the session's end stops the sign-in. */
+        const signIn = async (server, prompt, cancel) => {
             const connection = server.connection;
             const url = connection?.oauthUrl;
             if (!connection || !url)
                 return `MCP server "${server.entry.name}" does not use OAuth.`;
             const runtime = await loadMcpRuntime();
+            const signal = cancel ? AbortSignal.any([session.signal, cancel]) : session.signal;
             try {
-                await runtime.signInMcpServer({
+                await track(runtime.signInMcpServer({
                     serverUrl: url,
                     store: getCredentials(runtime).forServer(server.entry.name, url),
                     settings: connection.oauthSettings(),
                     challenge: connection.challenge,
                     prompt,
-                });
+                    signal,
+                }));
             }
             catch (error) {
                 if (error instanceof runtime.McpSignInCancelledError)
@@ -616,8 +651,14 @@ export function createMcpExtension(options = {}) {
             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) {
@@ -631,13 +672,24 @@ export function createMcpExtension(options = {}) {
                 return failed;
             if (!enabled) {
                 const connection = server.connection;
+                server.attempt = undefined;
                 server.connection = undefined;
+                if (connection) {
+                    server.closing = connection.close();
+                }
+                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;
         };
         const setExposure = (server, exposure) => {
@@ -787,8 +839,10 @@ export function createMcpExtension(options = {}) {
         };
         const signInWithUi = (ui, server) => {
             const title = `Sign in to ${server.entry.name}`;
+            const cancel = new AbortController();
+            const status = (message) => ui.status(title, message, () => cancel.abort());
             let authorizationUrl = "";
-            ui.status(title, "Contacting the authorization server…");
+            status("Contacting the authorization server…");
             return signIn(server, {
                 showAuthorizationUrl: (url) => {
                     authorizationUrl = url.href;
@@ -796,22 +850,36 @@ export function createMcpExtension(options = {}) {
                 },
                 promptForRedirectUrl: async (signal) => {
                     const value = await ui.redirectUrl(title, authorizationUrl, signal);
-                    ui.status(title, "Connecting…");
+                    status("Connecting…");
                     return value;
                 },
-            });
+            }, cancel.signal);
+        };
+        const runInBackground = (ctx, server, operation) => {
+            const ready = operation();
+            const attempt = server.attempt;
+            const { signal } = session;
+            const finish = (message) => {
+                if (signal.aborted || server.attempt !== attempt || !servers.includes(server))
+                    return;
+                server.message = message;
+                ensureDiscoveryActive(ctx);
+                emitChange();
+            };
+            void track(ready.then(finish, (error) => finish(errorMessage(error))));
         };
         const runAction = async (ui, ctx, server, action) => {
-            const { name } = server.entry;
+            const { signal } = session;
             let message;
             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);
+                    runInBackground(ctx, server, async () => {
+                        await reconnect(server);
+                        return undefined;
+                    });
                     break;
                 case "signout":
                     await signOut(server);
@@ -827,11 +895,13 @@ export function createMcpExtension(options = {}) {
                 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;
                 }
             }
+            // The session ended meanwhile (for example during a sign-in), which made ctx stale.
+            if (signal.aborted)
+                return;
             server.message = message;
             ensureDiscoveryActive(ctx);
             emitChange();
@@ -915,6 +985,7 @@ export function createMcpExtension(options = {}) {
                 ctx.ui.notify(`Signing in to MCP server "${name}" requires interactive mode.`, "error");
                 return;
             }
+            const { signal } = session;
             let failure;
             if (ctx.mode === "tui") {
                 let completed = false;
@@ -934,6 +1005,9 @@ export function createMcpExtension(options = {}) {
                     promptForRedirectUrl: (signal) => ctx.ui.input(`Waiting for sign-in to "${name}". If the browser cannot reach this machine, paste the URL it was redirected to.`, "http://127.0.0.1:.../callback?code=...", { signal }),
                 });
             }
+            // The session ended meanwhile, which cancelled the sign-in and made ctx stale.
+            if (signal.aborted)
+                return;
             if (failure) {
                 ctx.ui.notify(failure, failure === "Sign-in cancelled." ? "info" : "error");
                 return;
@@ -952,8 +1026,8 @@ export function createMcpExtension(options = {}) {
             waitedForStartup = false;
             sessionCwd = ctx.cwd;
             modelRegistry = ctx.modelRegistry;
-            const current = ++generation;
-            sessionActive = true;
+            session = new AbortController();
+            const { signal } = session;
             configuredEntries = loaded.servers;
             const registered = registeredServers();
             overridden = registered.overridden;
@@ -970,10 +1044,9 @@ export function createMcpExtension(options = {}) {
             // The MCP client loads only now, so sessions without servers never pay for it. Waiting one
             // 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())
+                if (!signal.aborted)
                     reportProblems(ctx);
             })
                 .catch((error) => {
@@ -1028,17 +1101,21 @@ export function createMcpExtension(options = {}) {
             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)
+            const readyServers = servers.filter((server) => isEnabled(server) && server.ready);
+            if (readyServers.length === 0)
                 return;
             let waiting = [];
             if (isCodemodeDiscoveryTool(tool)) {
                 const { code } = event.input;
                 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 (isToolSearchDiscoveryTool(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);
         });
@@ -1049,9 +1126,9 @@ export function createMcpExtension(options = {}) {
         });
         // Servers registered or unregistered during the session connect or disconnect right away.
         pi.onMcpServerContributionsChanged?.(() => {
-            if (!sessionActive)
+            const { signal } = session;
+            if (signal.aborted)
                 return;
-            const current = generation;
             const registered = registeredServers();
             overridden = registered.overridden;
             const next = new Map(registered.servers.map((server) => [server.entry.name, server]));
@@ -1067,22 +1144,21 @@ export function createMcpExtension(options = {}) {
                 ensureDiscoveryActive(latestContext);
             void Promise.all(removed.map((server) => server.connection?.close())).then(async () => {
                 const connecting = added.filter(isEnabled);
-                if (current !== generation || connecting.length === 0)
+                if (signal.aborted || connecting.length === 0)
                     return;
                 try {
-                    await Promise.all(connecting.map((server) => startConnection(server, () => current === generation)));
-                    if (current === generation)
+                    await Promise.all(connecting.map((server) => startConnection(server)));
+                    if (!signal.aborted)
                         reportProblems(latestContext, connecting);
                 }
                 catch (error) {
-                    if (current === generation)
+                    if (!signal.aborted)
                         latestContext?.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error");
                 }
             });
         });
         pi.on("session_shutdown", async () => {
-            sessionActive = false;
-            generation++;
+            session.abort();
             const closing = connections();
             // The retiring registry is discarded. Do not register hidden tools or call a
             // captured UI here: reload has already revoked this generation's authority.
@@ -1095,7 +1171,7 @@ export function createMcpExtension(options = {}) {
             servers = [];
             latestContext = undefined;
             subagentMcpDirectTools = undefined;
-            await Promise.all(closing.map((connection) => connection.close()));
+            await Promise.all([...backgroundActions, ...closing.map((connection) => connection.close())]);
         });
         pi.registerCommand("mcp", {
             description: "Manage MCP servers: sign in, reconnect, enable or disable, and change exposure",
@@ -1121,19 +1197,21 @@ export function createMcpExtension(options = {}) {
                 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
+                    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);