hud-python 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

hud/clients/fastmcp.py

@@ -0,0 +1,202 @@
+"""FastMCP-based client implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from contextlib import AsyncExitStack
+from typing import TYPE_CHECKING, Any
+
+from fastmcp import Client as FastMCPClient
+from mcp import Implementation, types
+from mcp.shared.exceptions import McpError
+
+from hud.types import MCPToolCall, MCPToolResult
+from hud.version import __version__ as hud_version
+
+from .base import BaseHUDClient
+from .utils.retry_transport import create_retry_httpx_client
+
+if TYPE_CHECKING:
+    from pydantic import AnyUrl
+
+logger = logging.getLogger(__name__)
+
+
+class FastMCPHUDClient(BaseHUDClient):
+    """FastMCP-based implementation of HUD MCP client."""
+
+    client_info = Implementation(
+        name="hud-fastmcp", title="hud FastMCP Client", version=hud_version
+    )
+
+    def __init__(self, mcp_config: dict[str, dict[str, Any]] | None = None, **kwargs: Any) -> None:
+        """
+        Initialize FastMCP client with retry support for HTTP transports.
+
+        Args:
+            mcp_config: MCP server configuration dict
+            **kwargs: Additional arguments passed to base class
+        """
+        super().__init__(mcp_config=mcp_config, **kwargs)
+
+        self._stack: AsyncExitStack | None = None
+        self._client: FastMCPClient | None = None
+
+    def _create_transport_with_retry(self, mcp_config: dict[str, dict[str, Any]]) -> Any:
+        """Create transport with retry support for HTTP-based servers."""
+        from fastmcp.client.transports import StreamableHttpTransport
+
+        # If single server with HTTP URL, create transport directly with retry
+        if len(mcp_config) == 1:
+            _, server_config = next(iter(mcp_config.items()))
+            url = server_config.get("url", "")
+
+            if url.startswith("http") and not url.endswith("/sse"):
+                headers = server_config.get("headers", {})
+
+                logger.debug("Enabling retry mechanism for HTTP transport to %s", url)
+                return StreamableHttpTransport(
+                    url=url,
+                    headers=headers,
+                    httpx_client_factory=create_retry_httpx_client,
+                )
+
+        # For all other cases, use standard config (no retry)
+        return {"mcpServers": mcp_config}
+
+    async def _connect(self, mcp_config: dict[str, dict[str, Any]]) -> None:
+        """Enter FastMCP context to establish connection."""
+        if self._client is not None:
+            logger.warning("Client is already connected, cannot connect again")
+            return
+
+        # Create FastMCP client with the custom transport
+        timeout = 5 * 60  # 5 minutes
+        os.environ["FASTMCP_CLIENT_INIT_TIMEOUT"] = str(timeout)
+
+        # Create custom transport with retry support for HTTP servers
+        transport = self._create_transport_with_retry(mcp_config)
+        self._client = FastMCPClient(transport, client_info=self.client_info, timeout=timeout)
+
+        if self._stack is None:
+            self._stack = AsyncExitStack()
+            await self._stack.enter_async_context(self._client)
+
+            # Configure validation for output schemas based on client setting
+            from mcp.client.session import ValidationOptions
+
+            if (
+                hasattr(self._client, "_session_state")
+                and self._client._session_state.session is not None
+            ):
+                self._client._session_state.session._validation_options = ValidationOptions(
+                    strict_output_validation=self._strict_validation
+                )
+
+            logger.info("FastMCP client connected")
+
+    async def list_tools(self) -> list[types.Tool]:
+        """List all available tools."""
+        if self._client is None:
+            raise ValueError("Client is not connected, call initialize() first")
+        return await self._client.list_tools()
+
+    async def _call_tool(self, tool_call: MCPToolCall) -> MCPToolResult:
+        """Execute a tool by name."""
+        if self._client is None:
+            raise ValueError("Client is not connected, call initialize() first")
+
+        # FastMCP returns a different result type, convert it
+        result = await self._client.call_tool(
+            name=tool_call.name,
+            arguments=tool_call.arguments or {},
+            raise_on_error=False,  # Don't raise, return error in result
+        )
+
+        # Convert FastMCP result to MCPToolResult
+        return MCPToolResult(
+            content=result.content,
+            isError=result.is_error,
+            structuredContent=result.structured_content,
+        )
+
+    async def list_resources(self) -> list[types.Resource]:
+        """List all available resources."""
+        if self._client is None:
+            raise ValueError("Client is not connected, call initialize() first")
+        return await self._client.list_resources()
+
+    async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult | None:
+        """Read a resource by URI."""
+        if self._client is None:
+            raise ValueError("Client is not connected, call initialize() first")
+        try:
+            contents = await self._client.read_resource(uri)
+            return types.ReadResourceResult(contents=contents)
+        except McpError as e:
+            if "telemetry://" in str(uri):
+                logger.debug("Telemetry resource not supported by server: %s", e)
+            elif self.verbose:
+                logger.debug("MCP resource error for '%s': %s", uri, e)
+            return None
+        except Exception as e:
+            if "telemetry://" in str(uri):
+                logger.debug("Failed to fetch telemetry: %s", e)
+            else:
+                logger.warning("Unexpected error reading resource '%s': %s", uri, e)
+            return None
+
+    async def _disconnect(self) -> None:
+        """Close the client connection, ensuring the underlying transport is terminated."""
+        if self._client is None:
+            logger.warning("Client is not connected, cannot disconnect")
+            return
+
+        # First close any active async context stack (this triggers client.__aexit__()).
+        if self._stack:
+            await self._stack.aclose()
+            self._stack = None
+
+        try:
+            # Close the FastMCP client - this calls transport.close()
+            await self._client.close()
+
+            # CRITICAL: Cancel any lingering transport tasks to ensure subprocess termination
+            # FastMCP's StdioTransport creates asyncio tasks that can outlive the client
+            # We need to handle nested transport structures (MCPConfigTransport -> StdioTransport)
+            transport = getattr(self._client, "transport", None)
+            if transport:
+                # If it's an MCPConfigTransport with a nested transport
+                if hasattr(transport, "transport"):
+                    transport = transport.transport
+
+                # Now check if it's a StdioTransport with a _connect_task
+                if (
+                    hasattr(transport, "_connect_task")
+                    and transport._connect_task
+                    and not transport._connect_task.done()
+                ):
+                    logger.debug("Canceling lingering StdioTransport connect task")
+                    transport._connect_task.cancel()
+                    try:
+                        await transport._connect_task
+                    except asyncio.CancelledError:
+                        logger.debug("Transport task cancelled successfully")
+                    except Exception as e:
+                        logger.debug("Error canceling transport task: %s", e)
+
+        except Exception as e:
+            logger.debug("Error while closing FastMCP client transport: %s", e)
+
+        logger.debug("FastMCP client closed")
+
+    async def __aenter__(self: Any) -> Any:
+        """Async context manager entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        """Async context manager exit."""
+        await self.shutdown()

