@office-ai/aioncli-core 0.2.2 → 0.8.1

package/dist/src/tools/mcp-client-manager.d.ts

@@ -3,11 +3,12 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { MCPServerConfig } from '../config/config.js';
-import { ToolRegistry } from './tool-registry.js';
-import { PromptRegistry } from '../prompts/prompt-registry.js';
+import type { Config, MCPServerConfig } from '../config/config.js';
+import type { ToolRegistry } from './tool-registry.js';
+import type { PromptRegistry } from '../prompts/prompt-registry.js';
 import { MCPDiscoveryState } from './mcp-client.js';
-import { WorkspaceContext } from '../utils/workspaceContext.js';
+import type { EventEmitter } from 'node:events';
+import type { WorkspaceContext } from '../utils/workspaceContext.js';
 /**
  * Manages the lifecycle of multiple MCP clients, including local child processes.
  * This class is responsible for starting, stopping, and discovering tools from
@@ -22,13 +23,14 @@ export declare class McpClientManager {
     private readonly debugMode;
     private readonly workspaceContext;
     private discoveryState;
-    constructor(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext);
+    private readonly eventEmitter?;
+    constructor(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext, eventEmitter?: EventEmitter);
     /**
      * Initiates the tool discovery process for all configured MCP servers.
      * It connects to each server, discovers its available tools, and registers
      * them with the `ToolRegistry`.
      */
-    discoverAllMcpTools(): Promise<void>;
+    discoverAllMcpTools(cliConfig: Config): Promise<void>;
     /**
      * Stops all running local MCP servers and closes all client connections.
      * This is the cleanup method to be called on application exit.

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

@@ -19,31 +19,40 @@ export class McpClientManager {
     debugMode;
     workspaceContext;
     discoveryState = MCPDiscoveryState.NOT_STARTED;
-    constructor(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext) {
+    eventEmitter;
+    constructor(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext, eventEmitter) {
         this.mcpServers = mcpServers;
         this.mcpServerCommand = mcpServerCommand;
         this.toolRegistry = toolRegistry;
         this.promptRegistry = promptRegistry;
         this.debugMode = debugMode;
         this.workspaceContext = workspaceContext;
+        this.eventEmitter = eventEmitter;
     }
     /**
      * Initiates the tool discovery process for all configured MCP servers.
      * It connects to each server, discovers its available tools, and registers
      * them with the `ToolRegistry`.
      */
