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

hud/server/server.py

@@ -13,13 +13,14 @@ 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
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Callable
 
-    from mcp.shared.context import RequestContext
+    from starlette.requests import Request
 
 __all__ = ["MCPServer"]
 
@@ -35,6 +36,31 @@ def _run_with_sigterm(coro_fn: Callable[..., Any], *args: Any, **kwargs: Any) ->
 
     sys.stderr.flush()
 
+    # Check if we're already in an event loop
+    try:
+        loop = asyncio.get_running_loop()
+        logger.warning(
+            "HUD server is running in an existing event loop. "
+            "SIGTERM handling may be limited. "
+            "Consider using await hub.run_async() instead of hub.run() in async contexts."
+        )
+
+        task = loop.create_task(coro_fn(*args, **kwargs))
+
+        # Try to handle SIGTERM if possible
+        if sys.platform != "win32":
+
+            def handle_sigterm(signum: Any, frame: Any) -> None:
+                logger.info("SIGTERM received in async context, cancelling task...")
+                loop.call_soon_threadsafe(task.cancel)
+
+            signal.signal(signal.SIGTERM, handle_sigterm)
+
+        return
+
+    except RuntimeError:
+        pass
+
     async def _runner() -> None:
         stop_evt: asyncio.Event | None = None
         if sys.platform != "win32" and os.getenv("FASTMCP_DISABLE_SIGTERM_HANDLER") != "1":
@@ -125,7 +151,11 @@ class MCPServer(FastMCP):
                     # Force flush logs to ensure they're visible
                     sys.stderr.flush()
 
-                    if self._shutdown_fn is not None and _sigterm_received:
+                    if (
+                        self._shutdown_fn is not None
+                        and _sigterm_received
+                        and not self._shutdown_has_run
+                    ):
                         logger.info("SIGTERM detected! Calling @mcp.shutdown handler...")
                         sys.stderr.flush()
                         try:
@@ -135,7 +165,9 @@ class MCPServer(FastMCP):
                         except Exception as e:
                             logger.error("Error during @mcp.shutdown: %s", e)
                             sys.stderr.flush()
