hud-python 0.4.45__py3-none-any.whl → 0.5.13__py3-none-any.whl

hud/environment/connection.py

@@ -0,0 +1,258 @@
+"""Connection management for MCP servers."""
+
+from __future__ import annotations
+
+import logging
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+import mcp.types as mcp_types
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from fastmcp.client import Client as FastMCPClient
+    from fastmcp.tools.tool import Tool
+
+__all__ = ["ConnectionConfig", "ConnectionType", "Connector"]
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionType(str, Enum):
+    """Type of connection - determines parallelization capability."""
+
+    LOCAL = "local"  # Stdio/Docker - single instance, not parallelizable
+    REMOTE = "remote"  # HTTP/URL - can spawn multiple instances
+
+
+class ConnectionConfig:
+    """Configuration for filtering/transforming tools from a remote connection."""
+
+    def __init__(
+        self,
+        *,
+        prefix: str | None = None,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        transform: Callable[[Tool], Tool | None] | None = None,
+    ) -> None:
+        self.prefix = prefix
+        self.include = include
+        self.exclude = exclude
+        self.transform = transform
+
+
+class Connector:
+    """Manages a connection to an MCP server with tool caching.
+
+    Client creation is deferred to connect() so that:
+    1. Each parallel trace gets fresh client instances
+    2. Connection happens inside trace context (for header injection)
+    """
+
+    def __init__(
+        self,
+        transport: Any,
+        config: ConnectionConfig,
+        name: str,
+        connection_type: ConnectionType,
+        *,
+        auth: str | None = None,
+    ) -> None:
+        # Store transport config - client created in connect()
+        self._transport = transport
+        self._auth = auth
+        self.config = config
+        self.name = name
+        self.connection_type = connection_type
+        self.client: FastMCPClient[Any] | None = None
+        self._tools_cache: list[mcp_types.Tool] | None = None
+        self._prompts_cache: list[mcp_types.Prompt] | None = None
+        self._resources_cache: list[mcp_types.Resource] | None = None
+
+    def copy(self) -> Connector:
+        """Create a copy of this connector with fresh (unconnected) state.
+
+        The copy shares transport config but has its own client instance,
+        allowing parallel execution without conflicts.
+        """
+        return Connector(
+            transport=self._transport,
+            config=self.config,
+            name=self.name,
+            connection_type=self.connection_type,
+            auth=self._auth,
+        )
+
+    @property
+    def is_local(self) -> bool:
+        """True if this is a local (non-parallelizable) connection."""
+        return self.connection_type == ConnectionType.LOCAL
+
+    @property
+    def is_remote(self) -> bool:
+        """True if this is a remote (parallelizable) connection."""
+        return self.connection_type == ConnectionType.REMOTE
+
+    @property
+    def is_connected(self) -> bool:
+        return self.client is not None and self.client.is_connected()
+
+    @property
+    def cached_tools(self) -> list[mcp_types.Tool]:
+        return self._tools_cache or []
+
+    @property
+    def cached_prompts(self) -> list[mcp_types.Prompt]:
+        return self._prompts_cache or []
+
+    @property
+    def cached_resources(self) -> list[mcp_types.Resource]:
+        return self._resources_cache or []
+
+    async def connect(self) -> None:
+        """Create FastMCP client and connect.
+
+        Client is created here (not in __init__) so that:
+        1. Each parallel trace gets fresh client instances
+        2. httpx auto-instrumentation can inject trace headers
+        """
+        from fastmcp.client import Client as FastMCPClient
+
+        self.client = FastMCPClient(
+            transport=self._transport,
+            auth=self._auth,
+        )
+        await self.client.__aenter__()
+
+    async def disconnect(self) -> None:
+        """Disconnect and clear all caches."""
+        if self.client is not None and self.is_connected:
+            await self.client.__aexit__(None, None, None)
+        self.client = None
+        self._tools_cache = None
+        self._prompts_cache = None
+        self._resources_cache = None
+
+    async def list_tools(self) -> list[mcp_types.Tool]:
+        """Fetch tools from server, apply filters/transforms/prefix, and cache.
+
+        Always fetches fresh data from the server (no caching check).
+        The result is cached for use by router.build() via cached_tools property.
+        """
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        tools = await self.client.list_tools()
+
+        result: list[mcp_types.Tool] = []
+        for tool in tools:
+            # Apply include/exclude filter
+            if self.config.include is not None and tool.name not in self.config.include:
+                continue
+            if self.config.exclude is not None and tool.name in self.config.exclude:
+                continue
+
+            # Apply transform
+            if self.config.transform is not None:
+                from fastmcp.tools.tool import Tool as FastMCPTool
+
+                fastmcp_tool = FastMCPTool.model_construct(
+                    name=tool.name,
+                    description=tool.description or "",
+                    parameters=tool.inputSchema,
+                )
+                transformed = self.config.transform(fastmcp_tool)
+                if transformed is None:
+                    continue
+                tool = mcp_types.Tool(
+                    name=transformed.name,
+                    description=transformed.description,
+                    inputSchema=transformed.parameters,
+                )
+
+            # Apply prefix
+            name = f"{self.config.prefix}_{tool.name}" if self.config.prefix else tool.name
+            result.append(
+                mcp_types.Tool(
+                    name=name,
+                    description=tool.description,
+                    inputSchema=tool.inputSchema,
+                )
+            )
+
+        self._tools_cache = result
+        return result
+
+    async def call_tool(
+        self, name: str, arguments: dict[str, Any] | None = None
+    ) -> mcp_types.CallToolResult:
+        """Call a tool, stripping prefix if needed."""
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        # Strip prefix when calling remote
+        if self.config.prefix and name.startswith(f"{self.config.prefix}_"):
+            name = name[len(self.config.prefix) + 1 :]
+        return await self.client.call_tool_mcp(name, arguments or {})
+
+    async def list_resources(self) -> list[mcp_types.Resource]:
+        """Fetch resources from server and cache.
+
+        Always fetches fresh data from the server (no caching check).
+        The result is cached for use by router.build_resources() via cached_resources property.
+
+        Note: resources/list is optional in the MCP spec. If the server doesn't
+        implement it, we return an empty list gracefully.
+        """
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        try:
+            self._resources_cache = await self.client.list_resources()
+        except Exception as e:
+            # Handle servers that don't implement resources/list (optional in MCP spec)
+            if "Method not found" in str(e):
+                logger.debug("Server %s does not support resources/list", self.name)
+                self._resources_cache = []
+            else:
+                raise
+        return self._resources_cache
+
+    async def list_prompts(self) -> list[mcp_types.Prompt]:
+        """Fetch prompts from server and cache.
+
+        Always fetches fresh data from the server (no caching check).
+        The result is cached for use by router.build_prompts() via cached_prompts property.
+
+        Note: prompts/list is optional in the MCP spec. If the server doesn't
+        implement it, we return an empty list gracefully.
+        """
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        try:
+            self._prompts_cache = await self.client.list_prompts()
+        except Exception as e:
+            # Handle servers that don't implement prompts/list (optional in MCP spec)
+            if "Method not found" in str(e):
+                logger.debug("Server %s does not support prompts/list", self.name)
+                self._prompts_cache = []
+            else:
+                raise
+        return self._prompts_cache
+
+    async def read_resource(
+        self, uri: str
+    ) -> list[mcp_types.TextResourceContents | mcp_types.BlobResourceContents]:
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        return await self.client.read_resource(uri)
+
+    async def get_prompt(
+        self, name: str, arguments: dict[str, Any] | None = None
+    ) -> mcp_types.GetPromptResult:
+        if self.client is None:
+            raise RuntimeError("Not connected - call connect() first")
+        return await self.client.get_prompt(name, arguments)
+
+    def __repr__(self) -> str:
+        t = self.connection_type.value
+        return f"Connector({self.name!r}, {t}, connected={self.is_connected})"

