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

hud/utils/mcp.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from hud.settings import settings
+
+
+class MCPConfigPatch(BaseModel):
+    """Patch for MCP config."""
+
+    headers: dict[str, Any] | None = Field(default_factory=dict, alias="headers")
+    meta: dict[str, Any] | None = Field(default_factory=dict, alias="meta")
+
+
+def patch_mcp_config(mcp_config: dict[str, dict[str, Any]], patch: MCPConfigPatch) -> None:
+    """Patch MCP config with additional values."""
+    hud_mcp_url = settings.hud_mcp_url
+
+    for server_cfg in mcp_config.values():
+        url = server_cfg.get("url", "")
+
+        # 1) HTTP header lane (only for hud MCP servers)
+        if hud_mcp_url in url and patch.headers:
+            for key, value in patch.headers.items():
+                headers = server_cfg.setdefault("headers", {})
+                headers.setdefault(key, value)
+
+        # 2) Metadata lane (for all servers)
+        if patch.meta:
+            for key, value in patch.meta.items():
+                meta = server_cfg.setdefault("meta", {})
+                meta.setdefault(key, value)
+
+
+def setup_hud_telemetry(mcp_config: dict[str, dict[str, Any]], auto_trace: bool = True) -> None:
+    """Setup telemetry for hud servers."""
+    if not mcp_config:
+        raise ValueError("Please run initialize() before setting up client-side telemetry")
+
+    from hud.otel import get_current_task_run_id
+    from hud.telemetry import trace
+
+    run_id = get_current_task_run_id()
+    if not run_id and auto_trace:
+        auto_trace_cm = trace("My Trace")
+        run_id = auto_trace_cm.__enter__()
+
+    # Patch HUD servers with run-id (works whether auto or user trace)
+    if run_id:
+        patch_mcp_config(
+            mcp_config,
+            MCPConfigPatch(headers={"Run-Id": run_id}, meta={"run_id": run_id}),
+        )

hud/utils/progress.py

