diff --git a/CHANGELOG.md b/CHANGELOG.md index 524cd157f..42f99a486 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,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. - workflow-streams: `WorkflowStream.onPoll` no longer serves a stale log index for a poll that was parked across a `truncate()` call, which could silently skip events. - Workflows no longer retain completion state when a child Workflow fails or is cancelled before starting. 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 });