@bastani/atomic 0.9.5-alpha.8 → 0.9.5

package/dist/builtin/mcp/init.ts

@@ -1,10 +1,10 @@
 import type { ExtensionAPI, ExtensionContext } from "@bastani/atomic";
-import type { McpExtensionState } from "./state.ts";
-import type { ToolMetadata } from "./types.ts";
+import type { McpExtensionState } from "./state.js";
+import type { ToolMetadata } from "./types.js";
 import { existsSync } from "node:fs";
-import { loadMcpConfig } from "./config.ts";
-import { ConsentManager } from "./consent-manager.ts";
-import { McpLifecycleManager } from "./lifecycle.ts";
+import { loadMcpConfig } from "./config.js";
+import { ConsentManager } from "./consent-manager.js";
+import { McpLifecycleManager } from "./lifecycle.js";
 import {
   computeServerHash,
   getMetadataCachePath,
@@ -15,13 +15,12 @@ import {
   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";
+} from "./metadata-cache.js";
+import { McpServerManager } from "./server-manager.js";
+import { buildToolMetadata, totalToolCount } from "./tool-metadata.js";
+import { UiResourceHandler } from "./ui-resource-handler.js";
+import { openUrl, parallelLimit } from "./utils.js";
+import { logger } from "./logger.js";
 
 const FAILURE_BACKOFF_MS = 60 * 1000;
 
@@ -33,96 +32,92 @@ export async function initializeMcp(
   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;
-  }
+  try {
+    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 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 idleSetting = typeof config.settings?.idleTimeout === "number" ? config.settings.idleTimeout : 10;
+    lifecycle.setGlobalIdleTimeout(idleSetting);
+
+    const cachePath = getMetadataCachePath();
+    const cacheFileExists = existsSync(cachePath);
+    let cache = loadMetadataCache();
+    if (!cacheFileExists) {
+      saveMetadataCache({ version: 1, servers: {} });
+    } else if (!cache) {
+      cache = { version: 1, servers: {} };
+      saveMetadataCache(cache);
+    }
 
-  const cachePath = getMetadataCachePath();
-  const cacheFileExists = existsSync(cachePath);
-  let cache = loadMetadataCache();
-  let bootstrapAll = false;
+    const prefix = config.settings?.toolPrefix ?? "server";
 
-  if (!cacheFileExists) {
-    bootstrapAll = true;
-    saveMetadataCache({ version: 1, servers: {} });
-  } else if (!cache) {
-    cache = { version: 1, servers: {} };
-    saveMetadataCache(cache);
-  }
+    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);
+      }
 
-  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);
+      }
     }
 
-    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";
-      });
+    const startupServers = 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...`);
-  }
+    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 };
+    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 };
@@ -160,39 +155,6 @@ export async function initializeMcp(
     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);
@@ -210,6 +172,14 @@ export async function initializeMcp(
   lifecycle.startHealthChecks();
 
   return state;
+  } catch (error) {
+    try {
+      await lifecycle.gracefulShutdown();
+    } catch (cleanupError) {
+      console.error("MCP: failed to clean resources after initialization failure", cleanupError);
+    }
+    throw error;
+  }
 }
 
 export function updateServerMetadata(state: McpExtensionState, serverName: string): void {