diff --git a/main/xiaozhi-server/core/mcp/MCPClient.py b/main/xiaozhi-server/core/mcp/MCPClient.py index 7996ff67..2ae34c29 100644 --- a/main/xiaozhi-server/core/mcp/MCPClient.py +++ b/main/xiaozhi-server/core/mcp/MCPClient.py @@ -1,94 +1,125 @@ +from __future__ import annotations + from datetime import timedelta -from typing import Optional +import asyncio, os, shutil, concurrent.futures from contextlib import AsyncExitStack -import os, shutil +from typing import Optional, List, Dict, Any + from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from mcp.client.sse import sse_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[str, Any]): self.logger = setup_logging() self.config = config - self.tolls = [] + + 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", []) + if self._worker_task: + return + self._worker_task = asyncio.create_task(self._worker(), name="MCPClientWorker") + await self._ready_evt.wait() - if "command" in self.config: - 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) - ) - elif "url" in self.config: - sse_transport = await self.exit_stack.enter_async_context(sse_client(self.config["url"])) - self.sse_read, self.sse_write = sse_transport - self.session = await self.exit_stack.enter_async_context( - ClientSession(read_stream=self.sse_read, write_stream=self.sse_write) - ) - else: - raise ValueError("MCPClient config must have 'command' or 'url'.") - - 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] - - 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 + 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 + 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"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): + if not self.session: + raise RuntimeError("MCPClient not initialized") + + loop = self._worker_task.get_loop() + coro = self.session.call_tool(name, args) + + 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) + + async def _worker(self): + 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: + sse_r, sse_w = await stack.enter_async_context(sse_client(self.config["url"])) + read_stream, write_stream = sse_r, sse_w + + else: + raise ValueError("MCPClient config must include 'command' or 'url'") + + self.session = await stack.enter_async_context( + ClientSession( + read_stream=read_stream, + write_stream=write_stream, + 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 a0e310fa..cc72228a 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()