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

hud/telemetry/tests/test_context.py

@@ -1,210 +0,0 @@
-from __future__ import annotations
-
-from unittest.mock import MagicMock
-
-from hud.telemetry.context import (
-    buffer_mcp_call,
-    flush_buffer,
-    get_current_task_run_id,
-    is_root_trace,
-    set_current_task_run_id,
-)
-from hud.telemetry.mcp_models import BaseMCPCall
-
-
-class TestTaskRunIdContext:
-    """Test task run ID context management."""
-
-    def test_get_current_task_run_id_initial(self):
-        """Test getting task run ID when none is set."""
-        # Reset context for clean test
-        set_current_task_run_id(None)
-        result = get_current_task_run_id()
-        assert result is None
-
-    def test_set_and_get_task_run_id(self):
-        """Test setting and getting task run ID."""
-        test_id = "test-task-run-id"
-        set_current_task_run_id(test_id)
-        result = get_current_task_run_id()
-        assert result == test_id
-
-    def test_task_run_id_isolation(self):
-        """Test that task run IDs are isolated per context."""
-        # This test simulates what would happen in different contexts
-        set_current_task_run_id("context-1")
-        assert get_current_task_run_id() == "context-1"
-
-        set_current_task_run_id("context-2")
-        assert get_current_task_run_id() == "context-2"
-
-        # Reset to None
-        set_current_task_run_id(None)
-        assert get_current_task_run_id() is None
-
-
-class TestRootTraceContext:
-    """Test root trace context management."""
-
-    def test_is_root_trace_initial(self):
-        """Test is_root_trace initial state."""
-        # The initial state may vary, so we just test that it returns a boolean
-        result = is_root_trace.get()
-        assert isinstance(result, bool)
-
-    def test_set_root_trace(self):
-        """Test setting root trace state."""
-        is_root_trace.set(True)
-        assert is_root_trace.get() is True
-
-        is_root_trace.set(False)
-        assert is_root_trace.get() is False
-
-
-class TestMCPCallBuffer:
-    """Test MCP call buffer management."""
-
-    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.reset_context()
-        result = flush_buffer()
-        assert result == []
-
-    def test_add_and_flush_mcp_call(self):
-        """Test adding and flushing MCP calls."""
-        self.reset_context()
-
-        # Set active task run ID
-        set_current_task_run_id("test-task")
-
-        # Create mock MCP call with required attributes
-        mock_call = MagicMock(spec=BaseMCPCall)
-        mock_call.model_dump.return_value = {"type": "test", "task_run_id": "test-task"}
-        mock_call.task_run_id = "test-task"
-
-        buffer_mcp_call(mock_call)
-
-        # Flush should return the call and clear buffer
-        result = flush_buffer()
-        assert len(result) == 1
-        assert result[0] == mock_call
-
-        # Buffer should be empty after flush
-        result2 = flush_buffer()
-        assert result2 == []
-
-    def test_add_multiple_mcp_calls(self):
-        """Test adding multiple MCP calls."""
-        self.reset_context()
-
-        # Set active task run ID
-        set_current_task_run_id("test-task")
-
-        # Create multiple mock calls
-        mock_calls = []
-        for i in range(3):
-            mock_call = MagicMock(spec=BaseMCPCall)
-            mock_call.model_dump.return_value = {"type": f"test_{i}", "task_run_id": "test-task"}
-            mock_call.task_run_id = "test-task"
-            mock_calls.append(mock_call)
-            buffer_mcp_call(mock_call)
-
-        # Flush should return all calls
-        result = flush_buffer()
-        assert len(result) == 3
-        assert result == mock_calls
-
-    def test_buffer_isolation_per_task(self):
-        """Test that MCP call buffers contain all calls regardless of task ID."""
-        self.reset_context()
-
-        # Set task run ID 1
-        set_current_task_run_id("task-1")
-        mock_call_1 = MagicMock(spec=BaseMCPCall)
-        mock_call_1.task_run_id = "task-1"
-        mock_call_1.model_dump.return_value = {"type": "test", "task_run_id": "task-1"}
-        buffer_mcp_call(mock_call_1)
-
-        # Set task run ID 2
-        set_current_task_run_id("task-2")
-        mock_call_2 = MagicMock(spec=BaseMCPCall)
-        mock_call_2.task_run_id = "task-2"
-        mock_call_2.model_dump.return_value = {"type": "test", "task_run_id": "task-2"}
-        buffer_mcp_call(mock_call_2)
-
-        # Flush should return all calls from both tasks
-        result = flush_buffer()
-        assert len(result) == 1
-        assert result[0] == mock_call_2
-
-        set_current_task_run_id("task-1")
-        result2 = flush_buffer()
-        assert len(result2) == 1
-        assert result2[0] == mock_call_1
-
-    def test_buffer_mcp_call_without_task_id(self):
-        """Test adding MCP call when no task run ID is set."""
-        self.reset_context()
-        set_current_task_run_id(None)
-
-        mock_call = MagicMock(spec=BaseMCPCall)
-        mock_call.task_run_id = None
-        buffer_mcp_call(mock_call)
-
-        # Should not buffer anything when no task ID is set
-        result = flush_buffer()
-        assert len(result) == 0
-
-
-class TestContextIntegration:
-    """Integration tests for context management."""
-
-    def test_context_lifecycle(self):
-        """Test complete context lifecycle."""
-        # Start with clean state
-        set_current_task_run_id(None)
-        flush_buffer()
-        is_root_trace.set(False)
-
-        # Set up trace context
-        task_id = "integration-test-task"
-        set_current_task_run_id(task_id)
-        is_root_trace.set(True)
-
-        # Add some MCP calls
-        mock_calls = []
-        for i in range(2):
-            mock_call = MagicMock(spec=BaseMCPCall)
-            mock_call.model_dump.return_value = {
-                "type": f"integration_test_{i}",
-                "task_run_id": task_id,
-            }
-            mock_call.task_run_id = task_id
-            mock_calls.append(mock_call)
-            buffer_mcp_call(mock_call)
-
-        # Verify context state
-        assert get_current_task_run_id() == task_id
-        assert is_root_trace.get() is True
-
-        # Flush and verify
-        result = flush_buffer()
-        assert len(result) == 2
-        assert result == mock_calls
-
-        # Clean up
-        set_current_task_run_id(None)
-        is_root_trace.set(False)
-
-        # Verify cleanup
-        assert get_current_task_run_id() is None
-        assert flush_buffer() == []

