mcp_manager.py

"""MCP (Model Context Protocol) manager — lifecycle, tool discovery & execution.

Manages one or more MCP servers as child processes.  Each server's tools are
discovered at startup and exposed to the agent under the naming convention
``mcp__<server_name>__<tool_name>``.

Architecture
------------
Because the MCP Python SDK is async (asyncio) but the Wcode agent runs
synchronously, this module runs a dedicated asyncio event loop in a daemon
thread.  All MCP operations are submitted to that loop via
``asyncio.run_coroutine_threadsafe``.

Adding a new MCP server
-----------------------
1. Install the server package (e.g. ``pip install mcp-server-fetch``).
2. Add an entry to ``MCP_SERVERS`` in ``config.py``:

   .. code-block:: python

       MCP_SERVERS = {
           "fetch": {
               "command": "python",
               "args": ["-m", "mcp_server_fetch"],
           },
           "my_tool": {
               "command": "my-mcp-server",
               "args": [],
               "env": {"API_KEY": "..."},   # optional
           },
       }

No other code changes are needed — tools are discovered automatically at
startup.
"""

from __future__ import annotations

import asyncio
import json
import threading
import time
from dataclasses import dataclass, field
from typing import Any

from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

from config import get_logger

logger = get_logger(__name__)

# ── Tool-name constants ──────────────────────────────────────────────────

MCP_TOOL_PREFIX = "mcp__"  # namespace all MCP tools to avoid collisions


def mcp_tool_name(server: str, tool: str) -> str:
    """Build the namespaced tool name: ``mcp__<server>__<tool>``."""
    return f"{MCP_TOOL_PREFIX}{server}__{tool}"


def parse_mcp_tool(full_name: str) -> tuple[str, str] | None:
    """Extract (server_name, tool_name) from a namespaced MCP tool name."""
    if not full_name.startswith(MCP_TOOL_PREFIX):
        return None
    inner = full_name[len(MCP_TOOL_PREFIX):]
    parts = inner.split("__", 1)
    if len(parts) != 2:
        return None
    return parts[0], parts[1]


# ── Async event-loop thread ──────────────────────────────────────────────


class _AsyncLoopThread:
    """Dedicated asyncio event loop running in a daemon thread.

    MCP operations are async, but the agent is sync — this bridges the two.
    """

    def __init__(self) -> None:
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None

    @property
    def loop(self) -> asyncio.AbstractEventLoop:
        if self._loop is None or self._loop.is_closed():
            raise RuntimeError("MCP event loop is not running")
        return self._loop

    def start(self) -> None:
        assert self._thread is None, "already started"
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        logger.info("[mcp] async event-loop thread started")

    def _run(self) -> None:
        asyncio.set_event_loop(self._loop)
        self._loop.run_forever()

    def run(self, coro: Any, timeout: float = 60.0) -> Any:
        """Submit a coroutine to the loop and block until it completes."""
        future = asyncio.run_coroutine_threadsafe(coro, self.loop)  # type: ignore[arg-type]
        return future.result(timeout=timeout)

    def stop(self) -> None:
        if self._loop is None:
            return
        self._loop.call_soon_threadsafe(self._loop.stop)
        if self._thread is not None:
            self._thread.join(timeout=5.0)
        logger.info("[mcp] async event-loop thread stopped")


# ── Single-server connection ─────────────────────────────────────────────


@dataclass
class _ServerConnection:
    """Holds the session and transport handles for one MCP server."""

    name: str
    session: ClientSession
    tools: list[dict[str, Any]] = field(default_factory=list)

    # Keep references so we can clean up the async context managers later.
    _read_stream: Any = None
    _write_stream: Any = None
    _transport_ctx: Any = None   # async context manager for stdio_client
    _session_ctx: Any = None     # async context manager for ClientSession


# ── MCP Manager ──────────────────────────────────────────────────────────


