feat(workflow): Support MCP and other toolset tools as workflow nodes - #6610
Open
jsilverhand48 wants to merge 1 commit into
Open
feat(workflow): Support MCP and other toolset tools as workflow nodes#6610jsilverhand48 wants to merge 1 commit into
jsilverhand48 wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
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):
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:
Please include a summary of passed
pytestresults.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
Additional context
Add any other context or screenshots about the feature request here.