diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index e42c96aea2..50f6a91c2e 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -318,9 +318,12 @@ class MCPTool: transport = await self._exit_stack.enter_async_context(self.get_mcp_client()) except Exception as ex: await self._exit_stack.aclose() - raise ToolException( - "Failed to connect to the MCP server. Please check your configuration.", inner_exception=ex - ) from ex + command = getattr(self, "command", None) + if command: + error_msg = f"Failed to start MCP server '{command}': {ex}" + else: + error_msg = f"Failed to connect to MCP server: {ex}" + raise ToolException(error_msg, inner_exception=ex) from ex try: session = await self._exit_stack.enter_async_context( ClientSession( @@ -335,9 +338,21 @@ class MCPTool: except Exception as ex: await self._exit_stack.aclose() raise ToolException( - message="Failed to create a session. Please check your configuration.", inner_exception=ex + message="Failed to create MCP session. Please check your configuration.", inner_exception=ex ) from ex - await session.initialize() + try: + await session.initialize() + except Exception as ex: + await self._exit_stack.aclose() + # Provide context about initialization failure + command = getattr(self, "command", None) + if command: + args_str = " ".join(getattr(self, "args", [])) + full_command = f"{command} {args_str}".strip() + error_msg = f"MCP server '{full_command}' failed to initialize: {ex}" + else: + error_msg = f"MCP server failed to initialize: {ex}" + raise ToolException(error_msg, inner_exception=ex) from ex self.session = session elif self.session._request_id == 0: # type: ignore[reportPrivateUsage] # If the session is not initialized, we need to reinitialize it diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index dc5c7935e2..e6dd1fd8a7 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -3,7 +3,7 @@ import os from contextlib import _AsyncGeneratorContextManager # type: ignore from typing import Any -from unittest.mock import AsyncMock, Mock +from unittest.mock import AsyncMock, Mock, patch import pytest from mcp import types @@ -33,7 +33,7 @@ from agent_framework._mcp import ( _mcp_type_to_ai_content, _normalize_mcp_name, ) -from agent_framework.exceptions import ToolExecutionException +from agent_framework.exceptions import ToolException, ToolExecutionException # Integration test skip condition skip_if_mcp_integration_tests_disabled = pytest.mark.skipif( @@ -776,3 +776,288 @@ async def test_streamable_http_integration(): result = await func.invoke(query="What is Agent Framework?") assert result[0].text is not None + + +async def test_mcp_tool_message_handler_notification(): + """Test that message_handler correctly processes tools/list_changed and prompts/list_changed + notifications.""" + tool = MCPStdioTool(name="test_tool", command="python") + + # Mock the load_tools and load_prompts methods + tool.load_tools = AsyncMock() + tool.load_prompts = AsyncMock() + + # Test tools list changed notification + tools_notification = Mock(spec=types.ServerNotification) + tools_notification.root = Mock() + tools_notification.root.method = "notifications/tools/list_changed" + + result = await tool.message_handler(tools_notification) + assert result is None + tool.load_tools.assert_called_once() + + # Reset mock + tool.load_tools.reset_mock() + + # Test prompts list changed notification + prompts_notification = Mock(spec=types.ServerNotification) + prompts_notification.root = Mock() + prompts_notification.root.method = "notifications/prompts/list_changed" + + result = await tool.message_handler(prompts_notification) + assert result is None + tool.load_prompts.assert_called_once() + + # Test unhandled notification + unknown_notification = Mock(spec=types.ServerNotification) + unknown_notification.root = Mock() + unknown_notification.root.method = "notifications/unknown" + + result = await tool.message_handler(unknown_notification) + assert result is None + + +async def test_mcp_tool_message_handler_error(): + """Test that message_handler gracefully handles exceptions by logging and returning None.""" + tool = MCPStdioTool(name="test_tool", command="python") + + # Test with exception message + test_exception = RuntimeError("Test error message") + + # The message handler should log the error and return None + result = await tool.message_handler(test_exception) + assert result is None + + +async def test_mcp_tool_sampling_callback_no_client(): + """Test sampling callback error path when no chat client is available.""" + tool = MCPStdioTool(name="test_tool", command="python") + + # Create minimal params mock + params = Mock() + params.messages = [] + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.ErrorData) + assert result.code == types.INTERNAL_ERROR + assert "No chat client available" in result.message + + +async def test_mcp_tool_sampling_callback_chat_client_exception(): + """Test sampling callback when chat client raises exception.""" + tool = MCPStdioTool(name="test_tool", command="python") + + # Mock chat client that raises exception + mock_chat_client = AsyncMock() + mock_chat_client.get_response.side_effect = RuntimeError("Chat client error") + + tool.chat_client = mock_chat_client + + # Create mock params + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.ErrorData) + assert result.code == types.INTERNAL_ERROR + assert "Failed to get chat message content: Chat client error" in result.message + + +async def test_mcp_tool_sampling_callback_no_valid_content(): + """Test sampling callback when response has no valid content types.""" + from agent_framework import ChatMessage, DataContent, Role + + tool = MCPStdioTool(name="test_tool", command="python") + + # Mock chat client with response containing only invalid content types + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [ + ChatMessage( + role=Role.ASSISTANT, + contents=[DataContent(uri="data:application/json;base64,e30K", media_type="application/json")], + ) + ] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.chat_client = mock_chat_client + + # Create mock params + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.ErrorData) + assert result.code == types.INTERNAL_ERROR + assert "Failed to get right content types from the response." in result.message + + +# Test error handling in connect() method + + +async def test_connect_session_creation_failure(): + """Test connect() raises ToolException when ClientSession creation fails.""" + tool = MCPStdioTool(name="test", command="test-command") + + # Mock successful transport creation + mock_transport = (Mock(), Mock()) # (read_stream, write_stream) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) + + # Mock ClientSession to raise an exception + with patch("agent_framework._mcp.ClientSession") as mock_session_class: + mock_session_class.side_effect = RuntimeError("Session creation failed") + + with pytest.raises(ToolException) as exc_info: + await tool.connect() + + assert "Failed to create MCP session" in str(exc_info.value) + assert "Session creation failed" in str(exc_info.value.__cause__) + + +async def test_connect_initialization_failure_http_no_command(): + """Test connect() when session.initialize() fails for HTTP tool (no command attribute).""" + tool = MCPStreamableHTTPTool(name="test", url="http://example.com") + + # Mock successful transport creation + mock_transport = (Mock(), Mock()) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) + + # Mock successful session creation but failed initialization + mock_session = Mock() + mock_session.initialize = AsyncMock(side_effect=ConnectionError("Server not ready")) + + with patch("agent_framework._mcp.ClientSession") as mock_session_class: + mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None) + + with pytest.raises(ToolException) as exc_info: + await tool.connect() + + # Should use generic error message since HTTP tool doesn't have command + assert "MCP server failed to initialize" in str(exc_info.value) + assert "Server not ready" in str(exc_info.value) + + +async def test_connect_cleanup_on_transport_failure(): + """Test that _exit_stack.aclose() is called when transport creation fails.""" + tool = MCPStdioTool(name="test", command="test-command") + + # Mock _exit_stack.aclose to verify it's called + tool._exit_stack.aclose = AsyncMock() + + # Mock get_mcp_client to raise an exception + tool.get_mcp_client = Mock(side_effect=RuntimeError("Transport failed")) + + with pytest.raises(ToolException): + await tool.connect() + + # Verify cleanup was called + tool._exit_stack.aclose.assert_called_once() + + +async def test_connect_cleanup_on_initialization_failure(): + """Test that _exit_stack.aclose() is called when initialization fails.""" + tool = MCPStdioTool(name="test", command="test-command") + + # Mock _exit_stack.aclose to verify it's called + tool._exit_stack.aclose = AsyncMock() + + # Mock successful transport creation + mock_transport = (Mock(), Mock()) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) + + # Mock successful session creation but failed initialization + mock_session = Mock() + mock_session.initialize = AsyncMock(side_effect=RuntimeError("Init failed")) + + with patch("agent_framework._mcp.ClientSession") as mock_session_class: + mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None) + + with pytest.raises(ToolException): + await tool.connect() + + # Verify cleanup was called + tool._exit_stack.aclose.assert_called_once() + + +def test_mcp_stdio_tool_get_mcp_client_with_env_and_kwargs(): + """Test MCPStdioTool.get_mcp_client() with environment variables and client kwargs.""" + env_vars = {"PATH": "/usr/bin", "DEBUG": "1"} + tool = MCPStdioTool(name="test", command="test-command", env=env_vars, custom_param="value1", another_param=42) + + with patch("agent_framework._mcp.stdio_client"), patch("agent_framework._mcp.StdioServerParameters") as mock_params: + tool.get_mcp_client() + + # Verify all parameters including custom kwargs were passed + mock_params.assert_called_once_with( + command="test-command", args=[], env=env_vars, custom_param="value1", another_param=42 + ) + + +def test_mcp_streamable_http_tool_get_mcp_client_all_params(): + """Test MCPStreamableHTTPTool.get_mcp_client() with all parameters.""" + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com", + headers={"Auth": "token"}, + timeout=30.0, + sse_read_timeout=10.0, + terminate_on_close=True, + custom_param="test", + ) + + with patch("agent_framework._mcp.streamablehttp_client") as mock_http_client: + tool.get_mcp_client() + + # Verify all parameters were passed + mock_http_client.assert_called_once_with( + url="http://example.com", + headers={"Auth": "token"}, + timeout=30.0, + sse_read_timeout=10.0, + terminate_on_close=True, + custom_param="test", + ) + + +def test_mcp_websocket_tool_get_mcp_client_with_kwargs(): + """Test MCPWebsocketTool.get_mcp_client() with client kwargs.""" + tool = MCPWebsocketTool( + name="test", url="wss://example.com", max_size=1024, ping_interval=30, compression="deflate" + ) + + with patch("agent_framework._mcp.websocket_client") as mock_ws_client: + tool.get_mcp_client() + + # Verify all kwargs were passed + mock_ws_client.assert_called_once_with( + url="wss://example.com", max_size=1024, ping_interval=30, compression="deflate" + )