bolt: 极速单文件 MCP 命令执行服务器

- 手写 JSON-RPC 2.0 + MCP (2025-06-18) 协议实现,仅依赖 serde_json + libc - 双传输:stdio(本地 AI 客户端)与 Streamable HTTP(远程连接) - 工具:run_command(超时/取消/进度通知/进程树杀除)与 shell 持久会话四件套 - 安全:命令白黑名单、Bearer 认证、Origin 校验 - LTO fat + codegen-units=1,ping 吞吐约 1.4 万 msg/s - 8 个单元测试 + stdio/HTTP 端到端集成测试(31 项)
2026-08-27 14:19:10 +08:00
commit 29b7dc3bbb
7 changed files with 2469 additions and 0 deletions
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""bolt stdio 模式端到端协议测试驱动"""
+import faulthandler, json, os, subprocess, sys, threading, time
+
+
+BOLT = "/root/bolt/target/release/bolt"
+proc = subprocess.Popen([BOLT], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0)
+sent = {}
+next_id = [0]
+
+def send(method, params=None):
+    next_id[0] += 1
+    msg = {"jsonrpc": "2.0", "id": next_id[0], "method": method}
+    if params is not None:
+        msg["params"] = params
+    proc.stdin.write((json.dumps(msg) + "\n").encode())
+    proc.stdin.flush()
+    return next_id[0]
+
+def send_notif(method, params=None):
+    msg = {"jsonrpc": "2.0", "method": method}
+    if params is not None:
+        msg["params"] = params
+    proc.stdin.write((json.dumps(msg) + "\n").encode())
+    proc.stdin.flush()
+
+def recv(timeout=30):
+    import select
+    end = time.time() + timeout
+    line = b""
+    while True:
+        r, _, _ = select.select([proc.stdout], [], [], max(0.0, end - time.time()))
+        if not r:
+            raise TimeoutError(f"timeout, got so far: {line[:200]}")
+        ch = proc.stdout.read(1)
+        if not ch:
+            raise TimeoutError("EOF")
+        if ch == b"\n":
+            return json.loads(line)
+        line += ch
+
+received = {}  # 乱序到达的响应缓存(id -> msg)
+def wait_id(i, timeout=30):
+    if i in received:
+        return received.pop(i)
+    deadline = time.time() + timeout
+    pend.setdefault("extra", []).append(f"awaiting {i}")
+    while time.time() < deadline:
+        try:
+            m = recv(timeout=min(3, deadline - time.time()))
+        except TimeoutError as e:
+            raise TimeoutError(f"no response for id {i}: {e}; extras: {json.dumps(pend.get('extra', []))[:800]}")
+        if m.get("id") == i:
+            return m
+        if m.get("id") is not None:
+            received[m["id"]] = m
+        else:
+            pend.setdefault("extra", []).append(m)
+    raise TimeoutError(f"no response for id {i}, extras: {json.dumps(pend.get('extra', []))[:800]}")
+
+pend = {}
+passed = failed = 0
+def check(name, cond, detail=""):
+    global passed, failed
+    if cond:
+        passed += 1
+        print(f"  PASS {name}")
+    else:
+        failed += 1
+        print(f"  FAIL {name} {detail}")
+
+# 1. initialize
+i = send("initialize", {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "test", "version": "1"}})
+r = wait_id(i)
+check("initialize", r["result"]["protocolVersion"] == "2025-06-18" and r["result"]["serverInfo"]["name"] == "bolt", str(r))
+send_notif("notifications/initialized")
+
+# 2. tools/list
+i = send("tools/list")
+r = wait_id(i)
+names = [t["name"] for t in r["result"]["tools"]]
+check("tools/list", names == ["run_command", "shell_open", "shell_write", "shell_read", "shell_close"], str(names))
+
+# 3. run_command 基本
+i = send("tools/call", {"name": "run_command", "arguments": {"command": "echo hello; printf 'oops' >&2; exit 7"}})
+r = wait_id(i)
+sc = r["result"]["structuredContent"]
+check("run basic", sc["exit_code"] == 7 and sc["stdout"] == "hello\n" and sc["stderr"] == "oops", str(sc))
+
+# 4. 超时
+t0 = time.time()
+i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30", "timeout_ms": 400}})
+r = wait_id(i)
+dt = time.time() - t0
+sc = r["result"]["structuredContent"]
+check("timeout kill", sc["timed_out"] and not sc["cancelled"] and dt < 5, f"{sc} dt={dt:.1f}")
+
+# 5. 取消:发一个长命令,再发 cancelled
+i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30 && echo done"}})
+time.sleep(0.4)
+send_notif("notifications/cancelled", {"requestId": i})
+r = wait_id(i)
+sc = r["result"]["structuredContent"]
+check("cancel kill", sc["cancelled"] and sc["exit_code"] is None, str(sc))
+
+# 6. 进程组连带杀:sh -c 'sleep 30 & sleep 30' 的子进程也应变僵尸
+i = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 30 & sleep 30", "timeout_ms": 400}})
+r = wait_id(i)
+check("kill whole tree", r["result"]["structuredContent"]["timed_out"])
+
+# 7. 进度通知(_meta.progressToken)
+i = send("tools/call", {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo tick-$x; sleep 0.3; done"}, "_meta": {"progressToken": "tok-1"}})
+got_progress = False
+deadline = time.time() + 10
+while time.time() < deadline:
+    m = recv(timeout=deadline - time.time())
+    if m.get("method") == "notifications/progress" and m.get("params", {}).get("progressToken") == "tok-1":
+        got_progress = True
+        break
+    if m.get("id") == i:
+        break
+check("progress notification", got_progress)
+# 排空管道:打印残留
+import select
+import os
+while True:
+    r, _, _ = select.select([proc.stdout], [], [], 0)
+    if not r: break
+    try:
+        m = recv(timeout=2)
+        print("  LEFT:", m.get("id"), m.get("method"))
+    except Exception:
+        break
+
+# 8. 并发:同时两个命令
+i1 = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 0.6; echo ONE"}})
+i2 = send("tools/call", {"name": "run_command", "arguments": {"command": "sleep 0.2; echo TWO"}})
+r1, r2 = wait_id(i1), wait_id(i2)
+check("concurrent", r1["result"]["structuredContent"]["stdout"].strip() == "ONE" and r2["result"]["structuredContent"]["stdout"].strip() == "TWO")
+
+# 9. shell 会话
+i = send("tools/call", {"name": "shell_open", "arguments": {}})
+r = wait_id(i)
+sid = r["result"]["session_id"]
+check("shell open", isinstance(sid, int))
+send("tools/call", {"name": "shell_write", "arguments": {"session_id": sid, "input": "cd /tmp && echo IN_TMP; echo DONE_$$; exit"}})
+# 输出先于退出标志到达,循环读取直到会话退出
+all_out = ""
+for _ in range(10):
+    i = send("tools/call", {"name": "shell_read", "arguments": {"session_id": sid, "timeout_ms": 3000}})
+    r = wait_id(i)
+    res = r["result"]
+    all_out += res["output"]
+    if res["exited"]:
+        break
+check("shell read output", "IN_TMP" in all_out and res["exited"], f"out={all_out!r}")
+i = send("tools/call", {"name": "shell_close", "arguments": {"session_id": sid}})
+r = wait_id(i)
+check("shell close", r["result"]["closed"] is True)
+
+# 10. stdio 关闭 → 进程退出
+send_notif("notifications/exit")
+try:
+    proc.wait(timeout=3)
+    check("graceful exit", True)
+except subprocess.TimeoutExpired:
+    check("graceful exit", False)
+
+err = proc.stderr.read().decode()
+print(f"\n=== {passed} passed, {failed} failed ===")
+if err.strip():
+    print("stderr:", err[:500])
+sys.exit(1 if failed else 0)