hud-python 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

hud/cli/utils/server.py

@@ -0,0 +1,250 @@
+"""Common server utilities for HUD CLI."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from fastmcp import FastMCP
+
+from hud.utils.design import HUDDesign
+
+from .docker import generate_container_name, remove_container
+
+
+class MCPServerManager:
+    """Manages MCP server lifecycle and configuration."""
+
+    def __init__(self, image: str, docker_args: list[str] | None = None) -> None:
+        """Initialize server manager.
+
+        Args:
+            image: Docker image name
+            docker_args: Additional Docker arguments
+        """
+        self.image = image
+        self.docker_args = docker_args or []
+        self.design = HUDDesign()
+        self.container_name = self._generate_container_name()
+
+    def _generate_container_name(self) -> str:
+        """Generate a unique container name from image."""
+        return generate_container_name(self.image)
+
+    def cleanup_container(self) -> None:
+        """Remove any existing container with the same name."""
+        remove_container(self.container_name)
+
+    def build_docker_command(
+        self,
+        extra_args: list[str] | None = None,
+        entrypoint: list[str] | None = None,
+    ) -> list[str]:
+        """Build Docker run command.
+
+        Args:
+            extra_args: Additional arguments to add before image
+            entrypoint: Custom entrypoint override
+
+        Returns:
+            Complete docker command as list
+        """
+        cmd = [
+            "docker",
+            "run",
+            "--rm",
+            "-i",
+            "--name",
+            self.container_name,
+        ]
+
+        # Add extra args (like volume mounts, env vars)
+        if extra_args:
+            cmd.extend(extra_args)
+
+        # Add user-provided docker args
+        cmd.extend(self.docker_args)
+
+        # Add entrypoint if specified
+        if entrypoint:
+            cmd.extend(["--entrypoint", entrypoint[0]])
+
+        # Add image
+        cmd.append(self.image)
+
+        # Add entrypoint args if specified
+        if entrypoint and len(entrypoint) > 1:
+            cmd.extend(entrypoint[1:])
+
+        return cmd
+
+    def create_mcp_config(self, docker_cmd: list[str]) -> dict[str, Any]:
+        """Create MCP configuration for stdio transport.
+
+        Args:
+            docker_cmd: Docker command to run
+
+        Returns:
+            MCP configuration dict
+        """
+        return {
+            "mcpServers": {
+                "default": {
+                    "command": docker_cmd[0],
+                    "args": docker_cmd[1:] if len(docker_cmd) > 1 else [],
+                    # transport defaults to stdio
+                }
+            }
+        }
+
+    def create_proxy(self, config: dict[str, Any], name: str | None = None) -> FastMCP:
+        """Create FastMCP proxy server.
+
+        Args:
+            config: MCP configuration
+            name: Optional server name
+
+        Returns:
+            FastMCP proxy instance
+        """
+        proxy_name = name or f"HUD Server - {self.image}"
+        return FastMCP.as_proxy(config, name=proxy_name)
+
+    async def run_http_server(
+        self,
+        proxy: FastMCP,
+        port: int,
+        verbose: bool = False,
+        path: str = "/mcp",
+    ) -> None:
+        """Run HTTP server with proper shutdown handling.
+
+        Args:
+            proxy: FastMCP proxy instance
+            port: Port to listen on
+            verbose: Enable verbose logging
+            path: URL path for MCP endpoint
+        """
+        # Set up logging
+        import logging
+        import os
+
+        os.environ["FASTMCP_DISABLE_BANNER"] = "1"
+
+        if not verbose:
+            logging.getLogger("fastmcp").setLevel(logging.ERROR)
+            logging.getLogger("mcp").setLevel(logging.ERROR)
+            logging.getLogger("uvicorn").setLevel(logging.ERROR)
+            logging.getLogger("uvicorn.access").setLevel(logging.ERROR)
+            logging.getLogger("uvicorn.error").setLevel(logging.ERROR)
+
+            import warnings
+
+            warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+        try:
+            await proxy.run_async(
+                transport="http",
+                host="0.0.0.0",  # noqa: S104
+                port=port,
+                path=path,
+                log_level="error" if not verbose else "info",
+                show_banner=False,
+            )
+        except asyncio.CancelledError:
+            pass  # Normal cancellation
+        except Exception as e:
+            if verbose:
+                self.design.error(f"Server error: {e}")
+            raise
+
+
+async def run_server_with_interactive(
+    server_manager: MCPServerManager,
+    port: int,
+    verbose: bool = False,
+) -> None:
+    """Run server with interactive testing mode.
+
+    Args:
+        server_manager: Server manager instance
+        port: Port to listen on
+        verbose: Enable verbose logging
+    """
+    from .interactive import run_interactive_mode
+    from .logging import find_free_port
+
+    design = HUDDesign()
+
+    # Find available port
+    actual_port = find_free_port(port)
+    if actual_port is None:
+        design.error(f"No available ports found starting from {port}")
+        return
+
+    if actual_port != port:
+        design.warning(f"Port {port} in use, using port {actual_port} instead")
+
+    # Clean up any existing container
+    server_manager.cleanup_container()
+
+    # Build docker command
+    docker_cmd = server_manager.build_docker_command()
+
+    # Create MCP config
+    config = server_manager.create_mcp_config(docker_cmd)
+
+    # Create proxy
+    proxy = server_manager.create_proxy(config, f"HUD Interactive - {server_manager.image}")
+
+    # Show header
+    design.info("")  # Empty line
+    design.header("HUD MCP Server - Interactive Mode", icon="🎮")
+
+    # Show configuration
+    design.section_title("Server Information")
+    design.info(f"Image: {server_manager.image}")
+    design.info(f"Port: {actual_port}")
+    design.info(f"URL: http://localhost:{actual_port}/mcp")
+    design.info(f"Container: {server_manager.container_name}")
+    design.info("")
+
+    # Create event to signal server is ready
+    server_ready = asyncio.Event()
+    server_task = None
+
+    async def start_server() -> None:
+        """Start the proxy server."""
+        nonlocal server_task
+        try:
+            # Signal that we're ready before starting
+            server_ready.set()
+            await server_manager.run_http_server(proxy, actual_port, verbose)
+        except asyncio.CancelledError:
+            pass
+
+    try:
+        # Start server in background
+        server_task = asyncio.create_task(start_server())
+
+        # Wait for server to be ready
+        await server_ready.wait()
+        await asyncio.sleep(0.5)  # Give it a moment to fully start
+
+        # Run interactive mode
+        server_url = f"http://localhost:{actual_port}/mcp"
+        await run_interactive_mode(server_url, verbose=verbose)
+
+    except KeyboardInterrupt:
+        design.info("\n👋 Shutting down...")
+    finally:
+        # Cancel server task
+        if server_task and not server_task.done():
+            server_task.cancel()
+            try:
+                await server_task
+            except asyncio.CancelledError:
+                design.error("Server task cancelled")
+
+        # Clean up container
+        server_manager.cleanup_container()

hud/clients/base.py

@@ -173,7 +173,7 @@ class BaseHUDClient(AgentMCPClient):
                 logger.warning("Failed to close auto-created trace: %s", e)
             finally:
                 self._auto_trace_cm = None
-        
+
         # Disconnect from server
         if self._initialized:
             await self._disconnect()

hud/clients/fastmcp.py

@@ -94,29 +94,21 @@ class FastMCPHUDClient(BaseHUDClient):
                         if "mcp.hud.so" in url:
                             raise RuntimeError(
                                 "Authentication failed for HUD API. "
-                                "Please ensure your HUD_API_KEY environment variable is set correctly. "
+                                "Please ensure your HUD_API_KEY environment variable is set correctly."  # noqa: E501
                                 "You can get an API key at https://app.hud.so"
                             ) from e
                     # Generic 401 error
                     raise RuntimeError(
-                        f"Authentication failed (401 Unauthorized). "
-                        f"Please check your credentials or API key."
+                        "Authentication failed (401 Unauthorized). "
+                        "Please check your credentials or API key."
                     ) from e
                 raise
 
             # Configure validation for output schemas based on client setting
             try:
-                from hud_mcp.client.session import ValidationOptions  # type: ignore[import-not-found]
-
-                if (
-                    hasattr(self._client, "_session_state")
-                    and self._client._session_state.session is not None
-                ):
-                    self._client._session_state.session._validation_options = ValidationOptions(
-                        strict_output_validation=self._strict_validation
-                    )
+                if hasattr(self._client, "_session_state") and self._client._session_state.session is not None:  # noqa: E501
+                    self._client._session_state.session._validate_structured_outputs = self._strict_validation  # noqa: E501
             except ImportError:
-                # ValidationOptions may not be available in some mcp versions
                 pass
 
             logger.info("FastMCP client connected")

hud/clients/mcp_use.py

@@ -16,12 +16,12 @@ from .base import BaseHUDClient
 
 if TYPE_CHECKING:
     from mcp import types
-    from mcp_use.client import MCPClient as MCPUseClient
-    from mcp_use.session import MCPSession as MCPUseSession
+    from mcp_use.client import MCPClient as MCPUseClient  # type: ignore[attr-defined]
+    from mcp_use.session import MCPSession as MCPUseSession  # type: ignore[attr-defined]
 
 try:
-    from mcp_use.client import MCPClient as MCPUseClient
-    from mcp_use.session import MCPSession as MCPUseSession
+    from mcp_use.client import MCPClient as MCPUseClient  # type: ignore[attr-defined]
+    from mcp_use.session import MCPSession as MCPUseSession  # type: ignore[attr-defined]
 except ImportError:
     MCPUseClient = None  # type: ignore[misc, assignment]
     MCPUseSession = None  # type: ignore[misc, assignment]
@@ -67,23 +67,19 @@ class MCPUseHUDClient(BaseHUDClient):
             raise ImportError("MCPUseClient is not available")
         self._client = MCPUseClient.from_dict(config)
         try:
-            assert self._client is not None  # For type checker
+            assert self._client is not None  # noqa: S101
             self._sessions = await self._client.create_all_sessions()
             logger.info("Created %d MCP sessions", len(self._sessions))
 
             # Configure validation for all sessions based on client setting
             try:
-                from hud_mcp.client.session import ValidationOptions  # type: ignore[import-not-found]
-
                 for session in self._sessions.values():
                     if (
                         hasattr(session, "connector")
                         and hasattr(session.connector, "client_session")
                         and session.connector.client_session is not None
                     ):
-                        session.connector.client_session._validation_options = ValidationOptions(
-                            strict_output_validation=self._strict_validation
-                        )
+                        session.connector.client_session._validate_structured_outputs = self._strict_validation  # noqa: E501
             except ImportError:
                 # ValidationOptions may not be available in some mcp versions
                 pass

hud/server/server.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 import os
 import signal
 import sys
-import contextlib
 from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
@@ -25,9 +25,13 @@ __all__ = ["MCPServer"]
 
 logger = logging.getLogger(__name__)
 
+# Global flag to track if shutdown was triggered by SIGTERM
+_sigterm_received = False
+
 
 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)."""
