/**
 * Built-in MCP integration.
 *
 * Connects the servers from `mcp.json` and the servers extensions register with
 * `pi.registerMcpServer()` when a session starts, and servers registered later right away. A server
 * in `mcp.json` takes precedence over a registered server of the same name. Connections run in the
 * background: the first prompt waits only for servers with `direct` tools, and codemode scripts,
 * `tool_search`, and the resource tools wait for the servers they need when they run. Tools are
 * registered as `mcp__<server>__<tool>`. By default (`"exposure": "codemode"`) the tools are only
 * callable from codemode scripts, which keeps MCP tools out of the model's tool declarations and
 * the codemode description: scripts find the tools with `searchTools()` and the server instructions
 * with `describeNamespace()`. The codemode tool is activated for that unless `autoEnableCodemode` is
 * false. `"deferred"` declares the tools to the model once the `tool_search` tool loads them, and
 * activates `tool_search` instead of codemode.
 * `"exposure": "direct"` declares them to the model right away, and `"hidden"` makes them
 * unreachable. `toolExposure` overrides the exposure of single tools. Servers with resources are
 * reached through Codex's `list_mcp_resources`, `list_mcp_resource_templates`, and
 * `read_mcp_resource` tools (resources.ts).
 *
 * Every call runs through pi's tool pipeline, so `tool_call`/`tool_result` hooks and permission
 * extensions apply to MCP tools the same way they do to built-in tools.
 *
 * Problems found at startup (config errors, failed connections, servers that need a sign-in) are
 * reported once. `/mcp` opens a manager to sign in, reconnect, enable or disable servers, and change
 * their exposure; the last two are saved to the `mcp.json` that defines the server, or apply to the
 * current session for registered servers.
 */
import { join, resolve } from "node:path";
import { createMcpToolName, mcpNamespace } from "@earendil-works/pi-mcp";
import { getAgentDir } from "../../config.js";
import { sliceHead } from "../../core/tools/truncate.js";
import { openBrowser } from "../../utils/open-browser.js";
import { CODEMODE_TOOL_NAME } from "../codemode/tool.js";
import { TOOL_SEARCH_TOOL_NAME } from "../tool-search/tool.js";
import { getMcpToolExposure, loadMcpConfig, updateMcpServerConfig, } from "./config.js";
import { createMcpResourceToolDefinitions, LIST_MCP_RESOURCE_TEMPLATES_TOOL, LIST_MCP_RESOURCES_TOOL, READ_MCP_RESOURCE_TOOL, } from "./resources.js";
import { loadMcpRuntime } from "./runtime.lazy.js";
import { createMcpToolDefinition, createMcpToolRenderers } from "./tools.js";
import { showMcpManager } from "./ui.js";
const DEFAULT_STARTUP_WAIT_MS = 10_000;
const RESOURCE_TOOL_NAMES = new Set([
    LIST_MCP_RESOURCES_TOOL,
    LIST_MCP_RESOURCE_TEMPLATES_TOOL,
    READ_MCP_RESOURCE_TOOL,
]);
const EXPOSURE_DESCRIPTIONS = {
    codemode: "called from codemode scripts, which find them with searchTools()",
    deferred: "not declared until tool_search loads them, then called directly; no codemode needed",
    direct: "declared to the model like built-in tools",
};
function errorMessage(error) {
    return error instanceof Error ? error.message : String(error);
}
function firstLine(text) {
    return text.split("\n", 1)[0] ?? "";
}
function isEnabled(server) {
    return server.entry.config.enabled !== false;
}
function exposureOf(entry) {
    return entry.config.exposure ?? "codemode";
}
/** Exposures the server's tools can have, known from its config before it connects. */
function configuredExposures(entry) {
    return new Set([exposureOf(entry), ...Object.values(entry.config.toolExposure ?? {})]);
}
/** Whether some of the server's tools are declared to the model, so the first prompt waits for them. */
function hasDirectTools(entry) {
    return configuredExposures(entry).has("direct");
}
/** Whether some of the server's tools are reached through codemode or tool_search. */
function hasIndirectTools(entry) {
    const exposures = configuredExposures(entry);
    return exposures.has("codemode") || exposures.has("deferred");
}
/** Name of the system prompt section that lists the servers whose tools are not declared. */
export const MCP_SERVERS_SECTION = "mcp_servers";
/** Characters of a server description in the section, as Codex allows for deferred namespaces. */
const MAX_SERVER_DESCRIPTION_CHARS = 250;
/**
 * Characters of the whole section. Descriptions shrink to fit; when the server lines alone do not
 * fit, the last servers are left out and counted in a closing line.
 */
