mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
0ccc43cf49
commit
9761fc6bbb
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user