@machina.ai/cell-cli-core 1.0.17-rc1 → 1.0.21-rc2

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

@@ -0,0 +1,74 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { McpClient, MCPDiscoveryState, populateMcpServerCommand, } from './mcp-client.js';
+import { getErrorMessage } from '../utils/errors.js';
+/**
+ * Manages the lifecycle of multiple MCP clients, including local child processes.
+ * This class is responsible for starting, stopping, and discovering tools from
+ * a collection of MCP servers defined in the configuration.
+ */
+export class McpClientManager {
+    clients = new Map();
+    mcpServers;
+    mcpServerCommand;
+    toolRegistry;
+    promptRegistry;
+    debugMode;
+    workspaceContext;
+    discoveryState = MCPDiscoveryState.NOT_STARTED;
+    constructor(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext) {
+        this.mcpServers = mcpServers;
+        this.mcpServerCommand = mcpServerCommand;
+        this.toolRegistry = toolRegistry;
+        this.promptRegistry = promptRegistry;
+        this.debugMode = debugMode;
+        this.workspaceContext = workspaceContext;
+    }
+    /**
+     * 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() {
+        await this.stop();
+        this.discoveryState = MCPDiscoveryState.IN_PROGRESS;
+        const servers = populateMcpServerCommand(this.mcpServers, this.mcpServerCommand);
+        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);
+            try {
+                await client.connect();
+                await client.discover();
+            }
+            catch (error) {
+                // Log the error but don't let a single failed server stop the others
+                console.error(`Error during discovery for server '${name}': ${getErrorMessage(error)}`);
+            }
+        });
+        await Promise.all(discoveryPromises);
+        this.discoveryState = MCPDiscoveryState.COMPLETED;
+    }
+    /**
+     * Stops all running local MCP servers and closes all client connections.
+     * This is the cleanup method to be called on application exit.
+     */
+    async stop() {
+        const disconnectionPromises = Array.from(this.clients.entries()).map(async ([name, client]) => {
+            try {
+                await client.disconnect();
+            }
+            catch (error) {
+                console.error(`Error stopping client '${name}': ${getErrorMessage(error)}`);
+            }
+        });
+        await Promise.all(disconnectionPromises);
+        this.clients.clear();
+    }
+    getDiscoveryState() {
+        return this.discoveryState;
+    }
+}
+//# sourceMappingURL=mcp-client-manager.js.map

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

@@ -0,0 +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"}

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

@@ -0,0 +1,6 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+export {};

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

@@ -0,0 +1,39 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import { McpClientManager } from './mcp-client-manager.js';
+import { McpClient } from './mcp-client.js';
+vi.mock('./mcp-client.js', async () => {
+    const originalModule = await vi.importActual('./mcp-client.js');
+    return {
+        ...originalModule,
+        McpClient: vi.fn(),
+        populateMcpServerCommand: vi.fn(() => ({
+            'test-server': {},
+        })),
+    };
+});
+describe('McpClientManager', () => {
+    afterEach(() => {
+        vi.restoreAllMocks();
+    });
+    it('should discover tools from all servers', 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();
+        expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
+        expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
+    });
+});
+//# sourceMappingURL=mcp-client-manager.test.js.map

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

@@ -0,0 +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"}

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

