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

hud/telemetry/mcp_models.py

@@ -28,16 +28,14 @@ class StatusType(str, Enum):
 
 
 class MCPCallType(str, Enum):
-    """Known MCP call types"""
+    """Enum for different types of MCP calls in telemetry."""
 
-    SEND_REQUEST = "mcp.shared.session.send_request"
-    SEND_NOTIFICATION = "mcp.shared.session.send_notification"
-    RECEIVE_RESPONSE = "mcp.shared.session.receive_response"
-    RECEIVE_REQUEST = "mcp.shared.session.receive_request"
-    STREAM_READ = "mcp.stream.read"
-    STREAM_WRITE = "mcp.stream.write"
-    HANDLE_INCOMING = "mcp.handle_incoming"
-    MANUAL_TEST = "manual.test"
+    # Requests and Notifications
+    SEND_REQUEST = "mcp.send_request"
+    SEND_NOTIFICATION = "mcp.send_notification"
+
+    # Responses
+    RECEIVE_RESPONSE = "mcp.receive_response"
 
 
 class BaseMCPCall(BaseModel):
@@ -87,6 +85,8 @@ class MCPRequestCall(BaseMCPCall):
     duration: float | None = None
     request_id: str | int | None = None
     request_data: dict[str, Any] | None = None
+    error: str | None = None
+    error_type: str | None = None
 
     @classmethod
     def from_jsonrpc_request(
@@ -191,6 +191,8 @@ class MCPNotificationCall(BaseMCPCall):
     end_time: float | None = None
     duration: float | None = None
     notification_data: dict[str, Any] | None = None
+    error: str | None = None
+    error_type: str | None = None
 
     @classmethod
     def from_jsonrpc_notification(
@@ -230,69 +232,6 @@ class MCPNotificationCall(BaseMCPCall):
         return None
 
 
-class MCPStreamEvent(BaseMCPCall):
-    """Record for an MCP stream event (read or write)"""
-
-    stream_event: bool = True
-    event_type: str = Field(..., description="Type of stream event: read or write")
-    item_type: str | None = None
-    is_response_or_error: bool = False
-    message_data: dict[str, Any] | None = None
-
-    @classmethod
-    def from_session_message(
-        cls, message: SessionMessage, task_run_id: str, event_type: str, **kwargs: Any
-    ) -> MCPStreamEvent:
-        """Create telemetry record for a stream event"""
-        method_name = "unknown_stream_operation"
-        is_response = False
-        item_type = "unknown"
-        message_data = None
-
-        if hasattr(message, "message") and hasattr(message.message, "root"):
-            msg_root = message.message.root
-            item_type = type(msg_root).__name__
-            message_data = msg_root.model_dump(exclude_none=True)
-
-            # Check type first before accessing attributes
-            if isinstance(msg_root, JSONRPCRequest | JSONRPCNotification) and hasattr(
-                msg_root, "method"
-            ):
-                method_name = msg_root.method
-            elif isinstance(msg_root, JSONRPCResponse | JSONRPCError) and hasattr(msg_root, "id"):
-                method_name = f"response_to_id_{msg_root.id}"
-                is_response = True
-
-        return cls(
-            task_run_id=task_run_id,
-            status=StatusType.COMPLETED,
-            method=method_name,
-            event_type=event_type,
-            item_type=item_type,
-            is_response_or_error=is_response,
-            message_data=message_data,
-            timestamp=datetime.now().timestamp(),
-            **kwargs,
-        )
-
-
-class MCPManualTestCall(BaseMCPCall):
-    """Record for a manual test record"""
-
-    call_type: str = MCPCallType.MANUAL_TEST
-    custom_data: dict[str, Any] = Field(default_factory=dict)
-
-    @classmethod
-    def create(cls, task_run_id: str, **custom_data: Any) -> MCPManualTestCall:
-        """Create a manual test record with custom data"""
-        return cls(
-            task_run_id=task_run_id,
-            status=StatusType.COMPLETED,
-            custom_data=custom_data,
-            timestamp=datetime.now().timestamp(),
-        )
-
-
 class MCPTelemetryRecord(BaseModel):
     """Container for a set of related MCP telemetry records"""
 
@@ -320,9 +259,9 @@ class MCPTelemetryRecord(BaseModel):
 
 
 class TrajectoryStep(BaseModel):
-    """Model representing a single step in a trajectory, for export."""
+    """Model for telemetry export format."""
 
-    type: str = Field(default="mcp-step")  # Default for MCP calls
+    type: str = Field(default="mcp-step")
     observation_url: str | None = None
     observation_text: str | None = None
     actions: list[dict[str, Any]] = Field(default_factory=list)

hud/telemetry/tests/test_context.py

@@ -64,21 +64,24 @@ class TestRootTraceContext:
 class TestMCPCallBuffer:
     """Test MCP call buffer management."""
 
-    def setUp(self):
+    def reset_context(self):
         """Clear buffer before each test."""
         # Flush any existing calls and reset context
+        set_current_task_run_id(None)
+        # Clear any existing buffers by setting a temporary task ID and flushing
+        set_current_task_run_id("temp-cleanup")
         flush_buffer()
         set_current_task_run_id(None)
 
     def test_flush_buffer_empty(self):
         """Test flushing empty buffer."""
-        self.setUp()
+        self.reset_context()
         result = flush_buffer()
         assert result == []
 
     def test_add_and_flush_mcp_call(self):
         """Test adding and flushing MCP calls."""
-        self.setUp()
+        self.reset_context()
 
         # Set active task run ID
         set_current_task_run_id("test-task")
@@ -101,7 +104,7 @@ class TestMCPCallBuffer:
 
     def test_add_multiple_mcp_calls(self):
         """Test adding multiple MCP calls."""
-        self.setUp()
+        self.reset_context()
 
         # Set active task run ID
         set_current_task_run_id("test-task")
@@ -122,7 +125,7 @@ class TestMCPCallBuffer:
 
     def test_buffer_isolation_per_task(self):
         """Test that MCP call buffers contain all calls regardless of task ID."""
-        self.setUp()
+        self.reset_context()
 
         # Set task run ID 1
         set_current_task_run_id("task-1")
@@ -150,7 +153,7 @@ class TestMCPCallBuffer:
 
     def test_buffer_mcp_call_without_task_id(self):
         """Test adding MCP call when no task run ID is set."""
-        self.setUp()
+        self.reset_context()
         set_current_task_run_id(None)
 
         mock_call = MagicMock(spec=BaseMCPCall)

hud/telemetry/tests/test_trace.py

@@ -8,21 +8,25 @@ import pytest
 
 from hud.telemetry._trace import (
     init_telemetry,
-    register_trace,
     trace,
+    trace_decorator,
+    trace_open,
 )
 from hud.telemetry.context import get_current_task_run_id as actual_get_current_task_run_id
 from hud.telemetry.context import is_root_trace as actual_is_root_trace
-from hud.telemetry.context import reset_context
 from hud.telemetry.context import set_current_task_run_id as actual_set_current_task_run_id
 
 
 @pytest.fixture(autouse=True)
 def reset_telemetry_context_fixture():
     """Ensures telemetry context is reset before and after each test in this file."""
-    reset_context()
+    # Reset context before test
+    actual_set_current_task_run_id(None)
+    actual_is_root_trace.set(False)
     yield
-    reset_context()
+    # Reset context after test
+    actual_set_current_task_run_id(None)
+    actual_is_root_trace.set(False)
 
 
 class TestInitTelemetry:
@@ -49,7 +53,7 @@ class TestTrace:
 
         initial_root_state = actual_is_root_trace.get()
 
-        with trace() as task_run_id:
+        with trace_open() as task_run_id:
             assert isinstance(task_run_id, str)
             uuid.UUID(task_run_id)
             assert actual_get_current_task_run_id() == task_run_id
@@ -73,7 +77,7 @@ class TestTrace:
         trace_name = "test_trace_with_data"
         attrs = {"key": "value", "number": 42}
 
-        with trace(name=trace_name, attributes=attrs) as task_run_id:
+        with trace_open(name=trace_name, attributes=attrs) as task_run_id:
             assert isinstance(task_run_id, str)
 
         mock_flush.assert_called_once()
@@ -101,7 +105,7 @@ class TestTrace:
         test_attrs = {"custom_attr": "test_val"}
         test_name = "mcp_export_test"
 
-        with trace(name=test_name, attributes=test_attrs) as task_run_id:
+        with trace_open(name=test_name, attributes=test_attrs) as task_run_id:
             pass
 
         mock_flush.assert_called_once()
@@ -113,10 +117,8 @@ class TestTrace:
         assert kwargs["mcp_calls"] == mock_mcp_calls
         assert kwargs["trace_attributes"]["trace_name"] == test_name
         assert kwargs["trace_attributes"]["custom_attr"] == "test_val"
-        assert "start_time" in kwargs["trace_attributes"]
-        assert "end_time" in kwargs["trace_attributes"]
-        assert "duration" in kwargs["trace_attributes"]
-        assert kwargs["trace_attributes"]["is_root"] is True
+        assert "duration_seconds" in kwargs["trace_attributes"]
+        assert kwargs["trace_attributes"]["is_root_trace"] is True
 
     def test_trace_nested(self, mocker):
         """Test nested traces, verifying context restoration and root trace logic."""
@@ -133,7 +135,7 @@ class TestTrace:
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() is False
 
-        with trace(name="outer") as outer_id:
+        with trace_open(name="outer") as outer_id:
             assert actual_get_current_task_run_id() == outer_id
             assert actual_is_root_trace.get() is True
             with trace(name="inner") as inner_id:
@@ -162,7 +164,10 @@ class TestTrace:
             "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
-        with pytest.raises(ValueError, match="Test exception"), trace(name="trace_with_exception"):
+        with (
+            pytest.raises(ValueError, match="Test exception"),
+            trace_open(name="trace_with_exception"),
+        ):
             assert actual_get_current_task_run_id() != initial_task_id_before_trace
             assert actual_is_root_trace.get() is False
             raise ValueError("Test exception")
@@ -173,98 +178,124 @@ class TestTrace:
         mock_submit_loop.assert_not_called()
 
 
-class TestRegisterTrace:
-    """Test the register_trace decorator."""
+class TestTraceSync:
+    """Test the trace_sync context manager."""
 
-    def test_register_trace_sync_function(self, mocker):
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_sync_basic(self, mocker):
+        """Test trace calls trace_open and flush."""
+        mock_flush = mocker.patch("hud.flush", autospec=True)
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open")
+        mock_trace_open.return_value.__enter__.return_value = "test-task-id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace(name="test_func_sync")
+        with trace(name="test_sync") as task_run_id:
+            assert task_run_id == "test-task-id"
+
+        mock_trace_open.assert_called_once_with(name="test_sync", attributes=None)
+        mock_flush.assert_called_once()
+
+    def test_trace_sync_with_attributes(self, mocker):
+        """Test trace passes attributes correctly."""
+        mock_flush = mocker.patch("hud.flush", autospec=True)
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open")
+        mock_trace_open.return_value.__enter__.return_value = "test-task-id"
+        mock_trace_open.return_value.__exit__.return_value = None
+        attrs = {"key": "value"}
+
+        with trace(name="test_sync", attributes=attrs):
+            pass
+
+        mock_trace_open.assert_called_once_with(name="test_sync", attributes=attrs)
+        mock_flush.assert_called_once()
+
+
+class TestTraceDecorator:
+    """Test the trace_decorator function decorator."""
+
+    def test_trace_decorator_sync_function(self, mocker):
+        """Test trace_decorator on synchronous functions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
+
+        @trace_decorator(name="test_func_sync")
         def sync_function(x, y):
             return x + y
 
         result = sync_function(1, 2)
         assert result == 3
-        mock_trace_context_manager.assert_called_once_with(name="test_func_sync", attributes=None)
+        mock_trace_open.assert_called_once_with(name="test_func_sync", attributes=None)
 
-    def test_register_trace_async_function(self, mocker):
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_async_function(self, mocker):
+        """Test trace_decorator on asynchronous functions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace(name="test_func_async")
+        @trace_decorator(name="test_func_async")
         async def async_function(x, y):
             return x + y
 
         async def run_test():
             result = await async_function(1, 2)
             assert result == 3
-            mock_trace_context_manager.assert_called_once_with(
-                name="test_func_async", attributes=None
-            )
+            mock_trace_open.assert_called_once_with(name="test_func_async", attributes=None)
 
         asyncio.run(run_test())
 
-    def test_register_trace_with_attributes(self, mocker):
-        """Test register_trace with attributes."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-
-        class _MockTraceContextManager:
-            def __enter__(self):
-                return "task_id"
-
-            def __exit__(self, exc_type, exc_value, traceback):
-                return None
-
-        mock_trace_context_manager.return_value = _MockTraceContextManager()
+    def test_trace_decorator_with_attributes(self, mocker):
+        """Test trace_decorator with attributes."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        attrs = {"operation": "add"}
+        attrs = {"operation": "multiply"}
 
-        @register_trace(name="test_func", attributes=attrs)
+        @trace_decorator(name="test_func", attributes=attrs)
         def func_with_attrs(x):
             return x * 2
 
         result = func_with_attrs(5)
         assert result == 10
-        mock_trace_context_manager.assert_called_once_with(name="test_func", attributes=attrs)
+        mock_trace_open.assert_called_once_with(name="test_func", attributes=attrs)
 
-    def test_register_trace_without_name(self, mocker):
-        """Test register_trace uses function name when name not provided."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_without_name(self, mocker):
+        """Test trace_decorator uses module.function name when name not provided."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace()
+        @trace_decorator()
         def my_function():
             return "result"
 
         result = my_function()
         assert result == "result"
-        mock_trace_context_manager.assert_called_once_with(name="my_function", attributes=None)
+        # Should use module.function name
+        expected_name = f"{my_function.__module__}.my_function"
+        mock_trace_open.assert_called_once_with(name=expected_name, attributes=None)
 
-    def test_register_trace_preserves_function_metadata(self):
-        """Test register_trace preserves original function metadata."""
+    def test_trace_decorator_preserves_function_metadata(self):
+        """Test trace_decorator preserves original function metadata."""
 
-        @register_trace(name="test")
+        @trace_decorator(name="test")
         def original_function():
             """Original docstring."""
 
         assert original_function.__name__ == "original_function"
         assert original_function.__doc__ == "Original docstring."
 
-    def test_register_trace_exception_propagation(self, mocker):
-        """Test register_trace propagates exceptions."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_exception_propagation(self, mocker):
+        """Test trace_decorator propagates exceptions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace()
+        @trace_decorator()
         def failing_function():
             raise RuntimeError("Test error")
 
         with pytest.raises(RuntimeError, match="Test error"):
             failing_function()
 
-        mock_trace_context_manager.assert_called_once()
+        mock_trace_open.assert_called_once()

hud/tools/__init__.py

@@ -0,0 +1,21 @@
+"""HUD tools for computer control, file editing, and bash commands."""
+
+from __future__ import annotations
+
+from .base import ToolError, ToolResult, tool_result_to_content_blocks
+from .bash import BashTool
+from .computer import AnthropicComputerTool, HudComputerTool, OpenAIComputerTool
+from .edit import EditTool
+from .playwright_tool import PlaywrightTool
+
+__all__ = [
+    "AnthropicComputerTool",
+    "BashTool",
+    "EditTool",
+    "HudComputerTool",
+    "OpenAIComputerTool",
+    "PlaywrightTool",
+    "ToolError",
+    "ToolResult",
+    "tool_result_to_content_blocks",
+]

hud/tools/base.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields, replace
+from typing import Any
+
+from mcp.types import ImageContent, TextContent
+
+
+@dataclass(kw_only=True, frozen=True)
+class ToolResult:
+    """Represents the result of a tool execution."""
+
+    output: str | None = None
+    error: str | None = None
+    base64_image: str | None = None
+    system: str | None = None
+
+    def __bool__(self) -> bool:
+        return any(getattr(self, field.name) for field in fields(self))
+
+    def __add__(self, other: ToolResult) -> ToolResult:
+        def combine_fields(
+            field: str | None, other_field: str | None, concatenate: bool = True
+        ) -> str | None:
+            if field and other_field:
+                if concatenate:
+                    return field + other_field
+                raise ValueError("Cannot combine tool results")
+            return field or other_field
+
+        return ToolResult(
+            output=combine_fields(self.output, other.output),
+            error=combine_fields(self.error, other.error),
+            base64_image=combine_fields(self.base64_image, other.base64_image, False),
+            system=combine_fields(self.system, other.system),
+        )
+
+    def replace(self, **kwargs: Any) -> ToolResult:
+        """Returns a new ToolResult with the given fields replaced."""
+        return replace(self, **kwargs)
+
+
+# Legacy alias for backward compatibility
+CLIResult = ToolResult
+
+
+class ToolError(Exception):
+    """An error raised by a tool."""
+
+
+# Legacy alias for backward compatibility
+CLIError = ToolError
+
+
+def tool_result_to_content_blocks(result: ToolResult) -> list[ImageContent | TextContent]:
+    """Convert a ToolResult to MCP content blocks."""
+    blocks = []
+
+    if result.output:
+        blocks.append(TextContent(text=result.output, type="text"))
+    if result.error:
+        blocks.append(TextContent(text=result.error, type="text"))
+    if result.base64_image:
+        blocks.append(ImageContent(data=result.base64_image, mimeType="image/png", type="image"))
+    return blocks

hud/tools/bash.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+from typing import Any
+
+from .base import CLIResult, ToolError, ToolResult
+
+
+class _BashSession:
+    """A session of a bash shell."""
+
+    _started: bool
+    _process: asyncio.subprocess.Process
+
+    command: str = "/bin/bash"
+    _output_delay: float = 0.2  # seconds
+    _timeout: float = 120.0  # seconds
+    _sentinel: str = "<<exit>>"
+
+    def __init__(self) -> None:
+        self._started = False
+        self._timed_out = False
+
+    async def start(self) -> None:
+        if self._started:
+            await asyncio.sleep(0)
+            return
+
+        # Platform-specific subprocess creation
+        kwargs = {
+            "shell": True,
+            "bufsize": 0,
+            "stdin": asyncio.subprocess.PIPE,
+            "stdout": asyncio.subprocess.PIPE,
+            "stderr": asyncio.subprocess.PIPE,
+        }
+
+        # Only use setsid on Unix-like systems
+        if sys.platform != "win32":
+            kwargs["preexec_fn"] = os.setsid
+
+        self._process = await asyncio.create_subprocess_shell(self.command, **kwargs)
+
+        self._started = True
+
+    def stop(self) -> None:
+        """Terminate the bash shell."""
+        if not self._started:
+            raise ToolError("Session has not started.")
+        if self._process.returncode is not None:
+            return
+        self._process.terminate()
+
+    async def run(self, command: str) -> CLIResult:
+        """Execute a command in the bash shell."""
+        if not self._started:
+            raise ToolError("Session has not started.")
+        if self._process.returncode is not None:
+            await asyncio.sleep(0)
+            return ToolResult(
+                system="tool must be restarted",
+                error=f"bash has exited with returncode {self._process.returncode}",
+            )
+        if self._timed_out:
+            raise ToolError(
+                f"timed out: bash did not return in {self._timeout} seconds and must be restarted",
+            ) from None
+
+        if self._process.stdin is None:
+            raise ToolError("stdin is None")
+        if self._process.stdout is None:
+            raise ToolError("stdout is None")
+        if self._process.stderr is None:
+            raise ToolError("stderr is None")
+
+        # Send command to the process
+        self._process.stdin.write(command.encode() + f"; echo '{self._sentinel}'\n".encode())
+        await self._process.stdin.drain()
+
+        # Read output from the process, until the sentinel is found
+        sentinel_line = f"{self._sentinel}\n"
+        sentinel_bytes = sentinel_line.encode()
+
+        try:
+            raw_out: bytes = await asyncio.wait_for(
+                self._process.stdout.readuntil(sentinel_bytes),
+                timeout=self._timeout,
+            )
+            output = raw_out.decode()[: -len(sentinel_line)]
+        except (TimeoutError, asyncio.LimitOverrunError):
+            self._timed_out = True
+            raise ToolError(
+                f"timed out: bash did not return in {self._timeout} seconds and must be restarted",
+            ) from None
+
+        # Attempt non-blocking stderr fetch (may return empty)
+        try:
+            error_bytes = await asyncio.wait_for(self._process.stderr.read(), timeout=0.01)
+            error = error_bytes.decode().rstrip("\n")
+        except TimeoutError:
+            error = ""
+
+        return CLIResult(output=output, error=error)
+
+
+class BashTool:
+    """
+    A tool that allows the agent to run bash commands.
+    The tool parameters are defined by Anthropic and are not editable.
+    """
+
+    _session: _BashSession | None
+
+    def __init__(self) -> None:
+        self._session = None
+
+    async def __call__(
+        self, command: str | None = None, restart: bool = False, **kwargs: Any
+    ) -> ToolResult:
+        if restart:
+            if self._session:
+                self._session.stop()
+            self._session = _BashSession()
+            await self._session.start()
+
+            return ToolResult(system="tool has been restarted.")
+
+        if self._session is None:
+            self._session = _BashSession()
+            await self._session.start()
+
+        if command is not None:
+            return await self._session.run(command)
+
+        raise ToolError("no command provided.")

hud/tools/computer/__init__.py

@@ -0,0 +1,13 @@
+"""Computer control tools for different agent APIs."""
+
+from __future__ import annotations
+
+from .anthropic import AnthropicComputerTool
+from .hud import HudComputerTool
+from .openai import OpenAIComputerTool
+
+__all__ = [
+    "AnthropicComputerTool",
+    "HudComputerTool",
+    "OpenAIComputerTool",
+]