hud/environment/connectors/__init__.py

@@ -0,0 +1,33 @@
+"""Connection connectors - methods for connecting to various sources."""
+
+from hud.environment.connectors.local import LocalConnectorMixin
+from hud.environment.connectors.openai import OpenAIConnectorMixin
+from hud.environment.connectors.remote import RemoteConnectorMixin
+
+__all__ = ["ConnectorsMixin"]
+
+
+class ConnectorsMixin(
+    RemoteConnectorMixin,
+    LocalConnectorMixin,
+    OpenAIConnectorMixin,
+):
+    """Combined connector mixin providing all connection methods.
+
+    Remote connections:
+        connect_hub(slug) - HUD Hub environment
+        connect_url(url) - MCP server via URL
+        connect_openapi(spec) - Mount OpenAPI spec as MCP server
+
+    Local connections (in-process):
+        connect_image(image) - Docker image via stdio
+        connect_fastapi(app) - Mount FastAPI app as MCP server
+        connect_server(server) - Mount MCPServer/FastMCP directly
+
+    MCP config:
+        connect_mcp(config) - Single mcp_config server (auto-detects local/remote)
+        connect_mcp_config(mcp_config) - Multiple mcp_config servers
+
+    Framework imports:
+        connect_function_tools(tools) - Import OpenAI Agents SDK FunctionTools
+    """