@@ -11,6 +11,7 @@ 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';
 export declare const MCP_DEFAULT_TIMEOUT_MSEC: number;
 export type DiscoveredMCPPrompt = Prompt & {
     serverName: string;
@@ -38,6 +39,45 @@ export declare enum MCPDiscoveryState {
     /** Discovery has completed (with or without errors) */
     COMPLETED = "completed"
 }
+/**
+ * A client for a single MCP server.
+ *
+ * This class is responsible for connecting to, discovering tools from, and
+ * managing the state of a single MCP server.
+ */
+export declare class McpClient {
+    private readonly serverName;
+    private readonly serverConfig;
+    private readonly toolRegistry;
+    private readonly promptRegistry;
+    private readonly workspaceContext;
+    private readonly debugMode;
+    private client;
+    private transport;
+    private status;
+    private isDisconnecting;
+    constructor(serverName: string, serverConfig: MCPServerConfig, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, workspaceContext: WorkspaceContext, debugMode: boolean);
+    /**
+     * Connects to the MCP server.
+     */
+    connect(): Promise<void>;
+    /**
+     * Discovers tools and prompts from the MCP server.
+     */
+    discover(): Promise<void>;
+    /**
+     * Disconnects from the MCP server.
+     */
+    disconnect(): Promise<void>;
+    /**
+     * Returns the current status of the client.
+     */
+    getStatus(): MCPServerStatus;
+    private updateStatus;
+    private createTransport;
+    private discoverTools;
+    private discoverPrompts;
+}
 /**
  * Map to track which MCP servers have been discovered to require OAuth
  */
@@ -54,6 +94,10 @@ export declare function addMCPStatusChangeListener(listener: StatusChangeListene
  * Remove a listener for MCP server status changes
  */
 export declare function removeMCPStatusChangeListener(listener: StatusChangeListener): void;
+/**
+ * Update the status of an MCP server
+ */
+export declare function updateMCPServerStatus(serverName: string, status: MCPServerStatus): void;
 /**
  * Get the current status of an MCP server
  */
@@ -76,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): Promise<void>;
+export declare function discoverMcpTools(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined, toolRegistry: ToolRegistry, promptRegistry: PromptRegistry, debugMode: boolean, workspaceContext: WorkspaceContext): Promise<void>;
 /** Visible for Testing */
 export declare function populateMcpServerCommand(mcpServers: Record<string, MCPServerConfig>, mcpServerCommand: string | undefined): Record<string, MCPServerConfig>;
 /**
@@ -89,7 +133,17 @@ 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): Promise<void>;
+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;
 /**
  * Discovers and sanitizes tools from a connected MCP client.
  * It retrieves function declarations from the client, filters out disabled tools,
@@ -130,7 +184,7 @@ export declare function invokeMcpPrompt(mcpServerName: string, mcpClient: Client
  * @returns A promise that resolves to a connected MCP `Client` instance.
  * @throws An error if the connection fails or the configuration is invalid.
  */
-export declare function connectToMcpServer(mcpServerName: string, mcpServerConfig: MCPServerConfig, debugMode: boolean): Promise<Client>;
+export declare function connectToMcpServer(mcpServerName: string, mcpServerConfig: MCPServerConfig, debugMode: boolean, workspaceContext: WorkspaceContext): Promise<Client>;
 /** Visible for Testing */
 export declare function createTransport(mcpServerName: string, mcpServerConfig: MCPServerConfig, debugMode: boolean): Promise<Transport>;
 /** Visible for testing */

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

@@ -7,7 +7,7 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
 import { SSEClientTransport, } from '@modelcontextprotocol/sdk/client/sse.js';
 import { StreamableHTTPClientTransport, } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { ListPromptsResultSchema, GetPromptResultSchema, } from '@modelcontextprotocol/sdk/types.js';
+import { ListPromptsResultSchema, GetPromptResultSchema, ListRootsRequestSchema, } from '@modelcontextprotocol/sdk/types.js';
 import { parse } from 'shell-quote';
 import { AuthProviderType } from '../config/config.js';
 import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js';
@@ -17,6 +17,8 @@ import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
 import { OAuthUtils } from '../mcp/oauth-utils.js';
 import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { getErrorMessage } from '../utils/errors.js';
+import { basename } from 'node:path';
+import { pathToFileURL } from 'node:url';
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 /**
  * Enum representing the connection status of an MCP server
@@ -42,6 +44,122 @@ export var MCPDiscoveryState;
     /** Discovery has completed (with or without errors) */
     MCPDiscoveryState["COMPLETED"] = "completed";
 })(MCPDiscoveryState || (MCPDiscoveryState = {}));
