Merge pull request #112 from hud-evals/l/task-server-cli-improvements

Multi-environment support

Author: lorenss-m

environments/remote_browser/hud.lock.yaml

@@ -7,6 +7,7 @@
 
 import asyncio
 import logging
+from datetime import datetime
 from typing import Dict, Any, Optional
 from hud.server.context import run_context_server
 
@@ -19,14 +20,12 @@ class RemoteBrowserContext:
     def __init__(self):
         """Initialize the remote browser context."""
         self.browser_provider = None
-        self.cdp_url: Optional[str] = None
         self.is_initialized = False
         self.provider_config: Optional[Dict[str, Any]] = None
         self.launch_options: Optional[Dict[str, Any]] = None
-        self.provider_name: Optional[str] = None
-        self.instance_id: Optional[str] = None
         self._startup_complete = False
         self.playwright_tool = None  # Store the playwright tool
+        self._telemetry: Optional[Dict[str, Any]] = None  # Store full telemetry data
 
         logger.info("[RemoteBrowserContext] Created new remote browser context")
 
@@ -55,13 +54,8 @@ def set_browser_provider(self, provider) -> None:
             logger.info(f"[RemoteBrowserContext] Set browser provider: {self.provider_name}")
 
     def get_cdp_url(self) -> Optional[str]:
-        """Get the CDP URL."""
-        return self.cdp_url
-
-    def set_cdp_url(self, url: str) -> None:
-        """Set the CDP URL."""
-        self.cdp_url = url
-        logger.info(f"[RemoteBrowserContext] Set CDP URL: {url}")
+        """Get the CDP URL from telemetry."""
+        return self._telemetry.get("cdp_url") if self._telemetry else None
 
     def get_is_initialized(self) -> bool:
         """Check if environment is initialized."""
@@ -99,38 +93,36 @@ def set_playwright_tool(self, tool) -> None:
         self.playwright_tool = tool
         logger.info(f"[RemoteBrowserContext] Set playwright tool")
 
+    def set_telemetry(self, telemetry: Dict[str, Any]) -> None:
+        """Set the full telemetry data."""
+        self._telemetry = telemetry
+        logger.info(f"[RemoteBrowserContext] Set telemetry: {telemetry}")
+
     def get_state_summary(self) -> Dict[str, Any]:
         """Get a summary of the current state."""
         return {
             "is_initialized": self.is_initialized,
             "startup_complete": self._startup_complete,
-            "provider_name": self.provider_name,
-            "has_cdp_url": self.cdp_url is not None,
+            "provider_name": self._telemetry.get("provider") if self._telemetry else None,
+            "has_cdp_url": self.get_cdp_url() is not None,
             "has_browser_provider": self.browser_provider is not None,
             "has_playwright_tool": self.playwright_tool is not None,
         }
 
     def get_telemetry(self) -> Dict[str, Any]:
         """Get telemetry data from the browser provider."""
-        # Return basic telemetry data without async calls
-        # The browser provider status check is skipped to avoid async issues
-
-        # Get live view URL if available
-        live_url = None
-        if self.browser_provider and hasattr(self.browser_provider, "get_live_view_url"):
-            try:
-                live_url = self.browser_provider.get_live_view_url()
-            except Exception as e:
-                logger.warning(f"Failed to get live view URL: {e}")
+        # If we have stored telemetry, return it
+        if self._telemetry:
+            return self._telemetry
 
+        # Otherwise return basic telemetry data
         return {
-            "provider": self.provider_name or "unknown",
-            "status": "running"
-            if self.browser_provider and self.is_initialized
-            else "not_initialized",
-            "live_url": live_url,
-            "cdp_url": self.cdp_url,
-            "instance_id": self.instance_id,
+            "provider": "unknown",
+            "status": "not_initialized",
+            "live_url": None,
+            "cdp_url": None,
+            "instance_id": None,
+            "timestamp": datetime.now().isoformat(),
         }
 
 

environments/remote_browser/src/hud_controller/context.py

@@ -79,7 +79,7 @@ async def get_telemetry_resource() -> Telemetry:
                 status=telemetry["status"],
                 live_url=telemetry["live_url"],
                 timestamp=datetime.now().isoformat(),
-                cdp_url=telemetry["cdp_url"],
+                cdp_url=None,
                 instance_id=telemetry["instance_id"],
             )
         except Exception as e:
@@ -235,8 +235,23 @@ async def send_progress(progress: int, message: str):
 
             # Create browser session
             cdp_url = await browser_provider.launch(**launch_options)
