hud-python 0.4.35__py3-none-any.whl → 0.4.36__py3-none-any.whl

hud/cli/utils/package_runner.py

@@ -0,0 +1,292 @@
+"""Run Python modules or commands as MCP servers.
+
+This module handles direct execution of MCP servers, including:
+- Python modules with an 'mcp' attribute
+- External commands via FastMCP proxy
+- Auto-reload functionality for development
+
+For Docker container execution, see hud dev command.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+import shlex
+import signal
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+from typing import Any
+
+from fastmcp import FastMCP
+
+logger = logging.getLogger(__name__)
+
+
+async def run_package_as_mcp(
+    command: str | list[str],
+    transport: str = "stdio",
+    port: int = 8765,
+    verbose: bool = False,
+    reload: bool = False,
+    watch_paths: list[str] | None = None,
+    server_attr: str = "mcp",
+    **extra_kwargs: Any,
+) -> None:
+    """Run a command as an MCP server.
+
+    Can run:
+    - Python modules: 'controller' (imports and looks for mcp attribute)
+    - Python -m commands: 'python -m controller'
+    - Docker commands: 'docker run -it my-mcp-server'
+    - Any executable: './my-mcp-binary'
+
+    Args:
+        command: Command to run (string or list)
+        transport: Transport type ("stdio" or "http")
+        port: Port for HTTP transport
+        verbose: Enable verbose logging
+        reload: Enable auto-reload on file changes
+        watch_paths: Paths to watch for changes (defaults to ['.'])
+        **extra_kwargs: Additional arguments
+    """
+    # Set up logging
+    if verbose:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+
+    # Handle reload mode
+    if reload:
+        if watch_paths is None:
+            watch_paths = ["."]
+
+        # Detect external command vs module reliably.
+        # If command is a string and contains spaces (e.g., "uv run python -m controller")
+        # treat as external command. Otherwise, detect common launchers or paths.
+        is_external_cmd = False
+        if isinstance(command, list):
+            is_external_cmd = True
+        elif isinstance(command, str):
+            stripped = command.strip()
+            if " " in stripped or any(
+                stripped.startswith(x)
+                for x in ["python", "uv ", "docker", "./", "/", ".\\", "C:\\"]
+            ):
+                is_external_cmd = True
+
+        if is_external_cmd:
+            # External command - pass command list directly
+            cmd_list = shlex.split(command) if isinstance(command, str) else command
+            run_with_reload(cmd_list, watch_paths, verbose)
+        else:
+            # Python module - use sys.argv approach
+            run_with_reload(None, watch_paths, verbose)
+        return
+
+    # Determine if it's a module import or a command
+    if isinstance(command, str) and not any(
+        command.startswith(x) for x in ["python", "docker", "./", "/", ".\\", "C:\\"]
+    ):
+        # Treat as Python module for backwards compatibility
+        logger.info("Importing module: %s", command)
+        module = importlib.import_module(command)
+
+        # Look for server attribute in the module
+        if not hasattr(module, server_attr):
+            logger.error(
+                "Module '%s' does not have an '%s' attribute (MCPServer instance)",
+                command,
+                server_attr,
+            )
+            sys.exit(1)
+
+        server = getattr(module, server_attr)
+
+        # Configure server options
+        run_kwargs = {
+            "transport": transport,
+            "show_banner": False,
+        }
+
+        if transport == "http":
+            # FastMCP expects port/path directly
+            run_kwargs["port"] = port
+            run_kwargs["path"] = "/mcp"
+
+        # Merge any extra kwargs
+        run_kwargs.update(extra_kwargs)
+
+        # Run the server
+        logger.info("Running %s on %s transport", server.name, transport)
+        await server.run_async(**run_kwargs)
+    else:
+        # Run as external command using shared proxy utility
+        # Parse command if string
+        cmd_list = shlex.split(command) if isinstance(command, str) else command
+
+        # Replace 'python' with the current interpreter to preserve venv
+        if cmd_list[0] == "python":
+            cmd_list[0] = sys.executable
+            logger.info("Replaced 'python' with: %s", sys.executable)
+
+        logger.info("Running command: %s", " ".join(cmd_list))
+
+        # Create MCP config for the command
+        config = {
+            "mcpServers": {
+                "default": {
+                    "command": cmd_list[0],
+                    "args": cmd_list[1:] if len(cmd_list) > 1 else [],
+                    # transport defaults to stdio
+                }
+            }
+        }
+
+        # Create proxy server
+        proxy = FastMCP.as_proxy(config, name=f"HUD Run - {cmd_list[0]}")
+
+        # Run the proxy
+        await proxy.run_async(
+            transport=transport if transport == "http" or transport == "stdio" else None,
+            port=port if transport == "http" else None,
+            show_banner=False,
+            **extra_kwargs,
+        )
+
+
+def run_with_reload(
+    target_func: Any,
+    watch_paths: list[str],
+    verbose: bool = False,
+) -> None:
+    """Run a function or command with file watching and auto-reload.
+
+    Args:
+        target_func: Function to run (sync) or command list
+        watch_paths: Paths to watch for changes
+        verbose: Enable verbose logging
+    """
+    try:
+        import watchfiles
+    except ImportError:
+        logger.error("watchfiles is required for --reload. Install with: pip install watchfiles")
+        sys.exit(1)
+
+    # Resolve watch paths
+    resolved_paths = []
+    for path_str in watch_paths:
+        path = Path(path_str).resolve()
+        if path.is_file():
+            # Watch the directory containing the file
+            resolved_paths.append(str(path.parent))
+        else:
+            resolved_paths.append(str(path))
+
+    def run_and_restart() -> None:
+        """Run the target function in a loop, restarting on file changes."""
+
+        process = None
+
+        def handle_signal(signum: int, frame: Any) -> None:
+            """Handle signals by terminating the subprocess."""
+            if process:
+                process.terminate()
+            sys.exit(0)
+
+        signal.signal(signal.SIGTERM, handle_signal)
+        signal.signal(signal.SIGINT, handle_signal)
+
+        stop_event = threading.Event()  # Define stop_event at the start
+
+        while True:
+            # Run the target function or command
+            if target_func is None:
+                # Use sys.argv approach for Python modules
+                child_args = [a for a in sys.argv[1:] if a != "--reload"]
+                # If first arg is already 'run', don't inject it again
+                if child_args and child_args[0] == "run":
+                    cmd = [sys.executable, "-m", "hud", *child_args]
+                else:
+                    cmd = [sys.executable, "-m", "hud", "run", *child_args]
+            elif isinstance(target_func, list):
+                # It's a command list
+                cmd = target_func
+            else:
+                # It's a callable - run it directly
+                target_func()
+                # Wait for file changes before restarting
+                stop_event.wait()
+                continue
+
+            if verbose:
+                logger.info("Starting process: %s", " ".join(cmd))
+
+            process = subprocess.Popen(cmd, env=os.environ)  # noqa: S603
+
+            # Watch for changes
+            try:
+                # Use a proper threading.Event for stop_event as required by watchfiles
+                stop_event = threading.Event()
+
+                def _wait_and_set(
+                    stop_event: threading.Event, process: subprocess.Popen[bytes]
+                ) -> None:
+                    try:
+                        if process is not None:
+                            process.wait()
+                    finally:
+                        stop_event.set()
+
+                threading.Thread(
+                    target=_wait_and_set, args=(stop_event, process), daemon=True
+                ).start()
+
+                for changes in watchfiles.watch(*resolved_paths, stop_event=stop_event):
+                    logger.info("Raw changes detected: %s", changes)
+                    # Filter for relevant file types
+                    relevant_changes = [
+                        (change_type, path)
+                        for change_type, path in changes
+                        if any(path.endswith(ext) for ext in [".py", ".json", ".toml", ".yaml"])
+                        and "__pycache__" not in path
+                        and not Path(path).name.startswith(".")
+                    ]
+
+                    if relevant_changes:
+                        logger.info("File changes detected, restarting server...")
+                        if verbose:
+                            for change_type, path in relevant_changes:
+                                logger.debug("  %s: %s", change_type, path)
+
+                        # Terminate the process
+                        if process is not None:
+                            process.terminate()
+                        try:
+                            if process is not None:
+                                process.wait(timeout=5)
+                        except subprocess.TimeoutExpired:
+                            if process is not None:
+                                process.kill()
+                                process.wait()
+
+                        # Brief pause before restart
+                        time.sleep(0.1)
+                        break
+                    else:
+                        logger.debug("Changes detected but filtered out: %s", changes)
+            except KeyboardInterrupt:
+                # Handle Ctrl+C gracefully
+                if process:
+                    process.terminate()
+                    process.wait()
+                break
+
+    # Always act as the parent. The child is launched without --reload,
+    # so it won't re-enter this function.
+
+    run_and_restart()

hud/cli/utils/remote_runner.py

@@ -293,7 +293,10 @@ def run_remote_server(
     if not api_key:
         api_key = settings.api_key
         if not api_key:
-            click.echo("❌ API key required. Set HUD_API_KEY env var or use --api-key", err=True)
+            click.echo(
+                "❌ API key required. Set HUD_API_KEY in your environment or run: hud set HUD_API_KEY=your-key-here",  # noqa: E501
+                err=True,
+            )
             sys.exit(1)
 
     # Build headers

hud/clients/mcp_use.py

@@ -5,19 +5,22 @@ from __future__ import annotations
 import logging
 import traceback
 from typing import Any
+from urllib.parse import urlparse
 
 from mcp import Implementation, types
 from mcp.shared.exceptions import McpError
 from mcp_use.client import MCPClient as MCPUseClient
 from mcp_use.session import MCPSession as MCPUseSession
+from mcp_use.types.http import HttpOptions
 from pydantic import AnyUrl
 
+from hud.settings import settings
 from hud.types import MCPToolCall, MCPToolResult
 from hud.utils.hud_console import HUDConsole
 from hud.version import __version__ as hud_version
 
 from .base import BaseHUDClient
-from .utils.mcp_use_retry import patch_all_sessions
+from .utils.retry_transport import create_retry_httpx_client
 
 logger = logging.getLogger(__name__)
 hud_console = HUDConsole(logger=logger)
@@ -30,7 +33,11 @@ class MCPUseHUDClient(BaseHUDClient):
         name="hud-mcp-use", title="hud MCP-use Client", version=hud_version
     )
 
-    def __init__(self, mcp_config: dict[str, dict[str, Any]] | None = None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        mcp_config: dict[str, dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> None:
         """
         Initialize MCP-use client.
 
