mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 03:01:58 +00:00
ThearchyHelios
342 lines
17 KiB
Python
342 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import time
|
||
from typing import Any
|
||
|
||
from langchain_core.tools import BaseTool
|
||
from sqlalchemy import select
|
||
|
||
from app.capabilities import secrets as vault
|
||
from app.db.base import SessionLocal
|
||
from app.db.models import McpServer
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _server_config(row: McpServer) -> dict[str, Any]:
|
||
"""连接配置。env、headers 的值存的是密文（A6 第 6 条），这里解开交给客户端；接口只回掩码。
|
||
绑定了登记端点的 http 服务，地址和认证头在连接前由 _endpoint_config 换成端点的。"""
|
||
if row.transport == "http":
|
||
config: dict[str, Any] = {"transport": "streamable_http", "url": row.url or ""}
|
||
if row.headers:
|
||
config["headers"] = vault.reveal(row.headers)
|
||
return config
|
||
return {
|
||
"transport": "stdio",
|
||
"command": row.command or "",
|
||
"args": list(row.args or []),
|
||
"env": vault.reveal(row.env),
|
||
}
|
||
|
||
|
||
def _config_key(row: McpServer) -> str:
|
||
"""一个服务的连接配置，缓存据此判断手上的工具是不是按现在这份配置建的。按存库的样子算（密文），明文不进缓存键。"""
|
||
return json.dumps({"transport": row.transport, "command": row.command, "args": list(row.args or []),
|
||
"env": dict(row.env or {}), "url": row.url, "headers": dict(row.headers or {}),
|
||
"endpoint_id": row.endpoint_id}, sort_keys=True, ensure_ascii=False)
|
||
|
||
|
||
async def check_url(row: McpServer) -> dict[str, Any]:
|
||
"""http 服务连接前的地址检查，交回连接配置。
|
||
|
||
没绑定端点的：地址整串过 SSRF 检查（tools/net.py），和模型给的地址一样对待——MCP 的地址是搭建者填的，搭建者不能
|
||
借它访问内网。绑定了登记端点的：地址是端点生产环境的 base_url 拼上这里填的路径，认证头是端点的加上这里的，
|
||
内网访问跟端点的开关走（元数据地址照样拦）。stdio 不经网络，原样交回。"""
|
||
config = _server_config(row)
|
||
if row.transport != "http":
|
||
return config
|
||
from app.tools.net import UnsafeUrlError, assert_safe_url
|
||
|
||
if not row.endpoint_id:
|
||
try:
|
||
await assert_safe_url(config["url"])
|
||
except UnsafeUrlError as e:
|
||
# 指向内网、本机的 MCP 服务要绑定登记的端点（管理员开内网访问），这里说清去哪改
|
||
raise UnsafeUrlError(f"{e}。内网或本机上的 MCP 服务请改为绑定登记的端点：由管理员在「工具 → 端点」登记并开启"
|
||
"内网访问，再在这个服务里选择该端点、地址只填路径") from e
|
||
return config
|
||
from app.db.models import Endpoint
|
||
|
||
async with SessionLocal() as session:
|
||
endpoint = await session.get(Endpoint, row.endpoint_id)
|
||
env = ((endpoint.environments or {}).get("prod") if endpoint else None) or {}
|
||
if not env.get("base_url"):
|
||
raise UnsafeUrlError("绑定的端点不存在或没有配置生产环境")
|
||
path = (row.url or "").strip()
|
||
if path.startswith(("http://", "https://")):
|
||
raise UnsafeUrlError("绑定端点时，地址只填路径（例如 /mcp），主机和协议由端点决定")
|
||
url = env["base_url"].rstrip("/") + (path if not path or path.startswith("/") else f"/{path}")
|
||
await assert_safe_url(url, allow_intranet=bool(env.get("allow_intranet")))
|
||
headers = {**vault.reveal(env.get("headers")), **vault.reveal(row.headers)}
|
||
return {"transport": "streamable_http", "url": url, **({"headers": headers} if headers else {})}
|
||
|
||
|
||
#: 升级后第一次启动时，存量的内网 MCP 服务改为绑定登记端点（adopt_intranet_servers）。审计里有这一条就是做过了
|
||
ADOPT_ACTION = "mcp.intranet_review"
|
||
#: 每个服务的地址检查最多等这么久（解析域名）：启动时做，不能让一个解析不了的域名拖住服务
|
||
_ADOPT_CHECK_S = 10
|
||
|
||
|
||
async def adopt_intranet_servers() -> list[dict[str, Any]]:
|
||
"""升级前建的 http MCP 服务没有绑定端点、地址也不过 SSRF 检查——升级之后没绑定端点的地址一律按模型给的地址对待
|
||
（check_url），指向 localhost、内网主机的存量服务就连不上了。这里在升级后第一次启动时把它们各自改成绑定一个登记的
|
||
端点：生产环境就是它原来的主机（开内网访问，元数据地址和高危端口照样拦），服务的地址改成路径，连接和以前一样。每一项
|
||
写审计（endpoint.adopt_mcp），管理员事后能在「工具 → 端点」里看到、收紧。
|
||
|
||
只做一次：做完写一条 ADOPT_ACTION 审计，之后的启动看到它就不再做。之后新建的服务照规则走（内网只有登记的端点能访问，
|
||
开内网访问是管理员的事）——不然搭建者先填一个公网域名、保存后把它改指向内网，下次重启就被自动收编了。
|
||
交回收编了的 [{server, endpoint, base_url}]。"""
|
||
from urllib.parse import urlsplit, urlunsplit
|
||
|
||
from app.auth.audit import record
|
||
from app.db.models import AuditEvent, Endpoint
|
||
from app.tools.net import UnsafeUrlError, assert_safe_url
|
||
|
||
async with SessionLocal() as session:
|
||
if (await session.execute(select(AuditEvent.id).where(AuditEvent.action == ADOPT_ACTION).limit(1))).first():
|
||
return []
|
||
servers = list((await session.execute(select(McpServer).where(
|
||
McpServer.transport == "http", McpServer.endpoint_id.is_(None)))).scalars())
|
||
taken = set((await session.execute(select(Endpoint.name))).scalars())
|
||
adopted: list[dict[str, Any]] = []
|
||
for server in servers:
|
||
url = (server.url or "").strip()
|
||
parts = urlsplit(url)
|
||
if parts.scheme not in ("http", "https") or not parts.netloc:
|
||
continue
|
||
try:
|
||
await asyncio.wait_for(assert_safe_url(url), _ADOPT_CHECK_S)
|
||
continue # 公网地址：照旧不绑定
|
||
except (UnsafeUrlError, asyncio.TimeoutError):
|
||
pass
|
||
try:
|
||
# 开了内网访问也照样拦的（元数据地址、高危端口、解析不了的）不收编：以前就连不上，或者本来就不该连
|
||
await asyncio.wait_for(assert_safe_url(url, allow_intranet=True), _ADOPT_CHECK_S)
|
||
except (UnsafeUrlError, asyncio.TimeoutError):
|
||
continue
|
||
base = f"{parts.scheme}://{parts.netloc}"
|
||
name = _free_name(f"mcp-{'-'.join(server.name.split()).replace('/', '-')}", taken)
|
||
taken.add(name)
|
||
endpoint = Endpoint(name=name, description=f"升级时为 MCP 服务「{server.name}」自动登记（原地址指向内网）",
|
||
environments={"prod": {"base_url": base, "headers": {}, "allow_intranet": True,
|
||
"verify_tls": True}},
|
||
updated_by=None)
|
||
session.add(endpoint)
|
||
await session.flush()
|
||
server.endpoint_id = endpoint.id
|
||
server.url = urlunsplit(("", "", parts.path, parts.query, parts.fragment))
|
||
item = {"server": server.name, "endpoint": name, "base_url": base}
|
||
await record(session, "endpoint.adopt_mcp", "endpoint", endpoint.id, actor=None, verified=False,
|
||
data={**item, "allow_intranet": True})
|
||
adopted.append(item)
|
||
await record(session, ADOPT_ACTION, "endpoint", None, actor=None, verified=False,
|
||
data={"adopted": adopted})
|
||
await session.commit()
|
||
if adopted:
|
||
mcp_manager.invalidate(*(a["server"] for a in adopted))
|
||
return adopted
|
||
|
||
|
||
def _free_name(wanted: str, taken: set[str]) -> str:
|
||
name, n = wanted[:100], 2
|
||
while name in taken:
|
||
name = f"{wanted[:95]}-{n}"
|
||
n += 1
|
||
return name
|
||
|
||
|
||
class McpManager:
|
||
"""管理外部 MCP server 的连接与工具缓存。
|
||
|
||
工具在图里用 `mcp:<server>/<tool>` 引用。连接是懒建立的，
|
||
并且按 server 缓存，避免每个节点都去重新拉起一次 stdio 子进程。
|
||
|
||
缓存按服务逐个补、逐个作废。以前只在「一个都没加载过」时整批建一次：改了
|
||
启动命令，运行照样用旧命令建出来的工具；删掉的服务照样调得到；缓存建好之后
|
||
新加的服务，工具列表里一直没有它，除非有人去点刷新。现在每次取工具都先对一遍
|
||
库里启用着的服务：缺的连上，配置变了的重连，删了、停用了的不再交出去。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._tools: dict[str, list[BaseTool]] = {} # server name -> tools
|
||
#: 每个服务的工具是按哪份连接配置建的（_config_key）
|
||
self._built_from: dict[str, str] = {}
|
||
self._lock = asyncio.Lock()
|
||
|
||
def invalidate(self, *names: str) -> None:
|
||
"""服务改了、删了：丢掉按旧配置建的工具，下次用到时再连。"""
|
||
for name in names:
|
||
self._tools.pop(name, None)
|
||
self._built_from.pop(name, None)
|
||
|
||
def cached_tools(self, name: str) -> list[BaseTool]:
|
||
"""这个服务已经连上过、缓存着的工具（不去连）。能力登记同步时取描述、参数和 readOnlyHint 用。"""
|
||
return list(self._tools.get(name, []))
|
||
|
||
def _fresh(self, row: McpServer) -> bool:
|
||
return row.name in self._tools and self._built_from.get(row.name) == _config_key(row)
|
||
|
||
async def _load_servers(self, names: list[str] | None = None) -> list[McpServer]:
|
||
async with SessionLocal() as session:
|
||
query = select(McpServer).where(McpServer.enabled.is_(True))
|
||
if names is not None:
|
||
query = query.where(McpServer.name.in_(names))
|
||
rows = await session.execute(query)
|
||
return list(rows.scalars())
|
||
|
||
async def _connect(self, rows: list[McpServer]) -> None:
|
||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||
|
||
for row in rows:
|
||
key = _config_key(row)
|
||
started = time.perf_counter()
|
||
try:
|
||
client = MultiServerMCPClient({row.name: await check_url(row)})
|
||
tools = list(await asyncio.wait_for(client.get_tools(), timeout=30))
|
||
status, error = "ok", None
|
||
except Exception as e: # noqa: BLE001 - 一个 server 挂了不影响其他的
|
||
logger.warning("MCP server %s 连接失败: %s", row.name, e)
|
||
tools, status, error = [], "error", _reason(e, row.command)
|
||
# 连不上也缓存成空：不然每次取工具都要再等一轮 30 秒超时。改配置或刷新时重连
|
||
self._tools[row.name] = tools
|
||
self._built_from[row.name] = key
|
||
await self._mark(row.id, key, status, error, [t.name for t in tools], started)
|
||
|
||
async def _mark(
|
||
self, server_id: str, key: str, status: str, error: str | None, tool_names: list[str],
|
||
started: float,
|
||
) -> None:
|
||
from app.api.health import record
|
||
|
||
async with SessionLocal() as session:
|
||
row = await session.get(McpServer, server_id)
|
||
# 连到一半服务被改了：这次连的是旧配置，它的结论不能盖掉改完之后的「未检查」
|
||
if row and _config_key(row) == key:
|
||
row.status = status
|
||
row.last_error = error
|
||
row.tools_cache = tool_names
|
||
# 加载工具本身就是一次连接：记下来，工具页不用再单独测一次才知道
|
||
row.last_check = record(
|
||
status == "ok", int((time.perf_counter() - started) * 1000), error)
|
||
await session.commit()
|
||
|
||
async def _ensure(self, rows: list[McpServer]) -> None:
|
||
"""rows 里还没连过、或者是按旧配置连的，连一次补上。"""
|
||
if all(self._fresh(r) for r in rows):
|
||
return
|
||
async with self._lock:
|
||
# 等锁的时候别人可能已经连好了
|
||
stale = [r for r in rows if not self._fresh(r)]
|
||
if stale:
|
||
await self._connect(stale)
|
||
|
||
async def refresh(self) -> dict[str, list[str]]:
|
||
"""全部重连一遍（服务端的工具变了，配置没变时用）。"""
|
||
async with self._lock:
|
||
rows = await self._load_servers()
|
||
self.invalidate(*(set(self._tools) - {r.name for r in rows}))
|
||
await self._connect(rows)
|
||
return {r.name: [t.name for t in self._tools.get(r.name, [])] for r in rows}
|
||
|
||
async def ensure_loaded(self) -> None:
|
||
await self._ensure(await self._load_servers())
|
||
|
||
async def list_tools(self) -> list[dict[str, Any]]:
|
||
rows = await self._load_servers()
|
||
await self._ensure(rows)
|
||
out: list[dict[str, Any]] = []
|
||
for row in rows:
|
||
server = row.name
|
||
for tool in self._tools.get(server, []):
|
||
out.append(
|
||
{
|
||
"id": f"mcp:{server}/{tool.name}",
|
||
"name": tool.name,
|
||
"server": server,
|
||
"description": tool.description or "",
|
||
"category": f"MCP · {server}",
|
||
"schema": tool.args_schema.model_json_schema()
|
||
if getattr(tool, "args_schema", None)
|
||
else {},
|
||
}
|
||
)
|
||
return out
|
||
|
||
async def get_tools(self, refs: list[str]) -> list[BaseTool]:
|
||
"""按 `mcp:<server>/<tool>` 取工具；`mcp:<server>/*` 表示该 server 全部工具。
|
||
|
||
只交出库里还在、并且启用着的服务的工具：删了、停用了的，缓存里就算还有也不给。
|
||
"""
|
||
parsed = []
|
||
for ref in refs:
|
||
body = ref[4:] if ref.startswith("mcp:") else ref
|
||
server, _, tool_name = body.partition("/")
|
||
parsed.append((server, tool_name))
|
||
rows = await self._load_servers(sorted({server for server, _ in parsed}))
|
||
await self._ensure(rows)
|
||
live = {r.name for r in rows}
|
||
wanted: list[BaseTool] = []
|
||
for server, tool_name in parsed:
|
||
tools = self._tools.get(server, []) if server in live else []
|
||
if tool_name in ("", "*"):
|
||
wanted.extend(_keyed(server, t) for t in tools)
|
||
continue
|
||
for tool in tools:
|
||
if tool.name == tool_name:
|
||
wanted.append(_keyed(server, tool))
|
||
break
|
||
return wanted
|
||
|
||
async def probe(self, row: McpServer) -> dict[str, Any]:
|
||
"""连一次看看，用于设置页的"测试连接"。"""
|
||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||
|
||
try:
|
||
client = MultiServerMCPClient({row.name: await check_url(row)})
|
||
tools = await asyncio.wait_for(client.get_tools(), timeout=30)
|
||
return {
|
||
"ok": True,
|
||
"tools": [
|
||
{"name": t.name, "description": t.description or ""} for t in tools
|
||
],
|
||
}
|
||
except Exception as e: # noqa: BLE001
|
||
from app.core.errors import raw
|
||
|
||
return {"ok": False, "error": _reason(e, row.command), "detail": raw(e)}
|
||
|
||
|
||
def _reason(e: BaseException, command: str | None = None) -> str:
|
||
"""连不上 MCP 服务的原因，写成工具页上看得懂的话。
|
||
|
||
最常见的两种：命令不存在（没装 npx / uvx）和超时（进程起来了但不讲 MCP）。
|
||
"""
|
||
from app.core.errors import explain
|
||
from app.tools.net import UnsafeUrlError
|
||
|
||
if isinstance(e, UnsafeUrlError):
|
||
return f"地址未通过安全检查：{e}"
|
||
if isinstance(e, FileNotFoundError):
|
||
# stdio 客户端抛的这个异常常常不带文件名，只剩「找不到命令 」半句：补上配置的命令
|
||
missing = e.filename or command or ""
|
||
return f"无法启动：找不到命令 {missing}。请确认已安装，且位于服务端的 PATH 中"
|
||
if isinstance(e, asyncio.TimeoutError):
|
||
return "已连接，但 30 秒内没有响应：请确认该命令启动的是 MCP 服务，并检查参数"
|
||
reason, hint = explain(e)
|
||
return f"{reason}。{hint}" if hint else reason
|
||
|
||
|
||
mcp_manager = McpManager()
|
||
|
||
|
||
def _keyed(server: str, tool: BaseTool) -> BaseTool:
|
||
"""标上信任三档用的 key（`mcp:<server>/<tool>`，和工具列表里的 id 一致）。
|
||
|
||
缓存里的工具对象是共用的，key 只由 server 和工具名决定，重复标同一个值不会串。
|
||
"""
|
||
tool.metadata = {**(tool.metadata or {}), "trust_key": f"mcp:{server}/{tool.name}"}
|
||
return tool
|