LZH-YS1998
300 lines
12 KiB
Python
300 lines
12 KiB
Python
"""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()
|