fix: Fix McpToolset hanging indefinitely

This fixes McpToolset from hanging indefinitely during a list_tools call (https://github.com/google/adk-python/issues/3084) by adding a timeout.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 829499276
This commit is contained in:
Kathy Wu
2025-11-07 10:50:38 -08:00
committed by Copybara-Service
parent 0ccc43cf49
commit 9761fc6bbb
3 changed files with 34 additions and 1 deletions
+12 -1
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import asyncio
import logging
import sys
from typing import Callable
@@ -177,7 +178,17 @@ class McpToolset(BaseToolset):
session = await self._mcp_session_manager.create_session(headers=headers)
# Fetch available tools from the MCP server
tools_response: ListToolsResult = await session.list_tools()
timeout_in_seconds = (
self._connection_params.timeout
if hasattr(self._connection_params, "timeout")
else None
)
try:
tools_response: ListToolsResult = await asyncio.wait_for(
session.list_tools(), timeout=timeout_in_seconds
)
except Exception as e:
raise ConnectionError("Failed to get tools from MCP server.") from e
# Apply filtering based on context and tool_filter
tools = []
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
from io import StringIO
import sys
import unittest
@@ -304,6 +305,26 @@ class TestMCPToolset:
assert "Warning: Error during McpToolset cleanup" in error_output
assert "Cleanup error" in error_output
@pytest.mark.asyncio
async def test_get_tools_with_timeout(self):
"""Test get_tools with timeout."""
stdio_params = StdioConnectionParams(
server_params=self.mock_stdio_params, timeout=0.01
)
toolset = MCPToolset(connection_params=stdio_params)
toolset._mcp_session_manager = self.mock_session_manager
async def long_running_list_tools():
await asyncio.sleep(0.1)
return MockListToolsResult([])
self.mock_session.list_tools = long_running_list_tools
with pytest.raises(
ConnectionError, match="Failed to get tools from MCP server."
):
await toolset.get_tools()
@pytest.mark.asyncio
async def test_get_tools_retry_decorator(self):
"""Test that get_tools has retry decorator applied."""
@@ -45,6 +45,7 @@ async def test_mcp_toolset_with_prefix():
"""Test that McpToolset correctly applies the tool_name_prefix."""
# Mock the connection parameters
mock_connection_params = MagicMock()
mock_connection_params.timeout = None
# Mock the MCPSessionManager and its create_session method
mock_session_manager = MagicMock()