hud-python 0.4.51__py3-none-any.whl → 0.4.53__py3-none-any.whl

hud/utils/task_tracking.py

@@ -0,0 +1,223 @@
+"""Task tracking for async telemetry operations.
+
+This module provides infrastructure to track async tasks created during
+telemetry operations (status updates, metric logging) to ensure they
+complete before process shutdown, preventing telemetry loss.
+
+The task tracker maintains strong references to tasks and explicitly cleans
+them up when they complete via callbacks. This ensures tasks are not garbage
+collected before they finish executing.
+
+Thread Safety:
+    Uses threading.Lock (not asyncio.Lock) because done callbacks run
+    synchronously and need to modify the task set safely.
+
+Race Condition Prevention:
+    The wait_all() method uses a multi-pass approach to catch tasks that
+    are created while waiting for existing tasks to complete.
+
+This is an internal module used by async context managers and cleanup
+routines. Users typically don't interact with it directly.
+"""
+
+import asyncio
+import contextlib
+import logging
+import threading
+from collections.abc import Coroutine
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Module exports
+__all__ = ["TaskTracker", "track_task", "wait_all_tasks"]
+
+# Global singleton task tracker
+_global_tracker: "TaskTracker | None" = None
+
+
+class TaskTracker:
+    """Tracks async tasks to ensure completion before shutdown.
+
+    Maintains a set of tasks with thread-safe access for both async code
+    and synchronous callbacks. Tasks are automatically removed when they
+    complete via done callbacks.
+    """
+
+    def __init__(self) -> None:
+        self._tasks: set[asyncio.Task] = set()
+        # Use threading.Lock for synchronous access from done callbacks
+        self._lock = threading.Lock()
+
+    def track_task(self, coro: Coroutine[Any, Any, Any], name: str = "task") -> asyncio.Task | None:
+        """Create and track an async task.
+
+        Args:
+            coro: The coroutine to run
+            name: Descriptive name for debugging and logging
+
+        Returns:
+            The created asyncio.Task, or None if no event loop is available
+        """
+        try:
+            task = asyncio.create_task(coro, name=name)
+
+            # Add task to tracking set (thread-safe)
+            with self._lock:
+                self._tasks.add(task)
+                task_count = len(self._tasks)
+
+            # Setup cleanup callback
+            def cleanup_callback(completed_task: asyncio.Task) -> None:
+                """Remove completed task from tracking set and log failures."""
+                with self._lock:
+                    self._tasks.discard(completed_task)
+
+                # Log exceptions outside lock to avoid blocking
+                with contextlib.suppress(Exception):
+                    if not completed_task.cancelled():
+                        with contextlib.suppress(Exception):
+                            exc = completed_task.exception()
+                            if exc:
+                                logger.warning("Task '%s' failed: %s", name, exc)
+
+            task.add_done_callback(cleanup_callback)
+            logger.debug("Tracking task '%s' (total active: %d)", name, task_count)
+            return task
+
+        except RuntimeError as e:
+            # No event loop - fall back to fire_and_forget
+            logger.warning("Cannot track task '%s': %s", name, e)
+            from hud.utils.async_utils import fire_and_forget
+
+            fire_and_forget(coro, name)
+            return None
+
+    async def wait_all(self, *, timeout_seconds: float = 30.0) -> int:
+        """Wait for all tracked tasks to complete.
+
+        Uses a multi-pass approach to handle race conditions where tasks are
+        added while waiting for existing tasks to complete. This ensures that
+        status updates created near the end of execution are still waited for.
+
+        Args:
+            timeout_seconds: Maximum time to wait in seconds
+
+        Returns:
+            Number of tasks that completed
+        """
+        total_completed = 0
+        time_remaining = timeout_seconds
+        max_passes = 10  # Prevent infinite loops if tasks keep spawning
+
+        for pass_num in range(max_passes):
+            # Get snapshot of pending tasks (thread-safe)
+            with self._lock:
+                pending = [t for t in self._tasks if not t.done()]
+
+            if not pending:
+                if pass_num == 0:
+                    logger.debug("No pending tasks to wait for")
+                else:
+                    logger.debug("All tasks completed after %d passes", pass_num)
+                break
+
+            # Log progress
+            if pass_num == 0:
+                logger.info("Waiting for %d pending tasks...", len(pending))
+            else:
+                logger.debug("Pass %d: Waiting for %d tasks", pass_num + 1, len(pending))
+
+            # Wait for this batch (max 5s per pass to check for new tasks)
+            batch_timeout = min(time_remaining, 5.0) if time_remaining > 0 else 5.0
+            start_time = asyncio.get_event_loop().time()
+
+            try:
+                done, still_pending = await asyncio.wait(
+                    pending, timeout=batch_timeout, return_when=asyncio.ALL_COMPLETED
+                )
+            except Exception as e:
+                logger.error("Error waiting for tasks: %s", e)
+                break
+
+            # Update timing
+            elapsed = asyncio.get_event_loop().time() - start_time
+            time_remaining -= elapsed
+            total_completed += len(done)
+
+            # Handle timeout
+            if still_pending:
+                if time_remaining <= 0:
+                    logger.warning(
+                        "%d tasks still pending after %ss timeout - cancelling",
+                        len(still_pending),
+                        timeout_seconds,
+                    )
+                    for task in still_pending:
+                        task.cancel()
+                    break
+                # Otherwise continue to next pass
+            else:
+                # All tasks from this batch completed, check for new ones
+                with self._lock:
+                    new_pending = [t for t in self._tasks if not t.done()]
+
+                if not new_pending:
+                    # No new tasks were added - we're done
+                    break
+                # Otherwise loop to wait for the new tasks
+
+        if total_completed > 0:
+            logger.info("Completed %d tasks", total_completed)
+
+        return total_completed
+
+    def get_pending_count(self) -> int:
+        """Get number of pending tasks (thread-safe)."""
+        with self._lock:
+            return sum(1 for t in self._tasks if not t.done())
+
+
+def get_global_tracker() -> TaskTracker:
+    """Get or create the global task tracker."""
+    global _global_tracker
+    if _global_tracker is None:
+        _global_tracker = TaskTracker()
+    return _global_tracker
+
+
+def track_task(coro: Coroutine[Any, Any, Any], name: str = "task") -> asyncio.Task | None:
+    """Create and track an async task for telemetry operations.
+
+    This is a convenience function that uses the global tracker to ensure
+    the task completes before shutdown. Used internally by async context
+    managers for status updates and metric logging.
+
+    Args:
+        coro: The coroutine to track
+        name: Descriptive name for debugging
+
+    Returns:
+        The created task, or None if no event loop is available
+    """
+    tracker = get_global_tracker()
+    return tracker.track_task(coro, name)
+
+
+async def wait_all_tasks(*, timeout_seconds: float = 30.0) -> int:
+    """Wait for all tracked telemetry tasks to complete.
+
+    Ensures that all async telemetry operations (status updates, logs)
+    complete before the calling function returns, preventing telemetry loss.
+
+    Uses a multi-pass approach to handle race conditions where status updates
+    are created while waiting for other tasks to complete.
+
+    Args:
+        timeout_seconds: Maximum time to wait for tasks in seconds
+
+    Returns:
+        Number of tasks that completed
+    """
+    tracker = get_global_tracker()
+    return await tracker.wait_all(timeout_seconds=timeout_seconds)

