@office-ai/aioncli-core 0.8.1 → 0.18.5

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

@@ -7,19 +7,20 @@ 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 { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { GetPromptResultSchema, ListPromptsResultSchema, ListRootsRequestSchema, } from '@modelcontextprotocol/sdk/types.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 { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { OAuthUtils } from '../mcp/oauth-utils.js';
 import { getErrorMessage } from '../utils/errors.js';
+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
@@ -28,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 */
@@ -61,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;
@@ -69,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) {
@@ -118,9 +103,7 @@ export class McpClient {
      * Discovers tools and prompts from the MCP server.
      */
     async discover(cliConfig) {
-        if (this.status !== MCPServerStatus.CONNECTED) {
-            throw new Error('Client is not connected.');
-        }
+        this.assertConnected();
         const prompts = await this.discoverPrompts();
         const tools = await this.discoverTools(cliConfig);
         if (prompts.length === 0 && tools.length === 0) {
@@ -129,16 +112,26 @@ export class McpClient {
         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);
     }
     /**
@@ -151,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(cliConfig) {
-        return discoverTools(this.serverName, this.serverConfig, this.client, 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
@@ -250,7 +250,7 @@ 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);
@@ -264,7 +264,7 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
             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
@@ -278,17 +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}'...`);
+        debugLogger.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
         const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
         await authProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
-        console.log(`OAuth authentication successful for server '${mcpServerName}'`);
+        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.
  *
@@ -299,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;
     }
 }
@@ -382,12 +409,12 @@ export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolReg
     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, cliConfig);
+        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.');
@@ -398,12 +425,13 @@ 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);
     }
 }
@@ -415,29 +443,29 @@ export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolReg
  * @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, cliConfig) {
+export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient, cliConfig, messageBus) {
     try {
-        const mcpCallableTool = mcpToTool(mcpClient, {
-            timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
-        });
-        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;
                 }
-                discoveredTools.push(new DiscoveredMCPTool(mcpCallableTool, mcpServerName, funcDecl.name, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, mcpServerConfig.trust, undefined, cliConfig));
+                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;
@@ -445,11 +473,69 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient, c
     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.
@@ -462,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,
@@ -477,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 [];
     }
@@ -493,19 +579,22 @@ 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;
     }
@@ -607,11 +696,11 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                         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}`);
                     }
                 }
@@ -622,7 +711,7 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
             let wwwAuthenticate = extractWWWAuthenticateHeader(errorString);
             // If we didn't get the header from the error string, try to get it from the server
             if (!wwwAuthenticate && hasNetworkTransport(mcpServerConfig)) {
-                console.log(`No www-authenticate header in error, trying to fetch it from server...`);
+                debugLogger.log(`No www-authenticate header in error, trying to fetch it from server...`);
                 try {
                     const urlToFetch = mcpServerConfig.httpUrl || mcpServerConfig.url;
                     const response = await fetch(urlToFetch, {
@@ -637,21 +726,21 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                     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 tokenStorage = new MCPOAuthTokenStorage();
@@ -666,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}'`);
                 }
             }
@@ -713,11 +792,11 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                             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}`);
                         }
                     }
@@ -725,80 +804,63 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                         `Please authenticate using: /mcp auth ${mcpServerName}`);
                 }
                 // For SSE/HTTP servers, try to discover OAuth configuration from the base URL
-                console.log(`🔍 Attempting OAuth discovery for '${mcpServerName}'...`);
+                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 {
-                        // 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 authServerUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
-                            console.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) {
-                                        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}'`);
-                                    }
+                    // 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.`);
                 }
             }
@@ -825,99 +887,57 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
 }
 /** Visible for Testing */
 export async function createTransport(mcpServerName, mcpServerConfig, debugMode) {
-    if (mcpServerConfig.authProviderType ===
-        AuthProviderType.SERVICE_ACCOUNT_IMPERSONATION) {
-        const provider = new ServiceAccountImpersonationProvider(mcpServerConfig);
-        const transportOptions = {
-            authProvider: provider,
-        };
-        if (mcpServerConfig.httpUrl) {
-            return new StreamableHTTPClientTransport(new URL(mcpServerConfig.httpUrl), transportOptions);
-        }
-        else if (mcpServerConfig.url) {
-            // Default to SSE if only url is provided
-            return new SSEClientTransport(new URL(mcpServerConfig.url), transportOptions);
+    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}`;
+            }
         }
-        throw new Error('No URL configured for ServiceAccountImpersonation MCP Server');
-    }
-    if (mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS) {
-        const provider = new GoogleCredentialProvider(mcpServerConfig);
         const transportOptions = {
-            authProvider: provider,
+            requestInit: createTransportRequestInit(mcpServerConfig, headers),
+            authProvider,
         };
-        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) {
-        const tokenStorage = new MCPOAuthTokenStorage();
-        const authProvider = new MCPOAuthProvider(tokenStorage);
-        accessToken = await authProvider.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 tokenStorage = new MCPOAuthTokenStorage();
-        const credentials = await tokenStorage.getCredentials(mcpServerName);
-        if (credentials) {
-            const authProvider = new MCPOAuthProvider(tokenStorage);
-            accessToken = await authProvider.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 = {};
-        // 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);
+        return createUrlTransport(mcpServerConfig, transportOptions);
     }
     if (mcpServerConfig.command) {
         const transport = new StdioClientTransport({
@@ -933,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;
@@ -943,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;