Files
bolt/docs/bolt_stdio_test.py
T
jocay 84653bec61 feat: v0.2.0 文件传输(MCP 工具 + HTTP 直传)与安全/正确性修复
文件功能:
- 新增 file_read / file_write / file_stat 工具(文本/base64/图片,分块读取,原子替换写入,SHA-256)
- 新增 file_link 工具与 GET/PUT /files 端点:HMAC 签名限时链接或 Bearer 认证,
  支持 Range、HEAD、chunked、Expect: 100-continue,大文件不经过模型上下文
- 新增 --file-root / --max-upload;配置白名单但未限定目录时文件功能自动禁用

安全:
- 白名单改为程序名精确匹配,shell 模式禁止元字符,shell_write 输入逐行受检;黑名单检查每一段
- HTTP 签发 Mcp-Session-Id,按会话隔离 shell 会话与取消,支持 DELETE 终止会话
- HTTP 客户端不能再用 notifications/exit 关停服务;token 常量时间比较

正确性:
- 后台进程占用输出管道时 run_command 不再挂起(detached=true)
- shell_read 改用绝对偏移,修复缓冲裁剪后错位;被信号杀死的会话能报告退出;阻塞读可取消
- 输出与分块读取不再切断 UTF-8 多字节字符
- shell_* 结果改为标准 content + structuredContent
- POST + SSE 的 tools/call 在同一响应流中推送进度与结果
- /mcp 支持 Expect: 100-continue(curl 发送大请求时不再白等 1 秒)
- CLI 参数优先于环境变量;SIGTERM 在 stdio 模式下也能清理退出;Windows 支持 Ctrl+C 清理与优雅退出

资源:并发工具调用 / HTTP 连接 / shell 会话数量上限,空闲会话自动回收

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-18 10:23:15 +08:00

217 lines
9.3 KiB
Python

#!/usr/bin/env python3
"""bolt stdio 模式端到端协议测试驱动"""
import faulthandler, json, os, subprocess, sys, threading, time
# 二进制路径:环境变量 BOLT_BIN 优先,默认为仓库内的 release 产物
BOLT = os.environ.get("BOLT_BIN") or os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "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", "file_read", "file_write", "file_stat"], 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 会话(结果为标准 MCP 格式:content + structuredContent)
i = send("tools/call", {"name": "shell_open", "arguments": {}})
r = wait_id(i)
sid = r["result"]["structuredContent"]["session_id"]
check("shell open", isinstance(sid, int) and isinstance(r["result"]["content"], list), str(r))
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"]["structuredContent"]
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"]["structuredContent"]["closed"] is True)
# 9b. shell_read 阻塞等待时可被取消
i = send("tools/call", {"name": "shell_open", "arguments": {}})
sid = wait_id(i)["result"]["structuredContent"]["session_id"]
i = send("tools/call", {"name": "shell_read", "arguments": {"session_id": sid, "timeout_ms": 0}})
time.sleep(0.3)
send_notif("notifications/cancelled", {"requestId": i})
r = wait_id(i, timeout=5)
check("shell_read cancel", r["result"]["structuredContent"]["cancelled"] is True, str(r))
wait_id(send("tools/call", {"name": "shell_close", "arguments": {"session_id": sid}}))
# 9c. 后台进程占用输出管道时不阻塞
t0 = time.time()
i = send("tools/call", {"name": "run_command", "arguments": {"command": "(sleep 5; echo late) & echo now"}})
r = wait_id(i)
sc = r["result"]["structuredContent"]
check("background detach", sc["stdout"] == "now\n" and sc["detached"] and time.time() - t0 < 3, f"{sc} dt={time.time() - t0:.1f}")
# 9d. 文件读写(文本 / base64 / 分块)
import base64, hashlib, tempfile
tmpdir = tempfile.mkdtemp(prefix="bolt-e2e-")
fp = os.path.join(tmpdir, "sub", "hello.txt")
i = send("tools/call", {"name": "file_write", "arguments": {"path": fp, "content": "你好, bolt\n", "mkdirs": True}})
sc = wait_id(i)["result"]["structuredContent"]
check("file_write text", sc["sha256"] == hashlib.sha256("你好, bolt\n".encode()).hexdigest(), str(sc))
i = send("tools/call", {"name": "file_read", "arguments": {"path": fp, "length": 4}})
content = wait_id(i)["result"]["content"]
meta = json.loads(content[0]["text"])
check("file_read chunk utf8", content[1]["text"] == "你" and meta["next_offset"] == 3 and not meta["eof"], str(content))
blob = bytes(range(256))
bp = os.path.join(tmpdir, "blob.bin")
i = send("tools/call", {"name": "file_write", "arguments": {"path": bp, "content": base64.b64encode(blob).decode(), "encoding": "base64"}})
wait_id(i)
i = send("tools/call", {"name": "file_read", "arguments": {"path": bp}})
content = wait_id(i)["result"]["content"]
check("file_read binary", json.loads(content[0]["text"])["encoding"] == "base64" and base64.b64decode(content[1]["text"]) == blob)
i = send("tools/call", {"name": "file_stat", "arguments": {"path": bp, "sha256": True}})
sc = wait_id(i)["result"]["structuredContent"]
check("file_stat sha256", sc["size"] == 256 and sc["sha256"] == hashlib.sha256(blob).hexdigest(), str(sc))
i = send("tools/call", {"name": "file_link", "arguments": {"path": bp}})
r = wait_id(i)
check("file_link stdio unavailable", "error" in r and r["error"]["code"] == -32602 or r["result"].get("isError") is True, str(r))
import shutil; shutil.rmtree(tmpdir, ignore_errors=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)