@office-ai/aioncli-core 0.2.3 → 0.18.4

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

@@ -4,21 +4,23 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.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 { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { 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 { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
-import { mcpToTool } from '@google/genai';
+import { basename } from 'node:path';
+import { pathToFileURL } from 'node:url';
 import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
-import { OAuthUtils } from '../mcp/oauth-utils.js';
 import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
+import { OAuthUtils } from '../mcp/oauth-utils.js';
 import { getErrorMessage } from '../utils/errors.js';
-import { basename } from 'node:path';
-import { pathToFileURL } from 'node:url';
+import { debugLogger } from '../utils/debugLogger.js';
+import { coreEvents } from '../utils/events.js';
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 /**
  * Enum representing the connection status of an MCP server
@@ -27,6 +29,8 @@ export var MCPServerStatus;
 (function (MCPServerStatus) {
     /** Server is disconnected or experiencing errors */
     MCPServerStatus["DISCONNECTED"] = "disconnected";
+    /** Server is actively disconnecting */
+    MCPServerStatus["DISCONNECTING"] = "disconnecting";
     /** Server is in the process of connecting */
     MCPServerStatus["CONNECTING"] = "connecting";
     /** Server is connected and ready to use */
@@ -60,7 +64,6 @@ export class McpClient {
     client;
     transport;
     status = MCPServerStatus.DISCONNECTED;
-    isDisconnecting = false;
     constructor(serverName, serverConfig, toolRegistry, promptRegistry, workspaceContext, debugMode) {
         this.serverName = serverName;
         this.serverConfig = serverConfig;
@@ -68,44 +71,27 @@ export class McpClient {
         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;
+        if (this.status !== MCPServerStatus.DISCONNECTED) {
+            throw new Error(`Can only connect when the client is disconnected, current state is ${this.status}`);
+        }
         this.updateStatus(MCPServerStatus.CONNECTING);
         try {
-            this.transport = await this.createTransport();
+            this.client = await connectToMcpServer(this.serverName, this.serverConfig, this.debugMode, this.workspaceContext);
+            const originalOnError = this.client.onerror;
             this.client.onerror = (error) => {
-                if (this.isDisconnecting) {
+                if (this.status !== MCPServerStatus.CONNECTED) {
                     return;
                 }
-                console.error(`MCP ERROR (${this.serverName}):`, error.toString());
+                if (originalOnError)
+                    originalOnError(error);
+                coreEvents.emitFeedback('error', `MCP ERROR (${this.serverName})`, error);
                 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) {
@@ -116,28 +102,36 @@ export class McpClient {
     /**
      * Discovers tools and prompts from the MCP server.
      */
-    async discover() {
-        if (this.status !== MCPServerStatus.CONNECTED) {
-            throw new Error('Client is not connected.');
-        }
+    async discover(cliConfig) {
+        this.assertConnected();
         const prompts = await this.discoverPrompts();
-        const tools = await this.discoverTools();
+        const tools = await this.discoverTools(cliConfig);
         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);
         }
+        this.toolRegistry.sortTools();
     }
     /**
      * Disconnects from the MCP server.
      */
     async disconnect() {
-        this.isDisconnecting = true;
+        if (this.status !== MCPServerStatus.CONNECTED) {
+            return;
+        }
+        this.toolRegistry.removeMcpToolsByServer(this.serverName);
+        this.promptRegistry.removePromptsByServer(this.serverName);
+        this.updateStatus(MCPServerStatus.DISCONNECTING);
+        const client = this.client;
+        this.client = undefined;
         if (this.transport) {
             await this.transport.close();
         }
-        this.client.close();
+        if (client) {
+            await client.close();
+        }
         this.updateStatus(MCPServerStatus.DISCONNECTED);
     }
     /**
@@ -150,15 +144,22 @@ export class McpClient {
         this.status = status;
         updateMCPServerStatus(this.serverName, status);
     }
-    async createTransport() {
-        return createTransport(this.serverName, this.serverConfig, this.debugMode);
+    assertConnected() {
+        if (this.status !== MCPServerStatus.CONNECTED) {
+            throw new Error(`Client is not connected, must connect before interacting with the server. Current state is ${this.status}`);
+        }
     }
-    async discoverTools() {
-        return discoverTools(this.serverName, this.serverConfig, this.client);
+    async discoverTools(cliConfig) {
+        this.assertConnected();
+        return discoverTools(this.serverName, this.serverConfig, this.client, cliConfig, this.toolRegistry.getMessageBus());
     }
     async discoverPrompts() {
+        this.assertConnected();
         return discoverPrompts(this.serverName, this.client, this.promptRegistry);
     }
+    getServerConfig() {
+        return this.serverConfig;
+    }
 }
 /**
  * Map to track the status of each MCP server within the core package
@@ -249,27 +250,21 @@ function extractWWWAuthenticateHeader(errorString) {
  */
 async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthenticate) {
     try {
-        console.log(`🔐 '${mcpServerName}' requires OAuth authentication`);
+        debugLogger.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}`;
+        else if (hasNetworkTransport(mcpServerConfig)) {
+            // Fallback: try to discover OAuth config from the base URL
+            const serverUrl = new URL(mcpServerConfig.httpUrl || mcpServerConfig.url);
+            const baseUrl = `${serverUrl.protocol}//${serverUrl.host}`;
             oauthConfig = await OAuthUtils.discoverOAuthConfig(baseUrl);
         }
         if (!oauthConfig) {
-            console.error(`❌ Could not configure OAuth for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
+            coreEvents.emitFeedback('error', `Could not configure OAuth for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
             return false;
         }
         // OAuth configuration discovered - proceed with authentication
@@ -283,16 +278,61 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
         // 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}'`);
+        debugLogger.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
+        const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
+        await authProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
+        debugLogger.log(`OAuth authentication successful for server '${mcpServerName}'`);
         return true;
     }
     catch (error) {
-        console.error(`Failed to handle automatic OAuth for server '${mcpServerName}': ${getErrorMessage(error)}`);
+        coreEvents.emitFeedback('error', `Failed to handle automatic OAuth for server '${mcpServerName}': ${getErrorMessage(error)}`, error);
         return false;
     }
 }
+/**
+ * Create RequestInit for TransportOptions.
+ *
+ * @param mcpServerConfig The MCP server configuration
+ * @param headers Additional headers
+ */
+function createTransportRequestInit(mcpServerConfig, headers) {
+    return {
+        headers: {
+            ...mcpServerConfig.headers,
+            ...headers,
+        },
+    };
+}
+/**
+ * Create an AuthProvider for the MCP Transport.
+ *
+ * @param mcpServerConfig The MCP server configuration
+ */
+function createAuthProvider(mcpServerConfig) {
+    if (mcpServerConfig.authProviderType ===
+        AuthProviderType.SERVICE_ACCOUNT_IMPERSONATION) {
+        return new ServiceAccountImpersonationProvider(mcpServerConfig);
+    }
+    if (mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS) {
+        return new GoogleCredentialProvider(mcpServerConfig);
+    }
+    return undefined;
+}
+/**
+ * Create a transport for URL based servers (remote servers).
+ *
+ * @param mcpServerConfig The MCP server configuration
+ * @param transportOptions The transport options
+ */
+function createUrlTransport(mcpServerConfig, transportOptions) {
+    if (mcpServerConfig.httpUrl) {
+        return new StreamableHTTPClientTransport(new URL(mcpServerConfig.httpUrl), transportOptions);
+    }
+    if (mcpServerConfig.url) {
+        return new SSEClientTransport(new URL(mcpServerConfig.url), transportOptions);
+    }
+    throw new Error('No URL configured for MCP Server');
+}
 /**
  * Create a transport with OAuth token for the given server configuration.
  *
@@ -303,33 +343,16 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
  */
 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;
+        const headers = {
+            Authorization: `Bearer ${accessToken}`,
+        };
+        const transportOptions = {
+            requestInit: createTransportRequestInit(mcpServerConfig, headers),
+        };
+        return createUrlTransport(mcpServerConfig, transportOptions);
     }
     catch (error) {
-        console.error(`Failed to create OAuth transport for server '${mcpServerName}': ${getErrorMessage(error)}`);
+        coreEvents.emitFeedback('error', `Failed to create OAuth transport for server '${mcpServerName}': ${getErrorMessage(error)}`, error);
         return null;
     }
 }
@@ -343,11 +366,11 @@ async function createTransportWithOAuth(mcpServerName, mcpServerConfig, accessTo
  * @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, promptRegistry, debugMode, workspaceContext) {
+export async function discoverMcpTools(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext, cliConfig) {
     mcpDiscoveryState = MCPDiscoveryState.IN_PROGRESS;
     try {
         mcpServers = populateMcpServerCommand(mcpServers, mcpServerCommand);
-        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext));
+        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext, cliConfig));
         await Promise.all(discoveryPromises);
     }
     finally {
@@ -380,18 +403,18 @@ 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, promptRegistry, debugMode, workspaceContext) {
+export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext, cliConfig) {
     updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTING);
     let mcpClient;
     try {
         mcpClient = await connectToMcpServer(mcpServerName, mcpServerConfig, debugMode, workspaceContext);
         mcpClient.onerror = (error) => {
-            console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
+            coreEvents.emitFeedback('error', `MCP ERROR (${mcpServerName}):`, error);
             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);
+        const tools = await discoverTools(mcpServerName, mcpServerConfig, mcpClient, cliConfig, toolRegistry.getMessageBus());
         // 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.');
@@ -402,68 +425,16 @@ export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolReg
         for (const tool of tools) {
             toolRegistry.registerTool(tool);
         }
+        toolRegistry.sortTools();
     }
     catch (error) {
         if (mcpClient) {
             mcpClient.close();
         }
-        console.error(`Error connecting to MCP server '${mcpServerName}': ${getErrorMessage(error)}`);
+        coreEvents.emitFeedback('error', `Error connecting to MCP server '${mcpServerName}': ${getErrorMessage(error)}`, 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,
@@ -472,33 +443,29 @@ export function hasValidTypes(schema) {
  * @param mcpServerName The name of the MCP server.
  * @param mcpServerConfig The configuration for the MCP server.
  * @param mcpClient The active MCP client instance.
+ * @param cliConfig The CLI configuration object.
+ * @param messageBus Optional message bus for policy engine integration.
  * @returns A promise that resolves to an array of discovered and enabled tools.
  * @throws An error if no enabled tools are found or if the server provides invalid function declarations.
  */
-export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
+export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient, cliConfig, messageBus) {
     try {
-        const mcpCallableTool = mcpToTool(mcpClient);
-        const tool = await mcpCallableTool.tool();
-        if (!Array.isArray(tool.functionDeclarations)) {
-            // This is a valid case for a prompt-only server
+        // Only request tools if the server supports them.
+        if (mcpClient.getServerCapabilities()?.tools == null)
             return [];
-        }
+        const response = await mcpClient.listTools({});
         const discoveredTools = [];
-        for (const funcDecl of tool.functionDeclarations) {
+        for (const toolDef of response.tools) {
             try {
-                if (!isEnabled(funcDecl, mcpServerName, mcpServerConfig)) {
+                if (!isEnabled(toolDef, 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));
+                const mcpCallableTool = new McpCallableTool(mcpClient, toolDef, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC);
+                const tool = new DiscoveredMCPTool(mcpCallableTool, mcpServerName, toolDef.name, toolDef.description ?? '', toolDef.inputSchema ?? { type: 'object', properties: {} }, mcpServerConfig.trust, undefined, cliConfig, mcpServerConfig.extension?.name, mcpServerConfig.extension?.id, messageBus);
+                discoveredTools.push(tool);
             }
             catch (error) {
-                console.error(`Error discovering tool: '${funcDecl.name}' from MCP server '${mcpServerName}': ${error.message}`);
+                coreEvents.emitFeedback('error', `Error discovering tool: '${toolDef.name}' from MCP server '${mcpServerName}': ${error.message}`, error);
             }
         }
         return discoveredTools;
@@ -506,11 +473,69 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
     catch (error) {
         if (error instanceof Error &&
             !error.message?.includes('Method not found')) {
-            console.error(`Error discovering tools from ${mcpServerName}: ${getErrorMessage(error)}`);
+            coreEvents.emitFeedback('error', `Error discovering tools from ${mcpServerName}: ${getErrorMessage(error)}`, error);
         }
         return [];
     }
 }
+class McpCallableTool {
+    client;
+    toolDef;
+    timeout;
+    constructor(client, toolDef, timeout) {
+        this.client = client;
+        this.toolDef = toolDef;
+        this.timeout = timeout;
+    }
+    async tool() {
+        return {
+            functionDeclarations: [
+                {
+                    name: this.toolDef.name,
+                    description: this.toolDef.description,
+                    parametersJsonSchema: this.toolDef.inputSchema,
+                },
+            ],
+        };
+    }
+    async callTool(functionCalls) {
+        // We only expect one function call at a time for MCP tools in this context
+        if (functionCalls.length !== 1) {
+            throw new Error('McpCallableTool only supports single function call');
+        }
+        const call = functionCalls[0];
+        try {
+            const result = await this.client.callTool({
+                name: call.name,
+                arguments: call.args,
+            }, undefined, { timeout: this.timeout });
+            return [
+                {
+                    functionResponse: {
+                        name: call.name,
+                        response: result,
+                    },
+                },
+            ];
+        }
+        catch (error) {
+            // Return error in the format expected by DiscoveredMCPTool
+            return [
+                {
+                    functionResponse: {
+                        name: call.name,
+                        response: {
+                            error: {
+                                message: error instanceof Error ? error.message : String(error),
+                                isError: true,
+                            },
+                        },
+                    },
+                },
+            ];
+        }
+    }
+}
 /**
  * Discovers and logs prompts from a connected MCP client.
  * It retrieves prompt declarations from the client and logs their names.
@@ -523,7 +548,7 @@ export async function discoverPrompts(mcpServerName, mcpClient, promptRegistry)
         // Only request prompts if the server supports them.
         if (mcpClient.getServerCapabilities()?.prompts == null)
             return [];
-        const response = await mcpClient.request({ method: 'prompts/list', params: {} }, ListPromptsResultSchema);
+        const response = await mcpClient.listPrompts({});
         for (const prompt of response.prompts) {
             promptRegistry.registerPrompt({
                 ...prompt,
@@ -538,7 +563,7 @@ export async function discoverPrompts(mcpServerName, mcpClient, promptRegistry)
         // 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)}`);
+            coreEvents.emitFeedback('error', `Error discovering prompts from ${mcpServerName}: ${getErrorMessage(error)}`, error);
         }
         return [];
     }
@@ -554,23 +579,35 @@ export async function discoverPrompts(mcpServerName, mcpClient, promptRegistry)
  */
 export async function invokeMcpPrompt(mcpServerName, mcpClient, promptName, promptParams) {
     try {
-        const response = await mcpClient.request({
-            method: 'prompts/get',
-            params: {
-                name: promptName,
-                arguments: promptParams,
-            },
-        }, GetPromptResultSchema);
+        const sanitizedParams = {};
+        for (const [key, value] of Object.entries(promptParams)) {
+            if (value !== undefined && value !== null) {
+                sanitizedParams[key] = String(value);
+            }
+        }
+        const response = await mcpClient.getPrompt({
+            name: promptName,
+            arguments: sanitizedParams,
+        });
         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)}`);
+            coreEvents.emitFeedback('error', `Error invoking prompt '${promptName}' from ${mcpServerName} ${promptParams}: ${getErrorMessage(error)}`, error);
         }
         throw error;
     }
 }
+/**
+ * @visiblefortesting
+ * Checks if the MCP server configuration has a network transport URL (SSE or HTTP).
+ * @param config The MCP server configuration.
+ * @returns True if a `url` or `httpUrl` is present, false otherwise.
+ */
+export function hasNetworkTransport(config) {
+    return !!(config.url || config.httpUrl);
+}
 /**
  * Creates and connects an MCP client to a server based on the provided configuration.
  * It determines the appropriate transport (Stdio, SSE, or Streamable HTTP) and
@@ -627,17 +664,6 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
         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) {
-        const origCallTool = mcpClient.callTool.bind(mcpClient);
-        mcpClient.callTool = function (params, resultSchema, options) {
-            return origCallTool(params, resultSchema, {
-                ...options,
-                timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
-            });
-        };
-    }
     try {
         const transport = await createTransport(mcpServerName, mcpServerConfig, debugMode);
         try {
@@ -654,26 +680,27 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
     catch (error) {
         // 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)) {
+        if (errorString.includes('401') && hasNetworkTransport(mcpServerConfig)) {
             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);
+                const tokenStorage = new MCPOAuthTokenStorage();
+                const credentials = await tokenStorage.getCredentials(mcpServerName);
                 if (credentials) {
-                    const hasStoredTokens = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                    const authProvider = new MCPOAuthProvider(tokenStorage);
+                    const hasStoredTokens = await authProvider.getValidToken(mcpServerName, {
                         // Pass client ID if available
                         clientId: credentials.clientId,
                     });
                     if (hasStoredTokens) {
-                        console.log(`Stored OAuth token for SSE server '${mcpServerName}' was rejected. ` +
+                        coreEvents.emitFeedback('error', `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. ` +
+                        coreEvents.emitFeedback('error', `401 error received for SSE server '${mcpServerName}' without OAuth configuration. ` +
                             `Please authenticate using: /mcp auth ${mcpServerName}`);
                     }
                 }
@@ -683,39 +710,44 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
             // 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...`);
+            if (!wwwAuthenticate && hasNetworkTransport(mcpServerConfig)) {
+                debugLogger.log(`No www-authenticate header in error, trying to fetch it from server...`);
                 try {
-                    const response = await fetch(mcpServerConfig.url, {
+                    const urlToFetch = mcpServerConfig.httpUrl || mcpServerConfig.url;
+                    const response = await fetch(urlToFetch, {
                         method: 'HEAD',
                         headers: {
-                            Accept: 'text/event-stream',
+                            Accept: mcpServerConfig.httpUrl
+                                ? 'application/json'
+                                : '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}`);
+                            debugLogger.log(`Found www-authenticate header from server: ${wwwAuthenticate}`);
                         }
                     }
                 }
                 catch (fetchError) {
-                    console.debug(`Failed to fetch www-authenticate header: ${getErrorMessage(fetchError)}`);
+                    debugLogger.debug(`Failed to fetch www-authenticate header: ${getErrorMessage(fetchError)}`);
                 }
             }
             if (wwwAuthenticate) {
-                console.log(`Received 401 with www-authenticate header: ${wwwAuthenticate}`);
+                debugLogger.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...`);
+                    debugLogger.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);
+                    const tokenStorage = new MCPOAuthTokenStorage();
+                    const credentials = await tokenStorage.getCredentials(mcpServerName);
                     if (credentials) {
-                        const accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                        const authProvider = new MCPOAuthProvider(tokenStorage);
+                        const accessToken = await authProvider.getValidToken(mcpServerName, {
                             // Pass client ID if available
                             clientId: credentials.clientId,
                         });
@@ -723,35 +755,25 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                             // 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;
-                                }
+                                await mcpClient.connect(oauthTransport, {
+                                    timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+                                });
+                                // Connection successful with OAuth
+                                return mcpClient;
                             }
                             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}'`);
                 }
             }
@@ -761,96 +783,84 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                 // 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);
+                    const tokenStorage = new MCPOAuthTokenStorage();
+                    const credentials = await tokenStorage.getCredentials(mcpServerName);
                     if (credentials) {
-                        const hasStoredTokens = await MCPOAuthProvider.getValidToken(mcpServerName, {
+                        const authProvider = new MCPOAuthProvider(tokenStorage);
+                        const hasStoredTokens = await authProvider.getValidToken(mcpServerName, {
                             // Pass client ID if available
                             clientId: credentials.clientId,
                         });
                         if (hasStoredTokens) {
-                            console.log(`Stored OAuth token for SSE server '${mcpServerName}' was rejected. ` +
+                            coreEvents.emitFeedback('error', `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. ` +
+                            coreEvents.emitFeedback('error', `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}'`);
-                                    }
+                // For SSE/HTTP servers, try to discover OAuth configuration from the base URL
+                debugLogger.log(`🔍 Attempting OAuth discovery for '${mcpServerName}'...`);
+                if (hasNetworkTransport(mcpServerConfig)) {
+                    const serverUrl = new URL(mcpServerConfig.httpUrl || mcpServerConfig.url);
+                    const baseUrl = `${serverUrl.protocol}//${serverUrl.host}`;
+                    // Try to discover OAuth configuration from the base URL
+                    const oauthConfig = await OAuthUtils.discoverOAuthConfig(baseUrl);
+                    if (oauthConfig) {
+                        debugLogger.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 authServerUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
+                        debugLogger.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
+                        const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
+                        await authProvider.authenticate(mcpServerName, oauthAuthConfig, authServerUrl);
+                        // Retry connection with OAuth token
+                        const tokenStorage = new MCPOAuthTokenStorage();
+                        const credentials = await tokenStorage.getCredentials(mcpServerName);
+                        if (credentials) {
+                            const authProvider = new MCPOAuthProvider(tokenStorage);
+                            const accessToken = await authProvider.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) {
+                                    await mcpClient.connect(oauthTransport, {
+                                        timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
+                                    });
+                                    // Connection successful with OAuth
+                                    return mcpClient;
                                 }
                                 else {
-                                    console.error(`Failed to get OAuth token for server '${mcpServerName}'`);
-                                    throw new Error(`Failed to get OAuth token for server '${mcpServerName}'`);
+                                    throw new Error(`Failed to create OAuth transport 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}'`);
+                                throw new Error(`Failed to get OAuth token 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}`);
+                            throw new Error(`Failed to get stored credentials for server '${mcpServerName}'`);
                         }
                     }
-                    catch (discoveryError) {
-                        console.error(`❌ OAuth discovery failed for '${mcpServerName}' - please authenticate manually with /mcp auth ${mcpServerName}`);
-                        throw discoveryError;
+                    else {
+                        throw new Error(`OAuth configuration failed for '${mcpServerName}'. Please authenticate manually with /mcp auth ${mcpServerName}`);
                     }
                 }
                 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.`);
                 }
             }
@@ -877,80 +887,57 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
 }
 /** Visible for Testing */
 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}'`);
+    const noUrl = !mcpServerConfig.url && !mcpServerConfig.httpUrl;
+    if (noUrl) {
+        if (mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS) {
+            throw new Error(`URL must be provided in the config for Google Credentials provider`);
+        }
+        if (mcpServerConfig.authProviderType ===
+            AuthProviderType.SERVICE_ACCOUNT_IMPERSONATION) {
+            throw new Error(`No URL configured for ServiceAccountImpersonation MCP Server`);
+        }
+    }
+    if (mcpServerConfig.httpUrl || mcpServerConfig.url) {
+        const authProvider = createAuthProvider(mcpServerConfig);
+        const headers = {};
+        if (authProvider === undefined) {
+            // Check if we have OAuth configuration or stored tokens
+            let accessToken = null;
+            let hasOAuthConfig = mcpServerConfig.oauth?.enabled;
+            if (hasOAuthConfig && mcpServerConfig.oauth) {
+                const tokenStorage = new MCPOAuthTokenStorage();
+                const mcpAuthProvider = new MCPOAuthProvider(tokenStorage);
+                accessToken = await mcpAuthProvider.getValidToken(mcpServerName, mcpServerConfig.oauth);
+                if (!accessToken) {
+                    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 tokenStorage = new MCPOAuthTokenStorage();
+                const credentials = await tokenStorage.getCredentials(mcpServerName);
+                if (credentials) {
+                    const mcpAuthProvider = new MCPOAuthProvider(tokenStorage);
+                    accessToken = await mcpAuthProvider.getValidToken(mcpServerName, {
+                        // Pass client ID if available
+                        clientId: credentials.clientId,
+                    });
+                    if (accessToken) {
+                        hasOAuthConfig = true;
+                        debugLogger.log(`Found stored OAuth token for server '${mcpServerName}'`);
+                    }
+                }
+            }
+            if (hasOAuthConfig && accessToken) {
+                headers['Authorization'] = `Bearer ${accessToken}`;
             }
         }
-    }
-    if (mcpServerConfig.httpUrl) {
-        const transportOptions = {};
-        // 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,
-            };
-        }
-        return new StreamableHTTPClientTransport(new URL(mcpServerConfig.httpUrl), transportOptions);
-    }
-    if (mcpServerConfig.url) {
-        const transportOptions = {};
-        // 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,
-            };
-        }
-        return new SSEClientTransport(new URL(mcpServerConfig.url), transportOptions);
+        const transportOptions = {
+            requestInit: createTransportRequestInit(mcpServerConfig, headers),
+            authProvider,
+        };
+        return createUrlTransport(mcpServerConfig, transportOptions);
     }
     if (mcpServerConfig.command) {
         const transport = new StdioClientTransport({
@@ -966,7 +953,7 @@ export async function createTransport(mcpServerName, mcpServerConfig, debugMode)
         if (debugMode) {
             transport.stderr.on('data', (data) => {
                 const stderrStr = data.toString().trim();
-                console.debug(`[DEBUG] [MCP STDERR (${mcpServerName})]: `, stderrStr);
+                debugLogger.debug(`[DEBUG] [MCP STDERR (${mcpServerName})]: `, stderrStr);
             });
         }
         return transport;
@@ -976,7 +963,7 @@ export async function createTransport(mcpServerName, mcpServerConfig, debugMode)
 /** Visible for testing */
 export function isEnabled(funcDecl, mcpServerName, mcpServerConfig) {
     if (!funcDecl.name) {
-        console.warn(`Discovered a function declaration without a name from MCP server '${mcpServerName}'. Skipping.`);
+        debugLogger.warn(`Discovered a function declaration without a name from MCP server '${mcpServerName}'. Skipping.`);
         return false;
     }
     const { includeTools, excludeTools } = mcpServerConfig;