@machina.ai/cell-cli-core 1.0.13-rc9 → 1.0.21-rc1

package/dist/src/tools/mcp-client.js

@@ -7,11 +7,18 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import { SSEClientTransport, } from '@modelcontextprotocol/sdk/client/sse.js';
 import { StreamableHTTPClientTransport, } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { ListPromptsResultSchema, GetPromptResultSchema, ListRootsRequestSchema, } from '@modelcontextprotocol/sdk/types.js';
 import { parse } from 'shell-quote';
+import { AuthProviderType } from '../config/config.js';
+import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
 import { mcpToTool } from '@google/genai';
-import { OpenFilesNotificationSchema, IDE_SERVER_NAME, ideContext, } from '../services/ideContext.js';
+import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
+import { OAuthUtils } from '../mcp/oauth-utils.js';
+import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { getErrorMessage } from '../utils/errors.js';
+import { basename } from 'node:path';
+import { pathToFileURL } from 'node:url';
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 /**
  * Enum representing the connection status of an MCP server
@@ -37,14 +44,134 @@ export var MCPDiscoveryState;
     /** Discovery has completed (with or without errors) */
     MCPDiscoveryState["COMPLETED"] = "completed";
 })(MCPDiscoveryState || (MCPDiscoveryState = {}));
+/**
+ * A client for a single MCP server.
+ *
+ * This class is responsible for connecting to, discovering tools from, and
+ * managing the state of a single MCP server.
+ */
+export class McpClient {
+    serverName;
+    serverConfig;
+    toolRegistry;
+    promptRegistry;
+    workspaceContext;
+    debugMode;
+    client;
+    transport;
+    status = MCPServerStatus.DISCONNECTED;
+    isDisconnecting = false;
+    constructor(serverName, serverConfig, toolRegistry, promptRegistry, workspaceContext, debugMode) {
+        this.serverName = serverName;
+        this.serverConfig = serverConfig;
+        this.toolRegistry = toolRegistry;
+        this.promptRegistry = promptRegistry;
+        this.workspaceContext = workspaceContext;
+        this.debugMode = debugMode;
+        this.client = new Client({
+            name: `gemini-cli-mcp-client-${this.serverName}`,
+            version: '0.0.1',
+        });
+    }
+    /**
+     * Connects to the MCP server.
+     */
+    async connect() {
+        this.isDisconnecting = false;
+        this.updateStatus(MCPServerStatus.CONNECTING);
+        try {
+            this.transport = await this.createTransport();
+            this.client.onerror = (error) => {
+                if (this.isDisconnecting) {
+                    return;
+                }
+                console.error(`MCP ERROR (${this.serverName}):`, error.toString());
+                this.updateStatus(MCPServerStatus.DISCONNECTED);
+            };
+            this.client.registerCapabilities({
+                roots: {},
+            });
+            this.client.setRequestHandler(ListRootsRequestSchema, async () => {
+                const roots = [];
+                for (const dir of this.workspaceContext.getDirectories()) {
+                    roots.push({
+                        uri: pathToFileURL(dir).toString(),
+                        name: basename(dir),
+                    });
+                }
+                return {
+                    roots,
+                };
+            });
+            await this.client.connect(this.transport, {
+                timeout: this.serverConfig.timeout,
+            });
+            this.updateStatus(MCPServerStatus.CONNECTED);
+        }
+        catch (error) {
+            this.updateStatus(MCPServerStatus.DISCONNECTED);
+            throw error;
+        }
+    }
+    /**
+     * Discovers tools and prompts from the MCP server.
+     */
+    async discover() {
+        if (this.status !== MCPServerStatus.CONNECTED) {
+            throw new Error('Client is not connected.');
+        }
+        const prompts = await this.discoverPrompts();
+        const tools = await this.discoverTools();
+        if (prompts.length === 0 && tools.length === 0) {
+            throw new Error('No prompts or tools found on the server.');
+        }
+        for (const tool of tools) {
+            this.toolRegistry.registerTool(tool);
+        }
+    }
+    /**
+     * Disconnects from the MCP server.
+     */
+    async disconnect() {
+        this.isDisconnecting = true;
+        if (this.transport) {
+            await this.transport.close();
+        }
+        this.client.close();
+        this.updateStatus(MCPServerStatus.DISCONNECTED);
+    }
+    /**
+     * Returns the current status of the client.
+     */
+    getStatus() {
+        return this.status;
+    }
+    updateStatus(status) {
+        this.status = status;
+        updateMCPServerStatus(this.serverName, status);
+    }
+    async createTransport() {
+        return createTransport(this.serverName, this.serverConfig, this.debugMode);
+    }
+    async discoverTools() {
+        return discoverTools(this.serverName, this.serverConfig, this.client);
+    }
+    async discoverPrompts() {
+        return discoverPrompts(this.serverName, this.client, this.promptRegistry);
+    }
+}
 /**
  * Map to track the status of each MCP server within the core package
  */
