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 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