fix(tools): Add proper cleanup for AgentTool to prevent MCP session errors #3411
Conversation
…rrors Clean up runner resources in AgentTool after sub-agent execution to ensure MCP sessions are closed in the correct async task context. Without this fix, MCP sessions were cleaned up during garbage collection in a different task, causing "Attempted to exit cancel scope in a different task" errors. This fix ensures that `runner.close()` is called immediately after the sub-agent finishes executing, properly closing all MCP sessions and other resources within the same async task context they were created in. Also adds two demo agents showing how to use AgentTool with MCP tools: - mcp_in_agent_tool_remote: Uses SSE mode (remote server connection) - mcp_in_agent_tool_stdio: Uses stdio mode (local subprocess) Both demos use Gemini 2.5 Flash and include zero-installation setup using uvx. Related: #1112, #929
Summary of ChangesHello @Jacksunwei, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical issue where AgentTool failed to properly clean up resources when interacting with Model Context Protocol (MCP) tools, leading to RuntimeError due to incorrect asynchronous task context. The core change involves explicitly calling the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses the RuntimeError related to MCP session cleanup by ensuring runner.close() is called within the correct asynchronous context in AgentTool. The fix is clean and directly targets the problem. The addition of two comprehensive demo agents for both SSE and stdio modes is a fantastic contribution, providing clear, runnable examples for developers. My review includes a couple of suggestions for the new sample agents to improve their maintainability by extracting hardcoded configuration values into constants.
| mcp_toolset = McpToolset( | ||
| connection_params=SseConnectionParams( | ||
| url="http://localhost:3000/sse", | ||
| timeout=10.0, | ||
| sse_read_timeout=300.0, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
To improve maintainability and readability, consider extracting hardcoded values like the URL, timeouts, and the model name ('gemini-2.5-flash') into module-level constants. This centralizes configuration, making it easier to manage and update.
For example:
# At the top of the file
_MCP_SSE_URL = "http://localhost:3000/sse"
_CONNECTION_TIMEOUT = 10.0
_SSE_READ_TIMEOUT = 300.0
_MODEL_NAME = "gemini-2.5-flash"
# In the McpToolset definition
mcp_toolset = McpToolset(
connection_params=SseConnectionParams(
url=_MCP_SSE_URL,
timeout=_CONNECTION_TIMEOUT,
sse_read_timeout=_SSE_READ_TIMEOUT,
)
)
# In agent definitions
sub_agent = Agent(
model=_MODEL_NAME,
# ...
)
root_agent = Agent(
model=_MODEL_NAME,
# ...
)| mcp_toolset = McpToolset( | ||
| connection_params=StdioConnectionParams( | ||
| server_params=StdioServerParameters( | ||
| command="uvx", | ||
| args=[ | ||
| "--from", | ||
| "git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool", | ||
| "mcp-simple-tool", | ||
| ], | ||
| ), | ||
| timeout=10.0, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
To improve readability and make the configuration easier to manage, consider extracting hardcoded values like the Git repository URL, command arguments, and model name into constants.
For example:
# At the top of the file
_MCP_SIMPLE_TOOL_REPO = "git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool"
_MCP_SIMPLE_TOOL_COMMAND = "mcp-simple-tool"
_CONNECTION_TIMEOUT = 10.0
_MODEL_NAME = "gemini-2.5-flash"
# In the McpToolset definition
mcp_toolset = McpToolset(
connection_params=StdioConnectionParams(
server_params=StdioServerParameters(
command="uvx",
args=[
"--from",
_MCP_SIMPLE_TOOL_REPO,
_MCP_SIMPLE_TOOL_COMMAND,
],
),
timeout=_CONNECTION_TIMEOUT,
)
)
# In agent definitions
sub_agent = Agent(
model=_MODEL_NAME,
# ...
)
root_agent = Agent(
model=_MODEL_NAME,
# ...
)
…rrors Merge #3411 ## Summary Fixes AgentTool cleanup to prevent MCP session errors by calling `runner.close()` after sub-agent execution. ## Problem When using AgentTool with MCP tools, the runner cleanup happened during garbage collection in a different async task context, causing: ``` RuntimeError: Attempted to exit cancel scope in a different task than it was entered in ``` ## Solution - Call `await runner.close()` immediately after sub-agent execution in AgentTool - This ensures MCP sessions and other resources are cleaned up in the correct async task context - Updated test mock to include the close() method ## Demo Agents Added two comprehensive demo agents showing how to use AgentTool with MCP tools: ### mcp_in_agent_tool_remote (SSE mode) - Uses HTTP/SSE connection to remote MCP server - Zero-installation setup with `uvx` - Demonstrates server-side MCP deployment pattern ### mcp_in_agent_tool_stdio (stdio mode) - Uses subprocess connection with automatic server launch - Fully automatic setup with `uvx` subdirectory syntax - Demonstrates embedded MCP deployment pattern Both demos: - Use Gemini 2.5 Flash - Include example prompts for JSON Schema exploration - Have comprehensive READMEs with architecture diagrams - Follow ADK agent structure conventions ## Testing - ✅ All existing unit tests pass - ✅ Manual testing with both SSE and stdio modes - ✅ Verified cleanup happens in correct async context - ✅ No more cancel scope errors with MCP tools ## Related - Fixes #1112 - Related to #929 Co-authored-by: Wei Sun (Jack) <weisun@google.com> COPYBARA_INTEGRATE_REVIEW=#3411 from google:fix/agent-tool-mcp-cleanup 9ae753b PiperOrigin-RevId: 828651896
|
Thank you @Jacksunwei for your contribution! 🎉 Your changes have been successfully imported and merged via Copybara in commit fa5c546. Closing this PR as the changes are now in the main branch. |
3 participants
Summary
Fixes AgentTool cleanup to prevent MCP session errors by calling
runner.close()after sub-agent execution.Problem
When using AgentTool with MCP tools, the runner cleanup happened during garbage collection in a different async task context, causing:
Solution
await runner.close()immediately after sub-agent execution in AgentToolDemo Agents
Added two comprehensive demo agents showing how to use AgentTool with MCP tools:
mcp_in_agent_tool_remote (SSE mode)
uvxmcp_in_agent_tool_stdio (stdio mode)
uvxsubdirectory syntaxBoth demos:
Testing
Related