mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-10-09 22:00:01 +08:00
33d6ea1335
检查bot是否运行正常 / bot check (push) Canceled after 0s
CodeQL Code Security Analysis / Analyze (python) (push) Canceled after 0s
Ruff Lint / Ruff Lint (push) Canceled after 0s
Sequential Lint and Type Check / ruff-call (push) Canceled after 0s
Release Drafter / Update Release Draft (push) Canceled after 0s
Force Sync to Aliyun / sync (push) Canceled after 0s
Update Version / update-version (push) Canceled after 0s
Sequential Lint and Type Check / pyright-call (push) Canceled after 0s
* ✨ feat(llm): 新增 Grok 适配器及 API 密钥重置指令 - 新增 Grok 适配器并支持原生推特检索工具 `x_search` - 优化 DeepSeek 适配器以迁移至 Responses 协议，增强多模态支持 - 在 LLM 管理插件中新增 `llm reset` 子命令以重置 Key 状态 - 增强 MCP 连接异常处理及流关闭容错能力 * 🚨 auto fix by pre-commit hooks --------- Co-authored-by: webjoin111 <455457521@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
672 lines
26 KiB
Python
672 lines
26 KiB
Python
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.context 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.ai.utils.logger import log_tool as 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
|
||
except BaseException as e:
|
||
logger.debug(
|
||
f"🔇 [MCP Stdout Filter] 捕获到底层流异常: {type(e).__name__}: {e}"
|
||
)
|
||
finally:
|
||
try:
|
||
await filtered_send.aclose()
|
||
except Exception:
|
||
pass
|
||
|
||
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",
|
||
):
|
||
"""
|
||
初始化远端 MCP 工具在本地的代理对象。
|
||
|
||
参数:
|
||
name: 工具在本地注册的唯一标识名称。
|
||
original_tool_name: 该工具在远端 MCP 服务器上的原始名称。
|
||
description: 工具的描述信息，用于指导大模型选用此工具。
|
||
parameters: 工具参数的 JSON 模式描述。
|
||
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,
|
||
headers: dict[str, str] | None = None,
|
||
env: dict | None = None,
|
||
cwd: str | None = None,
|
||
install_command: str | None = None,
|
||
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,
|
||
):
|
||
"""
|
||
初始化模型上下文协议 (MCP) 的工具箱封装。
|
||
|
||
参数:
|
||
server_name: MCP 服务器的唯一标识名称。
|
||
prefix: 工具名称的前缀，防止工具命名冲突。
|
||
transport: 连接 MCP 服务器的通信传输协议，可选 "stdio", "sse", "streamable-http", "sandbox_proxy"。
|
||
command: 用于 stdio 或 sandbox_proxy 模式启动服务器的可执行命令。
|
||
args: 启动服务器时附加的命令行参数。
|
||
url: 用于 sse 或 streamable-http 模式的服务器连接 URL。
|
||
headers: 静态配置的 HTTP 请求头。
|
||
env: 启动服务器时的进程环境变量。
|
||
cwd: 启动服务器时的进程工作目录。
|
||
install_command: 首次启动前执行的环境热装配/依赖安装命令。
|
||
timeout: 初始化和请求网络接口时的超时秒数，默认 30。
|
||
admin_level: 调用工具所需的管理权限等级，默认 0 (无限制)。
|
||
header_provider: 动态生成 HTTP 头部信息的工厂函数。
|
||
env_provider: 动态生成进程环境变量的工厂函数。
|
||
ttl: 闲置清理的生存时间 (秒)，默认 600。
|
||
sandbox_session_id: 当 transport 为 sandbox_proxy 时指定的沙箱会话 ID。
|
||
sandbox_blueprint: 用于沙箱自动装配的环境蓝图配置。
|
||
""" # noqa: E501
|
||
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.headers = headers or {}
|
||
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):
|
||
"""在本地运行预热依赖安装命令以准备 MCP 服务器环境"""
|
||
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:
|
||
"""获取或建立与 MCP 服务器的会话连接，并刷新其生存状态"""
|
||
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):
|
||
"""在后台异步拉起 MCP 传输协议通道并进行服务初始化"""
|
||
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 BaseException as e:
|
||
if isinstance(
|
||
e, asyncio.CancelledError | KeyboardInterrupt | SystemExit
|
||
):
|
||
raise 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}] 进程启动或"
|
||
f"底层流异常崩溃 ({type(e).__name__})，"
|
||
"疑似环境损坏。触发自愈机制 (准备重试)..."
|
||
)
|
||
self._init_exception = (
|
||
e if isinstance(e, Exception) else Exception(str(e))
|
||
)
|
||
self._shared_session = None
|
||
if not self._is_initialized:
|
||
self._is_initialized = False
|
||
continue
|
||
|
||
if not isinstance(e, Exception):
|
||
raise Exception(f"底层致命异常: {e}") from e
|
||
raise e
|
||
|
||
except BaseException as e:
|
||
if isinstance(e, asyncio.CancelledError | KeyboardInterrupt | SystemExit):
|
||
raise e
|
||
self._init_exception = e if isinstance(e, Exception) else Exception(str(e))
|
||
logger.error(
|
||
f"MCP 服务器 '{self.server_name}' "
|
||
f"初始化连接或运行期间发生底层故障（模式: {self.transport}）。"
|
||
f"错误类型: {type(e).__name__}, 原因: {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):
|
||
"""清空当前工具集并异步执行 MCP 服务端的全套初始化流程"""
|
||
if self._is_initialized:
|
||
return
|
||
|
||
self._tools.clear()
|
||
self._stop_event.clear()
|
||
self._ready_event.clear()
|
||
self._init_exception = None
|
||
|
||
dynamic_headers = self.headers.copy()
|
||
if self.header_provider and context:
|
||
dynamic_headers.update(self.header_provider(context))
|
||
|
||
dynamic_env = self.env.copy()
|
||
if self.env_provider and context:
|
||
dynamic_env.update(self.env_provider(context))
|
||
|
||
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):
|
||
"""彻底关闭 MCP 工具箱，注销生存期并释放底层 WebSocket/进程通道"""
|
||
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
|