fix(contrib): pool and idle-evict MCP connections in google_adk_agents#1664
fix(contrib): pool and idle-evict MCP connections in google_adk_agents#1664wankhede04 wants to merge 6 commits into
Conversation
TemporalMcpToolSetProvider's list-tools and call-tool activities called self._toolset_factory(...) on every invocation, constructing a brand-new McpToolset (and therefore a new MCPSessionManager/subprocess for stdio servers) each time with no cleanup. This leaked a spawned process per activity execution under sustained load. Pool one McpToolset per activity name, reused across calls and refcounted so idle eviction (default 5 minutes, overridable via the new mcp_connection_idle_timeout constructor parameter) only fires once no calls are in flight. A failed get_tools()/run_async() call evicts the connection so the next call reconnects instead of reusing a dead session. This brings google_adk_agents to parity with the pooling already shipped in the strands and google_genai contribs. Closes temporalio#1663
…viction in google_adk_agents Regression tests against a fake McpToolset asserting: N sequential call_tool executions against the same name reuse one toolset instead of creating N; list-tools and call-tool share a connection; idle connections close after the configured timeout but not while a call is in flight; and a failed call evicts the broken connection so the next call reconnects.
|
|
brianstrauch
left a comment
brianstrauch
left a comment
There was a problem hiding this comment.
Thanks for tackling this — the underlying bug is real (a fresh McpToolset created on every activity invocation and never close()d is a genuine subprocess/session leak for stdio servers, not just a missed optimization).
My concern is the pooling model, specifically its interaction with factory_argument, which I think makes this a breaking behavior change rather than a pure fix.
The problem: name-only, cross-workflow pooling silently breaks factory_argument
factory_argument is a per-call value threaded from the workflow into toolset_factory(args.factory_argument) on every invocation. The new pool is keyed by activity name and shared across all workflow executions on the worker, so the argument is only consulted the first time a connection opens. Every later call — including calls from a different workflow passing a different factory_argument — silently reuses that first connection. For any caller using factory_argument to select a tenant / credential / backend, that's silent mis-routing, which is worse than a hard error. That's a behavior change to a public (if experimental) API surface.
On the parity argument
The PR justifies ignoring factory_argument by matching strands / google_genai. But those two contribs have no per-call factory argument at all (Callable[[], ...]), so "one connection per name" is lossless there and lossy here — it isn't really parity. The contrib that actually shares this factory_argument API is openai_agents, and it deliberately avoids this exact problem by scoping connections so the argument never leaks:
- stateless (
StatelessMCPServerProvider): no pool — create + connect +cleanup()per call;factory_argumenthonored on every call. - stateful (
StatefulMCPServerProvider): pooled, but keyed per workflow run (name@run_id) and consumingfactory_argumentonce per run; never shared across runs.
Request
Please match the openai_agents stateless/stateful model here rather than the strands / google_genai worker-process cross-workflow pool. That still fixes the leak, but keeps factory_argument's per-call/per-run semantics intact and keeps this contrib consistent with the sibling that shares its API surface.
If there's a use case for cross-workflow connection sharing that I'm missing, happy to discuss — but as written, the factory_argument behavior change is the blocker.
…ctory_argument mis-routing) The previous fix pooled a single McpToolset per provider name in a worker-process-wide dict shared across all workflow executions, and only consulted factory_argument the first time a connection opened for that name. Every later call -- including calls from a different workflow run passing a different factory_argument -- silently reused that first connection, causing silent mis-routing for callers that use factory_argument to select a tenant/credential/backend. Remove the cross-workflow pool (_ConnectionRecord, _CONNECTIONS, get_connection, _evict_connection, and the mcp_connection_idle_timeout parameter) and make TemporalMcpToolSetProvider stateless, mirroring openai_agents' StatelessMCPServerProvider: each list-tools/call-tool activity builds a fresh McpToolset via toolset_factory(factory_argument), runs the operation, and always closes it in a finally block. This fixes the MCP session/subprocess leak from temporalio#1663 while honoring factory_argument on every call with zero cross-workflow sharing. Tests rewritten accordingly (test_mcp_pool.py -> test_mcp.py): prove a fresh toolset per call, close() on every path (success, error, no-match), and that a later call with a different factory_argument routes with its own argument. Fixes temporalio#1663
…k_agents
Add an opt-in TemporalStatefulMcpToolSetProvider (plus the workflow-side
TemporalStatefulMcpToolSet handle) mirroring openai_agents'
StatefulMCPServerProvider, for callers who genuinely need a persistent MCP
connection reused across tool calls within a single workflow run.
The workflow-side handle, used as an async context manager, starts a dedicated
{name}-server-session activity on a task queue scoped to the specific run
(name@run_id). That activity builds the McpToolset once via
toolset_factory(factory_argument), holds it open, and runs a nested Worker
(PollerBehaviorSimpleMaximum(1)) serving the run-scoped -list-tools/-call-tool
activities. The toolset is closed in a finally when the workflow cancels the
session handle on cleanup. A heartbeat loop lets the workflow detect a dead
dedicated worker; schedule-to-start and heartbeat timeouts surface as
ApplicationError(type="DedicatedWorkerFailure") via a _handle_worker_failure
decorator.
This honors factory_argument exactly once per run with zero cross-run sharing,
so it carries none of the silent mis-routing risk of a worker-wide pool, while
offering connection reuse the stateless provider intentionally forgoes. The
GoogleAdkPlugin now accepts either provider type.
Adds CI-safe integration tests driving the real connect -> dedicated-worker ->
get_tools path against an in-memory fake toolset (no subprocess): one toolset
per run, factory_argument consumed once, no cross-run sharing, teardown on
completion.
…-session activity The dedicated -server-session activity created its heartbeat task before the duplicate-connect guard, and only cancelled it inside the same finally block as toolset.close(). Two paths could leak the heartbeat task, leaving it calling activity.heartbeat() forever after the activity had already exited: 1. A duplicate connect() for an already-running server_id raised before the try/finally that cancels the heartbeat task was ever entered. 2. If toolset.close() itself raised during normal teardown, the subsequent heartbeat_task.cancel() in the same finally block was skipped. Move heartbeat_task creation after the duplicate-connect check (so the already-running case never creates one) and move its cancellation into an outermost finally that runs regardless of how the nested worker/toolset teardown exits.
|
Thanks for the detailed review — reworked this to match Pushed 3 new commits:
Tests rewritten accordingly (
|
3 participants
What was changed
TemporalMcpToolSetProvider's{name}-list-toolsand{name}-call-toolactivities intemporalio/contrib/google_adk_agents/_mcp.pynow pool and reuse a singleMcpToolsetconnection per activity name, instead of constructing a brand-new one on every single invocation._ConnectionRecord+ a module-level connection pool keyed by activity name.get_connection()reuses an existing live connection (refcounted) or lazily opens a new one via the existingtoolset_factory.mcp_connection_idle_timeoutparameter onTemporalMcpToolSetProvider.get_tools()/run_async()raises, the connection is evicted immediately so the next call gets a fresh one. A "no matching tool" business-logic error does not evict a healthy connection.tests/contrib/google_adk_agents/test_mcp_pool.pycovering connection reuse, list-tools/call-tool sharing a connection, idle eviction, in-flight calls blocking eviction, and error eviction.Why?
Every activity invocation called
self._toolset_factory(args.factory_argument), constructing a new ADKMcpToolset(and therefore a newMCPSessionManager) with no.close()/cleanup ever called on it. For stdio-transport MCP servers this spawns a new child process on every single activity execution and never cleans it up — a genuine subprocess leak under sustained load, not just a missed optimization.This brings
google_adk_agentsto parity with the pooling already implemented in thestrands(_temporal_mcp_client.py) andgoogle_genai(_mcp.py) contribs, both of which pool and idle-evict connections the same way.How tested
tests/contrib/google_adk_agents/test_mcp_pool.pyexercising the pool directly against a fakeMcpToolset(no real MCP server needed): connection reuse across N calls, list-tools/call-tool sharing one connection, idle eviction after timeout, in-flight calls blocking premature eviction, and eviction-then-reconnect on bothget_tools()andrun_async()failures.ruff check/ruff format --checkandmypy --namespace-packages --check-untyped-defsagainst the changed file — clean.tests/contrib/google_adk_agents/test_google_adk_agents.pyto confirm no regressions.Risks
factory_argumentis only consulted the first time a connection opens for a given activity name; a warm connection is reused regardless of later calls'factory_argumentvalues. This matches howstrands/google_genaipool a single connection per name (they have no per-callfactory_argumentconcept at all), but is a behavior change if any caller relied onfactory_argumentselecting a different backend on every call for the same provider name.strands/google_genaiimplementations directly, neither actually wires one in either (despite that being suggested in the tracking issue); idle + error eviction is the actual existing pattern, which this PR now matches.Closes #1663