-            persistent_ctx.set_cdp_url(cdp_url)
-            await send_progress(60, f"Browser launched, CDP URL: {cdp_url}")
+
+            # Build and store telemetry data
+            telemetry_data = {
+                "provider": provider_name,
+                "status": "running",
+                "live_url": browser_provider.get_live_view_url()
+                if hasattr(browser_provider, "get_live_view_url")
+                else None,
+                "cdp_url": cdp_url,
+                "instance_id": browser_provider._instance_id
+                if hasattr(browser_provider, "_instance_id")
+                else None,
+                "timestamp": datetime.now().isoformat(),
+            }
+            persistent_ctx.set_telemetry(telemetry_data)
+
+            await send_progress(60, f"Browser launched")
         else:
             # Reuse existing browser session
             await send_progress(20, "Reusing existing browser session...")
@@ -246,7 +261,7 @@ async def send_progress(progress: int, message: str):
             if not cdp_url:
                 raise ValueError("No CDP URL in persistent context")
 
-            await send_progress(60, f"Using existing CDP URL: {cdp_url}")
+            await send_progress(60, f"Using existing CDP URL")
 
         # Initialize PlaywrightToolWithMemory with CDP URL from context
         # This reconnects to the existing browser session on reloads

environments/remote_browser/src/hud_controller/server.py

@@ -20,3 +20,10 @@
     from .version import __version__
 except ImportError:
     __version__ = "unknown"
+
+try:
+    from .utils.pretty_errors import install_pretty_errors
+
+    install_pretty_errors()
+except Exception:  # noqa: S110
+    pass

hud/__init__.py

