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

hud/{mcp_agent → mcp}/claude.py

@@ -17,9 +17,15 @@ if TYPE_CHECKING:
         BetaToolResultBlockParam,
     )
 
+    from hud.datasets import TaskConfig
+
+import mcp.types as types
+from mcp.types import CallToolRequestParams as MCPToolCall
+from mcp.types import CallToolResult as MCPToolResult
+
 from hud.settings import settings
 
-from .base import BaseMCPAgent
+from .base import BaseMCPAgent, ModelResponse
 
 logger = logging.getLogger(__name__)
 
@@ -66,13 +72,13 @@ class ClaudeMCPAgent(BaseMCPAgent):
         Initialize Claude MCP agent.
 
         Args:
-            client: AsyncAnthropic client (created if not provided)
+            model_client: AsyncAnthropic client (created if not provided)
             model: Claude model to use
             max_tokens: Maximum tokens for response
             display_width_px: Display width for computer use tools
             display_height_px: Display height for computer use tools
             use_computer_beta: Whether to use computer-use beta features
-            **kwargs: Additional arguments passed to BaseMCPAgent
+            **kwargs: Additional arguments passed to BaseMCPAgent (including mcp_client)
         """
         super().__init__(**kwargs)
 
@@ -90,17 +96,19 @@ class ClaudeMCPAgent(BaseMCPAgent):
         self.display_height_px = display_height_px
         self.use_computer_beta = use_computer_beta
 
+        self.model_name = self.model
+
         # Track mapping from Claude tool names to MCP tool names
         self._claude_to_mcp_tool_map: dict[str, str] = {}
 
-    async def initialize(self) -> None:
+    async def initialize(self, task: str | TaskConfig | None = None) -> None:
         """Initialize the agent and build tool mappings."""
-        await super().initialize()
+        await super().initialize(task)
         # Build tool mappings after tools are discovered
         self._convert_tools_for_claude()
 
     async def create_initial_messages(
-        self, prompt: str, screenshot: str | None
+        self, prompt: str, screenshot: str | None = None
     ) -> list[BetaMessageParam]:
         """Create initial messages for Claude."""
         user_content: list[BetaImageBlockParam | BetaTextBlockParam] = []
@@ -123,9 +131,7 @@ class ClaudeMCPAgent(BaseMCPAgent):
             )
         ]
 
-    async def get_model_response(
-        self, messages: list[BetaMessageParam], step: int
-    ) -> dict[str, Any]:
+    async def get_model_response(self, messages: list[BetaMessageParam]) -> ModelResponse:
         """Get response from Claude including any tool calls."""
         # Get Claude tools
         claude_tools = self._convert_tools_for_claude()
@@ -166,7 +172,6 @@ class ClaudeMCPAgent(BaseMCPAgent):
                 else:
                     raise
 
-        # Add assistant response to messages (for next step)
         messages.append(
             cast(
                 "BetaMessageParam",
@@ -178,12 +183,7 @@ class ClaudeMCPAgent(BaseMCPAgent):
         )
 
         # Process response
-        result = {
-            "content": "",
-            "tool_calls": [],
-            "done": True,
-            "raw_response": response.model_dump(),  # For debugging
-        }
+        result = ModelResponse(content="", tool_calls=[], done=True)
 
         # Extract text content and reasoning
         text_content = ""
@@ -194,16 +194,16 @@ class ClaudeMCPAgent(BaseMCPAgent):
                 # Map Claude tool name back to MCP tool name
                 mcp_tool_name = self._claude_to_mcp_tool_map.get(block.name, block.name)
 
-                # Include the tool_use_id in the tool call for later reference
-                result["tool_calls"].append(
-                    {
-                        "name": mcp_tool_name,  # Use MCP tool name for execution
-                        "arguments": block.input,
-                        "tool_use_id": block.id,  # Claude-specific metadata
-                        "claude_name": block.name,  # Keep original Claude name for reference
-                    }
+                # Create MCPToolCall object with Claude metadata as extra fields
+                # Pyright will complain but the tool class accepts extra fields
+                tool_call = MCPToolCall(
+                    name=mcp_tool_name,
+                    arguments=block.input,
+                    tool_use_id=block.id,  # type: ignore
+                    claude_name=block.name,  # type: ignore
                 )
-                result["done"] = False
+                result.tool_calls.append(tool_call)
+                result.done = False
             elif block.type == "text":
                 text_content += block.text
             elif hasattr(block, "type") and block.type == "thinking":
@@ -211,41 +211,44 @@ class ClaudeMCPAgent(BaseMCPAgent):
 
         # Combine text and thinking for final content
         if thinking_content:
-            result["content"] = thinking_content + text_content
+            result.content = thinking_content + text_content
         else:
-            result["content"] = text_content
+            result.content = text_content
 
         return result
 
     async def format_tool_results(
-        self, processed_results: dict[str, Any], tool_calls: list[dict]
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
     ) -> list[BetaMessageParam]:
         """Format tool results into Claude messages."""
-        # Build a mapping of tool_name to tool_use_id from the original calls
-        tool_id_map = {}
-        for tool_call in tool_calls:
-            if "tool_use_id" in tool_call:
-                tool_id_map[tool_call["name"]] = tool_call["tool_use_id"]
-
         # Process each tool result
         user_content = []
 
-        for tool_name, content_blocks in processed_results["results"]:
-            # Get the tool_use_id for this tool
-            tool_use_id = tool_id_map.get(tool_name)
+        for tool_call, result in zip(tool_calls, tool_results, strict=True):
+            # Extract Claude-specific metadata from extra fields
+            tool_use_id = getattr(tool_call, "tool_use_id", None)
             if not tool_use_id:
-                logger.warning("No tool_use_id found for %s", tool_name)
+                logger.warning("No tool_use_id found for %s", tool_call.name)
                 continue
 
-            # Convert content blocks to Claude format
+            # Convert MCP tool results to Claude format
             claude_blocks = []
-            for block in content_blocks:
-                if block["type"] == "text":
-                    claude_blocks.append(text_to_content_block(block["text"]))
-                elif block["type"] == "error":
-                    claude_blocks.append(text_to_content_block(f"Error: {block['text']}"))
-                elif block["type"] == "image":
-                    claude_blocks.append(base64_to_content_block(block["data"]))
+
+            if result.isError:
+                # Extract error message from content
+                error_msg = "Tool execution failed"
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        error_msg = content.text
+                        break
+                claude_blocks.append(text_to_content_block(f"Error: {error_msg}"))
+            else:
+                # Process success content
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        claude_blocks.append(text_to_content_block(content.text))
+                    elif isinstance(content, types.ImageContent):
+                        claude_blocks.append(base64_to_content_block(content.data))
 
             # Add tool result
             user_content.append(tool_use_content_block(tool_use_id, claude_blocks))
@@ -282,7 +285,7 @@ class ClaudeMCPAgent(BaseMCPAgent):
                 }
                 # Map Claude's "computer" back to the actual MCP tool name
                 self._claude_to_mcp_tool_map["computer"] = tool.name
-            else:
+            elif tool.name not in self.lifecycle_tools:
                 # Convert regular tools
                 claude_tool = {
                     "name": tool.name,
@@ -295,6 +298,8 @@ class ClaudeMCPAgent(BaseMCPAgent):
                 }
                 # Direct mapping for non-computer tools
                 self._claude_to_mcp_tool_map[tool.name] = tool.name
+            else:
+                continue
 
             claude_tools.append(claude_tool)
 

hud/mcp/client.py

@@ -0,0 +1,312 @@
+"""MCP Client wrapper with automatic initialization and debugging capabilities."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+
+from mcp_use.client import MCPClient as MCPUseClient
+from pydantic import AnyUrl
+
+if TYPE_CHECKING:
+    from typing import Self
+
+    from mcp import types
+    from mcp_use.session import MCPSession as MCPUseSession
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClient:
+    """
+    High-level MCP client wrapper that handles initialization, tool discovery,
+    and provides debugging capabilities.
+    """
+
+    def __init__(
+        self,
+        mcp_config: dict[str, dict[str, Any]],
+        verbose: bool = False,
+    ) -> None:
+        """
+        Initialize the MCP client.
+
+        Args:
+            mcp_config: MCP server configuration dict (required)
+            verbose: Enable verbose logging of server communications
+            auto_initialize: Whether to automatically initialize on construction
+        """
+        self.verbose = verbose
+
+        # Initialize mcp_use client with proper config
+        # Use from_dict to properly initialize with config
+        config = {"mcpServers": mcp_config}
+        self._mcp_client = MCPUseClient.from_dict(config)
+
+        self._sessions: dict[str, MCPUseSession] = {}
+        self._available_tools: list[types.Tool] = []
+        self._tool_map: dict[str, tuple[str, types.Tool]] = {}
+        self._telemetry_data: dict[str, Any] = {}
+
+        # Set up verbose logging if requested
+        if self.verbose:
+            self._setup_verbose_logging()
+
+    def _setup_verbose_logging(self) -> None:
+        """Configure verbose logging for debugging."""
+        # Set MCP-related loggers to DEBUG
+        logging.getLogger("mcp").setLevel(logging.DEBUG)
+        logging.getLogger("mcp_use").setLevel(logging.DEBUG)
+        logging.getLogger("mcp.client.stdio").setLevel(logging.DEBUG)
+
+        # Add handler for server communications
+        if not any(isinstance(h, logging.StreamHandler) for h in logger.handlers):
+            handler = logging.StreamHandler()
+            handler.setFormatter(
+                logging.Formatter("[%(levelname)s] %(asctime)s - %(name)s - %(message)s")
+            )
+            logger.addHandler(handler)
+            logger.setLevel(logging.DEBUG)
+
+    async def initialize(self) -> None:
+        """Perform async initialization tasks."""
+        await self.create_sessions()
+        await self.discover_tools()
+        await self.fetch_telemetry()
+
+    async def create_sessions(self) -> dict[str, MCPUseSession]:
+        # Create all sessions at once
+        try:
+            self._sessions = await self._mcp_client.create_all_sessions()
+        except Exception as e:
+            # If session creation fails, try to get Docker logs
+            logger.error("Failed to create sessions: %s", e)
+            if self.verbose:
+                logger.info("Attempting to check Docker container status...")
+                # await self._check_docker_containers()
+            raise
+
+        # 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__)
+
+        return self._sessions
+
+    async def discover_tools(self) -> list[types.Tool]:
+        """Discover all available tools from connected servers."""
+        logger.info("Discovering available tools...")
+
+        self._available_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:
+                    self._available_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:")
+
+        logger.info("Total tools discovered: %d", len(self._available_tools))
+        return self._available_tools
+
+    async def fetch_telemetry(self) -> dict[str, Any]:
+        """Fetch telemetry resource from all servers that provide it."""
+        logger.info("Fetching telemetry resources...")
+
+        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
+
+                # Try to read telemetry resource
+                try:
+                    result = await session.connector.client_session.read_resource(
+                        AnyUrl("telemetry://live")
+                    )
+                    if result and result.contents and len(result.contents) > 0:
+                        telemetry_data = json.loads(result.contents[0].text)  # type: ignore
+                        self._telemetry_data[server_name] = telemetry_data
+
+                        logger.info("📡 Telemetry data from server '%s':", server_name)
+                        if "live_url" in telemetry_data:
+                            logger.info("   🖥️  Live URL: %s", telemetry_data["live_url"])
+                        if "status" in telemetry_data:
+                            logger.info("   📊 Status: %s", telemetry_data["status"])
+                        if "services" in telemetry_data:
+                            logger.info("   📋 Services:")
+                            for service, status in telemetry_data["services"].items():
+                                status_icon = "✅" if status == "running" else "❌"
+                                logger.info("      %s %s: %s", status_icon, service, status)
+
+                        if self.verbose:
+                            logger.debug(
+                                "Full telemetry data:\n%s", json.dumps(telemetry_data, indent=2)
+                            )
+
+                except Exception as e:
+                    # Resource might not exist, which is fine
+                    if self.verbose:
+                        logger.debug("No telemetry resource from '%s': %s", server_name, e)
+
+            except Exception as e:
+                logger.error("Error fetching telemetry from '%s': %s", server_name, e)
+
+        return self._telemetry_data
+
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any] | None = None
+    ) -> types.CallToolResult:
+        """
+        Call a tool by name with the given arguments.
+
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Tool arguments
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If tool not found
+        """
+        if tool_name not in self._tool_map:
+            raise ValueError(f"Tool '{tool_name}' not found")
+
+        server_name, tool = self._tool_map[tool_name]
+        session = self._sessions[server_name]
+
+        if self.verbose:
+            logger.debug(
+                "Calling tool '%s' on server '%s' with arguments: %s",
+                tool_name,
+                server_name,
+                json.dumps(arguments, indent=2) if arguments else "None",
+            )
+
+        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_name, arguments=arguments or {}
+        )
+
+        if self.verbose:
+            logger.debug("Tool '%s' result: %s", tool_name, result)
+
+        return result
+
+    async def read_resource(self, uri: AnyUrl) -> types.ReadResourceResult | None:
+        """
+        Read a resource by URI from any server that provides it.
+
+        Args:
+            uri: Resource URI (e.g., "telemetry://live")
+
+        Returns:
+            Resource contents or None if not found
+        """
+        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
+
+                result = await session.connector.client_session.read_resource(uri)
+
+                if self.verbose:
+                    logger.debug(
+                        "Successfully read resource '%s' from server '%s'", uri, server_name
+                    )
+
+                return result
+
+            except Exception as e:
+                if self.verbose:
+                    logger.debug(
+                        "Could not read resource '%s' from server '%s': %s", uri, server_name, e
+                    )
+                continue
+
+        return None
+
+    def get_available_tools(self) -> list[types.Tool]:
+        """Get list of all available tools."""
+        return self._available_tools
+
+    def get_tool_map(self) -> dict[str, tuple[str, types.Tool]]:
+        """Get mapping of tool names to (server_name, tool) tuples."""
+        return self._tool_map
+
+    def get_sessions(self) -> dict[str, MCPUseSession]:
+        """Get active MCP sessions."""
+        return self._sessions
+
+    def get_telemetry_data(self) -> dict[str, Any]:
+        """Get collected telemetry data from all servers."""
+        return self._telemetry_data
+
+    def get_all_active_sessions(self) -> dict[str, MCPUseSession]:
+        """Get all active sessions (compatibility method)."""
+        return self._sessions
+
+    async def close(self) -> None:
+        """Close all active sessions."""
+        await self._mcp_client.close_all_sessions()
+
+        self._sessions = {}
+        self._available_tools = []
+        self._tool_map = {}
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        """Async context manager exit."""
+        await self.close()

