improvement(mcp): dedupe create thunks into buildClient, harden dispose against liveness race

Author: waleedlatif1

apps/sim/lib/mcp/connection-pool.ts

@@ -17,9 +17,10 @@
  * disconnect the client themselves.
  *
  * Config-change invalidation is the caller's job via `evictServer` (wired to the
- * service's `clearCache`), not a per-acquire timestamp check — `updatedAt` is also
- * bumped by status telemetry, so it can't distinguish a config edit. Cross-process,
- * a config edit is bounded by max age (and self-heals when a stale connection errors).
+ * service's `evictServerConnections`), not a per-acquire timestamp check — `updatedAt`
+ * is also bumped by status telemetry, so it can't distinguish a config edit.
+ * Cross-process, a config edit is bounded by max age (self-heals when a stale
+ * connection errors).
  */
 
 import { createLogger } from '@sim/logger'
@@ -243,10 +244,13 @@ export class McpConnectionPool {
       clearInterval(this.idleCheckTimer)
       this.idleCheckTimer = null
     }
-    const clients = [...this.entries.values()].map((e) => e.client)
+    const entries = [...this.entries.values()]
+    // Mark retired before disconnecting so an in-flight liveness ping observes it
+    // and can't hand out a torn-down client (the retired-before-disconnect invariant).
+    for (const entry of entries) entry.retired = true
     this.entries.clear()
     this.pending.clear()
-    void Promise.allSettled(clients.map((client) => client.disconnect()))
+    void Promise.allSettled(entries.map((entry) => entry.client.disconnect()))
   }
 }
 

apps/sim/lib/mcp/service.ts

@@ -336,6 +336,30 @@ class McpService {
     return `${serverId}:${workspaceId ?? ''}:${userId ?? ''}`
   }
 
+  /**
+   * A `create` thunk for {@link withServerClient} that resolves env vars + SSRF-pins
+   * and connects. Deferred so a pool hit skips this work; `extraHeaders` (per-request)
+   * are merged in and force the caller to bypass the pool.
+   */
+  private buildClient(
+    config: McpServerConfig,
+    userId: string,
+    workspaceId: string,
+    extraHeaders?: Record<string, string>
+  ): () => Promise<McpClient> {
+    return async () => {
+      const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars(
+        config,
+        userId,
+        workspaceId
+      )
+      if (extraHeaders) {
+        resolvedConfig.headers = { ...resolvedConfig.headers, ...extraHeaders }
+      }
+      return this.createClient(resolvedConfig, resolvedIP, userId)
+    }
+  }
+
   /**
    * Run `fn` against a connected client. When `allowPool`, borrow from the warm
    * pool (`create` runs only on a miss, so a hit skips env resolution + DNS); a
@@ -399,25 +423,13 @@ class McpService {
         }
 
         const hasExtraHeaders = Boolean(extraHeaders && Object.keys(extraHeaders).length > 0)
-        const create = async () => {
-          const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars(
-            config,
-            userId,
-            workspaceId
-          )
-          if (hasExtraHeaders) {
-            resolvedConfig.headers = { ...resolvedConfig.headers, ...extraHeaders }
-          }
-          return this.createClient(resolvedConfig, resolvedIP, userId)
-        }
-
         const result = await this.withServerClient(
           {
             key: this.poolKey(serverId, workspaceId, userId),
             serverId,
             allowPool: !hasExtraHeaders,
           },
-          create,
+          this.buildClient(config, userId, workspaceId, hasExtraHeaders ? extraHeaders : undefined),
           (client) => client.callTool(toolCall)
         )
         logger.info(`[${requestId}] Successfully executed tool ${toolCall.name}`)
@@ -649,21 +661,13 @@ class McpService {
           }
 
           try {
-            const create = async () => {
-              const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars(
-                config,
-                userId,
-                workspaceId
-              )
-              return this.createClient(resolvedConfig, resolvedIP, userId)
-            }
             const tools = await this.withServerClient(
               {
                 key: this.poolKey(config.id, workspaceId, userId),
                 serverId: config.id,
                 allowPool: true,
               },
-              create,
+              this.buildClient(config, userId, workspaceId),
               (client) => client.listTools()
             )
             logger.debug(
@@ -870,21 +874,13 @@ class McpService {
         }
         authType = config.authType
 
-        const create = async () => {
-          const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars(
-            config,
-            userId,
-            workspaceId
-          )
-          return this.createClient(resolvedConfig, resolvedIP, userId)
-        }
         const tools = await this.withServerClient(
           {
             key: this.poolKey(serverId, workspaceId, userId),
             serverId,
             allowPool: true,
           },
-          create,
+          this.buildClient(config, userId, workspaceId),
           (client) => client.listTools()
         )
         logger.info(`[${requestId}] Discovered ${tools.length} tools from server ${config.name}`)
@@ -946,21 +942,13 @@ class McpService {
 
       for (const config of servers) {
         try {
-          const create = async () => {
-            const { config: resolvedConfig, resolvedIP } = await this.resolveConfigEnvVars(
-              config,
-              userId,
-              workspaceId
-            )
-            return this.createClient(resolvedConfig, resolvedIP, userId)
-          }
           const tools = await this.withServerClient(
             {
               key: this.poolKey(config.id, workspaceId, userId),
               serverId: config.id,
               allowPool: true,
             },
-            create,
+            this.buildClient(config, userId, workspaceId),
             (client) => client.listTools()
           )