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 92bf994dd..1f23961da 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 @@ -88,16 +100,28 @@ 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] + self, + name: str, + toolset_factory: Callable[[Any | None], McpToolset], ) -> None: """Initializes the toolset provider. Args: name: Name prefix for the generated activities. 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 @@ -108,42 +132,52 @@ def _get_activities(self) -> Sequence[Callable]: async def get_tools( args: _GetToolsArguments, ) -> list[_ToolResult]: + # 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) - 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 - ] + try: + 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: + await toolset.close() @activity.defn(name=self._name + "-call-tool") 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}" + try: + 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"Found multiple MCP tools with the same 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 ) - 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) + return _CallToolResult(result=result, tool_context=args.tool_context) + finally: + await toolset.close() return get_tools, call_tool @@ -281,3 +315,395 @@ 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"Found multiple MCP tools with the same 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: + 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." + ) + + # 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: + 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() + finally: + del self._toolsets[server_id] + finally: + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass + + 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_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_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 == {}