"""服务端MCP客户端"""
|
|
from __future__ import annotations
|
|
from datetime import timedelta
|
import asyncio
|
import os
|
import shutil
|
import concurrent.futures
|
from contextlib import AsyncExitStack
|
from typing import Optional, List, Dict, Any
|
|
from mcp import ClientSession, StdioServerParameters, Implementation
|
from mcp.client.session import SamplingFnT, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT
|
from mcp.client.stdio import stdio_client
|
from mcp.client.sse import sse_client
|
from mcp.client.streamable_http import streamablehttp_client
|
from mcp.shared.session import ProgressFnT
|
|
from config.logger import setup_logging
|
from core.utils.util import sanitize_tool_name
|
|
TAG = __name__
|
|
|
class ServerMCPClient:
|
"""服务端MCP客户端，用于连接和管理MCP服务"""
|
|
def __init__(self, config: Dict[str, Any]):
|
"""初始化服务端MCP客户端
|
|
Args:
|
config: MCP服务配置字典
|
"""
|
self.logger = setup_logging()
|
self.config = config
|
|
self._worker_task: Optional[asyncio.Task] = None
|
self._ready_evt = asyncio.Event()
|
self._shutdown_evt = asyncio.Event()
|
|
self.session: Optional[ClientSession] = None
|
self.tools: List = [] # 原始工具对象
|
self.tools_dict: Dict[str, Any] = {}
|
self.name_mapping: Dict[str, str] = {}
|
|
async def initialize(self, read_timeout_seconds: timedelta | None = None,
|
sampling_callback: SamplingFnT | None = None,
|
elicitation_callback: ElicitationFnT | None = None,
|
list_roots_callback: ListRootsFnT | None = None,
|
logging_callback: LoggingFnT | None = None,
|
message_handler: MessageHandlerFnT | None = None,
|
client_info: Implementation | None = None):
|
"""初始化MCP客户端连接"""
|
if self._worker_task:
|
return
|
|
self._worker_task = asyncio.create_task(
|
self._worker(read_timeout_seconds=read_timeout_seconds,
|
sampling_callback=sampling_callback,
|
elicitation_callback=elicitation_callback,
|
list_roots_callback=list_roots_callback,
|
logging_callback=logging_callback,
|
message_handler=message_handler,
|
client_info=client_info), name="ServerMCPClientWorker"
|
)
|
await self._ready_evt.wait()
|
|
self.logger.bind(tag=TAG).info(
|
f"服务端MCP客户端已连接，可用工具: {[name for name in self.name_mapping.values()]}"
|
)
|
|
async def cleanup(self):
|
"""清理MCP客户端资源"""
|
if not self._worker_task:
|
return
|
|
self._shutdown_evt.set()
|
try:
|
await asyncio.wait_for(self._worker_task, timeout=20)
|
except (asyncio.TimeoutError, Exception) as e:
|
self.logger.bind(tag=TAG).error(f"服务端MCP客户端关闭错误: {e}")
|
finally:
|
self._worker_task = None
|
|
def has_tool(self, name: str) -> bool:
|
"""检查是否包含指定工具
|
|
Args:
|
name: 工具名称
|
|
Returns:
|
bool: 是否包含该工具
|
"""
|
return name in self.tools_dict
|
|
def get_available_tools(self) -> List[Dict[str, Any]]:
|
"""获取所有可用工具的定义
|
|
Returns:
|
List[Dict[str, Any]]: 工具定义列表
|
"""
|
return [
|
{
|
"type": "function",
|
"function": {
|
"name": name,
|
"description": tool.description,
|
"parameters": tool.inputSchema,
|
},
|
}
|
for name, tool in self.tools_dict.items()
|
]
|
|
async def call_tool(self, name: str, arguments: dict, read_timeout_seconds: timedelta | None = None, progress_callback: ProgressFnT | None = None, *, meta: dict[str, Any] | None = None) -> Any:
|
"""调用指定工具
|
|
Args:
|
name: 工具名称
|
arguments: 工具参数
|
read_timeout_seconds:
|
progress_callback: 进度回调函数
|
meta:
|
|
Returns:
|
Any: 工具执行结果
|
|
Raises:
|
RuntimeError: 客户端未初始化时抛出
|
"""
|
if not self.session:
|
raise RuntimeError("服务端MCP客户端未初始化")
|
|
real_name = self.name_mapping.get(name, name)
|
loop = self._worker_task.get_loop()
|
coro = self.session.call_tool(real_name, arguments=arguments, read_timeout_seconds=read_timeout_seconds, progress_callback=progress_callback, meta=meta)
|
|
if loop is asyncio.get_running_loop():
|
return await coro
|
|
fut: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(coro, loop)
|
return await asyncio.wrap_future(fut)
|
|
def is_connected(self) -> bool:
|
"""检查MCP客户端是否连接正常
|
|
Returns:
|
bool: 如果客户端已连接并正常工作，返回True，否则返回False
|
"""
|
# 检查工作任务是否存在
|
if self._worker_task is None:
|
return False
|
|
# 检查工作任务是否已经完成或取消
|
if self._worker_task.done():
|
return False
|
|
# 检查会话是否存在
|
if self.session is None:
|
return False
|
|
# 所有检查都通过，连接正常
|
return True
|
|
async def _worker(self, read_timeout_seconds: timedelta | None = None,
|
sampling_callback: SamplingFnT | None = None,
|
elicitation_callback: ElicitationFnT | None = None,
|
list_roots_callback: ListRootsFnT | None = None,
|
logging_callback: LoggingFnT | None = None,
|
message_handler: MessageHandlerFnT | None = None,
|
client_info: Implementation | None = None):
|
"""MCP客户端工作协程"""
|
async with AsyncExitStack() as stack:
|
try:
|
# 建立 StdioClient
|
if "command" in self.config:
|
cmd = (
|
shutil.which("npx")
|
if self.config["command"] == "npx"
|
else self.config["command"]
|
)
|
env = {**os.environ, **self.config.get("env", {})}
|
params = StdioServerParameters(
|
command=cmd,
|
args=self.config.get("args", []),
|
env=env,
|
)
|
stdio_r, stdio_w = await stack.enter_async_context(
|
stdio_client(params)
|
)
|
read_stream, write_stream = stdio_r, stdio_w
|
|
# 建立SSEClient
|
elif "url" in self.config:
|
headers = dict(self.config.get("headers", {}))
|
# TODO 兼容旧版本
|
if "API_ACCESS_TOKEN" in self.config:
|
headers["Authorization"] = f"Bearer {self.config['API_ACCESS_TOKEN']}"
|
self.logger.bind(tag=TAG).warning(f"你正在使用旧过时的配置 API_ACCESS_TOKEN ，请在.mcp_server_settings.json中将API_ACCESS_TOKEN直接设置在headers中，例如 'Authorization': 'Bearer API_ACCESS_TOKEN'")
|
|
# 根据transport类型选择不同的客户端，默认为SSE
|
transport_type = self.config.get("transport", "sse")
|
|
if transport_type == "streamable-http" or transport_type == "http":
|
# 使用 Streamable HTTP 传输
|
http_r, http_w, get_session_id = await stack.enter_async_context(
|
streamablehttp_client(
|
url=self.config["url"],
|
headers=headers,
|
timeout=self.config.get("timeout", 30),
|
sse_read_timeout=self.config.get("sse_read_timeout", 60 * 5),
|
terminate_on_close=self.config.get("terminate_on_close", True)
|
)
|
)
|
read_stream, write_stream = http_r, http_w
|
else:
|
# 使用传统的 SSE 传输
|
sse_r, sse_w = await stack.enter_async_context(
|
sse_client(
|
url=self.config["url"],
|
headers=headers,
|
timeout=self.config.get("timeout", 5),
|
sse_read_timeout=self.config.get("sse_read_timeout", 60 * 5)
|
)
|
)
|
read_stream, write_stream = sse_r, sse_w
|
|
else:
|
raise ValueError("MCP客户端配置必须包含'command'或'url'")
|
|
self.session = await stack.enter_async_context(
|
ClientSession(
|
read_stream=read_stream,
|
write_stream=write_stream,
|
read_timeout_seconds=read_timeout_seconds,
|
sampling_callback=sampling_callback,
|
elicitation_callback=elicitation_callback,
|
list_roots_callback=list_roots_callback,
|
logging_callback=logging_callback,
|
message_handler=message_handler,
|
client_info=client_info
|
)
|
)
|
await self.session.initialize()
|
|
# 获取工具
|
self.tools = (await self.session.list_tools()).tools
|
for t in self.tools:
|
sanitized = sanitize_tool_name(t.name)
|
self.tools_dict[sanitized] = t
|
self.name_mapping[sanitized] = t.name
|
|
self._ready_evt.set()
|
|
# 挂起等待关闭
|
await self._shutdown_evt.wait()
|
|
except Exception as e:
|
self.logger.bind(tag=TAG).error(f"服务端MCP客户端工作协程错误: {e}")
|
self._ready_evt.set()
|
raise
|
