MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-18 22:10:52 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): multi-registry architecture and tool filtering for subagents (#22712)

Browse Source
This commit is contained in:
AK
2026-03-17 13:54:07 -07:00
committed by GitHub
parent c469d563f6
commit bffe9cbcfa
8 changed files with 586 additions and 342 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/tools/mcp-client-manager.ts
+122 -41
View File
@@ -13,6 +13,7 @@ import type { ToolRegistry } from './tool-registry.js';
import {
McpClient,
MCPDiscoveryState,
MCPServerStatus,
populateMcpServerCommand,
} from './mcp-client.js';
import { getErrorMessage, isAuthenticationError } from '../utils/errors.js';
@@ -20,6 +21,11 @@ import type { EventEmitter } from 'node:events';
import { coreEvents } from '../utils/events.js';
import { debugLogger } from '../utils/debugLogger.js';
import { createHash } from 'node:crypto';
import { stableStringify } from '../policy/stable-stringify.js';
import type { PromptRegistry } from '../prompts/prompt-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
/**
* Manages the lifecycle of multiple MCP clients, including local child processes.
* This class is responsible for starting, stopping, and discovering tools from
@@ -30,7 +36,6 @@ export class McpClientManager {
// Track all configured servers (including disabled ones) for UI display
private allServerConfigs: Map<string, MCPServerConfig> = new Map();
private readonly clientVersion: string;
private readonly toolRegistry: ToolRegistry;
private readonly cliConfig: Config;
// If we have ongoing MCP client discovery, this completes once that is done.
private discoveryPromise: Promise<void> | undefined;
@@ -42,6 +47,10 @@ export class McpClientManager {
extensionName: string;
}> = [];
private mainToolRegistry: ToolRegistry | undefined;
private mainPromptRegistry: PromptRegistry | undefined;
private mainResourceRegistry: ResourceRegistry | undefined;
/**
* Track whether the user has explicitly interacted with MCP in this session
* (e.g. by running an /mcp command).
@@ -66,16 +75,24 @@ export class McpClientManager {
constructor(
clientVersion: string,
toolRegistry: ToolRegistry,
cliConfig: Config,
eventEmitter?: EventEmitter,
) {
this.clientVersion = clientVersion;
this.toolRegistry = toolRegistry;
this.cliConfig = cliConfig;
this.eventEmitter = eventEmitter;
}
setMainRegistries(registries: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
}) {
this.mainToolRegistry = registries.toolRegistry;
this.mainPromptRegistry = registries.promptRegistry;
this.mainResourceRegistry = registries.resourceRegistry;
}
setUserInteractedWithMcp() {
this.userInteractedWithMcp = true;
}
@@ -147,6 +164,16 @@ export class McpClientManager {
return this.clients.get(serverName);
}
removeRegistries(registries: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
}): void {
for (const client of this.clients.values()) {
client.removeRegistries(registries);
}
}
/**
* For all the MCP servers associated with this extension:
*
@@ -236,16 +263,17 @@ export class McpClientManager {
return false;
}
private async disconnectClient(name: string, skipRefresh = false) {
const existing = this.clients.get(name);
private async disconnectClient(clientKey: string, skipRefresh = false) {
const existing = this.clients.get(clientKey);
if (existing) {
const serverName = existing.getServerName();
try {
this.clients.delete(name);
this.clients.delete(clientKey);
this.eventEmitter?.emit('mcp-client-update', this.clients);
await existing.disconnect();
} catch (error) {
debugLogger.warn(
`Error stopping client '${name}': ${getErrorMessage(error)}`,
`Error stopping client '${serverName}': ${getErrorMessage(error)}`,
);
} finally {
if (!skipRefresh) {
@@ -257,6 +285,16 @@ export class McpClientManager {
}
}
private getClientKey(name: string, config: MCPServerConfig): string {
const { extension, ...rest } = config;
const keyData = {
name,
config: rest,
extensionId: extension?.id,
};
return createHash('sha256').update(stableStringify(keyData)).digest('hex');
}
/**
* Merges two MCP configurations. The second configuration (override)
* takes precedence for scalar properties, but array properties are
@@ -305,6 +343,11 @@ export class McpClientManager {
async maybeDiscoverMcpServer(
name: string,
config: MCPServerConfig,
registries?: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
},
): Promise<void> {
const existingConfig = this.allServerConfigs.get(name);
if (
@@ -337,11 +380,27 @@ export class McpClientManager {
// Always track server config for UI display
this.allServerConfigs.set(name, finalConfig);
// Capture the existing client synchronously here before any asynchronous
// operations. This ensures that if multiple discovery turns happen
// concurrently, this turn only replaces/disconnects the client that was
// present when this specific configuration update request began.
const existing = this.clients.get(name);
const clientKey = this.getClientKey(name, finalConfig);
// If no registries are provided (main agent) and a server with this name already exists
// but with a different configuration, handle potential conflicts.
if (!registries) {
const existingSameName = Array.from(this.clients.values()).find(
(c) => c.getServerName() === name,
);
if (existingSameName) {
const existingConfigFromClient = existingSameName.getServerConfig();
const existingKey = this.getClientKey(name, existingConfigFromClient);
if (existingKey !== clientKey) {
// This is a configuration update (hot-reload).
// We should stop the old client before starting the new one.
await this.disconnectClient(existingKey, true);
}
}
}
const existing = this.clients.get(clientKey);
// If no connection details are provided, we can't discover this server.
// This often happens when a user provides only overrides (like excludeTools)
@@ -363,7 +422,7 @@ export class McpClientManager {
// User-disabled servers: disconnect if running, don't start
if (await this.isDisabledByUser(name)) {
if (existing) {
await this.disconnectClient(name);
await this.disconnectClient(clientKey);
}
return;
}
@@ -374,34 +433,48 @@ export class McpClientManager {
return;
}
const currentDiscoveryPromise = new Promise<void>((resolve, reject) => {
(async () => {
const currentDiscoveryPromise = new Promise<void>((resolve) => {
void (async () => {
try {
if (existing) {
this.clients.delete(name);
await existing.disconnect();
let client = existing;
if (!client) {
client = new McpClient(
name,
finalConfig,
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),
this.clientVersion,
async () => {
debugLogger.log(
`🔔 Refreshing context for server '${name}'...`,
);
await this.scheduleMcpContextRefresh();
},
);
this.clients.set(clientKey, client);
this.eventEmitter?.emit('mcp-client-update', this.clients);
}
const client = new McpClient(
name,
finalConfig,
this.toolRegistry,
this.cliConfig.getPromptRegistry(),
this.cliConfig.getResourceRegistry(),
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),
this.clientVersion,
async () => {
debugLogger.log(`🔔 Refreshing context for server '${name}'...`);
await this.scheduleMcpContextRefresh();
},
);
this.clients.set(name, client);
this.eventEmitter?.emit('mcp-client-update', this.clients);
const targetRegistries =
registries ??
(this.mainToolRegistry &&
this.mainPromptRegistry &&
this.mainResourceRegistry
? {
toolRegistry: this.mainToolRegistry,
promptRegistry: this.mainPromptRegistry,
resourceRegistry: this.mainResourceRegistry,
}
: undefined);
try {
await client.connect();
await client.discover(this.cliConfig);
if (client.getStatus() === MCPServerStatus.DISCONNECTED) {
await client.connect();
}
if (targetRegistries) {
await client.discoverInto(this.cliConfig, targetRegistries);
}
this.eventEmitter?.emit('mcp-client-update', this.clients);
} catch (error) {
this.eventEmitter?.emit('mcp-client-update', this.clients);
@@ -421,13 +494,13 @@ export class McpClientManager {
const errorMessage = getErrorMessage(error);
this.emitDiagnostic(
'error',
`Error initializing MCP server '${name}': ${errorMessage}`,
`Fatal error ensuring MCP server '${name}' is connected: ${errorMessage}`,
error,
);
} finally {
resolve();
}
})().catch(reject);
})();
});
if (this.discoveryPromise) {
@@ -510,6 +583,11 @@ export class McpClientManager {
* Restarts all MCP servers (including newly enabled ones).
*/
async restart(): Promise<void> {
const disconnectionPromises = Array.from(this.clients.keys()).map((key) =>
this.disconnectClient(key, true),
);
await Promise.all(disconnectionPromises);
await Promise.all(
Array.from(this.allServerConfigs.entries()).map(
async ([name, config]) => {
@@ -534,6 +612,8 @@ export class McpClientManager {
if (!config) {
throw new Error(`No MCP server registered with the name "${name}"`);
}
const clientKey = this.getClientKey(name, config);
await this.disconnectClient(clientKey, true);
await this.maybeDiscoverMcpServer(name, config);
await this.scheduleMcpContextRefresh();
}
@@ -578,11 +658,12 @@ export class McpClientManager {
getMcpInstructions(): string {
const instructions: string[] = [];
for (const [name, client] of this.clients) {
for (const client of this.clients.values()) {
const serverName = client.getServerName();
const clientInstructions = client.getInstructions();
if (clientInstructions) {
instructions.push(
`The following are instructions provided by the tool server '${name}':\n---[start of server instructions]---\n${clientInstructions}\n---[end of server instructions]---`,
`The following are instructions provided by the tool server '${serverName}':\n---[start of server instructions]---\n${clientInstructions}\n---[end of server instructions]---`,
);
}
}