hud-python 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

hud/clients/tests/test_fastmcp.py

@@ -0,0 +1,342 @@
+"""Tests for FastMCP client implementation."""
+
+from __future__ import annotations
+
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+import pytest
+from mcp import types
+from pydantic.networks import AnyUrl
+
+from hud.clients.fastmcp import FastMCPHUDClient
+from hud.types import MCPToolCall, MCPToolResult
+
+
+class TestFastMCPHUDClient:
+    """Test FastMCP HUD client."""
+
+    def test_initialization(self):
+        """Test client initialization."""
+        config = {"server1": {"command": "python", "args": ["server.py"]}}
+
+        # Client is just instantiated, not connected yet
+        client = FastMCPHUDClient(config)
+
+        # Check that the client has the config stored
+        assert client._mcp_config == config
+        assert client._client is None  # Not connected yet
+
+    @pytest.mark.asyncio
+    async def test_connect_creates_client(self):
+        """Test that _connect creates the FastMCP client."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            await client._connect(config)
+
+            # Check FastMCP client was created
+            mock_client_class.assert_called_once()
+
+            # Check it was created with correct transport and client info
+            call_args = mock_client_class.call_args
+            assert call_args[0][0] == {"mcpServers": config}
+            assert call_args[1]["client_info"].name == "hud-fastmcp"
+
+    @pytest.mark.asyncio
+    async def test_connect_logs_info(self):
+        """Test that connect logs info message."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+
+            with patch("hud.clients.fastmcp.logger") as mock_logger:
+                await client._connect(config)
+
+                # Check info was logged
+                mock_logger.info.assert_called_with("FastMCP client connected")
+
+    @pytest.mark.asyncio
+    async def test_list_tools(self):
+        """Test listing tools."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_tools = [
+                MagicMock(spec=types.Tool, name="tool1"),
+                MagicMock(spec=types.Tool, name="tool2"),
+            ]
+            mock_fastmcp.list_tools.return_value = mock_tools
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            client._initialized = True  # Skip initialization
+            client._client = mock_fastmcp  # Set the mock client
+
+            tools = await client.list_tools()
+
+            assert tools == mock_tools
+            mock_fastmcp.list_tools.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_call_tool(self):
+        """Test calling a tool."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+
+            # Mock FastMCP result
+            mock_result = MagicMock()
+            mock_result.content = [types.TextContent(type="text", text="result")]
+            mock_result.is_error = False
+            mock_result.structured_content = {"key": "value"}
+
+            mock_fastmcp.call_tool.return_value = mock_result
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            result = await client.call_tool(name="test_tool", arguments={"arg": "value"})
+
+            assert isinstance(result, MCPToolResult)
+            assert result.content == mock_result.content
+            assert result.isError is False
+            assert result.structuredContent == {"key": "value"}
+
+            mock_fastmcp.call_tool.assert_called_once_with(
+                name="test_tool",
+                arguments={"arg": "value"},
+                raise_on_error=False,
+            )
+
+    @pytest.mark.asyncio
+    async def test_call_tool_with_mcp_tool_call(self):
+        """Test calling a tool with MCPToolCall object."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+
+            # Mock FastMCP result
+            mock_result = MagicMock()
+            mock_result.content = [types.TextContent(type="text", text="result")]
+            mock_result.is_error = False
+            mock_result.structured_content = {"key": "value"}
+
+            mock_fastmcp.call_tool.return_value = mock_result
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            # Test with MCPToolCall object
+            tool_call = MCPToolCall(name="test_tool", arguments={"arg": "value"})
+            result = await client.call_tool(tool_call)
+
+            assert isinstance(result, MCPToolResult)
+            assert result.content == mock_result.content
+            assert result.isError is False
+            assert result.structuredContent == {"key": "value"}
+
+            mock_fastmcp.call_tool.assert_called_once_with(
+                name="test_tool",
+                arguments={"arg": "value"},
+                raise_on_error=False,
+            )
+
+    @pytest.mark.asyncio
+    async def test_call_tool_no_arguments(self):
+        """Test calling a tool without arguments."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_result = MagicMock()
+            mock_result.content = []
+            mock_result.is_error = True
+            mock_result.structured_content = None
+
+            mock_fastmcp.call_tool.return_value = mock_result
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            await client.call_tool(name="test_tool", arguments={})
+
+            # Should pass empty dict for arguments
+            mock_fastmcp.call_tool.assert_called_once_with(
+                name="test_tool",
+                arguments={},
+                raise_on_error=False,
+            )
+
+    @pytest.mark.asyncio
+    async def test_list_resources(self):
+        """Test listing resources."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_resources = [
+                MagicMock(spec=types.Resource, uri="file:///test1"),
+                MagicMock(spec=types.Resource, uri="file:///test2"),
+            ]
+            mock_fastmcp.list_resources.return_value = mock_resources
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            resources = await client.list_resources()
+
+            assert resources == mock_resources
+            mock_fastmcp.list_resources.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_read_resource_internal_success(self):
+        """Test reading a resource successfully."""
+        config = {"server1": {"command": "test"}}
+
+        # Create proper resource contents that ReadResourceResult expects
+        mock_contents = [
+            types.TextResourceContents(
+                uri=AnyUrl("file:///test"),
+                mimeType="text/plain",
+                text="resource content",
+            )
+        ]
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            # Create a mock FastMCP client
+            mock_fastmcp = AsyncMock()
+            mock_fastmcp.read_resource.return_value = mock_contents
+            mock_client_class.return_value = mock_fastmcp
+
+            # Now create the HUD client - it will use our mocked FastMCP client
+            client = FastMCPHUDClient(config)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            result = await client.read_resource("file:///test")
+
+            assert isinstance(result, types.ReadResourceResult)
+            assert result.contents == mock_contents
+            mock_fastmcp.read_resource.assert_called_once_with("file:///test")
+
+    @pytest.mark.asyncio
+    async def test_read_resource_internal_error_verbose(self):
+        """Test reading a resource with error in verbose mode."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_fastmcp.read_resource.side_effect = Exception("Read failed")
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config, verbose=True)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            with patch("hud.clients.fastmcp.logger") as mock_logger:
+                result = await client.read_resource("file:///bad")
+
+                assert result is None
+                mock_logger.warning.assert_called_with(
+                    "Unexpected error reading resource '%s': %s", "file:///bad", ANY
+                )
+
+    @pytest.mark.asyncio
+    async def test_read_resource_internal_error_not_verbose(self):
+        """Test reading a resource with error in non-verbose mode."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_fastmcp.read_resource.side_effect = Exception("Read failed")
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config, verbose=False)
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            with patch("hud.clients.fastmcp.logger") as mock_logger:
+                result = await client.read_resource("file:///bad")
+
+                assert result is None
+                # Should not log in non-verbose mode
+                mock_logger.debug.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_shutdown(self):
+        """Test shutting down the client."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+
+            # Set up stack and client
+            mock_stack = AsyncMock()
+            client._stack = mock_stack
+            client._initialized = True
+            client._client = mock_fastmcp  # Set the mock client
+
+            with patch("hud.clients.fastmcp.logger") as mock_logger:
+                await client.shutdown()
+
+                mock_stack.aclose.assert_called_once()
+                assert client._stack is None
+                assert client._initialized is False
+                mock_logger.debug.assert_called_with("FastMCP client closed")
+
+    @pytest.mark.asyncio
+    async def test_shutdown_no_stack(self):
+        """Test shutting down when no stack exists."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient"):
+            client = FastMCPHUDClient(config)
+            client._stack = None
+
+            # Should not raise error
+            await client.shutdown()
+
+            assert client._stack is None
+
+    @pytest.mark.asyncio
+    async def test_context_manager(self):
+        """Test using client as async context manager."""
+        config = {"server1": {"command": "test"}}
+
+        with patch("hud.clients.fastmcp.FastMCPClient") as mock_client_class:
+            mock_fastmcp = AsyncMock()
+            mock_client_class.return_value = mock_fastmcp
+
+            client = FastMCPHUDClient(config)
+
+            with (
+                patch.object(client, "initialize", new_callable=AsyncMock) as mock_init,
+                patch.object(client, "shutdown", new_callable=AsyncMock) as mock_close,
+            ):
+                async with client as ctx:
+                    assert ctx is client
+                    mock_init.assert_called_once()
+
+                mock_close.assert_called_once()

hud/clients/tests/test_protocol.py

@@ -0,0 +1,188 @@
+"""Tests for the MCP client protocol and implementations."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+from mcp import types
+
+from hud.clients.base import AgentMCPClient, BaseHUDClient
+from hud.clients.fastmcp import FastMCPHUDClient
+from hud.clients.mcp_use import MCPUseHUDClient
+from hud.types import MCPToolCall, MCPToolResult
+
+
+class MockClient(BaseHUDClient):
+    """Mock client for testing the base class."""
+
+    def __init__(self, **kwargs):
+        super().__init__(mcp_config={"test": {"url": "mock://test"}}, **kwargs)
+        self._connected = False
+        self._mock_tools = [
+            types.Tool(
+                name="test_tool",
+                description="A test tool",
+                inputSchema={"type": "object", "properties": {}},
+            )
+        ]
+
+    async def _connect(self, mcp_config: dict[str, dict[str, Any]]) -> None:
+        self._connected = True
+
+    async def list_tools(self) -> list[types.Tool]:
+        if not self._connected:
+            raise RuntimeError("Not connected")
+        return self._mock_tools
+
+    async def list_resources(self) -> list[types.Resource]:
+        """Minimal list_resources for protocol satisfaction in tests."""
+        return []
+
+    async def _call_tool(self, tool_call: MCPToolCall) -> MCPToolResult:
+        if tool_call.name == "test_tool":
+            return MCPToolResult(
+                content=[types.TextContent(type="text", text="Success")], isError=False
+            )
+        raise ValueError(f"Tool {tool_call.name} not found")
+
+    async def read_resource(self, uri: str) -> types.ReadResourceResult | None:
+        if uri == "telemetry://live":
+            from pydantic import AnyUrl
+
+            return types.ReadResourceResult(
+                contents=[
+                    types.TextResourceContents(
+                        uri=AnyUrl(uri),
+                        mimeType="application/json",
+                        text='{"status": "healthy", "services": {"api": "running"}}',
+                    )
+                ]
+            )
+        return None
+
+    async def _disconnect(self) -> None:
+        """Disconnect from the MCP server."""
+        self._connected = False
+
+
+class TestProtocol:
+    """Test that all clients implement the protocol correctly."""
+
+    def test_mock_client_implements_protocol(self):
+        """Test that our mock client implements the protocol."""
+        client = MockClient()
+        assert isinstance(client, AgentMCPClient)
+
+    def test_fastmcp_client_implements_protocol(self):
+        """Test that FastMCPHUDClient implements the protocol."""
+        client = FastMCPHUDClient({"test": {"url": "http://localhost"}})
+        assert isinstance(client, AgentMCPClient)
+
+    def test_mcp_use_client_implements_protocol(self):
+        """Test that MCPUseHUDClient implements the protocol."""
+        client = MCPUseHUDClient({"test": {"url": "http://localhost"}})
+        assert isinstance(client, AgentMCPClient)
+
+    @pytest.mark.asyncio
+    async def test_base_client_initialization(self):
+        """Test that base client initialization works correctly."""
+        client = MockClient()
+
+        # Not initialized yet
+        assert not client._initialized
+        # Can't call list_tools before initialization, it would raise an error
+
+        # Initialize
+        await client.initialize()
+
+        # Should be initialized with tools discovered
+        assert client._initialized
+        tools = await client.list_tools()
+        assert len(tools) == 1
+        assert tools[0].name == "test_tool"
+
+    @pytest.mark.asyncio
+    async def test_telemetry_fetching(self):
+        """Test that telemetry is fetched during initialization."""
+        client = MockClient()
+
+        # No telemetry before initialization
+        assert not hasattr(client, "_telemetry_data") or client._telemetry_data == {}
+
+        # Initialize
+        await client.initialize()
+
+        # Should have telemetry
+        assert hasattr(client, "_telemetry_data")
+        assert client._telemetry_data["status"] == "healthy"
+        assert client._telemetry_data["services"]["api"] == "running"
+
+    @pytest.mark.asyncio
+    async def test_context_manager(self):
+        """Test that clients work as context managers."""
+        client = MockClient()
+
+        async with client:
+            assert client._initialized
+            tools = await client.list_tools()
+            assert len(tools) == 1
+
+        # Should be closed after exiting context
+        assert not client._initialized
+
+    @pytest.mark.asyncio
+    async def test_tool_execution(self):
+        """Test tool execution through the protocol."""
+        client = MockClient()
+
+        await client.initialize()
+
+        # Execute a tool - test both call signatures
+        # Test with MCPToolCall
+        tool_call = MCPToolCall(name="test_tool", arguments={"arg": "value"})
+        result = await client.call_tool(tool_call)
+
+        assert isinstance(result, MCPToolResult)
+        assert not result.isError
+        from mcp.types import TextContent
+
+        assert isinstance(result.content[0], TextContent) and result.content[0].text == "Success"
+
+        # Test with name/arguments
+        result2 = await client.call_tool(name="test_tool", arguments={"arg": "value"})
+        assert isinstance(result2, MCPToolResult)
+        assert not result2.isError
+        assert isinstance(result2.content[0], TextContent) and result2.content[0].text == "Success"
+
+    @pytest.mark.asyncio
+    async def test_tool_not_found(self):
+        """Test error handling for missing tools."""
+        client = MockClient()
+
+        await client.initialize()
+
+        # Try to execute non-existent tool
+        with pytest.raises(ValueError, match="Tool unknown_tool not found"):
+            await client.call_tool(name="unknown_tool", arguments={})
+
+
+class TestClientCompatibility:
+    """Test that clients are compatible with agents."""
+
+    def test_protocol_satisfied(self):
+        """Test that all clients satisfy the protocol."""
+        # Test mock client
+        mock_client = MockClient()
+        assert isinstance(mock_client, AgentMCPClient)
+        assert hasattr(mock_client, "initialize")
+        assert hasattr(mock_client, "list_tools")
+        assert hasattr(mock_client, "call_tool")
+
+        # Test FastMCP client
+        fastmcp_client = FastMCPHUDClient({"test": {"url": "http://localhost"}})
+        assert isinstance(fastmcp_client, AgentMCPClient)
+
+        # Test MCP-use client
+        mcp_use_client = MCPUseHUDClient({"test": {"url": "http://localhost"}})
+        assert isinstance(mcp_use_client, AgentMCPClient)

