From 9909edc8f214e1e13eaaeadd23856f28b130ed2b Mon Sep 17 00:00:00 2001 From: qqqys Date: Tue, 17 Mar 2026 10:46:40 +0800 Subject: [PATCH] refactor(mcp): enhance reconnect logic and error handling --- packages/cli/src/commands/mcp/reconnect.ts | 66 +++++++++++---- packages/core/src/tools/mcp-tool.test.ts | 6 +- packages/core/src/tools/mcp-tool.ts | 97 +++++++++------------- 3 files changed, 92 insertions(+), 77 deletions(-) diff --git a/packages/cli/src/commands/mcp/reconnect.ts b/packages/cli/src/commands/mcp/reconnect.ts index b4d10cb71..56bb8f59b 100644 --- a/packages/cli/src/commands/mcp/reconnect.ts +++ b/packages/cli/src/commands/mcp/reconnect.ts @@ -15,16 +15,21 @@ import { import { isWorkspaceTrusted } from '../../config/trustedFolders.js'; import type { MCPServerConfig } from '@qwen-code/qwen-code-core'; -async function getMcpServersFromConfig(): Promise< - Record -> { +async function getMcpServersFromConfig( + extensionManager?: ExtensionManager, +): Promise> { const settings = loadSettings(); - const extensionManager = new ExtensionManager({ - isWorkspaceTrusted: !!isWorkspaceTrusted(settings.merged), - telemetrySettings: settings.merged.telemetry, - }); - await extensionManager.refreshCache(); - const extensions = extensionManager.getLoadedExtensions(); + const extManager = + extensionManager ?? + new ExtensionManager({ + isWorkspaceTrusted: !!isWorkspaceTrusted(settings.merged), + telemetrySettings: settings.merged.telemetry, + }); + + if (!extensionManager) { + await extManager.refreshCache(); + } + const extensions = extManager.getLoadedExtensions(); const mcpServers = { ...(settings.merged.mcpServers || {}) }; for (const extension of extensions) { if (extension.isActive) { @@ -64,14 +69,26 @@ async function createMinimalConfig(): Promise { return config; } +interface ReconnectError extends Error { + exitCode: number; +} + +function createReconnectError( + message: string, + exitCode: number = 1, +): ReconnectError { + const error = new Error(message) as ReconnectError; + error.exitCode = exitCode; + return error; +} + async function reconnectMcpServer(serverName: string): Promise { const mcpServers = await getMcpServersFromConfig(); if (!mcpServers[serverName]) { - writeStderrLine( + throw createReconnectError( `Error: Server "${serverName}" not found in configuration.`, ); - process.exit(1); } writeStdoutLine(`Reconnecting to server "${serverName}"...`); @@ -84,15 +101,21 @@ async function reconnectMcpServer(serverName: string): Promise { await config.shutdown(); } catch (error) { const message = error instanceof Error ? error.message : String(error); - writeStderrLine( + throw createReconnectError( `Failed to reconnect to server "${serverName}": ${message}`, ); - process.exit(1); } } async function reconnectAllMcpServers(): Promise { - const mcpServers = await getMcpServersFromConfig(); + const settings = loadSettings(); + const extensionManager = new ExtensionManager({ + isWorkspaceTrusted: !!isWorkspaceTrusted(settings.merged), + telemetrySettings: settings.merged.telemetry, + }); + await extensionManager.refreshCache(); + + const mcpServers = await getMcpServersFromConfig(extensionManager); const serverNames = Object.keys(mcpServers); if (serverNames.length === 0) { @@ -154,10 +177,17 @@ export const reconnectCommand: CommandModule = { const serverName = argv['server-name'] as string | undefined; const all = argv['all'] as boolean; - if (all) { - await reconnectAllMcpServers(); - } else if (serverName) { - await reconnectMcpServer(serverName); + try { + if (all) { + await reconnectAllMcpServers(); + } else if (serverName) { + await reconnectMcpServer(serverName); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const exitCode = (error as ReconnectError)?.exitCode ?? 1; + writeStderrLine(message); + process.exit(exitCode); } }, }; diff --git a/packages/core/src/tools/mcp-tool.test.ts b/packages/core/src/tools/mcp-tool.test.ts index 7098dec66..6837e0984 100644 --- a/packages/core/src/tools/mcp-tool.test.ts +++ b/packages/core/src/tools/mcp-tool.test.ts @@ -1219,7 +1219,7 @@ describe('DiscoveredMCPTool', () => { expect(discoverToolsForServer).not.toHaveBeenCalled(); }); - it('should not retry more than once', async () => { + it('should not retry after reconnection attempt fails', async () => { const params = { param: 'test' }; const mockMcpClient: McpDirectClient = { callTool: vi.fn(), @@ -1271,8 +1271,8 @@ describe('DiscoveredMCPTool', () => { ).rejects.toThrow('ECONNREFUSED'); expect(mockMcpClient.callTool).toHaveBeenCalledTimes(1); - expect(secondMockMcpClient.callTool).toHaveBeenCalledTimes(1); - expect(discoverToolsForServer).toHaveBeenCalledTimes(1); + expect(secondMockMcpClient.callTool).toHaveBeenCalledTimes(3); + expect(discoverToolsForServer).toHaveBeenCalledTimes(3); }); it('should detect various connection error patterns', async () => { diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index b5f520238..971c605bf 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -227,6 +227,45 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< } } + private async handleReconnectOnError( + error: unknown, + signal: AbortSignal, + updateOutput?: (output: ToolResultDisplay) => void, + ): Promise { + debugLogger.error(`MCP server error '${this.serverName}': ${error}`); + + if (this.retryCount < DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES) { + debugLogger.info( + `Reconnection attempt ${this.retryCount + 1}/${DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES} for MCP server '${this.serverName}'`, + ); + const newTool = await this.attemptReconnect(); + if (newTool) { + const newInvocation = new DiscoveredMCPToolInvocation( + newTool['mcpTool'], + this.serverName, + this.serverToolName, + this.displayName, + this.trust, + this.params, + this.cliConfig, + newTool['mcpClient'], + this.mcpTimeout, + this.annotations, + this.retryCount + 1, + ); + return newInvocation.execute(signal, updateOutput); + } + } else if ( + this.retryCount >= DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES + ) { + debugLogger.error( + `Max reconnection attempts (${DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES}) reached for MCP server '${this.serverName}'`, + ); + } + + throw error; + } + async execute( signal: AbortSignal, updateOutput?: (output: ToolResultDisplay) => void, @@ -305,34 +344,7 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< returnDisplay: getStringifiedResultForDisplay(rawResponseParts), }; } catch (error) { - debugLogger.error(`MCP server error '${this.serverName}': ${error}`); - - // Attempt reconnection with retry limit - if (this.retryCount < DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES) { - const newTool = await this.attemptReconnect(); - if (newTool) { - const newInvocation = new DiscoveredMCPToolInvocation( - newTool['mcpTool'], - this.serverName, - this.serverToolName, - this.displayName, - this.trust, - this.params, - this.cliConfig, - newTool['mcpClient'], - this.mcpTimeout, - this.annotations, - this.retryCount + 1, - ); - return newInvocation.execute(signal, updateOutput); - } - } else { - debugLogger.error( - `Max reconnection attempts (${DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES}) reached for MCP server '${this.serverName}'`, - ); - } - - throw error; + return this.handleReconnectOnError(error, signal, updateOutput); } } @@ -406,34 +418,7 @@ class DiscoveredMCPToolInvocation extends BaseToolInvocation< returnDisplay: getStringifiedResultForDisplay(rawResponseParts), }; } catch (error) { - debugLogger.error(`MCP server error '${this.serverName}': ${error}`); - - // Attempt reconnection with retry limit - if (this.retryCount < DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES) { - const newTool = await this.attemptReconnect(); - if (newTool) { - const newInvocation = new DiscoveredMCPToolInvocation( - newTool['mcpTool'], - this.serverName, - this.serverToolName, - this.displayName, - this.trust, - this.params, - this.cliConfig, - newTool['mcpClient'], - this.mcpTimeout, - this.annotations, - this.retryCount + 1, - ); - return newInvocation.execute(signal); - } - } else { - debugLogger.error( - `Max reconnection attempts (${DiscoveredMCPToolInvocation.MAX_RECONNECT_RETRIES}) reached for MCP server '${this.serverName}'`, - ); - } - - throw error; + return this.handleReconnectOnError(error, signal); } }