diff options
Diffstat (limited to 'packages/core/src/tools/mcp-client-manager.ts')
| -rw-r--r-- | packages/core/src/tools/mcp-client-manager.ts | 115 |
1 files changed, 115 insertions, 0 deletions
diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts new file mode 100644 index 00000000..c22afb8f --- /dev/null +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -0,0 +1,115 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { MCPServerConfig } from '../config/config.js'; +import { ToolRegistry } from './tool-registry.js'; +import { PromptRegistry } from '../prompts/prompt-registry.js'; +import { + McpClient, + MCPDiscoveryState, + populateMcpServerCommand, +} from './mcp-client.js'; +import { getErrorMessage } from '../utils/errors.js'; +import { WorkspaceContext } from '../utils/workspaceContext.js'; + +/** + * Manages the lifecycle of multiple MCP clients, including local child processes. + * This class is responsible for starting, stopping, and discovering tools from + * a collection of MCP servers defined in the configuration. + */ +export class McpClientManager { + private clients: Map<string, McpClient> = new Map(); + private readonly mcpServers: Record<string, MCPServerConfig>; + private readonly mcpServerCommand: string | undefined; + private readonly toolRegistry: ToolRegistry; + private readonly promptRegistry: PromptRegistry; + private readonly debugMode: boolean; + private readonly workspaceContext: WorkspaceContext; + private discoveryState: MCPDiscoveryState = MCPDiscoveryState.NOT_STARTED; + + constructor( + mcpServers: Record<string, MCPServerConfig>, + mcpServerCommand: string | undefined, + toolRegistry: ToolRegistry, + promptRegistry: PromptRegistry, + debugMode: boolean, + workspaceContext: WorkspaceContext, + ) { + this.mcpServers = mcpServers; + this.mcpServerCommand = mcpServerCommand; + this.toolRegistry = toolRegistry; + this.promptRegistry = promptRegistry; + this.debugMode = debugMode; + this.workspaceContext = workspaceContext; + } + + /** + * Initiates the tool discovery process for all configured MCP servers. + * It connects to each server, discovers its available tools, and registers + * them with the `ToolRegistry`. + */ + async discoverAllMcpTools(): Promise<void> { + await this.stop(); + this.discoveryState = MCPDiscoveryState.IN_PROGRESS; + const servers = populateMcpServerCommand( + this.mcpServers, + this.mcpServerCommand, + ); + + const discoveryPromises = Object.entries(servers).map( + async ([name, config]) => { + const client = new McpClient( + name, + config, + this.toolRegistry, + this.promptRegistry, + this.workspaceContext, + this.debugMode, + ); + this.clients.set(name, client); + try { + await client.connect(); + await client.discover(); + } catch (error) { + // Log the error but don't let a single failed server stop the others + console.error( + `Error during discovery for server '${name}': ${getErrorMessage( + error, + )}`, + ); + } + }, + ); + + await Promise.all(discoveryPromises); + this.discoveryState = MCPDiscoveryState.COMPLETED; + } + + /** + * Stops all running local MCP servers and closes all client connections. + * This is the cleanup method to be called on application exit. + */ + async stop(): Promise<void> { + const disconnectionPromises = Array.from(this.clients.entries()).map( + async ([name, client]) => { + try { + await client.disconnect(); + } catch (error) { + console.error( + `Error stopping client '${name}': ${getErrorMessage(error)}`, + ); + } + }, + ); + + await Promise.all(disconnectionPromises); + this.clients.clear(); + } + + getDiscoveryState(): MCPDiscoveryState { + return this.discoveryState; + } +} |