-const mcpServerStatusesInternal = new Map();
+const serverStatuses = new Map();
 /**
  * Track the overall MCP discovery state
  */
 let mcpDiscoveryState = MCPDiscoveryState.NOT_STARTED;
+/**
+ * Map to track which MCP servers have been discovered to require OAuth
+ */
+export const mcpServerRequiresOAuth = new Map();
 const statusChangeListeners = [];
 /**
  * Add a listener for MCP server status changes
@@ -64,8 +191,8 @@ export function removeMCPStatusChangeListener(listener) {
 /**
  * Update the status of an MCP server
  */
-function updateMCPServerStatus(serverName, status) {
-    mcpServerStatusesInternal.set(serverName, status);
+export function updateMCPServerStatus(serverName, status) {
+    serverStatuses.set(serverName, status);
     // Notify all listeners
     for (const listener of statusChangeListeners) {
         listener(serverName, status);
@@ -75,13 +202,13 @@ function updateMCPServerStatus(serverName, status) {
  * Get the current status of an MCP server
  */
 export function getMCPServerStatus(serverName) {
-    return (mcpServerStatusesInternal.get(serverName) || MCPServerStatus.DISCONNECTED);
+    return serverStatuses.get(serverName) || MCPServerStatus.DISCONNECTED;
 }
 /**
  * Get all MCP server statuses
  */
 export function getAllMCPServerStatuses() {
-    return new Map(mcpServerStatusesInternal);
+    return new Map(serverStatuses);
 }
 /**
  * Get the current MCP discovery state
@@ -89,6 +216,123 @@ export function getAllMCPServerStatuses() {
 export function getMCPDiscoveryState() {
     return mcpDiscoveryState;
 }
+/**
+ * Extract WWW-Authenticate header from error message string.
+ * This is a more robust approach than regex matching.
+ *
+ * @param errorString The error message string
+ * @returns The www-authenticate header value if found, null otherwise
+ */
+function extractWWWAuthenticateHeader(errorString) {
+    // Try multiple patterns to extract the header
+    const patterns = [
+        /www-authenticate:\s*([^\n\r]+)/i,
+        /WWW-Authenticate:\s*([^\n\r]+)/i,
+        /"www-authenticate":\s*"([^"]+)"/i,
+        /'www-authenticate':\s*'([^']+)'/i,
+    ];
+    for (const pattern of patterns) {
+        const match = errorString.match(pattern);
+        if (match) {
+            return match[1].trim();
+        }
+    }
+    return null;
+}
+/**
+ * Handle automatic OAuth discovery and authentication for a server.
+ *
+ * @param mcpServerName The name of the MCP server
+ * @param mcpServerConfig The MCP server configuration
+ * @param wwwAuthenticate The www-authenticate header value
+ * @returns True if OAuth was successfully configured and authenticated, false otherwise
+ */
+async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthenticate) {
+    try {
+        console.log(`🔐 '${mcpServerName}' requires OAuth authentication`);
+        // Always try to parse the resource metadata URI from the www-authenticate header
+        let oauthConfig;
+        const resourceMetadataUri = OAuthUtils.parseWWWAuthenticateHeader(wwwAuthenticate);
+        if (resourceMetadataUri) {
+            oauthConfig = await OAuthUtils.discoverOAuthConfig(resourceMetadataUri);
+        }
+        else if (mcpServerConfig.url) {
+            // Fallback: try to discover OAuth config from the base URL for SSE
+            const sseUrl = new URL(mcpServerConfig.url);
+            const baseUrl = `${sseUrl.protocol}//${sseUrl.host}`;
+            oauthConfig = await OAuthUtils.discoverOAuthConfig(baseUrl);
+        }
+        else if (mcpServerConfig.httpUrl) {
+            // Fallback: try to discover OAuth config from the base URL for HTTP
+            const httpUrl = new URL(mcpServerConfig.httpUrl);
+            const baseUrl = `${httpUrl.protocol}//${httpUrl.host}`;
+            oauthConfig = await OAuthUtils.discoverOAuthConfig(baseUrl);
+        }
+        if (!oauthConfig) {
+            console.error(`❌ Could not configure OAuth for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
+            return false;
+        }
+        // OAuth configuration discovered - proceed with authentication
+        // Create OAuth configuration for authentication
+        const oauthAuthConfig = {
+            enabled: true,
+            authorizationUrl: oauthConfig.authorizationUrl,
+            tokenUrl: oauthConfig.tokenUrl,
+            scopes: oauthConfig.scopes || [],
+        };
+        // Perform OAuth authentication
+        // Pass the server URL for proper discovery
+        const serverUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
+        console.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
+        await MCPOAuthProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
+        console.log(`OAuth authentication successful for server '${mcpServerName}'`);
+        return true;
+    }
+    catch (error) {
+        console.error(`Failed to handle automatic OAuth for server '${mcpServerName}': ${getErrorMessage(error)}`);
+        return false;
+    }
+}
+/**
+ * Create a transport with OAuth token for the given server configuration.
+ *
+ * @param mcpServerName The name of the MCP server
+ * @param mcpServerConfig The MCP server configuration
+ * @param accessToken The OAuth access token
+ * @returns The transport with OAuth token, or null if creation fails
+ */
+async function createTransportWithOAuth(mcpServerName, mcpServerConfig, accessToken) {
+    try {
+        if (mcpServerConfig.httpUrl) {
+            // Create HTTP transport with OAuth token
+            const oauthTransportOptions = {
+                requestInit: {
+                    headers: {
+                        ...mcpServerConfig.headers,
+                        Authorization: `Bearer ${accessToken}`,
+                    },
+                },
+            };
+            return new StreamableHTTPClientTransport(new URL(mcpServerConfig.httpUrl), oauthTransportOptions);
+        }
+        else if (mcpServerConfig.url) {
+            // Create SSE transport with OAuth token in Authorization header
+            return new SSEClientTransport(new URL(mcpServerConfig.url), {
+                requestInit: {
+                    headers: {
+                        ...mcpServerConfig.headers,
+                        Authorization: `Bearer ${accessToken}`,
+                    },
+                },
+            });
+        }
+        return null;
+    }
+    catch (error) {
+        console.error(`Failed to create OAuth transport for server '${mcpServerName}': ${getErrorMessage(error)}`);
+        return null;
+    }
+}
 /**
  * Discovers tools from all configured MCP servers and registers them with the tool registry.
  * It orchestrates the connection and discovery process for each server defined in the
@@ -99,11 +343,11 @@ export function getMCPDiscoveryState() {
  * @param toolRegistry The central registry where discovered tools will be registered.
  * @returns A promise that resolves when the discovery process has been attempted for all servers.
  */
-export async function discoverMcpTools(mcpServers, mcpServerCommand, toolRegistry, debugMode) {
+export async function discoverMcpTools(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext) {
     mcpDiscoveryState = MCPDiscoveryState.IN_PROGRESS;
     try {
         mcpServers = populateMcpServerCommand(mcpServers, mcpServerCommand);
-        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, debugMode));
+        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext));
         await Promise.all(discoveryPromises);
     }
     finally {
@@ -136,39 +380,90 @@ export function populateMcpServerCommand(mcpServers, mcpServerCommand) {
  * @param toolRegistry The registry to register discovered tools with
  * @returns Promise that resolves when discovery is complete
  */
-export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, debugMode) {
+export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext) {
     updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTING);
+    let mcpClient;
     try {
-        const mcpClient = await connectToMcpServer(mcpServerName, mcpServerConfig, debugMode);
-        try {
-            updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
-            mcpClient.onerror = (error) => {
-                console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
-                updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
-                if (mcpServerName === IDE_SERVER_NAME) {
-                    ideContext.clearOpenFilesContext();
-                }
-            };
-            if (mcpServerName === IDE_SERVER_NAME) {
-                mcpClient.setNotificationHandler(OpenFilesNotificationSchema, (notification) => {
-                    ideContext.setOpenFilesContext(notification.params);
-                });
-            }
-            const tools = await discoverTools(mcpServerName, mcpServerConfig, mcpClient);
-            for (const tool of tools) {
-                toolRegistry.registerTool(tool);
-            }
+        mcpClient = await connectToMcpServer(mcpServerName, mcpServerConfig, debugMode, workspaceContext);
+        mcpClient.onerror = (error) => {
+            console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
+            updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
+        };
+        // Attempt to discover both prompts and tools
+        const prompts = await discoverPrompts(mcpServerName, mcpClient, promptRegistry);
+        const tools = await discoverTools(mcpServerName, mcpServerConfig, mcpClient);
+        // If we have neither prompts nor tools, it's a failed discovery
+        if (prompts.length === 0 && tools.length === 0) {
+            throw new Error('No prompts or tools found on the server.');
         }
-        catch (error) {
-            mcpClient.close();
-            throw error;
+        // If we found anything, the server is connected
+        updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
+        // Register any discovered tools
+        for (const tool of tools) {
+            toolRegistry.registerTool(tool);
         }
     }
     catch (error) {
+        if (mcpClient) {
+            mcpClient.close();
+        }
         console.error(`Error connecting to MCP server '${mcpServerName}': ${getErrorMessage(error)}`);
         updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
     }
 }
+/**
+ * Recursively validates that a JSON schema and all its nested properties and
+ * items have a `type` defined.
+ *
+ * @param schema The JSON schema to validate.
+ * @returns `true` if the schema is valid, `false` otherwise.
+ *
+ * @visiblefortesting
+ */
+export function hasValidTypes(schema) {
+    if (typeof schema !== 'object' || schema === null) {
+        // Not a schema object we can validate, or not a schema at all.
+        // Treat as valid as it has no properties to be invalid.
+        return true;
+    }
+    const s = schema;
+    if (!s['type']) {
+        // These keywords contain an array of schemas that should be validated.
+        //
+        // If no top level type was given, then they must each have a type.
+        let hasSubSchema = false;
+        const schemaArrayKeywords = ['anyOf', 'allOf', 'oneOf'];
+        for (const keyword of schemaArrayKeywords) {
+            const subSchemas = s[keyword];
+            if (Array.isArray(subSchemas)) {
+                hasSubSchema = true;
+                for (const subSchema of subSchemas) {
+                    if (!hasValidTypes(subSchema)) {
+                        return false;
+                    }
+                }
+            }
+        }
+        // If the node itself is missing a type and had no subschemas, then it isn't valid.
+        if (!hasSubSchema)
+            return false;
+    }
+    if (s['type'] === 'object' && s['properties']) {
+        if (typeof s['properties'] === 'object' && s['properties'] !== null) {
+            for (const prop of Object.values(s['properties'])) {
+                if (!hasValidTypes(prop)) {
+                    return false;
+                }
+            }
+        }
+    }
+    if (s['type'] === 'array' && s['items']) {
+        if (!hasValidTypes(s['items'])) {
+            return false;
+        }
+    }
+    return true;
+}
 /**
  * Discovers and sanitizes tools from a connected MCP client.
  * It retrieves function declarations from the client, filters out disabled tools,
@@ -185,22 +480,95 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
         const mcpCallableTool = mcpToTool(mcpClient);
         const tool = await mcpCallableTool.tool();
         if (!Array.isArray(tool.functionDeclarations)) {
-            throw new Error(`Server did not return valid function declarations.`);
+            // This is a valid case for a prompt-only server
+            return [];
         }
         const discoveredTools = [];
         for (const funcDecl of tool.functionDeclarations) {
-            if (!isEnabled(funcDecl, mcpServerName, mcpServerConfig)) {
-                continue;
+            try {
+                if (!isEnabled(funcDecl, mcpServerName, mcpServerConfig)) {
+                    continue;
+                }
+                if (!hasValidTypes(funcDecl.parametersJsonSchema)) {
+                    console.warn(`Skipping tool '${funcDecl.name}' from MCP server '${mcpServerName}' ` +
+                        `because it has missing types in its parameter schema. Please file an ` +
+                        `issue with the owner of the MCP server.`);
+                    continue;
+                }
+                discoveredTools.push(new DiscoveredMCPTool(mcpCallableTool, mcpServerName, funcDecl.name, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, mcpServerConfig.trust));
+            }
+            catch (error) {
+                console.error(`Error discovering tool: '${funcDecl.name}' from MCP server '${mcpServerName}': ${error.message}`);
             }
-            discoveredTools.push(new DiscoveredMCPTool(mcpCallableTool, mcpServerName, funcDecl.name, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, mcpServerConfig.trust));
-        }
-        if (discoveredTools.length === 0) {
-            throw Error('No enabled tools found');
         }
         return discoveredTools;
     }
     catch (error) {
-        throw new Error(`Error discovering tools: ${error}`);
+        if (error instanceof Error &&
+            !error.message?.includes('Method not found')) {
+            console.error(`Error discovering tools from ${mcpServerName}: ${getErrorMessage(error)}`);
+        }
+        return [];
+    }
+}
+/**
+ * Discovers and logs prompts from a connected MCP client.
+ * It retrieves prompt declarations from the client and logs their names.
+ *
+ * @param mcpServerName The name of the MCP server.
+ * @param mcpClient The active MCP client instance.
+ */
+export async function discoverPrompts(mcpServerName, mcpClient, promptRegistry) {
+    try {
+        // Only request prompts if the server supports them.
+        if (mcpClient.getServerCapabilities()?.prompts == null)
+            return [];
+        const response = await mcpClient.request({ method: 'prompts/list', params: {} }, ListPromptsResultSchema);
+        for (const prompt of response.prompts) {
+            promptRegistry.registerPrompt({
+                ...prompt,
+                serverName: mcpServerName,
+                invoke: (params) => invokeMcpPrompt(mcpServerName, mcpClient, prompt.name, params),
+            });
+        }
+        return response.prompts;
+    }
+    catch (error) {
+        // It's okay if this fails, not all servers will have prompts.
+        // Don't log an error if the method is not found, which is a common case.
+        if (error instanceof Error &&
+            !error.message?.includes('Method not found')) {
+            console.error(`Error discovering prompts from ${mcpServerName}: ${getErrorMessage(error)}`);
+        }
+        return [];
+    }
+}
+/**
+ * Invokes a prompt on a connected MCP client.
+ *
+ * @param mcpServerName The name of the MCP server.
+ * @param mcpClient The active MCP client instance.
+ * @param promptName The name of the prompt to invoke.
+ * @param promptParams The parameters to pass to the prompt.
+ * @returns A promise that resolves to the result of the prompt invocation.
+ */
+export async function invokeMcpPrompt(mcpServerName, mcpClient, promptName, promptParams) {
+    try {
+        const response = await mcpClient.request({
+            method: 'prompts/get',
+            params: {
+                name: promptName,
+                arguments: promptParams,
+            },
+        }, GetPromptResultSchema);
+        return response;
+    }
+    catch (error) {
+        if (error instanceof Error &&
+            !error.message?.includes('Method not found')) {
+            console.error(`Error invoking prompt '${promptName}' from ${mcpServerName} ${promptParams}: ${getErrorMessage(error)}`);
+        }
+        throw error;
     }
 }
 /**
@@ -213,11 +581,52 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
  * @returns A promise that resolves to a connected MCP `Client` instance.
  * @throws An error if the connection fails or the configuration is invalid.
  */
-export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMode) {
+export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMode, workspaceContext) {
     const mcpClient = new Client({
         name: 'gemini-cli-mcp-client',
         version: '0.0.1',
     });
+    mcpClient.registerCapabilities({
+        roots: {
+            listChanged: true,
+        },
+    });
+    mcpClient.setRequestHandler(ListRootsRequestSchema, async () => {
+        const roots = [];
+        for (const dir of workspaceContext.getDirectories()) {
+            roots.push({
+                uri: pathToFileURL(dir).toString(),
+                name: basename(dir),
+            });
+        }
+        return {
+            roots,
+        };
+    });
+    let unlistenDirectories = workspaceContext.onDirectoriesChanged(async () => {
+        try {
+            await mcpClient.notification({
+                method: 'notifications/roots/list_changed',
+            });
+        }
+        catch (_) {
+            // If this fails, its almost certainly because the connection was closed
+            // and we should just stop listening for future directory changes.
+            unlistenDirectories?.();
+            unlistenDirectories = undefined;
+        }
+    });
+    // Attempt to pro-actively unsubscribe if the mcp client closes. This API is
+    // very brittle though so we don't have any guarantees, hence the try/catch
+    // above as well.
+    //
+    // Be a good steward and don't just bash over onclose.
+    const oldOnClose = mcpClient.onclose;
+    mcpClient.onclose = () => {
+        oldOnClose?.();
+        unlistenDirectories?.();
+        unlistenDirectories = undefined;
+    };
     // patch Client.callTool to use request timeout as genai McpCallTool.callTool does not do it
     // TODO: remove this hack once GenAI SDK does callTool with request options
     if ('callTool' in mcpClient) {
@@ -230,7 +639,7 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
         };
     }
     try {
-        const transport = createTransport(mcpServerName, mcpServerConfig, debugMode);
+        const transport = await createTransport(mcpServerName, mcpServerConfig, debugMode);
         try {
             await mcpClient.connect(transport, {
                 timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
@@ -243,29 +652,282 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
         }
     }
     catch (error) {
-        // Create a safe config object that excludes sensitive information
-        const safeConfig = {
-            command: mcpServerConfig.command,
-            url: mcpServerConfig.url,
-            httpUrl: mcpServerConfig.httpUrl,
-            cwd: mcpServerConfig.cwd,
-            timeout: mcpServerConfig.timeout,
-            trust: mcpServerConfig.trust,
-            // Exclude args, env, and headers which may contain sensitive data
-        };
-        let errorString = `failed to start or connect to MCP server '${mcpServerName}' ` +
-            `${JSON.stringify(safeConfig)}; \n${error}`;
-        if (process.env.SANDBOX) {
-            errorString += `\nMake sure it is available in the sandbox`;
+        // Check if this is a 401 error that might indicate OAuth is required
+        const errorString = String(error);
+        if (errorString.includes('401') &&
+            (mcpServerConfig.httpUrl || mcpServerConfig.url)) {
+            mcpServerRequiresOAuth.set(mcpServerName, true);
+            // Only trigger automatic OAuth discovery for HTTP servers or when OAuth is explicitly configured
+            // For SSE servers, we should not trigger new OAuth flows automatically
+            const shouldTriggerOAuth = mcpServerConfig.httpUrl || mcpServerConfig.oauth?.enabled;
+            if (!shouldTriggerOAuth) {
+                // For SSE servers without explicit OAuth config, if a token was found but rejected, report it accurately.
+                const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+                if (credentials) {
+                    const hasStoredTokens = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                        // Pass client ID if available
+                        clientId: credentials.clientId,
+                    });
+                    if (hasStoredTokens) {
+                        console.log(`Stored OAuth token for SSE server '${mcpServerName}' was rejected. ` +
+                            `Please re-authenticate using: /mcp auth ${mcpServerName}`);
+                    }
+                    else {
+                        console.log(`401 error received for SSE server '${mcpServerName}' without OAuth configuration. ` +
+                            `Please authenticate using: /mcp auth ${mcpServerName}`);
+                    }
+                }
+                throw new Error(`401 error received for SSE server '${mcpServerName}' without OAuth configuration. ` +
+                    `Please authenticate using: /mcp auth ${mcpServerName}`);
+            }
+            // Try to extract www-authenticate header from the error
+            let wwwAuthenticate = extractWWWAuthenticateHeader(errorString);
+            // If we didn't get the header from the error string, try to get it from the server
+            if (!wwwAuthenticate && mcpServerConfig.url) {
+                console.log(`No www-authenticate header in error, trying to fetch it from server...`);
+                try {
+                    const response = await fetch(mcpServerConfig.url, {
+                        method: 'HEAD',
+                        headers: {
+                            Accept: 'text/event-stream',
+                        },
+                        signal: AbortSignal.timeout(5000),
+                    });
+                    if (response.status === 401) {
+                        wwwAuthenticate = response.headers.get('www-authenticate');
+                        if (wwwAuthenticate) {
+                            console.log(`Found www-authenticate header from server: ${wwwAuthenticate}`);
+                        }
+                    }
+                }
+                catch (fetchError) {
+                    console.debug(`Failed to fetch www-authenticate header: ${getErrorMessage(fetchError)}`);
+                }
+            }
+            if (wwwAuthenticate) {
+                console.log(`Received 401 with www-authenticate header: ${wwwAuthenticate}`);
+                // Try automatic OAuth discovery and authentication
+                const oauthSuccess = await handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthenticate);
+                if (oauthSuccess) {
+                    // Retry connection with OAuth token
+                    console.log(`Retrying connection to '${mcpServerName}' with OAuth token...`);
+                    // Get the valid token - we need to create a proper OAuth config
+                    // The token should already be available from the authentication process
+                    const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+                    if (credentials) {
+                        const accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                            // Pass client ID if available
+                            clientId: credentials.clientId,
+                        });
+                        if (accessToken) {
+                            // Create transport with OAuth token
+                            const oauthTransport = await createTransportWithOAuth(mcpServerName, mcpServerConfig, accessToken);
+                            if (oauthTransport) {
+                                try {
+                                    await mcpClient.connect(oauthTransport, {
+                                        timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+                                    });
+                                    // Connection successful with OAuth
+                                    return mcpClient;
+                                }
+                                catch (retryError) {
+                                    console.error(`Failed to connect with OAuth token: ${getErrorMessage(retryError)}`);
+                                    throw retryError;
+                                }
+                            }
+                            else {
+                                console.error(`Failed to create OAuth transport for server '${mcpServerName}'`);
+                                throw new Error(`Failed to create OAuth transport for server '${mcpServerName}'`);
+                            }
+                        }
+                        else {
+                            console.error(`Failed to get OAuth token for server '${mcpServerName}'`);
+                            throw new Error(`Failed to get OAuth token for server '${mcpServerName}'`);
+                        }
+                    }
+                    else {
+                        console.error(`Failed to get credentials for server '${mcpServerName}' after successful OAuth authentication`);
+                        throw new Error(`Failed to get credentials for server '${mcpServerName}' after successful OAuth authentication`);
+                    }
+                }
+                else {
+                    console.error(`Failed to handle automatic OAuth for server '${mcpServerName}'`);
+                    throw new Error(`Failed to handle automatic OAuth for server '${mcpServerName}'`);
+                }
+            }
+            else {
+                // No www-authenticate header found, but we got a 401
+                // Only try OAuth discovery for HTTP servers or when OAuth is explicitly configured
+                // For SSE servers, we should not trigger new OAuth flows automatically
+                const shouldTryDiscovery = mcpServerConfig.httpUrl || mcpServerConfig.oauth?.enabled;
+                if (!shouldTryDiscovery) {
+                    const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+                    if (credentials) {
+                        const hasStoredTokens = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                            // Pass client ID if available
+                            clientId: credentials.clientId,
+                        });
+                        if (hasStoredTokens) {
+                            console.log(`Stored OAuth token for SSE server '${mcpServerName}' was rejected. ` +
+                                `Please re-authenticate using: /mcp auth ${mcpServerName}`);
+                        }
+                        else {
+                            console.log(`401 error received for SSE server '${mcpServerName}' without OAuth configuration. ` +
+                                `Please authenticate using: /mcp auth ${mcpServerName}`);
+                        }
+                    }
+                    throw new Error(`401 error received for SSE server '${mcpServerName}' without OAuth configuration. ` +
+                        `Please authenticate using: /mcp auth ${mcpServerName}`);
+                }
+                // For SSE servers, try to discover OAuth configuration from the base URL
+                console.log(`🔍 Attempting OAuth discovery for '${mcpServerName}'...`);
+                if (mcpServerConfig.url) {
+                    const sseUrl = new URL(mcpServerConfig.url);
+                    const baseUrl = `${sseUrl.protocol}//${sseUrl.host}`;
+                    try {
+                        // Try to discover OAuth configuration from the base URL
+                        const oauthConfig = await OAuthUtils.discoverOAuthConfig(baseUrl);
+                        if (oauthConfig) {
+                            console.log(`Discovered OAuth configuration from base URL for server '${mcpServerName}'`);
+                            // Create OAuth configuration for authentication
+                            const oauthAuthConfig = {
+                                enabled: true,
+                                authorizationUrl: oauthConfig.authorizationUrl,
+                                tokenUrl: oauthConfig.tokenUrl,
+                                scopes: oauthConfig.scopes || [],
+                            };
+                            // Perform OAuth authentication
+                            // Pass the server URL for proper discovery
+                            const serverUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
+                            console.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
+                            await MCPOAuthProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
+                            // Retry connection with OAuth token
+                            const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+                            if (credentials) {
+                                const accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                                    // Pass client ID if available
+                                    clientId: credentials.clientId,
+                                });
+                                if (accessToken) {
+                                    // Create transport with OAuth token
+                                    const oauthTransport = await createTransportWithOAuth(mcpServerName, mcpServerConfig, accessToken);
+                                    if (oauthTransport) {
+                                        try {
+                                            await mcpClient.connect(oauthTransport, {
+                                                timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+                                            });
+                                            // Connection successful with OAuth
+                                            return mcpClient;
+                                        }
+                                        catch (retryError) {
+                                            console.error(`Failed to connect with OAuth token: ${getErrorMessage(retryError)}`);
+                                            throw retryError;
+                                        }
+                                    }
+                                    else {
+                                        console.error(`Failed to create OAuth transport for server '${mcpServerName}'`);
+                                        throw new Error(`Failed to create OAuth transport for server '${mcpServerName}'`);
+                                    }
+                                }
+                                else {
+                                    console.error(`Failed to get OAuth token for server '${mcpServerName}'`);
+                                    throw new Error(`Failed to get OAuth token for server '${mcpServerName}'`);
+                                }
+                            }
+                            else {
+                                console.error(`Failed to get stored credentials for server '${mcpServerName}'`);
+                                throw new Error(`Failed to get stored credentials for server '${mcpServerName}'`);
+                            }
+                        }
+                        else {
+                            console.error(`❌ Could not configure OAuth for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
+                            throw new Error(`OAuth configuration failed for '${mcpServerName}'. Please authenticate manually with /mcp auth ${mcpServerName}`);
+                        }
+                    }
+                    catch (discoveryError) {
+                        console.error(`❌ OAuth discovery failed for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
+                        throw discoveryError;
+                    }
+                }
+                else {
+                    console.error(`❌ '${mcpServerName}' requires authentication but no OAuth configuration found`);
+                    throw new Error(`MCP server '${mcpServerName}' requires authentication. Please configure OAuth or check server settings.`);
+                }
+            }
+        }
+        else {
+            // Handle other connection errors
+            // Create a concise error message
+            const errorMessage = error.message || String(error);
+            const isNetworkError = errorMessage.includes('ENOTFOUND') ||
+                errorMessage.includes('ECONNREFUSED');
+            let conciseError;
+            if (isNetworkError) {
+                conciseError = `Cannot connect to '${mcpServerName}' - server may be down or URL incorrect`;
+            }
+            else {
+                conciseError = `Connection failed for '${mcpServerName}': ${errorMessage}`;
+            }
+            if (process.env['SANDBOX']) {
+                conciseError += ` (check sandbox availability)`;
+            }
+            throw new Error(conciseError);
         }
-        throw new Error(errorString);
     }
 }
 /** Visible for Testing */
