@bastani/atomic 0.7.17-1 → 0.8.0-0

package/dist/builtin/mcp/index.ts

@@ -0,0 +1,334 @@
+import type { ExtensionAPI, ToolInfo } from "@earendil-works/pi-coding-agent";
+import type { McpExtensionState } from "./state.ts";
+import { Type } from "typebox";
+import { showStatus, showTools, reconnectServers, authenticateServer, logoutServer, openMcpAuthPanel, openMcpPanel, openMcpSetup } from "./commands.ts";
+import { loadMcpConfig } from "./config.ts";
+import { buildProxyDescription, createDirectToolExecutor, getMissingConfiguredDirectToolServers, resolveDirectTools } from "./direct-tools.ts";
+import { flushMetadataCache, initializeMcp, updateStatusBar } from "./init.ts";
+import { loadMetadataCache } from "./metadata-cache.ts";
+import { executeCall, executeConnect, executeDescribe, executeList, executeSearch, executeStatus, executeUiMessages } from "./proxy-modes.ts";
+import { getConfigPathFromArgv, truncateAtWord } from "./utils.ts";
+import { initializeOAuth, shutdownOAuth } from "./mcp-auth-flow.ts";
+import { renderMcpToolResult } from "./tool-result-renderer.ts";
+
+export default function mcpAdapter(pi: ExtensionAPI) {
+  let state: McpExtensionState | null = null;
+  let initPromise: Promise<McpExtensionState> | null = null;
+  let lifecycleGeneration = 0;
+
+  async function shutdownState(currentState: McpExtensionState | null, reason: string): Promise<void> {
+    if (!currentState) return;
+
+    if (currentState.uiServer) {
+      currentState.uiServer.close(reason);
+      currentState.uiServer = null;
+    }
+
+    let flushError: unknown;
+    try {
+      flushMetadataCache(currentState);
+    } catch (error) {
+      flushError = error;
+    }
+
+    try {
+      await currentState.lifecycle.gracefulShutdown();
+    } catch (error) {
+      if (flushError) {
+        console.error("MCP: graceful shutdown failed after metadata flush error", error);
+      } else {
+        throw error;
+      }
+    }
+
+    if (flushError) {
+      throw flushError;
+    }
+  }
+
+  const earlyConfigPath = getConfigPathFromArgv();
+  const earlyConfig = loadMcpConfig(earlyConfigPath);
+  const earlyCache = loadMetadataCache();
+  const prefix = earlyConfig.settings?.toolPrefix ?? "server";
+
+  const envRaw = process.env.MCP_DIRECT_TOOLS;
+  const directSpecs = envRaw === "__none__"
+    ? []
+    : resolveDirectTools(
+        earlyConfig,
+        earlyCache,
+        prefix,
+        envRaw?.split(",").map(s => s.trim()).filter(Boolean),
+      );
+  const missingConfiguredDirectToolServers = getMissingConfiguredDirectToolServers(earlyConfig, earlyCache);
+  const shouldRegisterProxyTool =
+    earlyConfig.settings?.disableProxyTool !== true
+    || directSpecs.length === 0
+    || missingConfiguredDirectToolServers.length > 0;
+
+  for (const spec of directSpecs) {
+    (pi.registerTool as (tool: unknown) => unknown)({
+      name: spec.prefixedName,
+      label: `MCP: ${spec.originalName}`,
+      description: spec.description || "(no description)",
+      promptSnippet: truncateAtWord(spec.description, 100) || `MCP tool from ${spec.serverName}`,
+      parameters: Type.Unsafe((spec.inputSchema || { type: "object", properties: {} }) as never),
+      execute: createDirectToolExecutor(() => state, () => initPromise, spec),
+      renderResult: renderMcpToolResult,
+    });
+  }
+
+  const getPiTools = (): ToolInfo[] => pi.getAllTools();
+
+  pi.registerFlag("mcp-config", {
+    description: "Path to MCP config file",
+    type: "string",
+  });
+
+  pi.on("session_start", async (_event, ctx) => {
+    const generation = ++lifecycleGeneration;
+    const previousState = state;
+    state = null;
+    initPromise = null;
+
+    try {
+      await Promise.all([
+        shutdownState(previousState, "session_restart"),
+        shutdownOAuth(),
+      ]);
+    } catch (error) {
+      console.error("MCP: failed to shut down previous session state", error);
+    }
+
+    if (generation !== lifecycleGeneration) {
+      return;
+    }
+
+    await initializeOAuth().catch(() => {
+      console.error("MCP OAuth initialization failed");
+    });
+
+    const promise = initializeMcp(pi, ctx);
+    initPromise = promise;
+
+    promise.then(async (nextState) => {
+      if (generation !== lifecycleGeneration || initPromise !== promise) {
+        try {
+          await shutdownState(nextState, "stale_session_start");
+        } catch (error) {
+          console.error("MCP: failed to clean stale session state", error);
+        }
+        return;
+      }
+
+      state = nextState;
+      updateStatusBar(nextState);
+      initPromise = null;
+    }).catch(err => {
+      if (generation !== lifecycleGeneration) {
+        return;
+      }
+      if (initPromise !== promise && initPromise !== null) {
+        return;
+      }
+      console.error("MCP initialization failed:", err);
+      initPromise = null;
+    });
+  });
+
+  pi.on("session_shutdown", async () => {
+    ++lifecycleGeneration;
+    const currentState = state;
+    state = null;
+    initPromise = null;
+
+    try {
+      await Promise.all([
+        shutdownState(currentState, "session_shutdown"),
+        shutdownOAuth(),
+      ]);
+    } catch (error) {
+      console.error("MCP: session shutdown cleanup failed", error);
+    }
+  });
+
+  pi.registerCommand("mcp", {
+    description: "Show MCP server status",
+    handler: async (args, ctx) => {
+      if (!state && initPromise) {
+        try {
+          state = await initPromise;
+        } catch (error) {
+          const message = error instanceof Error ? error.message : String(error);
+          if (ctx.hasUI) ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
+          return;
+        }
+      }
+      if (!state) {
+        if (ctx.hasUI) ctx.ui.notify("MCP not initialized", "error");
+        return;
+      }
+
+      const parts = args?.trim()?.split(/\s+/) ?? [];
+      const subcommand = parts[0] ?? "";
+      const targetServer = parts[1];
+      const rest = parts.slice(1).join(" ");
+
+      switch (subcommand) {
+        case "reconnect":
+          await reconnectServers(state, ctx, targetServer);
+          break;
+        case "tools":
+          await showTools(state, ctx);
+          break;
+        case "setup": {
+          const result = await openMcpSetup(state, pi, ctx, earlyConfigPath, "setup");
+          if (result?.configChanged) {
+            await ctx.reload();
+            return;
+          }
+          break;
+        }
+        case "logout": {
+          const serverName = rest;
+          if (!serverName) {
+            if (ctx.hasUI) ctx.ui.notify("Usage: /mcp logout <server>", "error");
+            return;
+          }
+          await logoutServer(serverName, state, ctx);
+          break;
+        }
+        case "status":
+        case "":
+        default:
+          if (ctx.hasUI) {
+            const result = await openMcpPanel(state, pi, ctx, earlyConfigPath);
+            if (result?.configChanged) {
+              await ctx.reload();
+              return;
+            }
+          } else {
+            await showStatus(state, ctx);
+          }
+          break;
+      }
+    },
+  });
+
+  pi.registerCommand("mcp-auth", {
+    description: "Authenticate with an MCP server (OAuth)",
+    handler: async (args, ctx) => {
+      const serverName = args?.trim();
+      if (!serverName && !ctx.hasUI) {
+        return;
+      }
+
+      if (!state && initPromise) {
+        try {
+          state = await initPromise;
+        } catch (error) {
+          const message = error instanceof Error ? error.message : String(error);
+          if (ctx.hasUI) ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
+          return;
+        }
+      }
+      if (!state) {
+        if (ctx.hasUI) ctx.ui.notify("MCP not initialized", "error");
+        return;
+      }
+
+      if (!serverName) {
+        await openMcpAuthPanel(state, pi, ctx, earlyConfigPath);
+        return;
+      }
+
+      await authenticateServer(serverName, state.config, ctx);
+    },
+  });
+
+  if (shouldRegisterProxyTool) {
+    (pi.registerTool as (tool: unknown) => unknown)({
+      name: "mcp",
+      label: "MCP",
+      description: buildProxyDescription(earlyConfig, earlyCache, directSpecs),
+      promptSnippet: "MCP gateway - connect to MCP servers and call their tools",
+      parameters: Type.Object({
+        tool: Type.Optional(Type.String({ description: "Tool name to call (e.g., 'xcodebuild_list_sims')" })),
+        args: Type.Optional(Type.String({ description: "Arguments as JSON string (e.g., '{\"key\": \"value\"}')" })),
+        connect: Type.Optional(Type.String({ description: "Server name to connect (lazy connect + metadata refresh)" })),
+        describe: Type.Optional(Type.String({ description: "Tool name to describe (shows parameters)" })),
+        search: Type.Optional(Type.String({ description: "Search tools by name/description" })),
+        regex: Type.Optional(Type.Boolean({ description: "Treat search as regex (default: substring match)" })),
+        includeSchemas: Type.Optional(Type.Boolean({ description: "Include parameter schemas in search results (default: true)" })),
+        server: Type.Optional(Type.String({ description: "Filter to specific server (also disambiguates tool calls)" })),
+        action: Type.Optional(Type.String({ description: "Action: 'ui-messages' to retrieve prompts/intents from UI sessions" })),
+      }),
+      renderResult: renderMcpToolResult,
+      async execute(_toolCallId, params: {
+        tool?: string;
+        args?: string;
+        connect?: string;
+        describe?: string;
+        search?: string;
+        regex?: boolean;
+        includeSchemas?: boolean;
+        server?: string;
+        action?: string;
+      }, _signal, _onUpdate, _ctx) {
+        let parsedArgs: Record<string, unknown> | undefined;
+        if (params.args) {
+          try {
+            parsedArgs = JSON.parse(params.args);
+            if (typeof parsedArgs !== "object" || parsedArgs === null || Array.isArray(parsedArgs)) {
+              const gotType = Array.isArray(parsedArgs) ? "array" : parsedArgs === null ? "null" : typeof parsedArgs;
+              throw new Error(`Invalid args: expected a JSON object, got ${gotType}`);
+            }
+          } catch (error) {
+            if (error instanceof SyntaxError) {
+              throw new Error(`Invalid args JSON: ${error.message}`, { cause: error });
+            }
+            throw error;
+          }
+        }
+
+        if (!state && initPromise) {
+          try {
+            state = await initPromise;
+          } catch (error) {
+            const message = error instanceof Error ? error.message : String(error);
+            return {
+              content: [{ type: "text" as const, text: `MCP initialization failed: ${message}` }],
+              details: { error: "init_failed", message },
+            };
+          }
+        }
+        if (!state) {
+          return {
+            content: [{ type: "text" as const, text: "MCP not initialized" }],
+            details: { error: "not_initialized" },
+          };
+        }
+
+        if (params.action === "ui-messages") {
+          return executeUiMessages(state);
+        }
+        if (params.tool) {
+          return executeCall(state, params.tool, parsedArgs, params.server, getPiTools);
+        }
+        if (params.connect) {
+          return executeConnect(state, params.connect);
+        }
+        if (params.describe) {
+          return executeDescribe(state, params.describe);
+        }
+        if (params.search) {
+          return executeSearch(state, params.search, params.regex, params.server, params.includeSchemas);
+        }
+        if (params.server) {
+          return executeList(state, params.server);
+        }
+        return executeStatus(state);
+      },
+    });
+  }
+}

