Initial commit

opc/mcp_client.py

@@ -0,0 +1,299 @@
+"""MCP (Model Context Protocol) client supporting local (stdio) and remote
+(StreamableHTTP / SSE) servers.
+
+Connects to MCP servers, discovers their tools, and registers them into the
+OPC ToolRegistry so the agent loop can call them like any other built-in tool.
+Tool names are prefixed with the server name to prevent collisions when
+multiple servers expose identically-named tools.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import re
+from contextlib import AsyncExitStack
+from typing import Any
+
+from loguru import logger
+
+try:
+    from mcp import ClientSession
+    from mcp.client.stdio import stdio_client, StdioServerParameters
+    MCP_IMPORT_ERROR: Exception | None = None
+except ModuleNotFoundError as exc:  # pragma: no cover - exercised via import-time fallback
+    ClientSession = Any  # type: ignore[assignment]
+    stdio_client = None  # type: ignore[assignment]
+    StdioServerParameters = Any  # type: ignore[assignment]
+    MCP_IMPORT_ERROR = exc
+
+from opc.layer4_tools.registry import ToolDefinition
+
+
+def _sanitize(name: str) -> str:
+    """Replace non-alphanumeric/underscore chars with ``_``."""
+    return re.sub(r"[^a-zA-Z0-9_]", "_", name)
+
+
+def _require_mcp_dependency() -> None:
+    if MCP_IMPORT_ERROR is None:
+        return
+    raise RuntimeError(
+        "MCP support requires the optional `mcp` package to be installed."
+    ) from MCP_IMPORT_ERROR
+
+
+def _extract_content(content_blocks: list[Any]) -> str:
+    """Extract human-readable text from MCP tool result content blocks.
+
+    Handles ``text``, ``image``, and ``resource`` block types.
+    """
+    parts: list[str] = []
+    for block in content_blocks:
+        btype = getattr(block, "type", "")
+        if btype == "text":
+            parts.append(getattr(block, "text", ""))
+        elif btype == "image":
+            mime = getattr(block, "mimeType", "unknown")
+            data = getattr(block, "data", "")
+            if data:
+                parts.append(f"[image ({mime}): {len(base64.b64decode(data))} bytes]")
+            else:
+                parts.append(f"[image ({mime})]")
+        elif btype == "resource":
+            resource = getattr(block, "resource", None)
+            if resource:
+                uri = getattr(resource, "uri", "")
+                text = getattr(resource, "text", None)
+                if text:
+                    parts.append(text)
+                else:
+                    parts.append(f"[resource: {uri}]")
+    return "\n".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# Base connection interface
+# ---------------------------------------------------------------------------
+
+class _MCPConnectionBase:
+    """Shared interface for local and remote MCP connections."""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self._exit_stack = AsyncExitStack()
+        self._session: ClientSession | None = None
+        self._tools: list[dict[str, Any]] = []
+
+    async def start(self) -> None:
+        raise NotImplementedError
+
+    async def discover_tools(self) -> list[dict[str, Any]]:
+        assert self._session
+        result = await self._session.list_tools()
+        self._tools = [
+            {
+                "name": t.name,
+                "description": t.description or "",
+                "inputSchema": t.inputSchema if hasattr(t, "inputSchema") else {},
+            }
+            for t in result.tools
+        ]
+        logger.info(f"MCP server '{self.name}' exposes {len(self._tools)} tools")
+        return self._tools
+
+    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        assert self._session
+        result = await self._session.call_tool(tool_name, arguments)
+        is_error = getattr(result, "isError", False)
+        content_blocks = getattr(result, "content", [])
+        text = _extract_content(content_blocks)
+        if is_error:
+            return {"error": text or "MCP tool returned an error", "success": False}
+        return {"result": text, "success": True}
+
+    async def stop(self) -> None:
+        await self._exit_stack.aclose()
+        self._session = None
+        logger.info(f"MCP server '{self.name}' stopped")
+
+
+# ---------------------------------------------------------------------------
+# Local (stdio) connection
+# ---------------------------------------------------------------------------
+
+class MCPServerConnection(_MCPConnectionBase):
+    """A connection to a local MCP server process via stdio."""
+
+    def __init__(self, name: str, command: str, args: list[str], env: dict[str, str] | None = None) -> None:
+        super().__init__(name)
+        _require_mcp_dependency()
+        self.server_params = StdioServerParameters(command=command, args=args, env=env)
+
+    async def start(self) -> None:
+        transport = await self._exit_stack.enter_async_context(
+            stdio_client(self.server_params)
+        )
+        read_stream, write_stream = transport
+        self._session = await self._exit_stack.enter_async_context(
+            ClientSession(read_stream, write_stream)
+        )
+        result = await self._session.initialize()
+        server_info = getattr(result, "serverInfo", None) or {}
+        s_name = getattr(server_info, "name", "?") if not isinstance(server_info, dict) else server_info.get("name", "?")
+        s_ver = getattr(server_info, "version", "?") if not isinstance(server_info, dict) else server_info.get("version", "?")
+        logger.info(f"MCP server '{self.name}' initialized: {s_name} v{s_ver}")
+
+
+# ---------------------------------------------------------------------------
+# Remote (StreamableHTTP / SSE) connection
+# ---------------------------------------------------------------------------
+
+class MCPRemoteConnection(_MCPConnectionBase):
+    """A connection to a remote MCP server via StreamableHTTP or SSE fallback."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        super().__init__(name)
+        self.url = url
+        self.headers = headers or {}
+
+    async def start(self) -> None:
+        _require_mcp_dependency()
+        last_error: Exception | None = None
+
+        # Strategy 1: StreamableHTTP
+        try:
+            from mcp.client.streamable_http import streamablehttp_client
+            transport = await self._exit_stack.enter_async_context(
+                streamablehttp_client(self.url, headers=self.headers)
+            )
+            read_stream, write_stream = transport[0], transport[1]
+            self._session = await self._exit_stack.enter_async_context(
+                ClientSession(read_stream, write_stream)
+            )
+            result = await self._session.initialize()
+            server_info = getattr(result, "serverInfo", None) or {}
+            s_name = getattr(server_info, "name", "?") if not isinstance(server_info, dict) else server_info.get("name", "?")
+            s_ver = getattr(server_info, "version", "?") if not isinstance(server_info, dict) else server_info.get("version", "?")
+            logger.info(f"MCP remote '{self.name}' connected via StreamableHTTP: {s_name} v{s_ver}")
+            return
+        except Exception as exc:
+            last_error = exc
+            logger.debug(f"MCP remote '{self.name}' StreamableHTTP failed, trying SSE: {exc}")
+            # Reset exit stack for retry
+            await self._exit_stack.aclose()
+            self._exit_stack = AsyncExitStack()
+
+        # Strategy 2: SSE fallback
+        try:
+            from mcp.client.sse import sse_client
+            transport = await self._exit_stack.enter_async_context(
+                sse_client(self.url, headers=self.headers)
+            )
+            read_stream, write_stream = transport[0], transport[1]
+            self._session = await self._exit_stack.enter_async_context(
+                ClientSession(read_stream, write_stream)
+            )
+            result = await self._session.initialize()
+            server_info = getattr(result, "serverInfo", None) or {}
+            s_name = getattr(server_info, "name", "?") if not isinstance(server_info, dict) else server_info.get("name", "?")
+            s_ver = getattr(server_info, "version", "?") if not isinstance(server_info, dict) else server_info.get("version", "?")
+            logger.info(f"MCP remote '{self.name}' connected via SSE: {s_name} v{s_ver}")
+            return
+        except Exception as exc:
+            last_error = exc
+            logger.debug(f"MCP remote '{self.name}' SSE also failed: {exc}")
+
+        raise ConnectionError(
+            f"Failed to connect to remote MCP server '{self.name}' at {self.url}: {last_error}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Manager
+# ---------------------------------------------------------------------------
+
+class MCPManager:
+    """Manages multiple MCP server connections and registers their tools."""
+
+    def __init__(self) -> None:
+        self._servers: dict[str, _MCPConnectionBase] = {}
+
+    async def connect_local(
+        self,
+        name: str,
+        command: list[str],
+        env: dict[str, str] | None = None,
+        timeout: float = 30.0,
+    ) -> MCPServerConnection:
+        if len(command) < 1:
+            raise ValueError(f"MCP server '{name}' command is empty")
+        conn = MCPServerConnection(name, command=command[0], args=command[1:], env=env)
+        try:
+            await asyncio.wait_for(conn.start(), timeout=timeout)
+            self._servers[name] = conn
+            return conn
+        except Exception as exc:
+            logger.warning(f"Failed to start local MCP server '{name}': {exc}")
+            raise
+
+    # Backwards-compatible alias
+    connect = connect_local
+
+    async def connect_remote(
+        self,
+        name: str,
+        url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float = 30.0,
+    ) -> MCPRemoteConnection:
+        conn = MCPRemoteConnection(name, url=url, headers=headers)
+        try:
+            await asyncio.wait_for(conn.start(), timeout=timeout)
+            self._servers[name] = conn
+            return conn
+        except Exception as exc:
+            logger.warning(f"Failed to connect to remote MCP server '{name}': {exc}")
+            raise
+
+    async def register_tools(
+        self,
+        conn: _MCPConnectionBase,
+        tool_filter: set[str] | None = None,
+    ) -> list[ToolDefinition]:
+        raw_tools = await conn.discover_tools()
+        definitions: list[ToolDefinition] = []
+        server_prefix = _sanitize(conn.name)
+        for tool_spec in raw_tools:
+            original_name = tool_spec["name"]
+            if tool_filter and original_name not in tool_filter:
+                continue
+            prefixed_name = f"{server_prefix}_{_sanitize(original_name)}"
+            td = ToolDefinition(
+                name=prefixed_name,
+                description=f"[MCP:{conn.name}] {tool_spec.get('description', '')}",
+                parameters=tool_spec.get("inputSchema", {"type": "object", "properties": {}}),
+                func=self._make_caller(conn, original_name),
+                category="mcp",
+            )
+            definitions.append(td)
+        return definitions
+
+    @staticmethod
+    def _make_caller(conn: _MCPConnectionBase, tool_name: str) -> Any:
+        async def _call(**kwargs: Any) -> dict[str, Any]:
+            return await conn.call_tool(tool_name, kwargs)
+        return _call
+
+    async def shutdown(self) -> None:
+        for conn in self._servers.values():
+            try:
+                await conn.stop()
+            except Exception as exc:
+                logger.warning(f"Error stopping MCP server '{conn.name}': {exc}")
+        self._servers.clear()