hud/{mcp_agent → mcp}/langchain.py

@@ -5,15 +5,18 @@ from __future__ import annotations
 import logging
 from typing import TYPE_CHECKING, Any
 
+import mcp.types as types
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from mcp.types import CallToolRequestParams as MCPToolCall
+from mcp.types import CallToolResult as MCPToolResult
 from mcp_use.adapters.langchain_adapter import LangChainAdapter
 
 if TYPE_CHECKING:
     from langchain.schema.language_model import BaseLanguageModel
     from langchain_core.tools import BaseTool
-from .base import BaseMCPAgent
+from .base import BaseMCPAgent, ModelResponse
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +47,12 @@ class LangChainMCPAgent(BaseMCPAgent):
         self.adapter = LangChainAdapter(disallowed_tools=self.disallowed_tools)
         self._langchain_tools: list[BaseTool] | None = None
 
+        self.model_name = (
+            "langchain-" + self.llm.model_name  # type: ignore
+            if hasattr(self.llm, "model_name")
+            else "unknown"
+        )
+
     def _get_langchain_tools(self) -> list[BaseTool]:
         """Get or create LangChain tools from MCP tools."""
         if self._langchain_tools is not None:
@@ -86,7 +95,7 @@ class LangChainMCPAgent(BaseMCPAgent):
 
         return messages
 
