fix(tools): Add proper cleanup for AgentTool to prevent MCP session errors

Merge https://github.com/google/adk-python/pull/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=https://github.com/google/adk-python/pull/3411 from google:fix/agent-tool-mcp-cleanup 9ae753b5a4
PiperOrigin-RevId: 828651896
This commit is contained in:
Wei Sun (Jack)
2025-11-05 15:23:39 -08:00
committed by Copybara-Service
parent ee8106be77
commit fa5c546a55
8 changed files with 333 additions and 0 deletions
@@ -0,0 +1,74 @@
# AgentTool with MCP Demo (SSE Mode)
This demo shows how `AgentTool` works with MCP (Model Context Protocol) toolsets using **SSE mode**.
## SSE vs Stdio Mode
This demo uses **SSE (Server-Sent Events) mode** where the MCP server runs as a separate HTTP server:
- **Remote connection** - Connects to server via HTTP
- **Separate process** - Server must be started manually
- **Network communication** - Uses HTTP/SSE for messaging
For the **stdio (subprocess) version**, see [mcp_in_agent_tool_stdio](../mcp_in_agent_tool_stdio/).
## Setup
**Start the MCP simple-tool server in SSE mode** (in a separate terminal):
```bash
# Run the server using uvx (no installation needed)
# Port 3000 avoids conflict with adk web (which uses 8000)
uvx --from 'git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool' \
mcp-simple-tool --transport sse --port 3000
```
The server should be accessible at `http://localhost:3000/sse`.
## Running the Demo
```bash
adk web contributing/samples
```
Then select **mcp_in_agent_tool_remote** from the list and interact with the agent.
## Try These Prompts
This demo uses **Gemini 2.5 Flash** as the model. Try these prompts:
1. **Check available tools:**
```
What tools do you have access to?
```
2. **Fetch and summarize JSON Schema specification:**
```
Use the mcp_helper to fetch https://json-schema.org/specification and summarize the key features of JSON Schema
```
## Architecture
```
main_agent (root_agent)
└── AgentTool wrapping:
└── mcp_helper (sub_agent)
└── McpToolset (SSE connection)
└── http://localhost:3000/sse
└── MCP simple-tool server
└── Website Fetcher Tool
```
## Related
- **Issue:** [#1112 - Using agent as tool outside of adk web doesn't exit cleanly](https://github.com/google/adk-python/issues/1112)
- **Related Issue:** [#929 - LiteLLM giving error with OpenAI models and Grafana's MCP server](https://github.com/google/adk-python/issues/929)
- **Stdio Version:** [mcp_in_agent_tool_stdio](../mcp_in_agent_tool_stdio/) - Uses local subprocess connection
@@ -0,0 +1,15 @@
# 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.
from . import agent
@@ -0,0 +1,70 @@
# 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.
from google.adk.agents import Agent
from google.adk.tools import AgentTool
from google.adk.tools.mcp_tool import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
# Create MCP toolset
# This uses the simple-tool MCP server via SSE
# You need to start the MCP server separately (see README.md)
mcp_toolset = McpToolset(
connection_params=SseConnectionParams(
url="http://localhost:3000/sse",
timeout=10.0,
sse_read_timeout=300.0,
)
)
# Create sub-agent with MCP tools
# This agent has direct access to MCP tools
sub_agent = Agent(
name="mcp_helper",
model="gemini-2.5-flash",
description=(
"A helpful assistant with access to MCP tools for fetching websites."
),
instruction="""You are a helpful assistant with access to MCP tools.
When the user asks for help:
1. Explain what tools you have available (website fetching)
2. Use the appropriate tool if needed
3. Provide clear and helpful responses
You have access to a website fetcher tool via MCP. Use it to fetch and return website content.""",
tools=[mcp_toolset],
)
# Wrap sub-agent as an AgentTool
# This allows the main agent to delegate tasks to the sub-agent
# The sub-agent has access to MCP tools for fetching websites
mcp_agent_tool = AgentTool(agent=sub_agent)
# Create main agent
# This agent can delegate to the sub-agent via AgentTool
root_agent = Agent(
name="main_agent",
model="gemini-2.5-flash",
description="Main agent that can delegate to a sub-agent with MCP tools.",
instruction="""You are a helpful assistant. You have access to a sub-agent (mcp_helper)
that has MCP tools for fetching websites.
When the user asks for help:
- If they need to fetch a website, call the mcp_helper tool
- Otherwise, respond directly
Always be helpful and explain what you're doing.""",
tools=[mcp_agent_tool],
)
@@ -0,0 +1,74 @@
# AgentTool with MCP Demo (Stdio Mode)
This demo shows how `AgentTool` works with MCP (Model Context Protocol) toolsets using **stdio mode**.
## Stdio vs SSE Mode
This demo uses **stdio mode** where the MCP server runs as a subprocess:
- **Simpler setup** - No need to start a separate server
- **Auto-launched** - Server starts automatically when agent runs
- **Local process** - Uses stdin/stdout for communication
For the **SSE (remote server) version**, see [mcp_in_agent_tool_remote](../mcp_in_agent_tool_remote/).
## Setup
**No installation required!** The MCP server will be launched automatically using `uvx` when you run the agent.
The demo uses `uvx` to fetch and run the MCP simple-tool server directly from the GitHub repository's subdirectory:
```bash
uvx --from 'git+https://github.com/modelcontextprotocol/python-sdk.git#subdirectory=examples/servers/simple-tool' \
mcp-simple-tool
```
This happens automatically via the stdio connection when the agent starts.
## Running the Demo
```bash
adk web contributing/samples
```
Then select **mcp_in_agent_tool_stdio** from the list and interact with the agent.
## Try These Prompts
This demo uses **Gemini 2.5 Flash** as the model. Try these prompts:
1. **Check available tools:**
```
What tools do you have access to?
```
2. **Fetch and summarize JSON Schema specification:**
```
Use the mcp_helper to fetch https://json-schema.org/specification and summarize the key features of JSON Schema
```
## Architecture
```
main_agent (root_agent)
└── AgentTool wrapping:
└── mcp_helper (sub_agent)
└── McpToolset (stdio connection)
└── MCP Server (subprocess via uvx)
└── uvx --from git+...#subdirectory=... mcp-simple-tool
└── Website Fetcher Tool
```
## Related
- **Issue:** [#1112 - Using agent as tool outside of adk web doesn't exit cleanly](https://github.com/google/adk-python/issues/1112)
- **Related Issue:** [#929 - LiteLLM giving error with OpenAI models and Grafana's MCP server](https://github.com/google/adk-python/issues/929)
- **SSE Version:** [mcp_in_agent_tool_remote](../mcp_in_agent_tool_remote/) - Uses remote server connection
@@ -0,0 +1,15 @@
# 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.
from . import agent
@@ -0,0 +1,77 @@
# 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.
from google.adk.agents import Agent
from google.adk.tools import AgentTool
from google.adk.tools.mcp_tool import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from mcp import StdioServerParameters
# Create MCP toolset
# This uses the simple-tool MCP server via stdio
# The server will be launched automatically using uvx from the subdirectory
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,
)
)
# Create sub-agent with MCP tools
# This agent has direct access to MCP tools
sub_agent = Agent(
name="mcp_helper",
model="gemini-2.5-flash",
description=(
"A helpful assistant with access to MCP tools for fetching websites."
),
instruction="""You are a helpful assistant with access to MCP tools.
When the user asks for help:
1. Explain what tools you have available (website fetching)
2. Use the appropriate tool if needed
3. Provide clear and helpful responses
You have access to a website fetcher tool via MCP. Use it to fetch and return website content.""",
tools=[mcp_toolset],
)
# Wrap sub-agent as an AgentTool
# This allows the main agent to delegate tasks to the sub-agent
# The sub-agent has access to MCP tools for fetching websites
mcp_agent_tool = AgentTool(agent=sub_agent)
# Create main agent
# This agent can delegate to the sub-agent via AgentTool
root_agent = Agent(
name="main_agent",
model="gemini-2.5-flash",
description="Main agent that can delegate to a sub-agent with MCP tools.",
instruction="""You are a helpful assistant. You have access to a sub-agent (mcp_helper)
that has MCP tools for fetching websites.
When the user asks for help:
- If they need to fetch a website, call the mcp_helper tool
- Otherwise, respond directly
Always be helpful and explain what you're doing.""",
tools=[mcp_agent_tool],
)
+4
View File
@@ -164,6 +164,10 @@ class AgentTool(BaseTool):
if event.content:
last_content = event.content
# Clean up runner resources (especially MCP sessions)
# to avoid "Attempted to exit cancel scope in a different task" errors
await runner.close()
if not last_content:
return ''
merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
+4
View File
@@ -131,6 +131,10 @@ async def test_agent_tool_inherits_parent_app_name(monkeypatch):
)
return _empty_async_generator()
async def close(self):
"""Mock close method."""
pass
monkeypatch.setattr('google.adk.runners.Runner', StubRunner)
tool_agent = Agent(