diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 9dfb29932f..b7a1251668 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -248,6 +248,8 @@ def __init__( self.is_connected: bool = False self._tools_loaded: bool = False self._prompts_loaded: bool = False + self._ping_supported: bool = True + self._reconnecting: bool = False def __str__(self) -> str: return f"MCPTool(name={self.name}, description={self.description})" @@ -651,6 +653,7 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: await self._safe_close_exit_stack() self.session = None self.is_connected = False + self._ping_supported = True self._exit_stack = AsyncExitStack() if not self.session: try: @@ -1039,10 +1042,30 @@ async def _ensure_connected(self) -> None: Raises: ToolExecutionException: If reconnection fails. """ + from mcp.shared.exceptions import McpError + + # Avoid re-entrant reconnection when connect() calls load_tools()/load_prompts() + if self._reconnecting: + return + + # If the server previously indicated it doesn't support ping, skip the check. + # The actual list_tools/list_prompts call will surface transport errors if any. + if not self._ping_supported: + return + try: await self.session.send_ping() # type: ignore[union-attr] + except McpError as mcp_exc: + # McpError means the server received the request and responded—the + # transport is alive. If the error is "Method not found" (-32601), + # the server doesn't implement the optional ping method per MCP spec. + if mcp_exc.error.code == -32601: # JSON-RPC "Method not found" + self._ping_supported = False + logger.debug("MCP server does not support ping (method not found). Connection is valid.") + return except Exception: logger.info("MCP connection invalid or closed. Reconnecting...") + self._reconnecting = True try: await self.connect(reset=True) except Exception as ex: @@ -1050,6 +1073,8 @@ async def _ensure_connected(self) -> None: "Failed to establish MCP connection.", inner_exception=ex, ) from ex + finally: + self._reconnecting = False async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: """Call a tool with the given arguments. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 01cf1717bd..c67286788f 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -3662,6 +3662,124 @@ async def test_ensure_connected_wraps_reconnect_failure() -> None: await tool._ensure_connected() +async def test_ensure_connected_treats_ping_method_not_found_as_connected() -> None: + """When the MCP server does not support the optional ping method, treat as connected.""" + tool = MCPTool(name="test_tool") + tool.session = Mock( + send_ping=AsyncMock( + side_effect=McpError(types.ErrorData(code=-32601, message="Unsupported method: ping")) + ) + ) + + with patch.object(tool, "connect", AsyncMock()) as mock_connect: + await tool._ensure_connected() + + mock_connect.assert_not_awaited() + assert tool._ping_supported is False + + +async def test_ensure_connected_skips_ping_when_previously_unsupported() -> None: + """After ping is marked unsupported, _ensure_connected skips the ping call entirely.""" + tool = MCPTool(name="test_tool") + tool._ping_supported = False + tool.session = Mock(send_ping=AsyncMock()) + + await tool._ensure_connected() + + tool.session.send_ping.assert_not_awaited() + + +async def test_ensure_connected_does_not_reconnect_on_any_mcp_error() -> None: + """Any McpError from ping means the transport is alive—do not reconnect.""" + tool = MCPTool(name="test_tool") + tool.session = Mock( + send_ping=AsyncMock(side_effect=McpError(types.ErrorData(code=-32603, message="Internal error"))) + ) + + with patch.object(tool, "connect", AsyncMock()) as mock_connect: + await tool._ensure_connected() + + mock_connect.assert_not_awaited() + # Non -32601 errors don't cache ping as unsupported + assert tool._ping_supported is True + + +async def test_ensure_connected_no_infinite_loop_when_ping_unsupported() -> None: + """Verify that load_tools does not loop infinitely when the server lacks ping support.""" + tool = MCPTool(name="test_tool", load_tools=True, load_prompts=False) + mock_session = Mock(spec=ClientSession) + mock_session.send_ping = AsyncMock( + side_effect=McpError(types.ErrorData(code=-32601, message="Method not found: ping")) + ) + mock_session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="example_tool", + description="An example tool", + inputSchema={"type": "object", "properties": {}}, + ) + ] + ) + ) + tool.session = mock_session + tool.is_connected = True + + await tool.load_tools() + + assert len(tool._functions) == 1 + assert tool._functions[0].name == "example_tool" + # After first ping failure, subsequent pages skip ping entirely + assert tool._ping_supported is False + + +async def test_ensure_connected_recursion_guard_prevents_loop() -> None: + """When _reconnecting is True, _ensure_connected returns immediately (prevents recursion).""" + tool = MCPTool(name="test_tool") + tool._reconnecting = True + tool.session = Mock(send_ping=AsyncMock(side_effect=RuntimeError("should not be called"))) + + # Should return immediately without calling ping or connect + await tool._ensure_connected() + + tool.session.send_ping.assert_not_awaited() + + +async def test_ensure_connected_reconnect_resets_ping_supported() -> None: + """After a successful reconnect, _ping_supported is reset to True.""" + tool = MCPTool(name="test_tool") + tool._ping_supported = False + tool.session = Mock(send_ping=AsyncMock(side_effect=RuntimeError("closed"))) + + async def mock_connect(*, reset: bool = False) -> None: + # connect(reset=True) resets _ping_supported + if reset: + tool._ping_supported = True + tool.session = Mock(send_ping=AsyncMock()) + + # First, enable ping so _ensure_connected actually checks + tool._ping_supported = True + + with patch.object(tool, "connect", AsyncMock(side_effect=mock_connect)): + await tool._ensure_connected() + + assert tool._ping_supported is True + + +async def test_ensure_connected_wraps_reconnect_failure_on_transport_error() -> None: + """Transport error from ping + reconnect failure should raise ToolExecutionException.""" + tool = MCPTool(name="test_tool") + tool.session = Mock(send_ping=AsyncMock(side_effect=ConnectionError("connection lost"))) + + with ( + patch.object(tool, "connect", AsyncMock(side_effect=RuntimeError("still closed"))), + pytest.raises(ToolExecutionException, match="Failed to establish MCP connection"), + ): + await tool._ensure_connected() + + assert tool._reconnecting is False + + async def test_mcp_tool_filters_framework_kwargs(): """Test that call_tool filters out framework-specific kwargs before calling MCP session.