hud/environment/connectors/base.py

@@ -0,0 +1,68 @@
+"""Base connector mixin with shared helper."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from fastmcp.tools.tool import Tool
+
+    from hud.environment.connection import ConnectionType, Connector
+
+__all__ = ["BaseConnectorMixin"]
+
+
+class BaseConnectorMixin:
+    """Base mixin providing connection helper.
+
+    Requires:
+        _connections: dict[str, Connector]
+    """
+
+    _connections: dict[str, Connector]
+
+    def _add_connection(
+        self,
+        name: str,
+        transport: Any,
+        *,
+        connection_type: ConnectionType,
+        auth: str | None = None,
+        prefix: str | None = None,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        transform: Callable[[Tool], Tool | None] | None = None,
+    ) -> Any:
+        """Add a connection to the environment.
+
+        Args:
+            name: Connection name/alias.
+            transport: FastMCP transport (URL, config dict, etc.).
+            connection_type: LOCAL or REMOTE - determines parallelization.
+            auth: Authorization header value.
+            prefix: Prefix for tool names.
+            include: Only include these tools.
+            exclude: Exclude these tools.
+            transform: Transform function for tools.
+
+        Returns:
+            self for chaining.
+        """
+        from hud.environment.connection import ConnectionConfig, Connector
+
+        config = ConnectionConfig(
+            prefix=prefix,
+            include=include,
+            exclude=exclude,
+            transform=transform,
+        )
+        self._connections[name] = Connector(
+            transport,
+            config,
+            name,
+            connection_type=connection_type,
+            auth=auth,
+        )
+        return self

hud/environment/connectors/local.py

@@ -0,0 +1,177 @@
+"""Local connection connectors - Docker image, FastAPI, MCPServer."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from hud.environment.connectors.mcp_config import MCPConfigConnectorMixin
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from fastmcp.tools.tool import Tool
+
+__all__ = ["LocalConnectorMixin"]
+
+
+class LocalConnectorMixin(MCPConfigConnectorMixin):
+    """Mixin providing local connection methods.
+
+    Methods:
+        connect_image(image) - Run Docker image via stdio
+        connect_fastapi(app) - Mount FastAPI app as MCP server
+        connect_server(server) - Mount any MCPServer/FastMCP directly
+
+    Inherits connect_mcp() from MCPConfigConnectorMixin.
+
+    Note: include_router() is inherited from MCPServer (via FastMCP).
+    """
+
+    def connect_image(
+        self,
+        image: str,
+        *,
+        alias: str | None = None,
+        docker_args: list[str] | None = None,
+        env_vars: dict[str, str] | None = None,
+        prefix: str | None = None,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        transform: Callable[[Tool], Tool | None] | None = None,
+    ) -> Any:
+        """Connect to a Docker image via stdio.
+
+        Creates an MCP config that runs: docker run -i --rm {image}
+        Environment variables from `.env` files are auto-injected.
+
+        Example:
+            ```python
+            env = Environment("my-env")
+            env.connect_image("mcp/fetch")
+
+            async with env:
+                result = await env.call_tool("fetch", url="https://example.com")
+            ```
+        """
+        from hud.cli.utils.docker import create_docker_run_command
+
+        cmd = create_docker_run_command(
+            image=image,
+            docker_args=docker_args,
+            extra_env=env_vars,
+            interactive=True,
+            remove=True,
+        )
+
+        name = alias or image
+        mcp_config = {
+            name: {
+                "command": cmd[0],
+                "args": cmd[1:],
+            }
+        }
+        return self.connect_mcp(
+            mcp_config,
+            alias=name,
+            prefix=prefix,
+            include=include,
+            exclude=exclude,
+            transform=transform,
+        )
+
+    def connect_fastapi(
+        self,
+        app: Any,
+        *,
+        name: str | None = None,
+        prefix: str | None = None,
+        include_hidden: bool = True,
+    ) -> Any:
+        """Import a FastAPI application's routes as MCP tools.
+
+        Uses FastMCP's from_fastapi() to convert FastAPI endpoints to MCP tools,
+        then imports them synchronously so they're available immediately.
+
+        Args:
+            app: FastAPI application instance
+            name: Custom name for the server (defaults to app.title)
+            prefix: Optional prefix for tool names
+            include_hidden: If True (default), includes routes with include_in_schema=False
+
+        Example:
+            ```python
+            from fastapi import FastAPI
+
+            api = FastAPI()
+
+
+            @api.get("/users/{user_id}", operation_id="get_user")
+            def get_user(user_id: int):
+                return {"id": user_id, "name": "Alice"}
+
+
+            env = Environment("my-env")
+            env.connect_fastapi(api)
+
+            async with env:
+                result = await env.call_tool("get_user", user_id=1)
+            ```
+
+        Tip: Use operation_id in FastAPI decorators for cleaner tool names.
+        """
+        from fastmcp import FastMCP
+
+        # Temporarily enable hidden routes for OpenAPI generation
+        hidden_routes: list[Any] = []
+        if include_hidden:
+            for route in getattr(app, "routes", []):
+                if hasattr(route, "include_in_schema") and not route.include_in_schema:
+                    hidden_routes.append(route)
+                    route.include_in_schema = True
+            # Clear cached openapi schema so it regenerates
+            if hasattr(app, "openapi_schema"):
+                app.openapi_schema = None
+
+        try:
+            server_name = name or getattr(app, "title", None) or "fastapi"
+            mcp_server = FastMCP.from_fastapi(app=app, name=server_name)
+            # Use include_router for synchronous import (tools available immediately)
+            self.include_router(mcp_server, prefix=prefix)  # type: ignore
+        finally:
+            # Restore original states
+            for route in hidden_routes:
+                route.include_in_schema = False
+            if hidden_routes and hasattr(app, "openapi_schema"):
+                app.openapi_schema = None  # Clear cache again
+
+        return self
+
+    def connect_server(
+        self,
+        server: Any,
+        *,
+        prefix: str | None = None,
+    ) -> Any:
+        """Import an MCPServer or FastMCP instance's tools directly.
+
+        Example:
+            ```python
+            from fastmcp import FastMCP
+
+            tools = FastMCP("tools")
+
+
+            @tools.tool
+            def greet(name: str) -> str:
+                return f"Hello, {name}!"
+
+
+            env = Environment("my-env")
+            env.connect_server(tools)
+
+            async with env:
+                result = await env.call_tool("greet", name="World")
+            ```
+        """
+        self.include_router(server, prefix=prefix)  # type: ignore
+        return self

