1 files changed, 23 insertions, 20 deletions
diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts
index ff155679..90531742 100644
--- a/packages/core/src/tools/tool-registry.ts
+++ b/packages/core/src/tools/tool-registry.ts
@@ -16,7 +16,8 @@ import {
 import { Config } from '../config/config.js';
 import { spawn } from 'node:child_process';
 import { StringDecoder } from 'node:string_decoder';
-import { discoverMcpTools } from './mcp-client.js';
+import { connectAndDiscover } from './mcp-client.js';
+import { McpClientManager } from './mcp-client-manager.js';
 import { DiscoveredMCPTool } from './mcp-tool.js';
 import { parse } from 'shell-quote';
 
@@ -163,9 +164,18 @@ Signal: Signal number or \`(none)\` if no signal was received.
 export class ToolRegistry {
   private tools: Map<string, AnyDeclarativeTool> = new Map();
   private config: Config;
+  private mcpClientManager: McpClientManager;
 
   constructor(config: Config) {
     this.config = config;
+    this.mcpClientManager = new McpClientManager(
+      this.config.getMcpServers() ?? {},
+      this.config.getMcpServerCommand(),
+      this,
+      this.config.getPromptRegistry(),
+      this.config.getDebugMode(),
+      this.config.getWorkspaceContext(),
+    );
   }
 
   /**
@@ -220,14 +230,7 @@ export class ToolRegistry {
     await this.discoverAndRegisterToolsFromCommand();
 
     // discover tools using MCP servers, if configured
-    await discoverMcpTools(
-      this.config.getMcpServers() ?? {},
-      this.config.getMcpServerCommand(),
-      this,
-      this.config.getPromptRegistry(),
-      this.config.getDebugMode(),
-      this.config.getWorkspaceContext(),
-    );
+    await this.mcpClientManager.discoverAllMcpTools();
   }
 
   /**
@@ -242,14 +245,14 @@ export class ToolRegistry {
     this.config.getPromptRegistry().clear();
 
     // discover tools using MCP servers, if configured
-    await discoverMcpTools(
-      this.config.getMcpServers() ?? {},
-      this.config.getMcpServerCommand(),
-      this,
-      this.config.getPromptRegistry(),
-      this.config.getDebugMode(),
-      this.config.getWorkspaceContext(),
-    );
+    await this.mcpClientManager.discoverAllMcpTools();
+  }
+
+  /**
+   * Restarts all MCP servers and re-discovers tools.
+   */
+  async restartMcpServers(): Promise<void> {
+    await this.discoverMcpTools();
   }
 
   /**
@@ -269,9 +272,9 @@ export class ToolRegistry {
     const mcpServers = this.config.getMcpServers() ?? {};
     const serverConfig = mcpServers[serverName];
     if (serverConfig) {
-      await discoverMcpTools(
-        { [serverName]: serverConfig },
-        undefined,
+      await connectAndDiscover(
+        serverName,
+        serverConfig,
         this,
         this.config.getPromptRegistry(),
         this.config.getDebugMode(),