export const MAX_SERVERS_SECTION_CHARS = 4096;
/** The section's first line. It explains only the ways of reaching tools that the listed servers use. */
function serversSectionIntro(reaches) {
    let intro = "MCP servers whose tools are not declared to you.";
    if (reaches.has("codemode"))
        intro += " Call the tools of `codemode` servers from codemode scripts.";
    if (reaches.has("tool_search"))
        intro += " Load the tools of `tool_search` servers with `tool_search`.";
    return intro;
}
function truncate(text, max) {
    if (text.length <= max)
        return text;
    return max <= 1 ? "" : `${sliceHead(text, max - 1).trimEnd()}…`;
}
/** First line of the configured description, or of the server instructions once connected. */
function serverSummary(server) {
    const text = server.entry.config.description?.trim() || server.connection?.instructions || "";
    return (text.split("\n", 1)[0] ?? "").trim();
}
/**
 * The `mcp_servers` section: every enabled server with codemode or deferred tools, with how its
 * tools are reached and a one-line summary. The model learns of the servers from it, since neither
 * codemode nor tool_search lists them. Undefined when there are no such servers.
 */
export function renderServersSection(servers) {
    const listed = servers
        .filter((server) => isEnabled(server) && hasIndirectTools(server.entry))
        .sort((a, b) => a.entry.name.localeCompare(b.entry.name));
    if (listed.length === 0)
        return undefined;
    const reaches = listed.map((server) => configuredExposures(server.entry).has("codemode") ? "codemode" : "tool_search");
    const intro = serversSectionIntro(new Set(reaches));
    const heads = listed.map((server, index) => `- ${mcpNamespace(server.entry.name)} (${reaches[index]})`);
    const omitted = (count) => count > 0 ? [`- … ${count} more server${count === 1 ? "" : "s"}; find their tools with searchTools()`] : [];
    // Characters of the intro, the first `kept` server lines without descriptions, and the omission line.
    const size = (kept) => [intro, ...heads.slice(0, kept), ...omitted(listed.length - kept)].join("\n").length;
    let kept = listed.length;
    while (kept > 0 && size(kept) > MAX_SERVERS_SECTION_CHARS)
        kept--;
    // Each description also takes a ": " separator.
    const perServer = kept === 0
        ? 0
        : Math.min(MAX_SERVER_DESCRIPTION_CHARS, Math.floor((MAX_SERVERS_SECTION_CHARS - size(kept)) / kept) - 2);
    const lines = listed.slice(0, kept).map((server, index) => {
        const summary = perServer > 0 ? truncate(serverSummary(server), perServer) : "";
        return summary ? `${heads[index]}: ${summary}` : heads[index];
    });
    return [intro, ...lines, ...omitted(listed.length - kept)].join("\n");
}
/**
 * Whether a call that reaches `reach` (see `ToolDefinition.reachesTools`) needs the server: it
 * reaches every tool, or one of its words contains the server's namespace.
 */
