From e57c06fc6bb88b49a62ad7219e07e190c98ffa49 Mon Sep 17 00:00:00 2001 From: Mehmet Turac Date: Tue, 19 May 2026 13:00:36 +0300 Subject: [PATCH] fix: reconnect stale SSE MCP sessions --- core/context/mcp/MCPConnection.ts | 106 ++++++++++++++++++++--- core/context/mcp/MCPConnection.vitest.ts | 43 +++++++++ core/context/mcp/MCPManagerSingleton.ts | 2 +- core/tools/callTool.ts | 2 +- 4 files changed, 137 insertions(+), 16 deletions(-) diff --git a/core/context/mcp/MCPConnection.ts b/core/context/mcp/MCPConnection.ts index d25e9f6a0e7..d1953ecc58a 100644 --- a/core/context/mcp/MCPConnection.ts +++ b/core/context/mcp/MCPConnection.ts @@ -56,6 +56,18 @@ function is401Error(error: unknown) { ); } +function createMcpClient() { + return new Client( + { + name: "continue-client", + version: "1.0.0", + }, + { + capabilities: {}, + }, + ); +} + export type MCPExtras = { ide: IDE; }; @@ -85,19 +97,70 @@ class MCPConnection { // Don't construct transport in constructor to avoid blocking this.transport = {} as Transport; // Will be set in connectClient - this.client = new Client( - { - name: "continue-client", - version: "1.0.0", - }, - { - capabilities: {}, - }, - ); + this.client = createMcpClient(); this.abortController = new AbortController(); } + private async resetClientAndTransport() { + try { + await this.client.close(); + } catch { + // Ignore close errors while replacing stale clients/transports. + } + + try { + await this.transport.close?.(); + } catch { + // Ignore close errors while replacing stale clients/transports. + } + + this.client = createMcpClient(); + this.transport = {} as Transport; + } + + private shouldReconnectAfterError(error: unknown) { + if (this.options.type !== "sse" || this.status === "disabled") { + return false; + } + + const message = ( + error instanceof Error ? error.message : String(error) + ).toLowerCase(); + + const sessionError = + message.includes("session") && + (message.includes("invalid") || + message.includes("unknown") || + message.includes("expired") || + message.includes("not found") || + message.includes("valid") || + message.includes("missing")); + + return ( + sessionError || + message.includes("connection closed") || + message.includes("transport closed") + ); + } + + private async withSseReconnectRetry(operation: () => Promise) { + try { + return await operation(); + } catch (error) { + if (!this.shouldReconnectAfterError(error)) { + throw error; + } + + await this.connectClient(true, new AbortController().signal); + if (this.status !== "connected") { + throw error; + } + + return await operation(); + } + } + async disconnect(disable = false) { this.abortController.abort(); await this.client.close(); @@ -147,6 +210,7 @@ class MCPConnection { this.abortController.abort(); this.abortController = new AbortController(); + await this.resetClientAndTransport(); // currently support oauth for sse transports only if (this.options.type === "sse") { @@ -613,11 +677,25 @@ Org-level secrets can only be used for MCP by Background Agents (https://docs.co } async getResource(uri: string) { - return await this.client.readResource( - { uri }, - { - timeout: this.options.timeout, - }, + return await this.withSseReconnectRetry(() => + this.client.readResource( + { uri }, + { + timeout: this.options.timeout, + }, + ), + ); + } + + async getPrompt(...args: Parameters) { + return await this.withSseReconnectRetry(() => + this.client.getPrompt(...args), + ); + } + + async callTool(...args: Parameters) { + return await this.withSseReconnectRetry(() => + this.client.callTool(...args), ); } } diff --git a/core/context/mcp/MCPConnection.vitest.ts b/core/context/mcp/MCPConnection.vitest.ts index adb4b9c5493..14983545464 100644 --- a/core/context/mcp/MCPConnection.vitest.ts +++ b/core/context/mcp/MCPConnection.vitest.ts @@ -310,6 +310,49 @@ describe("MCPConnection", () => { expect(mockConnect).toHaveBeenCalled(); }); + it("should reconnect and retry SSE tool calls after stale session errors", async () => { + const conn = new MCPConnection({ + name: "test-mcp", + id: "test-id", + type: "sse", + url: "http://test.com/events", + }); + conn.status = "connected"; + + const mockCallTool = vi + .spyOn(Client.prototype, "callTool") + .mockRejectedValueOnce(new Error("Invalid session ID")) + .mockResolvedValueOnce({ content: [], isError: false } as any); + const mockReconnect = vi + .spyOn(conn, "connectClient") + .mockImplementation(async () => { + conn.status = "connected"; + }); + + const result = await conn.callTool({ name: "test-tool" } as any); + + expect(result).toEqual({ content: [], isError: false }); + expect(mockReconnect).toHaveBeenCalledWith(true, expect.any(AbortSignal)); + expect(mockCallTool).toHaveBeenCalledTimes(2); + }); + + it("should not retry non-SSE tool calls after stale session errors", async () => { + const conn = new MCPConnection(options); + conn.status = "connected"; + + const mockCallTool = vi + .spyOn(Client.prototype, "callTool") + .mockRejectedValue(new Error("Invalid session ID")); + const mockReconnect = vi.spyOn(conn, "connectClient"); + + await expect(conn.callTool({ name: "test-tool" } as any)).rejects.toThrow( + "Invalid session ID", + ); + + expect(mockReconnect).not.toHaveBeenCalled(); + expect(mockCallTool).toHaveBeenCalledTimes(1); + }); + it.skip("should include stderr output in error message when stdio command fails", async () => { // Clear any existing mocks to ensure we get real behavior vi.restoreAllMocks(); diff --git a/core/context/mcp/MCPManagerSingleton.ts b/core/context/mcp/MCPManagerSingleton.ts index cac57e1e68f..277624e95c2 100644 --- a/core/context/mcp/MCPManagerSingleton.ts +++ b/core/context/mcp/MCPManagerSingleton.ts @@ -196,7 +196,7 @@ export class MCPManagerSingleton { `Error getting prompt: MCP Connection ${serverName} not found`, ); } - return await connection.client.getPrompt({ + return await connection.getPrompt({ name: promptName, arguments: args, }); diff --git a/core/tools/callTool.ts b/core/tools/callTool.ts index 22e83bb46ab..57659e24053 100644 --- a/core/tools/callTool.ts +++ b/core/tools/callTool.ts @@ -99,7 +99,7 @@ async function callToolFromUri( args, extras.tool?.function?.parameters, ); - const response = await client.client.callTool( + const response = await client.callTool( { name: toolName, arguments: coercedArgs,