jocay
29b7dc3bbb
- 手写 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 项)
1907 lines
69 KiB
Rust
1907 lines
69 KiB
Rust
//! bolt — 极速单文件 MCP 命令执行服务器
|
|
//!
|
|
//! AI(Claude Desktop / Claude Code / Cursor 等 MCP 客户端)连接后即可在服务器上执行命令。
|
|
//!
|
|
//! 特性:
|
|
//! * 零框架、纯手写 JSON-RPC 2.0 + MCP 协议(仅依赖 serde_json,Unix 下加 libc)
|
|
//! * 双传输:stdio(本地 AI 客户端默认方式)与 Streamable HTTP(远程连接)
|
|
//! * 工具:
|
|
//! - run_command 执行 shell 命令(支持管道/重定向/超时/环境变量/输出截断/进度通知)
|
|
//! - shell_open 开启持久化交互式 shell 会话(状态保留,可连续对话)
|
|
//! - shell_write 向会话写入输入
|
|
//! - shell_read 读取会话输出(可阻塞等待)
|
|
//! - shell_close 关闭会话
|
|
//! * 取消:客户端发送 notifications/cancelled 可杀掉进行中的命令(Unix 下按进程组杀,连带子进程)
|
|
//! * 安全:命令 allow/deny 前缀白黑名单、HTTP Bearer Token、Origin 校验、默认本机绑定
|
|
//! * 性能:单线程事件循环 + 复用缓冲区、原子计数、无锁热路径、LTO fat + codegen-units=1
|
|
//!
|
|
//! 用法:
|
|
//! bolt # stdio 模式(默认,MCP 客户端拉起本进程)
|
|
//! bolt --http 127.0.0.1:8080 # Streamable HTTP 模式,远程 AI 客户端连接
|
|
//! bolt --http 0.0.0.0:8080 --token mysecret --allow "git,docker,curl"
|
|
//!
|
|
//! 环境变量(CLI 参数优先):BOLT_TOKEN / BOLT_ALLOW / BOLT_DENY / BOLT_TIMEOUT_MS /
|
|
//! BOLT_CWD / BOLT_MAX_OUTPUT / BOLT_ALLOW_ORIGIN
|
|
//!
|
|
//! 客户端配置示例:
|
|
//! Claude Code(本地): claude mcp add --transport stdio bolt -- /abs/path/bolt
|
|
//! Claude Code(远程): claude mcp add --transport http bolt http://host:8080/mcp
|
|
//! 通用 MCP 客户端: {"mcpServers": {"bolt": {"command": "/abs/path/bolt"}}}
|
|
|
|
use std::collections::HashMap;
|
|
use std::env;
|
|
use std::io::{self, BufRead, BufReader, BufWriter, Read, Write};
|
|
use std::net::{TcpListener, TcpStream};
|
|
use std::path::Path;
|
|
use std::process::{Command, ExitStatus, Stdio};
|
|
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
|
|
use std::sync::{mpsc, Arc, Condvar, Mutex, OnceLock};
|
|
use std::thread;
|
|
use std::time::{Duration, Instant};
|
|
|
|
#[cfg(unix)]
|
|
use std::os::unix::process::CommandExt;
|
|
|
|
use serde_json::{json, Value};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 常量
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const VERSION: &str = env!("CARGO_PKG_VERSION");
|
|
const PROTOCOL_VERSIONS: [&str; 3] = ["2024-11-05", "2025-03-26", "2025-06-18"];
|
|
const LATEST_PROTOCOL: &str = "2025-06-18";
|
|
|
|
const MAX_LINE_BYTES: usize = 16 * 1024 * 1024; // stdio 单条消息上限
|
|
const MAX_HTTP_HEAD: usize = 64 * 1024; // HTTP 请求头上限
|
|
const MAX_HTTP_BODY: usize = 4 * 1024 * 1024; // HTTP body 上限
|
|
const DEFAULT_TIMEOUT_MS: u64 = 30_000; // 命令默认超时
|
|
const MAX_TIMEOUT_MS: u64 = 6 * 3600 * 1000; // 命令超时上限(6h)
|
|
const DEFAULT_MAX_OUTPUT: usize = 1024 * 1024; // 单流默认捕获上限
|
|
const MAX_MAX_OUTPUT: usize = 64 * 1024 * 1024;
|
|
const SESSION_BUF_CAP: usize = 16 * 1024 * 1024; // 会话环形缓冲上限
|
|
const SESSION_WRITE_TIMEOUT: Duration = Duration::from_secs(10); // 会话写入超时
|
|
const SSE_HEARTBEAT: Duration = Duration::from_secs(10); // SSE 保活间隔
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 配置
|
|
// ---------------------------------------------------------------------------
|
|
|
|
enum Mode {
|
|
Stdio,
|
|
Http(String),
|
|
}
|
|
|
|
struct Config {
|
|
mode: Mode,
|
|
token: Option<String>,
|
|
allow: Vec<String>,
|
|
deny: Vec<String>,
|
|
timeout_ms: u64,
|
|
cwd: Option<String>,
|
|
max_output: usize,
|
|
allow_origins: Vec<String>,
|
|
}
|
|
|
|
fn split_list(s: &str) -> Vec<String> {
|
|
s.split(',')
|
|
.map(str::trim)
|
|
.filter(|x| !x.is_empty())
|
|
.map(str::to_string)
|
|
.collect()
|
|
}
|
|
|
|
/// 解析 "1024" / "64K" / "1M" / "2G" 形式的字节数
|
|
fn parse_size(s: &str) -> Option<usize> {
|
|
let s = s.trim();
|
|
let (num, mult) = match s.chars().last() {
|
|
Some('k') | Some('K') => (&s[..s.len() - 1], 1 << 10),
|
|
Some('m') | Some('M') => (&s[..s.len() - 1], 1 << 20),
|
|
Some('g') | Some('G') => (&s[..s.len() - 1], 1 << 30),
|
|
_ => (s, 1),
|
|
};
|
|
num.parse::<usize>().ok().map(|n| n.saturating_mul(mult))
|
|
}
|
|
|
|
fn parse_args() -> Option<Config> {
|
|
let mut mode = Mode::Stdio;
|
|
let mut token: Option<String> = None;
|
|
let mut allow: Vec<String> = Vec::new();
|
|
let mut deny: Vec<String> = Vec::new();
|
|
let mut timeout_ms = DEFAULT_TIMEOUT_MS;
|
|
let mut cwd: Option<String> = None;
|
|
let mut max_output = DEFAULT_MAX_OUTPUT;
|
|
let mut allow_origins: Vec<String> = Vec::new();
|
|
|
|
let args: Vec<String> = env::args().skip(1).collect();
|
|
let mut i = 0;
|
|
while i < args.len() {
|
|
let a = args[i].as_str();
|
|
let next = |i: &mut usize| -> Option<String> {
|
|
*i += 1;
|
|
args.get(*i).cloned()
|
|
};
|
|
match a {
|
|
"-h" | "--help" => return None,
|
|
"--http" => mode = Mode::Http(next(&mut i)?),
|
|
"--token" => token = next(&mut i),
|
|
"--allow" => allow = split_list(&next(&mut i)?),
|
|
"--deny" => deny = split_list(&next(&mut i)?),
|
|
"--timeout-ms" => timeout_ms = next(&mut i)?.parse().ok()?,
|
|
"--cwd" => cwd = next(&mut i),
|
|
"--max-output" => max_output = parse_size(&next(&mut i)?)?,
|
|
"--allow-origin" => allow_origins.push(next(&mut i)?),
|
|
_ => {
|
|
eprintln!("[bolt] 未知参数: {a}");
|
|
return None;
|
|
}
|
|
}
|
|
i += 1;
|
|
}
|
|
// 环境变量回退(CLI 优先)
|
|
if token.is_none() {
|
|
token = env::var("BOLT_TOKEN").ok().filter(|s| !s.is_empty());
|
|
}
|
|
if allow.is_empty() {
|
|
allow = split_list(&env::var("BOLT_ALLOW").unwrap_or_default());
|
|
}
|
|
if deny.is_empty() {
|
|
deny = split_list(&env::var("BOLT_DENY").unwrap_or_default());
|
|
}
|
|
if let Ok(v) = env::var("BOLT_TIMEOUT_MS") {
|
|
if let Ok(n) = v.parse() {
|
|
timeout_ms = n;
|
|
}
|
|
}
|
|
if cwd.is_none() {
|
|
cwd = env::var("BOLT_CWD").ok().filter(|s| !s.is_empty());
|
|
}
|
|
if let Ok(v) = env::var("BOLT_MAX_OUTPUT") {
|
|
if let Some(n) = parse_size(&v) {
|
|
max_output = n;
|
|
}
|
|
}
|
|
if allow_origins.is_empty() {
|
|
allow_origins = split_list(&env::var("BOLT_ALLOW_ORIGIN").unwrap_or_default());
|
|
}
|
|
Some(Config {
|
|
mode,
|
|
token,
|
|
allow,
|
|
deny,
|
|
timeout_ms,
|
|
cwd,
|
|
max_output,
|
|
allow_origins,
|
|
})
|
|
}
|
|
|
|
fn print_help() {
|
|
println!(
|
|
r#"bolt {VERSION} — 极速单文件 MCP 命令执行服务器
|
|
AI 通过 MCP(stdio 或 Streamable HTTP)连接后即可在服务器上执行命令。
|
|
|
|
用法:
|
|
bolt [选项]
|
|
|
|
选项:
|
|
--http <addr> 启用 Streamable HTTP 传输并监听(如 127.0.0.1:8080 或 :8080)
|
|
缺省为 stdio 模式(由 MCP 客户端拉起本进程)
|
|
--token <token> HTTP 模式要求 Authorization: Bearer <token>
|
|
--allow <a,b,...> 命令白名单(命令首词前缀匹配,非空时仅允许列出的命令)
|
|
--deny <a,b,...> 命令黑名单(前缀匹配)
|
|
--timeout-ms <n> 命令默认超时毫秒(0 = 无限,工具参数可覆盖;默认 30000)
|
|
--cwd <dir> 命令默认工作目录
|
|
--max-output <n> 单流输出捕获上限,支持 K/M/G 后缀(默认 1M)
|
|
--allow-origin <url> 额外允许的 HTTP Origin(默认仅 localhost/127.0.0.1)
|
|
-h, --help 显示帮助
|
|
|
|
环境变量(CLI 优先):BOLT_TOKEN BOLT_ALLOW BOLT_DENY BOLT_TIMEOUT_MS BOLT_CWD BOLT_MAX_OUTPUT BOLT_ALLOW_ORIGIN
|
|
|
|
客户端配置:
|
|
Claude Code 本地: claude mcp add --transport stdio bolt -- /abs/path/bolt
|
|
Claude Code 远程: claude mcp add --transport http bolt http://host:8080/mcp
|
|
通用 MCP 客户端: {{"mcpServers": {{"bolt": {{"command": "/abs/path/bolt"}}}}}}
|
|
"#
|
|
);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 出站通道:stdio 单写者锁 / HTTP 广播到 SSE 流
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct Hub {
|
|
map: Mutex<HashMap<u64, mpsc::Sender<String>>>,
|
|
next: AtomicU64,
|
|
}
|
|
|
|
impl Hub {
|
|
fn add(&self, tx: mpsc::Sender<String>) -> u64 {
|
|
let id = self.next.fetch_add(1, Ordering::Relaxed) + 1;
|
|
self.map.lock().unwrap().insert(id, tx);
|
|
id
|
|
}
|
|
fn remove(&self, id: u64) {
|
|
self.map.lock().unwrap().remove(&id);
|
|
}
|
|
/// 广播(发送失败即剔除死连接;mpsc 无界通道,不阻塞)
|
|
fn broadcast(&self, msg: &str) {
|
|
let mut map = self.map.lock().unwrap();
|
|
map.retain(|_, tx| tx.send(msg.to_string()).is_ok());
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
enum Sink {
|
|
Stdio(Arc<Mutex<io::Stdout>>),
|
|
Http(Arc<Hub>),
|
|
}
|
|
|
|
impl Sink {
|
|
#[inline]
|
|
fn send(&self, msg: &str) {
|
|
match self {
|
|
Sink::Stdio(o) => {
|
|
if let Ok(mut g) = o.lock() {
|
|
let _ = g.write_all(msg.as_bytes());
|
|
let _ = g.write_all(b"\n");
|
|
let _ = g.flush();
|
|
}
|
|
}
|
|
Sink::Http(h) => h.broadcast(msg),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 应用状态
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct App {
|
|
cfg: Config,
|
|
sink: Sink,
|
|
hub: Arc<Hub>,
|
|
exit_flag: Arc<AtomicBool>,
|
|
inflight: Mutex<HashMap<String, Arc<Killable>>>, // 进行中的命令(id -> 可杀句柄)
|
|
sessions: Mutex<HashMap<u64, Arc<Session>>>,
|
|
next_session: AtomicU64,
|
|
}
|
|
|
|
type Shared = Arc<App>;
|
|
|
|
impl App {
|
|
fn new(cfg: Config) -> App {
|
|
let hub = Arc::new(Hub {
|
|
map: Mutex::new(HashMap::new()),
|
|
next: AtomicU64::new(0),
|
|
});
|
|
let sink = match &cfg.mode {
|
|
Mode::Stdio => Sink::Stdio(Arc::new(Mutex::new(io::stdout()))),
|
|
Mode::Http(_) => Sink::Http(hub.clone()),
|
|
};
|
|
App {
|
|
cfg,
|
|
sink,
|
|
hub,
|
|
exit_flag: Arc::new(AtomicBool::new(false)),
|
|
inflight: Mutex::new(HashMap::new()),
|
|
sessions: Mutex::new(HashMap::new()),
|
|
next_session: AtomicU64::new(0),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 进程组杀掉 pid 对应的整棵进程树
|
|
#[cfg(unix)]
|
|
fn kill_tree(pid: u32) {
|
|
unsafe {
|
|
libc::kill(-(pid as i32), libc::SIGKILL);
|
|
}
|
|
}
|
|
#[cfg(windows)]
|
|
fn kill_tree(pid: u32) {
|
|
let _ = Command::new("taskkill")
|
|
.args(["/PID", &pid.to_string(), "/T", "/F"])
|
|
.stdin(Stdio::null())
|
|
.stdout(Stdio::null())
|
|
.stderr(Stdio::null())
|
|
.status();
|
|
}
|
|
|
|
/// 可取消句柄:注册 pid 前收到取消则标记,注册后立即补杀
|
|
struct Killable {
|
|
pid: AtomicU32,
|
|
want_kill: AtomicBool,
|
|
cancelled: AtomicBool,
|
|
}
|
|
|
|
impl Default for Killable {
|
|
fn default() -> Self {
|
|
Killable {
|
|
pid: AtomicU32::new(0),
|
|
want_kill: AtomicBool::new(false),
|
|
cancelled: AtomicBool::new(false),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Killable {
|
|
fn register_pid(&self, pid: u32) {
|
|
self.pid.store(pid, Ordering::SeqCst);
|
|
if self.want_kill.load(Ordering::SeqCst) {
|
|
kill_tree(pid);
|
|
}
|
|
}
|
|
fn kill(&self, mark_cancelled: bool) {
|
|
if mark_cancelled {
|
|
self.cancelled.store(true, Ordering::SeqCst);
|
|
}
|
|
self.want_kill.store(true, Ordering::SeqCst);
|
|
let p = self.pid.load(Ordering::SeqCst);
|
|
if p != 0 {
|
|
kill_tree(p);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn id_key(id: &Value) -> String {
|
|
serde_json::to_string(id).unwrap_or_default()
|
|
}
|
|
|
|
fn register_inflight(app: &Shared, id: &Value) -> Arc<Killable> {
|
|
let k = Arc::new(Killable::default());
|
|
app.inflight.lock().unwrap().insert(id_key(id), k.clone());
|
|
k
|
|
}
|
|
|
|
fn unregister_inflight(app: &Shared, id: &Value, k: &Arc<Killable>) {
|
|
let mut m = app.inflight.lock().unwrap();
|
|
if m.get(&id_key(id)).map(|x| Arc::ptr_eq(x, k)).unwrap_or(false) {
|
|
m.remove(&id_key(id));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// JSON-RPC / MCP 辅助
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn rpc_error(id: &Value, code: i64, message: &str) -> Value {
|
|
json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}})
|
|
}
|
|
|
|
const PARSE_ERROR: &str =
|
|
r#"{"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"parse error"}}"#;
|
|
|
|
/// 工具执行错误(非协议错误):结果中 isError: true
|
|
fn tool_error(msg: &str) -> Value {
|
|
json!({"content": [{"type": "text", "text": format!("error: {msg}")}], "isError": true})
|
|
}
|
|
|
|
/// 命令白/黑名单检查(first_token 为命令首词)
|
|
fn policy_check(cfg: &Config, first_token: &str) -> Result<(), String> {
|
|
if !cfg.deny.is_empty() && cfg.deny.iter().any(|d| first_token.starts_with(d)) {
|
|
return Err(format!("命令 '{first_token}' 被服务器策略禁止"));
|
|
}
|
|
if !cfg.allow.is_empty() && !cfg.allow.iter().any(|a| first_token.starts_with(a)) {
|
|
return Err(format!("命令 '{first_token}' 不在服务器白名单中"));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 工具:run_command 执行引擎
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct RunReq {
|
|
command: String,
|
|
args: Vec<String>,
|
|
shell: bool,
|
|
cwd: Option<String>,
|
|
env: Vec<(String, String)>,
|
|
timeout_ms: u64, // 0 = 无限
|
|
max_output: usize,
|
|
}
|
|
|
|
fn parse_run_params(p: &Value, cfg: &Config) -> Result<RunReq, String> {
|
|
let command = p
|
|
.get("command")
|
|
.and_then(Value::as_str)
|
|
.ok_or("'command' 必须是非空字符串")?
|
|
.to_string();
|
|
let shell = p.get("shell").and_then(Value::as_bool).unwrap_or(true);
|
|
let args: Vec<String> = match p.get("args") {
|
|
None => Vec::new(),
|
|
Some(Value::Array(a)) => a
|
|
.iter()
|
|
.map(|x| {
|
|
x.as_str()
|
|
.map(str::to_string)
|
|
.ok_or_else(|| "'args' 必须是字符串数组".to_string())
|
|
})
|
|
.collect::<Result<_, _>>()?,
|
|
Some(_) => return Err("'args' 必须是字符串数组".into()),
|
|
};
|
|
if shell && !args.is_empty() {
|
|
return Err("shell=true 时不能同时使用 'args'(shell 语法直接写在 command 里)".into());
|
|
}
|
|
let cwd = match p.get("cwd") {
|
|
None => cfg.cwd.clone(),
|
|
Some(v) => Some(v.as_str().ok_or("'cwd' 必须是字符串")?.to_string()),
|
|
};
|
|
let mut env_map: Vec<(String, String)> = Vec::new();
|
|
if let Some(e) = p.get("env") {
|
|
let obj = e.as_object().ok_or("'env' 必须是对象")?;
|
|
for (k, v) in obj {
|
|
let s = v.as_str().ok_or_else(|| format!("env 的值 '{k}' 必须是字符串"))?;
|
|
env_map.push((k.clone(), s.to_string()));
|
|
}
|
|
}
|
|
let timeout_ms = match p.get("timeout_ms") {
|
|
None => cfg.timeout_ms,
|
|
Some(v) => {
|
|
let n = v.as_u64().ok_or("'timeout_ms' 必须是整数")?;
|
|
if n > MAX_TIMEOUT_MS {
|
|
return Err(format!("'timeout_ms' 不能超过 {MAX_TIMEOUT_MS}"));
|
|
}
|
|
n
|
|
}
|
|
};
|
|
let max_output = match p.get("max_output_bytes") {
|
|
None => cfg.max_output,
|
|
Some(v) => {
|
|
let n = v.as_u64().ok_or("'max_output_bytes' 必须是整数")? as usize;
|
|
if !(1024..=MAX_MAX_OUTPUT).contains(&n) {
|
|
return Err(format!("'max_output_bytes' 必须在 1024..{MAX_MAX_OUTPUT}"));
|
|
}
|
|
n
|
|
}
|
|
};
|
|
Ok(RunReq {
|
|
command,
|
|
args,
|
|
shell,
|
|
cwd,
|
|
env: env_map,
|
|
timeout_ms,
|
|
max_output,
|
|
})
|
|
}
|
|
|
|
struct ExecOutcome {
|
|
exit_code: Option<i32>,
|
|
signal: Option<i32>,
|
|
timed_out: bool,
|
|
cancelled: bool,
|
|
duration_ms: u64,
|
|
stdout: String,
|
|
stderr: String,
|
|
trunc_out: bool,
|
|
trunc_err: bool,
|
|
}
|
|
|
|
struct OutBuf {
|
|
bytes: Vec<u8>,
|
|
truncated: bool,
|
|
}
|
|
|
|
impl Default for OutBuf {
|
|
fn default() -> Self {
|
|
OutBuf {
|
|
bytes: Vec::new(),
|
|
truncated: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
struct OutShared {
|
|
stdout: Arc<Mutex<OutBuf>>,
|
|
stderr: Arc<Mutex<OutBuf>>,
|
|
total: Arc<AtomicU64>,
|
|
}
|
|
|
|
/// 管道读取循环:读取到 cap 后继续排空(防止子进程写满管道阻塞),只标记截断
|
|
fn pipe_read_loop<R: Read + Send + 'static>(
|
|
mut pipe: R,
|
|
out: Arc<Mutex<OutBuf>>,
|
|
total: Arc<AtomicU64>,
|
|
cap: usize,
|
|
) {
|
|
let mut chunk = [0u8; 16 * 1024];
|
|
loop {
|
|
match pipe.read(&mut chunk) {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
total.fetch_add(n as u64, Ordering::Relaxed);
|
|
let mut b = out.lock().unwrap();
|
|
if b.bytes.len() < cap {
|
|
let take = (cap - b.bytes.len()).min(n);
|
|
b.bytes.extend_from_slice(&chunk[..take]);
|
|
if take < n {
|
|
b.truncated = true;
|
|
}
|
|
} else {
|
|
b.truncated = true;
|
|
}
|
|
}
|
|
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 进度通知循环:每 200ms 上报已读字节数
|
|
fn progress_loop(token: Value, total: Arc<AtomicU64>, done: Arc<AtomicBool>, sink: Sink) {
|
|
let mut last = 0u64;
|
|
while !done.load(Ordering::Relaxed) {
|
|
thread::sleep(Duration::from_millis(200));
|
|
let t = total.load(Ordering::Relaxed);
|
|
if t != last {
|
|
last = t;
|
|
let msg = json!({
|
|
"jsonrpc": "2.0",
|
|
"method": "notifications/progress",
|
|
"params": {
|
|
"progressToken": token,
|
|
"progress": t,
|
|
"message": format!("已读取 {t} 字节输出"),
|
|
}
|
|
});
|
|
sink.send(&msg.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
fn exec(req: &RunReq, killable: &Killable, progress: Option<Value>, app: &Shared) -> Result<ExecOutcome, String> {
|
|
// 构建命令
|
|
let mut cmd = if req.shell {
|
|
#[cfg(windows)]
|
|
let mut c = Command::new("cmd");
|
|
#[cfg(windows)]
|
|
c.arg("/C").arg(&req.command);
|
|
#[cfg(unix)]
|
|
let mut c = Command::new("sh");
|
|
#[cfg(unix)]
|
|
c.arg("-c").arg(&req.command);
|
|
c
|
|
} else {
|
|
let mut c = Command::new(&req.command);
|
|
c.args(&req.args);
|
|
c
|
|
};
|
|
cmd.stdin(Stdio::null())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
if let Some(cwd) = &req.cwd {
|
|
if !Path::new(cwd).is_dir() {
|
|
return Err(format!("工作目录不存在: {cwd}"));
|
|
}
|
|
cmd.current_dir(cwd);
|
|
}
|
|
cmd.envs(req.env.iter().map(|(k, v)| (k.as_str(), v.as_str())));
|
|
#[cfg(unix)]
|
|
unsafe {
|
|
// 子进程放入独立进程组,便于整树杀除
|
|
cmd.pre_exec(|| {
|
|
libc::setpgid(0, 0);
|
|
Ok(())
|
|
});
|
|
}
|
|
|
|
let start = Instant::now();
|
|
let mut child = cmd
|
|
.spawn()
|
|
.map_err(|e| format!("启动命令失败: {e}"))?;
|
|
killable.register_pid(child.id());
|
|
|
|
let out = Arc::new(OutShared {
|
|
stdout: Arc::new(Mutex::new(OutBuf::default())),
|
|
stderr: Arc::new(Mutex::new(OutBuf::default())),
|
|
total: Arc::new(AtomicU64::new(0)),
|
|
});
|
|
let done = Arc::new(AtomicBool::new(false));
|
|
|
|
let stdout_pipe = child.stdout.take().unwrap();
|
|
let stderr_pipe = child.stderr.take().unwrap();
|
|
let t_out = {
|
|
let buf = out.stdout.clone();
|
|
let total = out.total.clone();
|
|
let cap = req.max_output;
|
|
thread::Builder::new()
|
|
.name("bolt-out".into())
|
|
.spawn(move || pipe_read_loop(stdout_pipe, buf, total, cap))
|
|
.unwrap()
|
|
};
|
|
let t_err = {
|
|
let buf = out.stderr.clone();
|
|
let total = out.total.clone();
|
|
let cap = req.max_output;
|
|
thread::Builder::new()
|
|
.name("bolt-err".into())
|
|
.spawn(move || pipe_read_loop(stderr_pipe, buf, total, cap))
|
|
.unwrap()
|
|
};
|
|
let t_prog = progress.map(|tok| {
|
|
let total = out.total.clone();
|
|
let done = done.clone();
|
|
let sink = app.sink.clone();
|
|
thread::Builder::new()
|
|
.name("bolt-progress".into())
|
|
.spawn(move || progress_loop(tok, total, done, sink))
|
|
.unwrap()
|
|
});
|
|
|
|
// 收割线程独占 child.wait(),避免僵尸进程
|
|
let (tx, rx) = mpsc::channel();
|
|
let mut child2 = child;
|
|
let reaper = thread::Builder::new()
|
|
.name("bolt-reaper".into())
|
|
.spawn(move || {
|
|
let st = child2.wait();
|
|
let _ = tx.send(st);
|
|
})
|
|
.unwrap();
|
|
|
|
let mut timed_out = false;
|
|
let status: Option<ExitStatus> = if req.timeout_ms == 0 {
|
|
rx.recv().ok().and_then(Result::ok)
|
|
} else {
|
|
match rx.recv_timeout(Duration::from_millis(req.timeout_ms)) {
|
|
Ok(Ok(st)) => Some(st),
|
|
Ok(Err(_)) => None,
|
|
Err(mpsc::RecvTimeoutError::Timeout) => {
|
|
timed_out = true;
|
|
killable.kill(false); // 超时杀,不计为客户端取消
|
|
rx.recv().ok().and_then(Result::ok)
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) => None,
|
|
}
|
|
};
|
|
|
|
done.store(true, Ordering::Relaxed);
|
|
let _ = t_out.join();
|
|
let _ = t_err.join();
|
|
if let Some(t) = t_prog {
|
|
let _ = t.join();
|
|
}
|
|
let _ = reaper.join();
|
|
|
|
let stdout = String::from_utf8_lossy(&out.stdout.lock().unwrap().bytes).into_owned();
|
|
let stderr = String::from_utf8_lossy(&out.stderr.lock().unwrap().bytes).into_owned();
|
|
let trunc_out = out.stdout.lock().unwrap().truncated;
|
|
let trunc_err = out.stderr.lock().unwrap().truncated;
|
|
|
|
Ok(ExecOutcome {
|
|
exit_code: status.as_ref().and_then(ExitStatus::code),
|
|
signal: status.as_ref().map(exit_signal).unwrap_or(None),
|
|
timed_out,
|
|
cancelled: killable.cancelled.load(Ordering::SeqCst),
|
|
duration_ms: start.elapsed().as_millis() as u64,
|
|
stdout,
|
|
stderr,
|
|
trunc_out,
|
|
trunc_err,
|
|
})
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
fn exit_signal(st: &ExitStatus) -> Option<i32> {
|
|
use std::os::unix::process::ExitStatusExt;
|
|
st.signal()
|
|
}
|
|
#[cfg(not(unix))]
|
|
fn exit_signal(_st: &ExitStatus) -> Option<i32> {
|
|
None
|
|
}
|
|
|
|
fn run_result(o: &ExecOutcome, cwd: Option<&str>) -> Value {
|
|
let result = json!({
|
|
"exit_code": o.exit_code,
|
|
"signal": o.signal,
|
|
"timed_out": o.timed_out,
|
|
"cancelled": o.cancelled,
|
|
"duration_ms": o.duration_ms,
|
|
"stdout": o.stdout,
|
|
"stderr": o.stderr,
|
|
"truncated_stdout": o.trunc_out,
|
|
"truncated_stderr": o.trunc_err,
|
|
"cwd": cwd,
|
|
});
|
|
json!({
|
|
"content": [{"type": "text", "text": result.to_string()}],
|
|
"structuredContent": result,
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 工具:shell 持久会话
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct Session {
|
|
pid: u32,
|
|
stdin: Mutex<Option<std::process::ChildStdin>>,
|
|
out: Arc<SessionOut>,
|
|
}
|
|
|
|
struct SessionOut {
|
|
bytes: Mutex<Vec<u8>>,
|
|
off: Mutex<usize>, // 已读偏移(相对当前 Vec,裁剪后可能失效)
|
|
gen: AtomicU64,
|
|
cv: Condvar,
|
|
exited: Mutex<Option<i32>>,
|
|
dropped: AtomicU64,
|
|
}
|
|
|
|
fn session_read_loop<R: Read + Send + 'static>(mut pipe: R, out: Arc<SessionOut>) {
|
|
let mut chunk = [0u8; 16 * 1024];
|
|
loop {
|
|
match pipe.read(&mut chunk) {
|
|
Ok(0) => break,
|
|
Ok(n) => {
|
|
let mut b = out.bytes.lock().unwrap();
|
|
if b.len() >= SESSION_BUF_CAP {
|
|
let drop_n = b.len() / 2;
|
|
out.dropped.fetch_add(drop_n as u64, Ordering::Relaxed);
|
|
b.drain(..drop_n);
|
|
}
|
|
b.extend_from_slice(&chunk[..n]);
|
|
out.gen.fetch_add(1, Ordering::Relaxed);
|
|
out.cv.notify_all();
|
|
}
|
|
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn shell_open(app: &Shared, shell: Option<&str>, cwd: Option<&str>) -> Result<Value, String> {
|
|
let shell = shell.unwrap_or(if cfg!(windows) { "cmd" } else { "sh" });
|
|
policy_check(&app.cfg, shell)?;
|
|
let mut cmd = Command::new(shell);
|
|
cmd.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
if let Some(c) = cwd {
|
|
if !Path::new(c).is_dir() {
|
|
return Err(format!("工作目录不存在: {c}"));
|
|
}
|
|
cmd.current_dir(c);
|
|
}
|
|
#[cfg(unix)]
|
|
unsafe {
|
|
cmd.pre_exec(|| {
|
|
libc::setpgid(0, 0);
|
|
Ok(())
|
|
});
|
|
}
|
|
let mut child = cmd
|
|
.spawn()
|
|
.map_err(|e| format!("启动 shell '{shell}' 失败: {e}"))?;
|
|
let pid = child.id();
|
|
let stdin = child.stdin.take().unwrap();
|
|
let stdout_pipe = child.stdout.take().unwrap();
|
|
let stderr_pipe = child.stderr.take().unwrap();
|
|
let out = Arc::new(SessionOut {
|
|
bytes: Mutex::new(Vec::new()),
|
|
off: Mutex::new(0),
|
|
gen: AtomicU64::new(0),
|
|
cv: Condvar::new(),
|
|
exited: Mutex::new(None),
|
|
dropped: AtomicU64::new(0),
|
|
});
|
|
let r1 = {
|
|
let out1 = out.clone();
|
|
thread::Builder::new()
|
|
.name("bolt-sess-out".into())
|
|
.spawn(move || session_read_loop(stdout_pipe, out1))
|
|
.unwrap()
|
|
};
|
|
let r2 = {
|
|
let out2 = out.clone();
|
|
thread::Builder::new()
|
|
.name("bolt-sess-err".into())
|
|
.spawn(move || session_read_loop(stderr_pipe, out2))
|
|
.unwrap()
|
|
};
|
|
// 收割线程:先等两个 reader 结束(EOF 即输出全部读入),再置退出标志,
|
|
// 保证 shell_read 看到 exited 时所有输出都已进入缓冲区(消除丢失输出的竞态)
|
|
let out_r = out.clone();
|
|
let _ = {
|
|
thread::Builder::new()
|
|
.name("bolt-sess-reaper".into())
|
|
.spawn(move || {
|
|
let st = child.wait();
|
|
let _ = r1.join();
|
|
let _ = r2.join();
|
|
let code = st.ok().and_then(|s| s.code());
|
|
*out_r.exited.lock().unwrap() = code;
|
|
out_r.cv.notify_all();
|
|
})
|
|
.unwrap()
|
|
};
|
|
let id = app.next_session.fetch_add(1, Ordering::Relaxed) + 1;
|
|
app.sessions.lock().unwrap().insert(
|
|
id,
|
|
Arc::new(Session {
|
|
pid,
|
|
stdin: Mutex::new(Some(stdin)),
|
|
out,
|
|
}),
|
|
);
|
|
Ok(json!({"session_id": id, "shell": shell, "pid": pid}))
|
|
}
|
|
|
|
fn shell_write(session: &Session, input: &str, newline: bool) -> Result<Value, String> {
|
|
if session.out.exited.lock().unwrap().is_some() {
|
|
return Err("会话已退出".into());
|
|
}
|
|
let mut data = input.to_string();
|
|
if newline && !data.is_empty() && !data.ends_with('\n') {
|
|
data.push('\n');
|
|
}
|
|
let mut stdin_guard = session.stdin.lock().unwrap();
|
|
let mut stdin = stdin_guard.take().ok_or("会话标准输入已关闭")?;
|
|
let (tx, rx) = mpsc::channel();
|
|
let n = data.len();
|
|
thread::spawn(move || {
|
|
let r = stdin.write_all(data.as_bytes());
|
|
let _ = tx.send((r, stdin)); // 把 stdin 一起传回,便于后续复用
|
|
});
|
|
match rx.recv_timeout(SESSION_WRITE_TIMEOUT) {
|
|
Ok((Ok(()), stdin)) => {
|
|
*stdin_guard = Some(stdin);
|
|
Ok(json!({"ok": true, "bytes_written": n}))
|
|
}
|
|
Ok((Err(e), _)) => Err(format!("写入失败(会话可能已退出): {e}")),
|
|
Err(_) => Err("写入超时(进程未读取输入),可用 shell_close 关闭会话".into()),
|
|
}
|
|
}
|
|
|
|
fn shell_read(session: &Session, timeout_ms: u64, max_bytes: usize) -> Value {
|
|
let out = &session.out;
|
|
let deadline = (timeout_ms != 0).then(|| Instant::now() + Duration::from_millis(timeout_ms));
|
|
loop {
|
|
// 1) 有新输出则返回
|
|
{
|
|
let b = out.bytes.lock().unwrap();
|
|
let mut off = out.off.lock().unwrap();
|
|
if *off > b.len() {
|
|
*off = 0; // 缓冲被裁剪,重置偏移
|
|
}
|
|
if *off < b.len() {
|
|
let take = (b.len() - *off).min(max_bytes);
|
|
let s = String::from_utf8_lossy(&b[*off..*off + take]).into_owned();
|
|
*off += take;
|
|
let truncated = *off < b.len();
|
|
drop(off);
|
|
let exited = *out.exited.lock().unwrap();
|
|
let dropped = out.dropped.load(Ordering::Relaxed);
|
|
let prefix = if dropped > 0 {
|
|
format!("[…已丢弃 {dropped} 字节…]\n")
|
|
} else {
|
|
String::new()
|
|
};
|
|
return json!({
|
|
"output": prefix + &s,
|
|
"exited": exited.is_some(),
|
|
"exit_code": exited,
|
|
"truncated": truncated,
|
|
});
|
|
}
|
|
}
|
|
// 2) 无新输出但已退出
|
|
if let Some(code) = *out.exited.lock().unwrap() {
|
|
return json!({"output": "", "exited": true, "exit_code": code, "truncated": false});
|
|
}
|
|
// 3) 等待新输出或退出
|
|
let wait = match deadline {
|
|
Some(d) => d.saturating_duration_since(Instant::now()),
|
|
None => Duration::from_secs(365 * 24 * 3600),
|
|
};
|
|
if wait.is_zero() {
|
|
return json!({"output": "", "exited": false, "exit_code": null, "truncated": false});
|
|
}
|
|
let gen = out.gen.load(Ordering::Relaxed);
|
|
let (guard, tres) = out.cv.wait_timeout(out.bytes.lock().unwrap(), wait).unwrap();
|
|
drop(guard);
|
|
if out.gen.load(Ordering::Relaxed) != gen || out.exited.lock().unwrap().is_some() {
|
|
continue; // 有新事件,重新检查
|
|
}
|
|
if tres.timed_out() {
|
|
return json!({"output": "", "exited": false, "exit_code": null, "truncated": false});
|
|
}
|
|
}
|
|
}
|
|
|
|
fn shell_close(app: &Shared, id: u64) -> Result<Value, String> {
|
|
let s = app
|
|
.sessions
|
|
.lock()
|
|
.unwrap()
|
|
.remove(&id)
|
|
.ok_or_else(|| format!("会话不存在: {id}"))?;
|
|
kill_tree(s.pid);
|
|
drop(s.stdin.lock().unwrap().take()); // 关闭 stdin,子进程看到 EOF
|
|
Ok(json!({"closed": true, "session_id": id}))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 工具列表(静态)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
static TOOLS: OnceLock<Value> = OnceLock::new();
|
|
|
|
fn tools_list() -> &'static Value {
|
|
TOOLS.get_or_init(|| {
|
|
json!({
|
|
"tools": [
|
|
{
|
|
"name": "run_command",
|
|
"description": "在服务器上执行 shell 命令并捕获输出。shell=true(默认)时支持完整 shell 语法:管道、重定向、&&、环境变量等;shell=false 时配合 args 直接执行二进制(不经 shell)。输出超限自动截断,超时自动杀进程,客户端可通过 notifications/cancelled 取消。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {"type": "string", "description": "要执行的命令行(shell=true 时为 shell 语法)"},
|
|
"args": {"type": "array", "items": {"type": "string"}, "description": "shell=false 时的直接执行参数"},
|
|
"shell": {"type": "boolean", "description": "是否经系统 shell 执行(sh -c / cmd /C)。默认 true"},
|
|
"cwd": {"type": "string", "description": "工作目录,默认服务器配置目录"},
|
|
"env": {"type": "object", "additionalProperties": {"type": "string"}, "description": "附加环境变量"},
|
|
"timeout_ms": {"type": "integer", "description": "超时毫秒,0=无限,默认 30000,上限 21600000"},
|
|
"max_output_bytes": {"type": "integer", "description": "stdout/stderr 各自捕获上限,默认 1048576"}
|
|
},
|
|
"required": ["command"]
|
|
}
|
|
},
|
|
{
|
|
"name": "shell_open",
|
|
"description": "启动持久化交互式 shell 会话(cwd、环境变量等状态跨次保留)。与 shell_write / shell_read 配合使用。无 PTY,交互式程序可能表现为非交互模式。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"shell": {"type": "string", "description": "shell 程序,如 sh/bash/zsh/fish,默认 sh(windows 为 cmd)"},
|
|
"cwd": {"type": "string", "description": "会话工作目录"}
|
|
}
|
|
}
|
|
},
|
|
{
|
|
"name": "shell_write",
|
|
"description": "向会话写入一行输入(自动补换行)。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"session_id": {"type": "integer", "description": "shell_open 返回的会话 ID"},
|
|
"input": {"type": "string", "description": "要写入的内容"},
|
|
"newline": {"type": "boolean", "description": "是否自动补换行,默认 true"}
|
|
},
|
|
"required": ["session_id", "input"]
|
|
}
|
|
},
|
|
{
|
|
"name": "shell_read",
|
|
"description": "读取会话自上次读取以来的新输出;若无输出则阻塞至 timeout_ms(0 = 等到会话退出)。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"session_id": {"type": "integer"},
|
|
"timeout_ms": {"type": "integer", "description": "等待毫秒,默认 5000,0=等到退出"},
|
|
"max_bytes": {"type": "integer", "description": "本次最多返回字节数,默认 1048576"}
|
|
},
|
|
"required": ["session_id"]
|
|
}
|
|
},
|
|
{
|
|
"name": "shell_close",
|
|
"description": "关闭并杀掉会话进程树。",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"session_id": {"type": "integer"}
|
|
},
|
|
"required": ["session_id"]
|
|
}
|
|
}
|
|
]
|
|
})
|
|
})
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MCP 分发
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn initialize_result(params: &Value) -> Value {
|
|
let client_version = params.get("protocolVersion").and_then(Value::as_str).unwrap_or("");
|
|
let version = if PROTOCOL_VERSIONS.contains(&client_version) {
|
|
client_version.to_string()
|
|
} else {
|
|
LATEST_PROTOCOL.to_string()
|
|
};
|
|
json!({
|
|
"protocolVersion": version,
|
|
"capabilities": {"tools": {"listChanged": false}},
|
|
"serverInfo": {"name": "bolt", "version": VERSION},
|
|
"instructions": "你是连接到 bolt 命令执行服务器的 AI。你可以用 run_command 在服务器上执行 shell 命令(支持管道、重定向、超时与取消);需要保留状态(如 cd、环境变量)时,用 shell_open 开启持久会话,配合 shell_write/shell_read 交互。注意命令受服务器策略限制,输出可能被截断。"
|
|
})
|
|
}
|
|
|
|
/// 执行工具调用,返回工具结果(Ok)或协议错误(Err((code, msg)))
|
|
/// killable:外部可取消句柄(注册在 inflight 表中,供 notifications/cancelled 使用);
|
|
/// 为 None 时(测试或内部调用)自动创建本地句柄。
|
|
fn exec_tool_call(app: &Shared, name: &str, args: &Value, progress: Option<&Value>, killable: Option<&Killable>) -> Result<Value, (i64, String)> {
|
|
match name {
|
|
"run_command" => {
|
|
let req = parse_run_params(args, &app.cfg).map_err(|e| (-32602, e))?;
|
|
let token = req.command.split_whitespace().next().unwrap_or("");
|
|
if let Err(e) = policy_check(&app.cfg, token) {
|
|
return Ok(tool_error(&e));
|
|
}
|
|
let local_k;
|
|
let k = match killable {
|
|
Some(k) => k,
|
|
None => {
|
|
local_k = Killable::default();
|
|
&local_k
|
|
}
|
|
};
|
|
match exec(&req, k, progress.cloned(), app) {
|
|
Ok(o) => Ok(run_result(&o, req.cwd.as_deref())),
|
|
Err(e) => Ok(tool_error(&e)),
|
|
}
|
|
}
|
|
"shell_open" => {
|
|
let shell = args.get("shell").and_then(Value::as_str);
|
|
let cwd = args.get("cwd").and_then(Value::as_str);
|
|
match shell_open(app, shell, cwd) {
|
|
Ok(v) => Ok(v),
|
|
Err(e) => Ok(tool_error(&e)),
|
|
}
|
|
}
|
|
"shell_write" => {
|
|
let id = args
|
|
.get("session_id")
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?;
|
|
let input = args
|
|
.get("input")
|
|
.and_then(Value::as_str)
|
|
.ok_or_else(|| (-32602, "'input' 必须是字符串".to_string()))?;
|
|
let newline = args.get("newline").and_then(Value::as_bool).unwrap_or(true);
|
|
let session = match app.sessions.lock().unwrap().get(&id) {
|
|
Some(s) => s.clone(),
|
|
None => return Ok(tool_error(&format!("会话不存在: {id}"))),
|
|
};
|
|
match shell_write(&session, input, newline) {
|
|
Ok(v) => Ok(v),
|
|
Err(e) => Ok(tool_error(&e)),
|
|
}
|
|
}
|
|
"shell_read" => {
|
|
let id = args
|
|
.get("session_id")
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?;
|
|
let timeout_ms = args.get("timeout_ms").and_then(Value::as_u64).unwrap_or(5000);
|
|
let max_bytes = args
|
|
.get("max_bytes")
|
|
.and_then(Value::as_u64)
|
|
.map(|n| n as usize)
|
|
.unwrap_or(DEFAULT_MAX_OUTPUT);
|
|
let session = match app.sessions.lock().unwrap().get(&id) {
|
|
Some(s) => s.clone(),
|
|
None => return Ok(tool_error(&format!("会话不存在: {id}"))),
|
|
};
|
|
Ok(shell_read(&session, timeout_ms, max_bytes))
|
|
}
|
|
"shell_close" => {
|
|
let id = args
|
|
.get("session_id")
|
|
.and_then(Value::as_u64)
|
|
.ok_or_else(|| (-32602, "'session_id' 必须是整数".to_string()))?;
|
|
match shell_close(app, id) {
|
|
Ok(v) => Ok(v),
|
|
Err(e) => Ok(tool_error(&e)),
|
|
}
|
|
}
|
|
_ => Err((-32602, format!("未知工具: {name}"))),
|
|
}
|
|
}
|
|
|
|
// 处理 messages,返回需要发送的 JSON-RPC 响应(通知/异步工具调用返回 None)
|
|
fn process_message(v: &Value, app: &Shared, async_exec: bool) -> Option<Value> {
|
|
let obj = match v.as_object() {
|
|
Some(o) => o,
|
|
None => return Some(rpc_error(&Value::Null, -32600, "invalid request")),
|
|
};
|
|
let method = match obj.get("method").and_then(Value::as_str) {
|
|
Some(m) => m,
|
|
None => return Some(rpc_error(obj.get("id").unwrap_or(&Value::Null), -32600, "invalid request")),
|
|
};
|
|
let params = obj.get("params").cloned().unwrap_or_else(|| json!({}));
|
|
let id = obj.get("id").cloned();
|
|
|
|
if id.is_none() {
|
|
// 通知
|
|
match method {
|
|
"notifications/initialized" => {}
|
|
"notifications/cancelled" => {
|
|
if let Some(rid) = params.get("requestId") {
|
|
let key = id_key(rid);
|
|
if let Some(k) = app.inflight.lock().unwrap().get(&key).cloned() {
|
|
k.kill(true);
|
|
}
|
|
}
|
|
}
|
|
"notifications/exit" => {
|
|
app.exit_flag.store(true, Ordering::SeqCst);
|
|
}
|
|
_ => {}
|
|
}
|
|
return None;
|
|
}
|
|
|
|
let id = id.unwrap();
|
|
let resp = match method {
|
|
"initialize" => json!({"jsonrpc": "2.0", "id": id, "result": initialize_result(¶ms)}),
|
|
"ping" => json!({"jsonrpc": "2.0", "id": id, "result": {}}),
|
|
"tools/list" => json!({"jsonrpc": "2.0", "id": id, "result": tools_list()}),
|
|
"tools/call" => {
|
|
let name = params.get("name").and_then(Value::as_str).unwrap_or("");
|
|
let args = params.get("arguments").cloned().unwrap_or_else(|| json!({}));
|
|
let progress = params.get("_meta").and_then(|m| m.get("progressToken")).cloned();
|
|
if async_exec {
|
|
// stdio 模式:worker 线程执行,完成后经 sink 发响应(不阻塞主循环)
|
|
let app2 = app.clone();
|
|
let id2 = id.clone();
|
|
let name2 = name.to_string();
|
|
let args2 = args;
|
|
let tok2 = progress;
|
|
let k = register_inflight(app, &id);
|
|
thread::Builder::new()
|
|
.name("bolt-tool".into())
|
|
.spawn(move || {
|
|
let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
exec_tool_call(&app2, &name2, &args2, tok2.as_ref(), Some(&k))
|
|
}));
|
|
unregister_inflight(&app2, &id2, &k);
|
|
let resp = match r {
|
|
Ok(Ok(v)) => json!({"jsonrpc": "2.0", "id": id2, "result": v}),
|
|
Ok(Err((code, msg))) => rpc_error(&id2, code, &msg),
|
|
Err(_) => rpc_error(&id2, -32603, "internal error"),
|
|
};
|
|
app2.sink.send(&resp.to_string());
|
|
})
|
|
.unwrap();
|
|
return None;
|
|
} else {
|
|
let k = register_inflight(app, &id);
|
|
let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
exec_tool_call(app, name, &args, progress.as_ref(), Some(&k))
|
|
}));
|
|
unregister_inflight(app, &id, &k);
|
|
match r {
|
|
Ok(Ok(v)) => json!({"jsonrpc": "2.0", "id": id, "result": v}),
|
|
Ok(Err((code, msg))) => rpc_error(&id, code, &msg),
|
|
Err(_) => rpc_error(&id, -32603, "internal error"),
|
|
}
|
|
}
|
|
}
|
|
"shutdown" => json!({"jsonrpc": "2.0", "id": id, "result": {}}),
|
|
_ => rpc_error(&id, -32601, &format!("method not found: {method}")),
|
|
};
|
|
Some(resp)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// stdio 传输
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// 读取一行(以 \n 结尾,不含换行),上限 cap 字节;返回 false 表示 EOF
|
|
fn read_line_capped<R: BufRead>(r: &mut R, cap: usize, buf: &mut Vec<u8>) -> io::Result<bool> {
|
|
loop {
|
|
let (newline_at, n) = {
|
|
let available = r.fill_buf()?;
|
|
if available.is_empty() {
|
|
return Ok(!buf.is_empty());
|
|
}
|
|
let room = cap.saturating_sub(buf.len());
|
|
if room == 0 {
|
|
return Ok(true); // 超限,由调用方报 parse error
|
|
}
|
|
let upto = available.len().min(room);
|
|
let nl = available[..upto].iter().position(|&b| b == b'\n');
|
|
match nl {
|
|
Some(i) => {
|
|
buf.extend_from_slice(&available[..i]);
|
|
(true, i + 1)
|
|
}
|
|
None => {
|
|
buf.extend_from_slice(&available[..upto]);
|
|
(false, upto)
|
|
}
|
|
}
|
|
};
|
|
r.consume(n);
|
|
if newline_at {
|
|
return Ok(true);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn run_stdio(app: Shared) -> io::Result<()> {
|
|
let stdin = io::stdin();
|
|
let mut reader = BufReader::with_capacity(64 * 1024, stdin);
|
|
let mut buf: Vec<u8> = Vec::with_capacity(512);
|
|
loop {
|
|
// 顶部检查退出标志:exit 通知/信号处理后,即使 stdin 仍阻塞读取也能退出
|
|
if app.exit_flag.load(Ordering::SeqCst) {
|
|
break;
|
|
}
|
|
buf.clear();
|
|
match read_line_capped(&mut reader, MAX_LINE_BYTES, &mut buf) {
|
|
Ok(true) => {}
|
|
Ok(false) => break, // EOF:客户端关闭 stdin
|
|
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
|
Err(e) => {
|
|
eprintln!("[bolt] stdin 错误: {e}");
|
|
break;
|
|
}
|
|
}
|
|
// 去掉尾部 \r\n
|
|
let line = match buf.iter().rposition(|&b| b == b'\n') {
|
|
Some(i) => &buf[..i],
|
|
None => &buf[..],
|
|
};
|
|
let line = match line.iter().rposition(|&b| b == b'\r') {
|
|
Some(i) if i + 1 == line.len() => &line[..i],
|
|
_ => line,
|
|
};
|
|
if line.is_empty() {
|
|
continue;
|
|
}
|
|
let v = match serde_json::from_slice::<Value>(line) {
|
|
Ok(v) => v,
|
|
Err(_) => {
|
|
app.sink.send(PARSE_ERROR);
|
|
continue;
|
|
}
|
|
};
|
|
if let Some(resp) = process_message(&v, &app, true) {
|
|
app.sink.send(&resp.to_string());
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Streamable HTTP 传输
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct ReqHead {
|
|
method: String,
|
|
path: String,
|
|
conn_close: bool,
|
|
headers: HashMap<String, String>,
|
|
}
|
|
|
|
fn read_http_head<R: BufRead>(r: &mut R) -> io::Result<Option<Vec<u8>>> {
|
|
let mut buf = Vec::with_capacity(512);
|
|
loop {
|
|
if buf.len() > MAX_HTTP_HEAD {
|
|
return Err(io::Error::new(io::ErrorKind::InvalidData, "header too large"));
|
|
}
|
|
let n = r.read_until(b'\n', &mut buf)?;
|
|
if n == 0 {
|
|
return Ok(if buf.is_empty() { None } else { Some(buf) });
|
|
}
|
|
// 空行结束头部(\r\n\r\n 或 \n\n)
|
|
let len = buf.len();
|
|
if len >= 2 && buf[len - 1] == b'\n' && buf[len - 2] == b'\n' {
|
|
return Ok(Some(buf));
|
|
}
|
|
if len >= 4 && buf[len - 1] == b'\n' && buf[len - 2] == b'\r' && buf[len - 3] == b'\n' && buf[len - 4] == b'\r' {
|
|
return Ok(Some(buf));
|
|
}
|
|
}
|
|
}
|
|
|
|
fn parse_http_head(buf: &[u8]) -> Result<ReqHead, ()> {
|
|
let s = std::str::from_utf8(buf).map_err(|_| ())?;
|
|
let mut lines = s.lines();
|
|
let req_line = lines.next().ok_or(())?;
|
|
let mut parts = req_line.split_whitespace();
|
|
let method = parts.next().ok_or(())?.to_string();
|
|
let raw_path = parts.next().ok_or(())?.to_string();
|
|
let path = raw_path.split('?').next().unwrap_or("").to_string();
|
|
let mut headers = HashMap::new();
|
|
let mut conn_close = false;
|
|
for line in lines {
|
|
let line = line.trim_end_matches('\r');
|
|
if line.is_empty() {
|
|
continue;
|
|
}
|
|
if let Some((k, v)) = line.split_once(':') {
|
|
let key = k.trim().to_ascii_lowercase();
|
|
let val = v.trim().to_string();
|
|
if key == "connection" && val.eq_ignore_ascii_case("close") {
|
|
conn_close = true;
|
|
}
|
|
headers.insert(key, val);
|
|
}
|
|
}
|
|
Ok(ReqHead {
|
|
method,
|
|
path,
|
|
conn_close,
|
|
headers,
|
|
})
|
|
}
|
|
|
|
fn origin_allowed(cfg: &Config, origin: &str) -> bool {
|
|
if cfg.allow_origins.iter().any(|o| origin == o || origin.starts_with(&format!("{o}:"))) {
|
|
return true;
|
|
}
|
|
let host = origin
|
|
.split("://")
|
|
.nth(1)
|
|
.unwrap_or(origin)
|
|
.split(':')
|
|
.next()
|
|
.unwrap_or("");
|
|
host == "localhost" || host == "127.0.0.1" || host == "[::1]"
|
|
}
|
|
|
|
fn respond_plain(w: &mut impl Write, status: &str, body: &[u8]) -> io::Result<()> {
|
|
write!(
|
|
w,
|
|
"HTTP/1.1 {status}\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: {}\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n",
|
|
body.len()
|
|
)?;
|
|
w.write_all(body)?;
|
|
w.flush()
|
|
}
|
|
|
|
fn respond_json(w: &mut impl Write, status: &str, body: &str) -> io::Result<()> {
|
|
write!(
|
|
w,
|
|
"HTTP/1.1 {status}\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: {}\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n",
|
|
body.len()
|
|
)?;
|
|
w.write_all(body.as_bytes())?;
|
|
w.flush()
|
|
}
|
|
|
|
fn respond_options(w: &mut impl Write) -> io::Result<()> {
|
|
write!(
|
|
w,
|
|
"HTTP/1.1 204 No Content\r\nAccess-Control-Allow-Origin: *\r\nAccess-Control-Allow-Methods: GET, POST, OPTIONS\r\nAccess-Control-Allow-Headers: Content-Type, Authorization, Accept\r\nAccess-Control-Max-Age: 86400\r\nConnection: keep-alive\r\n\r\n"
|
|
)?;
|
|
w.flush()
|
|
}
|
|
|
|
fn respond_401(w: &mut impl Write) -> io::Result<()> {
|
|
write!(
|
|
w,
|
|
"HTTP/1.1 401 Unauthorized\r\nWWW-Authenticate: Bearer\r\nContent-Length: 0\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n"
|
|
)?;
|
|
w.flush()
|
|
}
|
|
|
|
/// POST 的单事件 SSE 响应(规范:响应发送后服务器 SHOULD 关闭流)
|
|
fn write_sse_event(w: &mut impl Write, data: &str) -> io::Result<()> {
|
|
write!(
|
|
w,
|
|
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\nConnection: close\r\nAccess-Control-Allow-Origin: *\r\n\r\nevent: message\r\ndata: {data}\r\n\r\n"
|
|
)?;
|
|
w.flush()
|
|
}
|
|
|
|
/// GET:长连接 SSE 流,接收服务器通知,15s 心跳保活
|
|
fn sse_stream(w: &mut impl Write, app: &Shared) {
|
|
let _ = write!(
|
|
w,
|
|
"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\nConnection: keep-alive\r\nAccess-Control-Allow-Origin: *\r\n\r\n"
|
|
);
|
|
let _ = w.flush();
|
|
let (tx, rx) = mpsc::channel::<String>();
|
|
let id = app.hub.add(tx);
|
|
loop {
|
|
match rx.recv_timeout(SSE_HEARTBEAT) {
|
|
Ok(msg) => {
|
|
let mut ev = String::with_capacity(msg.len() + 32);
|
|
ev.push_str("event: message\r\ndata: ");
|
|
ev.push_str(&msg);
|
|
ev.push_str("\r\n\r\n");
|
|
if w.write_all(ev.as_bytes()).and_then(|_| w.flush()).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Timeout) => {
|
|
if w.write_all(b": ping\r\n\r\n").and_then(|_| w.flush()).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
Err(mpsc::RecvTimeoutError::Disconnected) => break,
|
|
}
|
|
}
|
|
app.hub.remove(id);
|
|
}
|
|
|
|
fn conn_loop(stream: TcpStream, app: Shared) {
|
|
let _ = stream.set_read_timeout(Some(Duration::from_secs(120)));
|
|
let _ = stream.set_write_timeout(Some(Duration::from_secs(60)));
|
|
let reader_stream = match stream.try_clone() {
|
|
Ok(s) => s,
|
|
Err(_) => return,
|
|
};
|
|
let mut reader = BufReader::with_capacity(16 * 1024, reader_stream);
|
|
let mut writer = BufWriter::with_capacity(16 * 1024, stream);
|
|
|
|
loop {
|
|
let head = match read_http_head(&mut reader) {
|
|
Ok(Some(h)) => h,
|
|
Ok(None) => return, // 客户端关闭
|
|
Err(_) => {
|
|
let _ = respond_plain(&mut writer, "431 Request Header Fields Too Large", b"header too large");
|
|
return;
|
|
}
|
|
};
|
|
let req = match parse_http_head(&head) {
|
|
Ok(r) => r,
|
|
Err(_) => {
|
|
let _ = respond_plain(&mut writer, "400 Bad Request", b"bad request");
|
|
return;
|
|
}
|
|
};
|
|
if req.method == "OPTIONS" {
|
|
let _ = respond_options(&mut writer);
|
|
continue;
|
|
}
|
|
// 防 DNS rebinding:校验 Origin(规范 MUST)
|
|
if let Some(origin) = req.headers.get("origin") {
|
|
if !origin_allowed(&app.cfg, origin) {
|
|
let _ = respond_plain(&mut writer, "403 Forbidden", b"origin not allowed");
|
|
continue;
|
|
}
|
|
}
|
|
if let Some(tok) = &app.cfg.token {
|
|
let ok = req
|
|
.headers
|
|
.get("authorization")
|
|
.map(|a| a == &format!("Bearer {tok}"))
|
|
.unwrap_or(false);
|
|
if !ok {
|
|
let _ = respond_401(&mut writer);
|
|
continue;
|
|
}
|
|
}
|
|
if req.path != "/mcp" {
|
|
let _ = respond_plain(&mut writer, "404 Not Found", b"not found");
|
|
return;
|
|
}
|
|
match req.method.as_str() {
|
|
"GET" => {
|
|
sse_stream(&mut writer, &app);
|
|
return;
|
|
}
|
|
"POST" => {
|
|
let len: usize = req
|
|
.headers
|
|
.get("content-length")
|
|
.and_then(|v| v.parse().ok())
|
|
.unwrap_or(0);
|
|
if len > MAX_HTTP_BODY {
|
|
let _ = respond_plain(&mut writer, "413 Payload Too Large", b"body too large");
|
|
return;
|
|
}
|
|
let mut body = vec![0u8; len];
|
|
if reader.read_exact(&mut body).is_err() {
|
|
return;
|
|
}
|
|
let v: Value = match serde_json::from_slice(&body) {
|
|
Ok(v) => v,
|
|
Err(_) => {
|
|
let e = rpc_error(&Value::Null, -32700, "parse error");
|
|
let _ = respond_json(&mut writer, "400 Bad Request", &e.to_string());
|
|
return;
|
|
}
|
|
};
|
|
if !v.is_object() {
|
|
let e = rpc_error(&Value::Null, -32600, "invalid request");
|
|
let _ = respond_json(&mut writer, "400 Bad Request", &e.to_string());
|
|
return;
|
|
}
|
|
let wants_sse = req
|
|
.headers
|
|
.get("accept")
|
|
.map(|a| a.contains("text/event-stream"))
|
|
.unwrap_or(false);
|
|
match process_message(&v, &app, false) {
|
|
Some(resp) => {
|
|
let s = resp.to_string();
|
|
if wants_sse {
|
|
let _ = write_sse_event(&mut writer, &s);
|
|
return; // 规范:响应后关闭 SSE 流
|
|
}
|
|
let _ = respond_json(&mut writer, "200 OK", &s);
|
|
}
|
|
None => {
|
|
// 通知 → 202 Accepted,无 body
|
|
let _ = respond_plain(&mut writer, "202 Accepted", b"");
|
|
}
|
|
}
|
|
if req.conn_close {
|
|
return;
|
|
}
|
|
}
|
|
_ => {
|
|
let _ = respond_plain(&mut writer, "405 Method Not Allowed", b"method not allowed");
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(unix)]
|
|
fn run_http(addr: &str, app: Shared) -> io::Result<()> {
|
|
use std::os::unix::io::AsRawFd;
|
|
let listener = TcpListener::bind(addr)?;
|
|
eprintln!("[bolt] {VERSION} 监听 http://{addr}/mcp (Streamable HTTP)");
|
|
log_policy(&app);
|
|
let fd = listener.as_raw_fd();
|
|
let mut pollfd = libc::pollfd {
|
|
fd,
|
|
events: libc::POLLIN,
|
|
revents: 0,
|
|
};
|
|
loop {
|
|
if app.exit_flag.load(Ordering::SeqCst) {
|
|
break;
|
|
}
|
|
let r = unsafe { libc::poll(&mut pollfd, 1, 200) };
|
|
if r > 0 {
|
|
match listener.accept() {
|
|
Ok((s, _)) => {
|
|
let _ = s.set_nodelay(true);
|
|
let app2 = app.clone();
|
|
thread::Builder::new()
|
|
.name("bolt-http".into())
|
|
.spawn(move || conn_loop(s, app2))
|
|
.unwrap();
|
|
}
|
|
Err(e) => eprintln!("[bolt] accept 错误: {e}"),
|
|
}
|
|
} else if r < 0 {
|
|
let err = io::Error::last_os_error();
|
|
if err.kind() != io::ErrorKind::Interrupted {
|
|
eprintln!("[bolt] poll 错误: {err}");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(windows)]
|
|
fn run_http(addr: &str, app: Shared) -> io::Result<()> {
|
|
let listener = TcpListener::bind(addr)?;
|
|
eprintln!("[bolt] {VERSION} 监听 http://{addr}/mcp (Streamable HTTP)");
|
|
log_policy(&app);
|
|
loop {
|
|
match listener.accept() {
|
|
Ok((s, _)) => {
|
|
let _ = s.set_nodelay(true);
|
|
let app2 = app.clone();
|
|
thread::Builder::new()
|
|
.name("bolt-http".into())
|
|
.spawn(move || conn_loop(s, app2))
|
|
.unwrap();
|
|
}
|
|
Err(e) => eprintln!("[bolt] accept 错误: {e}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn log_policy(app: &Shared) {
|
|
if app.cfg.token.is_some() {
|
|
eprintln!("[bolt] 认证: Bearer Token 已启用");
|
|
}
|
|
if !app.cfg.allow.is_empty() {
|
|
eprintln!("[bolt] 白名单: {}", app.cfg.allow.join(", "));
|
|
}
|
|
if !app.cfg.deny.is_empty() {
|
|
eprintln!("[bolt] 黑名单: {}", app.cfg.deny.join(", "));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 信号与清理
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(unix)]
|
|
fn install_signal_flag(flag: Arc<AtomicBool>) {
|
|
unsafe {
|
|
let mut set: libc::sigset_t = std::mem::zeroed();
|
|
libc::sigemptyset(&mut set);
|
|
libc::sigaddset(&mut set, libc::SIGINT);
|
|
libc::sigaddset(&mut set, libc::SIGTERM);
|
|
libc::pthread_sigmask(libc::SIG_BLOCK, &set, std::ptr::null_mut());
|
|
thread::spawn(move || {
|
|
let mut sig: libc::c_int = 0;
|
|
libc::sigwait(&set, &mut sig);
|
|
flag.store(true, Ordering::SeqCst);
|
|
eprintln!("[bolt] 收到信号 {sig},正在清理并退出…");
|
|
});
|
|
}
|
|
}
|
|
|
|
#[cfg(not(unix))]
|
|
fn install_signal_flag(_flag: Arc<AtomicBool>) {}
|
|
|
|
/// 退出前杀掉所有进行中的命令与会话,避免孤儿进程
|
|
fn shutdown_cleanup(app: &Shared) {
|
|
let inflight: Vec<Arc<Killable>> = app.inflight.lock().unwrap().values().cloned().collect();
|
|
for k in inflight {
|
|
k.kill(true);
|
|
}
|
|
let sessions: Vec<Arc<Session>> = app.sessions.lock().unwrap().drain().map(|(_, s)| s).collect();
|
|
for s in sessions {
|
|
kill_tree(s.pid);
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 入口
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn main() {
|
|
let Some(cfg) = parse_args() else {
|
|
print_help();
|
|
return;
|
|
};
|
|
let app = Arc::new(App::new(cfg));
|
|
install_signal_flag(app.exit_flag.clone());
|
|
let res = match &app.cfg.mode {
|
|
Mode::Stdio => run_stdio(app.clone()),
|
|
Mode::Http(addr) => run_http(addr, app.clone()),
|
|
};
|
|
if let Err(e) = res {
|
|
eprintln!("[bolt] 致命错误: {e}");
|
|
}
|
|
shutdown_cleanup(&app);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 单元测试
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_parse_size() {
|
|
assert_eq!(parse_size("1024"), Some(1024));
|
|
assert_eq!(parse_size("64K"), Some(64 * 1024));
|
|
assert_eq!(parse_size("1M"), Some(1024 * 1024));
|
|
assert_eq!(parse_size("2G"), Some(2 * 1024 * 1024 * 1024));
|
|
assert_eq!(parse_size("x"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_policy() {
|
|
let cfg = Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec!["git".into(), "docker".into()],
|
|
deny: vec!["rm".into()],
|
|
timeout_ms: 1000,
|
|
cwd: None,
|
|
max_output: 1024,
|
|
allow_origins: vec![],
|
|
};
|
|
assert!(policy_check(&cfg, "git").is_ok());
|
|
assert!(policy_check(&cfg, "git status").is_ok());
|
|
assert!(policy_check(&cfg, "docker ps").is_ok());
|
|
assert!(policy_check(&cfg, "curl").is_err());
|
|
let cfg2 = Config {
|
|
allow: vec![],
|
|
deny: vec!["rm".into()],
|
|
..cfg
|
|
};
|
|
assert!(policy_check(&cfg2, "rm -rf /").is_err());
|
|
assert!(policy_check(&cfg2, "ls").is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_read_line_capped() {
|
|
let data = b"hello\nworld\r\n";
|
|
let mut r = BufReader::new(&data[..]);
|
|
let mut buf = Vec::new();
|
|
assert!(read_line_capped(&mut r, 1024, &mut buf).unwrap());
|
|
assert_eq!(buf, b"hello");
|
|
buf.clear();
|
|
assert!(read_line_capped(&mut r, 1024, &mut buf).unwrap());
|
|
assert_eq!(buf, b"world\r");
|
|
buf.clear();
|
|
assert!(!read_line_capped(&mut r, 1024, &mut buf).unwrap());
|
|
// 超限:返回 true 但截断
|
|
let mut r2 = BufReader::new(&b"abcdef"[..]);
|
|
let mut buf2 = Vec::new();
|
|
assert!(read_line_capped(&mut r2, 3, &mut buf2).unwrap());
|
|
assert_eq!(buf2.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_http_head_parse() {
|
|
let head = b"POST /mcp HTTP/1.1\r\nHost: x\r\nContent-Length: 10\r\nConnection: close\r\nAccept: application/json, text/event-stream\r\n\r\n";
|
|
let req = parse_http_head(head).unwrap();
|
|
assert_eq!(req.method, "POST");
|
|
assert_eq!(req.path, "/mcp");
|
|
assert!(req.conn_close);
|
|
assert_eq!(req.headers.get("content-length").unwrap(), "10");
|
|
assert!(req.headers.get("accept").unwrap().contains("text/event-stream"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_origin() {
|
|
let cfg = Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec![],
|
|
deny: vec![],
|
|
timeout_ms: 1000,
|
|
cwd: None,
|
|
max_output: 1024,
|
|
allow_origins: vec!["https://example.com".into()],
|
|
};
|
|
assert!(origin_allowed(&cfg, "http://localhost:3000"));
|
|
assert!(origin_allowed(&cfg, "https://127.0.0.1"));
|
|
assert!(origin_allowed(&cfg, "https://example.com"));
|
|
assert!(!origin_allowed(&cfg, "https://evil.com"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_dispatch_basic() {
|
|
let app = Arc::new(App::new(Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec![],
|
|
deny: vec![],
|
|
timeout_ms: 1000,
|
|
cwd: None,
|
|
max_output: 1024,
|
|
allow_origins: vec![],
|
|
}));
|
|
// ping
|
|
let v: Value = serde_json::from_str(r#"{"jsonrpc":"2.0","id":1,"method":"ping"}"#).unwrap();
|
|
let r = process_message(&v, &app, false).unwrap();
|
|
assert_eq!(r["result"], json!({}));
|
|
assert_eq!(r["id"], json!(1));
|
|
// initialize
|
|
let v: Value = serde_json::from_str(
|
|
r#"{"jsonrpc":"2.0","id":2,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}"#,
|
|
)
|
|
.unwrap();
|
|
let r = process_message(&v, &app, false).unwrap();
|
|
assert_eq!(r["result"]["protocolVersion"], "2024-11-05");
|
|
assert_eq!(r["result"]["serverInfo"]["name"], "bolt");
|
|
// 未知方法
|
|
let v: Value = serde_json::from_str(r#"{"jsonrpc":"2.0","id":3,"method":"nope"}"#).unwrap();
|
|
let r = process_message(&v, &app, false).unwrap();
|
|
assert_eq!(r["error"]["code"], -32601);
|
|
// 未知工具
|
|
let v: Value = serde_json::from_str(
|
|
r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"nope","arguments":{}}}"#,
|
|
)
|
|
.unwrap();
|
|
let r = process_message(&v, &app, false).unwrap();
|
|
assert_eq!(r["error"]["code"], -32602);
|
|
// 参数错误
|
|
let v: Value = serde_json::from_str(
|
|
r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"run_command","arguments":{}}}"#,
|
|
)
|
|
.unwrap();
|
|
let r = process_message(&v, &app, false).unwrap();
|
|
assert_eq!(r["error"]["code"], -32602);
|
|
// 解析失败走 PARSE_ERROR
|
|
assert!(PARSE_ERROR.contains("-32700"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_command_basic() {
|
|
let app = Arc::new(App::new(Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec![],
|
|
deny: vec![],
|
|
timeout_ms: 5000,
|
|
cwd: None,
|
|
max_output: 1024 * 1024,
|
|
allow_origins: vec![],
|
|
}));
|
|
let args = json!({"command": "echo hello && printf err >&2 && exit 3"});
|
|
let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap();
|
|
assert_eq!(r["isError"], Value::Null);
|
|
let sc = &r["structuredContent"];
|
|
assert_eq!(sc["exit_code"], 3);
|
|
assert_eq!(sc["stdout"], "hello\n");
|
|
assert_eq!(sc["stderr"], "err");
|
|
assert!(!sc["timed_out"].as_bool().unwrap());
|
|
// 超时
|
|
let args = json!({"command": "sleep 10", "timeout_ms": 300});
|
|
let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap();
|
|
assert_eq!(r["structuredContent"]["timed_out"], true);
|
|
// shell=false 直接执行
|
|
let args = json!({"command": "echo", "args": ["direct"], "shell": false});
|
|
let r = exec_tool_call(&app, "run_command", &args, None, None).unwrap();
|
|
assert_eq!(r["structuredContent"]["stdout"], "direct\n");
|
|
// 白名单拦截
|
|
let app2 = Arc::new(App::new(Config {
|
|
allow: vec!["echo".into()],
|
|
..Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec![],
|
|
deny: vec![],
|
|
timeout_ms: 5000,
|
|
cwd: None,
|
|
max_output: 1024,
|
|
allow_origins: vec![],
|
|
}
|
|
}));
|
|
let r = exec_tool_call(&app2, "run_command", &json!({"command": "ls"}), None, None).unwrap();
|
|
assert_eq!(r["isError"], true);
|
|
// 经 shell 执行不存在的命令:sh 启动成功,内部报错,exit 127
|
|
let r = exec_tool_call(&app, "run_command", &json!({"command": "definitely_not_a_cmd_xyz"}), None, None).unwrap();
|
|
assert_eq!(r["structuredContent"]["exit_code"], 127);
|
|
assert!(!r["structuredContent"]["stderr"].as_str().unwrap().is_empty());
|
|
// shell=false 直接执行不存在的程序:spawn 失败 → isError
|
|
let r = exec_tool_call(&app, "run_command", &json!({"command": "definitely_not_a_cmd_xyz", "shell": false}), None, None).unwrap();
|
|
assert_eq!(r["isError"], true);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shell_session() {
|
|
let app = Arc::new(App::new(Config {
|
|
mode: Mode::Stdio,
|
|
token: None,
|
|
allow: vec![],
|
|
deny: vec![],
|
|
timeout_ms: 5000,
|
|
cwd: None,
|
|
max_output: 1024,
|
|
allow_origins: vec![],
|
|
}));
|
|
let r = shell_open(&app, Some("sh"), None).unwrap();
|
|
let sid = r["session_id"].as_u64().unwrap();
|
|
assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo FIRST", true).is_ok());
|
|
let r0 = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 2000, 1 << 20);
|
|
assert!(r0["output"].as_str().unwrap().contains("FIRST"), "output: {}", r0);
|
|
// 状态保留:cd 后 pwd(非交互 sh 无回显,分两步写、一次读)
|
|
assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "cd /tmp", true).is_ok());
|
|
assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo PROBE && pwd", true).is_ok());
|
|
// 输出分块到达,循环读取直到拿到全部(工具语义:返回当前可用的新输出)
|
|
let mut probe_out = String::new();
|
|
for _ in 0..10 {
|
|
let r = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 500, 1 << 20);
|
|
probe_out.push_str(r["output"].as_str().unwrap());
|
|
if probe_out.contains("PROBE") && probe_out.contains("/tmp") {
|
|
break;
|
|
}
|
|
}
|
|
assert!(probe_out.contains("PROBE"), "output: {probe_out}");
|
|
assert!(probe_out.contains("/tmp"), "output: {probe_out}");
|
|
assert!(shell_write(app.sessions.lock().unwrap().get(&sid).unwrap(), "echo STATE_OK && exit", true).is_ok());
|
|
let mut out = String::new();
|
|
for _ in 0..20 {
|
|
let r = shell_read(app.sessions.lock().unwrap().get(&sid).unwrap(), 1000, 1 << 20);
|
|
out.push_str(r["output"].as_str().unwrap());
|
|
if r["exited"].as_bool().unwrap() {
|
|
break;
|
|
}
|
|
}
|
|
assert!(out.contains("STATE_OK"), "out: {out}");
|
|
let r = shell_close(&app, sid).unwrap();
|
|
assert_eq!(r["closed"], true);
|
|
}
|
|
}
|