fix: address critical issues in MCP server restart handling

- Add pendingRestarts processing after tool execution completes
- Add cleanup for activeToolExecutions with timeout mechanism to prevent memory leaks
- Refactor updateServerToolList to reduce code duplication using updateServerToolListInternal
- Fix race condition in watcher re-enable logic by checking disposal state
- Add comprehensive test coverage for all fixes

Author: roomote

src/services/mcp/McpHub.ts

@@ -152,6 +152,8 @@ export class McpHub {
 	private refCount: number = 0 // Reference counter for active clients
 	private configChangeDebounceTimers: Map<string, NodeJS.Timeout> = new Map()
 	private activeToolExecutions: Map<string, Set<string>> = new Map() // Track active tool executions per server
+	private pendingRestarts: Set<string> = new Set() // Track servers that need restart after tool execution
+	private toolExecutionTimeouts: Map<string, NodeJS.Timeout> = new Map() // Track timeouts for tool execution cleanup
 
 	constructor(provider: ClineProvider) {
 		this.providerRef = new WeakRef(provider)
@@ -1172,10 +1174,13 @@ export class McpHub {
 	async restartConnection(serverName: string, source?: "global" | "project"): Promise<void> {
 		// Check if there are active tool executions for this server
 		if (this.hasActiveToolExecutions(serverName, source)) {
-			console.log(`Skipping restart for ${serverName} - tools are currently executing`)
-			vscode.window.showWarningMessage(
-				t("mcp:errors.cannot_restart_tools_running", { serverName }) ||
-					`Cannot restart server "${serverName}" while tools are running. Please wait for tool execution to complete.`,
+			console.log(`Deferring restart for ${serverName} - tools are currently executing`)
+			// Add to pending restarts
+			const serverKey = `${serverName}:${source || "global"}`
+			this.pendingRestarts.add(serverKey)
+			vscode.window.showInformationMessage(
+				t("mcp:info.restart_deferred", { serverName }) ||
+					`Server "${serverName}" will be restarted after tool execution completes.`,
 			)
 			return
 		}
@@ -1636,6 +1641,15 @@ export class McpHub {
 		const executionId = `${toolName}:${Date.now()}`
 		this.activeToolExecutions.get(serverKey)!.add(executionId)
 
+		// Set up a timeout to clean up the execution tracking in case of hangs
+		const cleanupTimeout = setTimeout(() => {
+			this.cleanupToolExecution(serverKey, executionId)
+		}, timeout + 5000) // Add 5 seconds buffer after the tool timeout
+
+		// Store the timeout so we can clear it if the tool completes normally
+		const timeoutKey = `${serverKey}:${executionId}`
+		this.toolExecutionTimeouts.set(timeoutKey, cleanupTimeout)
+
 		try {
 			const result = await connection.client.request(
 				{
@@ -1651,23 +1665,60 @@ export class McpHub {
 				},
 			)
 
+			// Clear the cleanup timeout since the tool completed successfully
+			this.clearToolExecutionTimeout(timeoutKey)
+
 			// Remove from active executions on success
-			this.activeToolExecutions.get(serverKey)?.delete(executionId)
-			if (this.activeToolExecutions.get(serverKey)?.size === 0) {
-				this.activeToolExecutions.delete(serverKey)
-			}
+			this.cleanupToolExecution(serverKey, executionId)
 
 			return result
 		} catch (error) {
+			// Clear the cleanup timeout
+			this.clearToolExecutionTimeout(timeoutKey)
+
 			// Remove from active executions on error
-			this.activeToolExecutions.get(serverKey)?.delete(executionId)
-			if (this.activeToolExecutions.get(serverKey)?.size === 0) {
-				this.activeToolExecutions.delete(serverKey)
-			}
+			this.cleanupToolExecution(serverKey, executionId)
+
 			throw error
 		}
 	}
 
+	/**
+	 * Clean up tool execution tracking and process pending restarts if needed
+	 * @param serverKey The server key
+	 * @param executionId The execution ID to clean up
+	 */
+	private cleanupToolExecution(serverKey: string, executionId: string): void {
+		// Remove from active executions
+		this.activeToolExecutions.get(serverKey)?.delete(executionId)
+		if (this.activeToolExecutions.get(serverKey)?.size === 0) {
+			this.activeToolExecutions.delete(serverKey)
+
+			// Check if this server has pending restarts
+			if (this.pendingRestarts.has(serverKey)) {
+				this.pendingRestarts.delete(serverKey)
+				// Extract server name and source from the key
+				const [serverName, source] = serverKey.split(":")
+				// Process the pending restart
+				this.restartConnection(serverName, source as "global" | "project").catch((error) => {
+					console.error(`Failed to process pending restart for ${serverName}:`, error)
+				})
+			}
+		}
+	}
+
+	/**
+	 * Clear a tool execution timeout
+	 * @param timeoutKey The timeout key to clear
+	 */
+	private clearToolExecutionTimeout(timeoutKey: string): void {
+		const timeout = this.toolExecutionTimeouts.get(timeoutKey)
+		if (timeout) {
+			clearTimeout(timeout)
+			this.toolExecutionTimeouts.delete(timeoutKey)
+		}
+	}
+
 	/**
 	 * Check if any tools are currently executing for a specific server
 	 * @param serverName The name of the server to check
@@ -1708,6 +1759,27 @@ export class McpHub {
 		toolName: string,
 		listName: "alwaysAllow" | "disabledTools",
 		addTool: boolean,
+	): Promise<void> {
+		// Use the common implementation with restart enabled
+		await this.updateServerToolListInternal(serverName, source, toolName, listName, addTool, true)
+	}
+
+	/**
+	 * Internal implementation for updating server tool lists
+	 * @param serverName The name of the server to update
+	 * @param source Whether to update the global or project config
+	 * @param toolName The name of the tool to add or remove
+	 * @param listName The name of the list to modify ("alwaysAllow" or "disabledTools")
+	 * @param addTool Whether to add (true) or remove (false) the tool from the list
+	 * @param allowRestart Whether to allow file watchers to trigger restart
+	 */
+	private async updateServerToolListInternal(
+		serverName: string,
+		source: "global" | "project",
+		toolName: string,
+		listName: "alwaysAllow" | "disabledTools",
+		addTool: boolean,
+		allowRestart: boolean,
 	): Promise<void> {
 		// Find the connection with matching name and source
 		const connection = this.findConnection(serverName, source)
@@ -1731,7 +1803,6 @@ export class McpHub {
 		}
 
 		// Normalize path for cross-platform compatibility
-		// Use a consistent path format for both reading and writing
 		const normalizedPath = process.platform === "win32" ? configPath.replace(/\\/g, "/") : configPath
 
 		// Read the appropriate config file
@@ -1763,12 +1834,57 @@ export class McpHub {
 			targetList.splice(toolIndex, 1)
 		}
 
+		// Handle file watcher based on allowRestart flag
+		let watcherToRestore: vscode.FileSystemWatcher | undefined
+		if (!allowRestart) {
+			// Temporarily disable the watcher to prevent restart
+			if (source === "project" && this.projectMcpWatcher) {
+				watcherToRestore = this.projectMcpWatcher
+				this.projectMcpWatcher.dispose()
+				this.projectMcpWatcher = undefined
+			} else if (source === "global" && this.settingsWatcher) {
+				watcherToRestore = this.settingsWatcher
+				this.settingsWatcher.dispose()
+				this.settingsWatcher = undefined
+			}
+		}
+
 		await fs.writeFile(normalizedPath, JSON.stringify(config, null, 2))
 
 		if (connection) {
-			connection.server.tools = await this.fetchToolsList(serverName, source)
+			if (allowRestart) {
+				// Normal flow - fetch fresh tool list
+				connection.server.tools = await this.fetchToolsList(serverName, source)
+			} else {
+				// Update the in-memory tool state without restarting
+				const tools = connection.server.tools
+				if (tools) {
+					const tool = tools.find((t) => t.name === toolName)
+					if (tool) {
+						if (listName === "alwaysAllow") {
+							tool.alwaysAllow = addTool
+						} else if (listName === "disabledTools") {
+							tool.enabledForPrompt = !addTool
+						}
+					}
+				}
+			}
 			await this.notifyWebviewOfServerChanges()
 		}
+
+		// Re-enable the watcher after a short delay if it was disabled
+		if (!allowRestart && watcherToRestore) {
+			setTimeout(async () => {
+				// Check if not disposed during the timeout
+				if (!this.isDisposed) {
+					if (source === "project" && !this.projectMcpWatcher) {
+						await this.watchProjectMcpFile()
+					} else if (source === "global" && !this.settingsWatcher) {
+						await this.watchMcpSettingsFile()
+					}
+				}
+			}, 1000)
+		}
 	}
 
 	async toggleToolAlwaysAllow(
@@ -1780,7 +1896,7 @@ export class McpHub {
 		try {
 			// Check if there are active tool executions for this server
 			if (this.hasActiveToolExecutions(serverName, source)) {
-				console.log(`Skipping server restart for ${serverName} - tools are currently executing`)
+				console.log(`Deferring config update for ${serverName} - tools are currently executing`)
 				// Update the config file without triggering a restart
 				await this.updateServerToolListWithoutRestart(serverName, source, toolName, "alwaysAllow", shouldAllow)
 			} else {
@@ -1807,101 +1923,8 @@ export class McpHub {
 		listName: "alwaysAllow" | "disabledTools",
 		addTool: boolean,
 	): Promise<void> {
-		// Find the connection with matching name and source
-		const connection = this.findConnection(serverName, source)
-
-		if (!connection) {
-			throw new Error(`Server ${serverName} with source ${source} not found`)
-		}
-
-		// Determine the correct config path based on the source
-		let configPath: string
-		if (source === "project") {
-			// Get project MCP config path
-			const projectMcpPath = await this.getProjectMcpPath()
-			if (!projectMcpPath) {
-				throw new Error("Project MCP configuration file not found")
-			}
-			configPath = projectMcpPath
-		} else {
-			// Get global MCP settings path
-			configPath = await this.getMcpSettingsFilePath()
-		}
-
-		// Normalize path for cross-platform compatibility
-		const normalizedPath = process.platform === "win32" ? configPath.replace(/\\/g, "/") : configPath
-
-		// Read the appropriate config file
-		const content = await fs.readFile(normalizedPath, "utf-8")
-		const config = JSON.parse(content)
-
-		if (!config.mcpServers) {
-			config.mcpServers = {}
-		}
-
-		if (!config.mcpServers[serverName]) {
-			config.mcpServers[serverName] = {
-				type: "stdio",
-				command: "node",
-				args: [], // Default to an empty array; can be set later if needed
-			}
-		}
-
-		if (!config.mcpServers[serverName][listName]) {
-			config.mcpServers[serverName][listName] = []
-		}
-
-		const targetList = config.mcpServers[serverName][listName]
-		const toolIndex = targetList.indexOf(toolName)
-
-		if (addTool && toolIndex === -1) {
-			targetList.push(toolName)
-		} else if (!addTool && toolIndex !== -1) {
-			targetList.splice(toolIndex, 1)
-		}
-
-		// Write the config file directly without triggering file watcher
-		// We'll temporarily disable the watcher to prevent restart
-		const watcherKey = source === "project" ? this.projectMcpWatcher : this.settingsWatcher
-		const wasWatcherActive = !!watcherKey
-
-		if (wasWatcherActive && source === "project" && this.projectMcpWatcher) {
-			this.projectMcpWatcher.dispose()
-			this.projectMcpWatcher = undefined
-		} else if (wasWatcherActive && source === "global" && this.settingsWatcher) {
-			this.settingsWatcher.dispose()
-			this.settingsWatcher = undefined
-		}
-
-		await fs.writeFile(normalizedPath, JSON.stringify(config, null, 2))
-
-		// Update the in-memory tool state without restarting
-		if (connection) {
-			// Update the tool's alwaysAllow or enabledForPrompt status in memory
-			const tools = connection.server.tools
-			if (tools) {
-				const tool = tools.find((t) => t.name === toolName)
-				if (tool) {
-					if (listName === "alwaysAllow") {
-						tool.alwaysAllow = addTool
-					} else if (listName === "disabledTools") {
-						tool.enabledForPrompt = !addTool
-					}
-				}
-			}
-			await this.notifyWebviewOfServerChanges()
-		}
-
-		// Re-enable the watcher after a short delay
-		if (wasWatcherActive) {
-			setTimeout(async () => {
-				if (source === "project") {
-					await this.watchProjectMcpFile()
-				} else {
-					await this.watchMcpSettingsFile()
-				}
-			}, 1000)
-		}
+		// Use the common implementation
+		await this.updateServerToolListInternal(serverName, source, toolName, listName, addTool, false)
 	}
 
 	async toggleToolEnabledForPrompt(
@@ -1917,7 +1940,7 @@ export class McpHub {
 
 			// Check if there are active tool executions for this server
 			if (this.hasActiveToolExecutions(serverName, source)) {
-				console.log(`Skipping server restart for ${serverName} - tools are currently executing`)
+				console.log(`Deferring config update for ${serverName} - tools are currently executing`)
 				// Update the config file without triggering a restart
 				await this.updateServerToolListWithoutRestart(
 					serverName,
@@ -2004,6 +2027,15 @@ export class McpHub {
 		}
 		this.configChangeDebounceTimers.clear()
 
+		// Clear all tool execution timeouts
+		for (const timeout of this.toolExecutionTimeouts.values()) {
+			clearTimeout(timeout)
+		}
+		this.toolExecutionTimeouts.clear()
+
+		// Clear pending restarts
+		this.pendingRestarts.clear()
+
 		this.removeAllFileWatchers()
 		for (const connection of this.connections) {
 			try {