hud-python 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

hud/mcp/tests/test_claude.py

@@ -0,0 +1,294 @@
+"""Tests for Claude MCP Agent implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from anthropic import BadRequestError
+from mcp import types
+from mcp.types import CallToolRequestParams as MCPToolCall
+
+from hud.mcp.claude import (
+    ClaudeMCPAgent,
+    base64_to_content_block,
+    text_to_content_block,
+    tool_use_content_block,
+)
+
+if TYPE_CHECKING:
+    from anthropic.types.beta import BetaImageBlockParam, BetaMessageParam, BetaTextBlockParam
+
+
+class TestClaudeHelperFunctions:
+    """Test helper functions for Claude message formatting."""
+
+    def test_base64_to_content_block(self):
+        """Test base64 image conversion."""
+        base64_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="  # noqa: E501
+        result = base64_to_content_block(base64_data)
+
+        assert result["type"] == "image"
+        assert result["source"]["type"] == "base64"
+        assert result["source"]["media_type"] == "image/png"
+        assert result["source"]["data"] == base64_data
+
+    def test_text_to_content_block(self):
+        """Test text conversion."""
+        text = "Hello, world!"
+        result = text_to_content_block(text)
+
+        assert result["type"] == "text"
+        assert result["text"] == text
+
+    def test_tool_use_content_block(self):
+        """Test tool result content block creation."""
+        tool_use_id = "tool_123"
+        content: list[BetaTextBlockParam | BetaImageBlockParam] = [
+            text_to_content_block("Result text")
+        ]
+
+        result = tool_use_content_block(tool_use_id, content)
+
+        assert result["type"] == "tool_result"
+        assert result["tool_use_id"] == tool_use_id
+        assert result["content"] == content  # type: ignore
+
+
+class TestClaudeMCPAgent:
+    """Test ClaudeMCPAgent class."""
+
+    @pytest.fixture
+    def mock_mcp_client(self):
+        """Create a mock MCP client."""
+        mcp_client = MagicMock()
+        mcp_client.get_all_active_sessions = MagicMock(return_value={})
+        mcp_client.get_tool_map = MagicMock(return_value={})
+        return mcp_client
+
+    @pytest.fixture
+    def mock_anthropic(self):
+        """Create a mock Anthropic client."""
+        with patch("hud.mcp.claude.AsyncAnthropic") as mock:
+            client = AsyncMock()
+            # Add beta attribute with messages
+            client.beta = AsyncMock()
+            client.beta.messages = AsyncMock()
+            mock.return_value = client
+            yield client
+
+    @pytest.mark.asyncio
+    async def test_init(self, mock_mcp_client, mock_anthropic):
+        """Test agent initialization."""
+        # Test with provided model_client
+        mock_model_client = MagicMock()
+        agent = ClaudeMCPAgent(
+            mcp_client=mock_mcp_client,
+            model_client=mock_model_client,
+            model="claude-3-opus-20240229",
+            max_tokens=1000,
+        )
+
+        assert agent.model_name == "claude-3-opus-20240229"
+        assert agent.max_tokens == 1000
+        assert agent.anthropic_client == mock_model_client
+
+    @pytest.mark.asyncio
+    async def test_init_without_model_client(self, mock_mcp_client):
+        """Test agent initialization without model client."""
+        with patch("hud.mcp.claude.settings.anthropic_api_key", "test_key"):
+            agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model="claude-3-opus-20240229")
+
+            assert agent.model_name == "claude-3-opus-20240229"
+            assert agent.anthropic_client is not None
+
+    @pytest.mark.asyncio
+    async def test_create_initial_messages(self, mock_mcp_client):
+        """Test creating initial messages."""
+        mock_model_client = MagicMock()
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
+
+        # Test with text only
+        messages = await agent.create_initial_messages("Hello, Claude!")
+        assert len(messages) == 1
+        assert messages[0]["role"] == "user"
+        content = list(messages[0]["content"])
+        assert content[0]["type"] == "text"  # type: ignore
+        assert content[0]["text"] == "Hello, Claude!"  # type: ignore
+
+        # Test with screenshot
+        messages = await agent.create_initial_messages("Look at this", screenshot="base64data")
+        assert len(messages) == 1
+        assert messages[0]["role"] == "user"
+        content = list(messages[0]["content"])
+        assert len(content) == 2
+        # Claude puts text first, then image
+        assert content[0]["type"] == "text"  # type: ignore
+        assert content[0]["text"] == "Look at this"  # type: ignore
+        assert content[1]["type"] == "image"  # type: ignore
+
+    @pytest.mark.asyncio
+    async def test_format_tool_results_method(self, mock_mcp_client):
+        """Test the agent's format_tool_results method."""
+        mock_model_client = MagicMock()
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_model_client)
+
+        tool_calls = [
+            MCPToolCall(name="test_tool", arguments={}, tool_use_id="id1"),  # type: ignore
+        ]
+
+        tool_results = [
+            types.CallToolResult(
+                content=[types.TextContent(type="text", text="Success")], isError=False
+            ),
+        ]
+
+        messages = await agent.format_tool_results(tool_calls, tool_results)
+
+        # format_tool_results returns a single user message with tool result content
+        assert len(messages) == 1
+        assert messages[0]["role"] == "user"
+        # The content is wrapped in a tool result block
+        content = list(messages[0]["content"])
+        assert len(content) == 1
+        assert content[0]["type"] == "tool_result"  # type: ignore
+        assert content[0]["tool_use_id"] == "id1"  # type: ignore
+        # The actual content is nested inside
+        inner_content = list(content[0]["content"])  # type: ignore
+        assert inner_content[0]["type"] == "text"  # type: ignore
+        assert inner_content[0]["text"] == "Success"  # type: ignore
+
+    @pytest.mark.asyncio
+    async def test_get_model_response(self, mock_mcp_client, mock_anthropic):
+        """Test getting model response from Claude API."""
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+
+        # Mock the API response
+        mock_response = MagicMock()
+
+        # Create text block
+        text_block = MagicMock()
+        text_block.type = "text"
+        text_block.text = "Hello!"
+
+        # Create tool use block
+        tool_block = MagicMock()
+        tool_block.type = "tool_use"
+        tool_block.id = "tool_123"
+        tool_block.name = "test_tool"
+        tool_block.input = {"param": "value"}
+
+        mock_response.content = [text_block, tool_block]
+        mock_response.usage = MagicMock(input_tokens=10, output_tokens=20)
+        mock_anthropic.beta.messages.create = AsyncMock(return_value=mock_response)
+
+        messages = [
+            cast("BetaMessageParam", {"role": "user", "content": [{"type": "text", "text": "Hi"}]})
+        ]
+        response = await agent.get_model_response(messages)
+
+        assert response.content == "Hello!"
+        assert len(response.tool_calls) == 1
+        assert response.tool_calls[0].name == "test_tool"
+        assert response.tool_calls[0].arguments == {"param": "value"}
+        # The test was checking for Claude-specific attributes that aren't part of ModelResponse
+        # These would need to be accessed from the original Claude response if needed
+
+        # Verify API was called correctly
+        mock_anthropic.beta.messages.create.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_get_model_response_text_only(self, mock_mcp_client, mock_anthropic):
+        """Test getting text-only response."""
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+
+        mock_response = MagicMock()
+        # Create text block
+        text_block = MagicMock()
+        text_block.type = "text"
+        text_block.text = "Just text"
+        mock_response.content = [text_block]
+        mock_response.usage = MagicMock(input_tokens=5, output_tokens=10)
+        mock_anthropic.beta.messages.create = AsyncMock(return_value=mock_response)
+
+        messages = [
+            cast("BetaMessageParam", {"role": "user", "content": [{"type": "text", "text": "Hi"}]})
+        ]
+        response = await agent.get_model_response(messages)
+
+        assert response.content == "Just text"
+        assert response.tool_calls == []
+
+    @pytest.mark.asyncio
+    async def test_get_model_response_error(self, mock_mcp_client, mock_anthropic):
+        """Test handling API errors."""
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+
+        # Mock API error
+        mock_anthropic.beta.messages.create = AsyncMock(
+            side_effect=BadRequestError(
+                message="Invalid request",
+                response=MagicMock(status_code=400),
+                body={"error": {"message": "Invalid request"}},
+            )
+        )
+
+        messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}]
+
+        with pytest.raises(BadRequestError):
+            await agent.get_model_response(messages)  # type: ignore
+
+    @pytest.mark.asyncio
+    async def test_run_with_tools(self, mock_mcp_client, mock_anthropic):
+        """Test running agent with tool usage."""
+        agent = ClaudeMCPAgent(mcp_client=mock_mcp_client, model_client=mock_anthropic)
+
+        # Mock tool availability
+        agent._available_tools = [
+            types.Tool(name="calculator", description="Calculator", inputSchema={"type": "object"})
+        ]
+        agent._tool_map = {
+            "calculator": (
+                "server1",
+                types.Tool(
+                    name="calculator", description="Calculator", inputSchema={"type": "object"}
+                ),
+            )
+        }
+
+        # Mock initial response with tool use
+        initial_response = MagicMock()
+        # Create tool use block
+        tool_block = MagicMock()
+        tool_block.type = "tool_use"
+        tool_block.id = "calc_123"
+        tool_block.name = "calculator"
+        tool_block.input = {"operation": "add", "a": 2, "b": 3}
+        initial_response.content = [tool_block]
+        initial_response.usage = MagicMock(input_tokens=10, output_tokens=15)
+
+        # Mock follow-up response
+        final_response = MagicMock()
+        text_block = MagicMock()
+        text_block.type = "text"
+        text_block.text = "2 + 3 = 5"
+        final_response.content = [text_block]
+        final_response.usage = MagicMock(input_tokens=20, output_tokens=10)
+
+        mock_anthropic.beta.messages.create = AsyncMock(
+            side_effect=[initial_response, final_response]
+        )
+
+        # Mock tool execution
+        agent.mcp_client.call_tool = AsyncMock(
+            return_value=types.CallToolResult(
+                content=[types.TextContent(type="text", text="5")], isError=False
+            )
+        )
+
+        # Use a string prompt instead of a task
+        result = await agent.run("What is 2 + 3?")
+
+        assert result.content == "2 + 3 = 5"
+        assert result.done is True

