diff --git a/omlx/mcp/client.py b/omlx/mcp/client.py index bfae8c5d..5603609c 100644 --- a/omlx/mcp/client.py +++ b/omlx/mcp/client.py @@ -198,7 +198,10 @@ async def _connect_streamable_http(self): self._session = ClientSession(self._read, self._write) await self._session.__aenter__() except Exception: - if hasattr(self, '_streamable_http_client') and self._streamable_http_client: + if ( + hasattr(self, "_streamable_http_client") + and self._streamable_http_client + ): await self._streamable_http_client.__aexit__(None, None, None) self._streamable_http_client = None await self._http_client.__aexit__(None, None, None) @@ -232,7 +235,9 @@ async def _discover_tools(self): server_name=self.name, name=tool.name, description=tool.description or "", - input_schema=tool.inputSchema if hasattr(tool, 'inputSchema') else {}, + input_schema=tool.inputSchema + if hasattr(tool, "inputSchema") + else {}, ) self._tools.append(mcp_tool) logger.debug(f"Discovered tool: {mcp_tool.full_name}") @@ -248,11 +253,11 @@ async def _cleanup_resources(self): await self._session.__aexit__(None, None, None) self._session = None - if hasattr(self, '_stdio_client') and self._stdio_client: + if hasattr(self, "_stdio_client") and self._stdio_client: await self._stdio_client.__aexit__(None, None, None) self._stdio_client = None - if hasattr(self, '_sse_client') and self._sse_client: + if hasattr(self, "_sse_client") and self._sse_client: await self._sse_client.__aexit__(None, None, None) self._sse_client = None @@ -331,7 +336,7 @@ async def call_tool( return MCPToolResult( tool_name=tool_name, content=content, - is_error=result.isError if hasattr(result, 'isError') else False, + is_error=result.isError if hasattr(result, "isError") else False, ) except asyncio.TimeoutError: @@ -351,15 +356,18 @@ async def call_tool( def _extract_content(self, result) -> Any: """Extract content from MCP tool result.""" - if not hasattr(result, 'content') or not result.content: + if not hasattr(result, "content") or not result.content: + # Fall back to structuredContent if available (used by web search tools) + if hasattr(result, "structuredContent") and result.structuredContent: + return result.structuredContent return None # Handle list of content items contents = [] for item in result.content: - if hasattr(item, 'text'): + if hasattr(item, "text"): contents.append(item.text) - elif hasattr(item, 'data'): + elif hasattr(item, "data"): contents.append(item.data) else: contents.append(str(item)) diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 2979b091..917327c3 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -235,9 +235,13 @@ async def test_connect_disabled_server(self): async def test_connect_stdio_success(self, stdio_client: MCPClient): """Test successful stdio connection.""" # Mock the internal methods instead of trying to patch imports - with patch.object(stdio_client, "_connect_stdio", new_callable=AsyncMock) as mock_connect, \ - patch.object(stdio_client, "_initialize_session", new_callable=AsyncMock), \ - patch.object(stdio_client, "_discover_tools", new_callable=AsyncMock): + with ( + patch.object( + stdio_client, "_connect_stdio", new_callable=AsyncMock + ) as mock_connect, + patch.object(stdio_client, "_initialize_session", new_callable=AsyncMock), + patch.object(stdio_client, "_discover_tools", new_callable=AsyncMock), + ): mock_connect.return_value = None # Set up session to pass the check stdio_client._session = MagicMock() @@ -251,9 +255,11 @@ async def test_connect_stdio_success(self, stdio_client: MCPClient): @pytest.mark.asyncio async def test_connect_sse_success(self, sse_client: MCPClient): """Test successful SSE connection.""" - with patch.object(sse_client, "_connect_sse", new_callable=AsyncMock), \ - patch.object(sse_client, "_initialize_session", new_callable=AsyncMock), \ - patch.object(sse_client, "_discover_tools", new_callable=AsyncMock): + with ( + patch.object(sse_client, "_connect_sse", new_callable=AsyncMock), + patch.object(sse_client, "_initialize_session", new_callable=AsyncMock), + patch.object(sse_client, "_discover_tools", new_callable=AsyncMock), + ): sse_client._session = MagicMock() result = await sse_client.connect() @@ -433,9 +439,14 @@ async def test_connect_failure_cleans_up_resources(self): mock_stdio = AsyncMock() with ( - patch.object(client, "_connect_stdio", new_callable=AsyncMock) as mock_connect, - patch.object(client, "_initialize_session", new_callable=AsyncMock) as mock_init, + patch.object( + client, "_connect_stdio", new_callable=AsyncMock + ) as mock_connect, + patch.object( + client, "_initialize_session", new_callable=AsyncMock + ) as mock_init, ): + async def setup_partial_resources(): client._stdio_client = mock_stdio client._session = AsyncMock() @@ -465,8 +476,11 @@ async def test_streamable_http_partial_connect_cleanup(self): with ( patch("omlx.mcp.client.MCPClient._connect_streamable_http") as mock_connect, - patch.object(client, "_initialize_session", new_callable=AsyncMock) as mock_init, + patch.object( + client, "_initialize_session", new_callable=AsyncMock + ) as mock_init, ): + async def setup_and_fail(): client._http_client = mock_http_client client._streamable_http_client = mock_streamable @@ -543,6 +557,7 @@ async def test_call_tool_success(self, connected_client: MCPClient): @pytest.mark.asyncio async def test_call_tool_timeout(self, connected_client: MCPClient): """Test tool call timeout.""" + async def slow_call(*args, **kwargs): await asyncio.sleep(10) @@ -592,6 +607,23 @@ async def test_call_tool_multiple_content_items(self, connected_client: MCPClien assert result.is_error is False assert result.content == ["Line 1", "Line 2"] + @pytest.mark.asyncio + async def test_call_tool_structured_content(self, connected_client: MCPClient): + """Test tool call with structuredContent (used by web search tools).""" + mock_result = MagicMock() + mock_result.content = [] + mock_result.structuredContent = {"results": ["Result 1", "Result 2"]} + mock_result.isError = False + connected_client._session.call_tool.return_value = mock_result + + result = await connected_client.call_tool("web_search", {"query": "test"}) + + assert result.is_error is False + assert result.content == {"results": ["Result 1", "Result 2"]} + connected_client._session.call_tool.assert_called_with( + "web_search", {"query": "test"} + ) + @pytest.mark.asyncio async def test_call_tool_data_content(self, connected_client: MCPClient): """Test tool call with data content.""" diff --git a/tests/test_structured_content.py b/tests/test_structured_content.py new file mode 100644 index 00000000..1454d449 --- /dev/null +++ b/tests/test_structured_content.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Tests for structuredContent support in MCP tool results. + +This addresses issue #469: MCP web search returns no output. +Web search MCP servers return results in structuredContent field. +""" + +from unittest.mock import MagicMock + + +def test_extract_content_with_structured_content(): + """Test that _extract_content falls back to structuredContent when content is empty.""" + from omlx.mcp.client import MCPClient + from omlx.mcp.types import MCPServerConfig, MCPTransport + + config = MCPServerConfig( + name="test", + transport=MCPTransport.STDIO, + command="python", + ) + client = MCPClient(config) + + mock_result = MagicMock() + mock_result.content = [] + mock_result.structuredContent = {"results": ["Result 1", "Result 2"]} + + result = client._extract_content(mock_result) + + assert result == {"results": ["Result 1", "Result 2"]} + + +def test_extract_content_with_text_content(): + """Test that _extract_content still works with regular text content.""" + from omlx.mcp.client import MCPClient + from omlx.mcp.types import MCPServerConfig, MCPTransport + + config = MCPServerConfig( + name="test", + transport=MCPTransport.STDIO, + command="python", + ) + client = MCPClient(config) + + mock_result = MagicMock() + mock_result.content = [MagicMock(text="Tool output")] + + result = client._extract_content(mock_result) + + assert result == "Tool output" + + +def test_extract_content_empty_no_structured(): + """Test that _extract_content returns None when both content and structuredContent are empty.""" + from omlx.mcp.client import MCPClient + from omlx.mcp.types import MCPServerConfig, MCPTransport + + config = MCPServerConfig( + name="test", + transport=MCPTransport.STDIO, + command="python", + ) + client = MCPClient(config) + + mock_result = MagicMock() + mock_result.content = [] + mock_result.structuredContent = None + + result = client._extract_content(mock_result) + + assert result is None + + +def test_extract_content_priority_text_over_structured(): + """Test that text content takes priority over structuredContent.""" + from omlx.mcp.client import MCPClient + from omlx.mcp.types import MCPServerConfig, MCPTransport + + config = MCPServerConfig( + name="test", + transport=MCPTransport.STDIO, + command="python", + ) + client = MCPClient(config) + + mock_result = MagicMock() + mock_result.content = [MagicMock(text="Text content")] + mock_result.structuredContent = {"results": ["Structured"]} + + result = client._extract_content(mock_result) + + assert result == "Text content"