From 62cc7eb9a74d0f6eb9040b5df2fe4f6ece87752e Mon Sep 17 00:00:00 2001
From: tru3 <tru3@tru3.com>
Date: Wed, 22 Jul 2026 11:57:59 +0200
Subject: [PATCH] fix(mcp): fix resource leak and stale-header bug in MCP
 client

- Fix hashConfig to include headers for HTTP/SSE configs, preventing
  stale auth tokens when two MCP servers differ only by headers.

- Add closeMCPClient() export to clean up client connections and
  cached tool listings, fixing the memory/connection leak where
  runningClients and listToolsCache grew without bound.

- Integrate automatic cleanup in SDK: when an MCP tool call fails,
  the dead client is evicted so the next call reconnects with a
  fresh transport instead of reusing the broken connection.
---
 common/src/mcp/client.ts | 24 ++++++++++++++++++++++++
 sdk/src/run.ts           |  9 ++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/common/src/mcp/client.ts b/common/src/mcp/client.ts
index 5a5608d57f..e9bf89d994 100644
--- a/common/src/mcp/client.ts
+++ b/common/src/mcp/client.ts
@@ -65,6 +65,7 @@ function hashConfig(config: MCPConfig): string {
       type: 'http',
       url: config.url,
       params: config.params,
+      headers: config.headers,
     })
   }
   if (config.type === 'sse') {
@@ -72,6 +73,7 @@ function hashConfig(config: MCPConfig): string {
       type: 'sse',
       url: config.url,
       params: config.params,
+      headers: config.headers,
     })
   }
   config.type satisfies never
@@ -231,3 +233,25 @@ export async function callMCPTool(
     } satisfies ToolResultOutput
   })
 }
+
+/**
+ * Close an MCP client connection and clean up associated resources.
+ *
+ * Removes the client from the running clients registry and clears any cached
+ * tool listings. The underlying transport (stdio child process, HTTP/SSE
+ * connection) is also torn down via the MCP SDK's client.close() method.
+ *
+ * Safe to call on an already-closed or unknown client id (no-ops silently).
+ */
+export async function closeMCPClient(clientId: string): Promise<void> {
+  const client = runningClients[clientId]
+  if (client) {
+    try {
+      await client.close()
+    } catch {
+      // Ignore errors during close — the transport may already be dead.
+    }
+    delete runningClients[clientId]
+  }
+  delete listToolsCache[clientId]
+}
diff --git a/sdk/src/run.ts b/sdk/src/run.ts
index 4ed4f1321f..9b190deaf5 100644
--- a/sdk/src/run.ts
+++ b/sdk/src/run.ts
@@ -11,6 +11,7 @@ import {
   getMCPClient,
   listMCPTools,
   callMCPTool,
+  closeMCPClient,
 } from '@codebuff/common/mcp/client'
 import {
   COMPOSIO_META_TOOL_NAMES,
@@ -809,8 +810,9 @@ async function handleToolCall({
 
   // Handle MCP tool calls when mcpConfig is present
   if (action.mcpConfig) {
+    let mcpClientId: string | undefined
     try {
-      const mcpClientId = await getMCPClient(action.mcpConfig)
+      mcpClientId = await getMCPClient(action.mcpConfig)
       const result = await callMCPTool(
         mcpClientId,
         {
@@ -822,6 +824,11 @@ async function handleToolCall({
       )
       return { output: result }
     } catch (error) {
+      // Clean up the dead client so the next call reconnects with a fresh
+      // transport rather than hitting the same broken connection.
+      if (mcpClientId) {
+        closeMCPClient(mcpClientId).catch(() => {})
+      }
       return {
         output: [
           {