pi-mcp-adapter 2.3.5 → 2.4.1

package/direct-tools.ts

@@ -7,11 +7,50 @@ import { isServerCacheValid } from "./metadata-cache.js";
 import { formatSchema } from "./tool-metadata.js";
 import { transformMcpContent } from "./tool-registrar.js";
 import { maybeStartUiSession, type UiSessionRuntime } from "./ui-session.js";
-import { formatToolName } from "./types.js";
+import { formatToolName, isToolExcluded } from "./types.js";
 import { resourceNameToToolName } from "./resource-tools.js";
+import { authenticate, supportsOAuth } from "./mcp-auth-flow.js";
 
 const BUILTIN_NAMES = new Set(["read", "bash", "edit", "write", "grep", "find", "ls", "mcp"]);
 
+type DirectAutoAuthResult =
+  | { status: "skipped" }
+  | { status: "success" }
+  | { status: "failed"; message: string };
+
+async function attemptDirectAutoAuth(
+  state: McpExtensionState,
+  serverName: string,
+): Promise<DirectAutoAuthResult> {
+  if (state.config.settings?.autoAuth !== true) {
+    return { status: "skipped" };
+  }
+
+  const definition = state.config.mcpServers[serverName];
+  if (!definition || !supportsOAuth(definition) || !definition.url) {
+    return { status: "skipped" };
+  }
+
+  const grantType = definition.oauth?.grantType ?? "authorization_code";
+  if (!state.ui && grantType !== "client_credentials") {
+    return {
+      status: "failed",
+      message: `MCP server "${serverName}" requires OAuth authentication. Run /mcp-auth ${serverName} in an interactive session.`,
+    };
+  }
+
+  try {
+    await authenticate(serverName, definition.url, definition);
+    return { status: "success" };
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    return {
+      status: "failed",
+      message: `OAuth authentication failed for "${serverName}": ${message}. Run /mcp-auth ${serverName} first.`,
+    };
+  }
+}
+
 export function resolveDirectTools(
   config: McpConfig,
   cache: MetadataCache | null,
@@ -68,6 +107,7 @@ export function resolveDirectTools(
 
     for (const tool of serverCache.tools ?? []) {
       if (toolFilter !== true && !toolFilter.includes(tool.name)) continue;
+      if (isToolExcluded(tool.name, serverName, prefix, definition.excludeTools)) continue;
       const prefixedName = formatToolName(tool.name, serverName, prefix);
       if (BUILTIN_NAMES.has(prefixedName)) {
         console.warn(`MCP: skipping direct tool "${prefixedName}" (collides with builtin)`);
@@ -93,6 +133,7 @@ export function resolveDirectTools(
       for (const resource of serverCache.resources ?? []) {
         const baseName = `get_${resourceNameToToolName(resource.name)}`;
         if (toolFilter !== true && !toolFilter.includes(baseName)) continue;
+        if (isToolExcluded(baseName, serverName, prefix, definition.excludeTools)) continue;
         const prefixedName = formatToolName(baseName, serverName, prefix);
         if (BUILTIN_NAMES.has(prefixedName)) {
           console.warn(`MCP: skipping direct resource tool "${prefixedName}" (collides with builtin)`);
@@ -117,11 +158,35 @@ export function resolveDirectTools(
   return specs;
 }
 
+export function getMissingConfiguredDirectToolServers(
+  config: McpConfig,
+  cache: MetadataCache | null,
+): string[] {
+  const missing: string[] = [];
+  const globalDirect = config.settings?.directTools;
+
+  for (const [serverName, definition] of Object.entries(config.mcpServers)) {
+    const hasDirectTools = definition.directTools !== undefined
+      ? !!definition.directTools
+      : !!globalDirect;
+
+    if (!hasDirectTools) continue;
+
+    const serverCache = cache?.servers?.[serverName];
+    if (!serverCache || !isServerCacheValid(serverCache, definition)) {
+      missing.push(serverName);
+    }
+  }
+
+  return missing;
+}
+
 export function buildProxyDescription(
   config: McpConfig,
   cache: MetadataCache | null,
   directSpecs: DirectToolSpec[],
 ): string {
+  const prefix = config.settings?.toolPrefix ?? "server";
   let desc = `MCP gateway - connect to MCP servers and call their tools.\n`;
 
   const directByServer = new Map<string, number>();
@@ -139,8 +204,15 @@ export function buildProxyDescription(
   for (const serverName of Object.keys(config.mcpServers)) {
     const entry = cache?.servers?.[serverName];
     const definition = config.mcpServers[serverName];
-    const toolCount = entry?.tools?.length ?? 0;
-    const resourceCount = definition?.exposeResources !== false ? (entry?.resources?.length ?? 0) : 0;
+    const toolCount = (entry?.tools ?? []).filter(
+      (tool) => !isToolExcluded(tool.name, serverName, prefix, definition.excludeTools),
+    ).length;
+    const resourceCount = definition?.exposeResources !== false
+      ? (entry?.resources ?? []).filter((resource) => {
+          const baseName = `get_${resourceNameToToolName(resource.name)}`;
+          return !isToolExcluded(baseName, serverName, prefix, definition.excludeTools);
+        }).length
+      : 0;
     const totalItems = toolCount + resourceCount;
     if (totalItems === 0) continue;
     const directCount = directByServer.get(serverName) ?? 0;
@@ -196,13 +268,32 @@ export function createDirectToolExecutor(
       };
     }
 
-    const connected = await lazyConnect(state, spec.serverName);
+    let connected = await lazyConnect(state, spec.serverName);
+    let autoAuthAttempted = false;
+
+    if (!connected && state.manager.getConnection(spec.serverName)?.status === "needs-auth") {
+      autoAuthAttempted = true;
+      const autoAuth = await attemptDirectAutoAuth(state, spec.serverName);
+      if (autoAuth.status === "failed") {
+        return {
+          content: [{ type: "text" as const, text: autoAuth.message }],
+          details: { error: "auth_required", server: spec.serverName, message: autoAuth.message },
+        };
+      }
+      if (autoAuth.status === "success") {
+        await state.manager.close(spec.serverName);
+        state.failureTracker.delete(spec.serverName);
+        connected = await lazyConnect(state, spec.serverName);
+      }
+    }
+
     if (!connected) {
       const authConnection = state.manager.getConnection(spec.serverName);
       if (authConnection?.status === "needs-auth") {
+        const message = `MCP server "${spec.serverName}" requires OAuth authentication. Run /mcp-auth ${spec.serverName} first.`;
         return {
-          content: [{ type: "text" as const, text: `MCP server "${spec.serverName}" requires OAuth authentication. Run /mcp-auth ${spec.serverName} first.` }],
-          details: { error: "auth_required", server: spec.serverName },
+          content: [{ type: "text" as const, text: message }],
+          details: { error: "auth_required", server: spec.serverName, message, autoAuthAttempted },
         };
       }
       const failedAgo = getFailureAgeSeconds(state, spec.serverName);

package/index.ts

@@ -1,9 +1,9 @@
 import type { ExtensionAPI, ToolInfo } from "@mariozechner/pi-coding-agent";
 import type { McpExtensionState } from "./state.js";
 import { Type } from "@sinclair/typebox";
-import { showStatus, showTools, reconnectServers, authenticateServer, openMcpPanel } from "./commands.js";
+import { showStatus, showTools, reconnectServers, authenticateServer, openMcpPanel, openMcpSetup } from "./commands.js";
 import { loadMcpConfig } from "./config.js";
-import { buildProxyDescription, createDirectToolExecutor, resolveDirectTools } from "./direct-tools.js";
+import { buildProxyDescription, createDirectToolExecutor, getMissingConfiguredDirectToolServers, resolveDirectTools } from "./direct-tools.js";
 import { flushMetadataCache, initializeMcp, updateStatusBar } from "./init.js";
 import { loadMetadataCache } from "./metadata-cache.js";
 import { executeCall, executeConnect, executeDescribe, executeList, executeSearch, executeStatus, executeUiMessages } from "./proxy-modes.js";
@@ -59,6 +59,11 @@ export default function mcpAdapter(pi: ExtensionAPI) {
         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({
@@ -173,11 +178,23 @@ export default function mcpAdapter(pi: ExtensionAPI) {
         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 "status":
         case "":
         default:
           if (ctx.hasUI) {
-            await openMcpPanel(state, pi, ctx, earlyConfigPath);
+            const result = await openMcpPanel(state, pi, ctx, earlyConfigPath);
+            if (result?.configChanged) {
+              await ctx.reload();
+              return;
+            }
           } else {
             await showStatus(state, ctx);
           }
@@ -213,86 +230,88 @@ export default function mcpAdapter(pi: ExtensionAPI) {
     },
   });
 
-  pi.registerTool({
-    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" })),
-    }),
-    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 });
+  if (shouldRegisterProxyTool) {
+    pi.registerTool({
+      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" })),
+      }),
+      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;
           }
-          throw error;
         }
-      }
 
-      if (!state && initPromise) {
-        try {
-          state = await initPromise;
-        } catch (error) {
-          const message = error instanceof Error ? error.message : String(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 initialization failed: ${message}` }],
-            details: { error: "init_failed", message },
+            content: [{ type: "text" as const, text: "MCP not initialized" }],
+            details: { error: "not_initialized" },
           };
         }
-      }
-      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);
-      }
-      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, getPiTools);
-      }
-      if (params.server) {
-        return executeList(state, params.server);
-      }
-      return executeStatus(state);
-    },
-  });
+        if (params.action === "ui-messages") {
+          return executeUiMessages(state);
+        }
+        if (params.tool) {
+          return executeCall(state, params.tool, parsedArgs, params.server);
+        }
+        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, getPiTools);
+        }
+        if (params.server) {
+          return executeList(state, params.server);
+        }
+        return executeStatus(state);
+      },
+    });
+  }
 }

package/init.ts

@@ -21,6 +21,7 @@ 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";
+import { getMissingConfiguredDirectToolServers } from "./direct-tools.js";
 
 const FAILURE_BACKOFF_MS = 60 * 1000;
 
@@ -89,7 +90,7 @@ export async function initializeMcp(
     }
 
     if (cache?.servers?.[name] && isServerCacheValid(cache.servers[name], definition)) {
-      const metadata = reconstructToolMetadata(name, cache.servers[name], prefix, definition.exposeResources);
+      const metadata = reconstructToolMetadata(name, cache.servers[name], prefix, definition);
       toolMetadata.set(name, metadata);
     }
   }
@@ -151,18 +152,8 @@ export async function initializeMcp(
 
   const envDirect = process.env.MCP_DIRECT_TOOLS;
   if (envDirect !== "__none__") {
-    const missingCacheServers: string[] = [];
     const currentCache = loadMetadataCache();
-    for (const [name, definition] of serverEntries) {
-      const hasDirect = definition.directTools !== undefined
-        ? !!definition.directTools
-        : !!config.settings?.directTools;
-      if (!hasDirect) continue;
-      const entry = currentCache?.servers?.[name];
-      if (!entry || !isServerCacheValid(entry, definition)) {
-        missingCacheServers.push(name);
-      }
-    }
+    const missingCacheServers = getMissingConfiguredDirectToolServers(config, currentCache);
 
     if (missingCacheServers.length > 0) {
       const bootstrapResults = await parallelLimit(

package/mcp-auth-flow.ts

@@ -36,6 +36,9 @@ export type AuthStatus = "authenticated" | "expired" | "not_authenticated"
 // Track pending transports for auth completion
 const pendingTransports = new Map<string, StreamableHTTPClientTransport>()
 
+// Deduplicate concurrent authenticate() calls per server.
+const pendingAuthentications = new Map<string, Promise<AuthStatus>>()
+
 /**
  * Generate a cryptographically secure random state parameter.
  */
@@ -182,56 +185,74 @@ export async function authenticate(
   serverUrl: string,
   definition?: ServerEntry,
 ): Promise<AuthStatus> {
-  // Start auth flow
-  const { authorizationUrl } = await startAuth(serverName, serverUrl, definition)
-
-  // If no auth URL needed, already authenticated
-  if (!authorizationUrl) {
-    return "authenticated"
+  const inFlight = pendingAuthentications.get(serverName)
+  if (inFlight) {
+    return inFlight
   }
 
-  // Get the state that was already generated and stored in startAuth()
-  const oauthState = await getOAuthState(serverName)
-  if (!oauthState) {
-    throw new Error("OAuth state not found - this should not happen")
-  }
+  const operation = (async (): Promise<AuthStatus> => {
+    // Start auth flow
+    const { authorizationUrl } = await startAuth(serverName, serverUrl, definition)
 
-  // Register the callback BEFORE opening the browser
-  const callbackPromise = waitForCallback(oauthState)
+    // If no auth URL needed, already authenticated
+    if (!authorizationUrl) {
+      return "authenticated"
+    }
 
-  // Open browser
-  console.log(`MCP Auth: Opening browser for ${serverName}`)
-  try {
-    await open(authorizationUrl)
-  } catch (error) {
-    console.warn(`MCP Auth: Failed to open browser for ${serverName}`, { error })
-    throw new Error(
-      `Could not open browser. Please open this URL manually: ${authorizationUrl}`
-    )
-  }
+    // Get the state that was already generated and stored in startAuth()
+    const oauthState = await getOAuthState(serverName)
+    if (!oauthState) {
+      throw new Error("OAuth state not found - this should not happen")
+    }
 
-  try {
-    // Wait for callback
-    const code = await callbackPromise
+    // Register the callback BEFORE opening the browser
+    const callbackPromise = waitForCallback(oauthState)
+
+    // Open browser
+    console.log(`MCP Auth: Opening browser for ${serverName}`)
+    try {
+      await open(authorizationUrl)
+    } catch (error) {
+      console.warn(`MCP Auth: Failed to open browser for ${serverName}`, { error })
+      throw new Error(
+        `Could not open browser. Please open this URL manually: ${authorizationUrl}`,
+        { cause: error },
+      )
+    }
 
-    // Validate state
-    const storedState = await getOAuthState(serverName)
-    if (storedState !== oauthState) {
+    try {
+      // Wait for callback
+      const code = await callbackPromise
+
+      // Validate state
+      const storedState = await getOAuthState(serverName)
+      if (storedState !== oauthState) {
+        await clearOAuthState(serverName)
+        throw new Error("OAuth state mismatch - potential CSRF attack")
+      }
       await clearOAuthState(serverName)
-      throw new Error("OAuth state mismatch - potential CSRF attack")
+
+      // Complete the auth
+      return await completeAuth(serverName, code)
+    } catch (error) {
+      cancelPendingCallback(oauthState)
+      const pendingTransport = pendingTransports.get(serverName)
+      if (pendingTransport) {
+        pendingTransports.delete(serverName)
+        await pendingTransport.close().catch(() => {})
+      }
+      throw error
     }
-    await clearOAuthState(serverName)
+  })()
 
-    // Complete the auth
-    return await completeAuth(serverName, code)
-  } catch (error) {
-    cancelPendingCallback(oauthState)
-    const pendingTransport = pendingTransports.get(serverName)
-    if (pendingTransport) {
-      pendingTransports.delete(serverName)
-      await pendingTransport.close().catch(() => {})
+  pendingAuthentications.set(serverName, operation)
+
+  try {
+    return await operation
+  } finally {
+    if (pendingAuthentications.get(serverName) === operation) {
+      pendingAuthentications.delete(serverName)
     }
-    throw error
   }
 }