hud/utils/tests/test_agent_factories.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+
+def test_create_openai_agent():
+    from hud.agents.openai_chat_generic import GenericOpenAIChatAgent
+    from hud.utils.agent_factories import create_openai_agent
+
+    agent = create_openai_agent(
+        api_key="test_key", model_name="test_model", completion_kwargs={"temperature": 0.5}
+    )
+    assert isinstance(agent, GenericOpenAIChatAgent)
+    assert agent.model_name == "test_model"
+    assert agent.completion_kwargs["temperature"] == 0.5
+
+
+def test_create_grounded_agent():
+    with (
+        patch("hud.utils.agent_factories.AsyncOpenAI") as mock_async_openai,
+        patch("hud.utils.agent_factories.GrounderConfig"),
+        patch("hud.utils.agent_factories.GroundedOpenAIChatAgent") as mock_agent_class,
+    ):
+        mock_agent = MagicMock()
+        mock_agent_class.return_value = mock_agent
+
+        from hud.utils.agent_factories import create_grounded_agent
+
+        agent = create_grounded_agent(
+            api_key="test_key",
+            grounder_api_key="grounder_key",
+            model_name="test_model",
+        )
+
+        assert agent == mock_agent
+        mock_async_openai.assert_called_with(api_key="test_key", base_url=None)
+        mock_agent_class.assert_called_once()
+
+
+def test_create_grounded_agent_custom_grounder():
+    with (
+        patch("hud.utils.agent_factories.AsyncOpenAI"),
+        patch("hud.utils.agent_factories.GrounderConfig") as mock_grounder_config,
+        patch("hud.utils.agent_factories.GroundedOpenAIChatAgent"),
+    ):
+        from hud.utils.agent_factories import create_grounded_agent
+
+        create_grounded_agent(
+            api_key="test_key",
+            grounder_api_key="grounder_key",
+            model_name="test_model",
+            grounder_api_base="https://custom.api",
+            grounder_model="custom/model",
+        )
+
+        mock_grounder_config.assert_called_with(
+            api_base="https://custom.api",
+            model="custom/model",
+            api_key="grounder_key",
+        )

