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

hud/otel/tests/test_processors.py

@@ -0,0 +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])

hud/server/__init__.py

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

hud/server/context.py

@@ -0,0 +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/server/helper/__init__.py

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

hud/server/low_level.py

@@ -0,0 +1,132 @@
+"""Custom low-level MCP server that supports per-server initialization hooks.
+
+This duplicates the upstream `mcp.server.lowlevel.server.Server.run` logic so we
+can inject our own `InitSession` subtype without touching global state.
+"""
+
+from __future__ import annotations
+
+from contextlib import AsyncExitStack
+from typing import TYPE_CHECKING, Any
+
+import anyio
+import mcp.types as types
+from fastmcp.server.low_level import LowLevelServer as _BaseLL
+from mcp.server.lowlevel.server import (
+    logger,
+)
+from mcp.server.session import ServerSession
+from mcp.shared.context import RequestContext
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+    from mcp.server.models import InitializationOptions
+    from mcp.shared.message import SessionMessage
+    from mcp.shared.session import RequestResponder
+
+
+class InitSession(ServerSession):
+    """ServerSession that runs a one-time `init_fn(ctx)` on *initialize*."""
+
+    def __init__(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        init_opts: InitializationOptions,
+        *,
+        init_fn: Callable[[RequestContext], Awaitable[None]] | None = None,
+        stateless: bool = False,
+    ) -> None:
+        super().__init__(read_stream, write_stream, init_opts, stateless=stateless)
+        self._init_fn = init_fn
+        self._did_init = stateless  # skip when running stateless
+
+    # pylint: disable=protected-access  # we need to hook into internal method
+    async def _received_request(
+        self,
+        responder: RequestResponder[types.ClientRequest, types.ServerResult],
+    ) -> types.ServerResult | None:
+        # Intercept initialize
+        if (
+            isinstance(responder.request.root, types.InitializeRequest)
+            and not self._did_init
+            and self._init_fn is not None
+        ):
+            req = responder.request.root
+            ctx = RequestContext[
+                "ServerSession",
+                dict[str, Any],
+                types.InitializeRequest,
+            ](
+                request_id=req.id,  # type: ignore[attr-defined]
+                meta=req.params.meta,
+                session=self,
+                lifespan_context={},
+                request=req,
+            )
+            try:
+                await self._init_fn(ctx)
+            except Exception as exc:
+                token = getattr(req.params.meta, "progressToken", None)
+                if token is not None:
+                    await self.send_progress_notification(
+                        progress_token=token,
+                        progress=0,
+                        total=100,
+                        message=f"Initialization failed: {exc}",
+                    )
+                raise
+            finally:
+                self._did_init = True
+        # fall through to original behaviour
+        return await super()._received_request(responder)
+
+
+class LowLevelServerWithInit(_BaseLL):
+    """LowLevelServer that uses :class:`InitSession` instead of `ServerSession`."""
+
+    def __init__(
+        self,
+        *args: Any,
+        init_fn: Callable[[RequestContext], Awaitable[None]] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self._init_fn = init_fn
+
+    async def run(
+        self,
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception],
+        write_stream: MemoryObjectSendStream[SessionMessage],
+        initialization_options: InitializationOptions,
+        *,
+        raise_exceptions: bool = False,
+        stateless: bool = False,
+    ) -> None:
+        """Copy of upstream run with InitSession injected."""
+
+        async with AsyncExitStack() as stack:
+            lifespan_context = await stack.enter_async_context(self.lifespan(self))
+            session = await stack.enter_async_context(
+                InitSession(
+                    read_stream,
+                    write_stream,
+                    initialization_options,
+                    stateless=stateless,
+                    init_fn=self._init_fn,
+                )
+            )
+
+            async with anyio.create_task_group() as tg:
+                async for message in session.incoming_messages:
+                    logger.debug("Received message: %s", message)
+
+                    tg.start_soon(
+                        self._handle_message,
+                        message,
+                        session,
+                        lifespan_context,
+                        raise_exceptions,
+                    )

hud/server/server.py

