shuiyu67andWorkBuddy
749fb3f91c
- 会话数据改用 SQLite（chat_history/conversations.db，WAL）：sessions / snapshots 两表， 首次启动自动迁移旧 JSON，/chat 与 /resume 改走 SQLite，memory 搜索兼容数据库 - 去掉流式显示（_typeprint 改整段打印）与常驻加载动画（移除 show_animation 线程及逐帧输出） - 展示给用户的消息过长自动截断（仅展示层，底层 messages 完整保存） - 工具调用解析加固：_loads_json 四级兜底修复；新增 _extract_last_action_json 成对错括号扫描； 续轮也下发 Function Calling 工具，减少长文本格式错误浪费 token - 版本号统一更新至 8.0.4（README/about.txt/constants/TUI banner/LSP·MCP clientInfo） - 12 项 session 单测 + 8 项 JSON 提取边界用例全部通过 Co-Authored-By: WorkBuddy <workbuddy@tencent.com>
307 lines
11 KiB
Python
307 lines
11 KiB
Python
"""
|
||
小狸 MCP 客户端
|
||
===============
|
||
连接外部 MCP (Model Context Protocol) server，将其工具桥接进小狸的统一工具体系。
|
||
|
||
- transport = stdio（subprocess + 逐行 JSON-RPC 2.0），与项目根 mcp_server.py（暴露侧）对称。
|
||
- 同步 + 线程安全：契合小狸「同步主循环 + ThreadPoolExecutor」模型，不引入官方 mcp SDK
|
||
（其 asyncio 风格与现有同步调用链不搭）。
|
||
- 单个 server 失败不阻断其他 server，也不阻断主程序启动。
|
||
|
||
后续可在此文件扩展 SSE / HTTP transport（McpClient 子类即可）。
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
import threading
|
||
import subprocess
|
||
import logging
|
||
from typing import Dict, List, Any, Optional, Tuple
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
PROTOCOL_VERSION = "2024-11-05"
|
||
|
||
|
||
class McpClientError(RuntimeError):
|
||
"""MCP 客户端错误"""
|
||
|
||
|
||
class McpClient:
|
||
"""单个 MCP server 的 stdio 客户端（同步、线程安全）"""
|
||
|
||
def __init__(self, name: str, command: str, args: Optional[List[str]] = None,
|
||
env: Optional[Dict[str, str]] = None, cwd: Optional[str] = None,
|
||
timeout: float = 30.0):
|
||
self.name = name
|
||
self.command = command
|
||
self.args = list(args or [])
|
||
self.env = env or {}
|
||
self.cwd = cwd
|
||
self.timeout = timeout
|
||
|
||
self.proc = None
|
||
self._lock = threading.RLock()
|
||
self._next_id = 1
|
||
self._pending: Dict[int, Dict[str, Any]] = {}
|
||
self._reader_thread = None
|
||
self._started = False
|
||
|
||
# ── 生命周期 ──
|
||
def start(self) -> None:
|
||
"""启动子进程并开启 reader 线程"""
|
||
if self._started:
|
||
return
|
||
full_env = dict(os.environ)
|
||
full_env.update(self.env)
|
||
cmd = [self.command] + self.args
|
||
try:
|
||
self.proc = subprocess.Popen(
|
||
cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.DEVNULL,
|
||
cwd=self.cwd,
|
||
env=full_env,
|
||
text=True,
|
||
encoding="utf-8",
|
||
bufsize=1,
|
||
)
|
||
except Exception as e:
|
||
raise McpClientError(f"启动 MCP server '{self.name}' 失败: {e}")
|
||
|
||
self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
|
||
self._reader_thread.start()
|
||
self._started = True
|
||
|
||
def _reader_loop(self):
|
||
"""后台读线程：逐行解析 JSON-RPC 响应，按 id 投递到 _pending 表"""
|
||
try:
|
||
for raw in self.proc.stdout:
|
||
line = raw.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
msg = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
# 容错：跳过 server 可能混入 stdout 的非 JSON 噪声
|
||
continue
|
||
msg_id = msg.get("id")
|
||
if msg_id is None:
|
||
# 通知（如进度），忽略
|
||
continue
|
||
with self._lock:
|
||
pend = self._pending.pop(msg_id, None)
|
||
if pend is not None:
|
||
if "error" in msg:
|
||
pend["error"] = msg["error"]
|
||
else:
|
||
pend["value"] = msg.get("result")
|
||
pend["event"].set()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
# 连接断开：唤醒所有等待中的请求，避免永久阻塞
|
||
with self._lock:
|
||
pendings = list(self._pending.values())
|
||
self._pending.clear()
|
||
for pend in pendings:
|
||
pend["error"] = {"code": -1, "message": "MCP server 连接断开"}
|
||
pend["event"].set()
|
||
|
||
def _send(self, method: str, params: Any = None,
|
||
is_notification: bool = False) -> Optional[Any]:
|
||
"""发送一个 JSON-RPC 请求；非通知则等待响应（带超时）"""
|
||
if not self._started or self.proc is None:
|
||
raise McpClientError(f"MCP client '{self.name}' 未启动")
|
||
|
||
with self._lock:
|
||
if is_notification:
|
||
req = {"jsonrpc": "2.0", "method": method, "params": params or {}}
|
||
pend = None
|
||
req_id = None
|
||
else:
|
||
req_id = self._next_id
|
||
self._next_id += 1
|
||
pend = {"event": threading.Event(), "value": None, "error": None}
|
||
self._pending[req_id] = pend
|
||
req = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params or {}}
|
||
|
||
try:
|
||
self.proc.stdin.write(json.dumps(req, ensure_ascii=False) + "\n")
|
||
self.proc.stdin.flush()
|
||
except (BrokenPipeError, ValueError, OSError) as e:
|
||
raise McpClientError(f"向 MCP server '{self.name}' 写入失败: {e}")
|
||
|
||
if is_notification:
|
||
return None
|
||
|
||
if not pend["event"].wait(self.timeout):
|
||
with self._lock:
|
||
self._pending.pop(req_id, None)
|
||
raise McpClientError(f"MCP server '{self.name}' 响应超时: {method}")
|
||
if pend["error"]:
|
||
err = pend["error"]
|
||
raise McpClientError(
|
||
f"MCP server '{self.name}' 错误 [{err.get('code')}]: {err.get('message')}")
|
||
return pend["value"]
|
||
|
||
# ── MCP 协议方法 ──
|
||
def initialize(self) -> Dict[str, Any]:
|
||
"""握手：协商协议版本并通知初始化完成"""
|
||
result = self._send("initialize", {
|
||
"protocolVersion": PROTOCOL_VERSION,
|
||
"capabilities": {},
|
||
"clientInfo": {"name": "xiaoli-cli", "version": "8.0.4"},
|
||
})
|
||
# 通知 server 初始化完成（无响应，无需等待）
|
||
self._send("notifications/initialized", {}, is_notification=True)
|
||
return result or {}
|
||
|
||
def list_tools(self) -> List[Dict[str, Any]]:
|
||
"""列出 server 提供的工具定义"""
|
||
result = self._send("tools/list", {})
|
||
if isinstance(result, dict):
|
||
return result.get("tools", [])
|
||
return []
|
||
|
||
def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||
"""调用一个工具，返回文本结果"""
|
||
result = self._send("tools/call", {"name": tool_name, "arguments": arguments or {}})
|
||
return self._extract_content(result)
|
||
|
||
@staticmethod
|
||
def _extract_content(result: Any) -> str:
|
||
"""从 tools/call 的 result 提取文本结果"""
|
||
if result is None:
|
||
return ""
|
||
if isinstance(result, str):
|
||
return result
|
||
if isinstance(result, dict):
|
||
content = result.get("content")
|
||
if isinstance(content, list):
|
||
parts = []
|
||
for item in content:
|
||
if isinstance(item, dict):
|
||
if item.get("type") == "text" or "text" in item:
|
||
parts.append(str(item.get("text", "")))
|
||
return "\n".join(p for p in parts if p)
|
||
if isinstance(result.get("content"), str):
|
||
return result["content"]
|
||
# 部分 server 直接把结果放 result 顶层
|
||
if "isError" in result and result.get("isError"):
|
||
return f"[工具错误] {json.dumps(result, ensure_ascii=False)}"
|
||
return json.dumps(result, ensure_ascii=False)
|
||
return str(result)
|
||
|
||
def ping(self) -> bool:
|
||
try:
|
||
self._send("ping", {})
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def stop(self) -> None:
|
||
"""优雅关闭子进程，失败则强杀"""
|
||
if self.proc is None:
|
||
return
|
||
try:
|
||
self._send("shutdown", {})
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self._send("exit", {}, is_notification=True)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self.proc.stdin.close()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self.proc.wait(timeout=3)
|
||
except Exception:
|
||
try:
|
||
self.proc.kill()
|
||
except Exception:
|
||
pass
|
||
self._started = False
|
||
|
||
|
||
class McpClientManager:
|
||
"""管理多个外部 MCP server 客户端（按配置并发连接、聚合工具）"""
|
||
|
||
def __init__(self):
|
||
self.clients: Dict[str, McpClient] = {}
|
||
self._tool_index: List[Tuple[str, Dict[str, Any]]] = [] # (server_name, tool_def)
|
||
self._lock = threading.RLock()
|
||
self._cleaned = False
|
||
|
||
def connect_all(self, servers: Dict[str, Dict[str, Any]],
|
||
cwd: Optional[str] = None) -> None:
|
||
"""连接所有已声明且未禁用的 server。单个失败不阻断其他。
|
||
|
||
servers: {name: {"command", "args", "env", "cwd", "disabled"}}
|
||
"""
|
||
for name, spec in servers.items():
|
||
if not isinstance(spec, dict):
|
||
continue
|
||
if spec.get("disabled"):
|
||
logger.info(f"MCP server '{name}' 已禁用，跳过")
|
||
continue
|
||
command = spec.get("command")
|
||
if not command:
|
||
logger.warning(f"MCP server '{name}' 缺少 command，跳过")
|
||
continue
|
||
args = spec.get("args") or []
|
||
env = spec.get("env") or {}
|
||
server_cwd = spec.get("cwd") or cwd
|
||
client = McpClient(name, command, args, env, server_cwd)
|
||
try:
|
||
client.start()
|
||
client.initialize()
|
||
tools = client.list_tools()
|
||
with self._lock:
|
||
self.clients[name] = client
|
||
for tool in tools:
|
||
self._tool_index.append((name, tool))
|
||
logger.info(f"已连接 MCP server '{name}': {len(tools)} 个工具")
|
||
except Exception as e:
|
||
try:
|
||
client.stop()
|
||
except Exception:
|
||
pass
|
||
logger.error(f"连接 MCP server '{name}' 失败: {e}")
|
||
# 不向上抛，继续下一个 server
|
||
|
||
def get_all_tools(self) -> List[Tuple[str, Dict[str, Any]]]:
|
||
"""返回 [(server_name, tool_def), ...]"""
|
||
with self._lock:
|
||
return list(self._tool_index)
|
||
|
||
def get_servers(self) -> List[str]:
|
||
with self._lock:
|
||
return list(self.clients.keys())
|
||
|
||
def call_tool(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||
with self._lock:
|
||
client = self.clients.get(server_name)
|
||
if client is None:
|
||
raise McpClientError(f"未连接的 MCP server: {server_name}")
|
||
return client.call_tool(tool_name, arguments)
|
||
|
||
def cleanup(self) -> None:
|
||
"""幂等：终止所有子进程。每个桥接工具的 shutdown 都会调一次，故必须可重入。"""
|
||
with self._lock:
|
||
if self._cleaned:
|
||
return
|
||
self._cleaned = True
|
||
clients = list(self.clients.values())
|
||
self.clients.clear()
|
||
self._tool_index.clear()
|
||
for client in clients:
|
||
try:
|
||
client.stop()
|
||
except Exception:
|
||
pass
|