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

hud/otel/tests/test_processors.py

@@ -1,197 +1,197 @@
-"""Tests for OpenTelemetry processors."""
-
-from __future__ import annotations
-
-from unittest.mock import MagicMock, patch
-
-from hud.otel.processors import HudEnrichmentProcessor
-
-
-class TestHudEnrichmentProcessor:
-    """Test HudEnrichmentProcessor."""
-
-    def test_on_start_with_run_id(self):
-        """Test on_start with current task run ID."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Mock span
-        span = MagicMock()
-        span.set_attribute = MagicMock()
-        span.is_recording.return_value = True
-
-        # Mock baggage to return run ID
-        parent_context = {}
-        with patch("hud.otel.processors.baggage.get_baggage") as mock_get_baggage:
-            # Return run ID for task_run_id, None for job_id
-            mock_get_baggage.side_effect = (
-                lambda key, context: "test-run-123" if key == "hud.task_run_id" else None
-            )
-            processor.on_start(span, parent_context)
-
-        # Verify attribute was set
-        span.set_attribute.assert_called_with("hud.task_run_id", "test-run-123")
-
-    def test_on_start_no_run_id(self):
-        """Test on_start without current task run ID."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Mock span
-        span = MagicMock()
-        span.set_attribute = MagicMock()
-        span.is_recording.return_value = True
-        span.name = "test_span"
-
-        # Set up attributes to return None (not matching any step type)
-        span.attributes = {}
-
-        # Mock baggage to return None
-        parent_context = {}
-        with patch("hud.otel.processors.baggage.get_baggage", return_value=None):
-            processor.on_start(span, parent_context)
-
-        # Verify only step count attributes were set (no run_id or job_id)
-        calls = span.set_attribute.call_args_list
-        set_attrs = {call[0][0] for call in calls}
-
-        # Should have step counts but not run_id/job_id
-        assert "hud.task_run_id" not in set_attrs
-        assert "hud.job_id" not in set_attrs
-        assert "hud.base_mcp_steps" in set_attrs
-        assert "hud.mcp_tool_steps" in set_attrs
-        assert "hud.agent_steps" in set_attrs
-
-    def test_on_end(self):
-        """Test on_end does nothing."""
-
-        processor = HudEnrichmentProcessor()
-        span = MagicMock()
-
-        # Should not raise
-        processor.on_end(span)
-
-    def test_shutdown(self):
-        """Test shutdown does nothing."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Should not raise
-        processor.shutdown()
-
-    def test_force_flush(self):
-        """Test force_flush returns True."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Should return True
-        result = processor.force_flush()
-        assert result is True
-
-    def test_on_start_with_job_id(self):
-        """Test on_start with job ID in baggage."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Mock span
-        span = MagicMock()
-        span.set_attribute = MagicMock()
-        span.is_recording.return_value = True
-
-        # Mock baggage with job ID
-        parent_context = {}
-        with patch("hud.otel.processors.baggage.get_baggage") as mock_get_baggage:
-            # Return None for task_run_id, job-123 for job_id
-            mock_get_baggage.side_effect = (
-                lambda key, context: "job-123" if key == "hud.job_id" else None
-            )
-            processor.on_start(span, parent_context)
-
-        # Verify job ID attribute was set
-        span.set_attribute.assert_called_with("hud.job_id", "job-123")
-
-    def test_on_start_exception_handling(self):
-        """Test on_start handles exceptions gracefully."""
-
-        processor = HudEnrichmentProcessor()
-
-        # Mock span that raises exception
-        span = MagicMock()
-        span.is_recording.side_effect = Exception("Test error")
-
-        # Should not raise
-        processor.on_start(span, parent_context=None)
-
-    def test_on_start_exception_handling_extended(self):
-        """Test that exceptions in on_start are caught and logged."""
-        from hud.otel.processors import HudEnrichmentProcessor
-
-        processor = HudEnrichmentProcessor()
-
-        # Create a mock span that raises when setting attributes
-        mock_span = MagicMock()
-        mock_span.is_recording.return_value = True
-        mock_span.set_attribute.side_effect = RuntimeError("Attribute error")
-
-        parent_context = {}
-
-        # Patch logger and baggage to force an exception when setting attribute
-        with (
-            patch("hud.otel.processors.logger") as mock_logger,
-            patch("hud.otel.processors.baggage.get_baggage", return_value="test-id"),
-        ):
-            # Should not raise, exception should be caught
-            processor.on_start(mock_span, parent_context)
-
-            # Verify logger.debug was called with the exception
-            mock_logger.debug.assert_called_once()
-            args = mock_logger.debug.call_args[0]
-            assert "HudEnrichmentProcessor.on_start error" in args[0]
-            assert "Attribute error" in str(args[1])
-
-    def test_on_start_with_baggage_get_exception(self):
-        """Test exception handling when baggage.get_baggage fails for task_run_id."""
-        processor = HudEnrichmentProcessor()
-
-        mock_span = MagicMock()
-        mock_span.is_recording.return_value = True
-
-        parent_context = {}
-
-        # Make baggage.get_baggage raise an exception for task_run_id
-        with (
-            patch(
-                "hud.otel.processors.baggage.get_baggage",
-                side_effect=ValueError("Context error"),
-            ),
-            patch("hud.otel.processors.logger") as mock_logger,
-        ):
-            # Should not raise
-            processor.on_start(mock_span, parent_context)
-
-            # Verify logger.debug was called
-            mock_logger.debug.assert_called_once()
-            args = mock_logger.debug.call_args[0]
-            assert "Context error" in str(args[1])
-
-    def test_on_start_with_baggage_exception(self):
-        """Test exception handling when baggage.get_baggage fails."""
-        processor = HudEnrichmentProcessor()
-
-        mock_span = MagicMock()
-        mock_span.is_recording.return_value = True
-
-        parent_context = {}
-
-        # Make baggage.get_baggage raise an exception
-        with (
-            patch("hud.otel.processors.baggage.get_baggage", side_effect=KeyError("Baggage error")),
-            patch("hud.otel.processors.logger") as mock_logger,
-        ):
-            # Should not raise
-            processor.on_start(mock_span, parent_context)
-
-            # Verify logger.debug was called
-            mock_logger.debug.assert_called_once()
-            args = mock_logger.debug.call_args[0]
-            assert "Baggage error" in str(args[1])
+"""Tests for OpenTelemetry processors."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+from hud.otel.processors import HudEnrichmentProcessor
+
+
+class TestHudEnrichmentProcessor:
+    """Test HudEnrichmentProcessor."""
+
+    def test_on_start_with_run_id(self):
+        """Test on_start with current task run ID."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Mock span
+        span = MagicMock()
+        span.set_attribute = MagicMock()
+        span.is_recording.return_value = True
+
+        # Mock baggage to return run ID
+        parent_context = {}
+        with patch("hud.otel.processors.baggage.get_baggage") as mock_get_baggage:
+            # Return run ID for task_run_id, None for job_id
+            mock_get_baggage.side_effect = (
+                lambda key, context: "test-run-123" if key == "hud.task_run_id" else None
+            )
+            processor.on_start(span, parent_context)
+
+        # Verify attribute was set
+        span.set_attribute.assert_called_with("hud.task_run_id", "test-run-123")
+
+    def test_on_start_no_run_id(self):
+        """Test on_start without current task run ID."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Mock span
+        span = MagicMock()
+        span.set_attribute = MagicMock()
+        span.is_recording.return_value = True
+        span.name = "test_span"
+
+        # Set up attributes to return None (not matching any step type)
+        span.attributes = {}
+
+        # Mock baggage to return None
+        parent_context = {}
+        with patch("hud.otel.processors.baggage.get_baggage", return_value=None):
+            processor.on_start(span, parent_context)
+
+        # Verify only step count attributes were set (no run_id or job_id)
+        calls = span.set_attribute.call_args_list
+        set_attrs = {call[0][0] for call in calls}
+
+        # Should have step counts but not run_id/job_id
+        assert "hud.task_run_id" not in set_attrs
+        assert "hud.job_id" not in set_attrs
+        assert "hud.base_mcp_steps" in set_attrs
+        assert "hud.mcp_tool_steps" in set_attrs
+        assert "hud.agent_steps" in set_attrs
+
+    def test_on_end(self):
+        """Test on_end does nothing."""
+
+        processor = HudEnrichmentProcessor()
+        span = MagicMock()
+
+        # Should not raise
+        processor.on_end(span)
+
+    def test_shutdown(self):
+        """Test shutdown does nothing."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Should not raise
+        processor.shutdown()
+
+    def test_force_flush(self):
+        """Test force_flush returns True."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Should return True
+        result = processor.force_flush()
+        assert result is True
+
+    def test_on_start_with_job_id(self):
+        """Test on_start with job ID in baggage."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Mock span
+        span = MagicMock()
+        span.set_attribute = MagicMock()
+        span.is_recording.return_value = True
+
+        # Mock baggage with job ID
+        parent_context = {}
+        with patch("hud.otel.processors.baggage.get_baggage") as mock_get_baggage:
+            # Return None for task_run_id, job-123 for job_id
+            mock_get_baggage.side_effect = (
+                lambda key, context: "job-123" if key == "hud.job_id" else None
+            )
+            processor.on_start(span, parent_context)
+
+        # Verify job ID attribute was set
+        span.set_attribute.assert_called_with("hud.job_id", "job-123")
+
+    def test_on_start_exception_handling(self):
+        """Test on_start handles exceptions gracefully."""
+
+        processor = HudEnrichmentProcessor()
+
+        # Mock span that raises exception
+        span = MagicMock()
+        span.is_recording.side_effect = Exception("Test error")
+
+        # Should not raise
+        processor.on_start(span, parent_context=None)
+
+    def test_on_start_exception_handling_extended(self):
+        """Test that exceptions in on_start are caught and logged."""
+        from hud.otel.processors import HudEnrichmentProcessor
+
+        processor = HudEnrichmentProcessor()
+
+        # Create a mock span that raises when setting attributes
+        mock_span = MagicMock()
+        mock_span.is_recording.return_value = True
+        mock_span.set_attribute.side_effect = RuntimeError("Attribute error")
+
+        parent_context = {}
+
+        # Patch logger and baggage to force an exception when setting attribute
+        with (
+            patch("hud.otel.processors.logger") as mock_logger,
+            patch("hud.otel.processors.baggage.get_baggage", return_value="test-id"),
+        ):
+            # Should not raise, exception should be caught
+            processor.on_start(mock_span, parent_context)
+
+            # Verify logger.debug was called with the exception
+            mock_logger.debug.assert_called_once()
+            args = mock_logger.debug.call_args[0]
+            assert "HudEnrichmentProcessor.on_start error" in args[0]
+            assert "Attribute error" in str(args[1])
+
+    def test_on_start_with_baggage_get_exception(self):
+        """Test exception handling when baggage.get_baggage fails for task_run_id."""
+        processor = HudEnrichmentProcessor()
+
+        mock_span = MagicMock()
+        mock_span.is_recording.return_value = True
+
+        parent_context = {}
+
+        # Make baggage.get_baggage raise an exception for task_run_id
+        with (
+            patch(
+                "hud.otel.processors.baggage.get_baggage",
+                side_effect=ValueError("Context error"),
+            ),
+            patch("hud.otel.processors.logger") as mock_logger,
+        ):
+            # Should not raise
+            processor.on_start(mock_span, parent_context)
+
+            # Verify logger.debug was called
+            mock_logger.debug.assert_called_once()
+            args = mock_logger.debug.call_args[0]
+            assert "Context error" in str(args[1])
+
+    def test_on_start_with_baggage_exception(self):
+        """Test exception handling when baggage.get_baggage fails."""
+        processor = HudEnrichmentProcessor()
+
+        mock_span = MagicMock()
+        mock_span.is_recording.return_value = True
+
+        parent_context = {}
+
+        # Make baggage.get_baggage raise an exception
+        with (
+            patch("hud.otel.processors.baggage.get_baggage", side_effect=KeyError("Baggage error")),
+            patch("hud.otel.processors.logger") as mock_logger,
+        ):
+            # Should not raise
+            processor.on_start(mock_span, parent_context)
+
+            # Verify logger.debug was called
+            mock_logger.debug.assert_called_once()
+            args = mock_logger.debug.call_args[0]
+            assert "Baggage error" in str(args[1])

hud/server/__init__.py

@@ -1,5 +1,5 @@
-from __future__ import annotations
-
-from .server import MCPServer
-
-__all__ = ["MCPServer"]
+from __future__ import annotations
+
+from .server import MCPServer
+
+__all__ = ["MCPServer"]

hud/server/context.py

@@ -1,114 +1,114 @@
-"""
-HUD context helpers for persistent state across hot-reloads.
-
-Provides utilities for creating shared context servers that survive
-code reloads during development.
-
-Usage in your environment:
-    # In your context_server.py:
-    from hud.server.context import serve_context
-
-    class MyContext:
-        def __init__(self):
-            self.state = {}
-        def startup(self):
-            # Initialize resources
-            pass
-
-    if __name__ == "__main__":
-        serve_context(MyContext())
-
-    # In your MCP server:
-    from hud.server.context import attach_context
-    ctx = attach_context()  # Gets the persistent context
-"""
-
-from __future__ import annotations
-
-import asyncio
-import logging
-import os
-from multiprocessing.managers import BaseManager
-from typing import Any
-
-logger = logging.getLogger(__name__)
-# Default Unix socket path (can be overridden with HUD_CTX_SOCK)
-DEFAULT_SOCK_PATH = "/tmp/hud_ctx.sock"  # noqa: S108
-
-
-def serve_context(
-    context_instance: Any, sock_path: str | None = None, authkey: bytes = b"hud-context"
-) -> BaseManager:
-    """
-    Serve a context object via multiprocessing Manager.
-
-    Args:
-        context_instance: The context object to serve
-        sock_path: Unix socket path (defaults to HUD_CTX_SOCK env var or /tmp/hud_ctx.sock)
-        authkey: Authentication key for the manager
-
-    Returns:
-        The manager instance (can be used to shutdown)
-    """
-    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
-
-    class ContextManager(BaseManager):
-        pass
-
-    ContextManager.register("get_context", callable=lambda: context_instance)
-
-    manager = ContextManager(address=sock_path, authkey=authkey)
-    manager.start()
-
-    return manager
-
-
-def attach_context(sock_path: str | None = None, authkey: bytes = b"hud-context") -> Any:
-    """
-    Attach to a running context server.
-
-    Args:
-        sock_path: Unix socket path (defaults to HUD_CTX_SOCK env var or /tmp/hud_ctx.sock)
-        authkey: Authentication key for the manager
-
-    Returns:
-        The shared context object
-    """
-    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
-
-    class ContextManager(BaseManager):
-        pass
-
-    ContextManager.register("get_context")
-
-    manager = ContextManager(address=sock_path, authkey=authkey)
-    manager.connect()
-
-    return manager.get_context()  # type: ignore
-
-
-async def run_context_server(
-    context_instance: Any, sock_path: str | None = None, authkey: bytes = b"hud-context"
-) -> None:
-    """
-    Run a context server until interrupted.
-
-    Args:
-        context_instance: The context object to serve
-        sock_path: Unix socket path
-        authkey: Authentication key
-    """
-    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
-
-    logger.info("[Context Server] Starting on %s...", sock_path)
-
-    # Start the manager
-    manager = serve_context(context_instance, sock_path, authkey)
-    logger.info("[Context Server] Ready on %s", sock_path)
-
-    # Wait forever (until killed)
-    try:
-        await asyncio.Event().wait()
-    except KeyboardInterrupt:
-        logger.info("[Context Server] Shutting down...")
-        manager.shutdown()
+"""
+HUD context helpers for persistent state across hot-reloads.
+
+Provides utilities for creating shared context servers that survive
+code reloads during development.
+
+Usage in your environment:
+    # In your context_server.py:
+    from hud.server.context import serve_context
+
+    class MyContext:
+        def __init__(self):
+            self.state = {}
+        def startup(self):
+            # Initialize resources
+            pass
+
+    if __name__ == "__main__":
+        serve_context(MyContext())
+
+    # In your MCP server:
+    from hud.server.context import attach_context
+    ctx = attach_context()  # Gets the persistent context
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from multiprocessing.managers import BaseManager
+from typing import Any
+
+logger = logging.getLogger(__name__)
+# Default Unix socket path (can be overridden with HUD_CTX_SOCK)
+DEFAULT_SOCK_PATH = "/tmp/hud_ctx.sock"  # noqa: S108
+
+
+def serve_context(
+    context_instance: Any, sock_path: str | None = None, authkey: bytes = b"hud-context"
+) -> BaseManager:
+    """
+    Serve a context object via multiprocessing Manager.
+
+    Args:
+        context_instance: The context object to serve
+        sock_path: Unix socket path (defaults to HUD_CTX_SOCK env var or /tmp/hud_ctx.sock)
+        authkey: Authentication key for the manager
+
+    Returns:
+        The manager instance (can be used to shutdown)
+    """
+    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
+
+    class ContextManager(BaseManager):
+        pass
+
+    ContextManager.register("get_context", callable=lambda: context_instance)
+
+    manager = ContextManager(address=sock_path, authkey=authkey)
+    manager.start()
+
+    return manager
+
+
+def attach_context(sock_path: str | None = None, authkey: bytes = b"hud-context") -> Any:
+    """
+    Attach to a running context server.
+
+    Args:
+        sock_path: Unix socket path (defaults to HUD_CTX_SOCK env var or /tmp/hud_ctx.sock)
+        authkey: Authentication key for the manager
+
+    Returns:
+        The shared context object
+    """
+    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
+
+    class ContextManager(BaseManager):
+        pass
+
+    ContextManager.register("get_context")
+
+    manager = ContextManager(address=sock_path, authkey=authkey)
+    manager.connect()
+
+    return manager.get_context()  # type: ignore
+
+
+async def run_context_server(
+    context_instance: Any, sock_path: str | None = None, authkey: bytes = b"hud-context"
+) -> None:
+    """
+    Run a context server until interrupted.
+
+    Args:
+        context_instance: The context object to serve
+        sock_path: Unix socket path
+        authkey: Authentication key
+    """
+    sock_path = sock_path or os.getenv("HUD_CTX_SOCK", DEFAULT_SOCK_PATH)
+
+    logger.info("[Context Server] Starting on %s...", sock_path)
+
+    # Start the manager
+    manager = serve_context(context_instance, sock_path, authkey)
+    logger.info("[Context Server] Ready on %s", sock_path)
+
+    # Wait forever (until killed)
+    try:
+        await asyncio.Event().wait()
+    except KeyboardInterrupt:
+        logger.info("[Context Server] Shutting down...")
+        manager.shutdown()

hud/server/helper/__init__.py

@@ -1,5 +1,5 @@
-"""Helper sub-package: utilities, registration helpers, shims."""
-
-from __future__ import annotations
-
-__all__ = []
+"""Helper sub-package: utilities, registration helpers, shims."""
+
+from __future__ import annotations
+
+__all__ = []