148 lines
5.0 KiB
Python
148 lines
5.0 KiB
Python
"""
|
||
MaiSaka - 单个 MCP 服务器连接管理
|
||
封装单个 MCP 服务器的连接生命周期：连接 → 发现工具 → 调用工具 → 断开。
|
||
"""
|
||
|
||
from contextlib import AsyncExitStack
|
||
from typing import Any, Optional
|
||
|
||
from ..config import console
|
||
from .config import MCPServerConfig
|
||
|
||
# ──────────────────── MCP SDK 可选导入 ────────────────────
|
||
#
|
||
# mcp 是可选依赖。如果未安装，MCP_AVAILABLE = False，
|
||
# MCPManager.from_config() 会检测到并返回 None，不影响主程序运行。
|
||
|
||
try:
|
||
from mcp import ClientSession
|
||
|
||
try:
|
||
from mcp.client.stdio import StdioServerParameters
|
||
except ImportError:
|
||
from mcp import StdioServerParameters # type: ignore[attr-defined]
|
||
|
||
from mcp.client.stdio import stdio_client
|
||
|
||
MCP_AVAILABLE = True
|
||
except ImportError:
|
||
MCP_AVAILABLE = False
|
||
ClientSession = None # type: ignore[assignment,misc]
|
||
StdioServerParameters = None # type: ignore[assignment,misc]
|
||
stdio_client = None # type: ignore[assignment]
|
||
|
||
try:
|
||
from mcp.client.sse import sse_client
|
||
|
||
SSE_AVAILABLE = True
|
||
except ImportError:
|
||
SSE_AVAILABLE = False
|
||
sse_client = None # type: ignore[assignment]
|
||
|
||
|
||
class MCPConnection:
|
||
"""
|
||
管理单个 MCP 服务器的连接生命周期。
|
||
|
||
支持两种传输方式：
|
||
- Stdio: 启动子进程，通过 stdin/stdout 通信
|
||
- SSE: 连接远程 HTTP SSE 端点
|
||
"""
|
||
|
||
def __init__(self, config: MCPServerConfig):
|
||
self.config = config
|
||
self.session: Optional[Any] = None # mcp.ClientSession
|
||
self.tools: list = [] # mcp Tool objects
|
||
self._exit_stack = AsyncExitStack()
|
||
|
||
async def connect(self) -> bool:
|
||
"""
|
||
连接到 MCP 服务器并发现可用工具。
|
||
|
||
Returns:
|
||
True 表示连接成功，False 表示失败。
|
||
"""
|
||
if not MCP_AVAILABLE:
|
||
console.print("[warning]⚠️ 未安装 mcp SDK，请运行: pip install mcp[/warning]")
|
||
return False
|
||
|
||
try:
|
||
await self._exit_stack.__aenter__()
|
||
|
||
if self.config.transport_type == "stdio":
|
||
read_stream, write_stream = await self._connect_stdio()
|
||
elif self.config.transport_type == "sse":
|
||
read_stream, write_stream = await self._connect_sse()
|
||
else:
|
||
console.print(f"[warning]MCP '{self.config.name}': 未知传输类型[/warning]")
|
||
return False
|
||
|
||
# 创建并初始化 MCP 会话
|
||
self.session = await self._exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
|
||
await self.session.initialize()
|
||
|
||
# 发现工具
|
||
result = await self.session.list_tools()
|
||
self.tools = result.tools if hasattr(result, "tools") else []
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
console.print(f"[warning]⚠️ MCP 服务器 '{self.config.name}' 连接失败: {e}[/warning]")
|
||
await self.close()
|
||
return False
|
||
|
||
async def _connect_stdio(self):
|
||
"""建立 Stdio 传输连接。"""
|
||
params = StdioServerParameters(
|
||
command=self.config.command,
|
||
args=self.config.args,
|
||
env=self.config.env,
|
||
)
|
||
return await self._exit_stack.enter_async_context(stdio_client(params))
|
||
|
||
async def _connect_sse(self):
|
||
"""建立 SSE 传输连接。"""
|
||
if not SSE_AVAILABLE:
|
||
raise ImportError("SSE 传输需要额外依赖，请运行: pip install mcp[sse]")
|
||
return await self._exit_stack.enter_async_context(sse_client(url=self.config.url, headers=self.config.headers))
|
||
|
||
async def call_tool(self, tool_name: str, arguments: dict) -> str:
|
||
"""
|
||
调用 MCP 工具并返回结果文本。
|
||
|
||
Args:
|
||
tool_name: 工具名称
|
||
arguments: 工具参数字典
|
||
|
||
Returns:
|
||
工具执行结果的文本表示。
|
||
"""
|
||
if not self.session:
|
||
return f"MCP 服务器 '{self.config.name}' 未连接"
|
||
|
||
result = await self.session.call_tool(tool_name, arguments=arguments)
|
||
|
||
# 将结果内容转换为文本
|
||
parts: list[str] = []
|
||
for content in result.content:
|
||
if hasattr(content, "text"):
|
||
parts.append(content.text)
|
||
elif hasattr(content, "data"):
|
||
# 二进制/图片内容，展示类型信息
|
||
content_type = getattr(content, "mimeType", "unknown")
|
||
parts.append(f"[{content_type} 二进制内容]")
|
||
elif hasattr(content, "type"):
|
||
parts.append(f"[{content.type} 内容]")
|
||
|
||
return "\n".join(parts) if parts else "工具执行成功（无输出）"
|
||
|
||
async def close(self):
|
||
"""关闭连接并释放资源。"""
|
||
try:
|
||
await self._exit_stack.aclose()
|
||
except Exception:
|
||
pass
|
||
self.session = None
|
||
self.tools = []
|