From ffd5914333be215595712076384e31bee9270708 Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Tue, 21 Jul 2026 13:52:19 +0530 Subject: [PATCH 1/7] fix(contrib): pool and idle-evict MCP connections in google_adk_agents 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 #1663 --- temporalio/contrib/google_adk_agents/_mcp.py | 179 ++++++++++++++++--- 1 file changed, 159 insertions(+), 20 deletions(-) diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index 92bf994dd..8b2ec5269 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import Sequence from dataclasses import dataclass from datetime import timedelta @@ -17,6 +18,106 @@ from temporalio.exceptions import ApplicationError from temporalio.workflow import ActivityConfig +# Default time an idle pooled McpToolset connection stays open before being +# closed. The timer resets on every call that reuses the connection. +# Overridable via ``TemporalMcpToolSetProvider(mcp_connection_idle_timeout=...)``, +# matching the ``_MCP_CONNECTION_IDLE`` default used by the strands and +# google_genai contribs' equivalent MCP connection pools. +_MCP_CONNECTION_IDLE = timedelta(minutes=5) + +# Provider name -> live pooled connection held open in the activity worker +# process. Activities run in the worker process, so this module state is +# shared across activity invocations on the worker. +_CONNECTIONS: dict[str, "_ConnectionRecord"] = {} + + +class _ConnectionRecord: + """A single ``McpToolset`` instance held open and reused across calls. + + ``McpToolset`` lazily opens its MCP session on first use and manages + reconnection internally via its own ``MCPSessionManager``, so -- unlike + the raw ``mcp.ClientSession`` pooled by the strands/google_genai contribs + -- no dedicated owner task is needed here: the toolset instance itself is + the thing kept alive and reused across activity invocations. + """ + + def __init__(self, name: str, toolset: McpToolset, idle_timeout: timedelta) -> None: + self._name = name + self._toolset = toolset + self._idle_timeout = idle_timeout + self._idle_handle: asyncio.TimerHandle | None = None + self._inflight = 0 + + @property + def toolset(self) -> McpToolset: + return self._toolset + + def acquire(self) -> None: + """Mark a call in flight; pause idle eviction while calls are active.""" + self._inflight += 1 + if self._idle_handle is not None: + self._idle_handle.cancel() + self._idle_handle = None + + def release(self) -> None: + """Mark a call done; arm idle eviction once no calls remain in flight.""" + self._inflight -= 1 + # Only the record still cached under this name arms a timer; a record + # already evicted or never cached must not schedule one, or it could + # later evict a different, healthy connection for the same name. + if self._inflight == 0 and _CONNECTIONS.get(self._name) is self: + loop = asyncio.get_running_loop() + self._idle_handle = loop.call_later( + self._idle_timeout.total_seconds(), self._on_idle + ) + + def _on_idle(self) -> None: + asyncio.ensure_future(self._maybe_evict()) + + async def _maybe_evict(self) -> None: + # A call may have acquired the connection between the timer firing and + # this task running; only evict if it is still idle. + if self._inflight == 0: + await _evict_connection(self._name) + + async def aclose(self) -> None: + """Cancel any pending idle timer and close the underlying toolset.""" + if self._idle_handle is not None: + self._idle_handle.cancel() + self._idle_handle = None + await self._toolset.close() + + +async def get_connection( + name: str, + toolset_factory: Callable[[Any | None], McpToolset], + factory_argument: Any | None, + idle_timeout: timedelta, +) -> "_ConnectionRecord": + """Return the cached connection for ``name``, opening one lazily if needed. + + The returned record is acquired; the caller must ``release()`` it once the + call completes so idle eviction can resume. ``factory_argument`` is only + consulted the first time a connection is opened for ``name`` -- a warm + connection is reused as-is regardless of subsequent calls' ``factory_argument`` + values, matching how the strands/google_genai contribs pool a single + connection per name. + """ + record = _CONNECTIONS.get(name) + if record is None: + record = _ConnectionRecord( + name, toolset_factory(factory_argument), idle_timeout + ) + _CONNECTIONS[name] = record + record.acquire() + return record + + +async def _evict_connection(name: str) -> None: + record = _CONNECTIONS.pop(name, None) + if record is not None: + await record.aclose() + @dataclass class _GetToolsArguments: @@ -91,25 +192,46 @@ class TemporalMcpToolSetProvider: """ def __init__( - self, name: str, toolset_factory: Callable[[Any | None], McpToolset] + self, + name: str, + toolset_factory: Callable[[Any | None], McpToolset], + mcp_connection_idle_timeout: timedelta | None = None, ) -> None: """Initializes the toolset provider. Args: name: Name prefix for the generated activities. toolset_factory: Factory function that creates McpToolset instances. + mcp_connection_idle_timeout: How long a pooled MCP connection may sit + idle (no in-flight calls) before it is closed and evicted. + Defaults to 5 minutes, matching the strands/google_genai contribs. """ super().__init__() self._name = name self._toolset_factory = toolset_factory + self._idle_timeout = mcp_connection_idle_timeout or _MCP_CONNECTION_IDLE def _get_activities(self) -> Sequence[Callable]: @activity.defn(name=self._name + "-list-tools") async def get_tools( args: _GetToolsArguments, ) -> list[_ToolResult]: - toolset = self._toolset_factory(args.factory_argument) - tools = await toolset.get_tools() + record = await get_connection( + self._name, + self._toolset_factory, + args.factory_argument, + self._idle_timeout, + ) + try: + try: + tools = await record.toolset.get_tools() + except Exception: + # The underlying session may be broken; drop it so the + # next call reconnects instead of reusing a dead session. + await _evict_connection(self._name) + raise + finally: + record.release() return [ _ToolResult( tool.name, @@ -125,24 +247,41 @@ async def get_tools( async def call_tool( args: _CallToolArguments, ) -> _CallToolResult: - toolset = self._toolset_factory(args.factory_argument) - tools = await toolset.get_tools() - tool_match = [tool for tool in tools if tool.name == args.name] - if len(tool_match) == 0: - raise ApplicationError( - f"Unable to find matching mcp tool by name: {args.name}" - ) - if len(tool_match) > 1: - raise ApplicationError( - f"Unable too many matching mcp tools by name: {args.name}" - ) - tool = tool_match[0] - - # We cannot provide a full-fledged ToolContext so we need to provide only what is needed by the tool - result = await tool.run_async( - args=args.arguments, - tool_context=args.tool_context, # type:ignore + record = await get_connection( + self._name, + self._toolset_factory, + args.factory_argument, + self._idle_timeout, ) + try: + try: + tools = await record.toolset.get_tools() + except Exception: + await _evict_connection(self._name) + raise + + tool_match = [tool for tool in tools if tool.name == args.name] + if len(tool_match) == 0: + raise ApplicationError( + f"Unable to find matching mcp tool by name: {args.name}" + ) + if len(tool_match) > 1: + raise ApplicationError( + f"Unable too many matching mcp tools by name: {args.name}" + ) + tool = tool_match[0] + + try: + # We cannot provide a full-fledged ToolContext so we need to provide only what is needed by the tool + result = await tool.run_async( + args=args.arguments, + tool_context=args.tool_context, # type:ignore + ) + except Exception: + await _evict_connection(self._name) + raise + finally: + record.release() return _CallToolResult(result=result, tool_context=args.tool_context) return get_tools, call_tool From 11e7550e1009bc58342db4f3af6e155f71209ece Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Tue, 21 Jul 2026 13:52:27 +0530 Subject: [PATCH 2/7] test(contrib): cover MCP connection reuse, idle eviction, and error eviction 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. --- .../google_adk_agents/test_mcp_pool.py | 276 ++++++++++++++++++ 1 file changed, 276 insertions(+) create mode 100644 tests/contrib/google_adk_agents/test_mcp_pool.py diff --git a/tests/contrib/google_adk_agents/test_mcp_pool.py b/tests/contrib/google_adk_agents/test_mcp_pool.py new file mode 100644 index 000000000..976c1e48f --- /dev/null +++ b/tests/contrib/google_adk_agents/test_mcp_pool.py @@ -0,0 +1,276 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the pooled/idle-evicted MCP connection support in +``temporalio.contrib.google_adk_agents._mcp``. + +These exercise ``TemporalMcpToolSetProvider``'s generated activities directly +against a fake ``McpToolset``, so they don't require a real MCP server, a +Temporal test server, or a full ``Worker`` -- they're targeted regression +tests for the connection-leak fix (one underlying toolset per name, reused +across calls, evicted on idle or on error). +""" + +import asyncio +from datetime import timedelta +from typing import Any + +import pytest +from google.adk.events import EventActions + +from temporalio.contrib.google_adk_agents import _mcp +from temporalio.contrib.google_adk_agents._mcp import ( + TemporalMcpToolSetProvider, + TemporalToolContext, + _CallToolArguments, + _GetToolsArguments, +) + + +class _FakeTool: + def __init__(self, name: str, *, fail_run: bool = False) -> None: + self.name = name + self.description = "a fake tool" + self.is_long_running = False + self.custom_metadata: dict[str, Any] | None = None + self._fail_run = fail_run + + def _get_declaration(self) -> None: + return None + + async def run_async(self, *, args: dict[str, Any], tool_context: Any) -> Any: + if self._fail_run: + raise RuntimeError("tool call failed") + return {"echo": args} + + +class _FakeToolset: + """Stands in for ``google.adk.tools.mcp_tool.McpToolset``. + + Tracks whether ``close()`` was called so tests can assert eviction + actually tears the pooled connection down, not just drops it from the + cache. + """ + + def __init__(self, *, fail_get_tools: bool = False, fail_run: bool = False) -> None: + self.closed = False + self._fail_get_tools = fail_get_tools + self._fail_run = fail_run + + async def get_tools(self) -> list[_FakeTool]: + if self._fail_get_tools: + raise RuntimeError("get_tools failed") + return [_FakeTool("echo", fail_run=self._fail_run)] + + async def close(self) -> None: + self.closed = True + + +def _tool_context() -> TemporalToolContext: + return TemporalToolContext( + tool_confirmation=None, + function_call_id=None, + event_actions=EventActions(), + ) + + +def _call_tool_args(factory_argument: Any = None) -> _CallToolArguments: + return _CallToolArguments( + factory_argument=factory_argument, + name="echo", + arguments={"x": 1}, + tool_context=_tool_context(), + ) + + +@pytest.fixture(autouse=True) +def _clear_connection_pool(): + _mcp._CONNECTIONS.clear() + yield + _mcp._CONNECTIONS.clear() + + +async def test_call_tool_reuses_one_pooled_connection(): + """N sequential call_tool executions reuse one toolset, not N.""" + created: list[_FakeToolset] = [] + + def factory(_: Any) -> _FakeToolset: + toolset = _FakeToolset() + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("pool_reuse", factory) + _, call_tool = provider._get_activities() + + for _ in range(5): + result = await call_tool(_call_tool_args()) + assert result.result == {"echo": {"x": 1}} + + assert len(created) == 1 + assert not created[0].closed + + +async def test_get_tools_and_call_tool_share_one_connection(): + """list-tools and call-tool activities for the same name share a connection.""" + created: list[_FakeToolset] = [] + + def factory(_: Any) -> _FakeToolset: + toolset = _FakeToolset() + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("pool_shared", factory) + get_tools, call_tool = provider._get_activities() + + tools = await get_tools(_GetToolsArguments(factory_argument=None)) + assert [t.name for t in tools] == ["echo"] + + await call_tool(_call_tool_args()) + await call_tool(_call_tool_args()) + + assert len(created) == 1 + + +async def test_idle_eviction_closes_connection_after_timeout(): + def factory(_: Any) -> _FakeToolset: + return _FakeToolset() + + provider = TemporalMcpToolSetProvider( + "pool_idle", + factory, + mcp_connection_idle_timeout=timedelta(milliseconds=20), + ) + _, call_tool = provider._get_activities() + + await call_tool(_call_tool_args()) + + record = _mcp._CONNECTIONS["pool_idle"] + assert not record.toolset.closed + + for _ in range(100): + if "pool_idle" not in _mcp._CONNECTIONS: + break + await asyncio.sleep(0.01) + + assert "pool_idle" not in _mcp._CONNECTIONS + assert record.toolset.closed + + +async def test_idle_timer_only_arms_once_all_inflight_calls_release(): + """Two overlapping callers must not let the idle timer fire mid-call. + + Exercises ``_ConnectionRecord`` directly (rather than through the + activities) so the ordering of acquire/release calls -- which stands in + for two concurrent in-flight activity invocations sharing one pooled + connection -- is deterministic instead of depending on how the event loop + happens to interleave coroutines that never truly suspend. + """ + toolset = _FakeToolset() + record = _mcp._ConnectionRecord("pool_inflight", toolset, timedelta(milliseconds=1)) + _mcp._CONNECTIONS["pool_inflight"] = record + + record.acquire() + record.acquire() + assert record._idle_handle is None + + record.release() + # One caller is still in flight; the idle timer must stay disarmed. + assert record._idle_handle is None + assert "pool_inflight" in _mcp._CONNECTIONS + + record.release() + # Now that both callers are done, the idle timer arms. + assert record._idle_handle is not None + + for _ in range(100): + if "pool_inflight" not in _mcp._CONNECTIONS: + break + await asyncio.sleep(0.01) + assert "pool_inflight" not in _mcp._CONNECTIONS + assert toolset.closed + + +async def test_call_tool_error_evicts_and_next_call_reconnects(): + created: list[_FakeToolset] = [] + + def factory(_: Any) -> _FakeToolset: + # Only the first toolset's get_tools() fails. + toolset = _FakeToolset(fail_get_tools=(len(created) == 0)) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("pool_error", factory) + _, call_tool = provider._get_activities() + + with pytest.raises(RuntimeError, match="get_tools failed"): + await call_tool(_call_tool_args()) + + assert "pool_error" not in _mcp._CONNECTIONS + assert created[0].closed + + result = await call_tool(_call_tool_args()) + assert result.result == {"echo": {"x": 1}} + assert len(created) == 2 + assert not created[1].closed + + +async def test_call_tool_run_async_error_evicts_broken_connection(): + created: list[_FakeToolset] = [] + + def factory(_: Any) -> _FakeToolset: + toolset = _FakeToolset(fail_run=(len(created) == 0)) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("pool_run_error", factory) + _, call_tool = provider._get_activities() + + with pytest.raises(RuntimeError, match="tool call failed"): + await call_tool(_call_tool_args()) + + assert "pool_run_error" not in _mcp._CONNECTIONS + assert created[0].closed + + result = await call_tool(_call_tool_args()) + assert result.result == {"echo": {"x": 1}} + assert len(created) == 2 + + +async def test_call_tool_no_matching_tool_does_not_evict(): + """A business-logic ApplicationError isn't a broken session; keep pooling.""" + created: list[_FakeToolset] = [] + + def factory(_: Any) -> _FakeToolset: + toolset = _FakeToolset() + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("pool_no_match", factory) + _, call_tool = provider._get_activities() + + args = _CallToolArguments( + factory_argument=None, + name="does_not_exist", + arguments={}, + tool_context=_tool_context(), + ) + with pytest.raises(Exception, match="Unable to find matching mcp tool"): + await call_tool(args) + + assert "pool_no_match" in _mcp._CONNECTIONS + assert not created[0].closed + + # A subsequent successful call reuses the same toolset instance. + await call_tool(_call_tool_args()) + assert len(created) == 1 From 0996b6c19b32926f07b0f299126e6a0718f905be Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Wed, 22 Jul 2026 16:26:14 +0530 Subject: [PATCH 3/7] fix(contrib): make google_adk_agents MCP toolsets stateless (fixes factory_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 #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 #1663 --- temporalio/contrib/google_adk_agents/_mcp.py | 187 +++--------- tests/contrib/google_adk_agents/test_mcp.py | 222 ++++++++++++++ .../google_adk_agents/test_mcp_pool.py | 276 ------------------ 3 files changed, 257 insertions(+), 428 deletions(-) create mode 100644 tests/contrib/google_adk_agents/test_mcp.py delete mode 100644 tests/contrib/google_adk_agents/test_mcp_pool.py diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index 8b2ec5269..0e288be34 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -1,4 +1,3 @@ -import asyncio from collections.abc import Sequence from dataclasses import dataclass from datetime import timedelta @@ -18,106 +17,6 @@ from temporalio.exceptions import ApplicationError from temporalio.workflow import ActivityConfig -# Default time an idle pooled McpToolset connection stays open before being -# closed. The timer resets on every call that reuses the connection. -# Overridable via ``TemporalMcpToolSetProvider(mcp_connection_idle_timeout=...)``, -# matching the ``_MCP_CONNECTION_IDLE`` default used by the strands and -# google_genai contribs' equivalent MCP connection pools. -_MCP_CONNECTION_IDLE = timedelta(minutes=5) - -# Provider name -> live pooled connection held open in the activity worker -# process. Activities run in the worker process, so this module state is -# shared across activity invocations on the worker. -_CONNECTIONS: dict[str, "_ConnectionRecord"] = {} - - -class _ConnectionRecord: - """A single ``McpToolset`` instance held open and reused across calls. - - ``McpToolset`` lazily opens its MCP session on first use and manages - reconnection internally via its own ``MCPSessionManager``, so -- unlike - the raw ``mcp.ClientSession`` pooled by the strands/google_genai contribs - -- no dedicated owner task is needed here: the toolset instance itself is - the thing kept alive and reused across activity invocations. - """ - - def __init__(self, name: str, toolset: McpToolset, idle_timeout: timedelta) -> None: - self._name = name - self._toolset = toolset - self._idle_timeout = idle_timeout - self._idle_handle: asyncio.TimerHandle | None = None - self._inflight = 0 - - @property - def toolset(self) -> McpToolset: - return self._toolset - - def acquire(self) -> None: - """Mark a call in flight; pause idle eviction while calls are active.""" - self._inflight += 1 - if self._idle_handle is not None: - self._idle_handle.cancel() - self._idle_handle = None - - def release(self) -> None: - """Mark a call done; arm idle eviction once no calls remain in flight.""" - self._inflight -= 1 - # Only the record still cached under this name arms a timer; a record - # already evicted or never cached must not schedule one, or it could - # later evict a different, healthy connection for the same name. - if self._inflight == 0 and _CONNECTIONS.get(self._name) is self: - loop = asyncio.get_running_loop() - self._idle_handle = loop.call_later( - self._idle_timeout.total_seconds(), self._on_idle - ) - - def _on_idle(self) -> None: - asyncio.ensure_future(self._maybe_evict()) - - async def _maybe_evict(self) -> None: - # A call may have acquired the connection between the timer firing and - # this task running; only evict if it is still idle. - if self._inflight == 0: - await _evict_connection(self._name) - - async def aclose(self) -> None: - """Cancel any pending idle timer and close the underlying toolset.""" - if self._idle_handle is not None: - self._idle_handle.cancel() - self._idle_handle = None - await self._toolset.close() - - -async def get_connection( - name: str, - toolset_factory: Callable[[Any | None], McpToolset], - factory_argument: Any | None, - idle_timeout: timedelta, -) -> "_ConnectionRecord": - """Return the cached connection for ``name``, opening one lazily if needed. - - The returned record is acquired; the caller must ``release()`` it once the - call completes so idle eviction can resume. ``factory_argument`` is only - consulted the first time a connection is opened for ``name`` -- a warm - connection is reused as-is regardless of subsequent calls' ``factory_argument`` - values, matching how the strands/google_genai contribs pool a single - connection per name. - """ - record = _CONNECTIONS.get(name) - if record is None: - record = _ConnectionRecord( - name, toolset_factory(factory_argument), idle_timeout - ) - _CONNECTIONS[name] = record - record.acquire() - return record - - -async def _evict_connection(name: str) -> None: - record = _CONNECTIONS.pop(name, None) - if record is not None: - await record.aclose() - @dataclass class _GetToolsArguments: @@ -189,76 +88,64 @@ class TemporalMcpToolSetProvider: Manages the creation of toolset activities and handles tool execution within Temporal workflows. + + This provider is *stateless*: every ``list-tools``/``call-tool`` activity + invocation builds a fresh ``McpToolset`` via ``toolset_factory``, runs the + operation, and always closes the toolset in a ``finally`` block. This means + ``factory_argument`` is honored on every single call (no cross-workflow + connection sharing or silent mis-routing) and no MCP session or stdio + subprocess is leaked. State is not maintained across calls; if a persistent + connection is required, use :class:`TemporalStatefulMcpToolSetProvider`. """ def __init__( self, name: str, toolset_factory: Callable[[Any | None], McpToolset], - mcp_connection_idle_timeout: timedelta | None = None, ) -> None: """Initializes the toolset provider. Args: name: Name prefix for the generated activities. toolset_factory: Factory function that creates McpToolset instances. - mcp_connection_idle_timeout: How long a pooled MCP connection may sit - idle (no in-flight calls) before it is closed and evicted. - Defaults to 5 minutes, matching the strands/google_genai contribs. + It should return a new toolset each time so that no state is + shared between workflow runs. """ super().__init__() self._name = name self._toolset_factory = toolset_factory - self._idle_timeout = mcp_connection_idle_timeout or _MCP_CONNECTION_IDLE def _get_activities(self) -> Sequence[Callable]: @activity.defn(name=self._name + "-list-tools") async def get_tools( args: _GetToolsArguments, ) -> list[_ToolResult]: - record = await get_connection( - self._name, - self._toolset_factory, - args.factory_argument, - self._idle_timeout, - ) + # Build a fresh toolset per call, honoring this call's + # ``factory_argument``, and always close it so no MCP session + # (or stdio subprocess) leaks. See issue #1663. + toolset = self._toolset_factory(args.factory_argument) try: - try: - tools = await record.toolset.get_tools() - except Exception: - # The underlying session may be broken; drop it so the - # next call reconnects instead of reusing a dead session. - await _evict_connection(self._name) - raise + tools = await toolset.get_tools() + return [ + _ToolResult( + tool.name, + tool.description, + tool.is_long_running, + tool.custom_metadata, + tool._get_declaration(), + ) + for tool in tools + ] finally: - record.release() - return [ - _ToolResult( - tool.name, - tool.description, - tool.is_long_running, - tool.custom_metadata, - tool._get_declaration(), - ) - for tool in tools - ] + await toolset.close() @activity.defn(name=self._name + "-call-tool") async def call_tool( args: _CallToolArguments, ) -> _CallToolResult: - record = await get_connection( - self._name, - self._toolset_factory, - args.factory_argument, - self._idle_timeout, - ) + toolset = self._toolset_factory(args.factory_argument) try: - try: - tools = await record.toolset.get_tools() - except Exception: - await _evict_connection(self._name) - raise + tools = await toolset.get_tools() tool_match = [tool for tool in tools if tool.name == args.name] if len(tool_match) == 0: @@ -271,18 +158,14 @@ async def call_tool( ) tool = tool_match[0] - try: - # We cannot provide a full-fledged ToolContext so we need to provide only what is needed by the tool - result = await tool.run_async( - args=args.arguments, - tool_context=args.tool_context, # type:ignore - ) - except Exception: - await _evict_connection(self._name) - raise + # We cannot provide a full-fledged ToolContext so we need to provide only what is needed by the tool + result = await tool.run_async( + args=args.arguments, + tool_context=args.tool_context, # type:ignore + ) + return _CallToolResult(result=result, tool_context=args.tool_context) finally: - record.release() - return _CallToolResult(result=result, tool_context=args.tool_context) + await toolset.close() return get_tools, call_tool diff --git a/tests/contrib/google_adk_agents/test_mcp.py b/tests/contrib/google_adk_agents/test_mcp.py new file mode 100644 index 000000000..4261c0e58 --- /dev/null +++ b/tests/contrib/google_adk_agents/test_mcp.py @@ -0,0 +1,222 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the stateless MCP toolset support in +``temporalio.contrib.google_adk_agents._mcp``. + +These exercise ``TemporalMcpToolSetProvider``'s generated activities directly +against a fake ``McpToolset``, so they don't require a real MCP server, a +Temporal test server, or a full ``Worker``. They are targeted regression tests +for the connection-leak fix (issue #1663): every call builds its own toolset, +honors that call's ``factory_argument``, and always closes the toolset. +""" + +from typing import Any + +import pytest +from google.adk.events import EventActions + +from temporalio.contrib.google_adk_agents._mcp import ( + TemporalMcpToolSetProvider, + TemporalToolContext, + _CallToolArguments, + _GetToolsArguments, +) + + +class _FakeTool: + def __init__(self, name: str, *, fail_run: bool = False) -> None: + self.name = name + self.description = "a fake tool" + self.is_long_running = False + self.custom_metadata: dict[str, Any] | None = None + self._fail_run = fail_run + + def _get_declaration(self) -> None: + return None + + async def run_async(self, *, args: dict[str, Any], tool_context: Any) -> Any: + if self._fail_run: + raise RuntimeError("tool call failed") + return {"echo": args} + + +class _FakeToolset: + """Stands in for ``google.adk.tools.mcp_tool.McpToolset``. + + Tracks whether ``close()`` was called so tests can assert the stateless + path always tears the toolset down, and records the ``factory_argument`` + it was created with so tests can prove each call routes with its own + argument. + """ + + def __init__( + self, + factory_argument: Any = None, + *, + fail_get_tools: bool = False, + fail_run: bool = False, + ) -> None: + self.factory_argument = factory_argument + self.closed = False + self._fail_get_tools = fail_get_tools + self._fail_run = fail_run + + async def get_tools(self) -> list[_FakeTool]: + if self._fail_get_tools: + raise RuntimeError("get_tools failed") + return [_FakeTool("echo", fail_run=self._fail_run)] + + async def close(self) -> None: + self.closed = True + + +def _tool_context() -> TemporalToolContext: + return TemporalToolContext( + tool_confirmation=None, + function_call_id=None, + event_actions=EventActions(), + ) + + +def _call_tool_args(factory_argument: Any = None) -> _CallToolArguments: + return _CallToolArguments( + factory_argument=factory_argument, + name="echo", + arguments={"x": 1}, + tool_context=_tool_context(), + ) + + +async def test_call_tool_creates_and_closes_fresh_toolset_each_call(): + """Each call_tool builds its own toolset and closes it, every time.""" + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_reuse", factory) + _, call_tool = provider._get_activities() + + for _ in range(5): + result = await call_tool(_call_tool_args()) + assert result.result == {"echo": {"x": 1}} + + # One fresh toolset per call, and every one was closed. + assert len(created) == 5 + assert all(t.closed for t in created) + + +async def test_get_tools_creates_and_closes_fresh_toolset(): + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_list", factory) + get_tools, _ = provider._get_activities() + + tools = await get_tools(_GetToolsArguments(factory_argument=None)) + assert [t.name for t in tools] == ["echo"] + + assert len(created) == 1 + assert created[0].closed + + +async def test_factory_argument_honored_on_every_call(): + """The bug the reviewer flagged: a later call with a different + ``factory_argument`` must route with *its own* argument, never silently + reuse a connection opened for an earlier argument. + """ + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_routing", factory) + _, call_tool = provider._get_activities() + + await call_tool(_call_tool_args(factory_argument="tenant-a")) + await call_tool(_call_tool_args(factory_argument="tenant-b")) + + assert [t.factory_argument for t in created] == ["tenant-a", "tenant-b"] + assert all(t.closed for t in created) + + +async def test_call_tool_closes_toolset_on_error(): + """A failure mid-call still closes the toolset (no leak on the error path).""" + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg, fail_run=True) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_run_error", factory) + _, call_tool = provider._get_activities() + + with pytest.raises(RuntimeError, match="tool call failed"): + await call_tool(_call_tool_args()) + + assert len(created) == 1 + assert created[0].closed + + +async def test_get_tools_closes_toolset_on_error(): + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg, fail_get_tools=True) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_list_error", factory) + get_tools, _ = provider._get_activities() + + with pytest.raises(RuntimeError, match="get_tools failed"): + await get_tools(_GetToolsArguments(factory_argument=None)) + + assert len(created) == 1 + assert created[0].closed + + +async def test_call_tool_no_matching_tool_still_closes(): + """A business-logic ApplicationError still closes the fresh toolset.""" + created: list[_FakeToolset] = [] + + def factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg) + created.append(toolset) + return toolset + + provider = TemporalMcpToolSetProvider("stateless_no_match", factory) + _, call_tool = provider._get_activities() + + args = _CallToolArguments( + factory_argument=None, + name="does_not_exist", + arguments={}, + tool_context=_tool_context(), + ) + with pytest.raises(Exception, match="Unable to find matching mcp tool"): + await call_tool(args) + + assert len(created) == 1 + assert created[0].closed diff --git a/tests/contrib/google_adk_agents/test_mcp_pool.py b/tests/contrib/google_adk_agents/test_mcp_pool.py deleted file mode 100644 index 976c1e48f..000000000 --- a/tests/contrib/google_adk_agents/test_mcp_pool.py +++ /dev/null @@ -1,276 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for the pooled/idle-evicted MCP connection support in -``temporalio.contrib.google_adk_agents._mcp``. - -These exercise ``TemporalMcpToolSetProvider``'s generated activities directly -against a fake ``McpToolset``, so they don't require a real MCP server, a -Temporal test server, or a full ``Worker`` -- they're targeted regression -tests for the connection-leak fix (one underlying toolset per name, reused -across calls, evicted on idle or on error). -""" - -import asyncio -from datetime import timedelta -from typing import Any - -import pytest -from google.adk.events import EventActions - -from temporalio.contrib.google_adk_agents import _mcp -from temporalio.contrib.google_adk_agents._mcp import ( - TemporalMcpToolSetProvider, - TemporalToolContext, - _CallToolArguments, - _GetToolsArguments, -) - - -class _FakeTool: - def __init__(self, name: str, *, fail_run: bool = False) -> None: - self.name = name - self.description = "a fake tool" - self.is_long_running = False - self.custom_metadata: dict[str, Any] | None = None - self._fail_run = fail_run - - def _get_declaration(self) -> None: - return None - - async def run_async(self, *, args: dict[str, Any], tool_context: Any) -> Any: - if self._fail_run: - raise RuntimeError("tool call failed") - return {"echo": args} - - -class _FakeToolset: - """Stands in for ``google.adk.tools.mcp_tool.McpToolset``. - - Tracks whether ``close()`` was called so tests can assert eviction - actually tears the pooled connection down, not just drops it from the - cache. - """ - - def __init__(self, *, fail_get_tools: bool = False, fail_run: bool = False) -> None: - self.closed = False - self._fail_get_tools = fail_get_tools - self._fail_run = fail_run - - async def get_tools(self) -> list[_FakeTool]: - if self._fail_get_tools: - raise RuntimeError("get_tools failed") - return [_FakeTool("echo", fail_run=self._fail_run)] - - async def close(self) -> None: - self.closed = True - - -def _tool_context() -> TemporalToolContext: - return TemporalToolContext( - tool_confirmation=None, - function_call_id=None, - event_actions=EventActions(), - ) - - -def _call_tool_args(factory_argument: Any = None) -> _CallToolArguments: - return _CallToolArguments( - factory_argument=factory_argument, - name="echo", - arguments={"x": 1}, - tool_context=_tool_context(), - ) - - -@pytest.fixture(autouse=True) -def _clear_connection_pool(): - _mcp._CONNECTIONS.clear() - yield - _mcp._CONNECTIONS.clear() - - -async def test_call_tool_reuses_one_pooled_connection(): - """N sequential call_tool executions reuse one toolset, not N.""" - created: list[_FakeToolset] = [] - - def factory(_: Any) -> _FakeToolset: - toolset = _FakeToolset() - created.append(toolset) - return toolset - - provider = TemporalMcpToolSetProvider("pool_reuse", factory) - _, call_tool = provider._get_activities() - - for _ in range(5): - result = await call_tool(_call_tool_args()) - assert result.result == {"echo": {"x": 1}} - - assert len(created) == 1 - assert not created[0].closed - - -async def test_get_tools_and_call_tool_share_one_connection(): - """list-tools and call-tool activities for the same name share a connection.""" - created: list[_FakeToolset] = [] - - def factory(_: Any) -> _FakeToolset: - toolset = _FakeToolset() - created.append(toolset) - return toolset - - provider = TemporalMcpToolSetProvider("pool_shared", factory) - get_tools, call_tool = provider._get_activities() - - tools = await get_tools(_GetToolsArguments(factory_argument=None)) - assert [t.name for t in tools] == ["echo"] - - await call_tool(_call_tool_args()) - await call_tool(_call_tool_args()) - - assert len(created) == 1 - - -async def test_idle_eviction_closes_connection_after_timeout(): - def factory(_: Any) -> _FakeToolset: - return _FakeToolset() - - provider = TemporalMcpToolSetProvider( - "pool_idle", - factory, - mcp_connection_idle_timeout=timedelta(milliseconds=20), - ) - _, call_tool = provider._get_activities() - - await call_tool(_call_tool_args()) - - record = _mcp._CONNECTIONS["pool_idle"] - assert not record.toolset.closed - - for _ in range(100): - if "pool_idle" not in _mcp._CONNECTIONS: - break - await asyncio.sleep(0.01) - - assert "pool_idle" not in _mcp._CONNECTIONS - assert record.toolset.closed - - -async def test_idle_timer_only_arms_once_all_inflight_calls_release(): - """Two overlapping callers must not let the idle timer fire mid-call. - - Exercises ``_ConnectionRecord`` directly (rather than through the - activities) so the ordering of acquire/release calls -- which stands in - for two concurrent in-flight activity invocations sharing one pooled - connection -- is deterministic instead of depending on how the event loop - happens to interleave coroutines that never truly suspend. - """ - toolset = _FakeToolset() - record = _mcp._ConnectionRecord("pool_inflight", toolset, timedelta(milliseconds=1)) - _mcp._CONNECTIONS["pool_inflight"] = record - - record.acquire() - record.acquire() - assert record._idle_handle is None - - record.release() - # One caller is still in flight; the idle timer must stay disarmed. - assert record._idle_handle is None - assert "pool_inflight" in _mcp._CONNECTIONS - - record.release() - # Now that both callers are done, the idle timer arms. - assert record._idle_handle is not None - - for _ in range(100): - if "pool_inflight" not in _mcp._CONNECTIONS: - break - await asyncio.sleep(0.01) - assert "pool_inflight" not in _mcp._CONNECTIONS - assert toolset.closed - - -async def test_call_tool_error_evicts_and_next_call_reconnects(): - created: list[_FakeToolset] = [] - - def factory(_: Any) -> _FakeToolset: - # Only the first toolset's get_tools() fails. - toolset = _FakeToolset(fail_get_tools=(len(created) == 0)) - created.append(toolset) - return toolset - - provider = TemporalMcpToolSetProvider("pool_error", factory) - _, call_tool = provider._get_activities() - - with pytest.raises(RuntimeError, match="get_tools failed"): - await call_tool(_call_tool_args()) - - assert "pool_error" not in _mcp._CONNECTIONS - assert created[0].closed - - result = await call_tool(_call_tool_args()) - assert result.result == {"echo": {"x": 1}} - assert len(created) == 2 - assert not created[1].closed - - -async def test_call_tool_run_async_error_evicts_broken_connection(): - created: list[_FakeToolset] = [] - - def factory(_: Any) -> _FakeToolset: - toolset = _FakeToolset(fail_run=(len(created) == 0)) - created.append(toolset) - return toolset - - provider = TemporalMcpToolSetProvider("pool_run_error", factory) - _, call_tool = provider._get_activities() - - with pytest.raises(RuntimeError, match="tool call failed"): - await call_tool(_call_tool_args()) - - assert "pool_run_error" not in _mcp._CONNECTIONS - assert created[0].closed - - result = await call_tool(_call_tool_args()) - assert result.result == {"echo": {"x": 1}} - assert len(created) == 2 - - -async def test_call_tool_no_matching_tool_does_not_evict(): - """A business-logic ApplicationError isn't a broken session; keep pooling.""" - created: list[_FakeToolset] = [] - - def factory(_: Any) -> _FakeToolset: - toolset = _FakeToolset() - created.append(toolset) - return toolset - - provider = TemporalMcpToolSetProvider("pool_no_match", factory) - _, call_tool = provider._get_activities() - - args = _CallToolArguments( - factory_argument=None, - name="does_not_exist", - arguments={}, - tool_context=_tool_context(), - ) - with pytest.raises(Exception, match="Unable to find matching mcp tool"): - await call_tool(args) - - assert "pool_no_match" in _mcp._CONNECTIONS - assert not created[0].closed - - # A subsequent successful call reuses the same toolset instance. - await call_tool(_call_tool_args()) - assert len(created) == 1 From 4748a58b35e6f369cded93c4ae29c37734b65622 Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Wed, 22 Jul 2026 16:33:34 +0530 Subject: [PATCH 4/7] feat(contrib): add stateful pooled MCP toolset provider for google_adk_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. --- .../contrib/google_adk_agents/__init__.py | 4 + temporalio/contrib/google_adk_agents/_mcp.py | 399 +++++++++++++++++- .../contrib/google_adk_agents/_plugin.py | 15 +- .../google_adk_agents/test_stateful_mcp.py | 187 ++++++++ 4 files changed, 600 insertions(+), 5 deletions(-) create mode 100644 tests/contrib/google_adk_agents/test_stateful_mcp.py diff --git a/temporalio/contrib/google_adk_agents/__init__.py b/temporalio/contrib/google_adk_agents/__init__.py index 3f236516b..d4c969fb3 100644 --- a/temporalio/contrib/google_adk_agents/__init__.py +++ b/temporalio/contrib/google_adk_agents/__init__.py @@ -6,6 +6,8 @@ from temporalio.contrib.google_adk_agents._mcp import ( TemporalMcpToolSet, TemporalMcpToolSetProvider, + TemporalStatefulMcpToolSet, + TemporalStatefulMcpToolSetProvider, ) from temporalio.contrib.google_adk_agents._model import TemporalModel from temporalio.contrib.google_adk_agents._plugin import ( @@ -16,5 +18,7 @@ "GoogleAdkPlugin", "TemporalMcpToolSet", "TemporalMcpToolSetProvider", + "TemporalStatefulMcpToolSet", + "TemporalStatefulMcpToolSetProvider", "TemporalModel", ] diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index 0e288be34..0694f6a11 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -1,6 +1,9 @@ +import asyncio +import functools from collections.abc import Sequence from dataclasses import dataclass from datetime import timedelta +from types import TracebackType from typing import Any, Callable from google.adk.agents.readonly_context import ReadonlyContext @@ -14,8 +17,17 @@ from google.genai.types import FunctionDeclaration from temporalio import activity, workflow -from temporalio.exceptions import ApplicationError -from temporalio.workflow import ActivityConfig +from temporalio.api.enums.v1.workflow_pb2 import ( + TIMEOUT_TYPE_HEARTBEAT, + TIMEOUT_TYPE_SCHEDULE_TO_START, +) +from temporalio.exceptions import ( + ActivityError, + ApplicationError, + is_cancelled_exception, +) +from temporalio.worker import PollerBehaviorSimpleMaximum, Worker +from temporalio.workflow import ActivityConfig, ActivityHandle @dataclass @@ -303,3 +315,386 @@ async def get_tools( ) for tool_result in tool_results ] + + +def _handle_worker_failure(func: Callable) -> Callable: + """Surface dedicated-worker failures on the run-scoped task queue. + + A schedule-to-start timeout means the dedicated ``-server-session`` worker + never picked the activity up (it is gone); a heartbeat timeout means it + died mid-session. Either way Temporal cannot recreate the in-memory toolset + state, so we re-raise as an ``ApplicationError`` of type + ``"DedicatedWorkerFailure"`` for the caller to handle. + + Duplicated (rather than shared) from ``openai_agents._mcp`` on purpose: + these two contribs do not currently share internal code, and importing + across them would create an unwanted dependency. + """ + + @functools.wraps(func) + async def wrapper(*args: Any, **kwargs: Any): + try: + return await func(*args, **kwargs) + except ActivityError as e: + failure = e.failure + if failure: + cause = failure.cause + if cause: + if ( + cause.timeout_failure_info.timeout_type + == TIMEOUT_TYPE_SCHEDULE_TO_START + ): + raise ApplicationError( + "MCP Stateful Server Worker failed to schedule activity.", + type="DedicatedWorkerFailure", + ) from e + if ( + cause.timeout_failure_info.timeout_type + == TIMEOUT_TYPE_HEARTBEAT + ): + raise ApplicationError( + "MCP Stateful Server Worker failed to heartbeat.", + type="DedicatedWorkerFailure", + ) from e + raise e + + return wrapper + + +@dataclass +class _StatefulServerSessionArguments: + factory_argument: Any | None + + +@dataclass +class _StatefulCallToolArguments: + name: str + arguments: dict[str, Any] + tool_context: TemporalToolContext + + +class TemporalStatefulMcpToolSetProvider: + """Provider for a stateful, pooled MCP toolset backed by a dedicated worker. + + .. warning:: + This class is experimental and may change in future versions. + Use with caution in production environments. + + Unlike :class:`TemporalMcpToolSetProvider` (which builds and closes a fresh + ``McpToolset`` per activity call), this provider maintains a single + ``McpToolset`` for the lifetime of a workflow run. The workflow-side + :class:`TemporalStatefulMcpToolSet` starts a dedicated ``-server-session`` + activity on a task queue scoped to that specific run + (``name@run_id``); that activity builds the toolset once via + ``toolset_factory(factory_argument)``, holds it open, and runs a nested + :class:`~temporalio.worker.Worker` serving the ``-list-tools``/``-call-tool`` + activities off the same run-scoped queue. The toolset is closed when the + workflow cancels the session activity on cleanup. + + Because state lives in a dedicated worker's memory, the caller must handle + the case where that worker fails, as Temporal cannot seamlessly recreate the + lost toolset. Failure surfaces as an ``ApplicationError`` with + ``type="DedicatedWorkerFailure"``. Prefer the stateless + :class:`TemporalMcpToolSetProvider` unless a persistent connection is + genuinely required. + """ + + def __init__( + self, + name: str, + toolset_factory: Callable[[Any | None], McpToolset], + ) -> None: + """Initializes the stateful toolset provider. + + Args: + name: Name prefix for the generated activities. It is suffixed with + ``-stateful`` so it never collides with a stateless provider of + the same base name registered on the same worker. + toolset_factory: Factory function that creates McpToolset instances. + It should return a new toolset each time so that no state is + shared between workflow runs. + """ + super().__init__() + self._name = name + "-stateful" + self._toolset_factory = toolset_factory + self._toolsets: dict[str, McpToolset] = {} + + @property + def name(self) -> str: + """The activity-name prefix (base name with the ``-stateful`` suffix).""" + return self._name + + def _get_activities(self) -> Sequence[Callable]: + def _server_id() -> str: + return self._name + "@" + (activity.info().workflow_run_id or "") + + @activity.defn(name=self._name + "-list-tools") + async def get_tools() -> list[_ToolResult]: + toolset = self._toolsets[_server_id()] + tools = await toolset.get_tools() + return [ + _ToolResult( + tool.name, + tool.description, + tool.is_long_running, + tool.custom_metadata, + tool._get_declaration(), + ) + for tool in tools + ] + + @activity.defn(name=self._name + "-call-tool") + async def call_tool(args: _StatefulCallToolArguments) -> _CallToolResult: + toolset = self._toolsets[_server_id()] + tools = await toolset.get_tools() + + tool_match = [tool for tool in tools if tool.name == args.name] + if len(tool_match) == 0: + raise ApplicationError( + f"Unable to find matching mcp tool by name: {args.name}" + ) + if len(tool_match) > 1: + raise ApplicationError( + f"Unable too many matching mcp tools by name: {args.name}" + ) + tool = tool_match[0] + + # We cannot provide a full-fledged ToolContext so we need to provide only what is needed by the tool + result = await tool.run_async( + args=args.arguments, + tool_context=args.tool_context, # type:ignore + ) + return _CallToolResult(result=result, tool_context=args.tool_context) + + async def heartbeat_every(delay: float, *details: Any) -> None: + """Heartbeat every ``delay`` seconds until cancelled.""" + while True: + await asyncio.sleep(delay) + activity.heartbeat(*details) + + @activity.defn(name=self._name + "-server-session") + async def connect( + args: _StatefulServerSessionArguments | None = None, + ) -> None: + heartbeat_task = asyncio.create_task(heartbeat_every(30)) + + server_id = self._name + "@" + (activity.info().workflow_run_id or "") + if server_id in self._toolsets: + raise ApplicationError( + "Cannot connect to an already running toolset. Use a distinct " + "name if running multiple stateful toolsets in one workflow." + ) + toolset = self._toolset_factory(args.factory_argument if args else None) + try: + self._toolsets[server_id] = toolset + try: + worker = Worker( + activity.client(), + task_queue=server_id, + activities=[get_tools, call_tool], + activity_task_poller_behavior=PollerBehaviorSimpleMaximum(1), + ) + await worker.run() + finally: + await toolset.close() + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass + finally: + del self._toolsets[server_id] + + return (connect,) + + +class _TemporalStatefulTool(BaseTool): + def __init__( + self, + set_name: str, + config: ActivityConfig, + declaration: FunctionDeclaration | None, + *, + name: str, + description: str, + is_long_running: bool = False, + custom_metadata: dict[str, Any] | None = None, + ): + super().__init__( + name=name, + description=description, + is_long_running=is_long_running, + custom_metadata=custom_metadata, + ) + self._set_name = set_name + self._config = config + self._declaration = declaration + + def _get_declaration(self) -> types.FunctionDeclaration | None: + return self._declaration + + @_handle_worker_failure + async def run_async( + self, *, args: dict[str, Any], tool_context: ToolContext + ) -> Any: + result: _CallToolResult = await workflow.execute_activity( + self._set_name + "-call-tool", + _StatefulCallToolArguments( + self.name, + arguments=args, + tool_context=TemporalToolContext( + tool_confirmation=tool_context.tool_confirmation, + function_call_id=tool_context.function_call_id, + event_actions=tool_context._event_actions, + ), + ), + result_type=_CallToolResult, + **self._config, + ) + + # We need to propagate any event actions back to the main context + tool_context._event_actions = result.tool_context.event_actions + return result.result + + +class TemporalStatefulMcpToolSet(BaseToolset): + """Workflow-side handle for a stateful MCP toolset. + + .. warning:: + This class is experimental and may change in future versions. + Use with caution in production environments. + + Use it as an async context manager inside a workflow so the dedicated + ``-server-session`` worker is started on ``connect`` and torn down on + ``cleanup``:: + + async with TemporalStatefulMcpToolSet("my-tools") as toolset: + agent = Agent(..., tools=[toolset]) + await runner.run(...) + + The connection (and any MCP subprocess) lives for the duration of the + ``async with`` block and is reused across every tool call in that block, + then closed on exit. Callers must be prepared to handle + ``ApplicationError(type="DedicatedWorkerFailure")`` should the dedicated + worker die mid-run. + """ + + def __init__( + self, + name: str, + config: ActivityConfig | None = None, + server_session_config: ActivityConfig | None = None, + factory_argument: Any | None = None, + not_in_workflow_toolset: Callable[[Any | None], McpToolset] | None = None, + ): + """Initializes the stateful Temporal MCP toolset handle. + + Args: + name: Base name of the toolset. Must match the ``name`` passed to + the :class:`TemporalStatefulMcpToolSetProvider` (the + ``-stateful`` suffix is applied here automatically). + config: Optional activity configuration for the per-operation + (``-list-tools``/``-call-tool``) activities. A + ``schedule_to_start_timeout`` is used to detect a dead + dedicated worker. + server_session_config: Optional activity configuration for the + long-running ``-server-session`` activity. + factory_argument: Optional argument passed once to the toolset + factory when the session connects. + not_in_workflow_toolset: Optional factory that returns the + underlying ``McpToolset`` to use when this wrapper executes + outside ``workflow.in_workflow()``, such as local ADK runs. + """ + super().__init__() + self._name = name + "-stateful" + self._config = config or ActivityConfig( + start_to_close_timeout=timedelta(minutes=1), + schedule_to_start_timeout=timedelta(seconds=30), + ) + self._server_session_config = server_session_config or ActivityConfig( + start_to_close_timeout=timedelta(hours=1), + ) + self._factory_argument = factory_argument + self._not_in_workflow_toolset = not_in_workflow_toolset + self._connect_handle: ActivityHandle | None = None + + async def connect(self) -> None: + """Starts the dedicated ``-server-session`` activity for this run.""" + if not workflow.in_workflow(): + return + self._config["task_queue"] = self._name + "@" + workflow.info().run_id + self._connect_handle = workflow.start_activity( + self._name + "-server-session", + _StatefulServerSessionArguments(self._factory_argument), + **self._server_session_config, + ) + + async def cleanup(self) -> None: + """Cancels the dedicated session activity and awaits its teardown.""" + if self._connect_handle: + self._connect_handle.cancel() + try: + await self._connect_handle + except Exception as e: + if not is_cancelled_exception(e): + raise + finally: + self._connect_handle = None + + async def close(self) -> None: + """``BaseToolset`` teardown hook; delegates to :meth:`cleanup`.""" + await self.cleanup() + + async def __aenter__(self) -> "TemporalStatefulMcpToolSet": + await self.connect() + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.cleanup() + + @_handle_worker_failure + async def get_tools( + self, readonly_context: ReadonlyContext | None = None + ) -> list[BaseTool]: + """Retrieves available tools from the stateful MCP toolset.""" + # If executed outside a workflow, like when doing local adk runs, use the mcp server directly + if not workflow.in_workflow(): + if self._not_in_workflow_toolset is None: + raise ValueError( + "Attempted to use TemporalStatefulMcpToolSet outside a " + "workflow, but no not_in_workflow_toolset was provided. " + "Either use McpToolSet directly or pass a factory that " + "returns the underlying McpToolset for non-workflow execution." + ) + return await self._not_in_workflow_toolset( + self._factory_argument + ).get_tools(readonly_context) + + if not self._connect_handle: + raise ApplicationError( + "Stateful MCP toolset not connected. Use it as an async context " + "manager (async with ...) or call connect() first." + ) + + tool_results: list[_ToolResult] = await workflow.execute_activity( + self._name + "-list-tools", + result_type=list[_ToolResult], + **self._config, + ) + return [ + _TemporalStatefulTool( + set_name=self._name, + config=self._config, + declaration=tool_result.function_declaration, + name=tool_result.name, + description=tool_result.description, + is_long_running=tool_result.is_long_running, + custom_metadata=tool_result.custom_metadata, + ) + for tool_result in tool_results + ] diff --git a/temporalio/contrib/google_adk_agents/_plugin.py b/temporalio/contrib/google_adk_agents/_plugin.py index 7344485c8..15b6613e3 100644 --- a/temporalio/contrib/google_adk_agents/_plugin.py +++ b/temporalio/contrib/google_adk_agents/_plugin.py @@ -8,7 +8,10 @@ from typing import Any from temporalio import workflow -from temporalio.contrib.google_adk_agents._mcp import TemporalMcpToolSetProvider +from temporalio.contrib.google_adk_agents._mcp import ( + TemporalMcpToolSetProvider, + TemporalStatefulMcpToolSetProvider, +) from temporalio.contrib.google_adk_agents._model import ( invoke_model, invoke_model_streaming, @@ -72,12 +75,18 @@ class GoogleAdkPlugin(SimplePlugin): def __init__( self, - toolset_providers: list[TemporalMcpToolSetProvider] | None = None, + toolset_providers: list[ + TemporalMcpToolSetProvider | TemporalStatefulMcpToolSetProvider + ] + | None = None, ): """Initializes the Temporal ADK Plugin. Args: - toolset_providers: Optional list of toolset providers for MCP integration. + toolset_providers: Optional list of stateless + (:class:`TemporalMcpToolSetProvider`) or stateful + (:class:`TemporalStatefulMcpToolSetProvider`) toolset providers + for MCP integration. """ @asynccontextmanager diff --git a/tests/contrib/google_adk_agents/test_stateful_mcp.py b/tests/contrib/google_adk_agents/test_stateful_mcp.py new file mode 100644 index 000000000..40a49c35d --- /dev/null +++ b/tests/contrib/google_adk_agents/test_stateful_mcp.py @@ -0,0 +1,187 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests for the stateful (pooled) MCP toolset support in +``temporalio.contrib.google_adk_agents._mcp``. + +These drive the real ``connect`` -> dedicated-worker -> ``get_tools`` path +through a Temporal workflow, but against an in-memory fake ``McpToolset`` (no +subprocess, no real MCP server) so they run in CI. They prove the properties +the reviewer asked for: one toolset is built per workflow run, ``factory_argument`` +is consumed exactly once per run, distinct runs never share a toolset, and the +toolset is torn down when the workflow completes. +""" + +import uuid +from datetime import timedelta +from typing import Any + +from temporalio import workflow +from temporalio.client import Client +from temporalio.worker import Worker + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.google_adk_agents import ( + GoogleAdkPlugin, + TemporalStatefulMcpToolSet, + TemporalStatefulMcpToolSetProvider, + ) + +# Populated inside the activity worker process (same process as the test) each +# time the toolset factory runs, so tests can inspect what was built. +CREATED: list["_FakeToolset"] = [] + + +class _FakeTool: + def __init__(self, name: str) -> None: + self.name = name + self.description = "a fake tool" + self.is_long_running = False + self.custom_metadata: dict[str, Any] | None = None + + def _get_declaration(self) -> None: + return None + + async def run_async(self, *, args: dict[str, Any], tool_context: Any) -> Any: + return {"echo": args} + + +class _FakeToolset: + """In-memory stand-in for ``google.adk.tools.mcp_tool.McpToolset``.""" + + def __init__(self, factory_argument: Any = None) -> None: + self.factory_argument = factory_argument + self.get_tools_calls = 0 + self.closed = False + + async def get_tools(self) -> list[_FakeTool]: + self.get_tools_calls += 1 + return [_FakeTool("echo")] + + async def close(self) -> None: + self.closed = True + + +def _factory(arg: Any) -> _FakeToolset: + toolset = _FakeToolset(arg) + CREATED.append(toolset) + return toolset # type: ignore[return-value] + + +@workflow.defn +class StatefulMcpWorkflow: + @workflow.run + async def run(self, factory_argument: Any | None) -> list[str]: + async with TemporalStatefulMcpToolSet( + "stateful_set", + factory_argument=factory_argument, + ) as toolset: + # Two calls against the same persistent toolset -- exercises reuse. + tools_1 = await toolset.get_tools() + tools_2 = await toolset.get_tools() + return [t.name for t in tools_1] + [t.name for t in tools_2] + + +def _make_client( + client: Client, provider: TemporalStatefulMcpToolSetProvider +) -> Client: + new_config = client.config() + new_config["plugins"] = [GoogleAdkPlugin(toolset_providers=[provider])] + return Client(**new_config) + + +async def test_stateful_one_toolset_per_run_and_teardown(client: Client): + CREATED.clear() + provider = TemporalStatefulMcpToolSetProvider("stateful_set", _factory) + client = _make_client(client, provider) + + async with Worker( + client, + task_queue="adk-stateful-mcp", + workflows=[StatefulMcpWorkflow], + max_cached_workflows=0, + ): + result = await client.execute_workflow( + StatefulMcpWorkflow.run, + None, + id=f"stateful-mcp-{uuid.uuid4()}", + task_queue="adk-stateful-mcp", + execution_timeout=timedelta(seconds=60), + ) + + assert result == ["echo", "echo"] + # Exactly one toolset built for the whole run... + assert len(CREATED) == 1 + # ...reused across both get_tools calls (state maintained, not per-call)... + assert CREATED[0].get_tools_calls == 2 + # ...and closed when the workflow completed (no leak). + assert CREATED[0].closed + # The dedicated-worker toolset registry is empty after teardown. + assert provider._toolsets == {} + + +async def test_stateful_factory_argument_consumed_once(client: Client): + CREATED.clear() + provider = TemporalStatefulMcpToolSetProvider("stateful_set", _factory) + client = _make_client(client, provider) + + async with Worker( + client, + task_queue="adk-stateful-mcp-arg", + workflows=[StatefulMcpWorkflow], + max_cached_workflows=0, + ): + await client.execute_workflow( + StatefulMcpWorkflow.run, + {"tenant": "acme"}, + id=f"stateful-mcp-{uuid.uuid4()}", + task_queue="adk-stateful-mcp-arg", + execution_timeout=timedelta(seconds=60), + ) + + assert len(CREATED) == 1 + assert CREATED[0].factory_argument == {"tenant": "acme"} + + +async def test_stateful_no_cross_run_sharing(client: Client): + CREATED.clear() + provider = TemporalStatefulMcpToolSetProvider("stateful_set", _factory) + client = _make_client(client, provider) + + async with Worker( + client, + task_queue="adk-stateful-mcp-iso", + workflows=[StatefulMcpWorkflow], + max_cached_workflows=0, + ): + await client.execute_workflow( + StatefulMcpWorkflow.run, + "tenant-a", + id=f"stateful-mcp-a-{uuid.uuid4()}", + task_queue="adk-stateful-mcp-iso", + execution_timeout=timedelta(seconds=60), + ) + await client.execute_workflow( + StatefulMcpWorkflow.run, + "tenant-b", + id=f"stateful-mcp-b-{uuid.uuid4()}", + task_queue="adk-stateful-mcp-iso", + execution_timeout=timedelta(seconds=60), + ) + + # Two distinct runs -> two distinct toolsets, each with its own argument. + assert len(CREATED) == 2 + assert {t.factory_argument for t in CREATED} == {"tenant-a", "tenant-b"} + assert all(t.closed for t in CREATED) + assert provider._toolsets == {} From 57c91428610bcaa947886fea3323f062c34a1a58 Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Wed, 22 Jul 2026 17:48:51 +0530 Subject: [PATCH 5/7] fix(contrib): guarantee heartbeat task cleanup in stateful MCP server-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. --- temporalio/contrib/google_adk_agents/_mcp.py | 45 ++++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index 0694f6a11..feb71437d 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -476,34 +476,43 @@ async def heartbeat_every(delay: float, *details: Any) -> None: async def connect( args: _StatefulServerSessionArguments | None = None, ) -> None: - heartbeat_task = asyncio.create_task(heartbeat_every(30)) - server_id = self._name + "@" + (activity.info().workflow_run_id or "") if server_id in self._toolsets: raise ApplicationError( "Cannot connect to an already running toolset. Use a distinct " "name if running multiple stateful toolsets in one workflow." ) - toolset = self._toolset_factory(args.factory_argument if args else None) + + # Created only once the duplicate-connect check has passed, and + # cancelled in the outermost ``finally`` below so it is torn down + # on every exit path -- including a ``toolset.close()`` failure, + # which would otherwise leave it heartbeating forever after this + # activity has already exited. + heartbeat_task = asyncio.create_task(heartbeat_every(30)) try: - self._toolsets[server_id] = toolset + toolset = self._toolset_factory(args.factory_argument if args else None) try: - worker = Worker( - activity.client(), - task_queue=server_id, - activities=[get_tools, call_tool], - activity_task_poller_behavior=PollerBehaviorSimpleMaximum(1), - ) - await worker.run() - finally: - await toolset.close() - heartbeat_task.cancel() + self._toolsets[server_id] = toolset try: - await heartbeat_task - except asyncio.CancelledError: - pass + worker = Worker( + activity.client(), + task_queue=server_id, + activities=[get_tools, call_tool], + activity_task_poller_behavior=PollerBehaviorSimpleMaximum( + 1 + ), + ) + await worker.run() + finally: + await toolset.close() + finally: + del self._toolsets[server_id] finally: - del self._toolsets[server_id] + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass return (connect,) From 02f2cd768fa65f3a3246b5f2138f245da45f5b82 Mon Sep 17 00:00:00 2001 From: Brian Strauch Date: Wed, 22 Jul 2026 07:09:01 -0700 Subject: [PATCH 6/7] Update temporalio/contrib/google_adk_agents/_mcp.py --- temporalio/contrib/google_adk_agents/_mcp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index feb71437d..211ace69e 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -166,7 +166,7 @@ async def call_tool( ) if len(tool_match) > 1: raise ApplicationError( - f"Unable too many matching mcp tools by name: {args.name}" + f"Found multiple MCP tools with the same name: {args.name}" ) tool = tool_match[0] From 726469c5b50cfe2af66d28cd161cf3f986a02745 Mon Sep 17 00:00:00 2001 From: wankhede04 Date: Thu, 23 Jul 2026 12:13:44 +0530 Subject: [PATCH 7/7] fix(contrib): fix duplicate-tool error wording in stateful call_tool Applies the same wording fix @brianstrauch made in 02f2cd7 for the stateless call_tool to the stateful call_tool's identical duplicate-name error, which was missed since it's a separate code path. --- temporalio/contrib/google_adk_agents/_mcp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/temporalio/contrib/google_adk_agents/_mcp.py b/temporalio/contrib/google_adk_agents/_mcp.py index 211ace69e..1f23961da 100644 --- a/temporalio/contrib/google_adk_agents/_mcp.py +++ b/temporalio/contrib/google_adk_agents/_mcp.py @@ -455,7 +455,7 @@ async def call_tool(args: _StatefulCallToolArguments) -> _CallToolResult: ) if len(tool_match) > 1: raise ApplicationError( - f"Unable too many matching mcp tools by name: {args.name}" + f"Found multiple MCP tools with the same name: {args.name}" ) tool = tool_match[0]