-    async discoverAllMcpTools() {
+    async discoverAllMcpTools(cliConfig) {
+        if (!cliConfig.isTrustedFolder()) {
+            return;
+        }
         await this.stop();
-        this.discoveryState = MCPDiscoveryState.IN_PROGRESS;
         const servers = populateMcpServerCommand(this.mcpServers, this.mcpServerCommand);
+        this.discoveryState = MCPDiscoveryState.IN_PROGRESS;
+        this.eventEmitter?.emit('mcp-client-update', this.clients);
         const discoveryPromises = Object.entries(servers).map(async ([name, config]) => {
             const client = new McpClient(name, config, this.toolRegistry, this.promptRegistry, this.workspaceContext, this.debugMode);
             this.clients.set(name, client);
+            this.eventEmitter?.emit('mcp-client-update', this.clients);
             try {
                 await client.connect();
-                await client.discover();
+                await client.discover(cliConfig);
+                this.eventEmitter?.emit('mcp-client-update', this.clients);
             }
             catch (error) {
+                this.eventEmitter?.emit('mcp-client-update', this.clients);
                 // Log the error but don't let a single failed server stop the others
                 console.error(`Error during discovery for server '${name}': ${getErrorMessage(error)}`);
             }

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

@@ -1 +1 @@
-{"version":3,"file":"mcp-client-manager.js","sourceRoot":"","sources":["../../../src/tools/mcp-client-manager.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAKH,OAAO,EACL,SAAS,EACT,iBAAiB,EACjB,wBAAwB,GACzB,MAAM,iBAAiB,CAAC;AACzB,OAAO,EAAE,eAAe,EAAE,MAAM,oBAAoB,CAAC;AAGrD;;;;GAIG;AACH,MAAM,OAAO,gBAAgB;IACnB,OAAO,GAA2B,IAAI,GAAG,EAAE,CAAC;IACnC,UAAU,CAAkC;IAC5C,gBAAgB,CAAqB;IACrC,YAAY,CAAe;IAC3B,cAAc,CAAiB;IAC/B,SAAS,CAAU;IACnB,gBAAgB,CAAmB;IAC5C,cAAc,GAAsB,iBAAiB,CAAC,WAAW,CAAC;IAE1E,YACE,UAA2C,EAC3C,gBAAoC,EACpC,YAA0B,EAC1B,cAA8B,EAC9B,SAAkB,EAClB,gBAAkC;QAElC,IAAI,CAAC,UAAU,GAAG,UAAU,CAAC;QAC7B,IAAI,CAAC,gBAAgB,GAAG,gBAAgB,CAAC;QACzC,IAAI,CAAC,YAAY,GAAG,YAAY,CAAC;QACjC,IAAI,CAAC,cAAc,GAAG,cAAc,CAAC;QACrC,IAAI,CAAC,SAAS,GAAG,SAAS,CAAC;QAC3B,IAAI,CAAC,gBAAgB,GAAG,gBAAgB,CAAC;IAC3C,CAAC;IAED;;;;OAIG;IACH,KAAK,CAAC,mBAAmB;QACvB,MAAM,IAAI,CAAC,IAAI,EAAE,CAAC;QAClB,IAAI,CAAC,cAAc,GAAG,iBAAiB,CAAC,WAAW,CAAC;QACpD,MAAM,OAAO,GAAG,wBAAwB,CACtC,IAAI,CAAC,UAAU,EACf,IAAI,CAAC,gBAAgB,CACtB,CAAC;QAEF,MAAM,iBAAiB,GAAG,MAAM,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,GAAG,CACnD,KAAK,EAAE,CAAC,IAAI,EAAE,MAAM,CAAC,EAAE,EAAE;YACvB,MAAM,MAAM,GAAG,IAAI,SAAS,CAC1B,IAAI,EACJ,MAAM,EACN,IAAI,CAAC,YAAY,EACjB,IAAI,CAAC,cAAc,EACnB,IAAI,CAAC,gBAAgB,EACrB,IAAI,CAAC,SAAS,CACf,CAAC;YACF,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC;YAC/B,IAAI,CAAC;gBACH,MAAM,MAAM,CAAC,OAAO,EAAE,CAAC;gBACvB,MAAM,MAAM,CAAC,QAAQ,EAAE,CAAC;YAC1B,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,qEAAqE;gBACrE,OAAO,CAAC,KAAK,CACX,sCAAsC,IAAI,MAAM,eAAe,CAC7D,KAAK,CACN,EAAE,CACJ,CAAC;YACJ,CAAC;QACH,CAAC,CACF,CAAC;QAEF,MAAM,OAAO,CAAC,GAAG,CAAC,iBAAiB,CAAC,CAAC;QACrC,IAAI,CAAC,cAAc,GAAG,iBAAiB,CAAC,SAAS,CAAC;IACpD,CAAC;IAED;;;OAGG;IACH,KAAK,CAAC,IAAI;QACR,MAAM,qBAAqB,GAAG,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC,CAAC,GAAG,CAClE,KAAK,EAAE,CAAC,IAAI,EAAE,MAAM,CAAC,EAAE,EAAE;YACvB,IAAI,CAAC;gBACH,MAAM,MAAM,CAAC,UAAU,EAAE,CAAC;YAC5B,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO,CAAC,KAAK,CACX,0BAA0B,IAAI,MAAM,eAAe,CAAC,KAAK,CAAC,EAAE,CAC7D,CAAC;YACJ,CAAC;QACH,CAAC,CACF,CAAC;QAEF,MAAM,OAAO,CAAC,GAAG,CAAC,qBAAqB,CAAC,CAAC;QACzC,IAAI,CAAC,OAAO,CAAC,KAAK,EAAE,CAAC;IACvB,CAAC;IAED,iBAAiB;QACf,OAAO,IAAI,CAAC,cAAc,CAAC;IAC7B,CAAC;CACF"}
+{"version":3,"file":"mcp-client-manager.js","sourceRoot":"","sources":["../../../src/tools/mcp-client-manager.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAKH,OAAO,EACL,SAAS,EACT,iBAAiB,EACjB,wBAAwB,GACzB,MAAM,iBAAiB,CAAC;AACzB,OAAO,EAAE,eAAe,EAAE,MAAM,oBAAoB,CAAC;AAIrD;;;;GAIG;AACH,MAAM,OAAO,gBAAgB;IACnB,OAAO,GAA2B,IAAI,GAAG,EAAE,CAAC;IACnC,UAAU,CAAkC;IAC5C,gBAAgB,CAAqB;IACrC,YAAY,CAAe;IAC3B,cAAc,CAAiB;IAC/B,SAAS,CAAU;IACnB,gBAAgB,CAAmB;IAC5C,cAAc,GAAsB,iBAAiB,CAAC,WAAW,CAAC;IACzD,YAAY,CAAgB;IAE7C,YACE,UAA2C,EAC3C,gBAAoC,EACpC,YAA0B,EAC1B,cAA8B,EAC9B,SAAkB,EAClB,gBAAkC,EAClC,YAA2B;QAE3B,IAAI,CAAC,UAAU,GAAG,UAAU,CAAC;QAC7B,IAAI,CAAC,gBAAgB,GAAG,gBAAgB,CAAC;QACzC,IAAI,CAAC,YAAY,GAAG,YAAY,CAAC;QACjC,IAAI,CAAC,cAAc,GAAG,cAAc,CAAC;QACrC,IAAI,CAAC,SAAS,GAAG,SAAS,CAAC;QAC3B,IAAI,CAAC,gBAAgB,GAAG,gBAAgB,CAAC;QACzC,IAAI,CAAC,YAAY,GAAG,YAAY,CAAC;IACnC,CAAC;IAED;;;;OAIG;IACH,KAAK,CAAC,mBAAmB,CAAC,SAAiB;QACzC,IAAI,CAAC,SAAS,CAAC,eAAe,EAAE,EAAE,CAAC;YACjC,OAAO;QACT,CAAC;QACD,MAAM,IAAI,CAAC,IAAI,EAAE,CAAC;QAElB,MAAM,OAAO,GAAG,wBAAwB,CACtC,IAAI,CAAC,UAAU,EACf,IAAI,CAAC,gBAAgB,CACtB,CAAC;QAEF,IAAI,CAAC,cAAc,GAAG,iBAAiB,CAAC,WAAW,CAAC;QAEpD,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,mBAAmB,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;QAC3D,MAAM,iBAAiB,GAAG,MAAM,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,GAAG,CACnD,KAAK,EAAE,CAAC,IAAI,EAAE,MAAM,CAAC,EAAE,EAAE;YACvB,MAAM,MAAM,GAAG,IAAI,SAAS,CAC1B,IAAI,EACJ,MAAM,EACN,IAAI,CAAC,YAAY,EACjB,IAAI,CAAC,cAAc,EACnB,IAAI,CAAC,gBAAgB,EACrB,IAAI,CAAC,SAAS,CACf,CAAC;YACF,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC;YAE/B,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,mBAAmB,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;YAC3D,IAAI,CAAC;gBACH,MAAM,MAAM,CAAC,OAAO,EAAE,CAAC;gBACvB,MAAM,MAAM,CAAC,QAAQ,CAAC,SAAS,CAAC,CAAC;gBACjC,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,mBAAmB,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;YAC7D,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,mBAAmB,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;gBAC3D,qEAAqE;gBACrE,OAAO,CAAC,KAAK,CACX,sCAAsC,IAAI,MAAM,eAAe,CAC7D,KAAK,CACN,EAAE,CACJ,CAAC;YACJ,CAAC;QACH,CAAC,CACF,CAAC;QAEF,MAAM,OAAO,CAAC,GAAG,CAAC,iBAAiB,CAAC,CAAC;QACrC,IAAI,CAAC,cAAc,GAAG,iBAAiB,CAAC,SAAS,CAAC;IACpD,CAAC;IAED;;;OAGG;IACH,KAAK,CAAC,IAAI;QACR,MAAM,qBAAqB,GAAG,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC,CAAC,GAAG,CAClE,KAAK,EAAE,CAAC,IAAI,EAAE,MAAM,CAAC,EAAE,EAAE;YACvB,IAAI,CAAC;gBACH,MAAM,MAAM,CAAC,UAAU,EAAE,CAAC;YAC5B,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO,CAAC,KAAK,CACX,0BAA0B,IAAI,MAAM,eAAe,CAAC,KAAK,CAAC,EAAE,CAC7D,CAAC;YACJ,CAAC;QACH,CAAC,CACF,CAAC;QAEF,MAAM,OAAO,CAAC,GAAG,CAAC,qBAAqB,CAAC,CAAC;QACzC,IAAI,CAAC,OAAO,CAAC,KAAK,EAAE,CAAC;IACvB,CAAC;IAED,iBAAiB;QACf,OAAO,IAAI,CAAC,cAAc,CAAC;IAC7B,CAAC;CACF"}

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

@@ -31,9 +31,28 @@ describe('McpClientManager', () => {
         const manager = new McpClientManager({
             'test-server': {},
         }, '', {}, {}, false, {});
-        await manager.discoverAllMcpTools();
+        await manager.discoverAllMcpTools({
+            isTrustedFolder: () => true,
+        });
         expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
         expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
     });
+    it('should not discover tools if folder is not trusted', async () => {
+        const mockedMcpClient = {
+            connect: vi.fn(),
+            discover: vi.fn(),
+            disconnect: vi.fn(),
+            getStatus: vi.fn(),
+        };
+        vi.mocked(McpClient).mockReturnValue(mockedMcpClient);
+        const manager = new McpClientManager({
+            'test-server': {},
+        }, '', {}, {}, false, {});
+        await manager.discoverAllMcpTools({
+            isTrustedFolder: () => false,
+        });
+        expect(mockedMcpClient.connect).not.toHaveBeenCalled();
+        expect(mockedMcpClient.discover).not.toHaveBeenCalled();
+    });
 });
 //# sourceMappingURL=mcp-client-manager.test.js.map

package/dist/src/tools/mcp-client-manager.test.js.map

@@ -1 +1 @@
-{"version":3,"file":"mcp-client-manager.test.js","sourceRoot":"","sources":["../../../src/tools/mcp-client-manager.test.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAEH,OAAO,EAAE,SAAS,EAAE,QAAQ,EAAE,MAAM,EAAE,EAAE,EAAE,EAAE,EAAE,MAAM,QAAQ,CAAC;AAC7D,OAAO,EAAE,gBAAgB,EAAE,MAAM,yBAAyB,CAAC;AAC3D,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAK5C,EAAE,CAAC,IAAI,CAAC,iBAAiB,EAAE,KAAK,IAAI,EAAE;IACpC,MAAM,cAAc,GAAG,MAAM,EAAE,CAAC,YAAY,CAAC,iBAAiB,CAAC,CAAC;IAChE,OAAO;QACL,GAAG,cAAc;QACjB,SAAS,EAAE,EAAE,CAAC,EAAE,EAAE;QAClB,wBAAwB,EAAE,EAAE,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,CAAC;YACrC,aAAa,EAAE,EAAE;SAClB,CAAC,CAAC;KACJ,CAAC;AACJ,CAAC,CAAC,CAAC;AAEH,QAAQ,CAAC,kBAAkB,EAAE,GAAG,EAAE;IAChC,SAAS,CAAC,GAAG,EAAE;QACb,EAAE,CAAC,eAAe,EAAE,CAAC;IACvB,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,wCAAwC,EAAE,KAAK,IAAI,EAAE;QACtD,MAAM,eAAe,GAAG;YACtB,OAAO,EAAE,EAAE,CAAC,EAAE,EAAE;YAChB,QAAQ,EAAE,EAAE,CAAC,EAAE,EAAE;YACjB,UAAU,EAAE,EAAE,CAAC,EAAE,EAAE;YACnB,SAAS,EAAE,EAAE,CAAC,EAAE,EAAE;SACnB,CAAC;QACF,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC,eAAe,CAClC,eAAuC,CACxC,CAAC;QACF,MAAM,OAAO,GAAG,IAAI,gBAAgB,CAClC;YACE,aAAa,EAAE,EAAE;SAClB,EACD,EAAE,EACF,EAAkB,EAClB,EAAoB,EACpB,KAAK,EACL,EAAsB,CACvB,CAAC;QACF,MAAM,OAAO,CAAC,mBAAmB,EAAE,CAAC;QACpC,MAAM,CAAC,eAAe,CAAC,OAAO,CAAC,CAAC,oBAAoB,EAAE,CAAC;QACvD,MAAM,CAAC,eAAe,CAAC,QAAQ,CAAC,CAAC,oBAAoB,EAAE,CAAC;IAC1D,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}
+{"version":3,"file":"mcp-client-manager.test.js","sourceRoot":"","sources":["../../../src/tools/mcp-client-manager.test.ts"],"names":[],"mappings":"AAAA;;;;GAIG;AAEH,OAAO,EAAE,SAAS,EAAE,QAAQ,EAAE,MAAM,EAAE,EAAE,EAAE,EAAE,EAAE,MAAM,QAAQ,CAAC;AAC7D,OAAO,EAAE,gBAAgB,EAAE,MAAM,yBAAyB,CAAC;AAC3D,OAAO,EAAE,SAAS,EAAE,MAAM,iBAAiB,CAAC;AAM5C,EAAE,CAAC,IAAI,CAAC,iBAAiB,EAAE,KAAK,IAAI,EAAE;IACpC,MAAM,cAAc,GAAG,MAAM,EAAE,CAAC,YAAY,CAAC,iBAAiB,CAAC,CAAC;IAChE,OAAO;QACL,GAAG,cAAc;QACjB,SAAS,EAAE,EAAE,CAAC,EAAE,EAAE;QAClB,wBAAwB,EAAE,EAAE,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,CAAC;YACrC,aAAa,EAAE,EAAE;SAClB,CAAC,CAAC;KACJ,CAAC;AACJ,CAAC,CAAC,CAAC;AAEH,QAAQ,CAAC,kBAAkB,EAAE,GAAG,EAAE;IAChC,SAAS,CAAC,GAAG,EAAE;QACb,EAAE,CAAC,eAAe,EAAE,CAAC;IACvB,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,wCAAwC,EAAE,KAAK,IAAI,EAAE;QACtD,MAAM,eAAe,GAAG;YACtB,OAAO,EAAE,EAAE,CAAC,EAAE,EAAE;YAChB,QAAQ,EAAE,EAAE,CAAC,EAAE,EAAE;YACjB,UAAU,EAAE,EAAE,CAAC,EAAE,EAAE;YACnB,SAAS,EAAE,EAAE,CAAC,EAAE,EAAE;SACnB,CAAC;QACF,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC,eAAe,CAClC,eAAuC,CACxC,CAAC;QACF,MAAM,OAAO,GAAG,IAAI,gBAAgB,CAClC;YACE,aAAa,EAAE,EAAE;SAClB,EACD,EAAE,EACF,EAAkB,EAClB,EAAoB,EACpB,KAAK,EACL,EAAsB,CACvB,CAAC;QACF,MAAM,OAAO,CAAC,mBAAmB,CAAC;YAChC,eAAe,EAAE,GAAG,EAAE,CAAC,IAAI;SACP,CAAC,CAAC;QACxB,MAAM,CAAC,eAAe,CAAC,OAAO,CAAC,CAAC,oBAAoB,EAAE,CAAC;QACvD,MAAM,CAAC,eAAe,CAAC,QAAQ,CAAC,CAAC,oBAAoB,EAAE,CAAC;IAC1D,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,oDAAoD,EAAE,KAAK,IAAI,EAAE;QAClE,MAAM,eAAe,GAAG;YACtB,OAAO,EAAE,EAAE,CAAC,EAAE,EAAE;YAChB,QAAQ,EAAE,EAAE,CAAC,EAAE,EAAE;YACjB,UAAU,EAAE,EAAE,CAAC,EAAE,EAAE;YACnB,SAAS,EAAE,EAAE,CAAC,EAAE,EAAE;SACnB,CAAC;QACF,EAAE,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC,eAAe,CAClC,eAAuC,CACxC,CAAC;QACF,MAAM,OAAO,GAAG,IAAI,gBAAgB,CAClC;YACE,aAAa,EAAE,EAAE;SAClB,EACD,EAAE,EACF,EAAkB,EAClB,EAAoB,EACpB,KAAK,EACL,EAAsB,CACvB,CAAC;QACF,MAAM,OAAO,CAAC,mBAAmB,CAAC;YAChC,eAAe,EAAE,GAAG,EAAE,CAAC,KAAK;SACR,CAAC,CAAC;QACxB,MAAM,CAAC,eAAe,CAAC,OAAO,CAAC,CAAC,GAAG,CAAC,gBAAgB,EAAE,CAAC;QACvD,MAAM,CAAC,eAAe,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,gBAAgB,EAAE,CAAC;IAC1D,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC"}

package/dist/src/tools/mcp-client.d.ts

@@ -4,14 +4,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
-import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
-import { Prompt, GetPromptResult } from '@modelcontextprotocol/sdk/types.js';
-import { MCPServerConfig } from '../config/config.js';
+import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
+import type { GetPromptResult, Prompt } from '@modelcontextprotocol/sdk/types.js';
+import type { Config, MCPServerConfig } from '../config/config.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
-import { FunctionDeclaration } from '@google/genai';
-import { ToolRegistry } from './tool-registry.js';
-import { PromptRegistry } from '../prompts/prompt-registry.js';
-import { WorkspaceContext } from '../utils/workspaceContext.js';
+import type { FunctionDeclaration } from '@google/genai';
+import type { PromptRegistry } from '../prompts/prompt-registry.js';
+import type { WorkspaceContext } from '../utils/workspaceContext.js';
+import type { ToolRegistry } from './tool-registry.js';
 export declare const MCP_DEFAULT_TIMEOUT_MSEC: number;
 export type DiscoveredMCPPrompt = Prompt & {
     serverName: string;
@@ -64,7 +64,7 @@ export declare class McpClient {
     /**
      * Discovers tools and prompts from the MCP server.
      */
-    discover(): Promise<void>;
+    discover(cliConfig: Config): Promise<void>;
     /**
      * Disconnects from the MCP server.
      */
@@ -120,7 +120,7 @@ export declare function getMCPDiscoveryState(): MCPDiscoveryState;
  * @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 declare function discoverMcpTools(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext): Promise<void>;
+export declare function discoverMcpTools(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext, cliConfig: Config): Promise<void>;
 /** Visible for Testing */
 export declare function populateMcpServerCommand(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined): Record<string, MCPServerConfig>;
 /**
@@ -133,17 +133,7 @@ export declare function populateMcpServerCommand(mcpServers: Record<string, MCPS
  * @param toolRegistry The registry to register discovered tools with
  * @returns Promise that resolves when discovery is complete
  */
-export declare function connectAndDiscover(mcpServerName: string, mcpServerConfig: MCPServerConfig, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext): Promise<void>;
-/**
- * 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 declare function hasValidTypes(schema: unknown): boolean;
+export declare function connectAndDiscover(mcpServerName: string, mcpServerConfig: MCPServerConfig, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext, cliConfig: Config): Promise<void>;
 /**
  * Discovers and sanitizes tools from a connected MCP client.
  * It retrieves function declarations from the client, filters out disabled tools,
@@ -155,7 +145,7 @@ export declare function hasValidTypes(schema: unknown): boolean;
  * @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 declare function discoverTools(mcpServerName: string, mcpServerConfig: MCPServerConfig, mcpClient: Client): Promise<DiscoveredMCPTool[]>;
+export declare function discoverTools(mcpServerName: string, mcpServerConfig: MCPServerConfig, mcpClient: Client, cliConfig: Config): Promise<DiscoveredMCPTool[]>;
 /**
  * Discovers and logs prompts from a connected MCP client.
  * It retrieves prompt declarations from the client and logs their names.
@@ -174,6 +164,13 @@ export declare function discoverPrompts(mcpServerName: string, mcpClient: Client
  * @returns A promise that resolves to the result of the prompt invocation.
  */
 export declare function invokeMcpPrompt(mcpServerName: string, mcpClient: Client, promptName: string, promptParams: Record<string, unknown>): Promise<GetPromptResult>;
+/**
+ * @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 declare function hasNetworkTransport(config: MCPServerConfig): boolean;
 /**
  * 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

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

@@ -4,21 +4,22 @@
  * 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 { GetPromptResultSchema, ListPromptsResultSchema, 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';
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 /**
  * Enum representing the connection status of an MCP server
@@ -116,12 +117,12 @@ export class McpClient {
     /**
      * Discovers tools and prompts from the MCP server.
      */
-    async discover() {
+    async discover(cliConfig) {
         if (this.status !== MCPServerStatus.CONNECTED) {
             throw new Error('Client is not connected.');
         }
         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.');
         }
@@ -153,8 +154,8 @@ export class McpClient {
     async createTransport() {
         return createTransport(this.serverName, this.serverConfig, this.debugMode);
     }
-    async discoverTools() {
-        return discoverTools(this.serverName, this.serverConfig, this.client);
+    async discoverTools(cliConfig) {
+        return discoverTools(this.serverName, this.serverConfig, this.client, cliConfig);
     }
     async discoverPrompts() {
         return discoverPrompts(this.serverName, this.client, this.promptRegistry);
@@ -256,16 +257,10 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
         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) {
@@ -284,7 +279,8 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
         // 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);
+        const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
+        await authProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
         console.log(`OAuth authentication successful for server '${mcpServerName}'`);
         return true;
     }
@@ -343,11 +339,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,7 +376,7 @@ 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 {
@@ -391,7 +387,7 @@ export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolReg
         };
         // 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);
         // 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.');
@@ -411,59 +407,6 @@ export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolReg
         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,
@@ -475,9 +418,11 @@ export function hasValidTypes(schema) {
  * @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) {
     try {
-        const mcpCallableTool = mcpToTool(mcpClient);
+        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
@@ -489,13 +434,7 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
                 if (!isEnabled(funcDecl, mcpServerName, mcpServerConfig)) {
                     continue;
                 }
-                if (!hasValidTypes(funcDecl.parametersJsonSchema)) {
-                    console.warn(`Skipping tool '${funcDecl.name}' from MCP server '${mcpServerName}' ` +
-                        `because it has missing types in its parameter schema. Please file an ` +
-                        `issue with the owner of the MCP server.`);
-                    continue;
-                }
-                discoveredTools.push(new DiscoveredMCPTool(mcpCallableTool, mcpServerName, funcDecl.name, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, mcpServerConfig.trust));
+                discoveredTools.push(new DiscoveredMCPTool(mcpCallableTool, mcpServerName, funcDecl.name, funcDecl.description ?? '', funcDecl.parametersJsonSchema ?? { type: 'object', properties: {} }, mcpServerConfig.trust, undefined, cliConfig));
             }
             catch (error) {
                 console.error(`Error discovering tool: '${funcDecl.name}' from MCP server '${mcpServerName}': ${error.message}`);
@@ -571,6 +510,15 @@ export async function invokeMcpPrompt(mcpServerName, mcpClient, promptName, prom
         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 +575,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,17 +591,18 @@ 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,
                     });
@@ -683,13 +621,16 @@ 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) {
+            if (!wwwAuthenticate && hasNetworkTransport(mcpServerConfig)) {
                 console.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),
                     });
@@ -713,9 +654,11 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                     console.log(`Retrying connection to '${mcpServerName}' with OAuth token...`);
                     // Get the valid token - we need to create a proper OAuth config
                     // The token should already be available from the authentication process
-                    const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+                    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,
                         });
@@ -761,9 +704,11 @@ 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,
                         });
@@ -779,11 +724,11 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                     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
+                // For SSE/HTTP 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}`;
+                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);
@@ -798,13 +743,16 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                             };
                             // Perform OAuth authentication
                             // Pass the server URL for proper discovery
-                            const serverUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
+                            const authServerUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
                             console.log(`Starting OAuth authentication for server '${mcpServerName}'...`);
-                            await MCPOAuthProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
+                            const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
+                            await authProvider.authenticate(mcpServerName, oauthAuthConfig, authServerUrl);
                             // Retry connection with OAuth token
-                            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,
                                 });
@@ -877,6 +825,21 @@ 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);
+        }
+        throw new Error('No URL configured for ServiceAccountImpersonation MCP Server');
+    }
     if (mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS) {
         const provider = new GoogleCredentialProvider(mcpServerConfig);
         const transportOptions = {
@@ -894,7 +857,9 @@ export async function createTransport(mcpServerName, mcpServerConfig, debugMode)
     let accessToken = null;
     let hasOAuthConfig = mcpServerConfig.oauth?.enabled;
     if (hasOAuthConfig && mcpServerConfig.oauth) {
-        accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, 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.`);
@@ -904,9 +869,11 @@ export async function createTransport(mcpServerName, mcpServerConfig, debugMode)
     }
     else {
         // Check if we have stored OAuth tokens for this server (from previous authentication)
-        const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
+        const tokenStorage = new MCPOAuthTokenStorage();
+        const credentials = await tokenStorage.getCredentials(mcpServerName);
         if (credentials) {
-            accessToken = await MCPOAuthProvider.getValidToken(mcpServerName, {
+            const authProvider = new MCPOAuthProvider(tokenStorage);
+            accessToken = await authProvider.getValidToken(mcpServerName, {
                 // Pass client ID if available
                 clientId: credentials.clientId,
             });