mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: MCP Error Handling Fix + Added Unit Tests (#1621)
* mcp error fix * test docstring fixes
This commit is contained in:
Giles Odigwe
committed by
GitHub
Unverified
GitHub
parent
c408a2d8c3
commit
1f19a6da5c
@@ -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
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user