hud/clients/mcp_use.py

@@ -0,0 +1,278 @@
+"""MCP-use based client implementation (legacy)."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from mcp import Implementation
+from mcp.shared.exceptions import McpError
+from mcp_use.client import MCPClient as MCPUseClient
+from pydantic import AnyUrl
+
+from hud.types import MCPToolCall, MCPToolResult
+from hud.version import __version__ as hud_version
+
+from .base import BaseHUDClient
+
+if TYPE_CHECKING:
+    from mcp import types
+    from mcp_use.session import MCPSession as MCPUseSession
+
+logger = logging.getLogger(__name__)
+
+
+class MCPUseHUDClient(BaseHUDClient):
+    """MCP-use based implementation of HUD MCP client."""
+
+    client_info = Implementation(
+        name="hud-mcp-use", title="hud MCP-use Client", version=hud_version
+    )
+
+    def __init__(self, mcp_config: dict[str, dict[str, Any]] | None = None, **kwargs: Any) -> None:
+        """
+        Initialize MCP-use client.
+
+        Args:
+            mcp_config: MCP server configuration dict
+            **kwargs: Additional arguments passed to base class
+        """
+        super().__init__(mcp_config=mcp_config, **kwargs)
+
+        self._sessions: dict[str, MCPUseSession] = {}
+        self._tool_map: dict[str, tuple[str, types.Tool]] = {}
+        self._client: MCPUseClient | None = None
+
+    async def _connect(self, mcp_config: dict[str, dict[str, Any]]) -> None:
+        """Create all sessions for MCP-use client."""
+        if self._client is not None:
+            logger.warning("Client is already connected, cannot connect again")
+            return
+
+        config = {"mcpServers": mcp_config}
+        self._client = MCPUseClient.from_dict(config)
+        try:
+            self._sessions = await self._client.create_all_sessions()
+            logger.info("Created %d MCP sessions", len(self._sessions))
+
+            # Configure validation for all sessions based on client setting
+            from mcp.client.session import ValidationOptions
+
+            for session in self._sessions.values():
+                if (
+                    hasattr(session, "connector")
+                    and hasattr(session.connector, "client_session")
+                    and session.connector.client_session is not None
+                ):
+                    session.connector.client_session._validation_options = ValidationOptions(
+                        strict_output_validation=self._strict_validation
+                    )
+
+            # Log session details in verbose mode
+            if self.verbose and self._sessions:
+                for name, session in self._sessions.items():
+                    logger.debug("  - %s: %s", name, type(session).__name__)
+
+        except McpError as e:
+            # Protocol error - the server is reachable but rejecting our request
+            logger.error("MCP protocol error: %s", e)
+            logger.error("This typically means:")
+            logger.error("- Invalid or missing initialization parameters")
+            logger.error("- Incompatible protocol version")
+            logger.error("- Server-side configuration issues")
+            raise
+        except Exception as e:
+            # Transport or other errors
+            logger.error("Failed to create sessions: %s", e)
+            if self.verbose:
+                logger.info("Check that the MCP server is running and accessible")
+            raise
+
+    async def list_tools(self) -> list[types.Tool]:
+        """List all available tools from all sessions."""
+        if self._client is None or not self._sessions:
+            raise ValueError("Client is not connected, call initialize() first")
+
+        all_tools = []
+        self._tool_map = {}
+
+        for server_name, session in self._sessions.items():
+            try:
+                # Ensure session is initialized
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    await session.initialize()
+
+                if session.connector.client_session is None:
+                    logger.warning("Client session not initialized for %s", server_name)
+                    continue
+
+                # List tools
+                tools_result = await session.connector.client_session.list_tools()
+
+                logger.info(
+                    "Discovered %d tools from '%s': %s",
+                    len(tools_result.tools),
+                    server_name,
+                    [tool.name for tool in tools_result.tools],
+                )
+
+                # Add to collections
+                for tool in tools_result.tools:
+                    all_tools.append(tool)
+                    self._tool_map[tool.name] = (server_name, tool)
+
+                # Log detailed tool info in verbose mode
+                if self.verbose:
+                    for tool in tools_result.tools:
+                        description = tool.description or ""
+                        logger.debug(
+                            "  Tool '%s': %s",
+                            tool.name,
+                            description[:100] + "..." if len(description) > 100 else description,
+                        )
+
+            except Exception as e:
+                logger.error("Error discovering tools from '%s': %s", server_name, e)
+                if self.verbose:
+                    logger.exception("Full error details:")
+
+        return all_tools
+
+    async def _call_tool(self, tool_call: MCPToolCall) -> MCPToolResult:
+        """Execute a tool by name."""
+        if self._client is None or not self._initialized:
+            raise ValueError("Client is not connected, call initialize() first")
+
+        if tool_call.name not in self._tool_map:
+            raise ValueError(f"Tool '{tool_call.name}' not found")
+
+        server_name, _ = self._tool_map[tool_call.name]
+        session = self._sessions[server_name]
+
+        if self.verbose:
+            logger.debug(
+                "Calling tool '%s' on server '%s' with arguments: %s",
+                tool_call.name,
+                server_name,
+                tool_call.arguments,
+            )
+
+        if session.connector.client_session is None:
+            raise ValueError(f"Client session not initialized for {server_name}")
+
+        result = await session.connector.client_session.call_tool(
+            name=tool_call.name,
+            arguments=tool_call.arguments or {},
+        )
+
+        if self.verbose:
+            logger.debug("Tool '%s' result: %s", tool_call.name, result)
+
+        # MCP-use already returns the correct type, but we need to ensure it's MCPToolResult
+        return MCPToolResult(
+            content=result.content,
+            isError=result.isError,
+            structuredContent=result.structuredContent,
+        )
+
+    async def list_resources(self) -> list[types.Resource]:
+        """List all available resources."""
+        if self._client is None or not self._sessions:
+            raise ValueError("Client is not connected, call initialize() first")
+
+        for server_name, session in self._sessions.items():
+            try:
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    continue
+                if session.connector.client_session is None:
+                    continue
+                # Prefer standard method name if available
+                if hasattr(session.connector.client_session, "list_resources"):
+                    resources = await session.connector.client_session.list_resources()
+                else:
+                    # If the client doesn't support resource listing, skip
+                    continue
+                return resources.resources
+            except Exception as e:
+                if self.verbose:
+                    logger.debug("Could not list resources from server '%s': %s", server_name, e)
+                continue
+        return []
+
+    async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult | None:
+        """Read a resource by URI from any server that provides it."""
+        if self._client is None or not self._sessions:
+            raise ValueError("Client is not connected, call initialize() first")
+
+        for server_name, session in self._sessions.items():
+            try:
+                if not hasattr(session, "connector") or not hasattr(
+                    session.connector, "client_session"
+                ):
+                    continue
+
+                if session.connector.client_session is None:
+                    continue
+
+                # Convert str to AnyUrl if needed
+                resource_uri = AnyUrl(uri) if isinstance(uri, str) else uri
+                # Prefer read_resource; fall back to list_resources if needed
+                if hasattr(session.connector.client_session, "read_resource"):
+                    result = await session.connector.client_session.read_resource(resource_uri)
+                else:
+                    # Fallback path for older clients: not supported in strict typing
+                    raise AttributeError("read_resource not available")
+
+                if self.verbose:
+                    logger.debug(
+                        "Successfully read resource '%s' from server '%s'", uri, server_name
+                    )
+
+                return result
+
+            except McpError as e:
+                # McpError is expected for unsupported resources
+                if "telemetry://" in str(uri):
+                    logger.debug(
+                        "Telemetry resource not supported by server '%s': %s", server_name, e
+                    )
+                elif self.verbose:
+                    logger.debug(
+                        "MCP resource error for '%s' from server '%s': %s", uri, server_name, e
+                    )
+                continue
+            except Exception as e:
+                # Other errors might be more serious
+                if "telemetry://" in str(uri):
+                    logger.debug("Failed to fetch telemetry from server '%s': %s", server_name, e)
+                else:
+                    logger.warning(
+                        "Unexpected error reading resource '%s' from server '%s': %s",
+                        uri,
+                        server_name,
+                        e,
+                    )
+                continue
+
+        return None
+
+    async def _disconnect(self) -> None:
+        """Close all active sessions."""
+        if self._client is None:
+            logger.warning("Client is not connected, cannot close")
+            return
+
+        await self._client.close_all_sessions()
+        self._sessions = {}
+        self._tool_map = {}
+        self._initialized = False
+        logger.debug("MCP-use client disconnected")
+
+    # Legacy compatibility methods (limited; tests should not rely on these)
+    def get_sessions(self) -> dict[str, MCPUseSession]:
+        """Get active MCP sessions."""
+        return self._sessions

hud/clients/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for HUD MCP clients."""

