From 42a8f210cace01f64558f6d3723bd5553e7bee7c Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Sun, 1 Feb 2026 12:26:19 +0100 Subject: [PATCH 1/2] fix(mcp): ensure MCP transport is closed to prevent memory leaks Root cause: When MCP servers disconnect or encounter errors, the transport (Stdio/SSE/HTTP) was not being closed because: 1. `connectToMcpServer()` only returned the Client, not the transport 2. The transport reference was created locally and lost after connection 3. `McpClient.transport` was declared but never assigned This caused event listeners and buffers to run indefinitely, eventually causing "JavaScript heap out of memory" errors in long-running sessions. Changes: - `connectToMcpServer()` now returns `{client, transport}` instead of just `Client` - `McpClient.connect()` stores the transport reference - `McpClient.disconnect()` now properly closes the transport before the client - Error handlers in both `McpClient` and `connectAndDiscover` now close the transport to prevent leaks on connection errors - Helper functions `connectWithSSETransport()` and `retryWithOAuth()` now return the transport for proper lifecycle management Addresses the memory leak fix from PR #12801. --- packages/core/src/tools/mcp-client.test.ts | 16 ++--- packages/core/src/tools/mcp-client.ts | 76 +++++++++++++++++----- 2 files changed, 67 insertions(+), 25 deletions(-) diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index e4bbd7d756b..25b03c3c448 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -749,9 +749,9 @@ describe('mcp-client', () => { vi.mocked(ClientLib.Client).mockReturnValue( mockedClient as unknown as ClientLib.Client, ); - vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue( - {} as SdkClientStdioLib.StdioClientTransport, - ); + vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue({ + close: vi.fn(), + } as unknown as SdkClientStdioLib.StdioClientTransport); const mockedToolRegistry = { registerTool: vi.fn(), unregisterTool: vi.fn(), @@ -1853,7 +1853,7 @@ describe('connectToMcpServer with OAuth', () => { EMPTY_CONFIG, ); - expect(client).toBe(mockedClient); + expect(client.client).toBe(mockedClient); expect(mockedClient.connect).toHaveBeenCalledTimes(2); expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce(); @@ -1899,7 +1899,7 @@ describe('connectToMcpServer with OAuth', () => { EMPTY_CONFIG, ); - expect(client).toBe(mockedClient); + expect(client.client).toBe(mockedClient); expect(mockedClient.connect).toHaveBeenCalledTimes(2); expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce(); expect(OAuthUtils.discoverOAuthConfig).toHaveBeenCalledWith(serverUrl); @@ -1994,7 +1994,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { EMPTY_CONFIG, ); - expect(client).toBe(mockedClient); + expect(client.client).toBe(mockedClient); // First HTTP attempt fails, second SSE attempt succeeds expect(mockedClient.connect).toHaveBeenCalledTimes(2); }); @@ -2035,7 +2035,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { EMPTY_CONFIG, ); - expect(client).toBe(mockedClient); + expect(client.client).toBe(mockedClient); expect(mockedClient.connect).toHaveBeenCalledTimes(2); }); }); @@ -2120,7 +2120,7 @@ describe('connectToMcpServer - OAuth with transport fallback', () => { EMPTY_CONFIG, ); - expect(client).toBe(mockedClient); + expect(client.client).toBe(mockedClient); expect(mockedClient.connect).toHaveBeenCalledTimes(3); expect(mockAuthProvider.authenticate).toHaveBeenCalledOnce(); }); diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index e7aa866a09e..f6bbdc38fa4 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -144,7 +144,7 @@ export class McpClient { } this.updateStatus(MCPServerStatus.CONNECTING); try { - this.client = await connectToMcpServer( + const { client, transport } = await connectToMcpServer( this.clientVersion, this.serverName, this.serverConfig, @@ -152,11 +152,13 @@ export class McpClient { this.workspaceContext, this.cliConfig.sanitizationConfig, ); + this.client = client; + this.transport = transport; this.registerNotificationHandlers(); const originalOnError = this.client.onerror; - this.client.onerror = (error) => { + this.client.onerror = async (error) => { if (this.status !== MCPServerStatus.CONNECTED) { return; } @@ -167,6 +169,14 @@ export class McpClient { error, ); this.updateStatus(MCPServerStatus.DISCONNECTED); + // Close transport to prevent memory leaks + if (this.transport) { + try { + await this.transport.close(); + } catch { + // Ignore errors when closing transport on error + } + } }; this.updateStatus(MCPServerStatus.CONNECTED); } catch (error) { @@ -909,8 +919,9 @@ export async function connectAndDiscover( updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTING); let mcpClient: Client | undefined; + let transport: Transport | undefined; try { - mcpClient = await connectToMcpServer( + const result = await connectToMcpServer( clientVersion, mcpServerName, mcpServerConfig, @@ -918,10 +929,20 @@ export async function connectAndDiscover( workspaceContext, cliConfig.sanitizationConfig, ); + mcpClient = result.client; + transport = result.transport; - mcpClient.onerror = (error) => { + mcpClient.onerror = async (error) => { coreEvents.emitFeedback('error', `MCP ERROR (${mcpServerName}):`, error); updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED); + // Close transport to prevent memory leaks + if (transport) { + try { + await transport.close(); + } catch { + // Ignore errors when closing transport on error + } + } }; // Attempt to discover both prompts and tools @@ -952,6 +973,14 @@ export async function connectAndDiscover( } toolRegistry.sortTools(); } catch (error) { + // Close transport to prevent memory leaks + if (transport) { + try { + await transport.close(); + } catch { + // Ignore errors when closing transport + } + } if (mcpClient) { // eslint-disable-next-line @typescript-eslint/no-floating-promises mcpClient.close(); @@ -1302,16 +1331,18 @@ function createSSETransportWithAuth( * @param client The MCP client to connect * @param config The MCP server configuration * @param accessToken Optional OAuth access token for authentication + * @returns The transport used for connection */ async function connectWithSSETransport( client: Client, config: MCPServerConfig, accessToken?: string | null, -): Promise { +): Promise { const transport = createSSETransportWithAuth(config, accessToken); await client.connect(transport, { timeout: config.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, }); + return transport; } /** @@ -1341,6 +1372,7 @@ async function showAuthRequiredMessage(serverName: string): Promise { * @param config The MCP server configuration * @param accessToken The OAuth access token to use * @param httpReturned404 Whether the HTTP transport returned 404 (indicating SSE-only server) + * @returns The transport used for connection */ async function retryWithOAuth( client: Client, @@ -1348,17 +1380,21 @@ async function retryWithOAuth( config: MCPServerConfig, accessToken: string, httpReturned404: boolean, -): Promise { +): Promise { if (httpReturned404) { // HTTP returned 404, only try SSE debugLogger.log( `Retrying SSE connection to '${serverName}' with OAuth token...`, ); - await connectWithSSETransport(client, config, accessToken); + const transport = await connectWithSSETransport( + client, + config, + accessToken, + ); debugLogger.log( `Successfully connected to '${serverName}' using SSE with OAuth.`, ); - return; + return transport; } // HTTP returned 401, try HTTP with OAuth first @@ -1382,6 +1418,7 @@ async function retryWithOAuth( debugLogger.log( `Successfully connected to '${serverName}' using HTTP with OAuth.`, ); + return httpTransport; } catch (httpError) { await httpTransport.close(); @@ -1393,10 +1430,15 @@ async function retryWithOAuth( !config.httpUrl ) { debugLogger.log(`HTTP with OAuth returned 404, trying SSE with OAuth...`); - await connectWithSSETransport(client, config, accessToken); + const sseTransport = await connectWithSSETransport( + client, + config, + accessToken, + ); debugLogger.log( `Successfully connected to '${serverName}' using SSE with OAuth.`, ); + return sseTransport; } else { throw httpError; } @@ -1410,7 +1452,7 @@ async function retryWithOAuth( * * @param mcpServerName The name of the MCP server, used for logging and identification. * @param mcpServerConfig The configuration specifying how to connect to the server. - * @returns A promise that resolves to a connected MCP `Client` instance. + * @returns A promise that resolves to a connected MCP `Client` instance and its transport. * @throws An error if the connection fails or the configuration is invalid. */ export async function connectToMcpServer( @@ -1420,7 +1462,7 @@ export async function connectToMcpServer( debugMode: boolean, workspaceContext: WorkspaceContext, sanitizationConfig: EnvironmentSanitizationConfig, -): Promise { +): Promise<{ client: Client; transport: Transport }> { const mcpClient = new Client( { name: 'gemini-cli-mcp-client', @@ -1492,7 +1534,7 @@ export async function connectToMcpServer( await mcpClient.connect(transport, { timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, }); - return mcpClient; + return { client: mcpClient, transport }; } catch (error) { await transport.close(); firstAttemptError = error as Error; @@ -1523,7 +1565,7 @@ export async function connectToMcpServer( try { // Try SSE with stored OAuth token if available // This ensures that SSE fallback works for authenticated servers - await connectWithSSETransport( + const sseTransport = await connectWithSSETransport( mcpClient, mcpServerConfig, await getStoredOAuthToken(mcpServerName), @@ -1532,7 +1574,7 @@ export async function connectToMcpServer( debugLogger.log( `MCP server '${mcpServerName}': Successfully connected using SSE transport.`, ); - return mcpClient; + return { client: mcpClient, transport: sseTransport }; } catch (sseFallbackError) { sseError = sseFallbackError as Error; @@ -1639,14 +1681,14 @@ export async function connectToMcpServer( ); } - await retryWithOAuth( + const oauthTransport = await retryWithOAuth( mcpClient, mcpServerName, mcpServerConfig, accessToken, httpReturned404, ); - return mcpClient; + return { client: mcpClient, transport: oauthTransport }; } else { throw new Error( `Failed to handle automatic OAuth for server '${mcpServerName}'`, @@ -1727,7 +1769,7 @@ export async function connectToMcpServer( timeout: mcpServerConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC, }); // Connection successful with OAuth - return mcpClient; + return { client: mcpClient, transport: oauthTransport }; } else { throw new Error( `OAuth configuration failed for '${mcpServerName}'. Please authenticate manually with /mcp auth ${mcpServerName}`, From d3c5a883392d69c796e16c437d53c622b1e70e7b Mon Sep 17 00:00:00 2001 From: Chris Coutinho Date: Sun, 1 Feb 2026 16:20:28 +0100 Subject: [PATCH 2/2] fix(mcp): remove redundant transport.close() call in connectAndDiscover The SDK's Client.close() method already calls transport.close() internally, making the explicit transport.close() call redundant. This addresses the reviewer feedback about potentially calling close() twice on the transport. --- packages/core/src/tools/mcp-client.ts | 8 -------- 1 file changed, 8 deletions(-) diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index f6bbdc38fa4..7d39ef31c45 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -973,14 +973,6 @@ export async function connectAndDiscover( } toolRegistry.sortTools(); } catch (error) { - // Close transport to prevent memory leaks - if (transport) { - try { - await transport.close(); - } catch { - // Ignore errors when closing transport - } - } if (mcpClient) { // eslint-disable-next-line @typescript-eslint/no-floating-promises mcpClient.close();