nanobot/nanobot/agent/tools/mcp.py

"""MCP client: connects to MCP servers and wraps their tools as native nanobot tools."""

import asyncio
import json
import re
import os
from contextlib import AsyncExitStack
from typing import Any

from loguru import logger

from nanobot.agent.tools.base import Tool
from nanobot.agent.tools.registry import ToolRegistry


_SAFE_TOOL_NAME_RE = re.compile(r"[^A-Za-z0-9_]+")


def _normalize_tool_segment(segment: str) -> str:
    """
    Normalize MCP server/tool names into a safe function name segment.

    - Replace non [A-Za-z0-9_] with underscore
    - Collapse repeated underscores
    - Trim leading/trailing underscores
    - Ensure non-empty
    """
    s = _SAFE_TOOL_NAME_RE.sub("_", (segment or "").strip())
    s = re.sub(r"_+", "_", s).strip("_")
    return s or "tool"


def _render_mcp_content_blocks(blocks: list[Any]) -> str:
    """Render MCP content blocks into a stable, readable string."""
    from mcp import types

    parts: list[str] = []
    for block in blocks or []:
        if isinstance(block, types.TextContent):
            parts.append(block.text)
            continue

        # Prefer structured JSON for non-text blocks when possible.
        dump = getattr(block, "model_dump", None)
        if callable(dump):
            try:
                parts.append(json.dumps(dump(), ensure_ascii=False, indent=2))
                continue
            except Exception:
                pass

        parts.append(str(block))
    return "\n".join([p for p in parts if p is not None]).strip()


class MCPToolWrapper(Tool):
    """Wraps a single MCP server tool as a nanobot Tool."""

    def __init__(
        self,
        session,
        *,
        server_key: str,
        tool_def,
        registered_name: str,
        call_timeout_s: float = 30.0,
    ):
        self._session = session
        self._original_name = tool_def.name
        self._server_key = server_key
        self._name = registered_name
        self._description = tool_def.description or tool_def.name
        self._parameters = tool_def.inputSchema or {"type": "object", "properties": {}}
        self._call_timeout_s = call_timeout_s

    @property
    def name(self) -> str:
        return self._name

    @property
    def description(self) -> str:
        return self._description

    @property
    def parameters(self) -> dict[str, Any]:
        return self._parameters

    async def execute(self, **kwargs: Any) -> str:
        try:
            result = await asyncio.wait_for(
                self._session.call_tool(self._original_name, arguments=kwargs),
                timeout=self._call_timeout_s,
            )
        except asyncio.TimeoutError:
            return (
                f"Error: MCP tool timed out after {self._call_timeout_s:.0f}s "
                f"({self._server_key}:{self._original_name})"
            )

        output = _render_mcp_content_blocks(getattr(result, "content", []))
        if not output:
            return "(no output)"

        # If the tool returned JSON, normalize empty collections to a clearer message.
        try:
            parsed = json.loads(output)
            if parsed == [] or parsed == {}:
                return "No results found."
        except (json.JSONDecodeError, ValueError):
            pass  # Not JSON, continue with original output

        return output


async def connect_mcp_server(
    name: str, cfg: Any, registry: ToolRegistry, stack: AsyncExitStack
) -> None:
    """Connect one MCP server and register its tools (used for lazy profile-scoped connections)."""
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client

    def _expand_env(env: dict[str, str]) -> dict[str, str]:
        """
        Expand $VARS in cfg.env using the current process environment.

        This lets configs safely reference secrets that are already injected into the
        container environment (e.g. via .env.shared), without duplicating them in JSON:
          { "GITEA_ACCESS_TOKEN": "$NANOBOT_GITLE_TOKEN" }
        """
        if not env:
            return {}
        expanded: dict[str, str] = {}
        for k, v in env.items():
            if v is None:
                continue
            expanded[k] = os.path.expandvars(str(v))
        return expanded

    if cfg.command:
        params = StdioServerParameters(
            command=cfg.command, args=cfg.args, env=_expand_env(cfg.env) or None
        )
        read, write = await stack.enter_async_context(stdio_client(params))
    elif cfg.url:
        from mcp.client.streamable_http import streamable_http_client

        read, write, _ = await stack.enter_async_context(
            streamable_http_client(cfg.url)
        )
    else:
        logger.warning(f"MCP server '{name}': no command or url configured, skipping")
        return

    session = await stack.enter_async_context(ClientSession(read, write))
    await session.initialize()

    tools = await session.list_tools()
    for tool_def in tools.tools:
        safe_server = _normalize_tool_segment(name)
        safe_tool = _normalize_tool_segment(tool_def.name)
        base = f"mcp_{safe_server}_{safe_tool}"
        registered_name = base
        i = 2
        while registry.has(registered_name):
            registered_name = f"{base}_{i}"
            i += 1

        wrapper = MCPToolWrapper(
            session,
            server_key=name,
            tool_def=tool_def,
            registered_name=registered_name,
        )
        registry.register(wrapper)
        logger.debug(f"MCP: registered tool '{wrapper.name}' from server '{name}'")

    logger.info(f"MCP server '{name}': connected, {len(tools.tools)} tools registered")


async def connect_mcp_servers(
    mcp_servers: dict, registry: ToolRegistry, stack: AsyncExitStack
) -> None:
    """Connect to every configured MCP server and register their tools."""
    for name, cfg in mcp_servers.items():
        try:
            await connect_mcp_server(name, cfg, registry, stack)
        except Exception as e:
            logger.error(f"MCP server '{name}': failed to connect: {e}")