From 6d99113936c82a0dd1ee5f65c87d1484fc574363 Mon Sep 17 00:00:00 2001
From: Coco Sheng <cocosheng@google.com>
Date: Tue, 28 Apr 2026 18:46:17 -0400
Subject: [PATCH] fix(core): disconnect extension-backed MCP clients in
 stopExtension (#26136)

---
 .../core/src/tools/mcp-client-manager.test.ts | 35 +++++++++++++++++++
 packages/core/src/tools/mcp-client-manager.ts |  3 +-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts
index 83aa2b59a4..7d4602b8a5 100644
--- a/packages/core/src/tools/mcp-client-manager.test.ts
+++ b/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', () => {
diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts
index b109e2ac03..04563803cd 100644
--- a/packages/core/src/tools/mcp-client-manager.ts
+++ b/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();
       }),