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 项)
This commit is contained in:
jocay
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env python3
|
||||
"""bolt HTTP 模式端到端测试"""
|
||||
import http.client, json, subprocess, threading, time
|
||||
|
||||
BOLT = "/root/bolt/target/release/bolt"
|
||||
PORT = 18080
|
||||
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}")
|
||||
|
||||
proc = subprocess.Popen([BOLT, "--http", f"127.0.0.1:{PORT}"], stderr=subprocess.PIPE)
|
||||
time.sleep(0.5)
|
||||
|
||||
def post(obj, auth=None, origin=None, accept="application/json", path="/mcp"):
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=30)
|
||||
headers = {"Content-Type": "application/json", "Accept": accept}
|
||||
if auth: headers["Authorization"] = auth
|
||||
if origin: headers["Origin"] = origin
|
||||
body = json.dumps(obj)
|
||||
c.request("POST", path, body=body, headers=headers)
|
||||
r = c.getresponse()
|
||||
data = r.read()
|
||||
ct = r.getheader("Content-Type", "")
|
||||
c.close()
|
||||
return r.status, ct, data
|
||||
|
||||
# 1. initialize(JSON 响应)
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"}}})
|
||||
m = json.loads(data)
|
||||
check("init json", st == 200 and "application/json" in ct and m["result"]["serverInfo"]["name"] == "bolt", f"{st} {ct} {data[:100]}")
|
||||
|
||||
# 2. SSE 响应(accept 含 text/event-stream)
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 2, "method": "ping"}, accept="application/json, text/event-stream")
|
||||
sseline = data.split(b"data: ")[1].split(b"\r\n")[0].strip()
|
||||
m = json.loads(sseline)
|
||||
check("ping sse", st == 200 and "text/event-stream" in ct and m.get("result") == {}, f"{st} {ct} {data[:150]}")
|
||||
|
||||
# 3. 通知 → 202
|
||||
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/initialized"})
|
||||
check("notif 202", st == 202 and len(data) == 0, f"{st} {data[:50]}")
|
||||
|
||||
# 4. tools/call
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo HTTP_OK && exit 5"}}})
|
||||
m = json.loads(data)
|
||||
sc = m["result"]["structuredContent"]
|
||||
check("run over http", sc["exit_code"] == 5 and sc["stdout"].strip() == "HTTP_OK", str(sc))
|
||||
|
||||
# 5. GET SSE 流:开流收心跳,并发命令进度通知广播
|
||||
got_event = threading.Event()
|
||||
got_heartbeat = threading.Event()
|
||||
msgs = []
|
||||
def sse_get():
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=35)
|
||||
c.request("GET", "/mcp", headers={"Accept": "text/event-stream"})
|
||||
r = c.getresponse()
|
||||
assert r.status == 200 and "text/event-stream" in r.getheader("Content-Type")
|
||||
buf = b""
|
||||
end = time.time() + 30
|
||||
while time.time() < end:
|
||||
chunk = r.read(1)
|
||||
if not chunk: break
|
||||
buf += chunk
|
||||
if buf.endswith(b"\r\n\r\n"):
|
||||
msgs.append(buf.decode(errors="replace"))
|
||||
buf = b""
|
||||
if msgs and (": ping" in msgs[-1]):
|
||||
got_heartbeat.set()
|
||||
c.close()
|
||||
t = threading.Thread(target=sse_get)
|
||||
t.start()
|
||||
time.sleep(0.7)
|
||||
# 发一个带进度 token 的命令(约 1.2s)
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo h-$x; sleep 0.3; done"}, "_meta": {"progressToken": "h1"}}})
|
||||
m = json.loads(data)
|
||||
check("long cmd done", m["result"]["structuredContent"]["stdout"].count("h-") == 4, str(m)[:100])
|
||||
saw_progress = any("notifications/progress" in x and "h1" in x for x in msgs)
|
||||
check("sse got progress broadcast", saw_progress, str(msgs)[:200])
|
||||
t.join(timeout=25)
|
||||
saw_hb = got_heartbeat.is_set()
|
||||
check("sse heartbeat", saw_hb, str(msgs)[:200])
|
||||
|
||||
# 6. 取消:连接 A 长命令,连接 B 发 cancelled
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 5, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo nothing"}}}) # 预热(连接池无关,每次新连接)
|
||||
|
||||
def long_cmd():
|
||||
return post({"jsonrpc": "2.0", "id": 6, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "sleep 30"}}})
|
||||
res = {}
|
||||
t2 = threading.Thread(target=lambda: res.update(r=long_cmd()))
|
||||
t2.start()
|
||||
time.sleep(0.5)
|
||||
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 6}})
|
||||
check("cancel notif 202", st == 202)
|
||||
t2.join(timeout=10)
|
||||
m = json.loads(res["r"][2])
|
||||
check("http cancel works", m["result"]["structuredContent"]["cancelled"] is True, str(m)[:150])
|
||||
|
||||
# 7. 并发两个命令
|
||||
def run(cmd, rid):
|
||||
return post({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": cmd}}})
|
||||
results = {}
|
||||
ts = [threading.Thread(target=lambda c=cmd, i=i: results.update({i: run(c, i)})) for i, cmd in [(10, "sleep 0.5; echo A"), (11, "sleep 0.1; echo B")]]
|
||||
for t in ts: t.start()
|
||||
for t in ts: t.join(timeout=10)
|
||||
check("http concurrent", json.loads(results[10][2])["result"]["structuredContent"]["stdout"].strip() == "A" and json.loads(results[11][2])["result"]["structuredContent"]["stdout"].strip() == "B")
|
||||
|
||||
# 8. 错误路径
|
||||
st, _, _ = post({"jsonrpc": "2.0", "id": 99, "method": "nope"})
|
||||
check("unknown method", st == 200)
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
|
||||
c.request("POST", "/mcp", body=b"this is {not json", headers={"Content-Type": "application/json", "Accept": "application/json"})
|
||||
r = c.getresponse(); bd = r.read(); c.close()
|
||||
check("bad json -> 400", r.status == 400, str(r.status))
|
||||
st, _, data = post({"jsonrpc": "2.0", "id": 99, "method": "tools/call", "params": {"name": "nope", "arguments": {}}})
|
||||
check("unknown tool -32602", json.loads(data)["error"]["code"] == -32602)
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
|
||||
c.request("GET", "/other", headers={"Accept": "text/event-stream"})
|
||||
r = c.getresponse(); r.read(); c.close()
|
||||
check("404 other path", r.status == 404)
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
|
||||
c.request("PUT", "/mcp")
|
||||
r = c.getresponse(); r.read(); c.close()
|
||||
check("405 method", r.status == 405)
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
|
||||
c.request("OPTIONS", "/mcp")
|
||||
r = c.getresponse(); r.read(); c.close()
|
||||
check("options cors", r.status == 204 and r.getheader("Access-Control-Allow-Origin") == "*")
|
||||
|
||||
# 9. origin 校验
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 50, "method": "ping"}, origin="https://evil.example.com")
|
||||
check("evil origin 403", st == 403, str(st))
|
||||
st, ct, data = post({"jsonrpc": "2.0", "id": 51, "method": "ping"}, origin="http://localhost:4173")
|
||||
check("localhost origin ok", st == 200, str(st))
|
||||
|
||||
# 10. GET 无 Accept 头也返回 SSE(宽松)
|
||||
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
|
||||
c.request("GET", "/mcp")
|
||||
r = c.getresponse()
|
||||
c.close()
|
||||
check("get sse default", r.status == 200 and "text/event-stream" in r.getheader("Content-Type"))
|
||||
|
||||
proc.kill()
|
||||
err = proc.stderr.read().decode()
|
||||
print(f"\n=== HTTP: {passed} passed, {failed} failed ===")
|
||||
if err.strip():
|
||||
print("bolt stderr:", err[:400])
|
||||
sys_exit = 1 if failed else 0
|
||||
import sys; sys.exit(sys_exit)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user