@@ -51,6 +58,12 @@ class MCPUseHUDClient(BaseHUDClient):
             str, tuple[str, types.Tool, types.Tool]
         ] = {}  # server_name, original_tool, prefixed_tool
         self._client: Any | None = None  # Will be MCPUseClient when available
+        # Transport options for MCP-use (disable_sse_fallback, httpx_client_factory, etc.)
+        # Default to retry-enabled HTTPX client if factory not provided
+        self._http_options: HttpOptions = HttpOptions(
+            httpx_client_factory=create_retry_httpx_client,
+            disable_sse_fallback=True,
+        )
 
     async def _connect(self, mcp_config: dict[str, dict[str, Any]]) -> None:
         """Create all sessions for MCP-use client."""
@@ -58,19 +71,29 @@ class MCPUseHUDClient(BaseHUDClient):
             logger.warning("Client is already connected, cannot connect again")
             return
 
+        # If a server target matches HUD's MCP host and no auth is provided,
+        # inject the HUD API key as a Bearer token to avoid OAuth browser flow.
+        try:
+            hud_mcp_host = urlparse(settings.hud_mcp_url).netloc
+            if mcp_config and settings.api_key and hud_mcp_host:
+                for server_cfg in mcp_config.values():
+                    server_url = server_cfg.get("url")
+                    if not server_url:
+                        continue
+                    if urlparse(server_url).netloc == hud_mcp_host and not server_cfg.get("auth"):
+                        server_cfg["auth"] = settings.api_key
+        except Exception:
+            logger.warning("Failed to parse HUD MCP URL")
+
         config = {"mcpServers": mcp_config}
         if MCPUseClient is None:
             raise ImportError("MCPUseClient is not available")