@@ -0,0 +1,166 @@
+"""HUD server helpers."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import sys
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+import anyio
+from fastmcp.server.server import FastMCP, Transport
+
+from hud.server.low_level import LowLevelServerWithInit
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable
+
+    from mcp.shared.context import RequestContext
+
+__all__ = ["MCPServer"]
+
+logger = logging.getLogger(__name__)
+
+
+def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
+    """Run *coro_fn* via anyio.run() and cancel on SIGTERM or SIGINT (POSIX)."""
+
+    async def _runner() -> None:
+        stop_evt: asyncio.Event | None = None
+        if sys.platform != "win32" and os.getenv("FASTMCP_DISABLE_SIGTERM_HANDLER") != "1":
+            loop = asyncio.get_running_loop()
+            stop_evt = asyncio.Event()
+
+            # Handle both SIGTERM and SIGINT for graceful shutdown
+            if signal.getsignal(signal.SIGTERM) is signal.SIG_DFL:
+                loop.add_signal_handler(signal.SIGTERM, stop_evt.set)
+            if signal.getsignal(signal.SIGINT) is signal.SIG_DFL:
+                loop.add_signal_handler(signal.SIGINT, stop_evt.set)
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(coro_fn, *args, **kwargs)
+
+            if stop_evt is not None:
+
+                async def _watch() -> None:
+                    logger.info("Waiting for SIGTERM or SIGINT")
+                    if stop_evt is not None:
+                        await stop_evt.wait()
+                    logger.debug("Received shutdown signal, cancelling tasks...")
+                    tg.cancel_scope.cancel()
+
+                tg.start_soon(_watch)
+
+    anyio.run(_runner)
+
+
+class MCPServer(FastMCP):
+    """FastMCP wrapper that adds helpful functionality for dockerized environments.
+    This works with any MCP client, and adds just a few extra server-side features:
+    1. SIGTERM handling for graceful shutdown in container runtimes.
+    2. ``@MCPServer.initialize`` decorator that registers an async initializer
+       executed during the MCP *initialize* request. The initializer function receives
+       a single ``ctx`` parameter (RequestContext) from which you can access:
+       - ``ctx.session``: The MCP ServerSession
+       - ``ctx.meta.progressToken``: Token for progress notifications (if provided)
+       - ``ctx.session.client_params.clientInfo``: Client information
+    3. ``@MCPServer.shutdown`` decorator that registers a coroutine to run during
+       server teardown, after all lifespan contexts have exited.
+    4. Enhanced ``add_tool`` that accepts instances of
+       :class:`hud.tools.base.BaseTool` which are classes that implement the
+       FastMCP ``FunctionTool`` interface.
+    """
+
+    def __init__(self, *, name: str | None = None, **fastmcp_kwargs: Any) -> None:
+        # Store shutdown function placeholder before super().__init__
+        self._shutdown_fn: Callable | None = None
+
+        # Inject custom lifespan if user did not supply one
+        if "lifespan" not in fastmcp_kwargs:
+
+            @asynccontextmanager
+            async def _lifespan(_: Any) -> AsyncGenerator[dict[str, Any], None]:
+                try:
+                    yield {}
+                finally:
+                    if self._shutdown_fn is not None:
+                        await self._shutdown_fn()
+
+            fastmcp_kwargs["lifespan"] = _lifespan
+
+        super().__init__(name=name, **fastmcp_kwargs)
+        self._initializer_fn: Callable | None = None
+        self._did_init = False
+
+        # Replace FastMCP's low-level server with our version that supports
+        # per-server initialization hooks
+        def _run_init(ctx: RequestContext) -> Any:
+            if self._initializer_fn is not None and not self._did_init:
+                self._did_init = True
+                return self._initializer_fn(ctx)
+            return None
+
+        # Save the old server's handlers before replacing it
+        old_request_handlers = self._mcp_server.request_handlers
+        old_notification_handlers = self._mcp_server.notification_handlers
+
+        self._mcp_server = LowLevelServerWithInit(
+            name=self.name,
+            version=self.version,
+            instructions=self.instructions,
+            lifespan=self._mcp_server.lifespan,  # reuse the existing lifespan
+            init_fn=_run_init,
+        )
+
+        # Copy handlers from the old server to the new one
+        self._mcp_server.request_handlers = old_request_handlers
+        self._mcp_server.notification_handlers = old_notification_handlers
+
+    # Initializer decorator: runs on the initialize request
+    # The decorated function receives a RequestContext object with access to:
+    # - ctx.session: The MCP ServerSession
+    # - ctx.meta.progressToken: Progress token (if provided by client)
+    # - ctx.session.client_params.clientInfo: Client information
+    def initialize(self, fn: Callable | None = None) -> Callable | None:
+        def decorator(func: Callable) -> Callable:
+            self._initializer_fn = func
+            return func
+
+        return decorator(fn) if fn else decorator
+
+    # Shutdown decorator: runs after server stops
+    # Supports dockerized SIGTERM handling
+    def shutdown(self, fn: Callable | None = None) -> Callable | None:
+        def decorator(func: Callable) -> Callable:
+            self._shutdown_fn = func
+            return func
+
+        return decorator(fn) if fn else decorator
+
+    # Run with SIGTERM handling and custom initialization
+    def run(
+        self,
+        transport: Transport | None = None,
+        show_banner: bool = True,
+        **transport_kwargs: Any,
+    ) -> None:
+        if transport is None:
+            transport = "stdio"
+
+        async def _bootstrap() -> None:
+            await self.run_async(transport=transport, show_banner=show_banner, **transport_kwargs)  # type: ignore[arg-type]
+
+        _run_with_sigterm(_bootstrap)
+
+    # Tool registration helper -- appends BaseTool to FastMCP
+    def add_tool(self, obj: Any, **kwargs: Any) -> None:
+        from hud.tools.base import BaseTool
+
+        if isinstance(obj, BaseTool):
+            super().add_tool(obj.mcp, **kwargs)
+            return
+
+        super().add_tool(obj, **kwargs)

hud/server/tests/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+__all__ = []