fix(core): disconnect extension-backed MCP clients in stopExtension (#26136)

Author: cocosheng-g

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

@@ -17,6 +17,7 @@ import { McpClientManager } from './mcp-client-manager.js';
 import { McpClient, MCPDiscoveryState, MCPServerStatus } from './mcp-client.js';
 import type { ToolRegistry } from './tool-registry.js';
 import type { Config, GeminiCLIExtension } from '../config/config.js';
+import { MCPServerConfig } from '../config/config.js';
 import type { PromptRegistry } from '../prompts/prompt-registry.js';
 import type { ResourceRegistry } from '../resources/resource-registry.js';
 
@@ -726,6 +727,40 @@ describe('McpClientManager', () => {
         extensionName: 'test-extension',
       });
     });
+
+    it('should disconnect extension-backed MCP clients when stopping extension (#24050)', async () => {
+      const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
+      const extension: GeminiCLIExtension = {
+        id: 'test-ext-id',
+        name: 'test-extension',
+        isActive: true,
+        version: '1.0.0',
+        path: '/fake/path',
+        contextFiles: [],
+        mcpServers: {
+          'test-server': new MCPServerConfig('node', ['script.js']),
+        },
+      };
+
+      await manager.startExtension(extension);
+
+      // Wait for discovery to complete
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      while ((manager as any).discoveryPromise) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        await (manager as any).discoveryPromise;
+      }
+
+      // Verify it was connected
+      expect(mockedMcpClient.connect).toHaveBeenCalled();
+
+      // Stop the extension
+      await manager.stopExtension(extension);
+
+      // Verify disconnect was called on the client
+      expect(mockedMcpClient.disconnect).toHaveBeenCalled();
+      expect(manager.getClient('test-server')).toBeUndefined();
+    });
   });
 
   describe('diagnostic reporting', () => {

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

@@ -215,6 +215,7 @@ export class McpClientManager {
       Object.keys(extension.mcpServers ?? {}).map((name) => {
         const config = this.allServerConfigs.get(name);
         if (config?.extension?.id === extension.id) {
+          const clientKey = this.getClientKey(name, config);
           this.allServerConfigs.delete(name);
           // Also remove from blocked servers if present
           const index = this.blockedMcpServers.findIndex(
@@ -223,7 +224,7 @@ export class McpClientManager {
           if (index !== -1) {
             this.blockedMcpServers.splice(index, 1);
           }
-          return this.disconnectClient(name, true);
+          return this.disconnectClient(clientKey, true);
         }
         return Promise.resolve();
       }),