-    async def get_model_response(self, messages: list[BaseMessage], step: int) -> dict[str, Any]:
+    async def get_model_response(self, messages: list[BaseMessage]) -> ModelResponse:
         """Get response from LangChain model including any tool calls."""
         # Get LangChain tools (created lazily)
         langchain_tools = self._get_langchain_tools()
@@ -133,11 +142,7 @@ class LangChainMCPAgent(BaseMCPAgent):
                 break
 
         if not last_user_msg:
-            return {
-                "content": "No user message found",
-                "tool_calls": [],
-                "done": True,
-            }
+            return ModelResponse(content="No user message found", tool_calls=[], done=True)
 
         # Extract text from message content
         input_text = ""
@@ -175,54 +180,68 @@ class LangChainMCPAgent(BaseMCPAgent):
                 for action, _ in result["intermediate_steps"]:
                     if hasattr(action, "tool") and hasattr(action, "tool_input"):
                         tool_calls.append(
-                            {
-                                "name": action.tool,
-                                "arguments": action.tool_input,
-                            }
+                            MCPToolCall(
+                                name=action.tool,
+                                arguments=action.tool_input,
+                            )
                         )
 
-                return {
-                    "content": output,
-                    "tool_calls": tool_calls,
-                    "done": False,  # Continue if tools were called
-                }
+                return ModelResponse(content=output, tool_calls=tool_calls, done=False)
             else:
                 # No tools called, just text response
