feat: restart MCP servers on /mcp refresh (#5479)

Co-authored-by: Brian Ray <[email protected]>
Co-authored-by: N. Taylor Mullen <[email protected]>

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;
+  }
+}