onetool-mcp 1.0.0b1__py3-none-any.whl

ot/proxy/manager.py

@@ -0,0 +1,396 @@
+"""ProxyManager for connecting to external MCP servers using FastMCP Client.
+
+Manages connections to external MCP servers and routes tool calls
+through OneTool's single `run` tool interface.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import os
+from dataclasses import dataclass
+from typing import Any
+
+from fastmcp import Client
+from fastmcp.client.transports import StdioTransport
+from loguru import logger
+from mcp import types
+
+from ot.config.mcp import McpServerConfig, expand_secrets, expand_subprocess_env
+from ot.logging import LogSpan
+
+
+@dataclass
+class ProxyToolInfo:
+    """Information about a proxied tool."""
+
+    server: str
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+class ProxyManager:
+    """Manages connections to external MCP servers using FastMCP Client.
+
+    Connects to configured MCP servers at startup and provides
+    a unified interface for calling their tools.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the proxy manager."""
+        self._clients: dict[str, Client] = {}  # type: ignore[type-arg]
+        self._tools_by_server: dict[str, list[types.Tool]] = {}
+        self._initialized = False
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    @property
+    def servers(self) -> list[str]:
+        """List of connected server names."""
+        return list(self._clients.keys())
+
+    @property
+    def tool_count(self) -> int:
+        """Total number of proxied tools across all servers."""
+        return sum(len(tools) for tools in self._tools_by_server.values())
+
+    def get_connection(self, server: str) -> Client | None:  # type: ignore[type-arg]
+        """Get a client by server name."""
+        return self._clients.get(server)
+
+    def list_tools(self, server: str | None = None) -> list[ProxyToolInfo]:
+        """List available tools from proxied servers.
+
+        Args:
+            server: Optional server name to filter by. If None, returns all tools.
+
+        Returns:
+            List of ProxyToolInfo for available tools.
+        """
+        tools: list[ProxyToolInfo] = []
+
+        if server:
+            if server in self._tools_by_server:
+                for tool in self._tools_by_server[server]:
+                    tools.append(
+                        ProxyToolInfo(
+                            server=server,
+                            name=tool.name,
+                            description=tool.description or "",
+                            input_schema=tool.inputSchema,
+                        )
+                    )
+        else:
+            for srv_name, srv_tools in self._tools_by_server.items():
+                for tool in srv_tools:
+                    tools.append(
+                        ProxyToolInfo(
+                            server=srv_name,
+                            name=tool.name,
+                            description=tool.description or "",
+                            input_schema=tool.inputSchema,
+                        )
+                    )
+
+        return tools
+
+    async def call_tool(
+        self,
+        server: str,
+        tool: str,
+        arguments: dict[str, Any] | None = None,
+        timeout: float = 30.0,
+    ) -> str:
+        """Call a tool on a proxied MCP server.
+
+        Args:
+            server: Name of the server to call.
+            tool: Name of the tool to call.
+            arguments: Arguments to pass to the tool.
+            timeout: Timeout for the call in seconds.
+
+        Returns:
+            Text result from the tool.
+
+        Raises:
+            ValueError: If server is not connected.
+            RuntimeError: If the tool returns an error.
+            TimeoutError: If the call times out.
+        """
+        client = self._clients.get(server)
+        if not client:
+            available = ", ".join(self._clients.keys()) or "none"
+            raise ValueError(f"Server '{server}' not connected. Available: {available}")
+
+        arguments = arguments or {}
+
+        with LogSpan(span="proxy.tool.call", server=server, tool=tool) as span:
+            try:
+                result = await asyncio.wait_for(
+                    client.call_tool(tool, arguments),
+                    timeout=timeout,
+                )
+            except TimeoutError:
+                logger.error(
+                    f"Proxy tool timeout | server={server} | tool={tool} | timeout={timeout}s"
+                )
+                raise TimeoutError(
+                    f"Tool {server}.{tool} timed out after {timeout}s"
+                ) from None
+
+            # Extract text from result
+            text_parts: list[str] = []
+            for content in result.content:
+                if isinstance(content, types.TextContent):
+                    text_parts.append(content.text)
+                elif hasattr(content, "data"):
+                    text_parts.append(f"[Binary content: {type(content).__name__}]")
+
+            result_text = (
+                "\n".join(text_parts) if text_parts else "Tool returned empty response."
+            )
+            span.add("resultLength", len(result_text))
+            return result_text
+
+    def call_tool_sync(
+        self,
+        server: str,
+        tool: str,
+        arguments: dict[str, Any] | None = None,
+        timeout: float = 30.0,
+    ) -> str:
+        """Synchronously call a tool on a proxied MCP server.
+
+        This is a blocking wrapper around the async call_tool method,
+        suitable for use from sync code (like executed Python code).
+
+        Args:
+            server: Name of the server to call.
+            tool: Name of the tool to call.
+            arguments: Arguments to pass to the tool.
+            timeout: Timeout for the call in seconds.
+
+        Returns:
+            Text result from the tool.
+        """
+        if self._loop is None:
+            raise RuntimeError(
+                "Proxy manager not initialized - no event loop available"
+            )
+
+        future = asyncio.run_coroutine_threadsafe(
+            self.call_tool(server, tool, arguments, timeout),
+            self._loop,
+        )
+        return future.result(timeout=timeout + 5)
+
+    async def connect(self, configs: dict[str, McpServerConfig]) -> None:
+        """Connect to all enabled MCP servers.
+
+        Args:
+            configs: Dictionary of server name -> configuration.
+        """
+        if self._initialized:
+            return
+
+        self._loop = asyncio.get_running_loop()
+
+        enabled_configs = {name: cfg for name, cfg in configs.items() if cfg.enabled}
+
+        if not enabled_configs:
+            logger.debug("No MCP servers configured for proxying")
+            self._initialized = True
+            return
+
+        with LogSpan(span="proxy.init", serverCount=len(enabled_configs)) as span:
+            connected = 0
+            failed = 0
+
+            for name, config in enabled_configs.items():
+                try:
+                    await self._connect_server(name, config)
+                    connected += 1
+                except Exception as e:
+                    failed += 1
+                    logger.warning(f"Failed to connect to MCP server '{name}': {e}")
+
+            span.add("connected", connected)
+            span.add("failed", failed)
+            span.add("toolCount", self.tool_count)
+
+        self._initialized = True
+
+    async def _connect_server(self, name: str, config: McpServerConfig) -> None:
+        """Connect to a single MCP server using FastMCP Client."""
+        with LogSpan(span="proxy.connect", server=name, type=config.type) as span:
+            client = self._create_client(name, config)
+
+            # Enter the client context manager for persistent connection
+            await client.__aenter__()  # type: ignore[no-untyped-call]
+
+            try:
+                # List tools to verify connection and cache tool info
+                tools = await client.list_tools()
+
+                self._clients[name] = client
+                self._tools_by_server[name] = tools
+
+                span.add("toolCount", len(tools))
+                logger.info(
+                    f"Connected to {config.type} MCP server '{name}' with {len(tools)} tools"
+                )
+
+            except Exception:
+                # Clean up on failure
+                await client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+                raise
+
+    def _create_client(self, name: str, config: McpServerConfig) -> Client:  # type: ignore[type-arg]
+        """Create a FastMCP Client for the given configuration."""
+        if config.type == "http":
+            return self._create_http_client(name, config)
+        elif config.type == "stdio":
+            return self._create_stdio_client(name, config)
+        else:
+            raise ValueError(f"Unknown server type: {config.type}")
+
+    def _create_http_client(self, name: str, config: McpServerConfig) -> Client:  # type: ignore[type-arg]
+        """Create an HTTP/SSE client."""
+        if not config.url:
+            raise RuntimeError(f"Server {name}: HTTP server requires url")
+
+        # Auto-upgrade http:// to https://
+        url = config.url
+        if url.startswith("http://"):
+            url = "https://" + url[7:]
+            logger.debug(f"Upgraded {name} URL to HTTPS: {url}")
+
+        # Expand secrets in headers
+        headers = {}
+        for key, value in config.headers.items():
+            if "${" in value:
+                headers[key] = expand_secrets(value)
+            else:
+                headers[key] = value
+
+        return Client(url, headers=headers, timeout=float(config.timeout))
+
+    def _create_stdio_client(self, name: str, config: McpServerConfig) -> Client:  # type: ignore[type-arg]
+        """Create a stdio client."""
+        if not config.command:
+            raise RuntimeError(f"Server {name}: stdio server requires command")
+
+        # Build environment: PATH only + explicit config.env
+        env = {"PATH": os.environ.get("PATH", "")}
+        for key, value in config.env.items():
+            env[key] = expand_subprocess_env(value)
+
+        transport = StdioTransport(
+            command=config.command,
+            args=config.args,
+            env=env,
+        )
+
+        return Client(transport, timeout=float(config.timeout))
+
+    async def shutdown(self) -> None:
+        """Disconnect from all MCP servers."""
+        if not self._clients:
+            return
+
+        with LogSpan(span="proxy.shutdown", serverCount=len(self._clients)):
+            for name, client in list(self._clients.items()):
+                try:
+                    await client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+                    logger.debug(f"Disconnected from MCP server '{name}'")
+                except (Exception, asyncio.CancelledError) as e:
+                    logger.debug(f"Error disconnecting from '{name}': {e}")
+
+            self._clients.clear()
+            self._tools_by_server.clear()
+            self._initialized = False
+
+    async def reconnect(self, configs: dict[str, McpServerConfig]) -> None:
+        """Reconnect to all MCP servers.
+
+        Shuts down existing connections and reconnects with fresh config.
+
+        Args:
+            configs: Dictionary of server name -> configuration.
+        """
+        await self.shutdown()
+        await self.connect(configs)
+
+    def reconnect_sync(self, configs: dict[str, McpServerConfig]) -> None:
+        """Synchronously reconnect to all MCP servers.
+
+        Blocking wrapper for reconnect, suitable for calling from sync code.
+
+        Args:
+            configs: Dictionary of server name -> configuration.
+        """
+        loop = self._loop
+
+        # Try to get running loop if we don't have one stored
+        if loop is None:
+            with contextlib.suppress(RuntimeError):
+                loop = asyncio.get_running_loop()
+
+        if loop is None:
+            # No event loop available - just reset state, connect will happen on next use
+            self._clients.clear()
+            self._tools_by_server.clear()
+            self._initialized = False
+            return
+
+        future = asyncio.run_coroutine_threadsafe(
+            self.reconnect(configs),
+            loop,
+        )
+        try:
+            future.result(timeout=60)
+        except Exception as e:
+            logger.warning(f"Error during proxy reconnect: {e}")
+
+
+# Global proxy manager instance
+_proxy_manager: ProxyManager | None = None
+
+
+def get_proxy_manager() -> ProxyManager:
+    """Get or create the global proxy manager instance.
+
+    Returns:
+        ProxyManager instance.
+    """
+    global _proxy_manager
+    if _proxy_manager is None:
+        _proxy_manager = ProxyManager()
+    return _proxy_manager
+
+
+def reset_proxy_manager() -> None:
+    """Reset the global proxy manager (for testing)."""
+    global _proxy_manager
+    _proxy_manager = None
+
+
+def reconnect_proxy_manager() -> None:
+    """Reconnect the global proxy manager with fresh config.
+
+    Loads server configs from the current configuration and reconnects
+    all MCP proxy servers. Call this after modifying server config.
+    """
+    from ot.config.loader import get_config
+
+    proxy = get_proxy_manager()
+    cfg = get_config()
+
+    if cfg.servers:
+        proxy.reconnect_sync(cfg.servers)
+    else:
+        # No servers configured - just reset state
+        proxy._clients.clear()
+        proxy._tools_by_server.clear()
+        proxy._initialized = False

ot/registry/__init__.py

@@ -0,0 +1,189 @@
+"""Tool registry package with auto-discovery for user-defined Python tools.
+
+The registry scans the `src/ot_tools/` directory, extracts function signatures and
+docstrings using AST parsing, and provides formatted context for LLM code generation.
+"""
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Any
+
+from docstring_parser import parse as parse_docstring
+
+from .models import ArgInfo, ToolInfo
+from .registry import ToolRegistry
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from pathlib import Path
+
+__all__ = [
+    "ArgInfo",
+    "ToolInfo",
+    "ToolRegistry",
+    "describe_tool",
+    "get_registry",
+    "list_tools",
+]
+
+# Global registry instance
+_registry: ToolRegistry | None = None
+
+
+def _build_tool_info_from_callable(
+    name: str,
+    func: Callable[..., Any],
+    pack: str | None = None,
+) -> ToolInfo:
+    """Build ToolInfo from a callable using inspect.
+
+    Args:
+        name: Full tool name (e.g., "ot.tools").
+        func: The function object.
+        pack: Pack name if applicable.
+
+    Returns:
+        ToolInfo with extracted signature and docstring info.
+    """
+    # Get signature
+    try:
+        sig = inspect.signature(func)
+        signature = f"{name}{sig}"
+    except (ValueError, TypeError):
+        signature = f"{name}(...)"
+
+    # Parse docstring
+    doc = func.__doc__ or ""
+    parsed = parse_docstring(doc)
+
+    # Build args list
+    args: list[ArgInfo] = []
+    for param_name, param in sig.parameters.items():
+        if param_name in ("self", "cls"):
+            continue
+
+        # Get type annotation
+        if param.annotation != inspect.Parameter.empty:
+            param_type = (
+                param.annotation.__name__
+                if hasattr(param.annotation, "__name__")
+                else str(param.annotation)
+            )
+        else:
+            param_type = "Any"
+
+        # Get default value
+        default = None
+        if param.default != inspect.Parameter.empty:
+            default = repr(param.default)
+
+        # Get description from parsed docstring
+        description = ""
+        for doc_param in parsed.params:
+            if doc_param.arg_name == param_name:
+                description = doc_param.description or ""
+                break
+
+        args.append(
+            ArgInfo(
+                name=param_name,
+                type=param_type,
+                default=default,
+                description=description,
+            )
+        )
+
+    # Get return description
+    returns = (parsed.returns.description or "") if parsed.returns else ""
+
+    return ToolInfo(
+        name=name,
+        pack=pack,
+        module=func.__module__,
+        signature=signature,
+        description=parsed.short_description or "",
+        args=args,
+        returns=returns,
+    )
+
+
+def _register_ot_pack(registry: ToolRegistry) -> None:
+    """Register the ot pack tools in the registry.
+
+    The ot pack provides introspection functions that need parameter
+    shorthand support like other tools.
+    """
+    from ot.meta import PACK_NAME, get_ot_pack_functions
+
+    ot_functions = get_ot_pack_functions()
+
+    for func_name, func in ot_functions.items():
+        full_name = f"{PACK_NAME}.{func_name}"
+        tool_info = _build_tool_info_from_callable(full_name, func, pack=PACK_NAME)
+        registry.register_tool(tool_info)
+
+
+def get_registry(tools_path: Path | None = None, rescan: bool = False) -> ToolRegistry:
+    """Get or create the global tool registry.
+
+    Uses config's tools_dir glob patterns if available, otherwise falls back
+    to the provided tools_path or default 'src/ot_tools/' directory.
+
+    Args:
+        tools_path: Path to tools directory (fallback if no config).
+        rescan: If True, rescan even if registry exists.
+
+    Returns:
+        ToolRegistry instance with discovered tools.
+    """
+    from ot.config.loader import get_config
+
+    global _registry
+
+    if _registry is None:
+        _registry = ToolRegistry(tools_path)
+        # Use config's tool files if available
+        config = get_config()
+        tool_files = config.get_tool_files()
+        if tool_files:
+            _registry.scan_files(tool_files)
+        else:
+            _registry.scan_directory()
+        # Register ot pack tools for param shorthand support
+        _register_ot_pack(_registry)
+    elif rescan:
+        # Rescan using config's tool files
+        config = get_config()
+        tool_files = config.get_tool_files()
+        if tool_files:
+            _registry.scan_files(tool_files)
+        else:
+            _registry.scan_directory()
+        # Re-register ot pack tools after rescan
+        _register_ot_pack(_registry)
+
+    return _registry
+
+
+def list_tools() -> str:
+    """List all registered tools.
+
+    Returns:
+        Summary of all registered tools.
+    """
+    registry = get_registry(rescan=True)
+    return registry.format_summary()
+
+
+def describe_tool(name: str) -> str:
+    """Describe a specific tool.
+
+    Args:
+        name: Tool function name.
+
+    Returns:
+        Detailed tool description.
+    """
+    registry = get_registry()
+    return registry.describe_tool(name)

ot/registry/models.py

@@ -0,0 +1,57 @@
+"""Registry models - Pydantic models for tool and argument information."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class ArgInfo(BaseModel):
+    """Information about a function argument."""
+
+    name: str = Field(description="Argument name")
+    type: str = Field(default="Any", description="Type annotation")
+    default: str | None = Field(default=None, description="Default value if any")
+    description: str = Field(
+        default="", description="Argument description from docstring"
+    )
+
+
+class ToolInfo(BaseModel):
+    """Information about a registered tool function."""
+
+    name: str = Field(
+        description="Qualified tool name (e.g., 'code.search' or 'search')"
+    )
+    pack: str | None = Field(
+        default=None, description="Pack name if tool belongs to a pack"
+    )
+    module: str = Field(description="Module path (e.g., 'tools.gold_prices')")
+    signature: str = Field(description="Full function signature")
+    description: str = Field(
+        default="", description="Function description from docstring"
+    )
+    args: list[ArgInfo] = Field(default_factory=list, description="Function arguments")
+    returns: str = Field(
+        default="", description="Return type/description from docstring"
+    )
+    examples: list[str] = Field(
+        default_factory=list, description="Usage examples from @tool decorator"
+    )
+    tags: list[str] = Field(
+        default_factory=list, description="Categorization tags from @tool decorator"
+    )
+    enabled: bool = Field(default=True, description="Whether the tool is enabled")
+    deprecated: bool = Field(
+        default=False, description="Whether the tool is deprecated"
+    )
+    deprecated_message: str | None = Field(
+        default=None, description="Deprecation message"
+    )
+    config_schema: str | None = Field(
+        default=None,
+        description="Config class source code extracted via AST (class Config(BaseModel))",
+    )
+    requires: dict[str, list[tuple[str, ...] | dict[str, str] | str]] | None = Field(
+        default=None,
+        description="Dependencies from __ot_requires__ (cli and lib lists)",
+    )