hud/utils/tests/test_mcp.py

@@ -90,12 +90,10 @@ class TestPatchMCPConfig:
 class TestSetupHUDTelemetry:
     """Tests for setup_hud_telemetry function."""
 
-    def test_empty_config_raises_error(self):
-        """Test that empty config raises ValueError."""
-        with pytest.raises(
-            ValueError, match="Please run initialize\\(\\) before setting up client-side telemetry"
-        ):
-            setup_hud_telemetry({})
+    def test_empty_config_returns_none(self):
+        """Test that empty config returns None (no servers to set up telemetry for)."""
+        result = setup_hud_telemetry({})
+        assert result is None
 
     def test_none_config_raises_error(self):
         """Test that None config raises ValueError."""

hud/utils/tests/test_pretty_errors.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import sys
+from unittest.mock import MagicMock, patch
+
+from hud.utils.pretty_errors import (
+    _async_exception_handler,
+    _render_and_fallback,
+    install_pretty_errors,
+)
+
+
+def test_render_and_fallback_hud_exception():
+    """Test _render_and_fallback with HudException."""
+    from hud.shared.exceptions import HudException
+
+    exc = HudException("Test error")
+
+    with (
+        patch("sys.__excepthook__") as mock_excepthook,
+        patch("hud.utils.pretty_errors.hud_console") as mock_console,
+        patch("sys.stderr.flush"),
+    ):
+        _render_and_fallback(HudException, exc, None)
+
+        mock_excepthook.assert_called_once()
+        mock_console.render_exception.assert_called_once_with(exc)
+
+
+def test_render_and_fallback_non_hud_exception():
+    """Test _render_and_fallback with non-HudException."""
+    exc = ValueError("Test error")
+
+    with (
+        patch("sys.__excepthook__") as mock_excepthook,
+        patch("hud.utils.pretty_errors.hud_console") as mock_console,
+    ):
+        _render_and_fallback(ValueError, exc, None)
+
+        mock_excepthook.assert_called_once()
+        # Should not render for non-HudException
+        mock_console.render_exception.assert_not_called()
+
+
+def test_render_and_fallback_rendering_error():
+    """Test _render_and_fallback handles rendering errors gracefully."""
+    from hud.shared.exceptions import HudException
+
+    exc = HudException("Test error")
+
+    with (
+        patch("sys.__excepthook__") as mock_excepthook,
+        patch("hud.utils.pretty_errors.hud_console") as mock_console,
+    ):
+        mock_console.render_exception.side_effect = Exception("Render failed")
+
+        # Should not raise
+        _render_and_fallback(HudException, exc, None)
+
+        mock_excepthook.assert_called_once()
+
+
+def test_async_exception_handler_with_exception():
+    """Test _async_exception_handler with exception in context."""
+    mock_loop = MagicMock()
+    context = {"exception": ValueError("Test error")}
+
+    with patch("hud.utils.pretty_errors.hud_console") as mock_console:
+        _async_exception_handler(mock_loop, context)
+
+        mock_console.render_exception.assert_called_once()
+        mock_loop.default_exception_handler.assert_called_once_with(context)
+
+
+def test_async_exception_handler_with_message():
+    """Test _async_exception_handler with message only."""
+    mock_loop = MagicMock()
+    context = {"message": "Error message"}
+
+    with patch("hud.utils.pretty_errors.hud_console") as mock_console:
+        _async_exception_handler(mock_loop, context)
+
+        mock_console.error.assert_called_once_with("Error message")
+        mock_console.render_support_hint.assert_called_once()
+        mock_loop.default_exception_handler.assert_called_once()
+
+
+def test_async_exception_handler_rendering_error():
+    """Test _async_exception_handler handles rendering errors."""
+    mock_loop = MagicMock()
+    context = {"exception": ValueError("Test")}
+
+    with patch("hud.utils.pretty_errors.hud_console") as mock_console:
+        mock_console.render_exception.side_effect = Exception("Render failed")
+
+        # Should not raise, should call default handler
+        _async_exception_handler(mock_loop, context)
+
+        mock_loop.default_exception_handler.assert_called_once()
+
+
+def test_install_pretty_errors_with_running_loop():
+    """Test install_pretty_errors with a running event loop."""
+    mock_loop = MagicMock()
+
+    with patch("asyncio.get_running_loop", return_value=mock_loop):
+        install_pretty_errors()
+
+        assert sys.excepthook == _render_and_fallback
+        mock_loop.set_exception_handler.assert_called_once_with(_async_exception_handler)
+
+
+def test_install_pretty_errors_no_running_loop():
+    """Test install_pretty_errors without a running loop."""
+    with (
+        patch("asyncio.get_running_loop", side_effect=RuntimeError("No running loop")),
+        patch("asyncio.new_event_loop") as mock_new_loop,
+    ):
+        mock_loop = MagicMock()
+        mock_new_loop.return_value = mock_loop
+
+        install_pretty_errors()
+
+        assert sys.excepthook == _render_and_fallback
+        mock_loop.set_exception_handler.assert_called_once()
+
+
+def test_install_pretty_errors_new_loop_fails():
+    """Test install_pretty_errors when creating new loop fails."""
+    with (
+        patch("asyncio.get_running_loop", side_effect=RuntimeError("No running loop")),
+        patch("asyncio.new_event_loop", side_effect=Exception("Can't create loop")),
+    ):
+        # Should not raise
+        install_pretty_errors()
+
+        assert sys.excepthook == _render_and_fallback
+
+
+def test_install_pretty_errors_set_handler_fails():
+    """Test install_pretty_errors when set_exception_handler fails."""
+    mock_loop = MagicMock()
+    mock_loop.set_exception_handler.side_effect = Exception("Can't set handler")
+
+    with patch("asyncio.get_running_loop", return_value=mock_loop):
+        # Should not raise
+        install_pretty_errors()
+
+        assert sys.excepthook == _render_and_fallback
+
+
+def test_async_exception_handler_no_exception_or_message():
+    """Test _async_exception_handler with empty context."""
+    mock_loop = MagicMock()
+    context = {}
+
+    with patch("hud.utils.pretty_errors.hud_console") as mock_console:
+        _async_exception_handler(mock_loop, context)
+
+        mock_console.render_exception.assert_not_called()
+        mock_console.error.assert_not_called()
+        mock_loop.default_exception_handler.assert_called_once()
+
+
+def test_render_and_fallback_with_traceback():
+    """Test _render_and_fallback includes traceback."""
+    from hud.shared.exceptions import HudException
+
+    exc = HudException("Test error")
+
+    # Create a fake traceback
+    try:
+        raise exc
+    except HudException as e:
+        tb = e.__traceback__
+
+    with (
+        patch("sys.__excepthook__") as mock_excepthook,
+        patch("hud.utils.pretty_errors.hud_console"),
+        patch("sys.stderr.flush"),
+    ):
+        _render_and_fallback(HudException, exc, tb)
+
+        # Should call excepthook with traceback
+        call_args = mock_excepthook.call_args[0]
+        assert call_args[2] == tb