-export function createTransport(mcpServerName, mcpServerConfig, debugMode) {
+export async function createTransport(mcpServerName, mcpServerConfig, debugMode) {
+    if (mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS) {
+        const provider = new GoogleCredentialProvider(mcpServerConfig);
+        const transportOptions = {
+            authProvider: provider,
+        };
+        if (mcpServerConfig.httpUrl) {
+            return new StreamableHTTPClientTransport(new URL(mcpServerConfig.httpUrl), transportOptions);
+        }
+        else if (mcpServerConfig.url) {
+            return new SSEClientTransport(new URL(mcpServerConfig.url), transportOptions);
+        }
+        throw new Error('No URL configured for Google Credentials MCP server');
+    }
+    // Check if we have OAuth configuration or stored tokens
+    let accessToken = null;
+    let hasOAuthConfig = mcpServerConfig.oauth?.enabled;
+    if (hasOAuthConfig && mcpServerConfig.oauth) {
+        accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, mcpServerConfig.oauth);
+        if (!accessToken) {
+            console.error(`MCP server '${mcpServerName}' requires OAuth authentication. ` +
+                `Please authenticate using the /mcp auth command.`);
+            throw new Error(`MCP server '${mcpServerName}' requires OAuth authentication. ` +
+                `Please authenticate using the /mcp auth command.`);
+        }
+    }
+    else {
+        // Check if we have stored OAuth tokens for this server (from previous authentication)
+        const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+        if (credentials) {
+            accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                // Pass client ID if available
+                clientId: credentials.clientId,
+            });
+            if (accessToken) {
+                hasOAuthConfig = true;
+                console.log(`Found stored OAuth token for server '${mcpServerName}'`);
+            }
+        }
+    }
     if (mcpServerConfig.httpUrl) {
         const transportOptions = {};
-        if (mcpServerConfig.headers) {
+        // Set up headers with OAuth token if available
+        if (hasOAuthConfig && accessToken) {
+            transportOptions.requestInit = {
+                headers: {
+                    ...mcpServerConfig.headers,
+                    Authorization: `Bearer ${accessToken}`,
+                },
+            };
+        }
+        else if (mcpServerConfig.headers) {
             transportOptions.requestInit = {
                 headers: mcpServerConfig.headers,
             };
@@ -274,7 +936,16 @@ export function createTransport(mcpServerName, mcpServerConfig, debugMode) {
     }
     if (mcpServerConfig.url) {
         const transportOptions = {};
-        if (mcpServerConfig.headers) {
+        // Set up headers with OAuth token if available
+        if (hasOAuthConfig && accessToken) {
+            transportOptions.requestInit = {
+                headers: {
+                    ...mcpServerConfig.headers,
+                    Authorization: `Bearer ${accessToken}`,
+                },
+            };
+        }
+        else if (mcpServerConfig.headers) {
             transportOptions.requestInit = {
                 headers: mcpServerConfig.headers,
             };