From 0a765f4aaca98486f06bdc83122f8a5996b94ddd Mon Sep 17 00:00:00 2001 From: caixypromise Date: Sun, 4 May 2025 23:01:30 +0800 Subject: [PATCH] =?UTF-8?q?fix(MCP):=20=E9=87=8D=E6=9E=84MCPClient?= =?UTF-8?q?=E4=B8=BA=E5=90=8E=E5=8F=B0=E5=8D=8F=E7=A8=8B=20+=20AsyncExitSt?= =?UTF-8?q?ack=E7=AE=A1=E7=90=86=EF=BC=8C=E8=A7=A3=E5=86=B3=E8=BF=9B?= =?UTF-8?q?=E7=A8=8B=E9=80=80=E5=87=BA=E6=97=B6=E7=9A=84=E2=80=9CAttempted?= =?UTF-8?q?=20to=20exit=20cancel=20scope=20in=20a=20different=20task?= =?UTF-8?q?=E2=80=9D=E9=94=99=E8=AF=AF=20#=20=E5=8F=98=E6=9B=B4=20----=20-?= =?UTF-8?q?=20=E5=B0=86=E6=89=80=E6=9C=89stdio=5Fclient=E4=B8=8EClientSess?= =?UTF-8?q?ion=E7=9A=84=E5=88=9B=E5=BB=BA/=E9=94=80=E6=AF=81=E9=83=BD?= =?UTF-8?q?=E6=94=BE=E5=88=B0=E5=90=8C=E4=B8=80=E4=B8=AA=E5=90=8E=E5=8F=B0?= =?UTF-8?q?=20task=20=E4=B8=AD=20-=20=E4=BD=BF=E7=94=A8AsyncExitStack?= =?UTF-8?q?=E6=89=98=E7=AE=A1=E5=BC=82=E6=AD=A5=E8=B5=84=E6=BA=90=EF=BC=8C?= =?UTF-8?q?cleanup=E6=97=B6=E5=9C=A8=E5=90=8C=E4=B8=80task=E5=86=85?= =?UTF-8?q?=E6=89=A7=E8=A1=8Cexit=5Fstack.aclose()=20-=20=E5=A4=96?= =?UTF-8?q?=E9=83=A8=E5=8F=AA=E9=80=9A=E8=BF=87=E4=BA=8B=E4=BB=B6=E9=80=9A?= =?UTF-8?q?=E7=9F=A5=E5=90=8E=E5=8F=B0task=E9=80=80=E5=87=BA=EF=BC=8C?= =?UTF-8?q?=E9=81=BF=E5=85=8D=E8=B7=A8=E5=8D=8F=E7=A8=8B=E8=B0=83=E7=94=A8?= =?UTF-8?q?cancel-scope=E5=BC=82=E5=B8=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main/xiaozhi-server/core/mcp/MCPClient.py | 197 ++++++++++++++-------- main/xiaozhi-server/core/mcp/manager.py | 17 +- 2 files changed, 137 insertions(+), 77 deletions(-) diff --git a/main/xiaozhi-server/core/mcp/MCPClient.py b/main/xiaozhi-server/core/mcp/MCPClient.py index 103ac645..e2517f92 100644 --- a/main/xiaozhi-server/core/mcp/MCPClient.py +++ b/main/xiaozhi-server/core/mcp/MCPClient.py @@ -1,86 +1,145 @@ +from __future__ import annotations + +import asyncio, os, shutil, concurrent.futures from datetime import timedelta -from typing import Optional from contextlib import AsyncExitStack -import os, shutil +from typing import Optional, List + from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client - from config.logger import setup_logging TAG = __name__ + class MCPClient: - def __init__(self, config): - # Initialize session and client objects - self.session: Optional[ClientSession] = None - self.exit_stack = AsyncExitStack() + def __init__(self, config: dict): self.logger = setup_logging() self.config = config - self.tolls = [] + + # Back‑worker task & 状态同步 + 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 = [] async def initialize(self): - args = self.config.get("args", []) + """ + 启动后台 task，并等待其就绪（拿到 `tools`）。 + """ + if self._worker_task: + return # 已经 init 过 - command = ( - shutil.which("npx") - if self.config["command"] == "npx" - else self.config["command"] - ) - - env={**os.environ} - if self.config.get("env"): - env.update(self.config["env"]) - - server_params = StdioServerParameters( - command=command, - args=args, - env=env - ) - - stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params)) - self.stdio, self.write = stdio_transport - time_out_delta = timedelta(seconds=15) - self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream=self.stdio, write_stream=self.write, read_timeout_seconds=time_out_delta)) - - await self.session.initialize() - - # List available tools - response = await self.session.list_tools() - tools = response.tools - self.tools = tools - self.logger.bind(tag=TAG).info(f"Connected to server with tools:{[tool.name for tool in tools]}") - - def has_tool(self, tool_name): - return any(tool.name == tool_name for tool in self.tools) - - def get_available_tools(self): - available_tools = [{"type": "function", "function":{ - "name": tool.name, - "description": tool.description, - "parameters": tool.inputSchema - } } for tool in self.tools] + # 在当前 loop 创建后台 task + self._worker_task = asyncio.create_task(self._worker(), name="MCPClientWorker") + await self._ready_evt.wait() # 等待 worker 初始化完成 - return available_tools - - async def call_tool(self, tool_name: str, tool_args: dict): - self.logger.bind(tag=TAG).info(f"MCPClient Calling tool {tool_name} with args: {tool_args}") - try: - response = await self.session.call_tool(tool_name, tool_args) - except Exception as e: - self.logger.bind(tag=TAG).error(f"Error calling tool {tool_name}: {e}") - from types import SimpleNamespace - error_content = SimpleNamespace( - type='text', - text=f"Error calling tool {tool_name}: {e}" - ) - error_response = SimpleNamespace( - content=[error_content], - isError=True - ) - return error_response - self.logger.bind(tag=TAG).info(f"MCPClient Response from tool {tool_name}: {response}") - return response + # 此时 tools 已填充 + self.logger.bind(tag=TAG).info( + f"Connected, tools = {[t.name for t in self.tools]}" + ) async def cleanup(self): - """Clean up resources""" - await self.exit_stack.aclose() \ No newline at end of file + """ + 对外关闭接口： + · 只负责发出 “关机信号” + · 等待后台 task 正常退出 + 在任何 loop / task 调用都安全。 + """ + if not self._worker_task: + return + + self._shutdown_evt.set() # 发信号 + try: + await asyncio.wait_for(self._worker_task, timeout=15) + except (asyncio.TimeoutError, Exception) as e: + self.logger.bind(tag=TAG).error(f"worker shutdown err: {e}") + finally: + self._worker_task = None + + # ----------------------------- 工具接口 ----------------------------- + + def has_tool(self, name: str) -> bool: + return any(t.name == name for t in self.tools) + + def get_available_tools(self): + return [ + { + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.inputSchema, + }, + } + for t in self.tools + ] + + async def call_tool(self, name: str, args: dict): + """ + 转发到 session.call_tool。 + 若在 worker 之外的 task 调用，会通过 run_coroutine_threadsafe + 投递到 worker 所在 loop 中执行，保证线程安全。 + """ + if not self.session: # 尚未就绪 + raise RuntimeError("MCPClient not initialized") + + loop = self._worker_task.get_loop() + coro = self.session.call_tool(name, args) + + # 在同一个 loop ➜ 直接 await + if loop is asyncio.get_running_loop(): + return await coro + + # 跨 loop ➜ run_coroutine_threadsafe + fut: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(coro, loop) + return await asyncio.wrap_future(fut) + + # ----------------------------- 后台 task ----------------------------- + + async def _worker(self): + """ + 单线程协程： + 1. 创建所有异步资源 + 2. set_ready → 供外部使用 + 3. 等待 shutdown_evt + 4. 自动随 AsyncExitStack 退出而清理资源 + """ + async with AsyncExitStack() as stack: + try: + # ---------- 启动后端进程 ---------- + 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)) + + # ---------- 会话 ---------- + self.session = await stack.enter_async_context( + ClientSession( + read_stream=stdio_r, + write_stream=stdio_w, + read_timeout_seconds=timedelta(seconds=15), + ) + ) + await self.session.initialize() + + # ---------- 工具 ---------- + self.tools = (await self.session.list_tools()).tools + + # 初始化完成，放行外部 + self._ready_evt.set() + + # ---------- 挂起等待关闭 ---------- + await self._shutdown_evt.wait() + + except Exception as e: + self.logger.bind(tag=TAG).error(f"worker error: {e}") + self._ready_evt.set() # 确保外部不会卡死 + raise diff --git a/main/xiaozhi-server/core/mcp/manager.py b/main/xiaozhi-server/core/mcp/manager.py index f212b67a..0a35f357 100644 --- a/main/xiaozhi-server/core/mcp/manager.py +++ b/main/xiaozhi-server/core/mcp/manager.py @@ -1,5 +1,5 @@ """MCP服务管理器""" - +import asyncio import os, json from typing import Dict, Any, List from .MCPClient import MCPClient @@ -94,8 +94,8 @@ class MCPManager: """ for tool in self.tools: if ( - tool.get("function") != None - and tool["function"].get("name") == tool_name + tool.get("function") != None + and tool["function"].get("name") == tool_name ): return True return False @@ -120,12 +120,13 @@ class MCPManager: raise ValueError(f"Tool {tool_name} not found in any MCP server") async def cleanup_all(self) -> None: - for name, client in self.client.items(): + """依次关闭所有 MCPClient，不让异常阻断整体流程。""" + for name, client in list(self.client.items()): try: - await client.cleanup() - self.logger.bind(tag=TAG).info(f"Cleaned up MCP client: {name}") - except Exception as e: + await asyncio.wait_for(client.cleanup(), timeout=20) + self.logger.bind(tag=TAG).info(f"MCP client closed: {name}") + except (asyncio.TimeoutError, Exception) as e: self.logger.bind(tag=TAG).error( - f"Error cleaning up MCP client {name}: {e}" + f"Error closing MCP client {name}: {e}" ) self.client.clear()