-                        _sigterm_received = False
+                        finally:
+                            self._shutdown_has_run = True
+                            _sigterm_received = False
                     elif self._shutdown_fn is not None:
                         logger.info(
                             "No SIGTERM. This is a hot reload (SIGINT) or normal exit. Skipping @mcp.shutdown handler."  # noqa: E501
@@ -151,19 +183,53 @@ class MCPServer(FastMCP):
         self._initializer_fn: Callable | None = None
         self._did_init = False
         self._replaced_server = False
+        self._shutdown_has_run = False  # Guard against double-execution of shutdown hook
 
     def _replace_with_init_server(self) -> None:
         """Replace the low-level server with init version when needed."""
         if self._replaced_server:
             return
 
-        def _run_init(ctx: RequestContext | None = None) -> Any:
+        async def _run_init(ctx: object | None = None) -> None:
+            """Run the user initializer exactly once, with stdout redirected."""
             if self._initializer_fn is not None and not self._did_init:
                 self._did_init = True
-                # Redirect stdout to stderr during initialization to prevent
-                # any library prints from corrupting the MCP protocol
+                # Prevent stdout from polluting the MCP protocol on stdio/HTTP
                 with contextlib.redirect_stdout(sys.stderr):
-                    return self._initializer_fn(ctx)
+                    import inspect
+
+                    fn = self._initializer_fn
+                    sig = inspect.signature(fn)
+                    params = sig.parameters
+
+                    ctx_param = params.get("ctx") or params.get("_ctx")
+                    if ctx_param is not None:
+                        if ctx_param.kind == inspect.Parameter.KEYWORD_ONLY:
+                            result = fn(**{ctx_param.name: ctx})
+                        else:
+                            result = fn(ctx)
+                    else:
+                        required_params = [
+                            p
+                            for p in params.values()
+                            if p.default is inspect._empty
+                            and p.kind
+                            in (
+                                inspect.Parameter.POSITIONAL_ONLY,
+                                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                                inspect.Parameter.KEYWORD_ONLY,
+                            )
+                        ]
+                        if required_params:
+                            param_list = ", ".join(p.name for p in required_params)
+                            raise TypeError(
+                                "Initializer must accept no args or a single `ctx` argument; "
+                                f"received required parameters: {param_list}"
+                            )
+                        result = fn()
+                    if inspect.isawaitable(result):
+                        await result
+                    return None
             return None
 
         # Save the old server's handlers before replacing it
@@ -233,6 +299,38 @@ 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/*")
+
+        try:
+            await super().run_async(
+                transport=transport, show_banner=show_banner, **transport_kwargs
+            )
+        finally:
+            # Fallback: ensure SIGTERM-triggered shutdown runs even when a custom
+            # lifespan bypasses our default fastmcp shutdown path.
+            global _sigterm_received
+            if self._shutdown_fn is not None and _sigterm_received and not self._shutdown_has_run:
+                try:
+                    await self._shutdown_fn()
+                except Exception as e:  # pragma: no cover - defensive logging
+                    logger.error("Error during @mcp.shutdown (fallback): %s", e)
+                finally:
+                    self._shutdown_has_run = True
+                    _sigterm_received = False
+
     # Tool registration helper -- appends BaseTool to FastMCP
     def add_tool(self, obj: Any, **kwargs: Any) -> None:
         from hud.tools.base import BaseTool
@@ -242,3 +340,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
+                }
+            )

hud/server/tests/test_add_tool.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import sys
+import types
+from typing import Any, cast
+
+from hud.server import MCPServer
+
+
+def test_add_tool_accepts_base_tool(monkeypatch):
+    """If obj is BaseTool, its `.mcp` gets passed through to FastMCP.add_tool."""
+    # Stub hud.tools.base.BaseTool and capture FastMCP.add_tool calls
+    mod = types.ModuleType("hud.tools.base")
+
+    class FakeBaseTool:
+        """Stub type checked by isinstance() inside add_tool."""
+
+    # Tell the type checker we're mutating a dynamic module
+    mod_any = cast("Any", mod)
+    mod_any.BaseTool = FakeBaseTool
+    monkeypatch.setitem(sys.modules, "hud.tools.base", mod)
+
+    calls: dict[str, object | None] = {"obj": None, "kwargs": None}
+
+    def fake_super_add(self, obj: object, **kwargs: object) -> None:  # keep runtime the same
+        calls["obj"] = obj
+        calls["kwargs"] = kwargs
+
+    monkeypatch.setattr("hud.server.server.FastMCP.add_tool", fake_super_add, raising=True)
+
+    mcp = MCPServer(name="AddTool")
+    sentinel = object()
+
+    class MyTool(FakeBaseTool):
+        def __init__(self) -> None:
+            self.mcp = sentinel
+
+    mcp.add_tool(MyTool(), extra="yes")
+    assert calls["obj"] is sentinel
+    assert isinstance(calls["kwargs"], dict)
+    assert calls["kwargs"]["extra"] == "yes"
+
+
+def test_add_tool_plain_falls_back_to_super(monkeypatch):
+    """Non-BaseTool objects are passed unchanged to FastMCP.add_tool."""
+    calls = []
+
+    def fake_super_add(self, obj, **kwargs):
+        calls.append((obj, kwargs))
+
+    monkeypatch.setattr("hud.server.server.FastMCP.add_tool", fake_super_add, raising=True)
+
+    mcp = MCPServer(name="AddToolPlain")
+
+    async def fn():  # pragma: no cover - never awaited by FastMCP here
+        return "ok"
+
+    mcp.add_tool(fn, desc="x")
+    assert calls and calls[0][0] is fn
+    assert calls[0][1]["desc"] == "x"

hud/server/tests/test_context.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import os
+import sys
+
+try:
+    import multiprocessing.connection as _mp_conn
+
+    # Pull the exception dynamically; fall back to OSError if missing in stubs/runtime
+    MPAuthenticationError: type[BaseException] = getattr(_mp_conn, "AuthenticationError", OSError)
+except Exception:  # pragma: no cover
+    MPAuthenticationError = OSError
+
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from hud.server.context import attach_context, serve_context
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+pytestmark = pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="Context server uses UNIX domain sockets",
+)
+
+
+class CounterCtx:
+    def __init__(self) -> None:
+        self._n = 0
+
+    def inc(self) -> int:
+        self._n += 1
+        return self._n
+
+    def get(self) -> int:
+        return self._n
+
+
+def test_serve_and_attach_shared_state(tmp_path: Path) -> None:
+    sock = str(tmp_path / "hud_ctx.sock")
+
+    mgr = serve_context(CounterCtx(), sock_path=sock)
+    try:
+        c1 = attach_context(sock_path=sock)
+        assert c1.get() == 0
+        assert c1.inc() == 1
+
+        # Second attachment sees the same underlying object
+        c2 = attach_context(sock_path=sock)
+        assert c2.get() == 1
+        assert c2.inc() == 2
+        assert c1.get() == 2  # shared state
+    finally:
+        mgr.shutdown()
+
+
+def test_env_var_socket_path_overrides(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    sock = str(tmp_path / "env_ctx.sock")
+    monkeypatch.setenv("HUD_CTX_SOCK", sock)
+
+    mgr = serve_context(CounterCtx(), sock_path=None)
+    try:
+        c = attach_context(sock_path=None)
+        assert c.inc() == 1
+        assert c.get() == 1
+    finally:
+        mgr.shutdown()
+        monkeypatch.delenv("HUD_CTX_SOCK", raising=False)
+
+
+def test_wrong_authkey_rejected(tmp_path: Path) -> None:
+    sock = str(tmp_path / "auth_ctx.sock")
+    mgr = serve_context(CounterCtx(), sock_path=sock, authkey=b"correct")
+    try:
+        with pytest.raises(
+            (MPAuthenticationError, ConnectionRefusedError, BrokenPipeError, OSError)
+        ):
+            attach_context(sock_path=sock, authkey=b"wrong")
+    finally:
+        mgr.shutdown()
+
+
+def test_attach_nonexistent_raises(tmp_path: Path) -> None:
+    # ensure file truly doesn't exist
+    sock = str(tmp_path / "missing.sock")
+    if os.path.exists(sock):
+        os.unlink(sock)
+
+    with pytest.raises((FileNotFoundError, ConnectionRefusedError, OSError)):
+        attach_context(sock_path=sock)
+
+
+@pytest.mark.asyncio
+async def test_run_context_server_handles_keyboardinterrupt(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    """run_context_server should call manager.shutdown() when KeyboardInterrupt occurs."""
+    # Capture serve_context() and the returned manager
+    called = {"served": False, "shutdown": False, "addr": None}
+
+    class _Mgr:
+        def shutdown(self) -> None:
+            called["shutdown"] = True
+
+    def fake_serve(ctx, sock_path, authkey):
+        called["served"] = True
+        called["addr"] = sock_path
+        return _Mgr()
+
+    monkeypatch.setattr("hud.server.context.serve_context", fake_serve)
+
+    # Make asyncio.Event().wait() raise KeyboardInterrupt immediately
+    class _FakeEvent:
+        async def wait(self) -> None:
+            raise KeyboardInterrupt
+
+    monkeypatch.setattr("hud.server.context.asyncio.Event", lambda: _FakeEvent())
+
+    from hud.server.context import run_context_server
+
+    await run_context_server(object(), sock_path=str(tmp_path / "ctx.sock"))
+
+    assert called["served"] is True
+    assert called["shutdown"] is True
+    assert str(called["addr"]).endswith("ctx.sock")

hud/server/tests/test_mcp_server_handlers.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Any, cast
+
+from hud.server import MCPServer
+from hud.server.low_level import LowLevelServerWithInit
+
+
+def test_notification_handlers_preserved_on_replacement():
+    """When init server replaces low-level server, notification handlers must be kept."""
+    mcp = MCPServer(name="PreserveNotif")
+
+    # Seed a fake notification handler on the pre-replacement server
+    before = mcp._mcp_server
+    cast("dict[Any, Any]", before.notification_handlers)["foo/notify"] = object()
+
+    @mcp.initialize
+    async def _init(_ctx) -> None:
+        pass
+
+    after = mcp._mcp_server
+    assert isinstance(after, LowLevelServerWithInit)
+    assert after is not before, "low-level server should be replaced once"
+    # Must still contain our seeded handler (dict is copied over)
+    assert "foo/notify" in after.notification_handlers
+
+
+def test_init_server_replacement_is_idempotent():
+    """Second @initialize must NOT replace the low-level server again."""
+    mcp = MCPServer(name="InitIdempotent")
+
+    @mcp.initialize
+    async def _a(_ctx) -> None:
+        pass
+
+    first = mcp._mcp_server
+
+    @mcp.initialize
+    async def _b(_ctx) -> None:
+        # last initializer should win, but server object should not be replaced again
+        pass
+
+    second = mcp._mcp_server
+    assert first is second, "Server replacement should occur at most once per instance"