@@ -1,149 +1,149 @@
-from __future__ import annotations
-
-import time
-from collections import defaultdict
-
-
-class StepProgressTracker:
-    """
-    Tracks progress across potentially parallel async tasks based on steps completed.
-    Provides estimates assuming tasks run up to max_steps_per_task.
-    """
-
-    def __init__(self, total_tasks: int, max_steps_per_task: int) -> None:
-        """
-        Initialize the StepProgressTracker.
-
-        Args:
-            total_tasks: The total number of tasks to track.
-            max_steps_per_task: The maximum number of steps per task.
-
-        Raises:
-            ValueError: If total_tasks or max_steps_per_task is not positive.
-        """
-        if total_tasks <= 0:
-            raise ValueError("total_tasks must be positive")
-        if max_steps_per_task <= 0:
-            raise ValueError("max_steps_per_task must be positive")
-
-        self.total_tasks = total_tasks
-        self.max_steps_per_task = max_steps_per_task
-        self.total_potential_steps = total_tasks * max_steps_per_task
-
-        # Use asyncio.Lock for potentially concurrent updates/reads if needed,
-        # but start without for simplicity in single-threaded asyncio.
-        # self._lock = asyncio.Lock()
-        self._task_steps: dict[str, int] = defaultdict(int)
-        self._finished_tasks: dict[str, bool] = defaultdict(bool)
-        self._tasks_started = 0
-        self._tasks_finished = 0
-
-        self.start_time: float | None = None
-        self.current_total_steps = 0
-
-    def start_task(self, task_id: str) -> None:
-        # async with self._lock: # If using lock
-        if self.start_time is None:
-            self.start_time = time.monotonic()
-        self._task_steps[task_id] = 0
-        self._finished_tasks[task_id] = False
-        self._tasks_started += 1
-
-    def increment_step(self, task_id: str) -> None:
-        # async with self._lock:
-        if (
-            not self._finished_tasks[task_id]
-            and self._task_steps[task_id] < self.max_steps_per_task
-        ):
-            self._task_steps[task_id] += 1
-            # Update overall progress immediately
-            self._update_total_steps()
-
-    def finish_task(self, task_id: str) -> None:
-        # async with self._lock:
-        if not self._finished_tasks[task_id]:
-            # For calculation, consider a finished task as having completed max steps
-            self._task_steps[task_id] = self.max_steps_per_task
-            self._finished_tasks[task_id] = True
-            self._tasks_finished += 1
-            # Update overall progress
-            self._update_total_steps()
-
-    def _update_total_steps(self) -> None:
-        # This could be expensive if called extremely frequently.
-        # Called after increment or finish.
-        # async with self._lock:
-        self.current_total_steps = sum(self._task_steps.values())
-
-    def get_progress(self) -> tuple[int, int, float]:
-        """Returns (current_steps, total_potential_steps, percentage)."""
-        # async with self._lock:
-        # Recalculate here for safety, though _update_total_steps should keep it current
-        # current_steps = sum(self._task_steps.values())
-        current_steps = self.current_total_steps
-
-        percentage = 0.0
-        if self.total_potential_steps > 0:
-            percentage = (current_steps / self.total_potential_steps) * 100
-        return current_steps, self.total_potential_steps, percentage
-
-    def get_stats(self) -> tuple[float, float | None]:
-        """Returns (rate_steps_per_minute, eta_seconds_upper_bound)."""
-        # async with self._lock:
-        if self.start_time is None or self._tasks_started == 0:
-            return 0.0, None  # No rate or ETA yet
-
-        elapsed_time = time.monotonic() - self.start_time
-        current_steps = self.current_total_steps
-
-        rate_sec = 0.0
-        if elapsed_time > 0:
-            rate_sec = current_steps / elapsed_time
-
-        rate_min = rate_sec * 60  # Convert rate to steps per minute
-
-        eta = None
-        # ETA calculation still uses rate_sec (steps/second) for time estimation in seconds
-        if rate_sec > 0:
-            remaining_steps = self.total_potential_steps - current_steps
-            eta = remaining_steps / rate_sec if remaining_steps > 0 else 0.0
-
-        return rate_min, eta  # Return rate in steps/min
-
-    def is_finished(self) -> bool:
-        # async with self._lock:
-        return self._tasks_finished >= self.total_tasks
-
-    def display(self, bar_length: int = 40) -> str:
-        """Generates a progress string similar to tqdm."""
-        current_steps, total_steps, percentage = self.get_progress()
-        rate_min, eta = self.get_stats()  # Rate is now per minute
-
-        # Ensure valid values for display
-        current_steps = min(current_steps, total_steps)
-        percentage = max(0.0, min(100.0, percentage))
-
-        filled_length = int(bar_length * current_steps // total_steps) if total_steps else 0
-        bar = "█" * filled_length + "-" * (bar_length - filled_length)
-
-        # Format time
-        elapsed_str = "0:00"
-        eta_str = "??:??"
-        if self.start_time:
-            elapsed_seconds = int(time.monotonic() - self.start_time)
-            elapsed_str = f"{elapsed_seconds // 60}:{elapsed_seconds % 60:02d}"
-            if eta is not None:
-                eta_seconds = int(eta)
-                eta_str = f"{eta_seconds // 60}:{eta_seconds % 60:02d}"
-            elif self.is_finished():
-                eta_str = "0:00"
-
-        # Update rate string format
-        rate_str = f"{rate_min:.1f} steps/min" if rate_min > 0 else "?? steps/min"
-
-        # Format steps - use K/M for large numbers if desired, keep simple for now
-        steps_str = f"{current_steps}/{total_steps}"
-
-        # tasks_str = f" {self._tasks_finished}/{self.total_tasks} tasks" # Optional tasks counter
-
-        return f"{percentage:3.0f}%|{bar}| {steps_str} [{elapsed_str}<{eta_str}, {rate_str}]"
+from __future__ import annotations
+
+import time
+from collections import defaultdict
+
+
+class StepProgressTracker:
+    """
+    Tracks progress across potentially parallel async tasks based on steps completed.
+    Provides estimates assuming tasks run up to max_steps_per_task.
+    """
+
+    def __init__(self, total_tasks: int, max_steps_per_task: int) -> None:
+        """
+        Initialize the StepProgressTracker.
+
+        Args:
+            total_tasks: The total number of tasks to track.
+            max_steps_per_task: The maximum number of steps per task.
+
+        Raises:
+            ValueError: If total_tasks or max_steps_per_task is not positive.
+        """
+        if total_tasks <= 0:
+            raise ValueError("total_tasks must be positive")
+        if max_steps_per_task <= 0:
+            raise ValueError("max_steps_per_task must be positive")
+
+        self.total_tasks = total_tasks
+        self.max_steps_per_task = max_steps_per_task
+        self.total_potential_steps = total_tasks * max_steps_per_task
+
+        # Use asyncio.Lock for potentially concurrent updates/reads if needed,
+        # but start without for simplicity in single-threaded asyncio.
+        # self._lock = asyncio.Lock()
+        self._task_steps: dict[str, int] = defaultdict(int)
+        self._finished_tasks: dict[str, bool] = defaultdict(bool)
+        self._tasks_started = 0
+        self._tasks_finished = 0
+
+        self.start_time: float | None = None
+        self.current_total_steps = 0
+
+    def start_task(self, task_id: str) -> None:
+        # async with self._lock: # If using lock
+        if self.start_time is None:
+            self.start_time = time.monotonic()
+        self._task_steps[task_id] = 0
+        self._finished_tasks[task_id] = False
+        self._tasks_started += 1
+
+    def increment_step(self, task_id: str) -> None:
+        # async with self._lock:
+        if (
+            not self._finished_tasks[task_id]
+            and self._task_steps[task_id] < self.max_steps_per_task
+        ):
+            self._task_steps[task_id] += 1
+            # Update overall progress immediately
+            self._update_total_steps()
+
+    def finish_task(self, task_id: str) -> None:
+        # async with self._lock:
+        if not self._finished_tasks[task_id]:
+            # For calculation, consider a finished task as having completed max steps
+            self._task_steps[task_id] = self.max_steps_per_task
+            self._finished_tasks[task_id] = True
+            self._tasks_finished += 1
+            # Update overall progress
+            self._update_total_steps()
+
+    def _update_total_steps(self) -> None:
+        # This could be expensive if called extremely frequently.
+        # Called after increment or finish.
+        # async with self._lock:
+        self.current_total_steps = sum(self._task_steps.values())
+
+    def get_progress(self) -> tuple[int, int, float]:
+        """Returns (current_steps, total_potential_steps, percentage)."""
+        # async with self._lock:
+        # Recalculate here for safety, though _update_total_steps should keep it current
+        # current_steps = sum(self._task_steps.values())
+        current_steps = self.current_total_steps
+
+        percentage = 0.0
+        if self.total_potential_steps > 0:
+            percentage = (current_steps / self.total_potential_steps) * 100
+        return current_steps, self.total_potential_steps, percentage
+
+    def get_stats(self) -> tuple[float, float | None]:
+        """Returns (rate_steps_per_minute, eta_seconds_upper_bound)."""
+        # async with self._lock:
+        if self.start_time is None or self._tasks_started == 0:
+            return 0.0, None  # No rate or ETA yet
+
+        elapsed_time = time.monotonic() - self.start_time
+        current_steps = self.current_total_steps
+
+        rate_sec = 0.0
+        if elapsed_time > 0:
+            rate_sec = current_steps / elapsed_time
+
+        rate_min = rate_sec * 60  # Convert rate to steps per minute
+
+        eta = None
+        # ETA calculation still uses rate_sec (steps/second) for time estimation in seconds
+        if rate_sec > 0:
+            remaining_steps = self.total_potential_steps - current_steps
+            eta = remaining_steps / rate_sec if remaining_steps > 0 else 0.0
+
+        return rate_min, eta  # Return rate in steps/min
+
+    def is_finished(self) -> bool:
+        # async with self._lock:
+        return self._tasks_finished >= self.total_tasks
+
+    def display(self, bar_length: int = 40) -> str:
+        """Generates a progress string similar to tqdm."""
+        current_steps, total_steps, percentage = self.get_progress()
+        rate_min, eta = self.get_stats()  # Rate is now per minute
+
+        # Ensure valid values for display
+        current_steps = min(current_steps, total_steps)
+        percentage = max(0.0, min(100.0, percentage))
+
+        filled_length = int(bar_length * current_steps // total_steps) if total_steps else 0
+        bar = "█" * filled_length + "-" * (bar_length - filled_length)
+
+        # Format time
+        elapsed_str = "0:00"
+        eta_str = "??:??"
+        if self.start_time:
+            elapsed_seconds = int(time.monotonic() - self.start_time)
+            elapsed_str = f"{elapsed_seconds // 60}:{elapsed_seconds % 60:02d}"
+            if eta is not None:
+                eta_seconds = int(eta)
+                eta_str = f"{eta_seconds // 60}:{eta_seconds % 60:02d}"
+            elif self.is_finished():
+                eta_str = "0:00"
+
+        # Update rate string format
+        rate_str = f"{rate_min:.1f} steps/min" if rate_min > 0 else "?? steps/min"
+
+        # Format steps - use K/M for large numbers if desired, keep simple for now
+        steps_str = f"{current_steps}/{total_steps}"
+
+        # tasks_str = f" {self._tasks_finished}/{self.total_tasks} tasks" # Optional tasks counter
+
+        return f"{percentage:3.0f}%|{bar}| {steps_str} [{elapsed_str}<{eta_str}, {rate_str}]"

hud/utils/telemetry.py

@@ -1,66 +1,66 @@
-from __future__ import annotations
-
-import logging
-
-logger = logging.getLogger(__name__)
-
-
-def stream(live_url: str) -> str:
-    """
-    Display a stream in the HUD system.
-    """
-    from IPython.display import HTML, display
-
-    html_content = f"""
-    <div style="width: 960px; height: 540px; overflow: hidden;">
-        <div style="transform: scale(0.5); transform-origin: top left;">
-            <iframe src="{live_url}" width="1920" height="1080" style="border: 1px solid #ddd;">
-            </iframe>
-        </div>
-    </div>
-    """
-    try:
-        display(HTML(html_content))
-    except Exception as e:
-        logger.warning(e)
-
-    return html_content
-
-
-def display_screenshot(base64_image: str, width: int = 960, height: int = 540) -> str:
-    """
-    Display a base64-encoded screenshot image.
-
-    Args:
-        base64_image: Base64-encoded image string (without the data URI prefix)
-        width: Display width in pixels
-        height: Display height in pixels
-
-    Returns:
-        The HTML string used to display the image
-
-    Note:
-        This function will both display the image in IPython environments
-        and return the HTML string for other contexts.
-    """
-    from IPython.display import HTML, display
-
-    # Ensure the base64 image doesn't already have the data URI prefix
-    if base64_image.startswith("data:image"):
-        img_src = base64_image
-    else:
-        img_src = f"data:image/png;base64,{base64_image}"
-
-    html_content = f"""
-    <div style="width: {width}px; height: {height}px; overflow: hidden; margin: 10px 0; border: 1px solid #ddd;">
-        <img src="{img_src}" style="max-width: 100%; max-height: 100%;">
-    </div>
-    """  # noqa: E501
-
-    # Display in IPython environments
-    try:
-        display(HTML(html_content))
-    except Exception as e:
-        logger.warning(e)
-
-    return html_content
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def stream(live_url: str) -> str:
+    """
+    Display a stream in the HUD system.
+    """
+    from IPython.display import HTML, display
+
+    html_content = f"""
+    <div style="width: 960px; height: 540px; overflow: hidden;">
+        <div style="transform: scale(0.5); transform-origin: top left;">
+            <iframe src="{live_url}" width="1920" height="1080" style="border: 1px solid #ddd;">
+            </iframe>
+        </div>
+    </div>
+    """
+    try:
+        display(HTML(html_content))
+    except Exception as e:
+        logger.warning(e)
+
+    return html_content
+
+
+def display_screenshot(base64_image: str, width: int = 960, height: int = 540) -> str:
+    """
+    Display a base64-encoded screenshot image.
+
+    Args:
+        base64_image: Base64-encoded image string (without the data URI prefix)
+        width: Display width in pixels
+        height: Display height in pixels
+
+    Returns:
+        The HTML string used to display the image
+
+    Note:
+        This function will both display the image in IPython environments
+        and return the HTML string for other contexts.
+    """
+    from IPython.display import HTML, display
+
+    # Ensure the base64 image doesn't already have the data URI prefix
+    if base64_image.startswith("data:image"):
+        img_src = base64_image
+    else:
+        img_src = f"data:image/png;base64,{base64_image}"
+
+    html_content = f"""
+    <div style="width: {width}px; height: {height}px; overflow: hidden; margin: 10px 0; border: 1px solid #ddd;">
+        <img src="{img_src}" style="max-width: 100%; max-height: 100%;">
+    </div>
+    """  # noqa: E501
+
+    # Display in IPython environments
+    try:
+        display(HTML(html_content))
+    except Exception as e:
+        logger.warning(e)
+
+    return html_content

hud/utils/tests/test_async_utils.py

@@ -0,0 +1,173 @@
+"""Tests for async utilities."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from unittest.mock import patch
+
+import pytest
+
+from hud.utils.async_utils import fire_and_forget
+
+
+class TestFireAndForget:
+    """Test fire_and_forget function."""
+
+    @pytest.mark.asyncio
+    async def test_fire_and_forget_with_running_loop(self, caplog):
+        """Test fire_and_forget when event loop is already running."""
+        # Create a simple coroutine that sets a flag
+        flag = []
+
+        async def test_coro():
+            flag.append(True)
+
+        # Call fire_and_forget in async context
+        fire_and_forget(test_coro(), description="test task")
+
+        # Give it a moment to execute
+        await asyncio.sleep(0.1)
+
+        # Check that the coroutine was executed
+        assert flag == [True]
+
+    @pytest.mark.asyncio
+    async def test_fire_and_forget_with_exception(self, caplog):
+        """Test fire_and_forget handles exceptions gracefully."""
+
+        async def failing_coro():
+            raise ValueError("Test exception")
+
+        # This should not raise
+        fire_and_forget(failing_coro(), description="failing task")
+
+        # Give it a moment to execute
+        await asyncio.sleep(0.1)
+
+        # The exception should be handled silently
+
+    def test_fire_and_forget_no_event_loop(self):
+        """Test fire_and_forget when no event loop is running."""
+        # This test runs in sync context
+        flag = threading.Event()
+
+        async def test_coro():
+            flag.set()
+
+        # Call fire_and_forget in sync context
+        fire_and_forget(test_coro(), description="sync test")
+
+        # Wait for the thread to complete
+        assert flag.wait(timeout=2.0), "Coroutine did not execute in thread"
+
+    def test_fire_and_forget_thread_exception(self, caplog):
+        """Test fire_and_forget handles thread exceptions."""
+
+        async def failing_coro():
+            raise ValueError("Thread exception")
+
+        # Patch the logger to capture the debug call
+        from unittest.mock import patch
+
+        with patch("hud.utils.async_utils.logger") as mock_logger:
+            fire_and_forget(failing_coro(), description="thread fail")
+
+            # Give thread time to execute and log
+            import time
+
+            time.sleep(0.5)  # Wait for thread to complete
+
+            # Check that error was logged with correct format
+            mock_logger.debug.assert_called()
+            # Get the actual call arguments
+            calls = mock_logger.debug.call_args_list
+            assert any(
+                call[0][0] == "Error in threaded %s: %s"
+                and call[0][1] == "thread fail"
+                and "Thread exception" in str(call[0][2])
+                for call in calls
+            ), f"Expected log message not found in calls: {calls}"
+
+    def test_fire_and_forget_interpreter_shutdown(self, caplog):
+        """Test fire_and_forget handles interpreter shutdown gracefully."""
+
+        async def test_coro():
+            pass
+
+        # Mock the scenario where we get interpreter shutdown error
+        with patch("asyncio.get_running_loop") as mock_get_loop:
+            mock_get_loop.side_effect = RuntimeError("no running event loop")
+
+            with patch("threading.Thread") as mock_thread:
+                mock_thread.side_effect = RuntimeError(
+                    "cannot schedule new futures after interpreter shutdown"
+                )
+
+                with caplog.at_level(logging.DEBUG):
+                    # This should not raise or log
+                    fire_and_forget(test_coro(), description="shutdown test")
+
+                    # No error should be logged for interpreter shutdown
+                    assert not any(
+                        "Could not shutdown test" in record.message for record in caplog.records
+                    )
+
+    def test_fire_and_forget_other_thread_error(self, caplog):
+        """Test fire_and_forget logs non-shutdown thread errors."""
+
+        async def test_coro():
+            pass
+
+        # Mock the scenario where we get a different error
+        with patch("asyncio.get_running_loop") as mock_get_loop:
+            mock_get_loop.side_effect = RuntimeError("no running event loop")
+
+            with patch("threading.Thread") as mock_thread:
+                mock_thread.side_effect = RuntimeError("Some other error")
+
+                # Patch the logger to capture the debug call
+                with patch("hud.utils.async_utils.logger") as mock_logger:
+                    fire_and_forget(test_coro(), description="error test")
+
+                    # Check that error was logged with correct format
+                    mock_logger.debug.assert_called_once_with(
+                        "Could not %s - no event loop available: %s",
+                        "error test",
+                        mock_thread.side_effect,
+                    )
+
+    @pytest.mark.asyncio
+    async def test_fire_and_forget_cancelled_task(self):
+        """Test fire_and_forget handles cancelled tasks."""
+
+        cancel_event = asyncio.Event()
+
+        async def long_running_coro():
+            await cancel_event.wait()
+
+        # Get the current loop
+        loop = asyncio.get_running_loop()
+
+        # Patch create_task to capture the task
+        created_task = None
+        original_create_task = loop.create_task
+
+        def mock_create_task(coro):
+            nonlocal created_task
+            created_task = original_create_task(coro)
+            return created_task
+
+        with patch.object(loop, "create_task", side_effect=mock_create_task):
+            fire_and_forget(long_running_coro(), description="cancel test")
+
+        # Give it a moment to start
+        await asyncio.sleep(0.01)
+
+        # Cancel the task
+        assert created_task is not None
+        created_task.cancel()
+
+        # This should not raise any exceptions
+        await asyncio.sleep(0.01)