From ee3a2729d706a260054d01dd70e1c259906cdf90 Mon Sep 17 00:00:00 2001 From: jsilverhand48 <57627031+jsilverhand48@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:00:03 -0500 Subject: [PATCH] feat(workflow): Support MCP and other toolset tools as workflow nodes A BaseTool can be placed directly into a workflow's edges, but a tool served by a toolset cannot. McpToolset lists its tools over a live connection, so an individual tool only exists after awaiting get_tools(), and Workflow(...) is constructed synchronously. Listing eagerly does not help either: an MCP session is bound to the event loop that opened it, so a session created under asyncio.run() is already unusable by the time the graph runs. ToolsetNode holds the toolset and the tool's name, and resolves the tool while the node runs, on the runner's event loop: toolset = McpToolset(connection_params=StdioConnectionParams(...)) Workflow( name='research', edges=[ ('START', build_args), (build_args, ToolsetNode(toolset=toolset, tool_name='read_file')), (..., summarize), ], ) Resolution goes through BaseToolset.get_tools_with_prefix(), which caches per invocation, so several nodes sharing one toolset list its tools only once per run. The node is typed on BaseToolset rather than McpToolset, so the core workflow package gains no dependency on the optional mcp extra and the same node serves any toolset. Runner._collect_toolset now also walks a Workflow's graph nodes. A Workflow has neither tools nor sub_agents, so a toolset referenced only by a ToolsetNode was never closed and its MCP subprocess outlived the run. The tool invocation body of _ToolNode is extracted into shared helpers so both nodes coerce arguments and emit events identically. Closes google/adk-python#6533 --- .../workflows/mcp_toolset_node/README.md | 75 ++++++ .../workflows/mcp_toolset_node/agent.py | 73 ++++++ docs/guides/README.md | 1 + docs/guides/workflow/toolset_node/index.md | 98 ++++++++ src/google/adk/runners.py | 13 +- src/google/adk/workflow/__init__.py | 3 + src/google/adk/workflow/_tool_node.py | 104 ++++---- src/google/adk/workflow/_toolset_node.py | 149 ++++++++++++ tests/unittests/test_import_loading.py | 16 ++ tests/unittests/test_runners.py | 44 ++++ tests/unittests/workflow/test_toolset_node.py | 225 ++++++++++++++++++ 11 files changed, 760 insertions(+), 41 deletions(-) create mode 100644 contributing/samples/workflows/mcp_toolset_node/README.md create mode 100644 contributing/samples/workflows/mcp_toolset_node/agent.py create mode 100644 docs/guides/workflow/toolset_node/index.md create mode 100644 src/google/adk/workflow/_toolset_node.py create mode 100644 tests/unittests/workflow/test_toolset_node.py diff --git a/contributing/samples/workflows/mcp_toolset_node/README.md b/contributing/samples/workflows/mcp_toolset_node/README.md new file mode 100644 index 00000000000..6dd10765554 --- /dev/null +++ b/contributing/samples/workflows/mcp_toolset_node/README.md @@ -0,0 +1,75 @@ +# ADK Workflow MCP Toolset Node Sample + +## Overview + +This sample demonstrates calling a tool from an **MCP server** as a step in an +**ADK Workflow**, using `ToolsetNode`. + +A `BaseTool` can be dropped straight into a workflow's `edges`. An MCP tool +cannot, because it only exists after `await McpToolset.get_tools()` lists it over +a live connection, and `Workflow(...)` is constructed synchronously. +`ToolsetNode` closes that gap: it holds the toolset and the tool's name, and +resolves the tool when the node runs, on the runner's event loop. + +The workflow reads a file through the MCP filesystem server: + +1. `build_args` turns the user's message into the tool's arguments. +1. `ToolsetNode` resolves `read_file` from the toolset and calls it. +1. `summarize` formats the MCP response. + +Resolution goes through `BaseToolset.get_tools_with_prefix()`, which caches per +invocation, so several `ToolsetNode`s sharing one toolset list the server's tools +only once per run. + +## Prerequisites + +`npx` must be on your `PATH`; the sample launches +`@modelcontextprotocol/server-filesystem` as a stdio subprocess scoped to this +directory. Install the MCP extra with `pip install "google-adk[mcp]"`. + +## Sample Inputs + +- `README.md` + +- `agent.py` + +## Graph + +```mermaid +graph TD + START --> build_args + build_args --> read_file + read_file --> summarize +``` + +## How To + +1. **Declare the toolset normally.** No `await` is needed at module scope; the + server is contacted only while the workflow runs. + + ```python + filesystem_toolset = McpToolset( + connection_params=StdioConnectionParams(...), + tool_filter=['read_file', 'list_directory'], + ) + ``` + +1. **Name the tool you want as a node.** The node's input is the tool's argument + dict (or a JSON object string, or `None` for no arguments), and the tool's + response becomes the node's output. + + ```python + ToolsetNode(toolset=filesystem_toolset, tool_name='read_file') + ``` + +1. **Name the node explicitly when the tool's name is not a Python + identifier.** Node names must be identifiers, so a tool named `read-file` + becomes a node named `read_file` by default. Pass `name=` to choose your own. + + ```python + ToolsetNode(toolset=toolset, tool_name='read-file', name='reader') + ``` + +The `Runner` closes the toolset when it shuts down, including when the toolset +is referenced only by a `ToolsetNode` inside a workflow, so the MCP subprocess +does not outlive the run. diff --git a/contributing/samples/workflows/mcp_toolset_node/agent.py b/contributing/samples/workflows/mcp_toolset_node/agent.py new file mode 100644 index 00000000000..433ddf3ee67 --- /dev/null +++ b/contributing/samples/workflows/mcp_toolset_node/agent.py @@ -0,0 +1,73 @@ +# Copyright 2026 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. + + +import os + +from google.adk import Workflow +from google.adk.tools.mcp_tool import StdioConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset +from google.adk.workflow import ToolsetNode +from mcp import StdioServerParameters + +_allowed_path = os.path.dirname(os.path.abspath(__file__)) + +# The MCP server is only contacted while the workflow runs, so the toolset can +# be declared here even though listing its tools requires a live connection. +filesystem_toolset = McpToolset( + connection_params=StdioConnectionParams( + server_params=StdioServerParameters( + command='npx', + args=[ + '-y', + '@modelcontextprotocol/server-filesystem', + _allowed_path, + ], + ), + timeout=15, + ), + tool_filter=['read_file', 'list_directory'], +) + + +def build_args(node_input: str): + """Turns the user's message into arguments for the MCP tool.""" + filename = node_input.strip() or 'README.md' + return {'path': os.path.join(_allowed_path, filename)} + + +def summarize(node_input: dict): + """Formats the MCP tool's response for display.""" + if node_input.get('isError'): + return f'The MCP server reported an error: {node_input}' + texts = [ + part.get('text', '') + for part in node_input.get('content', []) + if part.get('type') == 'text' + ] + body = '\n'.join(texts) + return f'Read {len(body)} characters from the MCP server:\n\n{body}' + + +root_agent = Workflow( + name='mcp_toolset_node_sample', + edges=[ + ( + 'START', + build_args, + ToolsetNode(toolset=filesystem_toolset, tool_name='read_file'), + summarize, + ), + ], +) diff --git a/docs/guides/README.md b/docs/guides/README.md index 0ee1513566b..66d3de8977b 100644 --- a/docs/guides/README.md +++ b/docs/guides/README.md @@ -25,6 +25,7 @@ This directory contains specific developer guides for the ADK Python implementat * [Workflow Graphs](workflow/graph/index.md) - Understanding nodes, edges, and graph structures in workflows. * [Function Nodes](workflow/function_node/index.md) - Wrapping Python functions and generators as workflow nodes. * [JoinNode](workflow/join_node/index.md) - Synchronizing parallel execution paths in workflows. +* [ToolsetNode](workflow/toolset_node/index.md) - Running a tool from a toolset, such as an MCP server, as a workflow node. * [RetryConfig](workflow/retry_config/index.md) - Configuring retry policies for resilient workflow nodes. * [ParallelWorker](workflow/parallel_worker/index.md) - Processing lists of items concurrently in workflows. * [Dynamic Nodes](workflow/dynamic_nodes/index.md) - Scheduling and executing nodes dynamically at runtime. diff --git a/docs/guides/workflow/toolset_node/index.md b/docs/guides/workflow/toolset_node/index.md new file mode 100644 index 00000000000..745e01f54af --- /dev/null +++ b/docs/guides/workflow/toolset_node/index.md @@ -0,0 +1,98 @@ +# ToolsetNode + +`ToolsetNode` is a built-in workflow node that runs a named tool from a toolset, resolving the tool when the node runs rather than when the graph is built. + +## Introduction + +A `BaseTool` can be placed directly into a workflow's `edges`, and the framework wraps it as a node. That works because the tool object already exists when the graph is constructed. + +Tools served by a toolset are different. `McpToolset`, for example, lists its tools over a live connection, so an individual tool only exists after `await McpToolset.get_tools()`. `Workflow(name=..., edges=[...])` is constructed synchronously, leaving nowhere to await that call. Listing the tools eagerly at import time does not help either: the MCP session is bound to the event loop that created it, so a session opened under `asyncio.run(...)` is unusable by the time the runner executes the graph. + +`ToolsetNode` closes that gap. It holds the toolset and the name of the tool you want, and resolves the tool while the node runs, on the runner's event loop and inside the live invocation. + +Key features: +- **Lazy resolution**: The toolset is only contacted while the workflow runs. +- **Cached per invocation**: Resolution goes through `BaseToolset.get_tools_with_prefix()`, so several `ToolsetNode`s sharing one toolset list its tools only once per run. +- **Toolset-agnostic**: Works with any `BaseToolset`, not only `McpToolset`. + +## Get started + +The following example reads a file through the MCP filesystem server as one step of a workflow. + +```python +import os + +from google.adk import Workflow +from google.adk.tools.mcp_tool import StdioConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset +from google.adk.workflow import ToolsetNode +from mcp import StdioServerParameters + +filesystem_toolset = McpToolset( + connection_params=StdioConnectionParams( + server_params=StdioServerParameters( + command="npx", + args=["-y", "@modelcontextprotocol/server-filesystem", os.getcwd()], + ), + timeout=15, + ), + tool_filter=["read_file", "list_directory"], +) + + +def build_args(node_input: str) -> dict: + return {"path": node_input.strip()} + + +def summarize(node_input: dict) -> str: + texts = [ + part.get("text", "") + for part in node_input.get("content", []) + if part.get("type") == "text" + ] + return "\n".join(texts) + + +root_agent = Workflow( + name="root_agent", + edges=[( + "START", + build_args, + ToolsetNode(toolset=filesystem_toolset, tool_name="read_file"), + summarize, + )], +) +``` + +## How it works + +1. **Resolution**: When the node runs, it builds a `ReadonlyContext` from the invocation and calls `toolset.get_tools_with_prefix(readonly_context)`, then selects the tool whose name equals `tool_name`. If no tool matches, it raises a `ValueError` listing the names that were available. +2. **Argument coercion**: The node input becomes the tool's arguments. A dict is used as-is, a JSON object string is parsed, and `None` or an empty string means no arguments. Any other input raises a `TypeError`. +3. **Execution**: The resolved tool is called via `run_async`, and its response becomes the node's output. State the tool writes to its context is propagated to downstream nodes. + +## Configuration options + +| Option | Description | +|---|---| +| `toolset` | The `BaseToolset` to resolve the tool from. Required. | +| `tool_name` | The name of the tool to run. Matched against the names the toolset reports, so it includes the toolset's `tool_name_prefix` if one is set. Required. | +| `name` | The node's name. Defaults to `tool_name` with any character that is not valid in a Python identifier replaced by an underscore. | +| `description` | A human-readable description of what the node does. | +| `retry_config` | Configuration for retrying the node on failure. See [RetryConfig](../retry_config/index.md). | +| `timeout` | Maximum time in seconds for the node to complete. | + +### Node names + +Node names must be valid Python identifiers, but tool names are not constrained that way; MCP servers commonly use dashes. A tool named `read-file` therefore becomes a node named `read_file` by default. Pass `name=` to choose your own, which you will need to do if two tools would otherwise sanitize to the same node name. + +```python +ToolsetNode(toolset=toolset, tool_name="read-file", name="reader") +``` + +## Lifecycle + +A `Runner` closes the toolsets it finds on the agent it runs, including toolsets referenced only by a `ToolsetNode` inside a workflow, so an MCP server subprocess does not outlive the run. If you drive a workflow without a `Runner`, call `await toolset.close()` yourself. + +## Limitations + +A tool that needs to interrupt the run cannot do so through a `ToolsetNode`. If the toolset is configured with an `auth_scheme`, or its tools are configured with `require_confirmation`, the tool's request for credentials or confirmation is not surfaced to the client, and the node emits the tool's placeholder response instead. Use a [FunctionNode](../function_node/index.md) with `auth_config` when a workflow step needs user authentication. diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index dbccb4a89c4..0c55a0ceaaf 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -65,6 +65,7 @@ if TYPE_CHECKING: from .apps.app import App from .apps.app import ResumabilityConfig + from .workflow._base_node import BaseNode logger = logging.getLogger('google_adk.' + __name__) @@ -2229,8 +2230,12 @@ async def _handle_new_message( state_delta=state_delta, ) - def _collect_toolset(self, agent: BaseAgent) -> set[BaseToolset]: + def _collect_toolset(self, agent: BaseAgent | BaseNode) -> set[BaseToolset]: + from .workflow._toolset_node import ToolsetNode + toolsets = set() + if isinstance(agent, ToolsetNode): + toolsets.add(agent.toolset) if hasattr(agent, 'tools'): for tool_union in agent.tools: if isinstance(tool_union, BaseToolset): @@ -2238,6 +2243,12 @@ def _collect_toolset(self, agent: BaseAgent) -> set[BaseToolset]: if hasattr(agent, 'sub_agents'): for sub_agent in agent.sub_agents: toolsets.update(self._collect_toolset(sub_agent)) + # A Workflow holds its nodes in a graph rather than in sub_agents, so its + # toolsets are only reachable this way. + graph = getattr(agent, 'graph', None) + if graph is not None: + for graph_node in graph.nodes: + toolsets.update(self._collect_toolset(graph_node)) return toolsets async def _cleanup_toolsets(self, toolsets_to_close: set[BaseToolset]): diff --git a/src/google/adk/workflow/__init__.py b/src/google/adk/workflow/__init__.py index b18156f281b..0e885d7515e 100644 --- a/src/google/adk/workflow/__init__.py +++ b/src/google/adk/workflow/__init__.py @@ -29,6 +29,7 @@ from ._node import Node from ._node import node from ._retry_config import RetryConfig + from ._toolset_node import ToolsetNode from ._workflow import Workflow _LAZY_MEMBERS: dict[str, str] = { @@ -41,6 +42,7 @@ 'NodeTimeoutError': '._errors', 'RetryConfig': '._retry_config', 'START': '._base_node', + 'ToolsetNode': '._toolset_node', 'Workflow': '._workflow', 'node': '._node', } @@ -54,6 +56,7 @@ 'NodeTimeoutError', 'RetryConfig', 'START', + 'ToolsetNode', 'Workflow', 'node', ] diff --git a/src/google/adk/workflow/_tool_node.py b/src/google/adk/workflow/_tool_node.py index 68626e51e58..5216ee6b374 100644 --- a/src/google/adk/workflow/_tool_node.py +++ b/src/google/adk/workflow/_tool_node.py @@ -35,6 +35,68 @@ from ._retry_config import RetryConfig +def _coerce_tool_args(node_input: Any) -> dict[str, Any]: + """Coerces a node input into a dict of tool arguments. + + ``types.Content`` is reduced to its text, a JSON string is parsed, and an + empty or whitespace-only string (like ``None``) means no arguments. + + Raises: + TypeError: If the input cannot be interpreted as a dict of arguments. + """ + args = node_input + if isinstance(args, types.Content): + args = extract_text_from_content(args) + + if isinstance(args, str): + args = args.strip() + if not args: + args = None + else: + try: + args = json.loads(args) + except json.JSONDecodeError: + pass + + if args is None: + return {} + if not isinstance(args, dict): + raise TypeError( + 'The input to ToolNode must be a dictionary of tool arguments or' + f' None, but got {type(args)}.' + ) + return args + + +async def _run_tool( + tool: BaseTool, + *, + ctx: Context, + node_input: Any, +) -> AsyncGenerator[Event, None]: + """Runs a tool with the node input as its arguments and yields its output.""" + tool_context = ToolContext( + invocation_context=ctx.get_invocation_context(), + function_call_id=str(uuid.uuid4()), + ) + + args = _coerce_tool_args(node_input) + + response = await tool.run_async(args=args, tool_context=tool_context) + state_delta = ( + dict(tool_context.actions.state_delta) + if tool_context.actions.state_delta + else None + ) + if response is not None: + yield Event( + output=response, + state=state_delta, + ) + elif state_delta: + yield Event(state=state_delta) + + class _ToolNode(BaseNode): """A node that wraps an ADK Tool.""" @@ -64,43 +126,5 @@ async def _run_impl( ctx: Context, node_input: Any, ) -> AsyncGenerator[Any, None]: - tool_context = ToolContext( - invocation_context=ctx.get_invocation_context(), - function_call_id=str(uuid.uuid4()), - ) - - args = node_input - if isinstance(args, types.Content): - args = extract_text_from_content(args) - - if isinstance(args, str): - args = args.strip() - if not args: - args = None - else: - try: - args = json.loads(args) - except json.JSONDecodeError: - pass - - if args is None: - args = {} - elif not isinstance(args, dict): - raise TypeError( - 'The input to ToolNode must be a dictionary of tool arguments or' - f' None, but got {type(args)}.' - ) - - response = await self.tool.run_async(args=args, tool_context=tool_context) - state_delta = ( - dict(tool_context.actions.state_delta) - if tool_context.actions.state_delta - else None - ) - if response is not None: - yield Event( - output=response, - state=state_delta, - ) - elif state_delta: - yield Event(state=state_delta) + async for event in _run_tool(self.tool, ctx=ctx, node_input=node_input): + yield event diff --git a/src/google/adk/workflow/_toolset_node.py b/src/google/adk/workflow/_toolset_node.py new file mode 100644 index 00000000000..46b6bf9748b --- /dev/null +++ b/src/google/adk/workflow/_toolset_node.py @@ -0,0 +1,149 @@ +# Copyright 2026 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. + +from __future__ import annotations + +"""A node that runs a named tool from a toolset.""" + +from collections.abc import AsyncGenerator +import re +from typing import Any + +from pydantic import ConfigDict +from pydantic import Field +from typing_extensions import override + +from ..agents.context import Context +from ..agents.readonly_context import ReadonlyContext +from ..tools.base_tool import BaseTool +from ..tools.base_toolset import BaseToolset +from ._base_node import BaseNode +from ._retry_config import RetryConfig +from ._tool_node import _run_tool + + +def _to_node_name(tool_name: str) -> str: + """Converts a tool name into a valid Python identifier. + + Tool names from external servers are not constrained to Python identifiers + (MCP servers commonly use dashes), but node names are. + """ + name = re.sub(r'\W', '_', tool_name) + if not name or name[0].isdigit(): + name = f'_{name}' + return name + + +class ToolsetNode(BaseNode): + """A node that runs a named tool from a toolset. + + Unlike passing a ``BaseTool`` directly into a workflow's edges, the tool is + resolved lazily when the node runs, rather than when the graph is built. This + is what makes toolsets whose tools are only discoverable asynchronously -- + such as ``McpToolset``, which lists them over a live connection -- usable as + workflow nodes:: + + toolset = McpToolset(connection_params=StdioConnectionParams(...)) + + workflow = Workflow( + name='research', + edges=[ + (START, build_query), + (build_query, ToolsetNode(toolset=toolset, tool_name='search')), + (ToolsetNode(...), summarize), + ], + ) + + Resolution goes through ``BaseToolset.get_tools_with_prefix()``, which caches + per invocation, so several ``ToolsetNode``s sharing one toolset only list its + tools once per run. + + The node input must be a dict of tool arguments, a JSON object string, or + ``None`` for no arguments. The tool's response becomes the node's output. + + Closing the toolset remains the caller's responsibility, except when the + workflow is run by a ``Runner``, which closes the toolsets it finds on the + agent it runs. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + toolset: BaseToolset = Field(...) + """The toolset to resolve the tool from.""" + + tool_name: str = Field(...) + """The name of the tool to run. + + This is matched against the tool names the toolset reports, so it includes + the toolset's ``tool_name_prefix`` if one is set. + """ + + def __init__( + self, + *, + toolset: BaseToolset, + tool_name: str, + name: str | None = None, + description: str = '', + retry_config: RetryConfig | None = None, + timeout: float | None = None, + ): + """Initializes the ToolsetNode. + + Args: + toolset: The toolset to resolve the tool from. + tool_name: The name of the tool to run. + name: The node's name. Defaults to ``tool_name`` with any character that + is not valid in a Python identifier replaced by an underscore. + description: A human-readable description of what this node does. + retry_config: Configuration for retrying the node on failure. + timeout: Maximum time in seconds for this node to complete. + """ + super().__init__( + toolset=toolset, + tool_name=tool_name, + name=name or _to_node_name(tool_name), + description=description, + rerun_on_resume=False, + retry_config=retry_config, + timeout=timeout, + ) + + async def _resolve_tool(self, ctx: Context) -> BaseTool: + """Finds the named tool in the toolset. + + Raises: + ValueError: If the toolset does not offer a tool by that name. + """ + readonly_context = ReadonlyContext(ctx.get_invocation_context()) + tools = await self.toolset.get_tools_with_prefix(readonly_context) + for tool in tools: + if tool.name == self.tool_name: + return tool + available = ', '.join(sorted(tool.name for tool in tools)) or '' + raise ValueError( + f"Tool '{self.tool_name}' was not found in" + f' {type(self.toolset).__name__}. Available tools: {available}.' + ) + + @override + async def _run_impl( + self, + *, + ctx: Context, + node_input: Any, + ) -> AsyncGenerator[Any, None]: + tool = await self._resolve_tool(ctx) + async for event in _run_tool(tool, ctx=ctx, node_input=node_input): + yield event diff --git a/tests/unittests/test_import_loading.py b/tests/unittests/test_import_loading.py index 24193917f93..4a377a2d90e 100644 --- a/tests/unittests/test_import_loading.py +++ b/tests/unittests/test_import_loading.py @@ -77,6 +77,7 @@ 'google.adk.workflow._function_node', 'google.adk.workflow._join_node', 'google.adk.workflow._node', + 'google.adk.workflow._toolset_node', 'google.adk.workflow._workflow', ), ), @@ -119,6 +120,21 @@ def test_constructing_agent_defers_optional_mcp_server_stack(): ) +def test_toolset_node_defers_optional_mcp_stack(): + """ToolsetNode serves McpToolset without depending on the mcp extra.""" + if importlib.util.find_spec('mcp') is None: + pytest.skip('MCP import-boundary check requires the declared test extra.') + + assert_modules_unloaded( + """ +from google.adk.workflow import ToolsetNode + +assert ToolsetNode.__name__ == 'ToolsetNode' +""", + ('mcp',), + ) + + def test_lazy_packages_support_star_imports(): """Every lazy package still resolves through Python's public import syntax.""" result = run_isolated(f""" diff --git a/tests/unittests/test_runners.py b/tests/unittests/test_runners.py index b76953eab61..de1054bf845 100644 --- a/tests/unittests/test_runners.py +++ b/tests/unittests/test_runners.py @@ -41,6 +41,9 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.sessions.session import Session from google.adk.tools.base_toolset import BaseToolset +from google.adk.workflow import START +from google.adk.workflow import ToolsetNode +from google.adk.workflow import Workflow from google.genai import types import pytest @@ -1271,6 +1274,47 @@ async def close(self) -> None: assert toolset.close_cancelled is False assert toolset.close_finished.is_set() + @pytest.mark.asyncio + async def test_runner_close_closes_toolset_held_by_a_workflow_node(self): + """A toolset reachable only through a ToolsetNode is still closed.""" + + class RecordingToolset(BaseToolset): + + def __init__(self): + super().__init__() + self.closed = False + + async def get_tools(self, readonly_context=None): + del readonly_context + return [] + + async def close(self) -> None: + self.closed = True + + toolset = RecordingToolset() + + def start_node(): + return {} + + # The toolset is referenced by the node, not by any agent's `tools`. + workflow = Workflow( + name="wf", + edges=[ + (START, start_node), + (start_node, ToolsetNode(toolset=toolset, tool_name="search")), + ], + ) + runner = Runner( + app_name="test_app", + agent=workflow, + session_service=self.session_service, + artifact_service=self.artifact_service, + ) + + await runner.close() + + assert toolset.closed is True + @pytest.mark.asyncio async def test_runner_passes_plugin_close_timeout(self): """Test that runner passes plugin_close_timeout to PluginManager.""" diff --git a/tests/unittests/workflow/test_toolset_node.py b/tests/unittests/workflow/test_toolset_node.py new file mode 100644 index 00000000000..8e68669b996 --- /dev/null +++ b/tests/unittests/workflow/test_toolset_node.py @@ -0,0 +1,225 @@ +# Copyright 2026 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. + +"""Tests for ToolsetNode, which runs a named tool resolved from a toolset.""" + +from typing import Any +from typing import Optional + +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.events.event import Event +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.base_toolset import BaseToolset +from google.adk.workflow import START +from google.adk.workflow import ToolsetNode +from google.adk.workflow._workflow import Workflow +import pytest + +from . import workflow_testing_utils +from .. import testing_utils + + +class _EchoTool(BaseTool): + """A tool that returns the args it was called with.""" + + async def run_async(self, *, args: dict[str, Any], tool_context) -> Any: + return args + + +class _RecordingToolset(BaseToolset): + """A toolset that serves fixed tools and counts how often it is listed.""" + + def __init__(self, *tool_names: str, tool_name_prefix: Optional[str] = None): + super().__init__(tool_name_prefix=tool_name_prefix) + self._tool_names = tool_names + self.get_tools_call_count = 0 + + async def get_tools( + self, readonly_context: Optional[ReadonlyContext] = None + ) -> list[BaseTool]: + self.get_tools_call_count += 1 + return [ + _EchoTool(name=name, description=f'Echoes for {name}') + for name in self._tool_names + ] + + +async def _run(wf: Workflow) -> list[Any]: + """Runs a workflow and returns its simplified events.""" + app_instance = testing_utils.App(name='test_app', root_agent=wf) + runner = testing_utils.InMemoryRunner(app=app_instance) + events = await runner.run_async('start') + return workflow_testing_utils.simplify_events_with_node(events) + + +@pytest.mark.asyncio +async def test_named_tool_receives_the_node_input_as_arguments(): + """The tool named by the node is resolved and called with the node input.""" + toolset = _RecordingToolset('search', 'fetch') + args = {'query': 'adk'} + + def start_node(): + return Event(output=args) + + simplified = await _run( + Workflow( + name='wf', + edges=[ + (START, start_node), + (start_node, ToolsetNode(toolset=toolset, tool_name='search')), + ], + ) + ) + + assert ('wf@1/search@1', {'output': args}) in simplified + + +@pytest.mark.asyncio +async def test_unknown_tool_name_reports_the_available_tools(): + """Naming a tool the toolset does not serve fails with a usable message.""" + toolset = _RecordingToolset('search', 'fetch') + + def start_node(): + return Event(output={}) + + wf = Workflow( + name='wf', + edges=[ + (START, start_node), + (start_node, ToolsetNode(toolset=toolset, tool_name='missing')), + ], + ) + + with pytest.raises( + ValueError, match=r"'missing'.*Available tools: fetch, search" + ): + await _run(wf) + + +@pytest.mark.asyncio +async def test_tool_name_matches_the_toolsets_prefixed_name(): + """A toolset's tool_name_prefix is part of the name the node matches.""" + toolset = _RecordingToolset('search', tool_name_prefix='web') + + def start_node(): + return Event(output={'query': 'adk'}) + + simplified = await _run( + Workflow( + name='wf', + edges=[ + (START, start_node), + ( + start_node, + ToolsetNode(toolset=toolset, tool_name='web_search'), + ), + ], + ) + ) + + assert ('wf@1/web_search@1', {'output': {'query': 'adk'}}) in simplified + + +@pytest.mark.asyncio +async def test_nodes_sharing_a_toolset_list_its_tools_once_per_run(): + """Resolution is cached per invocation, so one run lists tools once.""" + toolset = _RecordingToolset('search', 'fetch') + + def start_node(): + return Event(output={}) + + await _run( + Workflow( + name='wf', + edges=[ + (START, start_node), + ( + start_node, + ToolsetNode(toolset=toolset, tool_name='search'), + ToolsetNode(toolset=toolset, tool_name='fetch'), + ), + ], + ) + ) + + assert toolset.get_tools_call_count == 1 + + +@pytest.mark.asyncio +async def test_tool_name_that_is_not_an_identifier_becomes_a_valid_node_name(): + """A dashed tool name, common for MCP servers, yields a usable node name.""" + toolset = _RecordingToolset('read-file') + + def start_node(): + return Event(output={'path': '/tmp/x'}) + + simplified = await _run( + Workflow( + name='wf', + edges=[ + (START, start_node), + (start_node, ToolsetNode(toolset=toolset, tool_name='read-file')), + ], + ) + ) + + assert ('wf@1/read_file@1', {'output': {'path': '/tmp/x'}}) in simplified + + +def test_explicit_name_overrides_the_derived_node_name(): + """An explicit name wins over the name derived from tool_name.""" + node = ToolsetNode( + toolset=_RecordingToolset('read-file'), + tool_name='read-file', + name='reader', + ) + + assert node.name == 'reader' + + +@pytest.mark.asyncio +async def test_state_written_by_the_tool_reaches_later_nodes(): + """State the tool sets on its context is persisted for downstream nodes.""" + + class _StatefulTool(BaseTool): + + async def run_async(self, *, args, tool_context): + tool_context.state['tool_key'] = 'tool_value' + return {'status': 'ok'} + + class _StatefulToolset(BaseToolset): + + async def get_tools(self, readonly_context=None) -> list[BaseTool]: + return [_StatefulTool(name='stateful', description='Sets state')] + + def start_node(): + return Event(output={}) + + def read_state(tool_key: str) -> str: + return f'tool_key={tool_key}' + + tool_node = ToolsetNode(toolset=_StatefulToolset(), tool_name='stateful') + + simplified = await _run( + Workflow( + name='wf', + edges=[ + (START, start_node), + (start_node, tool_node), + (tool_node, read_state), + ], + ) + ) + + assert ('wf@1/read_state@1', {'output': 'tool_key=tool_value'}) in simplified