hud/clients/tests/test_client_integration.py

@@ -0,0 +1,111 @@
+"""Integration tests for MCP clients."""
+
+from __future__ import annotations
+
+import pytest
+
+from hud.clients import FastMCPHUDClient
+from hud.clients.base import AgentMCPClient
+from hud.clients.mcp_use import MCPUseHUDClient
+
+
+class TestClientIntegration:
+    """Test that clients work with real configurations."""
+
+    def test_fastmcp_client_creation(self):
+        """Test that FastMCP client can be created with various configs."""
+        # HTTP config
+        config = {"server": {"url": "http://localhost:8080"}}
+        client = FastMCPHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+        assert client.is_connected is False
+
+        # Stdio config
+        config = {"server": {"command": "python", "args": ["server.py"]}}
+        client = FastMCPHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+
+        # Multi-server config
+        config = {
+            "server1": {"url": "http://localhost:8080"},
+            "server2": {"command": "python", "args": ["server.py"]},
+        }
+        client = FastMCPHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+
+    def test_mcp_use_client_creation(self):
+        """Test that MCP-use client can be created with various configs."""
+        # HTTP config
+        config = {"server": {"url": "http://localhost:8080"}}
+        client = MCPUseHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+        assert client.is_connected is False
+
+        # Stdio config
+        config = {"server": {"command": "python", "args": ["server.py"]}}
+        client = MCPUseHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+
+        # Multi-server config
+        config = {
+            "server1": {"url": "http://localhost:8080"},
+            "server2": {"command": "python", "args": ["server.py"]},
+        }
+        client = MCPUseHUDClient(config)
+        assert isinstance(client, AgentMCPClient)
+
+    def test_client_switching(self):
+        """Test that clients can be switched without changing agent code."""
+        config = {"server": {"url": "http://localhost:8080"}}
+
+        # Both clients should satisfy the protocol
+        fastmcp_client = FastMCPHUDClient(config)
+        mcp_use_client = MCPUseHUDClient(config)
+
+        # Both implement the same protocol
+        assert isinstance(fastmcp_client, AgentMCPClient)
+        assert isinstance(mcp_use_client, AgentMCPClient)
+
+        # Both have the same essential methods
+        for method in ["initialize", "list_tools", "call_tool"]:
+            assert hasattr(fastmcp_client, method)
+            assert hasattr(mcp_use_client, method)
+            assert callable(getattr(fastmcp_client, method))
+            assert callable(getattr(mcp_use_client, method))
+
+    @pytest.mark.asyncio
+    async def test_context_manager_usage(self):
+        """Test that both clients work as context managers."""
+        from unittest.mock import AsyncMock, patch
+
+        config = {"server": {"url": "http://localhost:8080"}}
+
+        # Test FastMCP client with mocked initialization
+        fastmcp_client = FastMCPHUDClient(config)
+        assert not fastmcp_client.is_connected
+
+        with (
+            patch.object(fastmcp_client, "initialize", new_callable=AsyncMock) as mock_init,
+            patch.object(fastmcp_client, "shutdown", new_callable=AsyncMock) as mock_shutdown,
+        ):
+            async with fastmcp_client:
+                # Verify initialization was called
+                mock_init.assert_called_once()
+
+            # Verify shutdown was called
+            mock_shutdown.assert_called_once()
+
+        # Test MCP-use client with mocked initialization
+        mcp_use_client = MCPUseHUDClient(config)
+        assert not mcp_use_client.is_connected
+
+        with (
+            patch.object(mcp_use_client, "initialize", new_callable=AsyncMock) as mock_init,
+            patch.object(mcp_use_client, "shutdown", new_callable=AsyncMock) as mock_shutdown,
+        ):
+            async with mcp_use_client:
+                # Verify initialization was called
+                mock_init.assert_called_once()
+
+            # Verify shutdown was called
+            mock_shutdown.assert_called_once()