Python: MCP Error Handling Fix + Added Unit Tests

python/packages/core/agent_framework/_mcp.py

@@ -318,9 +318,12 @@ async def connect(self) -> None:
                 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 @@ async def connect(self) -> None:
             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

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 @@
     _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"
+        )