# 🚀 Sync Gemini CLI v0.2.1 - Major Feature Update (#483)

packages/core/src/tools/mcp-client.ts

@@ -36,7 +36,7 @@ import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
 import { getErrorMessage } from '../utils/errors.js';
 import { basename } from 'node:path';
 import { pathToFileURL } from 'node:url';
-import { WorkspaceContext } from '../utils/workspaceContext.js';
+import { Unsubscribe, WorkspaceContext } from '../utils/workspaceContext.js';
 
 export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
 
@@ -69,6 +69,134 @@ export enum MCPDiscoveryState {
   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 class McpClient {
+  private client: Client;
+  private transport: Transport | undefined;
+  private status: MCPServerStatus = MCPServerStatus.DISCONNECTED;
+  private isDisconnecting = false;
+
+  constructor(
+    private readonly serverName: string,
+    private readonly serverConfig: MCPServerConfig,
+    private readonly toolRegistry: ToolRegistry,
+    private readonly promptRegistry: PromptRegistry,
+    private readonly workspaceContext: WorkspaceContext,
+    private readonly debugMode: boolean,
+  ) {
+    this.client = new Client({
+      name: `gemini-cli-mcp-client-${this.serverName}`,
+      version: '0.0.1',
+    });
+  }
+
+  /**
+   * Connects to the MCP server.
+   */
+  async connect(): Promise<void> {
+    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(): Promise<void> {
+    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(): Promise<void> {
+    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(): MCPServerStatus {
+    return this.status;
+  }
+
+  private updateStatus(status: MCPServerStatus): void {
+    this.status = status;
+    updateMCPServerStatus(this.serverName, status);
+  }
+
+  private async createTransport(): Promise<Transport> {
+    return createTransport(this.serverName, this.serverConfig, this.debugMode);
+  }
+
+  private async discoverTools(): Promise<DiscoveredMCPTool[]> {
+    return discoverTools(this.serverName, this.serverConfig, this.client);
+  }
+
+  private async discoverPrompts(): Promise<Prompt[]> {
+    return discoverPrompts(this.serverName, this.client, this.promptRegistry);
+  }
+}
+
 /**
  * Map to track the status of each MCP server within the core package
  */
@@ -117,7 +245,7 @@ export function removeMCPStatusChangeListener(
 /**
  * Update the status of an MCP server
  */
-function updateMCPServerStatus(
+export function updateMCPServerStatus(
   serverName: string,
   status: MCPServerStatus,
 ): void {
@@ -227,10 +355,16 @@ async function handleAutomaticOAuth(
     };
 
     // 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}'`,
@@ -442,7 +576,7 @@ export function hasValidTypes(schema: unknown): boolean {
 
   const s = schema as Record<string, unknown>;
 
-  if (!s.type) {
+  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.
@@ -464,9 +598,9 @@ export function hasValidTypes(schema: unknown): boolean {
     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 (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;
         }
@@ -474,8 +608,8 @@ export function hasValidTypes(schema: unknown): boolean {
     }
   }
 
-  if (s.type === 'array' && s.items) {
-    if (!hasValidTypes(s.items)) {
+  if (s['type'] === 'array' && s['items']) {
+    if (!hasValidTypes(s['items'])) {
       return false;
     }
   }
@@ -671,7 +805,9 @@ export async function connectToMcpServer(
   });
 
   mcpClient.registerCapabilities({
-    roots: {},
+    roots: {
+      listChanged: true,
+    },
   });
 
   mcpClient.setRequestHandler(ListRootsRequestSchema, async () => {
@@ -687,6 +823,32 @@ export async function connectToMcpServer(
     };
   });
 
+  let unlistenDirectories: Unsubscribe | undefined =
+    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) {
@@ -933,12 +1095,15 @@ export async function connectToMcpServer(
               };
 
               // Perform OAuth authentication
+              // Pass the server URL for proper discovery
+              const serverUrl = mcpServerConfig.httpUrl || mcpServerConfig.url;
               console.log(
                 `Starting OAuth authentication for server '${mcpServerName}'...`,
               );
               await MCPOAuthProvider.authenticate(
                 mcpServerName,
                 oauthAuthConfig,
+                serverUrl,
               );
 
               // Retry connection with OAuth token
@@ -1037,7 +1202,7 @@ export async function connectToMcpServer(
         conciseError = `Connection failed for '${mcpServerName}': ${errorMessage}`;
       }
 
-      if (process.env.SANDBOX) {
+      if (process.env['SANDBOX']) {
         conciseError += ` (check sandbox availability)`;
       }