+/**
+ * A client for a single MCP server.
+ *
+ * This class is responsible for connecting to, discovering tools from, and
+ * managing the state of a single MCP server.
+ */
+export class McpClient {
+    serverName;
+    serverConfig;
+    toolRegistry;
+    promptRegistry;
+    workspaceContext;
+    debugMode;
+    client;
+    transport;
+    status = MCPServerStatus.DISCONNECTED;
+    isDisconnecting = false;
+    constructor(serverName, serverConfig, toolRegistry, promptRegistry, workspaceContext, debugMode) {
+        this.serverName = serverName;
+        this.serverConfig = serverConfig;
+        this.toolRegistry = toolRegistry;
+        this.promptRegistry = promptRegistry;
+        this.workspaceContext = workspaceContext;
+        this.debugMode = debugMode;
+        this.client = new Client({
+            name: `gemini-cli-mcp-client-${this.serverName}`,
+            version: '0.0.1',
+        });
+    }
+    /**
+     * Connects to the MCP server.
+     */
+    async connect() {
+        this.isDisconnecting = false;
+        this.updateStatus(MCPServerStatus.CONNECTING);
+        try {
+            this.transport = await this.createTransport();
+            this.client.onerror = (error) => {
+                if (this.isDisconnecting) {
+                    return;
+                }
+                console.error(`MCP ERROR (${this.serverName}):`, error.toString());
+                this.updateStatus(MCPServerStatus.DISCONNECTED);
+            };
+            this.client.registerCapabilities({
+                roots: {},
+            });
+            this.client.setRequestHandler(ListRootsRequestSchema, async () => {
+                const roots = [];
+                for (const dir of this.workspaceContext.getDirectories()) {
+                    roots.push({
+                        uri: pathToFileURL(dir).toString(),
+                        name: basename(dir),
+                    });
+                }
+                return {
+                    roots,
+                };
+            });
+            await this.client.connect(this.transport, {
+                timeout: this.serverConfig.timeout,
+            });
+            this.updateStatus(MCPServerStatus.CONNECTED);
+        }
+        catch (error) {
+            this.updateStatus(MCPServerStatus.DISCONNECTED);
+            throw error;
+        }
+    }
+    /**
+     * Discovers tools and prompts from the MCP server.
+     */
+    async discover() {
+        if (this.status !== MCPServerStatus.CONNECTED) {
+            throw new Error('Client is not connected.');
+        }
+        const prompts = await this.discoverPrompts();
+        const tools = await this.discoverTools();
+        if (prompts.length === 0 && tools.length === 0) {
+            throw new Error('No prompts or tools found on the server.');
+        }
+        for (const tool of tools) {
+            this.toolRegistry.registerTool(tool);
+        }
+    }
+    /**
+     * Disconnects from the MCP server.
+     */
+    async disconnect() {
+        this.isDisconnecting = true;
+        if (this.transport) {
+            await this.transport.close();
+        }
+        this.client.close();
+        this.updateStatus(MCPServerStatus.DISCONNECTED);
+    }
+    /**
+     * Returns the current status of the client.
+     */
+    getStatus() {
+        return this.status;
+    }
+    updateStatus(status) {
+        this.status = status;
+        updateMCPServerStatus(this.serverName, status);
+    }
+    async createTransport() {
+        return createTransport(this.serverName, this.serverConfig, this.debugMode);
+    }
+    async discoverTools() {
+        return discoverTools(this.serverName, this.serverConfig, this.client);
+    }
+    async discoverPrompts() {
+        return discoverPrompts(this.serverName, this.client, this.promptRegistry);
+    }
+}
 /**
  * Map to track the status of each MCP server within the core package
  */