hud/utils/tests/test_tasks.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from hud.types import Task
+from hud.utils.tasks import load_tasks
+
+
+def test_load_tasks_from_list():
+    """Test loading tasks from a list of dictionaries."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    tasks = load_tasks(task_dicts)
+
+    assert len(tasks) == 2
+    assert all(isinstance(t, Task) for t in tasks)
+    assert tasks[0].prompt == "Test task 1"  # type: ignore
+    assert tasks[1].prompt == "Test task 2"  # type: ignore
+
+
+def test_load_tasks_from_list_raw():
+    """Test loading tasks from a list in raw mode."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    tasks = load_tasks(task_dicts, raw=True)
+
+    assert len(tasks) == 2
+    assert all(isinstance(t, dict) for t in tasks)
+    assert tasks[0]["prompt"] == "Test task 1"  # type: ignore
+
+
+def test_load_tasks_from_json_file():
+    """Test loading tasks from a JSON file."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
+        json.dump(task_dicts, f)
+        temp_path = f.name
+
+    try:
+        tasks = load_tasks(temp_path)
+
+        assert len(tasks) == 2
+        assert all(isinstance(t, Task) for t in tasks)
+        assert tasks[0].prompt == "Test task 1"  # type: ignore
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_from_json_file_raw():
+    """Test loading tasks from a JSON file in raw mode."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
+        json.dump(task_dicts, f)
+        temp_path = f.name
+
+    try:
+        tasks = load_tasks(temp_path, raw=True)
+
+        assert len(tasks) == 2
+        assert all(isinstance(t, dict) for t in tasks)
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_from_jsonl_file():
+    """Test loading tasks from a JSONL file."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".jsonl", delete=False, encoding="utf-8"
+    ) as f:
+        for task_dict in task_dicts:
+            f.write(json.dumps(task_dict) + "\n")
+        temp_path = f.name
+
+    try:
+        tasks = load_tasks(temp_path)
+
+        assert len(tasks) == 2
+        assert all(isinstance(t, Task) for t in tasks)
+        assert tasks[0].prompt == "Test task 1"  # type: ignore
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_from_jsonl_file_with_empty_lines():
+    """Test loading tasks from a JSONL file with empty lines."""
+    task_dicts = [
+        {"id": "1", "prompt": "Test task 1", "mcp_config": {}},
+        {"id": "2", "prompt": "Test task 2", "mcp_config": {}},
+    ]
+
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".jsonl", delete=False, encoding="utf-8"
+    ) as f:
+        f.write(json.dumps(task_dicts[0]) + "\n")
+        f.write("\n")  # Empty line
+        f.write(json.dumps(task_dicts[1]) + "\n")
+        temp_path = f.name
+
+    try:
+        tasks = load_tasks(temp_path)
+
+        assert len(tasks) == 2
+        assert all(isinstance(t, Task) for t in tasks)
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_from_jsonl_file_with_list():
+    """Test loading tasks from a JSONL file where a line contains a list."""
+    task_dict = {"id": "1", "prompt": "Test task 1", "mcp_config": {}}
+
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".jsonl", delete=False, encoding="utf-8"
+    ) as f:
+        f.write(json.dumps([task_dict, task_dict]) + "\n")
+        temp_path = f.name
+
+    try:
+        tasks = load_tasks(temp_path)
+
+        assert len(tasks) == 2
+        assert all(isinstance(t, Task) for t in tasks)
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_json_not_array_error():
+    """Test that loading from JSON file with non-array raises error."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
+        json.dump({"not": "an array"}, f)
+        temp_path = f.name
+
+    try:
+        with pytest.raises(ValueError, match="JSON file must contain an array"):
+            load_tasks(temp_path)
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_invalid_jsonl_format():
+    """Test that loading from JSONL with invalid format raises error."""
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".jsonl", delete=False, encoding="utf-8"
+    ) as f:
+        f.write(json.dumps("invalid") + "\n")
+        temp_path = f.name
+
+    try:
+        with pytest.raises(ValueError, match="Invalid JSONL format"):
+            load_tasks(temp_path)
+    finally:
+        Path(temp_path).unlink()
+
+
+def test_load_tasks_invalid_input_type():
+    """Test that invalid input type raises TypeError."""
+    with pytest.raises(TypeError, match="tasks_input must be str or list"):
+        load_tasks(123)  # type: ignore
+
+
+def test_load_tasks_nonexistent_file():
+    """Test that loading from nonexistent file raises error."""
+    with pytest.raises(ValueError, match="neither a file path nor a HuggingFace dataset"):
+        load_tasks("nonexistent_file_without_slash")