Skip to content

feat(workflow): Support MCP and other toolset tools as workflow nodes - #6610

Open
jsilverhand48 wants to merge 1 commit into
google:mainfrom
jsilverhand48:feat/workflow-toolset-node
Open

feat(workflow): Support MCP and other toolset tools as workflow nodes#6610
jsilverhand48 wants to merge 1 commit into
google:mainfrom
jsilverhand48:feat/workflow-toolset-node

Conversation

@jsilverhand48

Copy link
Copy Markdown

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

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

  • Closes: #issue_number
  • Related: #issue_number

2. Or, if no issue exists, describe the change:

If applicable, please follow the issue templates to provide as much detail as
possible.

Problem:
A clear and concise description of what the problem is.

Solution:
A clear and concise description of what you want to happen and why you choose
this solution.

Testing Plan

Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes.

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Please include a summary of passed pytest results.

Manual End-to-End (E2E) Tests:

Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Add any other context or screenshots about the feature request here.

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#6533
@adk-bot adk-bot added the mcp [Component] This issues is related to MCP support label Aug 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mcp [Component] This issues is related to MCP support

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support MCP and other toolset tools as graph workflow nodes

3 participants