hud/environment/connectors/mcp_config.py

@@ -0,0 +1,137 @@
+"""MCP config connection connectors."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from hud.environment.connectors.base import BaseConnectorMixin
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from fastmcp.tools.tool import Tool
+
+__all__ = ["MCPConfigConnectorMixin"]
+
+
+class MCPConfigConnectorMixin(BaseConnectorMixin):
+    """Mixin providing mcp_config connection methods."""
+
+    def connect_mcp(
+        self,
+        config: dict[str, dict[str, Any]],
+        *,
+        alias: str | None = None,
+        prefix: str | None = None,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        transform: Callable[[Tool], Tool | None] | None = None,
+    ) -> Any:
+        """Connect using an mcp_config dictionary (single server).
+
+        Auto-detects LOCAL (stdio) vs REMOTE (URL) based on config.
+
+        Example:
+            ```python
+            env = Environment("my-env")
+
+            # Stdio server
+            env.connect_mcp(
+                {
+                    "filesystem": {
+                        "command": "npx",
+                        "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
+                    }
+                }
+            )
+
+            async with env:
+                await env.call_tool("read_file", path="/tmp/test.txt")
+            ```
+        """
+        from hud.environment.connection import ConnectionType
+        from hud.settings import settings
+
+        name = alias or next(iter(config.keys()), "mcp")
+        server_config = next(iter(config.values()), {})
+
+        is_local = "command" in server_config or "args" in server_config
+        conn_type = ConnectionType.LOCAL if is_local else ConnectionType.REMOTE
+
+        transport: Any = config
+        if not is_local and "url" in server_config:
+            max_request_timeout = 840
+            server_config.setdefault(
+                "sse_read_timeout",
+                min(settings.client_timeout, max_request_timeout)
+                if settings.client_timeout > 0
+                else max_request_timeout,
+            )
+            transport = _build_transport(server_config)
+
+        return self._add_connection(
+            name,
+            transport,
+            connection_type=conn_type,
+            prefix=prefix,
+            include=include,
+            exclude=exclude,
+            transform=transform,
+        )
+
+    def connect_mcp_config(
+        self,
+        mcp_config: dict[str, dict[str, Any]],
+        **kwargs: Any,
+    ) -> Any:
+        """Connect multiple servers from an mcp_config dictionary.
+
+        Example:
+            ```python
+            env = Environment("my-env")
+
+            # Claude Desktop style config
+            env.connect_mcp_config(
+                {
+                    "filesystem": {
+                        "command": "npx",
+                        "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
+                    },
+                    "github": {
+                        "command": "npx",
+                        "args": ["-y", "@modelcontextprotocol/server-github"],
+                        "env": {"GITHUB_TOKEN": "..."},
+                    },
+                }
+            )
+
+            async with env:
+                await env.call_tool("read_file", path="/tmp/test.txt")
+                await env.call_tool("search_repositories", query="mcp")
+            ```
+        """
+        # Store mcp_config for serialization (v4 format)
+        # Merge with existing if called multiple times
+        if not hasattr(self, "_mcp_config") or self._mcp_config is None:
+            self._mcp_config = {}
+        self._mcp_config.update(mcp_config)
+
+        for server_name, server_config in mcp_config.items():
+            self.connect_mcp({server_name: server_config}, alias=server_name, **kwargs)
+        return self
+
+
+def _build_transport(server_config: dict[str, Any]) -> Any:
+    from fastmcp.client.transports import SSETransport, StreamableHttpTransport
+    from fastmcp.mcp_config import infer_transport_type_from_url
+
+    url = server_config["url"]
+    transport_type = server_config.get("transport") or infer_transport_type_from_url(url)
+    transport_cls = SSETransport if transport_type == "sse" else StreamableHttpTransport
+
+    return transport_cls(
+        url=url,
+        headers=server_config.get("headers"),
+        auth=server_config.get("auth"),
+        sse_read_timeout=server_config.get("sse_read_timeout"),
+    )