fix(mcp): retry auth failures in executeTool; simplify acquire re-check and clear generations on dispose

Author: waleedlatif1

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

@@ -54,6 +54,11 @@ function params(key: string, create: () => Promise<McpClient>): AcquireParams {
   return { key, serverId: key.split(':')[0], create }
 }
 
+/** Drain the microtask queue so an in-flight acquire fully settles before the next step. */
+async function flushMicrotasks(turns = 50): Promise<void> {
+  for (let i = 0; i < turns; i++) await Promise.resolve()
+}
+
 describe('McpConnectionPool', () => {
   let pool: McpConnectionPool
 
@@ -310,7 +315,7 @@ describe('McpConnectionPool', () => {
 
     // A's ping fails → A retires stale, rebuilds `fresh`, pools it, clears pending.
     releasePing[0](false)
-    for (let i = 0; i < 20; i++) await Promise.resolve()
+    await flushMicrotasks()
     // B's ping fails → B must reuse the pooled `fresh`, not create a third client.
     releasePing[1](false)
 

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

@@ -9,7 +9,7 @@
  * ping-cached; reconnect is demand-driven (callers' retries re-acquire).
  *
  * Connections are ref-counted. Eviction *retires* an entry — removes it from the
- * pool so no new borrower can take it — but the socket is closed only once the
+ * pool so later acquires stop reusing it — but the socket is closed only once the
  * last in-flight borrower releases, so a sibling failure, idle sweep, or config
  * change never tears down a connection mid-request. Retirement is identity-checked,
  * so a stale `onClose` can't disconnect a replacement stored under the same key.
@@ -78,30 +78,29 @@ export class McpConnectionPool {
       const client = await params.create()
       return { client, release: () => client.disconnect().catch(() => {}) }
     }
-    const entry = await this.resolveEntry(params)
-    // A concurrent evict/release/idle could have retired + started disconnecting
-    // this entry during the resolve's `await` gap; `closing` (not `retired`, which
-    // a usable one-shot also sets) means the socket is going away — get a fresh one.
-    if (entry.closing) return this.acquire(params)
+    // A concurrent evict/release/idle could have started disconnecting the resolved
+    // entry during an `await` gap; `closing` (not `retired`, which a usable one-shot
+    // also sets) means the socket is going away, so resolve again.
+    let entry = await this.resolveEntry(params)
+    while (entry.closing) entry = await this.resolveEntry(params)
     entry.borrowers++
     entry.lastActivityAt = Date.now()
     return { client: entry.client, release: (poison) => this.release(entry, poison ?? false) }
   }
 
   private async resolveEntry(params: AcquireParams): Promise<PoolEntry> {
-    const pending = this.pending.get(params.key)
-    if (pending) return pending
+    while (true) {
+      const pending = this.pending.get(params.key)
+      if (pending) return pending
+
+      const current = this.entries.get(params.key)
+      if (!current) return this.createEntry(params)
 
-    const current = this.entries.get(params.key)
-    if (current) {
       const reusable = await this.tryReuse(current)
       if (reusable) return reusable
-      // tryReuse awaited a ping and retired `current`; a concurrent acquire may
-      // have pooled a replacement meanwhile, so re-resolve rather than blindly
-      // creating (which would overwrite and leak that replacement).
-      return this.resolveEntry(params)
+      // tryReuse awaited a ping and retired `current`; loop to re-check pending/entries —
+      // a concurrent acquire may have pooled a replacement we must reuse, not overwrite.
     }
-    return this.createEntry(params)
   }
 
   /** Return `entry` if in-age, connected, and live; else retire it and return null. */
@@ -258,6 +257,7 @@ export class McpConnectionPool {
     for (const entry of entries) entry.retired = true
     this.entries.clear()
     this.pending.clear()
+    this.serverGenerations.clear()
     void Promise.allSettled(entries.map((entry) => entry.client.disconnect()))
   }
 }

apps/sim/lib/mcp/service-pool.test.ts

@@ -180,4 +180,23 @@ describe('McpService connection reuse wiring', () => {
 
     expect(mockRelease).toHaveBeenCalledWith(true)
   })
+
+  it('retries and recovers when a rotated credential causes a one-off auth failure', async () => {
+    mockCallTool
+      .mockRejectedValueOnce(new UnauthorizedError('stale key'))
+      .mockResolvedValueOnce({ content: [] })
+
+    const result = await mcpService.executeTool(
+      USER_ID,
+      'server-1',
+      { name: 'do', arguments: {} },
+      WORKSPACE_ID
+    )
+
+    expect(result).toEqual({ content: [] })
+    expect(mockCallTool).toHaveBeenCalledTimes(2)
+    // First attempt poisoned the stale lease; the retry re-acquired a fresh one.
+    expect(mockRelease).toHaveBeenCalledWith(true)
+    expect(mockAcquire).toHaveBeenCalledTimes(2)
+  })
 })

apps/sim/lib/mcp/service.ts

@@ -125,6 +125,18 @@ function isDeadConnectionError(error: unknown): boolean {
   )
 }
 
+/**
+ * An auth failure (401) from a rotated/revoked credential. Safe for `executeTool`
+ * to retry — auth is rejected *before* the tool runs, so re-acquiring on a fresh
+ * connection (which re-resolves the credential) can't double-execute a tool.
+ */
+function isAuthError(error: unknown): boolean {
+  return (
+    error instanceof UnauthorizedError ||
+    (error instanceof StreamableHTTPError && error.code === 401)
+  )
+}
+
 /** Transient failures a read-only `tools/list` may safely retry (idempotent, unlike `tools/call`); excludes OAuth and terminal 4xx. */
 function isRetryableDiscoveryError(error: unknown): boolean {
   if (isTimeoutError(error)) return true
@@ -435,9 +447,12 @@ class McpService {
         logger.info(`[${requestId}] Successfully executed tool ${toolCall.name}`)
         return result
       } catch (error) {
-        if (this.isSessionError(error) && attempt < maxRetries - 1) {
+        // A stale session (400/404) or a rotated/revoked credential (401) is rejected
+        // before the tool runs, so retrying on a fresh connection is safe and recovers
+        // the request. Timeouts/resets are NOT retried — the tool may have executed.
+        if ((this.isSessionError(error) || isAuthError(error)) && attempt < maxRetries - 1) {
           logger.warn(
-            `[${requestId}] Session error executing tool ${toolCall.name}, retrying (attempt ${attempt + 1}):`,
+            `[${requestId}] Retryable connection error executing tool ${toolCall.name}, retrying (attempt ${attempt + 1}):`,
             error
           )
           await sleep(100)