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

hud/server/low_level.py

@@ -1,132 +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,
-                    )
+"""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

@@ -1,166 +1,170 @@
-"""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 helpers."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import sys
+import contextlib
+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
+                # 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)
+            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

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