@@ -111,10 +111,12 @@ def __init__(
         # Initialize these here so methods can be called before initialize()
         self._available_tools: list[types.Tool] = []
         self._tool_map: dict[str, types.Tool] = {}  # Simplified: just name to tool
-        self.screenshot_history: list[str] = []
+        self.response_tool_name = None
+        self.initialization_complete = False
+
+        # Trace
         self._auto_trace = auto_trace
         self._auto_trace_cm: Any | None = None  # Store auto-created trace context manager
-        self.initialization_complete = False
 
         # Response agent to automatically interact with the model
         self.response_agent = response_agent
@@ -530,6 +532,9 @@ async def _filter_tools(self) -> None:
         self._available_tools = []
         self._tool_map = {}
 
+        # Track response tools by server
+        response_tools_by_server: dict[str, str] = {}  # server_name -> tool_name
+
         for tool in all_tools:
             # Check if tool should be included
             if self.allowed_tools and tool.name not in self.allowed_tools:
@@ -541,10 +546,36 @@ async def _filter_tools(self) -> None:
             # Simplified mapping - just tool name to tool
             self._tool_map[tool.name] = tool
 
-            # Auto-detect response tool as a lifecycle tool
-            if tool.name == "response" and "response" not in self.lifecycle_tools:
-                self.design.debug("Auto-detected 'response' tool as a lifecycle tool")
-                self.lifecycle_tools.append("response")
+            # Track response tools
+            if "response" in tool.name or tool.name == "response":
+                # Extract server name from tool name (e.g., "grader_response" -> "grader")
+                if "_" in tool.name:
+                    server_name = tool.name.split("_", 1)[0]
+                    response_tools_by_server[server_name] = tool.name
+                else:
+                    response_tools_by_server["_default"] = tool.name
+
+        # Find the response tool to use (prioritize last server in config)
+        if response_tools_by_server and hasattr(self.mcp_client, "mcp_config"):
+            # Get server names in order from mcp_config
+            server_names = list(self.mcp_client.mcp_config.keys())
+
+            # Try to find response tool from last server first
+            response_tool_name = None
+            for server_name in reversed(server_names):
+                if server_name in response_tools_by_server:
+                    response_tool_name = response_tools_by_server[server_name]
+                    break
+
+            # Fallback to any response tool
+            if not response_tool_name and response_tools_by_server:
+                response_tool_name = next(iter(response_tools_by_server.values()))
+
+            # Add to lifecycle tools if found
+            if response_tool_name and response_tool_name not in self.lifecycle_tools:
+                self.design.debug(f"Auto-detected '{response_tool_name}' tool as a lifecycle tool")
+                self.response_tool_name = response_tool_name
+                self.lifecycle_tools.append(response_tool_name)
 
         # Check if all required tools are available
         if self.required_tools:
@@ -565,13 +596,12 @@ async def _maybe_submit_response(self, response: AgentResponse, messages: list[A
             response: The agent's response
             messages: The current message history (will be modified in-place)
         """
-        # Check if we have a response lifecycle tool
-        if "response" in self.lifecycle_tools and "response" in self._tool_map:
-            self.design.debug("Calling response lifecycle tool")
+        if self.response_tool_name:
+            self.design.debug(f"Calling response lifecycle tool: {self.response_tool_name}")
             try:
                 # Call the response tool with the agent's response
                 response_tool_call = MCPToolCall(
-                    name="response", arguments={"response": response.content, "messages": messages}
+                    name=self.response_tool_name, arguments={"response": response.content}
                 )
                 response_results = await self.call_tools(response_tool_call)
 

hud/agents/base.py

@@ -306,19 +306,20 @@ def _convert_tools_for_claude(self) -> list[dict]:
         """Convert MCP tools to Claude tool format."""
         claude_tools = []
         self._claude_to_mcp_tool_map = {}  # Reset mapping
-        
+
         # Find computer tool by priority
         computer_tool_priority = ["anthropic_computer", "computer_anthropic", "computer"]
         selected_computer_tool = None
-        
+
         for priority_name in computer_tool_priority:
             for tool in self._available_tools:
-                if tool.name == priority_name:
+                # Check both exact match and suffix match (for prefixed tools)
+                if tool.name == priority_name or tool.name.endswith(f"_{priority_name}"):
                     selected_computer_tool = tool
                     break
             if selected_computer_tool:
                 break
-        
+
         # Add the selected computer tool if found
         if selected_computer_tool:
             claude_tool = {
@@ -330,14 +331,18 @@ def _convert_tools_for_claude(self) -> list[dict]:
             # Map Claude's "computer" back to the actual MCP tool name
             self._claude_to_mcp_tool_map["computer"] = selected_computer_tool.name
             claude_tools.append(claude_tool)
-            logger.debug(f"Using {selected_computer_tool.name} as computer tool for Claude")
-        
+            logger.debug("Using %s as computer tool for Claude", selected_computer_tool.name)
+
         # Add other non-computer tools
         for tool in self._available_tools:
             # Skip computer tools (already handled) and lifecycle tools
-            if tool.name in computer_tool_priority or tool.name in self.lifecycle_tools:
+            is_computer_tool = any(
+                tool.name == priority_name or tool.name.endswith(f"_{priority_name}")
+                for priority_name in computer_tool_priority
+            )
+            if is_computer_tool or tool.name in self.lifecycle_tools:
                 continue
-                
+
             claude_tool = {
                 "name": tool.name,
                 "description": tool.description or f"Execute {tool.name}",

hud/agents/claude.py

@@ -33,29 +33,6 @@ def mock_mcp_use_client(self):
         with patch("mcp_use.client.MCPClient.from_dict", return_value=mock_instance):
             yield mock_instance
 
-    @pytest.mark.asyncio
-    async def test_init_with_config(self, mock_telemetry):
-        """Test client initialization with config dictionary."""
-        mcp_config = {
-            "test_server": {
-                "command": "python",
-                "args": ["-m", "test_server"],
-                "env": {"TEST": "true"},
-            }
-        }
-
-        with patch("mcp_use.client.MCPClient.from_dict") as mock_from_dict:
-            mock_instance = MagicMock()
-            mock_instance.create_all_sessions = AsyncMock(return_value={})
-            mock_from_dict.return_value = mock_instance
-            client = MCPClient(mcp_config=mcp_config, verbose=True)
-            # Initialize to trigger connection
-            await client.initialize()
-
-            assert client.verbose is True
-            # Verify MCPUseClient.from_dict was called with proper config
-            mock_from_dict.assert_called_once_with({"mcpServers": mcp_config})
-
     @pytest.mark.asyncio
     async def test_connect_single_server(self, mock_telemetry, mock_mcp_use_client):
         """Test connecting to a single server."""
@@ -146,10 +123,10 @@ async def mock_list_tools2():
         # Verify sessions were created
         mock_mcp_use_client.create_all_sessions.assert_called_once()
 
-        # Check tools from both servers
+        # Check tools from both servers - should be prefixed with server names
         tools = await client.list_tools()
         names = {t.name for t in tools}
-        assert names == {"tool1", "tool2"}
+        assert names == {"server1_tool1", "server2_tool2"}
 
     @pytest.mark.asyncio
     async def test_call_tool(self, mock_telemetry, mock_mcp_use_client):
@@ -220,8 +197,10 @@ async def mock_list_tools():
 
         await client.initialize()
 
-        with pytest.raises(ValueError, match="Tool 'nonexistent' not found"):
-            await client.call_tool(name="nonexistent", arguments={})
+        # Calling a non-existent tool should return an error result
+        result = await client.call_tool(name="nonexistent", arguments={})
+        assert result.isError is True
+        assert "Tool 'nonexistent' not found" in result.content[0].text
 
     @pytest.mark.asyncio
     async def test_get_telemetry_data(self, mock_telemetry, mock_mcp_use_client):