hud/mcp/tests/test_client.py

@@ -0,0 +1,324 @@
+"""Tests for MCP Client implementation."""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from mcp import types
+from pydantic import AnyUrl
+
+from hud.mcp.client import MCPClient
+
+
+class TestMCPClient:
+    """Test MCPClient class."""
+
+    @pytest.fixture
+    def mock_mcp_use_client(self):
+        """Create a mock MCPUseClient (the internal mcp_use client)."""
+        # Create a mock instance that will be returned by from_dict
+        mock_instance = MagicMock()
+        mock_instance.create_session = AsyncMock()
+        mock_instance.create_all_sessions = AsyncMock(return_value={})
+        mock_instance.close_all_sessions = AsyncMock()
+        mock_instance.get_all_active_sessions = MagicMock(return_value={})
+
+        # Patch MCPUseClient that's imported in hud.mcp.client
+        with patch("hud.mcp.client.MCPUseClient") as mock_class:
+            mock_class.from_dict = MagicMock(return_value=mock_instance)
+            yield mock_instance
+
+    @pytest.mark.asyncio
+    async def test_init_with_config(self):
+        """Test client initialization with config dictionary."""
+        mcp_config = {
+            "test_server": {
+                "command": "python",
+                "args": ["-m", "test_server"],
+                "env": {"TEST": "true"},
+            }
+        }
+
+        with patch("hud.mcp.client.MCPUseClient") as mock_use_client:
+            client = MCPClient(mcp_config=mcp_config, verbose=True)
+
+            assert client.verbose is True
+            # Verify MCPUseClient.from_dict was called with proper config
+            mock_use_client.from_dict.assert_called_once_with({"mcpServers": mcp_config})
+
+    @pytest.mark.asyncio
+    async def test_connect_single_server(self, mock_mcp_use_client):
+        """Test connecting to a single server."""
+        config = {"test_server": {"command": "python", "args": ["-m", "test_server"]}}
+
+        # Create the MCPClient - the fixture already patches MCPUseClient
+        client = MCPClient(mcp_config=config, verbose=True)
+
+        # Verify internal client was created properly
+        assert client._mcp_client == mock_mcp_use_client
+
+        # Mock session
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        # Mock list_tools response
+        async def mock_list_tools():
+            return types.ListToolsResult(
+                tools=[
+                    types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
+                    types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
+                ]
+            )
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+
+        # Mock create_all_sessions to return a dict with our session
+        mock_mcp_use_client.create_all_sessions = AsyncMock(
+            return_value={"test_server": mock_session}
+        )
+
+        # Initialize the client (creates sessions and discovers tools)
+        await client.initialize()
+
+        # Verify session was created
+        mock_mcp_use_client.create_all_sessions.assert_called_once()
+
+        # Check tools were discovered
+        assert len(client._available_tools) == 2
+        assert len(client._tool_map) == 2
+        assert "tool1" in client._tool_map
+        assert "tool2" in client._tool_map
+
+    @pytest.mark.asyncio
+    async def test_connect_multiple_servers(self, mock_mcp_use_client):
+        """Test connecting to multiple servers."""
+        config = {
+            "server1": {"command": "python", "args": ["-m", "server1"]},
+            "server2": {"command": "node", "args": ["server2.js"]},
+        }
+
+        client = MCPClient(mcp_config=config)
+
+        # Mock sessions
+        mock_session1 = MagicMock()
+        mock_session1.connector = MagicMock()
+        mock_session1.connector.client_session = MagicMock()
+
+        mock_session2 = MagicMock()
+        mock_session2.connector = MagicMock()
+        mock_session2.connector.client_session = MagicMock()
+
+        # Mock tools for each server
+        async def mock_list_tools1():
+            return types.ListToolsResult(
+                tools=[
+                    types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"})
+                ]
+            )
+
+        async def mock_list_tools2():
+            return types.ListToolsResult(
+                tools=[
+                    types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"})
+                ]
+            )
+
+        mock_session1.connector.client_session.list_tools = mock_list_tools1
+        mock_session2.connector.client_session.list_tools = mock_list_tools2
+
+        # Mock create_all_sessions to return both sessions
+        mock_mcp_use_client.create_all_sessions = AsyncMock(
+            return_value={"server1": mock_session1, "server2": mock_session2}
+        )
+
+        await client.initialize()
+
+        # Verify sessions were created
+        mock_mcp_use_client.create_all_sessions.assert_called_once()
+
+        # Check tools from both servers
+        assert len(client._tool_map) == 2
+        assert "tool1" in client._tool_map
+        assert "tool2" in client._tool_map
+
+    @pytest.mark.asyncio
+    async def test_call_tool(self, mock_mcp_use_client):
+        """Test calling a tool."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        # Setup mock session
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        # Mock tool
+        tool = types.Tool(
+            name="calculator", description="Calculator", inputSchema={"type": "object"}
+        )
+
+        async def mock_list_tools():
+            return types.ListToolsResult(tools=[tool])
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+
+        # Mock tool execution
+        mock_result = types.CallToolResult(
+            content=[types.TextContent(type="text", text="Result: 42")], isError=False
+        )
+
+        mock_session.connector.client_session.call_tool = AsyncMock(return_value=mock_result)
+
+        mock_mcp_use_client.create_all_sessions = AsyncMock(return_value={"test": mock_session})
+
+        await client.initialize()
+
+        # Call the tool
+        result = await client.call_tool("calculator", {"operation": "add", "a": 20, "b": 22})
+
+        assert result == mock_result
+        mock_session.connector.client_session.call_tool.assert_called_once_with(
+            name="calculator", arguments={"operation": "add", "a": 20, "b": 22}
+        )
+
+    @pytest.mark.asyncio
+    async def test_call_tool_not_found(self, mock_mcp_use_client):
+        """Test calling a non-existent tool."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        async def mock_list_tools():
+            return types.ListToolsResult(tools=[])
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+        mock_mcp_use_client.create_all_sessions = AsyncMock(return_value={"test": mock_session})
+
+        await client.initialize()
+
+        with pytest.raises(ValueError, match="Tool 'nonexistent' not found"):
+            await client.call_tool("nonexistent", {})
+
+    @pytest.mark.asyncio
+    async def test_get_telemetry_data(self, mock_mcp_use_client):
+        """Test getting telemetry data."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        # Mock tools
+        async def mock_list_tools():
+            return types.ListToolsResult(tools=[])
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+
+        # Mock telemetry resource
+        mock_telemetry = types.ReadResourceResult(
+            contents=[
+                types.TextResourceContents(
+                    uri=AnyUrl("telemetry://live"),
+                    mimeType="application/json",
+                    text='{"events": [{"type": "test", "data": "value"}]}',
+                )
+            ]
+        )
+
+        mock_session.connector.client_session.read_resource = AsyncMock(return_value=mock_telemetry)
+
+        mock_mcp_use_client.create_all_sessions = AsyncMock(return_value={"test": mock_session})
+
+        await client.initialize()
+
+        telemetry_data = client.get_telemetry_data()
+
+        assert "test" in telemetry_data
+        assert telemetry_data["test"]["events"][0]["type"] == "test"
+
+    @pytest.mark.asyncio
+    async def test_close(self, mock_mcp_use_client):
+        """Test closing client connections."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        async def mock_list_tools():
+            return types.ListToolsResult(tools=[])
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+        mock_mcp_use_client.create_session = AsyncMock(return_value=mock_session)
+        mock_mcp_use_client.close_all_sessions = AsyncMock()
+
+        await client.initialize()
+        await client.close()
+
+        mock_mcp_use_client.close_all_sessions.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_context_manager(self, mock_mcp_use_client):
+        """Test using client as context manager."""
+        mock_session = MagicMock()
+        mock_session.connector = MagicMock()
+        mock_session.connector.client_session = MagicMock()
+
+        async def mock_list_tools():
+            return types.ListToolsResult(tools=[])
+
+        mock_session.connector.client_session.list_tools = mock_list_tools
+        mock_mcp_use_client.create_all_sessions = AsyncMock(return_value={"test": mock_session})
+        mock_mcp_use_client.close_all_sessions = AsyncMock()
+
+        config = {"test": {"command": "test"}}
+        # The fixture already patches MCPUseClient
+        async with MCPClient(mcp_config=config) as client:
+            assert client._mcp_client is not None
+            # Verify that the client uses our mock
+            assert client._mcp_client == mock_mcp_use_client
+
+        # Verify cleanup was called
+        mock_mcp_use_client.close_all_sessions.assert_called_once()
+
+    def test_get_available_tools(self, mock_mcp_use_client):
+        """Test getting available tools."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        # Manually set tools
+        client._available_tools = [
+            types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"}),
+            types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"}),
+        ]
+
+        tools = client.get_available_tools()
+        assert len(tools) == 2
+        assert tools[0].name == "tool1"
+        assert tools[1].name == "tool2"
+
+    def test_get_tool_map(self, mock_mcp_use_client):
+        """Test getting tool map."""
+        config = {"test": {"command": "test"}}
+        client = MCPClient(mcp_config=config)
+
+        # Manually set tool map
+        tool1 = types.Tool(name="tool1", description="Tool 1", inputSchema={"type": "object"})
+        tool2 = types.Tool(name="tool2", description="Tool 2", inputSchema={"type": "object"})
+
+        client._tool_map = {
+            "tool1": ("server1", tool1),
+            "tool2": ("server2", tool2),
+        }
+
+        tool_map = client.get_tool_map()
+        assert len(tool_map) == 2
+        assert tool_map["tool1"][0] == "server1"
+        assert tool_map["tool2"][0] == "server2"