Skip to content

Commit 6f96d4a

Browse files
authored
fix(mcp): retire pooled connections after consecutive request timeouts (#5821)
* fix(mcp): retire pooled connections after consecutive request timeouts Circuit breaker for the half-open-transport gap left by #5817: a lone timeout still keeps the session warm (retiring on every timeout caused the connect/stall/reconnect churn), but two consecutive timeouts with no healthy request in between retire the connection so a genuinely half-open transport can't serve repeated 30s failures until the liveness ping catches it. A healthy release resets the count. Matches the LibreChat pattern of counting connection-level failures rather than reacting to one. * fix(mcp): classify TimeoutError-named aborts as timeouts for the circuit breaker AbortSignal.timeout / undici surface a DOMException named TimeoutError whose message lacks 'timed out'; without this the breaker treated it as a healthy release and reset the streak. Adds name-based detection (incl. cause) + test.
1 parent 9fb9ac5 commit 6f96d4a

4 files changed

Lines changed: 108 additions & 12 deletions

File tree

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,50 @@ describe('McpConnectionPool', () => {
8383
expect(client.disconnect).not.toHaveBeenCalled()
8484
})
8585

86+
it('keeps the connection after a single timeout release', async () => {
87+
const client = makeFakeClient()
88+
const create = vi.fn(async () => client)
89+
90+
const lease = await pool.acquire(params('s1:w1:u1', create))
91+
await lease.release(false, true)
92+
await borrow(pool, params('s1:w1:u1', create))
93+
94+
expect(create).toHaveBeenCalledTimes(1)
95+
expect(client.disconnect).not.toHaveBeenCalled()
96+
})
97+
98+
it('retires the connection after consecutive timeouts (circuit breaker)', async () => {
99+
const client = makeFakeClient()
100+
const replacement = makeFakeClient()
101+
const create = vi.fn<() => Promise<McpClient>>()
102+
create.mockResolvedValueOnce(client)
103+
create.mockResolvedValue(replacement)
104+
105+
const l1 = await pool.acquire(params('s1:w1:u1', create))
106+
await l1.release(false, true)
107+
const l2 = await pool.acquire(params('s1:w1:u1', create))
108+
await l2.release(false, true)
109+
110+
expect(client.disconnect).toHaveBeenCalledTimes(1)
111+
const l3 = await pool.acquire(params('s1:w1:u1', create))
112+
expect(l3.client).toBe(replacement)
113+
await l3.release()
114+
})
115+
116+
it('resets the timeout count on a healthy release', async () => {
117+
const client = makeFakeClient()
118+
const create = vi.fn(async () => client)
119+
120+
const l1 = await pool.acquire(params('s1:w1:u1', create))
121+
await l1.release(false, true)
122+
await borrow(pool, params('s1:w1:u1', create)) // healthy — resets the count
123+
const l3 = await pool.acquire(params('s1:w1:u1', create))
124+
await l3.release(false, true) // first of a new streak, not the second strike
125+
126+
expect(client.disconnect).not.toHaveBeenCalled()
127+
expect(create).toHaveBeenCalledTimes(1)
128+
})
129+
86130
it('dedups concurrent creates into a single connect (single-flight)', async () => {
87131
let resolveCreate: ((client: McpClient) => void) | undefined
88132
const created = new Promise<McpClient>((resolve) => {

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

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,14 @@ const logger = createLogger('McpConnectionPool')
3232
const MAX_POOL_SIZE = 100
3333
/** Max lifetime; on expiry the pinned SSRF `resolvedIP` re-resolves. */
3434
const MAX_CONNECTION_AGE_MS = 10 * 60 * 1000
35+
/**
36+
* Circuit breaker: retire a connection after this many CONSECUTIVE request
37+
* timeouts. One timeout is a slow request and keeps the session (retiring on
38+
* every timeout causes connect/stall/reconnect churn); repeated timeouts with
39+
* no success in between indicate a half-open transport the liveness ping
40+
* hasn't caught yet. A healthy release resets the count.
41+
*/
42+
const MAX_CONSECUTIVE_TIMEOUTS = 2
3543
const LIVENESS_TTL_MS = 60 * 1000
3644
const LIVENESS_PING_TIMEOUT_MS = 5 * 1000
3745
const IDLE_TIMEOUT_MS = 5 * 60 * 1000
@@ -46,10 +54,14 @@ export interface AcquireParams {
4654
create: () => Promise<McpClient>
4755
}
4856

49-
/** A borrowed connection. `release(poison)` must be called exactly once; pass `poison: true` to retire on failure. */
57+
/**
58+
* A borrowed connection. `release(poison, sawTimeout)` must be called exactly once:
59+
* `poison: true` retires immediately (dead-connection error); `sawTimeout: true`
60+
* counts toward the consecutive-timeout circuit breaker without retiring on its own.
61+
*/
5062
export interface ConnectionLease {
5163
client: McpClient
52-
release: (poison?: boolean) => Promise<void>
64+
release: (poison?: boolean, sawTimeout?: boolean) => Promise<void>
5365
}
5466

5567
interface PoolEntry {
@@ -60,6 +72,7 @@ interface PoolEntry {
6072
lastActivityAt: number
6173
lastLivenessCheckAt: number
6274
borrowers: number
75+
consecutiveTimeouts: number
6376
retired: boolean
6477
closing: boolean
6578
}
@@ -95,7 +108,10 @@ export class McpConnectionPool {
95108
while (entry.closing) entry = await this.resolveEntry(params)
96109
entry.borrowers++
97110
entry.lastActivityAt = Date.now()
98-
return { client: entry.client, release: (poison) => this.release(entry, poison ?? false) }
111+
return {
112+
client: entry.client,
113+
release: (poison, sawTimeout) => this.release(entry, poison ?? false, sawTimeout ?? false),
114+
}
99115
}
100116

101117
private async resolveEntry(params: AcquireParams): Promise<PoolEntry> {
@@ -162,6 +178,7 @@ export class McpConnectionPool {
162178
lastActivityAt: now,
163179
lastLivenessCheckAt: now,
164180
borrowers: 0,
181+
consecutiveTimeouts: 0,
165182
retired: false,
166183
closing: false,
167184
}
@@ -194,9 +211,17 @@ export class McpConnectionPool {
194211
return record.promise
195212
}
196213

197-
private async release(entry: PoolEntry, poison: boolean): Promise<void> {
214+
private async release(entry: PoolEntry, poison: boolean, sawTimeout: boolean): Promise<void> {
198215
entry.borrowers = Math.max(0, entry.borrowers - 1)
199216
entry.lastActivityAt = Date.now()
217+
if (sawTimeout) {
218+
entry.consecutiveTimeouts++
219+
if (entry.consecutiveTimeouts >= MAX_CONSECUTIVE_TIMEOUTS) {
220+
this.retire(entry, `${entry.consecutiveTimeouts} consecutive request timeouts`)
221+
}
222+
} else if (!poison) {
223+
entry.consecutiveTimeouts = 0
224+
}
200225
if (poison) this.retire(entry, 'poisoned by failed operation')
201226
await this.closeIfIdle(entry)
202227
}

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

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ describe('McpService connection reuse wiring', () => {
134134
expect.objectContaining({ key: `server-1:${WORKSPACE_ID}:${USER_ID}`, serverId: 'server-1' })
135135
)
136136
expect(mockCallTool).toHaveBeenCalledTimes(1)
137-
expect(mockRelease).toHaveBeenCalledWith(false)
137+
expect(mockRelease).toHaveBeenCalledWith(false, false)
138138
expect(poolClient.disconnect).not.toHaveBeenCalled()
139139
// A pool hit must not re-resolve env vars (acquire never invoked `create`).
140140
expect(mockResolveEnvVars).not.toHaveBeenCalled()
@@ -158,7 +158,7 @@ describe('McpService connection reuse wiring', () => {
158158
mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID)
159159
).rejects.toThrow()
160160

161-
expect(mockRelease).toHaveBeenCalledWith(true)
161+
expect(mockRelease).toHaveBeenCalledWith(true, false)
162162
})
163163

164164
it('keeps the pooled connection warm on a benign (non-connection) tool error', async () => {
@@ -168,20 +168,37 @@ describe('McpService connection reuse wiring', () => {
168168
mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID)
169169
).rejects.toThrow()
170170

171-
expect(mockRelease).toHaveBeenCalledWith(false)
171+
expect(mockRelease).toHaveBeenCalledWith(false, false)
172172
})
173173

174174
it('keeps the pooled connection warm on a request timeout (does not retire the session)', async () => {
175175
// A streamable-HTTP request timeout aborts only that request's stream; the session stays
176176
// healthy for the next request, so a timeout must NOT poison the lease (matches every
177-
// production MCP client and avoids a connect/stall/reconnect churn loop).
177+
// production MCP client and avoids a connect/stall/reconnect churn loop). It DOES report
178+
// sawTimeout so the pool's consecutive-timeout circuit breaker can retire a half-open
179+
// transport after repeated strikes.
178180
mockCallTool.mockRejectedValue(new Error('Request timed out'))
179181

180182
await expect(
181183
mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID)
182184
).rejects.toThrow()
183185

184-
expect(mockRelease).not.toHaveBeenCalledWith(true)
186+
expect(mockRelease).toHaveBeenCalledWith(false, true)
187+
})
188+
189+
it('classifies an AbortSignal.timeout-shaped TimeoutError as a timeout for the breaker', async () => {
190+
// DOMException name 'TimeoutError' with a message that lacks "timed out".
191+
mockCallTool.mockRejectedValue(
192+
Object.assign(new Error('The operation was aborted due to timeout'), {
193+
name: 'TimeoutError',
194+
})
195+
)
196+
197+
await expect(
198+
mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID)
199+
).rejects.toThrow()
200+
201+
expect(mockRelease).toHaveBeenCalledWith(false, true)
185202
})
186203

187204
it('poisons the lease on an auth failure so a rotated credential is re-resolved', async () => {
@@ -191,7 +208,7 @@ describe('McpService connection reuse wiring', () => {
191208
mcpService.executeTool(USER_ID, 'server-1', { name: 'do', arguments: {} }, WORKSPACE_ID)
192209
).rejects.toThrow()
193210

194-
expect(mockRelease).toHaveBeenCalledWith(true)
211+
expect(mockRelease).toHaveBeenCalledWith(true, false)
195212
})
196213

197214
it('retries and recovers when a rotated credential causes a one-off auth failure', async () => {
@@ -209,7 +226,7 @@ describe('McpService connection reuse wiring', () => {
209226
expect(result).toEqual({ content: [] })
210227
expect(mockCallTool).toHaveBeenCalledTimes(2)
211228
// First attempt poisoned the stale lease; the retry re-acquired a fresh one.
212-
expect(mockRelease).toHaveBeenCalledWith(true)
229+
expect(mockRelease).toHaveBeenCalledWith(true, false)
213230
expect(mockAcquire).toHaveBeenCalledTimes(2)
214231
})
215232
})

apps/sim/lib/mcp/service.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,12 @@ function isTimeoutError(error: unknown): boolean {
9393
if (error instanceof McpError && error.code === ErrorCode.RequestTimeout) {
9494
return true
9595
}
96+
// AbortSignal.timeout / undici surface a DOMException named TimeoutError whose
97+
// message ("The operation was aborted due to timeout") lacks "timed out".
98+
const e = error as { name?: string; cause?: { name?: string } } | null
99+
if (e?.name === 'TimeoutError' || e?.cause?.name === 'TimeoutError') {
100+
return true
101+
}
96102
return getErrorMessage(error, '').toLowerCase().includes('timed out')
97103
}
98104

@@ -425,13 +431,17 @@ class McpService {
425431
create,
426432
})
427433
let poison = false
434+
let sawTimeout = false
428435
try {
429436
return await fn(lease.client)
430437
} catch (error) {
431438
poison = isDeadConnectionError(error)
439+
// A lone timeout keeps the session; the pool's circuit breaker retires it
440+
// after consecutive timeouts with no healthy request in between.
441+
sawTimeout = isTimeoutError(error)
432442
throw error
433443
} finally {
434-
await lease.release(poison)
444+
await lease.release(poison, sawTimeout)
435445
}
436446
}
437447

0 commit comments

Comments
 (0)