+    global _sigterm_received
 
     async def _runner() -> None:
         stop_evt: asyncio.Event | None = None
@@ -35,9 +39,16 @@ def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) ->
             loop = asyncio.get_running_loop()
             stop_evt = asyncio.Event()
 
+            # Handle SIGTERM for production shutdown
+            def handle_sigterm() -> None:
+                global _sigterm_received
+                _sigterm_received = True
+                logger.info("Received SIGTERM signal")
+                stop_evt.set()
+
             # 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)
+                loop.add_signal_handler(signal.SIGTERM, handle_sigterm)
             if signal.getsignal(signal.SIGINT) is signal.SIG_DFL:
                 loop.add_signal_handler(signal.SIGINT, stop_evt.set)
 
@@ -62,6 +73,7 @@ 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.
+       Note: SIGINT (Ctrl+C) is not handled, allowing normal hot reload behavior.
     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:
@@ -69,7 +81,7 @@ class MCPServer(FastMCP):
        - ``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.
+       server teardown ONLY when SIGTERM is received (not on hot reload/SIGINT).
     4. Enhanced ``add_tool`` that accepts instances of
        :class:`hud.tools.base.BaseTool` which are classes that implement the
        FastMCP ``FunctionTool`` interface.
@@ -84,11 +96,17 @@ class MCPServer(FastMCP):
 
             @asynccontextmanager
             async def _lifespan(_: Any) -> AsyncGenerator[dict[str, Any], None]:
