//! 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(&params)}),
        "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);
    }
}