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

This commit is contained in:
Coco Sheng
2026-04-28 18:46:17 -04:00
committed by GitHub
parent fbd8aaad57
commit 6d99113936
2 changed files with 37 additions and 1 deletions
@@ -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', () => {
@@ -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();
}),