1195 lines
45 KiB
Python
1195 lines
45 KiB
Python
"""MCP 客户端管理：工具发现、调用与 OpenAI 工具映射。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import select
|
||
import shutil
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING
|
||
|
||
import httpx
|
||
|
||
from config import MCP_PROTOCOL_VERSION, MCP_DEFAULT_TIMEOUT_SECONDS
|
||
from modules.mcp_server_registry import MCPServerRegistry
|
||
|
||
if TYPE_CHECKING:
|
||
from modules.user_container_manager import ContainerHandle
|
||
|
||
|
||
class MCPClientError(Exception):
|
||
"""MCP 客户端异常。"""
|
||
|
||
|
||
@dataclass
|
||
class MCPToolBinding:
|
||
alias: str
|
||
server_id: str
|
||
server_name: str
|
||
remote_name: str
|
||
transport: str
|
||
timeout_seconds: int
|
||
|
||
|
||
@dataclass
|
||
class MCPClientPoolEntry:
|
||
"""长连接 MCP 客户端池条目。"""
|
||
|
||
server_id: str
|
||
signature: str
|
||
client: Any
|
||
timeout_seconds: int
|
||
created_at: float
|
||
last_used_at: float
|
||
|
||
|
||
class _StdioMCPClient:
|
||
"""最小可用 stdio MCP 客户端（JSON-RPC over newline-delimited JSON）。"""
|
||
|
||
CONTAINER_PREFIXES = (
|
||
"/workspace",
|
||
"/usr/",
|
||
"/bin/",
|
||
"/sbin/",
|
||
"/lib/",
|
||
"/lib64/",
|
||
"/opt/",
|
||
"/etc/",
|
||
"/tmp/",
|
||
"/var/",
|
||
"/home/",
|
||
"/root/",
|
||
)
|
||
PATH_OPTION_KEYS = {
|
||
"--path",
|
||
"--dir",
|
||
"--directory",
|
||
"--root",
|
||
"--cwd",
|
||
"--workspace",
|
||
}
|
||
|
||
def __init__(
|
||
self,
|
||
server: Dict[str, Any],
|
||
timeout_seconds: int,
|
||
protocol_version: str,
|
||
container_session: Optional["ContainerHandle"] = None,
|
||
):
|
||
self.server = server
|
||
self.timeout_seconds = max(3, int(timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS))
|
||
self.protocol_version = protocol_version
|
||
self.container_session = container_session
|
||
self.process: Optional[subprocess.Popen] = None
|
||
self._seq = 0
|
||
self._initialized = False
|
||
self._initialize_result: Dict[str, Any] = {}
|
||
self._negotiated_protocol_version: Optional[str] = None
|
||
|
||
def _next_id(self) -> int:
|
||
self._seq += 1
|
||
return self._seq
|
||
|
||
@staticmethod
|
||
def _is_windows_abs_path(value: str) -> bool:
|
||
return len(value) > 2 and value[1] == ":" and value[0].isalpha()
|
||
|
||
@staticmethod
|
||
def _normalize_mount_path(value: str) -> str:
|
||
mount = str(value or "/workspace").strip()
|
||
mount = mount.rstrip("/")
|
||
return mount or "/workspace"
|
||
|
||
def _normalize_workspace_path(self, value: Any) -> str:
|
||
raw = str(value or "").strip()
|
||
if not raw:
|
||
return ""
|
||
try:
|
||
return str(Path(raw).expanduser().resolve())
|
||
except Exception:
|
||
return raw
|
||
|
||
def _map_workspace_path(self, raw: str, workspace_host: str, mount_path: str) -> str:
|
||
if not raw or not workspace_host:
|
||
return raw
|
||
if raw.startswith(mount_path):
|
||
return raw
|
||
if self._is_windows_abs_path(raw):
|
||
return raw
|
||
if not raw.startswith("/"):
|
||
return raw
|
||
try:
|
||
normalized = str(Path(raw).expanduser().resolve())
|
||
except Exception:
|
||
normalized = raw
|
||
if normalized == workspace_host:
|
||
return mount_path
|
||
prefix = workspace_host + os.sep
|
||
if normalized.startswith(prefix):
|
||
rel = normalized[len(prefix):].replace("\\", "/")
|
||
return f"{mount_path}/{rel}" if rel else mount_path
|
||
return raw
|
||
|
||
def _rewrite_path_like_value_for_container(
|
||
self,
|
||
value: str,
|
||
workspace_host: str,
|
||
mount_path: str,
|
||
) -> str:
|
||
raw = str(value or "").strip()
|
||
if not raw:
|
||
return raw
|
||
mapped = self._map_workspace_path(raw, workspace_host, mount_path)
|
||
return mapped if mapped != raw else raw
|
||
|
||
@staticmethod
|
||
def _command_basename(value: str) -> str:
|
||
text = str(value or "").strip().replace("\\", "/").rstrip("/")
|
||
if not text:
|
||
return ""
|
||
if "/" in text:
|
||
return text.split("/")[-1]
|
||
return text
|
||
|
||
def _rewrite_command_for_container(self, command: str, workspace_host: str, mount_path: str) -> str:
|
||
raw = str(command or "").strip()
|
||
if not raw:
|
||
return raw
|
||
mapped = self._map_workspace_path(raw, workspace_host, mount_path)
|
||
if mapped != raw:
|
||
return mapped
|
||
# command 若是宿主机绝对路径（如 /opt/homebrew/bin/npx），
|
||
# 在容器内通常应退化为命令名（npx/python3），避免无意义路径校验阻断。
|
||
if raw.startswith("/") or self._is_windows_abs_path(raw) or "\\" in raw:
|
||
candidate = self._command_basename(raw)
|
||
return candidate or raw
|
||
return raw
|
||
|
||
def _rewrite_arg_for_container(
|
||
self,
|
||
arg: str,
|
||
*,
|
||
index: int,
|
||
workspace_host: str,
|
||
mount_path: str,
|
||
) -> str:
|
||
raw = str(arg or "")
|
||
stripped = raw.strip()
|
||
if not stripped:
|
||
return raw
|
||
if "=" in stripped:
|
||
left, right = stripped.split("=", 1)
|
||
if left in self.PATH_OPTION_KEYS and right:
|
||
rewritten = self._rewrite_path_like_value_for_container(
|
||
right,
|
||
workspace_host=workspace_host,
|
||
mount_path=mount_path,
|
||
)
|
||
return f"{left}={rewritten}"
|
||
return self._rewrite_path_like_value_for_container(
|
||
stripped,
|
||
workspace_host=workspace_host,
|
||
mount_path=mount_path,
|
||
)
|
||
|
||
def _prepare_launch(self) -> Tuple[List[str], Optional[str], Dict[str, str]]:
|
||
command_raw = str(self.server.get("command") or "").strip()
|
||
args_raw = [str(item) for item in (self.server.get("args") or [])]
|
||
env_raw = dict(self.server.get("env") or {})
|
||
cwd_raw = str(self.server.get("cwd") or "").strip() or None
|
||
if not command_raw:
|
||
raise MCPClientError("stdio 服务缺少 command")
|
||
|
||
session = self.container_session
|
||
use_container = bool(session and getattr(session, "mode", None) == "docker")
|
||
if not use_container:
|
||
env = dict(os.environ)
|
||
env.update(env_raw)
|
||
return [command_raw, *args_raw], cwd_raw, env
|
||
|
||
container_name = str(getattr(session, "container_name", "") or "").strip()
|
||
if not container_name:
|
||
raise MCPClientError("docker 模式下 stdio MCP 缺少 container_name")
|
||
docker_bin = str(getattr(session, "sandbox_bin", "") or "docker").strip() or "docker"
|
||
docker_path = shutil.which(docker_bin) or docker_bin
|
||
mount_path = self._normalize_mount_path(getattr(session, "mount_path", "/workspace"))
|
||
workspace_host = self._normalize_workspace_path(getattr(session, "workspace_path", ""))
|
||
|
||
command = self._rewrite_command_for_container(
|
||
command_raw,
|
||
workspace_host=workspace_host,
|
||
mount_path=mount_path,
|
||
)
|
||
args: List[str] = []
|
||
for idx, item in enumerate(args_raw):
|
||
args.append(
|
||
self._rewrite_arg_for_container(
|
||
item,
|
||
index=idx,
|
||
workspace_host=workspace_host,
|
||
mount_path=mount_path,
|
||
)
|
||
)
|
||
cwd = (
|
||
self._rewrite_path_like_value_for_container(
|
||
cwd_raw,
|
||
workspace_host=workspace_host,
|
||
mount_path=mount_path,
|
||
)
|
||
if cwd_raw
|
||
else None
|
||
)
|
||
|
||
cmd: List[str] = [docker_path, "exec", "-i"]
|
||
if cwd:
|
||
cmd += ["-w", cwd]
|
||
|
||
exec_env = {
|
||
"PYTHONIOENCODING": "utf-8",
|
||
"TERM": "xterm-256color",
|
||
}
|
||
for key, value in env_raw.items():
|
||
if value is not None:
|
||
exec_env[str(key)] = str(value)
|
||
for key, value in exec_env.items():
|
||
cmd += ["-e", f"{key}={value}"]
|
||
cmd += [container_name, command, *args]
|
||
|
||
return cmd, None, dict(os.environ)
|
||
|
||
def start(self) -> None:
|
||
if self.process and self.process.poll() is None:
|
||
return
|
||
launch_cmd, cwd, env = self._prepare_launch()
|
||
try:
|
||
self.process = subprocess.Popen(
|
||
launch_cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
bufsize=1,
|
||
cwd=cwd or None,
|
||
env=env,
|
||
)
|
||
except Exception as exc:
|
||
raise MCPClientError(f"启动 stdio MCP 服务失败: {exc}") from exc
|
||
self._initialized = False
|
||
self._initialize_result = {}
|
||
self._negotiated_protocol_version = None
|
||
|
||
def _send(self, payload: Dict[str, Any]) -> None:
|
||
if not self.process or not self.process.stdin:
|
||
raise MCPClientError("stdio 进程未启动")
|
||
raw = json.dumps(payload, ensure_ascii=False)
|
||
try:
|
||
self.process.stdin.write(raw + "\n")
|
||
self.process.stdin.flush()
|
||
except Exception as exc:
|
||
raise MCPClientError(f"写入 stdio 消息失败: {exc}") from exc
|
||
|
||
def _read_message(self, timeout_seconds: Optional[float] = None) -> Dict[str, Any]:
|
||
if not self.process or not self.process.stdout:
|
||
raise MCPClientError("stdio 进程未启动")
|
||
timeout = self.timeout_seconds if timeout_seconds is None else max(0.1, float(timeout_seconds))
|
||
deadline = time.monotonic() + timeout
|
||
stdout = self.process.stdout
|
||
last_text_line = ""
|
||
while True:
|
||
if self.process.poll() is not None:
|
||
stderr_text = ""
|
||
stdout_text = ""
|
||
try:
|
||
if self.process.stderr:
|
||
stderr_text = self.process.stderr.read() or ""
|
||
except Exception:
|
||
pass
|
||
try:
|
||
if self.process.stdout:
|
||
stdout_text = self.process.stdout.read() or ""
|
||
except Exception:
|
||
pass
|
||
detail = (stderr_text or stdout_text or last_text_line or "")[:400]
|
||
raise MCPClientError(f"stdio 进程已退出（code={self.process.returncode}）{detail}")
|
||
|
||
remain = deadline - time.monotonic()
|
||
if remain <= 0:
|
||
raise MCPClientError("读取 stdio MCP 响应超时")
|
||
|
||
readable, _, _ = select.select([stdout], [], [], min(0.3, remain))
|
||
if not readable:
|
||
continue
|
||
line = stdout.readline()
|
||
if not line:
|
||
continue
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
return json.loads(line)
|
||
except Exception:
|
||
last_text_line = line
|
||
continue
|
||
|
||
@staticmethod
|
||
def _build_server_request_result(method: str, params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
|
||
"""处理服务端发起的请求（Server -> Client Request）。"""
|
||
name = str(method or "").strip()
|
||
if name == "ping":
|
||
return True, {}
|
||
if name == "roots/list":
|
||
return True, {"roots": []}
|
||
if name == "completion/complete":
|
||
return False, {"code": -32601, "message": "completion/complete not supported by client"}
|
||
if name == "sampling/createMessage":
|
||
return False, {"code": -32601, "message": "sampling not supported by client"}
|
||
if name == "elicitation/create":
|
||
return False, {"code": -32601, "message": "elicitation not supported by client"}
|
||
return False, {"code": -32601, "message": f"Client method not implemented: {name}"}
|
||
|
||
@staticmethod
|
||
def _handle_server_notification(method: str, params: Dict[str, Any]) -> None:
|
||
"""处理服务端通知（Server -> Client Notification）。当前策略：静默接收。"""
|
||
_ = method, params
|
||
|
||
def _respond_server_request(self, msg: Dict[str, Any]) -> None:
|
||
msg_id = msg.get("id")
|
||
method = str(msg.get("method") or "").strip()
|
||
params = msg.get("params")
|
||
if not isinstance(params, dict):
|
||
params = {}
|
||
try:
|
||
ok, payload = self._build_server_request_result(method, params)
|
||
if ok:
|
||
self._send({"jsonrpc": "2.0", "id": msg_id, "result": payload})
|
||
else:
|
||
self._send({"jsonrpc": "2.0", "id": msg_id, "error": payload})
|
||
except Exception as exc:
|
||
self._send(
|
||
{
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"error": {"code": -32603, "message": f"Client internal error: {exc}"},
|
||
}
|
||
)
|
||
|
||
def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
if not self.process or self.process.poll() is not None:
|
||
self.start()
|
||
req_id = self._next_id()
|
||
payload = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}}
|
||
self._send(payload)
|
||
while True:
|
||
msg = self._read_message()
|
||
|
||
# 收到目标响应
|
||
if msg.get("id") == req_id and "method" not in msg:
|
||
if "error" in msg:
|
||
error = msg.get("error") or {}
|
||
raise MCPClientError(f"{method} 调用失败: {error}")
|
||
return msg.get("result") or {}
|
||
|
||
# 服务端请求：回包后继续等待目标响应
|
||
if "id" in msg and "method" in msg:
|
||
self._respond_server_request(msg)
|
||
continue
|
||
|
||
# 服务端通知：消费后继续
|
||
if "method" in msg:
|
||
self._handle_server_notification(str(msg.get("method") or ""), msg.get("params") or {})
|
||
continue
|
||
|
||
def notify(self, method: str, params: Optional[Dict[str, Any]] = None) -> None:
|
||
if not self.process or self.process.poll() is not None:
|
||
self.start()
|
||
self._send({"jsonrpc": "2.0", "method": method, "params": params or {}})
|
||
|
||
def initialize(self) -> Dict[str, Any]:
|
||
if self._initialized:
|
||
return dict(self._initialize_result)
|
||
if not self.process or self.process.poll() is not None:
|
||
self.start()
|
||
result = self.request(
|
||
"initialize",
|
||
{
|
||
"protocolVersion": self.protocol_version,
|
||
"capabilities": {
|
||
"roots": {"listChanged": False},
|
||
},
|
||
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
|
||
},
|
||
)
|
||
self.notify("notifications/initialized")
|
||
self._initialize_result = dict(result or {})
|
||
self._initialized = True
|
||
negotiated = result.get("protocolVersion") if isinstance(result, dict) else None
|
||
if negotiated:
|
||
self._negotiated_protocol_version = str(negotiated)
|
||
return dict(self._initialize_result)
|
||
|
||
def close(self) -> None:
|
||
if not self.process:
|
||
return
|
||
try:
|
||
if self.process.stdin:
|
||
self.process.stdin.close()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
if self.process.poll() is None:
|
||
self.process.terminate()
|
||
self.process.wait(timeout=1.5)
|
||
except Exception:
|
||
try:
|
||
self.process.kill()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
try:
|
||
if self.process.stdout:
|
||
self.process.stdout.close()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
if self.process.stderr:
|
||
self.process.stderr.close()
|
||
except Exception:
|
||
pass
|
||
self.process = None
|
||
self._initialized = False
|
||
self._initialize_result = {}
|
||
self._negotiated_protocol_version = None
|
||
|
||
|
||
class _StreamableHTTPMCPClient:
|
||
"""最小可用 streamable HTTP MCP 客户端。"""
|
||
|
||
def __init__(self, server: Dict[str, Any], timeout_seconds: int, protocol_version: str):
|
||
self.server = server
|
||
self.url = str(server.get("url") or "").strip()
|
||
self.timeout_seconds = max(3, int(timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS))
|
||
self.protocol_version = protocol_version
|
||
self._session_id: Optional[str] = None
|
||
self._seq = 0
|
||
self._client = httpx.Client(timeout=self.timeout_seconds)
|
||
self._initialized = False
|
||
self._initialize_result: Dict[str, Any] = {}
|
||
self._negotiated_protocol_version: Optional[str] = None
|
||
|
||
def _next_id(self) -> int:
|
||
self._seq += 1
|
||
return self._seq
|
||
|
||
@staticmethod
|
||
def _header_pick(headers: httpx.Headers, key: str) -> Optional[str]:
|
||
for hk, hv in headers.items():
|
||
if hk.lower() == key.lower():
|
||
return hv
|
||
return None
|
||
|
||
@staticmethod
|
||
def _infer_mcp_name(method: str, params: Optional[Dict[str, Any]]) -> Optional[str]:
|
||
if not isinstance(params, dict):
|
||
params = {}
|
||
method = str(method or "").strip()
|
||
if method == "tools/call":
|
||
value = params.get("name")
|
||
return str(value).strip() if value is not None else None
|
||
if method == "resources/read":
|
||
value = params.get("uri")
|
||
return str(value).strip() if value is not None else None
|
||
if method == "prompts/get":
|
||
value = params.get("name")
|
||
return str(value).strip() if value is not None else None
|
||
return None
|
||
|
||
def _build_headers(self, *, method: Optional[str] = None, params: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
|
||
headers = {
|
||
"Accept": "application/json, text/event-stream",
|
||
"Content-Type": "application/json",
|
||
"MCP-Protocol-Version": self._negotiated_protocol_version or self.protocol_version,
|
||
}
|
||
if self._session_id:
|
||
headers["Mcp-Session-Id"] = self._session_id
|
||
if method:
|
||
headers["Mcp-Method"] = str(method)
|
||
name = self._infer_mcp_name(str(method), params)
|
||
if name:
|
||
headers["Mcp-Name"] = name
|
||
headers.update(self.server.get("headers") or {})
|
||
return headers
|
||
|
||
@staticmethod
|
||
def _parse_sse_json_events(raw_text: str) -> List[Dict[str, Any]]:
|
||
events: List[Dict[str, Any]] = []
|
||
blocks = raw_text.split("\n\n")
|
||
for block in blocks:
|
||
data_lines: List[str] = []
|
||
for line in block.splitlines():
|
||
if line.startswith("data:"):
|
||
data_lines.append(line[5:].lstrip())
|
||
if not data_lines:
|
||
continue
|
||
payload = "\n".join(data_lines).strip()
|
||
if not payload or payload == "[DONE]":
|
||
continue
|
||
try:
|
||
parsed = json.loads(payload)
|
||
except Exception:
|
||
continue
|
||
if isinstance(parsed, dict):
|
||
events.append(parsed)
|
||
return events
|
||
|
||
@staticmethod
|
||
def _build_server_request_result(method: str, params: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
|
||
name = str(method or "").strip()
|
||
if name == "ping":
|
||
return True, {}
|
||
if name == "roots/list":
|
||
return True, {"roots": []}
|
||
if name == "completion/complete":
|
||
return False, {"code": -32601, "message": "completion/complete not supported by client"}
|
||
if name == "sampling/createMessage":
|
||
return False, {"code": -32601, "message": "sampling not supported by client"}
|
||
if name == "elicitation/create":
|
||
return False, {"code": -32601, "message": "elicitation not supported by client"}
|
||
return False, {"code": -32601, "message": f"Client method not implemented: {name}"}
|
||
|
||
def _post_jsonrpc_message(self, payload: Dict[str, Any]) -> None:
|
||
try:
|
||
resp = self._client.post(
|
||
self.url,
|
||
headers=self._build_headers(),
|
||
json=payload,
|
||
)
|
||
except Exception as exc:
|
||
raise MCPClientError(f"HTTP MCP 回写失败: {exc}") from exc
|
||
if resp.status_code >= 400:
|
||
raise MCPClientError(f"HTTP MCP 回写状态异常: {resp.status_code} {resp.text[:400]}")
|
||
|
||
def _handle_server_request_message(self, message: Dict[str, Any]) -> None:
|
||
msg_id = message.get("id")
|
||
method = str(message.get("method") or "").strip()
|
||
params = message.get("params")
|
||
if not isinstance(params, dict):
|
||
params = {}
|
||
try:
|
||
ok, payload = self._build_server_request_result(method, params)
|
||
if ok:
|
||
self._post_jsonrpc_message({"jsonrpc": "2.0", "id": msg_id, "result": payload})
|
||
else:
|
||
self._post_jsonrpc_message({"jsonrpc": "2.0", "id": msg_id, "error": payload})
|
||
except Exception as exc:
|
||
self._post_jsonrpc_message(
|
||
{
|
||
"jsonrpc": "2.0",
|
||
"id": msg_id,
|
||
"error": {"code": -32603, "message": f"Client internal error: {exc}"},
|
||
}
|
||
)
|
||
|
||
def _request(
|
||
self,
|
||
method: str,
|
||
params: Optional[Dict[str, Any]] = None,
|
||
notification: bool = False,
|
||
*,
|
||
_allow_session_retry: bool = True,
|
||
) -> Dict[str, Any]:
|
||
if not self.url:
|
||
raise MCPClientError("streamable_http 服务缺少 url")
|
||
req_id = None if notification else self._next_id()
|
||
payload: Dict[str, Any] = {"jsonrpc": "2.0", "method": method, "params": params or {}}
|
||
if req_id is not None:
|
||
payload["id"] = req_id
|
||
try:
|
||
resp = self._client.post(
|
||
self.url,
|
||
headers=self._build_headers(method=method, params=params or {}),
|
||
json=payload,
|
||
)
|
||
except Exception as exc:
|
||
raise MCPClientError(f"HTTP MCP 请求失败: {exc}") from exc
|
||
|
||
session_id = self._header_pick(resp.headers, "Mcp-Session-Id")
|
||
if session_id:
|
||
self._session_id = session_id
|
||
|
||
# 会话失效（服务端返回 404）时，按规范重新建立 session 并重试一次
|
||
if (
|
||
resp.status_code == 404
|
||
and self._session_id
|
||
and method != "initialize"
|
||
and _allow_session_retry
|
||
):
|
||
self._session_id = None
|
||
self._initialized = False
|
||
self._initialize_result = {}
|
||
self._negotiated_protocol_version = None
|
||
self.initialize()
|
||
return self._request(
|
||
method,
|
||
params,
|
||
notification=notification,
|
||
_allow_session_retry=False,
|
||
)
|
||
|
||
if resp.status_code >= 400:
|
||
raise MCPClientError(f"HTTP MCP 状态异常: {resp.status_code} {resp.text[:400]}")
|
||
|
||
if notification:
|
||
return {}
|
||
|
||
content_type = (resp.headers.get("content-type") or "").lower()
|
||
message: Optional[Dict[str, Any]] = None
|
||
if "application/json" in content_type:
|
||
parsed = resp.json()
|
||
if isinstance(parsed, dict):
|
||
message = parsed
|
||
elif "text/event-stream" in content_type:
|
||
for event in self._parse_sse_json_events(resp.text):
|
||
if not isinstance(event, dict):
|
||
continue
|
||
if "method" in event:
|
||
if "id" in event:
|
||
self._handle_server_request_message(event)
|
||
# 通知直接消费
|
||
continue
|
||
if event.get("id") == req_id:
|
||
message = event
|
||
break
|
||
else:
|
||
try:
|
||
parsed = resp.json()
|
||
if isinstance(parsed, dict):
|
||
message = parsed
|
||
except Exception:
|
||
pass
|
||
|
||
if not message:
|
||
raise MCPClientError(f"HTTP MCP 未返回可解析响应（method={method}）")
|
||
if "error" in message:
|
||
raise MCPClientError(f"{method} 调用失败: {message.get('error')}")
|
||
return message.get("result") or {}
|
||
|
||
def initialize(self) -> Dict[str, Any]:
|
||
if self._initialized:
|
||
return dict(self._initialize_result)
|
||
result = self._request(
|
||
"initialize",
|
||
{
|
||
"protocolVersion": self.protocol_version,
|
||
"capabilities": {
|
||
"roots": {"listChanged": False},
|
||
},
|
||
"clientInfo": {"name": "agents-mcp-bridge", "version": "1.0.0"},
|
||
},
|
||
)
|
||
self._request("notifications/initialized", {}, notification=True)
|
||
self._initialize_result = dict(result or {})
|
||
self._initialized = True
|
||
negotiated = result.get("protocolVersion") if isinstance(result, dict) else None
|
||
if negotiated:
|
||
self._negotiated_protocol_version = str(negotiated)
|
||
return dict(self._initialize_result)
|
||
|
||
def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
return self._request(method, params, notification=False)
|
||
|
||
def close(self) -> None:
|
||
try:
|
||
if self._session_id and self.url:
|
||
self._client.delete(self.url, headers=self._build_headers())
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self._client.close()
|
||
except Exception:
|
||
pass
|
||
self._session_id = None
|
||
self._initialized = False
|
||
self._initialize_result = {}
|
||
self._negotiated_protocol_version = None
|
||
|
||
|
||
class MCPClientManager:
|
||
"""MCP 客户端主协调器。"""
|
||
|
||
def __init__(
|
||
self,
|
||
registry: MCPServerRegistry,
|
||
protocol_version: str = MCP_PROTOCOL_VERSION,
|
||
container_session: Optional["ContainerHandle"] = None,
|
||
):
|
||
self.registry = registry
|
||
self.protocol_version = str(protocol_version or MCP_PROTOCOL_VERSION)
|
||
self.container_session = container_session
|
||
self._latest_alias_map: Dict[str, MCPToolBinding] = {}
|
||
self._client_pool: Dict[str, MCPClientPoolEntry] = {}
|
||
self._pool_lock = threading.RLock()
|
||
|
||
def __del__(self):
|
||
try:
|
||
self.close_all_clients()
|
||
except Exception:
|
||
pass
|
||
|
||
def set_container_session(self, session: Optional["ContainerHandle"]) -> None:
|
||
if session is self.container_session:
|
||
return
|
||
self.close_all_clients()
|
||
self.container_session = session
|
||
|
||
@staticmethod
|
||
def _server_signature(server: Dict[str, Any]) -> str:
|
||
payload = {
|
||
"id": server.get("id"),
|
||
"transport": server.get("transport"),
|
||
"command": server.get("command"),
|
||
"args": server.get("args") or [],
|
||
"cwd": server.get("cwd"),
|
||
"env": server.get("env") or {},
|
||
"url": server.get("url"),
|
||
"headers": server.get("headers") or {},
|
||
"timeout_seconds": server.get("timeout_seconds"),
|
||
}
|
||
raw = json.dumps(payload, ensure_ascii=False, sort_keys=True)
|
||
return hashlib.sha1(raw.encode("utf-8")).hexdigest()
|
||
|
||
@staticmethod
|
||
def _is_client_alive(client: Any) -> bool:
|
||
process = getattr(client, "process", None)
|
||
if process is not None:
|
||
return bool(process.poll() is None)
|
||
# HTTP 客户端无法直接探活，按可用处理，失败时走重连兜底
|
||
return True
|
||
|
||
@staticmethod
|
||
def _safe_close_client(client: Any) -> None:
|
||
try:
|
||
if client:
|
||
client.close()
|
||
except Exception:
|
||
pass
|
||
|
||
def _close_pool_entry_locked(self, server_id: str) -> None:
|
||
entry = self._client_pool.pop(server_id, None)
|
||
if entry:
|
||
self._safe_close_client(entry.client)
|
||
|
||
def _ensure_client_locked(self, server: Dict[str, Any], timeout_seconds: int, *, force_recreate: bool = False):
|
||
sid = str(server.get("id") or "").strip()
|
||
if not sid:
|
||
raise MCPClientError("MCP 服务缺少 id")
|
||
signature = self._server_signature(server)
|
||
entry = self._client_pool.get(sid)
|
||
needs_recreate = force_recreate or (entry is None)
|
||
if entry and not needs_recreate:
|
||
if entry.signature != signature:
|
||
needs_recreate = True
|
||
elif not self._is_client_alive(entry.client):
|
||
needs_recreate = True
|
||
|
||
if needs_recreate:
|
||
if entry:
|
||
self._safe_close_client(entry.client)
|
||
client = self._make_client(server, timeout_seconds)
|
||
if hasattr(client, "start"):
|
||
client.start()
|
||
init_result = client.initialize()
|
||
now = time.time()
|
||
entry = MCPClientPoolEntry(
|
||
server_id=sid,
|
||
signature=signature,
|
||
client=client,
|
||
timeout_seconds=timeout_seconds,
|
||
created_at=now,
|
||
last_used_at=now,
|
||
)
|
||
self._client_pool[sid] = entry
|
||
return client, init_result
|
||
|
||
# 复用连接：动态刷新超时，并确保处于 initialized 状态
|
||
entry.timeout_seconds = timeout_seconds
|
||
try:
|
||
if hasattr(entry.client, "timeout_seconds"):
|
||
entry.client.timeout_seconds = max(3, int(timeout_seconds))
|
||
if hasattr(entry.client, "_client"):
|
||
entry.client._client.timeout = httpx.Timeout(max(3, int(timeout_seconds)))
|
||
except Exception:
|
||
pass
|
||
init_result = {}
|
||
if hasattr(entry.client, "initialize"):
|
||
init_result = entry.client.initialize() or {}
|
||
entry.last_used_at = time.time()
|
||
return entry.client, init_result
|
||
|
||
def _run_with_client(
|
||
self,
|
||
server: Dict[str, Any],
|
||
timeout_seconds: int,
|
||
operation: Callable[[Any], Any],
|
||
) -> Tuple[Any, Dict[str, Any]]:
|
||
last_exc: Optional[Exception] = None
|
||
for attempt in range(2):
|
||
with self._pool_lock:
|
||
client, init_result = self._ensure_client_locked(
|
||
server,
|
||
timeout_seconds,
|
||
force_recreate=attempt > 0,
|
||
)
|
||
sid = str(server.get("id") or "").strip()
|
||
try:
|
||
output = operation(client)
|
||
entry = self._client_pool.get(sid)
|
||
if entry:
|
||
entry.last_used_at = time.time()
|
||
return output, init_result
|
||
except Exception as exc: # pragma: no cover - 仅异常链路
|
||
last_exc = exc
|
||
self._close_pool_entry_locked(sid)
|
||
continue
|
||
if last_exc:
|
||
raise last_exc
|
||
raise MCPClientError("MCP 客户端执行失败")
|
||
|
||
def close_all_clients(self) -> None:
|
||
with self._pool_lock:
|
||
entries = list(self._client_pool.values())
|
||
self._client_pool.clear()
|
||
for entry in entries:
|
||
self._safe_close_client(entry.client)
|
||
|
||
@staticmethod
|
||
def _sanitize_input_schema(schema: Any) -> Dict[str, Any]:
|
||
if not isinstance(schema, dict):
|
||
return {"type": "object", "properties": {}}
|
||
output = dict(schema)
|
||
if output.get("type") != "object":
|
||
output["type"] = "object"
|
||
if not isinstance(output.get("properties"), dict):
|
||
output["properties"] = {}
|
||
if "required" in output and not isinstance(output.get("required"), list):
|
||
output.pop("required", None)
|
||
return output
|
||
|
||
@staticmethod
|
||
def _compact_tool_entry(tool: Dict[str, Any]) -> Dict[str, Any]:
|
||
return {
|
||
"name": str(tool.get("name") or "").strip(),
|
||
"description": str(tool.get("description") or ""),
|
||
"title": str(tool.get("title") or ""),
|
||
"inputSchema": tool.get("inputSchema") if isinstance(tool.get("inputSchema"), dict) else {"type": "object", "properties": {}},
|
||
}
|
||
|
||
@staticmethod
|
||
def _build_call_message(result: Dict[str, Any]) -> str:
|
||
content = result.get("content")
|
||
if isinstance(content, list):
|
||
lines: List[str] = []
|
||
for idx, item in enumerate(content, start=1):
|
||
if isinstance(item, dict) and item.get("type") == "text":
|
||
text = str(item.get("text") or "").strip()
|
||
if text:
|
||
lines.append(text)
|
||
continue
|
||
if not isinstance(item, dict):
|
||
continue
|
||
ctype = str(item.get("type") or "").strip().lower()
|
||
if ctype in {"image", "audio"}:
|
||
mime = str(item.get("mimeType") or item.get("mime_type") or "").strip()
|
||
if not mime:
|
||
mime = "image/png" if ctype == "image" else "audio/wav"
|
||
data_len = len(str(item.get("data") or ""))
|
||
lines.append(f"[MCP {ctype} #{idx}] {mime} ({data_len} base64 chars)")
|
||
continue
|
||
if ctype in {"resource_link", "resourcelink"}:
|
||
uri = str(item.get("uri") or item.get("url") or "").strip()
|
||
name = str(item.get("name") or item.get("title") or "").strip()
|
||
if uri and name:
|
||
lines.append(f"[MCP 资源] {name} -> {uri}")
|
||
elif uri:
|
||
lines.append(f"[MCP 资源] {uri}")
|
||
continue
|
||
if ctype == "resource":
|
||
resource = item.get("resource")
|
||
if isinstance(resource, dict):
|
||
text = str(resource.get("text") or "").strip()
|
||
if text:
|
||
lines.append(text)
|
||
continue
|
||
uri = str(resource.get("uri") or "").strip()
|
||
if uri:
|
||
lines.append(f"[MCP 资源] {uri}")
|
||
continue
|
||
uri = str(item.get("uri") or item.get("url") or "").strip()
|
||
if uri:
|
||
lines.append(f"[MCP 资源] {uri}")
|
||
if lines:
|
||
return "\n".join(lines)[:500]
|
||
structured = result.get("structuredContent")
|
||
if structured is not None:
|
||
try:
|
||
return json.dumps(structured, ensure_ascii=False)[:500]
|
||
except Exception:
|
||
return str(structured)[:500]
|
||
return ""
|
||
|
||
@staticmethod
|
||
def _content_preview(content_items: Any) -> str:
|
||
if not isinstance(content_items, list):
|
||
return ""
|
||
previews: List[str] = []
|
||
for item in content_items:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
ctype = item.get("type")
|
||
if ctype == "text":
|
||
text = str(item.get("text") or "").strip()
|
||
if text:
|
||
previews.append(text[:300])
|
||
elif ctype:
|
||
previews.append(f"[{ctype}]")
|
||
if len(previews) >= 4:
|
||
break
|
||
return "\n".join(previews).strip()
|
||
|
||
def _make_client(self, server: Dict[str, Any], timeout_seconds: int):
|
||
transport = server.get("transport")
|
||
if transport == "stdio":
|
||
return _StdioMCPClient(
|
||
server,
|
||
timeout_seconds,
|
||
self.protocol_version,
|
||
container_session=self.container_session,
|
||
)
|
||
if transport == "streamable_http":
|
||
return _StreamableHTTPMCPClient(server, timeout_seconds, self.protocol_version)
|
||
raise MCPClientError(f"不支持的 MCP transport: {transport}")
|
||
|
||
def _paged_list_tools(self, client, max_pages: int = 20) -> List[Dict[str, Any]]:
|
||
cursor: Optional[str] = None
|
||
tools: List[Dict[str, Any]] = []
|
||
for _ in range(max_pages):
|
||
params = {"cursor": cursor} if cursor else {}
|
||
result = client.request("tools/list", params)
|
||
page_tools = result.get("tools") if isinstance(result, dict) else []
|
||
if isinstance(page_tools, list):
|
||
for item in page_tools:
|
||
if isinstance(item, dict):
|
||
compact = self._compact_tool_entry(item)
|
||
if compact.get("name"):
|
||
tools.append(compact)
|
||
cursor = (result or {}).get("nextCursor") if isinstance(result, dict) else None
|
||
if not cursor:
|
||
break
|
||
return tools
|
||
|
||
def discover_tools_for_server(self, server_id: str, *, timeout_seconds: Optional[int] = None) -> Dict[str, Any]:
|
||
server = self.registry.get_server(server_id)
|
||
if not server:
|
||
return {"success": False, "error": f"未找到 MCP 服务: {server_id}"}
|
||
timeout = int(timeout_seconds or server.get("timeout_seconds") or MCP_DEFAULT_TIMEOUT_SECONDS)
|
||
started = time.time()
|
||
try:
|
||
tools, init_result = self._run_with_client(
|
||
server,
|
||
timeout,
|
||
lambda client: self._paged_list_tools(client),
|
||
)
|
||
self.registry.update_tools_cache(server_id, tools=tools, error="", touched=True)
|
||
elapsed_ms = int((time.time() - started) * 1000)
|
||
return {
|
||
"success": True,
|
||
"server_id": server_id,
|
||
"protocol_version": init_result.get("protocolVersion") if isinstance(init_result, dict) else None,
|
||
"tools_count": len(tools),
|
||
"tools": tools,
|
||
"elapsed_ms": elapsed_ms,
|
||
}
|
||
except Exception as exc:
|
||
self.registry.update_tools_cache(server_id, error=str(exc), touched=True)
|
||
return {"success": False, "server_id": server_id, "error": str(exc)}
|
||
|
||
def sync_servers(self, server_id: Optional[str] = None) -> Dict[str, Any]:
|
||
# 先从磁盘重载一次，避免「文件已改但进程内缓存未刷新」导致同步遗漏，
|
||
# 以及后续 update_tools_cache 把旧缓存反写覆盖新配置。
|
||
self.registry.reload()
|
||
with self._pool_lock:
|
||
enabled_ids = {
|
||
str(item.get("id") or "").strip()
|
||
for item in self.registry.list_enabled_servers()
|
||
if str(item.get("id") or "").strip()
|
||
}
|
||
stale_ids = [sid for sid in list(self._client_pool.keys()) if sid not in enabled_ids]
|
||
for sid in stale_ids:
|
||
self._close_pool_entry_locked(sid)
|
||
if server_id:
|
||
return self.discover_tools_for_server(server_id)
|
||
results = []
|
||
for server in self.registry.list_enabled_servers():
|
||
sid = server.get("id")
|
||
if not sid:
|
||
continue
|
||
results.append(self.discover_tools_for_server(sid))
|
||
ok_count = sum(1 for item in results if item.get("success"))
|
||
overall_success = (not results) or (ok_count == len(results))
|
||
return {
|
||
"success": overall_success,
|
||
"servers_total": len(results),
|
||
"servers_success": ok_count,
|
||
"servers_failed": len(results) - ok_count,
|
||
"results": results,
|
||
}
|
||
|
||
def _iter_visible_tools(self, server: Dict[str, Any]) -> Iterable[Dict[str, Any]]:
|
||
include_tools = set(server.get("include_tools") or [])
|
||
exclude_tools = set(server.get("exclude_tools") or [])
|
||
for item in server.get("tools_cache") or []:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
name = str(item.get("name") or "").strip()
|
||
if not name:
|
||
continue
|
||
if include_tools and name not in include_tools:
|
||
continue
|
||
if name in exclude_tools:
|
||
continue
|
||
yield item
|
||
|
||
def build_openai_tools(
|
||
self,
|
||
*,
|
||
ensure_discovery: bool = True,
|
||
discovery_timeout_seconds: int = 12,
|
||
) -> Tuple[List[Dict[str, Any]], Dict[str, MCPToolBinding]]:
|
||
tools: List[Dict[str, Any]] = []
|
||
alias_map: Dict[str, MCPToolBinding] = {}
|
||
|
||
servers = self.registry.list_enabled_servers()
|
||
for server in servers:
|
||
sid = str(server.get("id") or "").strip()
|
||
if not sid:
|
||
continue
|
||
if ensure_discovery and not (server.get("tools_cache") or []):
|
||
self.discover_tools_for_server(sid, timeout_seconds=discovery_timeout_seconds)
|
||
server = self.registry.get_server(sid) or server
|
||
|
||
for item in self._iter_visible_tools(server):
|
||
remote_name = str(item.get("name") or "").strip()
|
||
alias_base = self.registry.make_tool_alias(sid, remote_name)
|
||
alias = alias_base
|
||
if alias in alias_map:
|
||
suffix = hashlib.sha1(f"{sid}:{remote_name}".encode("utf-8")).hexdigest()[:6]
|
||
max_len = max(8, 64 - len(suffix) - 2)
|
||
alias = f"{alias_base[:max_len]}__{suffix}"
|
||
|
||
description = str(item.get("description") or "").strip()
|
||
if description:
|
||
description = f"[MCP:{sid}] {description}"
|
||
else:
|
||
description = f"[MCP:{sid}] 通过 MCP 服务调用远程工具 {remote_name}"
|
||
|
||
schema = self._sanitize_input_schema(item.get("inputSchema"))
|
||
tools.append(
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": alias,
|
||
"description": description,
|
||
"parameters": schema,
|
||
},
|
||
}
|
||
)
|
||
alias_map[alias] = MCPToolBinding(
|
||
alias=alias,
|
||
server_id=sid,
|
||
server_name=str(server.get("name") or sid),
|
||
remote_name=remote_name,
|
||
transport=str(server.get("transport") or ""),
|
||
timeout_seconds=int(server.get("timeout_seconds") or MCP_DEFAULT_TIMEOUT_SECONDS),
|
||
)
|
||
|
||
self._latest_alias_map = alias_map
|
||
return tools, alias_map
|
||
|
||
def _resolve_binding(
|
||
self,
|
||
tool_alias: str,
|
||
alias_map: Optional[Dict[str, MCPToolBinding]] = None,
|
||
) -> Optional[MCPToolBinding]:
|
||
if alias_map and tool_alias in alias_map:
|
||
return alias_map[tool_alias]
|
||
if tool_alias in self._latest_alias_map:
|
||
return self._latest_alias_map[tool_alias]
|
||
|
||
# 容错：从缓存动态重建一遍
|
||
_, rebuilt = self.build_openai_tools(ensure_discovery=False)
|
||
return rebuilt.get(tool_alias)
|
||
|
||
def call_tool_by_alias(
|
||
self,
|
||
tool_alias: str,
|
||
arguments: Dict[str, Any],
|
||
*,
|
||
alias_map: Optional[Dict[str, MCPToolBinding]] = None,
|
||
) -> Dict[str, Any]:
|
||
binding = self._resolve_binding(tool_alias, alias_map)
|
||
if not binding:
|
||
return {"success": False, "error": f"未找到 MCP 工具映射: {tool_alias}"}
|
||
server = self.registry.get_server(binding.server_id)
|
||
if not server:
|
||
return {"success": False, "error": f"未找到 MCP 服务: {binding.server_id}"}
|
||
if not server.get("enabled", True):
|
||
return {"success": False, "error": f"MCP 服务已禁用: {binding.server_id}"}
|
||
|
||
timeout = int(server.get("timeout_seconds") or binding.timeout_seconds or MCP_DEFAULT_TIMEOUT_SECONDS)
|
||
started = time.time()
|
||
try:
|
||
call_result, _ = self._run_with_client(
|
||
server,
|
||
timeout,
|
||
lambda client: client.request(
|
||
"tools/call",
|
||
{"name": binding.remote_name, "arguments": arguments or {}},
|
||
),
|
||
)
|
||
is_error = bool((call_result or {}).get("isError"))
|
||
content_items = (call_result or {}).get("content") or []
|
||
preview = self._content_preview(content_items)
|
||
elapsed_ms = int((time.time() - started) * 1000)
|
||
message = self._build_call_message(call_result)
|
||
|
||
return {
|
||
"success": not is_error,
|
||
"server_id": binding.server_id,
|
||
"server_name": binding.server_name,
|
||
"tool_alias": binding.alias,
|
||
"tool_name": binding.remote_name,
|
||
"transport": binding.transport,
|
||
"elapsed_ms": elapsed_ms,
|
||
"is_error": is_error,
|
||
"message": message,
|
||
"preview": preview,
|
||
"content": content_items,
|
||
"structured_content": (call_result or {}).get("structuredContent"),
|
||
"raw_result": call_result,
|
||
}
|
||
except Exception as exc:
|
||
return {
|
||
"success": False,
|
||
"server_id": binding.server_id,
|
||
"server_name": binding.server_name,
|
||
"tool_alias": binding.alias,
|
||
"tool_name": binding.remote_name,
|
||
"transport": binding.transport,
|
||
"error": str(exc),
|
||
}
|