import asyncio import base64 from collections.abc import AsyncGenerator, Callable from contextlib import AsyncExitStack, asynccontextmanager import re from typing import Any, Literal, cast import anyio from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream import httpx from mcp import ClientSession # type: ignore from mcp.client.sse import sse_client from mcp.client.stdio import StdioServerParameters, stdio_client from mcp.client.streamable_http import streamable_http_client from mcp.shared.message import SessionMessage from pydantic import ValidationError from zhenxun.services.ai.core.models import ToolDefinition from zhenxun.services.ai.core.stream_events import UserCustomEvent from zhenxun.services.ai.run import RunContext from zhenxun.services.ai.sandbox.addons.base import BaseMcpProxyExtension from zhenxun.services.ai.sandbox.models import SandboxBlueprint from zhenxun.services.ai.tools.core.tool import BaseTool from zhenxun.services.ai.tools.core.toolkit import BaseToolkit from zhenxun.services.ai.tools.models import ToolkitConfig, ToolResult from zhenxun.services.log import logger from zhenxun.utils.lifespan import LifespanManager from zhenxun.utils.pydantic_compat import model_dump _MESSAGE_START_CHARS = {"{", "["} _LITERAL_PREFIXES: tuple[str, ...] = ("true", "false", "null") def _should_ignore_exception(exc: Exception) -> bool: """ 判断该异常是否是由非 JSON 的脏数据标准输出引起的。 如果是脏数据，则可以安全忽略。 """ if not isinstance(exc, ValidationError): return False errors = exc.errors() first = next(iter(errors), None) if not first or first.get("type") != "json_invalid": return False input_value = first.get("input") if not isinstance(input_value, str): return False stripped = input_value.strip() if not stripped: return True first_char = stripped[0] lowered = stripped.lower() if first_char in _MESSAGE_START_CHARS or any( lowered.startswith(prefix) for prefix in _LITERAL_PREFIXES ): return False return True @asynccontextmanager async def filtered_stdio_client( server_name: str, server: StdioServerParameters ) -> AsyncGenerator[ tuple[ MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage], ], None, ]: """ 包裹官方的 stdio_client，拦截并过滤掉非 JSON 格式的标准输出噪音。 """ async with stdio_client(server=server) as (read_stream, write_stream): filtered_send, filtered_recv = anyio.create_memory_object_stream[ SessionMessage | Exception ](0) async def _forward_stdout() -> None: try: async with read_stream: async for item in read_stream: if isinstance(item, Exception) and _should_ignore_exception( item ): if isinstance(item, ValidationError): err_input = item.errors()[0].get("input", "") logger.debug( f"🔇 [MCP Stdout Filter] {server_name} 忽略脏数据: " f"{str(err_input)[:100].strip()}..." ) continue await filtered_send.send(item) except anyio.ClosedResourceError: pass finally: await filtered_send.aclose() async with anyio.create_task_group() as tg: tg.start_soon(_forward_stdout) try: yield filtered_recv, write_stream finally: tg.cancel_scope.cancel() class MCPRemoteTool(BaseTool): """远端 MCP 工具在本地的代理对象""" def __init__( self, name: str, original_tool_name: str, description: str, parameters: dict, toolkit: "MCPToolkit", ): super().__init__(name=name, description=description) self.original_tool_name = original_tool_name self.parameters = parameters self.toolkit = toolkit self.parent_toolkit = toolkit self.args_schema = None self.metadata = ( {"admin_level": toolkit.admin_level} if toolkit.admin_level > 0 else {} ) @property def effective_ttl(self) -> float: """动态获取当前最新的 TTL 配置，支持无缝热重载""" try: from zhenxun.services.ai.config import get_llm_config mcp_ttl = get_llm_config().agent_settings.mcp_cleanup_timeout return 31536000.0 if mcp_ttl <= 0 else float(mcp_ttl) except Exception: return 31536000.0 if self.toolkit.ttl <= 0 else float(self.toolkit.ttl) async def get_definition( self, context: RunContext | None = None ) -> ToolDefinition | None: if hasattr(self, "_dynamic_def") and self._dynamic_def is not None: return self._dynamic_def tool_def = ToolDefinition( name=self.name, description=self.description, parameters=self.parameters, metadata=self.metadata or {}, ) if context and self.settings.capabilities: from zhenxun.services.ai.capabilities import ( CombinedCapability, ) combined_cap = CombinedCapability(self.settings.capabilities) defs = await combined_cap.prepare_tools(context, [tool_def]) if not defs: return None tool_def = defs[0] return tool_def async def execute(self, context: RunContext | None = None, **kwargs) -> ToolResult: max_retries = 3 last_error = None for attempt in range(max_retries): await self.toolkit.lifespan_manager.touch( self.toolkit.server_name, self.toolkit.effective_ttl ) session = await self.toolkit.get_session(context) if not session: return ToolResult(output="MCP 连接错误").as_error() try: result = await session.call_tool(self.original_tool_name, kwargs) break except Exception as e: last_error = e logger.warning( f"MCP 工具 '{self.name}' 执行失败 " f"(尝试 {attempt + 1}/{max_retries}): {e}。" "正在尝试自动重连自愈..." ) if attempt < max_retries - 1: await self.toolkit.close() await asyncio.sleep(2**attempt) else: logger.error( f"MCP 工具 '{self.name}' " f"在重试 {max_retries} 次后仍然执行失败: {last_error}" ) return ToolResult(output=f"MCP 错误: {last_error}").as_error() if result.isError: return ToolResult(output=str(result.content)).as_error() from zhenxun.services.ai.core.messages import ImagePart, TextPart output_content = [] img_count = 0 for item in result.content: item_type = getattr(item, "type", "text") if item_type == "image": b64_data = getattr(item, "data", "") mime_type = getattr(item, "mimeType", "image/png") if b64_data: try: img_bytes = base64.b64decode(b64_data) output_content.append( ImagePart(raw=img_bytes, mime_type=mime_type) ) img_count += 1 continue except Exception as e: logger.warning(f"MCP 图片 Base64 解码失败: {e}") output_content.append(TextPart(text="[图片解码失败]")) elif item_type == "text": text = getattr(item, "text", str(item)) output_content.append(TextPart(text=text)) md_images = re.findall(r"!\[.*?\]\((https?://[^\)]+)\)", text) for img_url in md_images: output_content.append(ImagePart(url=img_url)) img_count += 1 else: dumped = model_dump(item) if hasattr(item, "model_dump") else str(item) output_content.append(TextPart(text=str(dumped))) tool_result = ToolResult(output=output_content) logger.info( f"获取到 {len(output_content)} 条返回数据，提取了 {img_count} 张图片" ) if img_count > 0 and context and context.run.event_bus: await context.run.event_bus.emit(UserCustomEvent(display=output_content)) return tool_result class MCPToolkit(BaseToolkit): """模型上下文协议 (MCP) 的工具箱封装 (支持声明式挂载与动态隔离)""" def __init__( self, server_name: str, prefix: str | None = None, transport: Literal[ "stdio", "sse", "streamable-http", "sandbox_proxy" ] = "stdio", command: str | None = None, args: list[str] | None = None, url: str | None = None, env: dict | None = None, cwd: str | None = None, install_command: str | None = None, isolation: Literal["shared", "per_session"] = "shared", timeout: int = 30, admin_level: int = 0, header_provider: Callable[[RunContext], dict[str, str]] | None = None, env_provider: Callable[[RunContext], dict[str, str]] | None = None, ttl: int = 600, sandbox_session_id: str | None = None, sandbox_blueprint: SandboxBlueprint | None = None, ): super().__init__( config=ToolkitConfig(prefix=prefix) if prefix is not None else None ) self.server_name = server_name self.transport = transport self.command = command self.args = args or [] self.url = url self.env = env or {} self.cwd = cwd self.install_command = install_command self.timeout = timeout self.admin_level = admin_level self.header_provider = header_provider self.env_provider = env_provider self.sandbox_session_id = sandbox_session_id self.sandbox_blueprint = sandbox_blueprint self.ttl = ttl self._shared_session: ClientSession | None = None self._tools: list[BaseTool] = [] self._is_initialized = False self.lifespan_manager = LifespanManager() self._shared_task: asyncio.Task | None = None self._stop_event: asyncio.Event = asyncio.Event() self._ready_event: asyncio.Event = asyncio.Event() self._init_exception: Exception | None = None @property def effective_ttl(self) -> float: """动态获取当前最新的 TTL 配置，支持无缝热重载""" try: from zhenxun.services.ai.config import get_llm_config mcp_ttl = get_llm_config().agent_settings.mcp_cleanup_timeout return 31536000.0 if mcp_ttl <= 0 else float(mcp_ttl) except Exception: return 31536000.0 if self.ttl <= 0 else float(self.ttl) async def _run_install_command(self, force: bool = False): """执行环境预安装（热加载依赖）""" if not self.install_command or not self.cwd: return from pathlib import Path marker_file = Path(self.cwd) / ".zx_installed" heuristic_missing = False if ( "npm" in self.install_command and not (Path(self.cwd) / "node_modules").exists() ): heuristic_missing = True if not force and not heuristic_missing and marker_file.exists(): return reason = "强制触发自愈重装" if force else "首次启动或检测到环境缺失" logger.info( f"🔧 [{self.server_name}] {reason}，正在执行: `{self.install_command}` ..." ) try: proc = await asyncio.create_subprocess_shell( self.install_command, cwd=self.cwd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await proc.communicate() if proc.returncode != 0: err_msg = ( stderr.decode("utf-8", errors="ignore") if stderr else stdout.decode("utf-8", errors="ignore") ) logger.error( f"❌ [{self.server_name}] 环境安装失败!\n错误输出:\n{err_msg}" ) raise RuntimeError(f"安装命令执行失败: {self.install_command}") else: marker_file.touch() logger.info(f"✅ [{self.server_name}] 环境装配成功！") except Exception as e: logger.error(f"执行预热安装异常: {e}") raise e async def get_session( self, context: RunContext | None = None ) -> ClientSession | None: await self.lifespan_manager.register( self.server_name, ttl=self.effective_ttl, cleanup_callback=self.release_resource, ) if not self._is_initialized: await self.initialize(context) return self._shared_session async def _spawn_session_task(self, dynamic_headers: dict, dynamic_env: dict): """生成后台连接任务""" max_attempts = 2 if (self.transport == "stdio" and self.install_command) else 1 try: for attempt in range(1, max_attempts + 1): try: if self.transport == "stdio": await self._run_install_command(force=(attempt > 1)) async with AsyncExitStack() as stack: read_stream: Any = None write_stream: Any = None if self.transport == "stdio": if not self.command: raise ValueError("stdio requires 'command'") params = StdioServerParameters( command=self.command, args=self.args, env=dynamic_env, cwd=self.cwd, ) transport_ctx = filtered_stdio_client( self.server_name, params ) elif self.transport == "sse": if not self.url: raise ValueError("sse requires 'url'") transport_ctx = sse_client( url=self.url, headers=dynamic_headers, timeout=self.timeout, ) elif self.transport == "streamable-http": if not self.url: raise ValueError("streamable-http requires 'url'") http_client = httpx.AsyncClient( headers=dynamic_headers, timeout=self.timeout ) await stack.enter_async_context(http_client) transport_ctx = streamable_http_client( url=self.url, http_client=http_client ) elif self.transport == "sandbox_proxy": from zhenxun.services.ai.sandbox.manager import ( sandbox_manager, ) target_session_id = ( self.sandbox_session_id or "mcp_global_session" ) bp = self.sandbox_blueprint or SandboxBlueprint() bp.enable_network = True driver = await sandbox_manager.get_or_create_session( target_session_id, blueprint=bp, ) plugin_name = "universal_mcp" await driver.mount_extension(plugin_name) mcp_plugin = cast( BaseMcpProxyExtension, driver.get_extension(plugin_name) ) if not self.command: raise ValueError("sandbox_proxy requires 'command'") streams = await stack.enter_async_context( mcp_plugin.connect_mcp( self.command, self.args, dynamic_env ) ) read_stream, write_stream = streams[0], streams[1] transport_ctx = None else: raise ValueError(f"Unknown transport: {self.transport}") if transport_ctx: transport = await stack.enter_async_context(transport_ctx) read_stream, write_stream = transport[0], transport[1] client_session = await stack.enter_async_context( ClientSession(read_stream, write_stream) ) await client_session.initialize() self._shared_session = client_session if not self._is_initialized: mcp_tools_res = await client_session.list_tools() tools_list = list(mcp_tools_res.tools) cursor = getattr(mcp_tools_res, "nextCursor", None) while cursor: mcp_tools_res = await client_session.list_tools( cursor=cursor ) tools_list.extend(mcp_tools_res.tools) cursor = getattr(mcp_tools_res, "nextCursor", None) for t in tools_list: t_name = t.name.replace("-", "_") final_name = ( f"{self.config.prefix}{t_name}" if self.config.prefix else t_name ) self._tools.append( MCPRemoteTool( name=final_name, original_tool_name=t.name, description=t.description or "", parameters=t.inputSchema, toolkit=self, ) ) self._is_initialized = True self._ready_event.set() logger.info( f"成功连接 MCP 服务器: {self.server_name}, " f"获取了 {len(self._tools)} 个工具" ) await self._stop_event.wait() return except Exception as e: from zhenxun.services.ai.core.exceptions import SandboxFatalError if isinstance(e, SandboxFatalError): raise e if attempt < max_attempts: logger.warning( f"⚠️ [{self.server_name}] 进程启动或运行异常崩溃，" "疑似环境损坏。触发自愈机制 (准备重试)..." ) self._init_exception = e self._shared_session = None if not self._is_initialized: self._is_initialized = False continue raise e except Exception as e: self._init_exception = e logger.error( f"MCP 服务器 '{self.server_name}' " f"初始化连接失败（模式: {self.transport}）。" f"错误原因: {e}" ) if "Connection closed" in str(e): logger.error( "提示：若是使用沙箱隧道，请进入 Docker Desktop 检查容器的日志。" "多半是因为 npx 命令执行出错" "（例如：镜像中缺少 Node 环境或国内网络无法连接 npm 等）。" ) finally: if not self._is_initialized: self._is_initialized = False self._shared_session = None self._ready_event.set() async def initialize(self, context: RunContext | None = None): if self._is_initialized: return self._tools.clear() self._stop_event.clear() self._ready_event.clear() self._init_exception = None dynamic_headers = {} dynamic_env = self.env.copy() self._shared_task = asyncio.create_task( self._spawn_session_task(dynamic_headers, dynamic_env) ) await self._ready_event.wait() if self._init_exception: raise self._init_exception async def get_tools(self, context: RunContext | None = None) -> dict[str, BaseTool]: await self.lifespan_manager.register( self.server_name, ttl=self.effective_ttl, cleanup_callback=self.release_resource, ) if not self._is_initialized: await self.initialize(context) tools_dict = {} for t in self._tools: tools_dict[t.name] = t return tools_dict async def release_resource(self, resource_id: str): await self.close() async def close(self): current_task = asyncio.current_task() if ( self.lifespan_manager._watchdog_task and self.lifespan_manager._watchdog_task is not current_task ): await self.lifespan_manager.stop() self._stop_event.set() self._is_initialized = False self._shared_session = None self._tools.clear() if ( self._shared_task and not self._shared_task.done() and self._shared_task is not current_task ): try: await asyncio.wait_for(self._shared_task, timeout=5.0) except asyncio.TimeoutError: logger.warning( f"MCP 服务器 '{self.server_name}' 的任务执行超时，正在取消。" ) self._shared_task.cancel() self._shared_task = None