@@ -73,7 +191,7 @@ export function removeMCPStatusChangeListener(listener) {
 /**
  * Update the status of an MCP server
  */
-function updateMCPServerStatus(serverName, status) {
+export function updateMCPServerStatus(serverName, status) {
     serverStatuses.set(serverName, status);
     // Notify all listeners
     for (const listener of statusChangeListeners) {
@@ -163,8 +281,10 @@ async function handleAutomaticOAuth(mcpServerName, mcpServerConfig, wwwAuthentic
             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);
+        await MCPOAuthProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
         console.log(`OAuth authentication successful for server '${mcpServerName}'`);
         return true;
     }
@@ -223,11 +343,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) {
+export async function discoverMcpTools(mcpServers, mcpServerCommand, toolRegistry, promptRegistry, debugMode, workspaceContext) {
     mcpDiscoveryState = MCPDiscoveryState.IN_PROGRESS;
     try {
         mcpServers = populateMcpServerCommand(mcpServers, mcpServerCommand);
-        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode));
+        const discoveryPromises = Object.entries(mcpServers).map(([mcpServerName, mcpServerConfig]) => connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext));
         await Promise.all(discoveryPromises);
     }
     finally {
@@ -260,11 +380,11 @@ 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) {
+export async function connectAndDiscover(mcpServerName, mcpServerConfig, toolRegistry, promptRegistry, debugMode, workspaceContext) {
     updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTING);
     let mcpClient;
     try {
-        mcpClient = await connectToMcpServer(mcpServerName, mcpServerConfig, debugMode);
+        mcpClient = await connectToMcpServer(mcpServerName, mcpServerConfig, debugMode, workspaceContext);
         mcpClient.onerror = (error) => {
             console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
             updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
@@ -291,6 +411,59 @@ 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,
@@ -316,6 +489,12 @@ 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));
             }
             catch (error) {
@@ -341,6 +520,9 @@ export async function discoverTools(mcpServerName, mcpServerConfig, mcpClient) {
  */
 export async function discoverPrompts(mcpServerName, mcpClient, promptRegistry) {
     try {
+        // Only request prompts if the server supports them.
+        if (mcpClient.getServerCapabilities()?.prompts == null)
+            return [];
         const response = await mcpClient.request({ method: 'prompts/list', params: {} }, ListPromptsResultSchema);
         for (const prompt of response.prompts) {
             promptRegistry.registerPrompt({
@@ -399,11 +581,52 @@ export async function invokeMcpPrompt(mcpServerName, mcpClient, promptName, prom
  * @returns A promise that resolves to a connected MCP `Client` instance.
  * @throws An error if the connection fails or the configuration is invalid.
  */
-export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMode) {
+export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMode, workspaceContext) {
     const mcpClient = new Client({
         name: 'gemini-cli-mcp-client',
         version: '0.0.1',
     });
+    mcpClient.registerCapabilities({
+        roots: {
+            listChanged: true,
+        },
+    });
+    mcpClient.setRequestHandler(ListRootsRequestSchema, async () => {
+        const roots = [];
+        for (const dir of workspaceContext.getDirectories()) {
+            roots.push({
+                uri: pathToFileURL(dir).toString(),
+                name: basename(dir),
+            });
+        }
+        return {
+            roots,
+        };
+    });
+    let unlistenDirectories = workspaceContext.onDirectoriesChanged(async () => {
+        try {
+            await mcpClient.notification({
+                method: 'notifications/roots/list_changed',
+            });
+        }
+        catch (_) {
+            // If this fails, its almost certainly because the connection was closed
+            // and we should just stop listening for future directory changes.
+            unlistenDirectories?.();
+            unlistenDirectories = undefined;
+        }
+    });
+    // Attempt to pro-actively unsubscribe if the mcp client closes. This API is
+    // very brittle though so we don't have any guarantees, hence the try/catch
+    // above as well.
+    //
+    // Be a good steward and don't just bash over onclose.
+    const oldOnClose = mcpClient.onclose;
+    mcpClient.onclose = () => {
+        oldOnClose?.();
+        unlistenDirectories?.();
+        unlistenDirectories = undefined;
+    };
     // patch Client.callTool to use request timeout as genai McpCallTool.callTool does not do it
     // TODO: remove this hack once GenAI SDK does callTool with request options
     if ('callTool' in mcpClient) {
@@ -574,8 +797,10 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
                                 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);
+                            await MCPOAuthProvider.authenticate(mcpServerName, oauthAuthConfig, serverUrl);
                             // Retry connection with OAuth token
                             const credentials = await MCPOAuthTokenStorage.getToken(mcpServerName);
                             if (credentials) {
@@ -643,7 +868,7 @@ export async function connectToMcpServer(mcpServerName, mcpServerConfig, debugMo
             else {
                 conciseError = `Connection failed for '${mcpServerName}': ${errorMessage}`;
             }
-            if (process.env.SANDBOX) {
+            if (process.env['SANDBOX']) {
                 conciseError += ` (check sandbox availability)`;
             }
             throw new Error(conciseError);