function reachesServer(reach, server) {
    if (reach === "all")
        return true;
    const namespace = mcpNamespace(server);
    return reach.some((word) => word.includes(namespace));
}
/** Short state for lists and the startup report. `withError` appends the first line of a failure. */
function describeState(server, withError = true) {
    if (!isEnabled(server))
        return "disabled";
    const connection = server.connection;
    if (!connection)
        return "starting";
    switch (connection.state) {
        case "needs-auth":
            return "needs sign-in";
        case "failed":
            return withError ? `failed: ${firstLine(connection.error ?? "unknown error")}` : "failed";
        case "connected": {
            const { tools, hasResources } = connection;
            return `connected · ${tools.length} tool${tools.length === 1 ? "" : "s"}${hasResources ? " · resources" : ""}`;
        }
        case "connecting":
            return "connecting…";
        default:
            return connection.state;
    }
}
/** Servers that need the user first. */
function attentionRank(server) {
    if (!isEnabled(server))
        return 5;
    switch (server.connection?.state) {
        case "needs-auth":
            return 0;
        case "failed":
            return 1;
        case "disconnected":
            return 2;
        case "connected":
            return 4;
        default:
            return 3;
    }
}
function describeTransport(entry) {
    const { config } = entry;
    if ("url" in config)
        return config.url;
    return [config.command, ...(config.args ?? [])].join(" ");
}
const MCP_USAGE = "Usage: /mcp, /mcp login [server], /mcp logout [server], /mcp reconnect [server]";
export function createMcpExtension(options = {}) {
    return (pi) => {
        let servers = [];
        /** Servers from `mcp.json`, which take precedence over registered servers of the same name. */
        let configuredEntries = [];
        let configErrors = [];
        /** The trusted project's `mcp.json`, where `/mcp` saves project overrides of global servers. */
        let projectConfig;
        /** Registered servers that `mcp.json` overrides, shown in `/mcp`. */
        let overridden = [];
        let autoEnableCodemode = true;
        /** Whether the "codemode tools unreachable" warning was shown since the session started. */
        let warnedUnreachable = false;
        let pending;
        /** Whether a prompt already waited for the startup connections since the session started. */
        let waitedForStartup = false;
        const startupWaitMs = options.startupWaitMs ?? DEFAULT_STARTUP_WAIT_MS;
        /** Working directory of the session, for stdio servers. */
        let sessionCwd = process.cwd();
        let credentials = options.credentials;
        /** The session's model registry, which resolves `auth.provider` tokens. */
        let modelRegistry;
        let serverLog;
        const openUrl = options.openUrl ?? openBrowser;
        const updateConfig = options.updateConfig ??
            ((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)
                listener();
        };
        const subscribe = (listener) => {
            listeners.add(listener);
            return () => {
                listeners.delete(listener);
            };
        };
        const connections = () => servers.flatMap((server) => (server.connection ? [server.connection] : []));
        const findServer = (name) => servers.find((server) => server.entry.name === name);
        /** Servers extensions registered, except names `mcp.json` defines, which take precedence. */
        const registeredServers = () => {
            const registered = [];
            const overriddenNames = [];
            for (const { name, config, extensionPath } of pi.getMcpServers()) {
                const configured = configuredEntries.find((entry) => mcpNamespace(entry.name) === mcpNamespace(name));
                if (configured) {
                    overriddenNames.push(`"${name}" registered by ${extensionPath} is overridden by "${configured.name}" in ${configured.source}`);
                    continue;
                }
                registered.push({
                    entry: { name, config, source: extensionPath, scope: "extension" },
                    registeredConfig: JSON.stringify(config),
                });
            }
            return { servers: registered, overridden: overriddenNames };
        };
        const getCredentials = (runtime) => {
            credentials ??= new runtime.McpOAuthCredentialStore();
            return credentials;
        };
        const getServerLog = (runtime) => {
            serverLog ??= new runtime.McpServerLog(options.logPath ?? join(getAgentDir(), "mcp.log"));
            return serverLog;
        };
        /** pi tool name to the `<server>\0<tool>` it was assigned to, so names stay unique and stable. */
        const toolOwners = new Map();
        /** Tool names currently offered by each server. */
        const serverTools = new Map();
        /** Last definition registered under each tool name, to re-register withdrawn tools as hidden. */
        const definitions = new Map();
        // A resumed session renders calls to MCP tools before their server connected, if it ever does.
        pi.registerToolRenderer((toolName, next) => {
            const match = /^mcp__(.+?)__(.+)$/.exec(toolName);
            return next() ?? (match ? createMcpToolRenderers(`${match[1]}/${match[2]}`) : undefined);
        });
        const registerTools = (connection) => {
            const server = connection.entry.name;
            const entry = findServer(server)?.entry ?? connection.entry;
            const description = entry.config.description?.trim();
            const namespace = {
                name: mcpNamespace(server),
                // `describeNamespace("docs")` finds `mcp__docs`.
                aliases: [server],
                ...(description ? { description } : {}),
                ...(connection.instructions ? { instructions: connection.instructions } : {}),
            };
            const previous = serverTools.get(server) ?? new Set();
            const current = new Set();
            // Like Codex, all tools whose names sanitize to the same name get the hash suffix, so which
            // one would keep the plain name does not depend on the order of the list.
            const plain = [...new Set(connection.tools.map((tool) => tool.name))].map((tool) => createMcpToolName(server, tool));
            const assignName = (tool, owner) => {
                const name = createMcpToolName(server, tool, (candidate) => {
                    const existing = toolOwners.get(candidate);
                    return ((existing !== undefined && existing !== owner) ||
                        current.has(candidate) ||
                        plain.indexOf(candidate) !== plain.lastIndexOf(candidate));
                });
                toolOwners.set(name, owner);
                current.add(name);
                return name;
            };
            for (const tool of connection.tools) {
                const definition = createMcpToolDefinition({
                    server,
                    tool,
                    name: assignName(tool.name, `${server}\0${tool.name}`),
                    exposure: getMcpToolExposure(entry.config, tool.name),
                    namespace,
                    timeoutMs: connection.timeoutMs,
                    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);
            }
            serverTools.set(server, current);
            // Tools cannot be unregistered, so tools the server dropped are re-registered as hidden. When
            // the server offers them again they are registered with their configured exposure above.
            for (const name of previous) {
                const definition = definitions.get(name);
                if (!current.has(name) && definition)
                    pi.registerTool({ ...definition, exposure: "hidden" });
            }
            syncResourceTools();
        };
        /** Make a disabled server's tools unreachable. */
        const hideTools = (server) => {
            for (const name of serverTools.get(server) ?? []) {
                const definition = definitions.get(name);
                if (definition)
                    pi.registerTool({ ...definition, exposure: "hidden" });
            }
            serverTools.set(server, new Set());
            syncResourceTools();
        };
        /** Enabled servers with resources whose exposure is not `hidden`, which the resource tools reach. */
        const serversWithResources = () => servers.filter((server) => server.connection?.hasResources && isEnabled(server) && exposureOf(server.entry) !== "hidden");
        const resourceServers = () => serversWithResources().flatMap((server) => (server.connection ? [server.connection] : []));
        /** Exposure the resource tools were last registered with; undefined until a server has resources. */
        let resourceToolsExposure;
        /**
         * Register the resource tools with the widest exposure of the servers they reach: `direct` when
         * one of them is direct, and so on. They are hidden when no server has resources.
         */
        const syncResourceTools = () => {
            const exposures = new Set(serversWithResources().map((server) => exposureOf(server.entry)));
            const exposure = ["direct", "codemode", "deferred"].find((candidate) => exposures.has(candidate));
            const next = exposure ?? "hidden";
            if (next === resourceToolsExposure || (resourceToolsExposure === undefined && next === "hidden"))
                return;
            const wasDirect = resourceToolsExposure === "direct";
            resourceToolsExposure = next;
            const resourceDefinitions = createMcpResourceToolDefinitions({ exposure: next, servers: resourceServers });
            for (const definition of resourceDefinitions)
                pi.registerTool(definition);
            if (wasDirect) {
                const names = new Set(resourceDefinitions.map((definition) => definition.name));
                pi.setActiveTools(pi.getActiveTools().filter((name) => !names.has(name)));
            }
        };
        /**
         * Tools that are not declared to the model are reached through the codemode tool (scripts call
         * them) or the tool_search tool (it declares them). Either reaches every such tool. Activate the
         * one the tools' exposure asks for: codemode for `codemode` unless `autoEnableCodemode` is false,
         * tool_search for `deferred`.
         */
        const ensureDiscoveryActive = (ctx) => {
            // From the config, so the tool is active before the servers connect. Resource tools share
            // their server's exposure.
            const exposures = new Set();
            for (const server of servers) {
                if (isEnabled(server))
                    for (const exposure of configuredExposures(server.entry))
                        exposures.add(exposure);
            }
            const needsCodemode = exposures.has("codemode");
            const needsToolSearch = exposures.has("deferred");
            if (!needsCodemode && !needsToolSearch)
                return;
            // Only a tool of these names that declares `reachesTools` can reach MCP tools; another
            // extension's tool of the same name without it is never activated.
            const tools = pi.getAllTools();
            const hasReachingTool = (name) => tools.some((tool) => tool.name === name && tool.reachesTools);
            const hasCodemode = hasReachingTool(CODEMODE_TOOL_NAME);
            const hasToolSearch = hasReachingTool(TOOL_SEARCH_TOOL_NAME);
            const active = pi.getActiveTools();
            const activate = [];
            if (needsCodemode && hasCodemode && autoEnableCodemode && !active.includes(CODEMODE_TOOL_NAME)) {
                activate.push(CODEMODE_TOOL_NAME);
            }
            if (needsToolSearch && hasToolSearch && !active.includes(TOOL_SEARCH_TOOL_NAME)) {
                activate.push(TOOL_SEARCH_TOOL_NAME);
            }
            if (activate.length > 0)
                pi.setActiveTools([...active, ...activate]);
            const reachable = [...active, ...activate];
            if (hasCodemode && reachable.includes(CODEMODE_TOOL_NAME))
                return;
            if (hasToolSearch && reachable.includes(TOOL_SEARCH_TOOL_NAME))
                return;
            if (warnedUnreachable)
                return;
            warnedUnreachable = true;
            const reason = needsCodemode && hasCodemode && !autoEnableCodemode ? " (autoEnableCodemode is false)" : "";
            ctx.ui.notify(`MCP tools are only reachable from the codemode or tool_search tool, but neither is active${reason}; they cannot be called.`, "warning");
        };
        /**
         * Stored tokens of servers waiting for a sign-in, as they were when the sign-in was needed.
         * `pi mcp login` in another process (for example run by the agent) changes them.
         */
        const tokensAtSignIn = new Map();
        const storedTokens = (connection) => {
            const url = connection.oauthUrl;
            return url && credentials ? JSON.stringify(credentials.tokens(connection.name, url) ?? null) : "null";
        };
        const onConnectionChange = (connection) => {
            if (connection.state !== "needs-auth")
                tokensAtSignIn.delete(connection);
            else if (!tokensAtSignIn.has(connection))
                tokensAtSignIn.set(connection, storedTokens(connection));
            emitChange();
        };
        /** Reconnect servers that need a sign-in when their credentials were stored since. */
        const reconnectSignedIn = async (ctx) => {
            const signedIn = [...tokensAtSignIn].filter(([connection, tokens]) => storedTokens(connection) !== tokens);
            if (signedIn.length === 0)
                return;
            for (const [connection] of signedIn)
                tokensAtSignIn.delete(connection);
            await Promise.allSettled(signedIn.map(([connection]) => connection.reconnect()));
            ensureDiscoveryActive(ctx);
        };
        /** Create the server's connection, loading the MCP runtime on first use. */
        const createConnection = async (server, isCurrent) => {
            const runtime = await loadMcpRuntime();
            if (!isCurrent())
                return undefined;
            const connection = new runtime.McpServerConnection({
                entry: server.entry,
                cwd: sessionCwd,
                createTransport: options.createTransport ?? runtime.createDefaultTransport,
                credentials: getCredentials(runtime),
                providerToken: async (provider) => modelRegistry?.getApiKeyForProvider(provider),
                log: getServerLog(runtime),
                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`, unless disabled, replaced, or shut down
         * meanwhile. The returned promise rejects when the MCP runtime cannot be loaded.
         */
        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 (!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 the latest connection attempts, including ones queued while waiting, or an abort. */
        const waitForServers = async (waiting, signal) => {
            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
         * that connected later.
         */
        const reportProblems = (ctx, only) => {
            const lines = only ? [] : configErrors.map((error) => `config: ${error}`);
            for (const server of only ?? servers) {
                const state = server.connection?.state;
                if (state === "needs-auth" || state === "failed")
                    lines.push(`${server.entry.name}: ${describeState(server)}`);
            }
            if (lines.length === 0)
                return;
            ctx.ui.notify(`MCP servers need attention:\n${lines.map((line) => `  ${line}`).join("\n")}\nRun /mcp to fix.`, "warning");
        };
        /**
         * Save a config change; returns an error message when the file could not be updated. Changes to
         * registered servers only apply to the current session. `inProject` adds a project override.
         */
        const saveConfig = (server, patch, inProject = false) => {
            const override = inProject ? projectConfig : server.entry.override;
            const entry = override ? { ...server.entry, override } : server.entry;
            if (entry.scope !== "extension") {
                try {
                    updateConfig(entry, patch);
                }
                catch (error) {
                    return `Could not update ${entry.override ?? entry.source}: ${errorMessage(error)}`;
                }
            }
            server.entry = { ...entry, config: { ...entry.config, ...patch } };
            return undefined;
        };
        /** 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 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)
                    return "Sign-in cancelled.";
                return `Sign-in failed: ${errorMessage(error)}`;
            }
            // The challenge that asked for this sign-in (for example for more scope) is answered.
            connection.challenge = undefined;
            try {
                await connection.reconnect();
            }
            catch (error) {
                return `Signed in, but ${errorMessage(error)}`;
            }
            return undefined;
        };
        const signOut = async (server) => {
            const connection = server.connection;
            const url = connection?.oauthUrl;
            if (!connection || !url)
                return false;
            const removed = getCredentials(await loadMcpRuntime()).remove(server.entry.name, url);
            await connection.signOut();
            return removed;
        };
        const reconnect = async (server) => {
            const connection = server.connection;
            if (!connection)
                return `MCP server "${server.entry.name}" is disabled.`;
            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);
            }
        };
        /** Returns an error message when the config could not be saved; connection errors show in the state. */
        const setEnabled = async (server, enabled, inProject = false) => {
            const failed = saveConfig(server, { enabled }, inProject);
            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();
                try {
                    await closing;
                }
                finally {
                    if (server.closing === closing)
                        server.closing = undefined;
                }
                return undefined;
            }
            await startConnection(server, server.closing);
            return undefined;
        };
        const setExposure = (server, exposure) => {
            const failed = saveConfig(server, { exposure });
            if (failed)
                return failed;
            if (server.connection?.state === "connected")
                registerTools(server.connection);
            syncResourceTools();
            // Tools no longer exposed directly leave the declared set; direct tools are activated on registration.
            const indirect = new Set(pi
                .getAllTools()
                .filter((tool) => tool.exposure !== "direct")
                .map((tool) => tool.name));
            const tools = serverTools.get(server.entry.name) ?? new Set();
            pi.setActiveTools(pi.getActiveTools().filter((name) => !tools.has(name) || !indirect.has(name)));
            emitChange();
            return undefined;
        };
        // ---------------------------------------------------------------------------------------
        // Manager (`/mcp` in the TUI)
        // ---------------------------------------------------------------------------------------
        const notices = () => [
            ...configErrors.map((error) => `config: ${error}`),
            ...overridden.map((line) => `overridden: ${line}`),
        ];
        const serversMenu = () => ({
            title: "MCP servers",
            error: notices().join("\n") || undefined,
            items: [...servers]
                .sort((a, b) => attentionRank(a) - attentionRank(b) || a.entry.name.localeCompare(b.entry.name))
                .map((server) => ({
                value: server.entry.name,
                label: server.entry.name,
                description: `${describeState(server)} · ${exposureOf(server.entry)} · ${server.entry.override ? "global, project override" : (server.entry.scope ?? server.entry.source)}`,
            })),
            empty: `No MCP servers configured. Add them to ${resolve(getAgentDir(), "mcp.json")} or .pi/mcp.json.`,
            confirmLabel: "manage",
            cancelLabel: "close",
        });
        const serverMenu = (name) => {
            const server = findServer(name);
            if (!server) {
                return {
                    title: name,
                    items: [],
                    empty: "This server is no longer configured.",
                    confirmLabel: "",
                    cancelLabel: "back",
                };
            }
            const { entry, connection } = server;
            const saved = entry.scope === "extension"
                ? "for this session"
                : entry.override
                    ? "saved to the project mcp.json"
                    : entry.scope
                        ? `saved to the ${entry.scope} mcp.json`
                        : "saved to mcp.json";
            // Global servers without an override can be turned on or off for the trusted project alone.
            const inProject = entry.scope === "global" && !entry.override && projectConfig !== undefined;
            const inProjectSaved = "saved to the project mcp.json";
            const items = [];
            if (!isEnabled(server)) {
                items.push({ value: "enable", label: "Enable", description: saved });
                if (inProject) {
                    items.push({ value: "enable-project", label: "Enable in this project", description: inProjectSaved });
                }
            }
            else {
                const state = connection?.state;
                if (state === "needs-auth")
                    items.push({ value: "signin", label: "Sign in", description: "opens the browser" });
                if (state === "connected" && connection) {
                    items.push({ value: "tools", label: "Tools", description: `${connection.tools.length} offered` });
                }
                if (state === "failed" || state === "disconnected" || state === "connected" || state === "needs-auth") {
                    items.push({ value: "reconnect", label: "Reconnect" });
                }
                if (state === "connected" && connection?.oauthUrl) {
                    items.push({ value: "signout", label: "Sign out", description: "deletes the stored credentials" });
                }
                items.push({ value: "exposure", label: "Exposure", description: exposureOf(entry) });
                items.push({ value: "disable", label: "Disable", description: saved });
                if (inProject) {
                    items.push({ value: "disable-project", label: "Disable in this project", description: inProjectSaved });
                }
            }
            const details = [
                describeTransport(entry),
                `${entry.scope ?? "config"}: ${entry.source}`,
                ...(entry.override ? [`project override: ${entry.override}`] : []),
                `State: ${describeState(server, false)}`,
            ];
            const error = [server.message, connection?.state === "connected" ? undefined : connection?.error]
                .filter((line) => line !== undefined)
                .join("\n");
            return {
                title: `MCP server ${name}`,
                details: details.join("\n"),
                error: error || undefined,
                items,
                selected: items[0]?.value,
                confirmLabel: "select",
                cancelLabel: "back",
            };
        };
        const showTools = async (ui, server) => {
            const exposure = exposureOf(server.entry);
            const overridden = Object.keys(server.entry.config.toolExposure ?? {}).length > 0;
            await ui.menu(() => ({
                title: `Tools of ${server.entry.name}`,
                details: `Exposure ${exposure}: ${exposure === "hidden" ? "unreachable" : EXPOSURE_DESCRIPTIONS[exposure]}${overridden ? "\nSome tools override it with toolExposure." : ""}`,
                items: (server.connection?.tools ?? []).map((tool) => {
                    const toolExposure = getMcpToolExposure(server.entry.config, tool.name);
                    const description = firstLine(tool.description ?? "");
                    return {
                        value: tool.name,
                        label: tool.name,
                        description: toolExposure === exposure ? description : `[${toolExposure}] ${description}`,
                    };
                }),
                empty: "The server offers no tools.",
                confirmLabel: "back",
                cancelLabel: "back",
            }));
        };
        const chooseExposure = async (ui, server) => {
            const current = exposureOf(server.entry);
            const choice = await ui.menu(() => ({
                title: `Exposure of ${server.entry.name}`,
                details: server.entry.scope === "extension"
                    ? `Applies to this session; the server is registered by ${server.entry.source}.`
                    : `Saved to ${server.entry.override ?? server.entry.source}.`,
                items: Object.keys(EXPOSURE_DESCRIPTIONS).map((exposure) => ({
                    value: exposure,
                    label: `${exposure === current ? "✓ " : "  "}${exposure}`,
                    description: EXPOSURE_DESCRIPTIONS[exposure],
                })),
                selected: current,
                confirmLabel: "save",
                cancelLabel: "back",
            }));
            if (!choice || choice === current)
                return undefined;
            return setExposure(server, choice);
        };
        /** Sign in with the manager view's sign-in screen, which shows the URL with a copy key and cancels with Esc. */
        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 = "";
            status("Contacting the authorization server…");
            return signIn(server, {
                showAuthorizationUrl: (url) => {
                    authorizationUrl = url.href;
                    openUrl(url.href);
                },
                promptForRedirectUrl: async (signal) => {
                    const value = await ui.redirectUrl(title, authorizationUrl, signal);
                    status("Connecting…");
                    return value;
                },
            }, cancel.signal);
        };
        /** Keep the subscribed menu usable while a connection opens or closes. */
        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 { signal } = session;
            let message;
            switch (action) {
                case "signin":
                    message = await signInWithUi(ui, server);
                    break;
                case "reconnect":
                    // 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);
                    break;
                case "tools":
                    await showTools(ui, server);
                    break;
                case "exposure":
                    message = await chooseExposure(ui, server);
                    break;
                case "enable":
                case "disable":
                case "enable-project":
                case "disable-project": {
                    const enable = action.startsWith("enable");
                    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();
        };
        const manage = async (ui, ctx) => {
            for (;;) {
                const name = await ui.menu(serversMenu, subscribe);
                if (!name)
                    return;
                for (;;) {
                    const action = await ui.menu(() => serverMenu(name), subscribe);
                    const server = findServer(name);
                    if (!action || !server)
                        break;
                    await runAction(ui, ctx, server, action);
                }
            }
        };
        // ---------------------------------------------------------------------------------------
        // Subcommands and plain status (no TUI)
        // ---------------------------------------------------------------------------------------
        const formatStatus = () => {
            if (servers.length === 0 && configErrors.length === 0 && overridden.length === 0) {
                return `No MCP servers configured. Add them to ${resolve(getAgentDir(), "mcp.json")} or .pi/mcp.json.`;
            }
            const lines = servers.map((server) => {
                const { name } = server.entry;
                const exposure = exposureOf(server.entry);
                const connection = server.connection;
                if (connection?.state === "needs-auth")
                    return `${name}: needs sign-in, run /mcp login ${name} (${exposure})`;
                const tools = connection?.state === "connected" ? `, ${connection.tools.length} tools` : "";
                const state = !isEnabled(server)
                    ? "disabled"
                    : connection?.state === "disconnected"
                        ? "disconnected, reconnects on next call"
                        : (connection?.state ?? "starting");
                const error = connection?.error && connection.state !== "connected"
                    ? `\n    ${connection.error.split("\n").join("\n    ")}`
                    : "";
                return `${name}: ${state}${tools} (${exposure})${error}`;
            });
            for (const error of configErrors)
                lines.push(`config error: ${error}`);
            for (const line of overridden)
                lines.push(`overridden: ${line}`);
            return lines.join("\n");
        };
        /** Resolve the server for a subcommand, asking when the name is omitted and ambiguous. */
        const pickServer = async (name, ctx, options) => {
            if (name) {
                const server = findServer(name);
                if (!server)
                    ctx.ui.notify(`No MCP server named "${name}".`, "error");
                else if (!options.eligible(server))
                    ctx.ui.notify(options.none, "error");
                return server && options.eligible(server) ? server : undefined;
            }
            const candidates = servers.filter(options.eligible);
            if (candidates.length === 0) {
                ctx.ui.notify(options.none, "info");
                return undefined;
            }
            const preferred = candidates.filter(options.preferred);
            if (candidates.length === 1)
                return candidates[0];
            if (preferred.length === 1)
                return preferred[0];
            const choice = await ctx.ui.select("MCP server", candidates.map((server) => server.entry.name));
            return candidates.find((server) => server.entry.name === choice);
        };
        const usesOAuth = (server) => server.connection?.oauthUrl !== undefined;
        const oauthPick = {
            eligible: usesOAuth,
            preferred: (server) => server.connection?.state === "needs-auth",
            none: "No enabled MCP server uses OAuth. Only HTTP servers without an Authorization header do.",
        };
        const loginCommand = async (server, ctx) => {
            const { name } = server.entry;
            if (!ctx.hasUI) {
                ctx.ui.notify(`Signing in to MCP server "${name}" requires interactive mode.`, "error");
                return;
            }
            const { signal } = session;
            let failure;
            if (ctx.mode === "tui") {
                await showMcpManager(ctx, async (ui) => {
                    failure = await signInWithUi(ui, server);
                });
            }
            else {
                failure = await signIn(server, {
                    showAuthorizationUrl: (url) => {
                        ctx.ui.notify(`Sign in to MCP server "${name}" in your browser:\n${url.href}`, "info");
                        openUrl(url.href);
                    },
                    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;
            }
            ensureDiscoveryActive(ctx);
            ctx.ui.notify(`Signed in to MCP server "${name}" (${server.connection?.tools.length ?? 0} tools).`, "info");
        };
        pi.on("session_start", (_event, ctx) => {
            const loaded = (options.loadConfig ?? defaultLoadConfig)(ctx);
            configErrors = loaded.errors;
            projectConfig = loaded.projectConfig;
            autoEnableCodemode = loaded.autoEnableCodemode ?? true;
            warnedUnreachable = false;
            waitedForStartup = false;
            sessionCwd = ctx.cwd;
            modelRegistry = ctx.modelRegistry;
            session = new AbortController();
            const { signal } = session;
            configuredEntries = loaded.servers;
            const registered = registeredServers();
            overridden = registered.overridden;
            servers = [...loaded.servers.map((entry) => ({ entry })), ...registered.servers];
            emitChange();
            // Codemode or tool_search is activated from the config: the first prompt does not wait for
            // servers whose tools are not declared to the model, and scripts or searches wait for them.
            ensureDiscoveryActive(ctx);
            const enabled = servers.filter(isEnabled);
            if (enabled.length === 0) {
                reportProblems(ctx);
                return;
            }
            // 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());
            pending = Promise.all(enabled.map((server) => startConnection(server, runtime)))
                .then(() => {
                if (!signal.aborted)
                    reportProblems(ctx);
            })
                .catch((error) => {
                // The session may have been disposed meanwhile, which makes ctx stale.
                try {
                    ctx.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error");
                }
                catch { }
            });
        });
        // The first prompt waits for servers whose tools are declared to the model, so they are declared
        // in its first request, but not indefinitely: a slow or hanging server must not hold up the
        // prompt. Other servers are waited for when a script or search needs them (below).
        const waitForDirectServers = async (ctx) => {
            if (waitedForStartup)
                return;
            waitedForStartup = true;
            const ready = servers.flatMap((server) => isEnabled(server) && hasDirectTools(server.entry) && server.ready ? [server.ready] : []);
            if (ready.length === 0)
                return;
            let timer;
            const finished = await Promise.race([
                Promise.all(ready).then(() => true),
                new Promise((resolve) => {
                    timer = setTimeout(() => resolve(false), startupWaitMs);
                }),
            ]);
            clearTimeout(timer);
            if (!finished) {
                ctx.ui.notify("MCP servers are still connecting; their tools become available once connected.", "info");
            }
        };
        // Every prompt lists the servers in the `mcp_servers` section as they are when it starts. Pi
        // appends the section to the conversation when it changed, for example after a server connected.
        pi.on("before_agent_start", async (event, ctx) => {
            await waitForDirectServers(ctx);
            const { sections } = event.systemPromptOptions;
            const section = renderServersSection(servers);
            if (section)
                sections[MCP_SERVERS_SECTION] = section;
            else
                delete sections[MCP_SERVERS_SECTION];
        });
        // A tool that reaches other tools (`reachesTools`) waits for the servers it may reach, so their
        // tools are registered before it runs: a codemode script for the servers it names, or for every
        // server when it searches or enumerates tools, and tool_search for every server. The 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;
            // 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 = [];
            const reach = tool.reachesTools?.(event.input);
            if (reach !== undefined) {
                waiting = readyServers.filter((server) => reachesServer(reach, server.entry.name));
            }
            else if (RESOURCE_TOOL_NAMES.has(tool.name)) {
                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);
        });
        // Pick up sign-ins done outside the session, such as `pi mcp login` run by the agent.
        pi.on("turn_start", async (_event, ctx) => {
            if (tokensAtSignIn.size > 0)
                await reconnectSignedIn(ctx);
        });
        // Servers registered or unregistered during the session connect or disconnect right away.
        pi.on("mcp_servers_change", async (_event, ctx) => {
            const { signal } = session;
            if (signal.aborted)
                return;
            const registered = registeredServers();
            overridden = registered.overridden;
            const next = new Map(registered.servers.map((server) => [server.entry.name, server]));
            // Unregistered servers and re-registered ones with a new config are dropped; the latter come back below.
            const removed = servers.filter((server) => server.entry.scope === "extension" &&
                next.get(server.entry.name)?.registeredConfig !== server.registeredConfig);
            servers = servers.filter((server) => !removed.includes(server));
            for (const server of removed)
                hideTools(server.entry.name);
            const added = registered.servers.filter((server) => !findServer(server.entry.name));
            servers.push(...added);
            emitChange();
            ensureDiscoveryActive(ctx);
            await Promise.all(removed.map((server) => server.connection?.close()));
            const connecting = added.filter(isEnabled);
            if (signal.aborted || connecting.length === 0)
                return;
            try {
                await Promise.all(connecting.map((server) => startConnection(server)));
                if (signal.aborted) {
                    await Promise.all(connecting.map((server) => server.connection?.close()));
                    return;
                }
            }
            catch (error) {
                ctx.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error");
                return;
            }
            if (signal.aborted)
                return;
            reportProblems(ctx, connecting);
        });
        pi.on("session_shutdown", async () => {
            session.abort();
            const closing = connections();
            servers = [];
            emitChange();
            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",
            getArgumentCompletions: (prefix) => {
                const [action, server, ...rest] = prefix.trimStart().split(/\s+/);
                if (rest.length > 0)
                    return null;
                if (server === undefined) {
                    return ["login", "logout", "reconnect"]
                        .filter((item) => item.startsWith(action ?? ""))
                        .map((item) => ({ value: `${item} `, label: item }));
                }
                if (action !== "login" && action !== "logout" && action !== "reconnect")
                    return null;
                const items = servers
                    .filter((candidate) => action === "reconnect" ? candidate.connection !== undefined : usesOAuth(candidate))
                    .filter((candidate) => candidate.entry.name.startsWith(server))
                    .map((candidate) => ({
                    value: `${action} ${candidate.entry.name}`,
                    label: candidate.entry.name,
                    description: describeState(candidate),
                }));
                return items.length > 0 ? items : null;
            },
            handler: async (args, ctx) => {
                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 {
                        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);
                        if (server)
                            await loginCommand(server, ctx);
                        return;
                    }
                    case "logout": {
                        const server = await pickServer(name, ctx, oauthPick);
                        if (!server)
                            return;
                        const removed = await signOut(server);
                        ctx.ui.notify(removed
                            ? `Signed out of MCP server "${server.entry.name}".`
                            : `No stored credentials for MCP server "${server.entry.name}".`, "info");
                        return;
                    }
                    case "reconnect": {
                        const server = await pickServer(name, ctx, {
                            eligible: (candidate) => candidate.connection !== undefined,
                            preferred: (candidate) => candidate.connection?.state === "failed" || candidate.connection?.state === "disconnected",
                            none: "No enabled MCP server to reconnect.",
                        });
                        if (!server)
                            return;
                        const failure = await reconnect(server);
                        if (failure)
                            ctx.ui.notify(failure, "error");
                        else {
                            ensureDiscoveryActive(ctx);
                            ctx.ui.notify(`Reconnected to MCP server "${server.entry.name}" (${describeState(server)}).`, "info");
                        }
                        return;
                    }
                    default:
                        ctx.ui.notify(MCP_USAGE, "warning");
                }
            },
        });
    };
}
function defaultLoadConfig(ctx) {
    return loadMcpConfig({ agentDir: getAgentDir(), cwd: ctx.cwd, projectTrusted: ctx.isProjectTrusted() });
}
export default createMcpExtension();
//# sourceMappingURL=index.js.map