hud/telemetry/tests/test_trace.py

@@ -1,312 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import uuid
-from unittest.mock import MagicMock
-
-import pytest
-
-from hud.telemetry._trace import (
-    init_telemetry,
-    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 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 before test
-    actual_set_current_task_run_id(None)
-    actual_is_root_trace.set(False)
-    yield
-    # Reset context after test
-    actual_set_current_task_run_id(None)
-    actual_is_root_trace.set(False)
-
-
-class TestInitTelemetry:
-    """Test telemetry initialization."""
-
-    def test_init_telemetry(self, mocker):
-        """Test telemetry initialization calls registry.install_all."""
-        mock_registry = mocker.patch("hud.telemetry._trace.registry", autospec=True)
-        init_telemetry()
-        mock_registry.install_all.assert_called_once()
-
-
-class TestTrace:
-    """Test the trace context manager."""
-
-    def test_trace_basic(self, mocker):
-        """Test basic trace functionality and context setting."""
-        mock_flush = mocker.patch(
-            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
-        )
-        mock_submit_loop = mocker.patch(
-            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
-        )
-
-        initial_root_state = actual_is_root_trace.get()
-
-        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
-            assert actual_is_root_trace.get() is True
-
-        assert actual_get_current_task_run_id() is None
-        assert actual_is_root_trace.get() == initial_root_state
-        mock_flush.assert_called_once()
-        # submit_to_worker_loop is now called for status updates
-        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
-
-    def test_trace_with_name_and_attributes(self, mocker):
-        """Test trace with name and attributes, checking they are passed on."""
-        mock_mcp_calls = [MagicMock()]
-        mock_flush = mocker.patch(
-            "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
-        )
-        mock_submit_loop = mocker.patch(
-            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
-        )
-
-        trace_name = "test_trace_with_data"
-        attrs = {"key": "value", "number": 42}
-
-        with trace_open(name=trace_name, attributes=attrs) as task_run_id:
-            assert isinstance(task_run_id, str)
-
-        mock_flush.assert_called_once()
-        # submit_to_worker_loop is now called for status updates
-        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
-
-    @pytest.mark.asyncio
-    async def test_trace_with_mcp_calls_exports(self, mocker):
-        """Test trace with MCP calls exports telemetry with correct data."""
-        mock_mcp_calls = [MagicMock(), MagicMock()]
-        mock_flush = mocker.patch(
-            "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
-        )
-        mock_submit_loop = mocker.patch(
-            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
-        )
-
-        async def mock_export(*args, **kwargs):
-            return None
-
-        mocker.patch(
-            "hud.telemetry.exporter.export_telemetry",
-            side_effect=mock_export,
-        )
-
-        test_attrs = {"custom_attr": "test_val"}
-        test_name = "mcp_export_test"
-
-        with trace_open(name=test_name, attributes=test_attrs) as task_run_id:
-            pass
-
-        mock_flush.assert_called_once()
-        # submit_to_worker_loop is now called for status updates and export
-        # The exact count may vary depending on whether export_incremental is called
-        assert mock_submit_loop.call_count >= 2  # At least INITIALIZING and COMPLETED
-
-        # With the new export flow, export_telemetry is submitted to worker loop
-        # so we can't directly assert on it being called synchronously
-        # Instead, verify that the trace completed successfully
-        assert task_run_id is not None
-
-    def test_trace_nested(self, mocker):
-        """Test nested traces, verifying context restoration and root trace logic."""
-        actual_set_current_task_run_id(None)
-        actual_is_root_trace.set(False)
-
-        mock_flush_internal = mocker.patch(
-            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
-        )
-        mock_submit_loop_internal = mocker.patch(
-            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
-        )
-
-        assert actual_get_current_task_run_id() is None
-        assert actual_is_root_trace.get() is False
-
-        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:
-                assert actual_get_current_task_run_id() == inner_id
-                assert actual_is_root_trace.get() is False
-                assert outer_id != inner_id
-            assert actual_get_current_task_run_id() == outer_id
-            assert actual_is_root_trace.get() is True
-
-        assert actual_get_current_task_run_id() is None
-        assert actual_is_root_trace.get() is False
-        assert mock_flush_internal.call_count == 2
-        # submit_to_worker_loop is now called for status updates
-        assert mock_submit_loop_internal.call_count == 2  # Only outer trace sends status updates
-
-    def test_trace_exception_handling(self, mocker):
-        """Test trace handles exceptions properly and restores context."""
-        initial_task_id_before_trace = "pre_existing_id_123"
-        initial_root_state_before_trace = True
-        actual_set_current_task_run_id(initial_task_id_before_trace)
-        actual_is_root_trace.set(initial_root_state_before_trace)
-
-        mock_flush = mocker.patch(
-            "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
-        )
-        mock_submit_loop = mocker.patch(
-            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
-        )
-
-        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")
-
-        mock_flush.assert_called_once()
-        assert actual_get_current_task_run_id() == initial_task_id_before_trace
-        assert actual_is_root_trace.get() == initial_root_state_before_trace
-        mock_submit_loop.assert_not_called()
-
-
-class TestTraceSync:
-    """Test the trace_sync context manager."""
-
-    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
-
-        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", agent_model=None, 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", agent_model=None, 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_open.assert_called_once_with(
-            name="test_func_sync", agent_model=None, attributes=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
-
-        @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_open.assert_called_once_with(
-                name="test_func_async", agent_model=None, attributes=None
-            )
-
-        asyncio.run(run_test())
-
-    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": "multiply"}
-
-        @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_open.assert_called_once_with(
-            name="test_func", agent_model=None, attributes=attrs
-        )
-
-    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
-
-        @trace_decorator()
-        def my_function():
-            return "result"
-
-        result = my_function()
-        assert result == "result"
-        # Should use module.function name
-        expected_name = f"{my_function.__module__}.my_function"
-        mock_trace_open.assert_called_once_with(
-            name=expected_name, agent_model=None, attributes=None
-        )
-
-    def test_trace_decorator_preserves_function_metadata(self):
-        """Test trace_decorator preserves original function metadata."""
-
-        @trace_decorator(name="test")
-        def original_function():
-            """Original docstring."""
-
-        assert original_function.__name__ == "original_function"
-        assert original_function.__doc__ == "Original docstring."
-
-    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
-
-        @trace_decorator()
-        def failing_function():
-            raise RuntimeError("Test error")
-
-        with pytest.raises(RuntimeError, match="Test error"):
-            failing_function()
-
-        mock_trace_open.assert_called_once()

hud/tools/helper/README.md

@@ -1,56 +0,0 @@
-# HUD Helper Package
-
-This sub-package bundles utilities that make it trivial to expose HUD
-Python tool classes as **Model Context Protocol (MCP)** tools.
-
-## Contents
-
-| File | Purpose |
-|------|---------|
-| `utils.py` | `register_instance_tool` – wrap a class instance into a FastMCP tool with auto-generated JSON schema |
-| `mcp_server.py` | CLI server (stdio/HTTP). Tool names: `computer`, `computer_anthropic`, `computer_openai`, `bash`, `edit_file` |
-
-## Quick start
-
-### 1 — Run a server (stdio)
-```bash
-python -m hud.tools.helper.mcp_server               # exposes all tools on stdio
-```
-
-### 2 — Run a server (HTTP)
-```bash
-python -m hud.tools.helper.mcp_server http --port 8040 \
-       --tools computer bash   # expose only two tools
-```
-This starts a Streamable-HTTP MCP server at `http://localhost:8040/mcp`.
-
-### 3 — From a client
-```python
-from mcp import ClientSession
-from mcp.client.streamable_http import streamablehttp_client
-
-async with streamablehttp_client("http://localhost:8040/mcp") as (r, w, _):
-    async with ClientSession(r, w) as sess:
-        await sess.initialize()
-        res = await sess.call_tool("bash", {"command": "echo hi"})
-        print(res.content[0].text)
-```
-
-## Advanced: registering custom tools
-
-```python
-from mcp.server.fastmcp import FastMCP
-from hud.tools.helper import register_instance_tool
-
-class MyTool:
-    async def __call__(self, name: str) -> str:   # type-hints generate schema!
-        return f"Hello {name}!"
-
-mcp = FastMCP("Custom")
-register_instance_tool(mcp, "my_tool", MyTool())
-
-mcp.run(transport="stdio")
-```
-
-The helper inspects `MyTool.__call__`, removes `*args/**kwargs`, and FastMCP
-automatically derives an input schema and registers the tool. 

hud/tools/helper/__init__.py

@@ -1,9 +0,0 @@
-from __future__ import annotations
-
-from .server_initialization import mcp_intialize_wrapper
-from .utils import register_instance_tool
-
-__all__ = [
-    "mcp_intialize_wrapper",
-    "register_instance_tool",
-]

hud/tools/helper/mcp_server.py

@@ -1,78 +0,0 @@
-#!/usr/bin/env python3
-"""Parameterised FastMCP server for HUD tools.
-
-Usage
------
-Run with default (stdio, all tools):
-
-    python -m hud.tools.helper.mcp_server
-
-Streamable HTTP on :8040 exposing computer + bash only:
-
-    python -m hud.tools.helper.mcp_server http --tools computer bash
-
-Arguments
-~~~~~~~~~
-transport   stdio (default) | http
---tools     list of tool names to expose (default = all)
---port      HTTP port (default 8040)
-"""
-
-from __future__ import annotations
-
-import argparse
-
-from mcp.server.fastmcp import FastMCP
-
-from hud.tools.bash import BashTool
-from hud.tools.computer.anthropic import AnthropicComputerTool
-from hud.tools.computer.hud import HudComputerTool
-from hud.tools.computer.openai import OpenAIComputerTool
-from hud.tools.edit import EditTool
-
-from .utils import register_instance_tool
-
-TOOL_MAP = {
-    "computer": HudComputerTool,
-    "computer_anthropic": AnthropicComputerTool,
-    "computer_openai": OpenAIComputerTool,
-    "bash": BashTool,
-    "edit_file": EditTool,
-}
-
-
-def build_server(
-    names: list[str] | None = None,
-    *,
-    port: int = 8040,
-    host: str = "0.0.0.0",  # noqa: S104
-) -> FastMCP:
-    server = FastMCP("HUD", port=port, host=host)
-    selected = names or list(TOOL_MAP.keys())
-
-    for name in selected:
-        cls = TOOL_MAP.get(name)
-        if cls is None:
-            raise SystemExit(f"Unknown tool '{name}'. Choices: {list(TOOL_MAP)}")
-        register_instance_tool(server, name, cls())
-    return server
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(prog="hud-mcp", description="Run HUD FastMCP server")
-    parser.add_argument("transport", nargs="?", choices=["stdio", "http"], default="stdio")
-    parser.add_argument("--tools", nargs="*", help="Tool names to expose (default: all)")
-    parser.add_argument("--port", type=int, default=8040, help="HTTP port (default 8040)")
-    parser.add_argument("--host", type=str, default="0.0.0.0", help="HTTP host (default 0.0.0.0)")  # noqa: S104
-    args = parser.parse_args()
-
-    mcp = build_server(args.tools, port=args.port, host=args.host)
-
-    if args.transport == "http":
-        mcp.run(transport="streamable-http")
-    else:
-        mcp.run(transport="stdio")
-
-
-if __name__ == "__main__":
-    main()