hud/clients/utils/__init__.py

@@ -0,0 +1 @@
+"""HUD MCP client utilities."""

hud/clients/utils/retry_transport.py

@@ -0,0 +1,160 @@
+"""Custom HTTPX transport with retry logic for HTTP errors."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any
+
+import httpx
+from httpx._transports.default import AsyncHTTPTransport
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from httpx._models import Request, Response
+
+
+class RetryTransport(AsyncHTTPTransport):
+    """
+    Custom HTTPX transport that retries on specific HTTP status codes.
+
+    This transport wraps the standard AsyncHTTPTransport and adds
+    retry logic with exponential backoff for gateway errors (502, 503, 504).
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        max_retries: int = 3,
+        retry_status_codes: set[int] | None = None,
+        retry_delay: float = 1.0,
+        backoff_factor: float = 2.0,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize retry transport.
+
+        Args:
+            max_retries: Maximum number of retry attempts
+            retry_status_codes: HTTP status codes to retry (default: 502, 503, 504)
+            retry_delay: Initial delay between retries in seconds
+            backoff_factor: Multiplier for exponential backoff
+            *args, **kwargs: Passed to AsyncHTTPTransport
+        """
+        super().__init__(*args, **kwargs)
+        self.max_retries = max_retries
+        self.retry_status_codes = retry_status_codes or {502, 503, 504}
+        self.retry_delay = retry_delay
+        self.backoff_factor = backoff_factor
+
+    async def handle_async_request(self, request: Request) -> Response:
+        """
+        Handle request with retry logic.
+
+        Retries the request if it fails with a retryable status code,
+        using exponential backoff between attempts.
+        """
+        last_exception = None
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                response = await super().handle_async_request(request)
+
+                # Check if we should retry based on status code
+                if response.status_code in self.retry_status_codes and attempt < self.max_retries:
+                    delay = self.retry_delay * (self.backoff_factor**attempt)
+                    logger.warning(
+                        "Got %d from %s, retrying in %.1fs (attempt %d/%d)",
+                        response.status_code,
+                        request.url,
+                        delay,
+                        attempt + 1,
+                        self.max_retries,
+                    )
+                    # Important: Close the response to free resources
+                    await response.aclose()
+                    await asyncio.sleep(delay)
+                    continue
+
+                return response
+
+            except (httpx.ConnectError, httpx.TimeoutException) as e:
+                last_exception = e
+                if attempt < self.max_retries:
+                    delay = self.retry_delay * (self.backoff_factor**attempt)
+                    logger.warning(
+                        "%s for %s, retrying in %.1fs (attempt %d/%d)",
+                        type(e).__name__,
+                        request.url,
+                        delay,
+                        attempt + 1,
+                        self.max_retries,
+                    )
+                    await asyncio.sleep(delay)
+                    continue
+                raise
+
+        # If we get here, we've exhausted retries
+        if last_exception:
+            raise last_exception
+        else:
+            # This shouldn't happen, but just in case
+            raise httpx.HTTPStatusError(
+                "Max retries exceeded",
+                request=request,
+                response=response,
+            )
+
+
+def create_retry_httpx_client(
+    headers: dict[str, str] | None = None,
+    timeout: httpx.Timeout | None = None,
+    auth: httpx.Auth | None = None,
+    max_retries: int = 3,
+    retry_status_codes: set[int] | None = None,
+) -> httpx.AsyncClient:
+    """
+    Create an HTTPX AsyncClient with HTTP error retry support.
+
+    This factory creates an HTTPX client with a custom transport that
+    retries on specific HTTP status codes (502, 503, 504 by default).
+
+    Args:
+        headers: Optional headers to include with all requests
+        timeout: Request timeout (defaults to 600s)
+        auth: Optional authentication handler
+        max_retries: Maximum retry attempts (default: 3)
+        retry_status_codes: Status codes to retry (default: {502, 503, 504})
+
+    Returns:
+        Configured httpx.AsyncClient with retry transport
+    """
+    if timeout is None:
+        timeout = httpx.Timeout(600.0)  # 10 minutes
+
+    # Use higher connection limits for concurrent operations
+    # These match HUD server's configuration for consistency
+    limits = httpx.Limits(
+        max_connections=1000,
+        max_keepalive_connections=1000,
+        keepalive_expiry=20.0,
+    )
+
+    # Create our custom retry transport
+    transport = RetryTransport(
+        max_retries=max_retries,
+        retry_status_codes=retry_status_codes,
+        # Connection-level retries (in addition to HTTP retries)
+        retries=3,
+        limits=limits,
+    )
+
+    return httpx.AsyncClient(
+        transport=transport,
+        headers=headers,
+        timeout=timeout,
+        auth=auth,
+        follow_redirects=True,
+        limits=limits,
+    )