package/dist/builtin/mcp/init.ts

@@ -0,0 +1,336 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { McpExtensionState } from "./state.ts";
+import type { ToolMetadata } from "./types.ts";
+import { existsSync } from "node:fs";
+import { loadMcpConfig } from "./config.ts";
+import { ConsentManager } from "./consent-manager.ts";
+import { McpLifecycleManager } from "./lifecycle.ts";
+import {
+  computeServerHash,
+  getMetadataCachePath,
+  isServerCacheValid,
+  loadMetadataCache,
+  reconstructToolMetadata,
+  saveMetadataCache,
+  serializeResources,
+  serializeTools,
+  type ServerCacheEntry,
+} from "./metadata-cache.ts";
+import { McpServerManager } from "./server-manager.ts";
+import { buildToolMetadata, totalToolCount } from "./tool-metadata.ts";
+import { UiResourceHandler } from "./ui-resource-handler.ts";
+import { openUrl, parallelLimit } from "./utils.ts";
+import { logger } from "./logger.ts";
+import { getMissingConfiguredDirectToolServers } from "./direct-tools.ts";
+
+const FAILURE_BACKOFF_MS = 60 * 1000;
+
+export async function initializeMcp(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext
+): Promise<McpExtensionState> {
+  const configPath = pi.getFlag("mcp-config") as string | undefined;
+  const config = loadMcpConfig(configPath, ctx.cwd);
+
+  const manager = new McpServerManager();
+  const samplingAutoApprove = config.settings?.samplingAutoApprove === true;
+  if (config.settings?.sampling !== false && (ctx.hasUI || samplingAutoApprove)) {
+    manager.setSamplingConfig({
+      autoApprove: samplingAutoApprove,
+      ui: ctx.hasUI ? ctx.ui : undefined,
+      modelRegistry: ctx.modelRegistry,
+      getCurrentModel: () => ctx.model,
+      getSignal: () => ctx.signal,
+    });
+  }
+  const lifecycle = new McpLifecycleManager(manager);
+  const toolMetadata = new Map<string, ToolMetadata[]>();
+  const failureTracker = new Map<string, number>();
+  const uiResourceHandler = new UiResourceHandler(manager);
+  const consentManager = new ConsentManager("once-per-server");
+  const ui = ctx.hasUI ? ctx.ui : undefined;
+  const state: McpExtensionState = {
+    manager,
+    lifecycle,
+    toolMetadata,
+    config,
+    failureTracker,
+    uiResourceHandler,
+    consentManager,
+    uiServer: null,
+    completedUiSessions: [],
+    openBrowser: (url: string) => openUrl(pi, url, process.env.BROWSER),
+    ui,
+    sendMessage: (message, options) => pi.sendMessage(message as unknown as Parameters<typeof pi.sendMessage>[0], options),
+  };
+
+  const serverEntries = Object.entries(config.mcpServers);
+  if (serverEntries.length === 0) {
+    return state;
+  }
+
+  const idleSetting = typeof config.settings?.idleTimeout === "number" ? config.settings.idleTimeout : 10;
+  lifecycle.setGlobalIdleTimeout(idleSetting);
+
+  const cachePath = getMetadataCachePath();
+  const cacheFileExists = existsSync(cachePath);
+  let cache = loadMetadataCache();
+  let bootstrapAll = false;
+
+  if (!cacheFileExists) {
+    bootstrapAll = true;
+    saveMetadataCache({ version: 1, servers: {} });
+  } else if (!cache) {
+    cache = { version: 1, servers: {} };
+    saveMetadataCache(cache);
+  }
+
+  const prefix = config.settings?.toolPrefix ?? "server";
+
+  for (const [name, definition] of serverEntries) {
+    const lifecycleMode = definition.lifecycle ?? "lazy";
+    const idleOverride = definition.idleTimeout ?? (lifecycleMode === "eager" ? 0 : undefined);
+    lifecycle.registerServer(
+      name,
+      definition,
+      idleOverride !== undefined ? { idleTimeout: idleOverride } : undefined
+    );
+    if (lifecycleMode === "keep-alive") {
+      lifecycle.markKeepAlive(name, definition);
+    }
+
+    if (cache?.servers?.[name] && isServerCacheValid(cache.servers[name], definition)) {
+      const metadata = reconstructToolMetadata(name, cache.servers[name], prefix, definition);
+      toolMetadata.set(name, metadata);
+    }
+  }
+
+  const startupServers = bootstrapAll
+    ? serverEntries
+    : serverEntries.filter(([, definition]) => {
+        const mode = definition.lifecycle ?? "lazy";
+        return mode === "keep-alive" || mode === "eager";
+      });
+
+  if (ctx.hasUI && startupServers.length > 0) {
+    ctx.ui.setStatus("mcp", `MCP: connecting to ${startupServers.length} servers...`);
+  }
+
+  const results = await parallelLimit(startupServers, 10, async ([name, definition]) => {
+    try {
+      const connection = await manager.connect(name, definition);
+      if (connection.status === "needs-auth") {
+        return { name, definition, connection: null, error: `OAuth authentication required. Run /mcp-auth ${name}.` };
+      }
+      return { name, definition, connection, error: null };
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      return { name, definition, connection: null, error: message };
+    }
+  });
+
+  for (const { name, definition, connection, error } of results) {
+    if (error || !connection) {
+      if (ctx.hasUI) {
+        ctx.ui.notify(`MCP: Failed to connect to ${name}: ${error}`, "error");
+      }
+      console.error(`MCP: Failed to connect to ${name}: ${error}`);
+      continue;
+    }
+
+    const { metadata, failedTools } = buildToolMetadata(connection.tools, connection.resources, definition, name, prefix);
+    toolMetadata.set(name, metadata);
+    updateMetadataCache(state, name);
+
+    if (failedTools.length > 0 && ctx.hasUI) {
+      ctx.ui.notify(
+        `MCP: ${name} - ${failedTools.length} tools skipped`,
+        "warning"
+      );
+    }
+  }
+
+  const connectedCount = results.filter(r => r.connection).length;
+  const failedCount = results.filter(r => r.error).length;
+  if (ctx.hasUI && connectedCount > 0) {
+    const totalTools = totalToolCount(state);
+    const msg = failedCount > 0
+      ? `MCP: ${connectedCount}/${startupServers.length} servers connected (${totalTools} tools)`
+      : `MCP: ${connectedCount} servers connected (${totalTools} tools)`;
+    ctx.ui.notify(msg, "info");
+  }
+
+  const envDirect = process.env.MCP_DIRECT_TOOLS;
+  if (envDirect !== "__none__") {
+    const currentCache = loadMetadataCache();
+    const missingCacheServers = getMissingConfiguredDirectToolServers(config, currentCache);
+
+    if (missingCacheServers.length > 0) {
+      const bootstrapResults = await parallelLimit(
+        missingCacheServers.filter(name => !results.some(r => r.name === name && r.connection)),
+        10,
+        async (name) => {
+          const definition = config.mcpServers[name];
+          try {
+            const connection = await manager.connect(name, definition);
+            if (connection.status === "needs-auth") {
+              return { name, ok: false };
+            }
+            const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, name, prefix);
+            toolMetadata.set(name, metadata);
+            updateMetadataCache(state, name);
+            return { name, ok: true };
+          } catch (error) {
+            const message = error instanceof Error ? error.message : String(error);
+            logger.debug(`MCP: direct-tools bootstrap failed for ${name}: ${message}`);
+            return { name, ok: false };
+          }
+        },
+      );
+      const bootstrapped = bootstrapResults.filter(r => r.ok).map(r => r.name);
+      if (bootstrapped.length > 0 && ctx.hasUI) {
+        ctx.ui.notify(`MCP: direct tools for ${bootstrapped.join(", ")} will be available after restart`, "info");
+      }
+    }
+  }
+
+  lifecycle.setReconnectCallback((serverName) => {
+    updateServerMetadata(state, serverName);
+    updateMetadataCache(state, serverName);
+    state.failureTracker.delete(serverName);
+    updateStatusBar(state);
+  });
+
+  lifecycle.setIdleShutdownCallback((serverName) => {
+    const idleMinutes = getEffectiveIdleTimeoutMinutes(state, serverName);
+    logger.debug(`${serverName} shut down (idle ${idleMinutes}m)`);
+    updateStatusBar(state);
+  });
+
+  lifecycle.startHealthChecks();
+
+  return state;
+}
+
+export function updateServerMetadata(state: McpExtensionState, serverName: string): void {
+  const connection = state.manager.getConnection(serverName);
+  if (!connection || connection.status !== "connected") return;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return;
+
+  const prefix = state.config.settings?.toolPrefix ?? "server";
+
+  const { metadata } = buildToolMetadata(connection.tools, connection.resources, definition, serverName, prefix);
+  state.toolMetadata.set(serverName, metadata);
+}
+
+export function updateMetadataCache(state: McpExtensionState, serverName: string): void {
+  const connection = state.manager.getConnection(serverName);
+  if (!connection || connection.status !== "connected") return;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return;
+
+  const configHash = computeServerHash(definition);
+  const existing = loadMetadataCache();
+  const existingEntry = existing?.servers?.[serverName];
+
+  const tools = serializeTools(connection.tools);
+  let resources = definition.exposeResources === false ? [] : serializeResources(connection.resources);
+
+  if (
+    definition.exposeResources !== false &&
+    resources.length === 0 &&
+    existingEntry?.resources?.length &&
+    existingEntry.configHash === configHash
+  ) {
+    resources = existingEntry.resources;
+  }
+
+  const entry: ServerCacheEntry = {
+    configHash,
+    tools,
+    resources,
+    cachedAt: Date.now(),
+  };
+
+  saveMetadataCache({ version: 1, servers: { [serverName]: entry } });
+}
+
+export function flushMetadataCache(state: McpExtensionState): void {
+  for (const [name, connection] of state.manager.getAllConnections()) {
+    if (connection.status === "connected") {
+      updateMetadataCache(state, name);
+    }
+  }
+}
+
+export function updateStatusBar(state: McpExtensionState): void {
+  const ui = state.ui;
+  if (!ui) return;
+  const total = Object.keys(state.config.mcpServers).length;
+  if (total === 0) {
+    ui.setStatus("mcp", undefined);
+    return;
+  }
+  const connectedCount = state.manager.getAllConnections().size;
+  ui.setStatus("mcp", ui.theme.fg("accent", `MCP: ${connectedCount}/${total} servers`));
+}
+
+export function getFailureAgeSeconds(state: McpExtensionState, serverName: string): number | null {
+  const failedAt = state.failureTracker.get(serverName);
+  if (!failedAt) return null;
+  const ageMs = Date.now() - failedAt;
+  if (ageMs > FAILURE_BACKOFF_MS) return null;
+  return Math.round(ageMs / 1000);
+}
+
+export async function lazyConnect(state: McpExtensionState, serverName: string): Promise<boolean> {
+  const connection = state.manager.getConnection(serverName);
+  if (connection?.status === "needs-auth") {
+    return false;
+  }
+  if (connection?.status === "connected") {
+    updateServerMetadata(state, serverName);
+    return true;
+  }
+
+  const failedAgo = getFailureAgeSeconds(state, serverName);
+  if (failedAgo !== null) return false;
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) return false;
+
+  try {
+    if (state.ui) {
+      state.ui.setStatus("mcp", `MCP: connecting to ${serverName}...`);
+    }
+    const newConnection = await state.manager.connect(serverName, definition);
+    if (newConnection.status === "needs-auth") {
+      return false;
+    }
+    state.failureTracker.delete(serverName);
+    updateServerMetadata(state, serverName);
+    updateMetadataCache(state, serverName);
+    updateStatusBar(state);
+    return true;
+  } catch (error) {
+    state.failureTracker.set(serverName, Date.now());
+    const message = error instanceof Error ? error.message : String(error);
+    logger.debug(`MCP: lazy connect failed for ${serverName}: ${message}`);
+    updateStatusBar(state);
+    return false;
+  }
+}
+
+function getEffectiveIdleTimeoutMinutes(state: McpExtensionState, serverName: string): number {
+  const definition = state.config.mcpServers[serverName];
+  if (!definition) {
+    return typeof state.config.settings?.idleTimeout === "number" ? state.config.settings.idleTimeout : 10;
+  }
+  if (typeof definition.idleTimeout === "number") return definition.idleTimeout;
+  const mode = definition.lifecycle ?? "lazy";
+  if (mode === "eager") return 0;
+  return typeof state.config.settings?.idleTimeout === "number" ? state.config.settings.idleTimeout : 10;
+}