class MCPManager:
    """Central registry for MCP server connections.

    Usage::

        manager = MCPManager()
        manager.start_all(MCP_SERVERS)       # launch servers, discover tools
        tools = manager.get_all_tool_defs()  # expose to agent
        ...
        result = manager.call_tool("mcp__fetch__fetch", {"url": "..."})
        manager.shutdown()
    """

    def __init__(self) -> None:
        self._loop_thread = _AsyncLoopThread()
        self._connections: dict[str, _ServerConnection] = {}
        self._lock = threading.Lock()

    # ── Startup / shutdown ───────────────────────────────────────────

    def start_all(self, server_configs: dict[str, dict[str, Any]]) -> list[str]:
        """Launch every configured MCP server and discover their tools.

        Parameters
        ----------
        server_configs : dict
            Mapping of server name → kwargs for :class:`StdioServerParameters`.

        Returns
        -------
        list[str]
            Names of servers that were started successfully.
        """
        self._loop_thread.start()
        started: list[str] = []

        for name, cfg in server_configs.items():
            try:
                self._start_one(name, cfg)
                started.append(name)
            except Exception:
                logger.exception(f"[mcp] failed to start server '{name}'")

        logger.info(
            f"[mcp] %d/%d servers started",
            len(started),
            len(server_configs),
        )
        return started

    def _start_one(self, name: str, cfg: dict[str, Any]) -> None:
        """Start a single MCP server and discover its tools."""
        params = StdioServerParameters(
            command=cfg["command"],
            args=cfg.get("args", []),
            env=cfg.get("env"),
        )
        conn = _ServerConnection(name=name, session=None)  # type: ignore[arg-type]

        # Create the stdio transport + session inside the event loop.
        async def _init() -> _ServerConnection:
            # Enter stdio_client context manager
            transport_ctx = stdio_client(params)
            read, write = await transport_ctx.__aenter__()
            conn._transport_ctx = transport_ctx
            conn._read_stream = read
            conn._write_stream = write

            # Enter ClientSession context manager
            session_ctx = ClientSession(read, write)
            session = await session_ctx.__aenter__()
            conn._session_ctx = session_ctx
            conn.session = session

            init_result = await session.initialize()
            logger.info(
                "[mcp] connected to '%s' (%s v%s, protocol %s)",
                name,
                init_result.serverInfo.name,
                init_result.serverInfo.version,
                init_result.protocolVersion,
            )

            # Discover tools
            tools_result = await session.list_tools()
            conn.tools = []
            for t in tools_result.tools:
                tool_info = {
                    "name": t.name,
                    "description": t.description or "",
                    "inputSchema": t.inputSchema or {},
                }
                conn.tools.append(tool_info)
            logger.info(
                "[mcp] '%s' exposes %d tool(s): %s",
                name,
                len(conn.tools),
                [t["name"] for t in conn.tools],
            )
            return conn

        conn = self._loop_thread.run(_init(), timeout=30.0)
        with self._lock:
            self._connections[name] = conn

    def shutdown(self) -> None:
        """Gracefully shut down all MCP server connections.

        Best-effort: errors during cleanup are logged but never fatal,
        because the OS reclaims child processes when we exit anyway.
        """
        async def _close_all() -> None:
            for name, conn in list(self._connections.items()):
                try:
                    if conn._session_ctx is not None:
                        await conn._session_ctx.__aexit__(None, None, None)
                except Exception:
                    logger.debug(f"[mcp] session close error for '{name}'")
                try:
                    if conn._transport_ctx is not None:
                        await conn._transport_ctx.__aexit__(None, None, None)
                except Exception:
                    logger.debug(f"[mcp] transport close error for '{name}'")

        try:
            self._loop_thread.run(_close_all(), timeout=10.0)
        except (RuntimeError, TimeoutError, Exception) as exc:
            # Context-manager cleanup can fail with RuntimeError when
            # the calling task differs from the entering task — this is
            # harmless; the process tree will be cleaned up on exit.
            logger.debug("[mcp] shutdown cleanup (non-critical): %s", exc)

        with self._lock:
            self._connections.clear()
        self._loop_thread.stop()

    # ── Tool definitions (OpenAI / litellm format) ───────────────────

    def get_all_tool_defs(self) -> list[dict[str, Any]]:
        """Return OpenAI-format tool definitions for every discovered tool.

        Each definition uses the namespaced name ``mcp__<server>__<tool>``
        so the agent can route calls back through :meth:`call_tool`.
        """
        defs: list[dict[str, Any]] = []
        with self._lock:
            for srv_name, conn in self._connections.items():
                for t in conn.tools:
                    func_def: dict[str, Any] = {
                        "name": mcp_tool_name(srv_name, t["name"]),
                        "description": (
                            f"[MCP:{srv_name}] {t['description']}".strip()
                        ),
                    }
                    schema = t.get("inputSchema")
                    if schema:
                        func_def["parameters"] = {
                            k: v
                            for k, v in schema.items()
                            if k not in ("$schema",)
                        }
                    defs.append(
                        {"type": "function", "function": func_def}
                    )
        return defs

    def get_tool_names(self) -> list[str]:
        """Return the namespaced names of all available MCP tools."""
        names: list[str] = []
        with self._lock:
            for srv_name, conn in self._connections.items():
                for t in conn.tools:
                    names.append(mcp_tool_name(srv_name, t["name"]))
        return names

    # ── Tool execution ───────────────────────────────────────────────

    def call_tool(
        self, full_name: str, arguments: dict[str, Any], timeout: float = 60.0
    ) -> str:
        """Execute an MCP tool and return its text output.

        Parameters
        ----------
        full_name : str
            Namespaced tool name (``mcp__<server>__<tool>``).
        arguments : dict
            Tool arguments as a flat dict.
        timeout : float
            Maximum seconds to wait for the tool result.

        Returns
        -------
        str
            The aggregated text content from the tool result, or an error
            message string.
        """
        parsed = parse_mcp_tool(full_name)
        if parsed is None:
            return f"Error: invalid MCP tool name '{full_name}'"

        server_name, tool_name = parsed
        with self._lock:
            conn = self._connections.get(server_name)
        if conn is None:
            return f"Error: MCP server '{server_name}' is not connected"

        async def _call() -> str:
            try:
                result = await conn.session.call_tool(
                    tool_name, arguments=arguments
                )
            except Exception as exc:
                logger.exception(
                    "[mcp] tool call failed: %s/%s", server_name, tool_name
                )
                return f"Error: MCP tool '{tool_name}' failed: {exc}"

            # Aggregate text blocks from the result content.
            parts: list[str] = []
            for block in result.content:
                if hasattr(block, "text"):
                    parts.append(block.text)
                elif hasattr(block, "data"):
                    parts.append(f"[binary data: {len(block.data)} bytes]")
                else:
                    parts.append(str(block))
            return "\n".join(parts) if parts else "(no output)"

        return self._loop_thread.run(_call(), timeout=timeout)


# ── Hooks for tools.py ───────────────────────────────────────────────────


def register_mcp_tools_to(
    handler: dict[str, Any],
    tools_list: list[dict[str, Any]],
    manager: MCPManager,
) -> None:
    """Register MCP tool definitions into a shared TOOLS list / handler dict.

    Idempotent — skips tools that are already registered.

    Args:
        handler: The name→callable mapping to populate.
        tools_list: The OpenAI-format tool definitions list to append to.
        manager: The :class:`MCPManager` instance to read tools from.
    """
    mcp_defs = manager.get_all_tool_defs()

    existing_names = {
        t.get("function", {}).get("name", "")
        for t in tools_list
    }

    for mcp_def in mcp_defs:
        name = mcp_def["function"]["name"]
        if name not in existing_names:
            tools_list.append(mcp_def)

    for name in manager.get_tool_names():
        if name in handler:
            continue  # already registered
        def _make_handler(tool_name: str):
            def _handler(**kwargs) -> str:
                return manager.call_tool(tool_name, kwargs)
            return _handler
        handler[name] = _make_handler(name)