zhenxun_bot/zhenxun/services/ai/tools/providers/mcp/toolkit.py

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