+                global _sigterm_received
                 try:
                     yield {}
                 finally:
-                    if self._shutdown_fn is not None:
+                    # Only call shutdown handler if SIGTERM was received
+                    if self._shutdown_fn is not None and _sigterm_received:
+                        logger.info("SIGTERM received, calling shutdown handler")
                         await self._shutdown_fn()
+                        _sigterm_received = False
+                    elif self._shutdown_fn is not None:
+                        logger.debug("Normal shutdown (hot reload), skipping shutdown handler")
 
             fastmcp_kwargs["lifespan"] = _lifespan
 
@@ -98,7 +116,7 @@ class MCPServer(FastMCP):
 
         # Replace FastMCP's low-level server with our version that supports
         # per-server initialization hooks
-        def _run_init(ctx: RequestContext) -> Any:
+        def _run_init(ctx: RequestContext | None = None) -> Any:
             if self._initializer_fn is not None and not self._did_init:
                 self._did_init = True
                 # Redirect stdout to stderr during initialization to prevent
@@ -138,6 +156,18 @@ class MCPServer(FastMCP):
     # Shutdown decorator: runs after server stops
     # Supports dockerized SIGTERM handling
     def shutdown(self, fn: Callable | None = None) -> Callable | None:
+        """Register a shutdown handler that runs ONLY on SIGTERM.
+
+        This handler will be called when the server receives a SIGTERM signal
+        (e.g., during container shutdown). It will NOT be called on:
+        - SIGINT (Ctrl+C or hot reload)
+        - Normal client disconnects
+        - Other graceful shutdowns
+
+        This ensures that persistent resources (like browser sessions) are only
+        cleaned up during actual termination, not during development hot reloads.
+        """
+
         def decorator(func: Callable) -> Callable:
             self._shutdown_fn = func
             return func

hud/shared/exceptions.py

@@ -18,6 +18,17 @@ class HudException(Exception):
     Consumers should be able to catch this exception to handle any HUD-related error.
     """
 
+    def __init__(self, message: str, response_json: dict[str, Any] | None = None) -> None:
+        super().__init__(message)
+        self.message = message
+        self.response_json = response_json
+
+    def __str__(self) -> str:
+        parts = [self.message]
+        if self.response_json:
+            parts.append(f"Response: {self.response_json}")
+        return " | ".join(parts)
+
 
 class HudRequestError(Exception):
     """Any request to the HUD API can raise this exception."""

hud/shared/tests/test_exceptions.py

@@ -155,3 +155,25 @@ class TestGymMakeException:
         assert "env_id" in error_str
         assert "test-env" in error_str
         assert "invalid config" in error_str
+
+
+class TestHudException:
+    """Test base HudException class."""
+
+    def test_str_with_response_json(self):
+        """Test HudException string representation with response_json."""
+        response_data = {"error": "test error", "code": 42}
+        error = HudException("Base error message", response_data)
+
+        error_str = str(error)
+        assert "Base error message" in error_str
+        assert "error" in error_str
+        assert "test error" in error_str
+
+    def test_str_without_response_json(self):
+        """Test HudException string representation without response_json."""
+        error = HudException("Just a message")
+
+        error_str = str(error)
+        assert error_str == "Just a message"
+        assert "Response:" not in error_str

hud/telemetry/tests/test_replay.py

@@ -0,0 +1,40 @@
+"""Tests for telemetry replay functionality."""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from hud.telemetry.replay import clear_trace, get_trace
+
+
+class TestReplayAPI:
+    """Tests for replay API functions."""
+
+    def test_get_trace_calls_internal(self):
+        """Test that get_trace calls the internal _get_trace function."""
+        with patch("hud.telemetry.replay._get_trace") as mock_get:
+            mock_get.return_value = None
+
+            result = get_trace("test-task-id")
+
+            mock_get.assert_called_once_with("test-task-id")
+            assert result is None
+
+    def test_clear_trace_calls_internal(self):
+        """Test that clear_trace calls the internal _clear_trace function."""
+        with patch("hud.telemetry.replay._clear_trace") as mock_clear:
+            clear_trace("test-task-id")
+
+            mock_clear.assert_called_once_with("test-task-id")
+
+    def test_get_trace_with_data(self):
+        """Test get_trace with mock data."""
+        mock_trace = {"trace": [{"step": 1}], "task_run_id": "test-123"}
+
+        with patch("hud.telemetry.replay._get_trace") as mock_get:
+            mock_get.return_value = mock_trace
+
+            result = get_trace("test-123")
+
+            assert result == mock_trace
+            mock_get.assert_called_once_with("test-123")

hud/telemetry/tests/test_trace.py

@@ -0,0 +1,63 @@
+"""Tests for telemetry trace functionality."""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from hud.telemetry.trace import trace
+
+
+class TestTraceAPI:
+    """Tests for trace API function."""
+
+    def test_trace_with_disabled_telemetry_and_no_api_key(self):
+        """Test trace behavior when telemetry is disabled and no API key."""
+        # Mock settings to disable telemetry and remove API key
+        mock_settings = type("Settings", (), {"telemetry_enabled": False, "api_key": None})()
+
+        with (
+            patch("hud.settings.get_settings", return_value=mock_settings),
+            patch("hud.telemetry.trace.OtelTrace") as mock_otel_trace,
+        ):
+            mock_otel_trace.return_value.__enter__.return_value = "custom-otlp-trace"
+
+            with trace("test-trace") as task_run_id:
+                # Should use placeholder ID for custom backends
+                assert task_run_id == "custom-otlp-trace"
+
+    def test_trace_with_enabled_telemetry_and_api_key(self):
+        """Test trace behavior when telemetry is enabled with API key."""
+        mock_settings = type("Settings", (), {"telemetry_enabled": True, "api_key": "test-key"})()
+
+        with (
+            patch("hud.settings.get_settings", return_value=mock_settings),
+            patch("hud.telemetry.trace.OtelTrace") as mock_otel_trace,
+            patch("hud.telemetry.trace.uuid.uuid4") as mock_uuid,
+        ):
+            mock_uuid.return_value = "mock-uuid-123"
+            mock_otel_trace.return_value.__enter__.return_value = "mock-uuid-123"
+
+            with trace("test-trace") as task_run_id:
+                # Should use generated UUID
+                assert task_run_id == "mock-uuid-123"
+
+    def test_trace_with_no_api_key(self):
+        """Test trace behavior with no API key (custom backend scenario)."""
+        mock_settings = type(
+            "Settings",
+            (),
+            {
+                "telemetry_enabled": True,  # Enabled but no API key
+                "api_key": None,
+            },
+        )()
+
+        with (
+            patch("hud.settings.get_settings", return_value=mock_settings),
+            patch("hud.telemetry.trace.OtelTrace") as mock_otel_trace,
+        ):
+            mock_otel_trace.return_value.__enter__.return_value = "custom-otlp-trace"
+
+            with trace("test-trace") as task_run_id:
+                # Should use custom backend placeholder
+                assert task_run_id == "custom-otlp-trace"

hud/tools/base.py

@@ -38,6 +38,7 @@ class BaseTool(ABC):
         name: str | None = None,
         title: str | None = None,
         description: str | None = None,
+        meta: dict[str, Any] | None = None,
     ) -> None:
         """Initialize the tool.
 
@@ -50,11 +51,13 @@ class BaseTool(ABC):
             name: Tool name for MCP registration (auto-generated from class name if not provided)
             title: Human-readable display name for the tool (auto-generated from class name)
             description: Tool description (auto-generated from docstring if not provided)
+            meta: Metadata to include in MCP tool listing (e.g., resolution info)
         """
         self.env = env
         self.name = name or self.__class__.__name__.lower().replace("tool", "")
         self.title = title or self.__class__.__name__.replace("Tool", "").replace("_", " ").title()
         self.description = description or (self.__doc__.strip() if self.__doc__ else None)
+        self.meta = meta
 
         # Expose attributes FastMCP expects when registering an instance directly
         self.__name__ = self.name  # FastMCP uses fn.__name__ if name param omitted
@@ -93,6 +96,7 @@ class BaseTool(ABC):
                 name=self.name,
                 title=self.title,
                 description=self.description,
+                meta=self.meta,
             )
         return self._mcp_tool
 
@@ -113,6 +117,7 @@ class BaseHub(FastMCP):
         env: Any | None = None,
         title: str | None = None,
         description: str | None = None,
+        meta: dict[str, Any] | None = None,
     ) -> None:
         """Create a new BaseHub.
 
@@ -124,6 +129,8 @@ class BaseHub(FastMCP):
             Optional long-lived environment object. Stored on the server
             instance (``layer.env``) and therefore available to every request
             via ``ctx.fastmcp.env``.
+        meta:
+            Metadata to include in MCP tool listing.
         """
 
         # Naming scheme for hidden objects
@@ -141,7 +148,7 @@ class BaseHub(FastMCP):
         async def _dispatch(  # noqa: ANN202
             name: str,
             arguments: dict | str | None = None,
-            ctx=None,  # noqa: ANN001
+            ctx: Any | None = None,
         ):
             """Gateway to hidden tools.
 
@@ -176,6 +183,7 @@ class BaseHub(FastMCP):
             title=dispatcher_title,
             description=dispatcher_desc,
             tags=set(),
+            meta=meta,
         )
         self._tool_manager.add_tool(dispatcher_tool)
 
@@ -344,8 +352,17 @@ class BaseHub(FastMCP):
                             "enum": [t[0] for t in sorted(internal_tools)],
                         },
                         "arguments": {
-                            "type": "object",
-                            "description": "Arguments to pass to the internal tool. See description for details on each tool's parameters.",  # noqa: E501
+                            "anyOf": [
+                                {
+                                    "type": "object",
+                                    "description": "Arguments object to pass to the internal tool",
+                                },
+                                {
+                                    "type": "string",
+                                    "description": "JSON string of arguments to pass to the internal tool",  # noqa: E501
+                                },
+                            ],
+                            "description": "Arguments to pass to the internal tool. Can be an object or JSON string. See description for details on each tool's parameters.",  # noqa: E501
                         },
                     },
                     "required": ["name", "arguments"],