-        self._client = MCPUseClient.from_dict(config)
+        self._client = MCPUseClient.from_dict(config, http_options=self._http_options)
         try:
             assert self._client is not None  # noqa: S101
             self._sessions = await self._client.create_all_sessions()
             hud_console.info(f"Created {len(self._sessions)} MCP sessions")
 
-            # Patch all sessions with retry logic
-            patch_all_sessions(self._sessions)
-            hud_console.debug("Applied retry logic to all MCP sessions")
-
             # Configure validation for all sessions based on client setting
             try:
                 for session in self._sessions.values():

hud/datasets/parallel.py

@@ -115,7 +115,9 @@ def _process_worker(
                     task_name = task_dict.get("prompt") or f"Task {index}"
 
                     # Use the job_id to group all tasks under the same job
-                    with hud.trace(task_name, job_id=job_id, task_id=task_dict.get("id")):
+                    raw_task_id = task_dict.get("id")
+                    safe_task_id = str(raw_task_id) if raw_task_id is not None else None
+                    with hud.trace(task_name, job_id=job_id, task_id=safe_task_id):
                         # Convert dict to Task
                         task = Task(**task_dict)
 

hud/datasets/runner.py

@@ -104,7 +104,10 @@ async def run_dataset(
                 task_name = task_dict.get("prompt") or f"Task {index}"
                 if custom_system_prompt and "system_prompt" not in task_dict:
                     task_dict["system_prompt"] = custom_system_prompt
-                with hud.trace(task_name, job_id=job_obj.id, task_id=task_dict.get("id")):
+                # Ensure task_id is a string for baggage propagation
+                raw_task_id = task_dict.get("id")
+                safe_task_id = str(raw_task_id) if raw_task_id is not None else None
+                with hud.trace(task_name, job_id=job_obj.id, task_id=safe_task_id):
                     # Convert dict to Task here, at trace level
                     task = Task(**task_dict)
 

hud/otel/context.py

@@ -239,8 +239,25 @@ async def _update_task_status_async(
 
     try:
         data: dict[str, Any] = {"status": status}
-        if job_id:
-            data["job_id"] = job_id
+
+        # Resolve effective job_id from explicit param, OTel baggage, or current job context
+        effective_job_id: str | None = job_id
+        if not effective_job_id:
+            bj = baggage.get_baggage("hud.job_id")
+            if isinstance(bj, str) and bj:
+                effective_job_id = bj
+        if not effective_job_id:
+            try:
+                from hud.telemetry.job import get_current_job  # Local import to avoid cycles
+
+                current_job = get_current_job()
+                if current_job:
+                    effective_job_id = current_job.id
+            except Exception:
+                effective_job_id = None
+
+        if effective_job_id:
+            data["job_id"] = effective_job_id
         if error_message:
             data["error_message"] = error_message
 
@@ -302,8 +319,25 @@ def _update_task_status_sync(
 
     try:
         data: dict[str, Any] = {"status": status}
-        if job_id:
-            data["job_id"] = job_id
+
+        # Resolve effective job_id from explicit param, OTel baggage, or current job context
+        effective_job_id: str | None = job_id
+        if not effective_job_id:
+            bj = baggage.get_baggage("hud.job_id")
+            if isinstance(bj, str) and bj:
+                effective_job_id = bj
+        if not effective_job_id:
+            try:
+                from hud.telemetry.job import get_current_job  # Local import to avoid cycles
+
+                current_job = get_current_job()
+                if current_job:
+                    effective_job_id = current_job.id
+            except Exception:
+                effective_job_id = None
+
+        if effective_job_id:
+            data["job_id"] = effective_job_id
         if error_message:
             data["error_message"] = error_message
 

hud/rl/buffer.py

@@ -155,6 +155,9 @@ class DatasetBuffer(Buffer[Task]):
                 f"This is because the number of training steps ({self.training_steps}) is not a multiple of the dataset size ({self.dataset_size})"  # noqa: E501
             )
 
+        if config.verbose:
+            hud_console.info(f"Sample task: {tasks[0]}")
+
         self.add_fill(tasks, self.number_of_tasks, config.training.shuffle_dataset)
 
     def _validate_tasks(self, tasks: list[Task]) -> list[Task]:

hud/rl/tests/test_learner.py

@@ -163,7 +163,7 @@ def test_skip_update_when_zero_adv(monkeypatch, learner_stub: GRPOLearner):
     # Return a zero scalar loss that *depends* on params so backward works,
     # but has zero gradients (no update signal).
     def _zero_loss(self, sample) -> torch.Tensor:
-        return sum(p.sum() for p in self.policy.parameters()) * 0.0
+        return sum(p.sum() for p in self.policy.parameters()) * 0.0  # type: ignore
 
     monkeypatch.setattr(GRPOLearner, "compute_loss", _zero_loss, raising=True)
 

hud/server/server.py

@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any
 
 import anyio
 from fastmcp.server.server import FastMCP, Transport
+from starlette.responses import JSONResponse, Response
 
 from hud.server.low_level import LowLevelServerWithInit
 
@@ -20,6 +21,7 @@ if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable
 
     from mcp.shared.context import RequestContext
+    from starlette.requests import Request
 
 __all__ = ["MCPServer"]
 
@@ -163,7 +165,15 @@ class MCPServer(FastMCP):
                 # Redirect stdout to stderr during initialization to prevent
                 # any library prints from corrupting the MCP protocol
                 with contextlib.redirect_stdout(sys.stderr):
-                    return self._initializer_fn(ctx)
+                    # Check if function accepts ctx parameter
+                    import inspect
+
+                    sig = inspect.signature(self._initializer_fn)
+                    if "ctx" in sig.parameters:
+                        return self._initializer_fn(ctx)
+                    else:
+                        # Call without ctx for simpler usage
+                        return self._initializer_fn()
             return None
 
         # Save the old server's handlers before replacing it
@@ -233,6 +243,23 @@ class MCPServer(FastMCP):
 
         _run_with_sigterm(_bootstrap)
 
+    async def run_async(
+        self,
+        transport: Transport | None = None,
+        show_banner: bool = True,
+        **transport_kwargs: Any,
+    ) -> None:
+        """Run the server with HUD enhancements."""
+        if transport is None:
+            transport = "stdio"
+
+        # Register HTTP helpers for HTTP transport
+        if transport in ("http", "sse"):
+            self._register_hud_helpers()
+            logger.info("Registered HUD helper endpoints at /hud/*")
+
+        await super().run_async(transport=transport, show_banner=show_banner, **transport_kwargs)
+
     # Tool registration helper -- appends BaseTool to FastMCP
     def add_tool(self, obj: Any, **kwargs: Any) -> None:
         from hud.tools.base import BaseTool
@@ -242,3 +269,132 @@ class MCPServer(FastMCP):
             return
 
         super().add_tool(obj, **kwargs)
+
+    # Override to keep original callables when used as a decorator
+    def tool(self, name_or_fn: Any = None, **kwargs: Any) -> Any:  # type: ignore[override]
+        """Register a tool but return the original function in decorator form.
+
+        - Decorator usage (@mcp.tool, @mcp.tool("name"), @mcp.tool(name="name"))
+          registers with FastMCP and returns the original function for composition.
+        - Call-form (mcp.tool(fn, ...)) behaves the same but returns fn.
+        """
+        # Accept BaseTool / FastMCP Tool instances or callables in call-form
+        if name_or_fn is not None and not isinstance(name_or_fn, str):
+            try:
+                from hud.tools.base import BaseTool  # lazy import
+            except Exception:
+                BaseTool = tuple()  # type: ignore[assignment]
+            try:
+                from fastmcp.tools.tool import Tool as _FastMcpTool
+            except Exception:
+                _FastMcpTool = tuple()  # type: ignore[assignment]
+
+            # BaseTool instance → add underlying FunctionTool
+            if isinstance(name_or_fn, BaseTool):
+                super().add_tool(name_or_fn.mcp, **kwargs)
+                return name_or_fn
+            # FastMCP Tool/FunctionTool instance → add directly
+            if isinstance(name_or_fn, _FastMcpTool):
+                super().add_tool(name_or_fn, **kwargs)
+                return name_or_fn
+            # Callable function → register via FastMCP.tool and return original fn
+            if callable(name_or_fn):
+                super().tool(name_or_fn, **kwargs)
+                return name_or_fn
+
+        # Decorator form: get FastMCP's decorator, register, then return original fn
+        base_decorator = super().tool(name_or_fn, **kwargs)
+
+        def _wrapper(fn: Any) -> Any:
+            base_decorator(fn)
+            return fn
+
+        return _wrapper
+
+    def _register_hud_helpers(self) -> None:
+        """Register HUD helper HTTP routes.
+
+        This adds:
+        - GET /hud - Overview of available endpoints
+        - GET /hud/tools - List all registered tools with their schemas
+        - GET /hud/resources - List all registered resources
+        - GET /hud/prompts - List all registered prompts
+        """
+
+        @self.custom_route("/hud/tools", methods=["GET"])
+        async def list_tools(request: Request) -> Response:
+            """List all registered tools with their names, descriptions, and schemas."""
+            tools = []
+            # _tools is a mapping of tool_name -> FunctionTool/Tool instance
+            for tool_key, tool in self._tool_manager._tools.items():
+                tool_data = {"name": tool_key}
+                try:
+                    # Prefer converting to MCP model for consistent fields
+                    mcp_tool = tool.to_mcp_tool()
+                    tool_data["description"] = getattr(mcp_tool, "description", "")
+                    if hasattr(mcp_tool, "inputSchema") and mcp_tool.inputSchema:
+                        tool_data["input_schema"] = mcp_tool.inputSchema  # type: ignore[assignment]
+                    if hasattr(mcp_tool, "outputSchema") and mcp_tool.outputSchema:
+                        tool_data["output_schema"] = mcp_tool.outputSchema  # type: ignore[assignment]
+                except Exception:
+                    # Fallback to direct attributes on FunctionTool
+                    tool_data["description"] = getattr(tool, "description", "")
+                    params = getattr(tool, "parameters", None)
+                    if params:
+                        tool_data["input_schema"] = params
+                tools.append(tool_data)
+
+            return JSONResponse({"server": self.name, "tools": tools, "count": len(tools)})
+
+        @self.custom_route("/hud/resources", methods=["GET"])
+        async def list_resources(request: Request) -> Response:
+            """List all registered resources."""
+            resources = []
+            for resource_key, resource in self._resource_manager._resources.items():
+                resource_data = {
+                    "uri": resource_key,
+                    "name": resource.name,
+                    "description": resource.description,
+                    "mimeType": resource.mime_type,
+                }
+                resources.append(resource_data)
+
+            return JSONResponse(
+                {"server": self.name, "resources": resources, "count": len(resources)}
+            )
+
+        @self.custom_route("/hud/prompts", methods=["GET"])
+        async def list_prompts(request: Request) -> Response:
+            """List all registered prompts."""
+            prompts = []
+            for prompt_key, prompt in self._prompt_manager._prompts.items():
+                prompt_data = {
+                    "name": prompt_key,
+                    "description": prompt.description,
+                }
+                # Check if it has arguments
+                if hasattr(prompt, "arguments") and prompt.arguments:
+                    prompt_data["arguments"] = [
+                        {"name": arg.name, "description": arg.description, "required": arg.required}
+                        for arg in prompt.arguments
+                    ]
+                prompts.append(prompt_data)
+
+            return JSONResponse({"server": self.name, "prompts": prompts, "count": len(prompts)})
+
+        @self.custom_route("/hud", methods=["GET"])
+        async def hud_info(request: Request) -> Response:
+            """Show available HUD helper endpoints."""
+            base_url = str(request.base_url).rstrip("/")
+            return JSONResponse(
+                {
+                    "name": "HUD MCP Development Helpers",
+                    "server": self.name,
+                    "endpoints": {
+                        "tools": f"{base_url}/hud/tools",
+                        "resources": f"{base_url}/hud/resources",
+                        "prompts": f"{base_url}/hud/prompts",
+                    },
+                    "description": "These endpoints help you inspect your MCP server during development.",  # noqa: E501
+                }
+            )