From 19e0cde1d8df1e76457e16c3b31a1a44433eb00a Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Fri, 24 Jul 2026 15:06:18 -0700 Subject: [PATCH] fix(strands): don't evict an MCP connection with a call in flight The worker-process MCP connection cache armed its idle timer in `getConnection`, i.e. at the *start* of each call, and nothing extended it while the call ran. A `callTool` or `listTools` activity that outlived `mcpConnectionIdleTimeout` (5 minutes by default) had its transport disconnected underneath it, failing the call; since each retry was cut off the same way, an MCP tool slower than the idle window could never succeed. Track in-flight uses per connection, as the Python SDK already does: acquiring clears the pending idle timer, and releasing re-arms it only once no calls remain and the record is still the cached one, so the window measures genuine idleness. Eviction now also removes the entry from the cache before awaiting `disconnect()`, so a call arriving during teardown opens a fresh connection instead of joining the dying one. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 1 + contrib/strands/src/__tests__/test-strands.ts | 51 ++++++++ contrib/strands/src/plugin.ts | 9 +- contrib/strands/src/temporal-mcp-client.ts | 121 ++++++++++++------ 4 files changed, 140 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 38d1d6805..eb2b2f93f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ to docs, or any other relevant information. ### Fixed - strands: Declare `zod` as a peer dependency. +- strands: MCP connections are no longer disconnected while being used by a `callTool` or `listTools` activity. ## [1.21.1] - 2026-07-23 diff --git a/contrib/strands/src/__tests__/test-strands.ts b/contrib/strands/src/__tests__/test-strands.ts index ba8a5bea4..8e006ef61 100644 --- a/contrib/strands/src/__tests__/test-strands.ts +++ b/contrib/strands/src/__tests__/test-strands.ts @@ -418,6 +418,57 @@ test('mcpConnectionIdleTimeout evicts the cached connection while the worker run }); }); +test('mcpConnectionIdleTimeout does not evict a connection with a call in flight', async (t) => { + const { createWorker, executeWorkflow } = helpers(t); + + const counters = { disconnects: 0, disconnectsDuringCall: 0 }; + let callInFlight = false; + class SlowMcpClient extends StubMcpClient { + override async disconnect(): Promise { + counters.disconnects++; + if (callInFlight) counters.disconnectsDuringCall++; + } + override async callTool(tool: { name: string }, args: JSONValue): Promise { + callInFlight = true; + try { + // Outlast the idle window several times over. + await new Promise((resolve) => setTimeout(resolve, 500)); + return await super.callTool(tool, args); + } finally { + callInFlight = false; + } + } + } + const factory = (): McpClient => + new SlowMcpClient([ + { name: 'listFiles', description: 'List files', inputSchema: { type: 'object', properties: {} } }, + ]); + + const worker = await createWorker({ + plugins: [ + new StrandsPlugin({ + models: { + test: () => new StubModel([toolCallTurn('listFiles', 'call_1', { path: '/' }), textTurn('done')]), + }, + // Distinct server name so the worker-process connection cache can't be + // shared with the sibling MCP tests running in the same process. + mcpClients: { slowServer: factory }, + // Idle window far shorter than the tool call it must not interrupt. + mcpConnectionIdleTimeout: '50ms', + }), + ], + }); + + await worker.runUntil(async () => { + // The tool call outlives the idle window, so an idle timer measured from the + // start of the call (rather than from the call completing) would disconnect + // the transport underneath it and fail the call. + const result = await executeWorkflow(mcpAgent, { args: ['list files', 'slowServer'] }); + t.is(result, 'done'); + t.is(counters.disconnectsDuringCall, 0, 'idle timer disconnected a connection with a call in flight'); + }); +}); + test('TemporalMCPClient re-lists tools each turn by default and picks up mid-run additions', async (t) => { const { createWorker, executeWorkflow } = helpers(t); diff --git a/contrib/strands/src/plugin.ts b/contrib/strands/src/plugin.ts index cc8f3742f..7ed5f4e40 100644 --- a/contrib/strands/src/plugin.ts +++ b/contrib/strands/src/plugin.ts @@ -33,10 +33,11 @@ import { * true` to list once at the start of the workflow instead. * * - `mcpConnectionIdleTimeout` — how long a worker-process MCP connection is - * kept open between `callTool` activities before it's disconnected. The timer - * resets on every reuse. Accepts a millisecond number or a duration string - * (e.g. `'5 minutes'`), like `startToCloseTimeout`. Defaults to - * {@link MCP_CONNECTION_IDLE_MS} (5 minutes). + * kept open between `callTool` activities before it's disconnected. The window + * only starts once no activity is using the connection, so it never cuts off + * an in-flight call, however long it runs. Accepts a millisecond number or a + * duration string (e.g. `'5 minutes'`), like `startToCloseTimeout`. Defaults + * to {@link MCP_CONNECTION_IDLE_MS} (5 minutes). */ export interface StrandsPluginOptions { models?: Record Model>; diff --git a/contrib/strands/src/temporal-mcp-client.ts b/contrib/strands/src/temporal-mcp-client.ts index c4169ce89..c3af85fe2 100644 --- a/contrib/strands/src/temporal-mcp-client.ts +++ b/contrib/strands/src/temporal-mcp-client.ts @@ -161,8 +161,9 @@ export function buildListToolsActivity( ): (input: { server: string }) => Promise { const name = listToolsActivityName(server); const fn = async (_input: { server: string }): Promise => { - const client = await getConnection(server, factory, idleMs); + const record = acquireConnection(server, factory); try { + const client = await record.client; const tools = await client.listTools(); return tools.map((t) => ({ name: t.name, @@ -174,6 +175,10 @@ export function buildListToolsActivity( // the next call reconnects to the current deployment. await _evictConnection(server); throw err; + } finally { + // No more in-flight call on this connection; let idle eviction resume + // (no-op if the connection was just evicted above). + releaseConnection(server, record, idleMs); } }; Object.defineProperty(fn, 'name', { value: name }); @@ -181,61 +186,96 @@ export function buildListToolsActivity( } /** - * Default for how long an idle MCP connection is kept open before it's - * disconnected. The timer resets on every {@link callToolActivityName | - * callTool} that reuses the connection. Override per worker via - * `StrandsPlugin({ mcpConnectionIdleMs })`. + * Default for how long an *idle* MCP connection is kept open before it's + * disconnected. The window only starts once no {@link callToolActivityName | + * callTool} or {@link listToolsActivityName | listTools} activity is using the + * connection, so a call that runs longer than this is never cut off. Override + * per worker via `StrandsPlugin({ mcpConnectionIdleTimeout })`. */ export const MCP_CONNECTION_IDLE_MS = 5 * 60 * 1000; -// Worker-process caches, keyed by server name. Activities run in the worker's -// Node process, so this module state is shared across every activity invocation on the worker. -const CONNECTIONS: Map> = new Map(); -const IDLE_TIMERS: Map = new Map(); - -function resetIdleTimer(server: string, idleMs: number): void { - const existing = IDLE_TIMERS.get(server); - if (existing !== undefined) clearTimeout(existing); - const timer = setTimeout(() => void _evictConnection(server), idleMs); - // Don't let an idle MCP connection keep the worker process alive. - timer.unref?.(); - IDLE_TIMERS.set(server, timer); +/** A worker-process MCP connection plus the state governing its idle eviction. */ +interface ConnectionRecord { + /** Resolves to the connected client, or rejects if the handshake failed. */ + client: Promise; + /** Number of activities currently using this connection. */ + inflight: number; + /** Armed only while `inflight` is 0. */ + idleTimer?: NodeJS.Timeout; } +// Worker-process cache, keyed by server name. Activities run in the worker's +// Node process, so this module state is shared across every activity invocation on the worker. +const CONNECTIONS: Map = new Map(); + export async function _evictConnection(server: string): Promise { - const timer = IDLE_TIMERS.get(server); - if (timer !== undefined) { - clearTimeout(timer); - IDLE_TIMERS.delete(server); - } - const entry = CONNECTIONS.get(server); - if (entry === undefined) return; + const record = CONNECTIONS.get(server); + if (record === undefined) return; + // Delete before awaiting: a call arriving during `disconnect` must open a + // fresh connection rather than join the one being torn down. CONNECTIONS.delete(server); + if (record.idleTimer !== undefined) { + clearTimeout(record.idleTimer); + record.idleTimer = undefined; + } try { - const client = await entry; + const client = await record.client; await client.disconnect(); } catch { // Best-effort; the session may already be broken. } } -function getConnection(server: string, factory: () => McpClient, idleMs: number): Promise { - let entry = CONNECTIONS.get(server); - if (entry === undefined) { - entry = (async () => { - const client = factory(); - await client.connect(); - return client; - })(); - CONNECTIONS.set(server, entry); +/** + * Return the cached connection for `server`, opening one lazily if needed, and + * mark a call in flight so idle eviction can't disconnect underneath it. + * Concurrent first-callers dedupe onto a single connect handshake by awaiting + * the same record. Every caller must {@link releaseConnection} once its call + * completes. + */ +function acquireConnection(server: string, factory: () => McpClient): ConnectionRecord { + let record = CONNECTIONS.get(server); + if (record === undefined) { + const entry: ConnectionRecord = { + inflight: 0, + client: (async () => { + const client = factory(); + await client.connect(); + return client; + })(), + }; // If the connect handshake fails, drop the rejected promise so the next // call retries instead of caching the failure forever. - entry.catch(() => { + entry.client.catch(() => { if (CONNECTIONS.get(server) === entry) CONNECTIONS.delete(server); }); + CONNECTIONS.set(server, entry); + record = entry; + } + record.inflight += 1; + if (record.idleTimer !== undefined) { + clearTimeout(record.idleTimer); + record.idleTimer = undefined; } - resetIdleTimer(server, idleMs); - return entry; + return record; +} + +/** Mark a call done; arm idle eviction once no calls remain in flight. */ +function releaseConnection(server: string, record: ConnectionRecord, idleMs: number): void { + record.inflight -= 1; + // Only the record still cached under this server arms a timer; a record + // already evicted must not schedule one, or it could later evict a + // different, healthy connection for the same server. + if (record.inflight > 0 || CONNECTIONS.get(server) !== record) return; + const timer = setTimeout(() => { + // `acquireConnection` clears this timer synchronously, so a pending timer + // implies the connection is still idle; re-checked to keep that invariant + // local. + if (record.inflight === 0) void _evictConnection(server); + }, idleMs); + // Don't let an idle MCP connection keep the worker process alive. + timer.unref?.(); + record.idleTimer = timer; } /** @@ -255,18 +295,23 @@ export function buildCallToolActivity( ActivityContext.current().log.debug(`Calling MCP tool ${input.toolName} on server ${server}`, { activityId: activityInfo().activityId, }); - const client = await getConnection(server, factory, idleMs); + const record = acquireConnection(server, factory); // Dispatch by name without re-listing tools on the connection: the agent // already has the schema from `listTools`, and `McpClient.callTool` only // reads `tool.name` to build the request. A minimal `{ name }` matches the // by-name dispatch the Python SDK does via `session.call_tool`. const tool = { name: input.toolName } as unknown as Parameters[0]; try { + const client = await record.client; return await client.callTool(tool, input.args); } catch (err) { // The session may be broken; drop it so the next call reconnects. await _evictConnection(server); throw err; + } finally { + // No more in-flight call on this connection; let idle eviction resume + // (no-op if the connection was just evicted above). + releaseConnection(server, record, idleMs); } }; Object.defineProperty(fn, 'name', { value: name });