-                return {
-                    "content": output,
-                    "tool_calls": [],
-                    "done": True,
-                }
+                return ModelResponse(content=output, tool_calls=[], done=True)
 
         except Exception as e:
             logger.error("Agent execution failed: %s", e)
-            return {
-                "content": f"Error: {e!s}",
-                "tool_calls": [],
-                "done": True,
-            }
+            return ModelResponse(content=f"Error: {e!s}", tool_calls=[], done=True)
 
     async def format_tool_results(
-        self, processed_results: dict[str, Any], tool_calls: list[dict]
+        self, tool_calls: list[MCPToolCall], tool_results: list[MCPToolResult]
     ) -> list[BaseMessage]:
         """Format tool results into LangChain messages."""
         # Create an AI message with the tool calls and results
         messages = []
 
         # First add an AI message indicating tools were called
-        tool_names = [tc["name"] for tc in tool_calls]
+        tool_names = [tc.name for tc in tool_calls]
         ai_content = f"I'll use the following tools: {', '.join(tool_names)}"
         messages.append(AIMessage(content=ai_content))
 
-        # Then add a human message with the tool results
-        result_text = processed_results["text"]
-        screenshot = processed_results.get("screenshot")
+        # Build result text from tool results
+        text_parts = []
+        latest_screenshot = None
+
+        for tool_call, result in zip(tool_calls, tool_results, strict=False):
+            if result.isError:
+                error_text = "Tool execution failed"
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        error_text = content.text
+                        break
+                text_parts.append(f"Error - {tool_call.name}: {error_text}")
+            else:
+                # Process success content
+                tool_output = []
+                for content in result.content:
+                    if isinstance(content, types.TextContent):
+                        tool_output.append(content.text)
+                    elif isinstance(content, types.ImageContent):
+                        latest_screenshot = content.data
 
-        if screenshot:
+                if tool_output:
+                    text_parts.append(f"{tool_call.name}: " + " ".join(tool_output))
+
+        result_text = "\n".join(text_parts) if text_parts else "No output from tools"
+
+        # Then add a human message with the tool results
+        if latest_screenshot:
             # Include screenshot in multimodal format
             content = [
                 {"type": "text", "text": f"Tool results:\n{result_text}"},
-                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{screenshot}"}},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/png;base64,{latest_screenshot}"},
+                },
             ]
             messages.append(HumanMessage(content=content))
         else: