iAs/src/tools/executor.rs

//! # 工具执行器
//!
//! `execute_loop()` spawn 工具进程，解析 stdout JSON 消息：
//! - `{"type":"http",...}` — 请求 HTTP
//! - `{"type":"db",...}`    — 请求数据库操作
//! - `{"type":"result","content":"..."}` — 最终结果
//!
//! 所有消息必须是合法 JSON。
//!
//! ## 安全边界（授权层）
//! **工具级审批在入口统一执行**：高风险工具在 spawn 之前必须获得用户确认，
//! 与工具后续输出的消息类型（result/http/db）无关 —— 杜绝 stdio→result 绕过。

use std::sync::Arc;
use std::time::Duration;

use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
use tokio::process::Command;

use crate::context::MemoryStore;
use crate::tools::approval::{ApprovalDecision, ApprovalManager};
use crate::tools::spec::ToolSpec;
use tracing::{error, info, warn};

struct HttpResponse {
    status: u16,
    body: Value,
}

// ─── 结构化结果 ───

/// 工具执行错误类别（供调用方区分，替代字符串前缀匹配）
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolErrorKind {
    SpawnFailed,
    StdoutReadFailed,
    StdoutTimeout,
    ProcessTimeout,
    ProcessAbnormal,
    NonJsonOutput,
    UnknownMessageType,
    UnsupportedMethod,
    HttpFailed,
    ApprovalFailed,
    ApprovalTimeout,
    ApprovalRejected,
    MaxRoundsExceeded,
    UnknownTool,
    MissingParams,
}

/// 工具执行返回值：成功带内容，失败带类别与消息
#[derive(Debug, Clone)]
pub enum ToolRunOutcome {
    Ok(String),
    Err {
        kind: ToolErrorKind,
        message: String,
    },
}

impl ToolRunOutcome {
    pub fn err(kind: ToolErrorKind, message: impl Into<String>) -> Self {
        ToolRunOutcome::Err {
            kind,
            message: message.into(),
        }
    }

    pub fn is_err(&self) -> bool {
        matches!(self, ToolRunOutcome::Err { .. })
    }

    /// 取出最终展示给 LLM 的消息文本
    pub fn into_message(self) -> String {
        match self {
            ToolRunOutcome::Ok(s) => s,
            ToolRunOutcome::Err { message, .. } => message,
        }
    }
}

pub async fn execute_loop(
    spec: &ToolSpec,
    params: &Value,
    user_id: &str,
    approved: bool,
    approval: &Arc<ApprovalManager>,
    memory: &Arc<MemoryStore>,
    wechat_send: &(
         dyn Fn(&str, &str) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>
             + Send
             + Sync
     ),
) -> ToolRunOutcome {
    // ── 统一授权边界：高风险工具在 spawn 前必须通过用户审批 ──
    // 不依赖后续消息类型，覆盖 stdio/http/db 所有路径
    if spec.is_high_risk() && !approved {
        let (code, rx) = match approval.create(user_id, &spec.name).await {
            Ok((c, r)) => (c, r),
            Err(e) => return ToolRunOutcome::err(ToolErrorKind::ApprovalFailed, format!("审批失败: {e}")),
        };
        wechat_send(
            user_id,
            &format!("⚠️ {}\n确认码: {}\n回复确认码继续，回复0取消", spec.desc, code),
        )
        .await;
        match rx.await.unwrap_or(ApprovalDecision::Expired) {
            ApprovalDecision::Approved => { /* 授权通过，继续执行 */ }
            ApprovalDecision::Rejected => {
                return ToolRunOutcome::err(ToolErrorKind::ApprovalRejected, "操作已被用户取消")
            }
            ApprovalDecision::Expired => {
                return ToolRunOutcome::err(ToolErrorKind::ApprovalTimeout, "审批超时，操作已取消")
            }
        }
    }

    let start = std::time::Instant::now();
    let binary_path = spec.path.as_deref().unwrap_or(&spec.tool);
    const MAX_ROUNDS: u32 = 20;

    let mut response: Option<HttpResponse> = None;
    let mut db_response: Option<Value> = None;
    let mut round = 0u32;

    loop {
        round += 1;
        if round > MAX_ROUNDS {
            error!(target: "ias::err", "[loop] 工具 '{}' 超过最大轮次 {}", spec.name, MAX_ROUNDS);
            return ToolRunOutcome::err(
                ToolErrorKind::MaxRoundsExceeded,
                format!("工具 '{}' 执行轮次过多（>{})", spec.name, MAX_ROUNDS),
            );
        }

        // ── 信封 ──
        let mut envelope = json!({"params":params,"approved":approved,"user_id":user_id});
        if let Some(ref r) = response {
            envelope["response"] = json!({"status":r.status,"body":r.body});
        }
        if let Some(ref d) = db_response {
            envelope["db_response"] = d.clone();
        }

        // ── spawn ──
        // 执行隔离边界：启用沙箱且工具 profile 非 local 时，通过 docker exec
        // 在长期容器内执行；否则本地 spawn（env_clear + 白名单注入）。
        // 沙箱降级（ensure 失败或容器崩溃未恢复）时自动回退本地执行。
        let use_docker = if crate::tools::sandbox::should_use_docker(&spec.sandbox.profile) {
            // 执行前确保容器就绪：未运行则重建，重建失败则降级
            crate::tools::sandbox::ensure_ready_or_degrade("tools").await
        } else {
            false
        };
        let mut child = match spawn_child(spec, binary_path, use_docker).await {
            Ok(c) => c,
            Err(e) => {
                error!(target: "ias::err","[spawn] {e}");
                return ToolRunOutcome::err(ToolErrorKind::SpawnFailed, format!("启动失败: {e}"));
            }
        };

        {
            let mut stdin = child.stdin.take().unwrap();
            let _ = stdin
                .write_all(
                    serde_json::to_string(&envelope)
                        .unwrap_or_default()
                        .as_bytes(),
                )
                .await;
        }

        let timeout = Duration::from_secs(spec.timeout_secs);

        // ── 并发读取 stdout + stderr（避免管道缓冲区死锁）──
        let stdout_reader = tokio::io::BufReader::new(child.stdout.take().unwrap());
        let stderr_reader = child.stderr.take().unwrap();

        let (line_result, stderr_result) = tokio::join!(
            tokio::time::timeout(timeout, async {
                let mut r = stdout_reader;
                let mut l = String::new();
                r.read_line(&mut l).await.map(|_| l)
            }),
            tokio::time::timeout(timeout, async {
                let mut s = stderr_reader;
                let mut b = vec![];
                s.read_to_end(&mut b)
                    .await
                    .map(|_| String::from_utf8_lossy(&b).to_string())
            }),
        );

        let line = match line_result {
            Ok(Ok(l)) => l,
            Ok(Err(e)) => {
                let _ = child.start_kill();
                return ToolRunOutcome::err(ToolErrorKind::StdoutReadFailed, format!("读取stdout失败: {e}"));
            }
            Err(_) => {
                let _ = child.start_kill();
                return ToolRunOutcome::err(
                    ToolErrorKind::StdoutTimeout,
                    format!("stdout超时({}s)", spec.timeout_secs),
                );
            }
        };

        let stderr_output = match stderr_result {
            Ok(Ok(s)) => s,
            Ok(Err(_)) => String::new(),
            Err(_) => String::new(),
        };
        let stderr_trimmed = stderr_output.trim();
        if !stderr_trimmed.is_empty() {
            warn!(target: "ias::err","[stderr] {}", stderr_trimmed);
        }

        // ── 等待进程退出 ──
        let status = match tokio::time::timeout(timeout, child.wait()).await {
            Ok(Ok(s)) => s,
            Ok(Err(e)) => {
                error!(target: "ias::err","[process] {e}");
                return ToolRunOutcome::err(ToolErrorKind::ProcessAbnormal, format!("进程异常: {e}"));
            }
            Err(_) => {
                let _ = child.start_kill();
                error!(target: "ias::err","[timeout] {}s", spec.timeout_secs);
                return ToolRunOutcome::err(
                    ToolErrorKind::ProcessTimeout,
                    format!("进程超时({}s)", spec.timeout_secs),
                );
            }
        };

        if !status.success() {
            error!(target: "ias::err","[exit] code={} stderr: {}", status.code().unwrap_or(-1), stderr_trimmed);
            let reason = if stderr_trimmed.is_empty() {
                format!("进程退出码 {}", status.code().unwrap_or(-1))
            } else {
                stderr_trimmed.to_string()
            };
            return ToolRunOutcome::err(
                ToolErrorKind::ProcessAbnormal,
                format!("异常退出({}): {}", status.code().unwrap_or(-1), reason),
            );
        }

        // ── 解析 ──
        let trimmed = line.trim();
        let msg: Value = match serde_json::from_str(trimmed) {
            Ok(v) => v,
            Err(_) => {
                warn!(target: "ias::err","[non-json] {:.100}", trimmed);
                return ToolRunOutcome::err(
                    ToolErrorKind::NonJsonOutput,
                    format!("工具输出非JSON: {:.100}", trimmed),
                );
            }
        };

        match msg["type"].as_str() {
            Some("result") => {
                let content = msg["content"].as_str().unwrap_or("");
                let elapsed = start.elapsed();
                info!(target: "ias::tool", "[done] {} 完成, 耗时 {:?}, {} 轮", spec.name, elapsed, round);
                return ToolRunOutcome::Ok(content.to_string());
            }
            Some("http") => {
                let method = msg["method"].as_str().unwrap_or("GET");
                let url = msg["url"].as_str().unwrap_or("");

                let client = reqwest::Client::builder()
                    .timeout(Duration::from_secs(25))
                    .build()
                    .unwrap();
                let body_str = msg["body"].as_str();
                let resp = match (method, body_str) {
                    ("GET", _) => client.get(url).send().await,
                    ("POST", Some(b)) => {
                        client
                            .post(url)
                            .header("Content-Type", "application/json")
                            .body(b.to_string())
                            .send()
                            .await
                    }
                    ("POST", None) => client.post(url).send().await,
                    (o, _) => {
                        return ToolRunOutcome::err(ToolErrorKind::UnsupportedMethod, format!("不支持的方法: {o}"))
                    }
                };
                match resp {
                    Ok(r) => {
                        let s = r.status().as_u16();
                        let b: Value = r.json().await.unwrap_or(Value::Null);
                        response = Some(HttpResponse { status: s, body: b });
                    }
                    Err(e) => {
                        error!(target: "ias::err","[http] {e}");
                        return ToolRunOutcome::err(ToolErrorKind::HttpFailed, format!("HTTP失败: {e}"));
                    }
                }
            }
            Some("db") => {
                let op = msg["operation"].as_str().unwrap_or("");
                let op_params = msg.get("params").cloned().unwrap_or(Value::Null);

                let result = match op {
                    "read_memories" => memory.read_for(user_id, None).await,
                    "write_memory" => {
                        let c = op_params["content"].as_str().unwrap_or("");
                        if c.is_empty() {
                            "缺少 content".into()
                        } else {
                            memory.write_for(user_id, c).await
                        }
                    }
                    "delete_memory" => {
                        let id = op_params["memory_id"].as_i64().unwrap_or(-1) as i32;
                        memory.delete_for(user_id, id).await
                    }
                    "update_memory" => {
                        let id = op_params["memory_id"].as_i64().unwrap_or(-1) as i32;
                        let c = op_params["content"].as_str().unwrap_or("");
                        memory.update_for(user_id, id, c).await
                    }
                    _ => format!("未知DB操作: {op}"),
                };
                db_response = Some(json!({"result": result}));
            }
            _ => {
                warn!(target: "ias::err","[unknown-type] {}", msg["type"]);
                return ToolRunOutcome::err(
                    ToolErrorKind::UnknownMessageType,
                    format!("未知消息类型: {}", msg["type"]),
                );
            }
        }
    }
}

/// 本地执行：env_clear + PATH + spec.env 白名单注入
fn build_local_cmd(spec: &ToolSpec, binary_path: &str) -> Command {
    let mut cmd = Command::new(std::path::Path::new(binary_path));
    if let Some(c) = &spec.command {
        cmd.arg("--command").arg(c);
    }
    cmd.env_clear();
    if let Ok(path) = std::env::var("PATH") {
        cmd.env("PATH", path);
    }
    for key in &spec.env {
        match std::env::var(key) {
            Ok(v) => {
                cmd.env(key, &v);
            }
            Err(_) => {
                warn!(target: "ias::err", "[env] 工具 '{}' 需要的环境变量 '{}' 未设置", spec.name, key);
            }
        }
    }
    cmd
}

/// spawn 工具进程。docker exec 模式下若首次 spawn 失败，尝试检测容器崩溃并重建后重试一次；
/// 重建失败则降级为本地执行（保留可用性，但仍受审批约束）。
async fn spawn_child(
    spec: &ToolSpec,
    binary_path: &str,
    use_docker: bool,
) -> Result<tokio::process::Child, std::io::Error> {
    let make_cmd = || {
        let mut cmd = if use_docker {
            crate::tools::sandbox::build_exec_cmd(spec)
        } else {
            build_local_cmd(spec, binary_path)
        };
        cmd.stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .kill_on_drop(true);
        cmd
    };

    if let Ok(child) = make_cmd().spawn() {
        return Ok(child);
    }

    // 首次 spawn 失败。docker 模式下可能是容器崩溃，尝试重建后重试一次
    if use_docker && crate::tools::sandbox::recover_if_crashed("tools").await {
        tracing::info!(target: "ias::sandbox", "沙箱容器已重建，重试执行");
        if let Ok(child) = make_cmd().spawn() {
            return Ok(child);
        }
        // 重建后仍失败，降级本地执行
        tracing::warn!(target: "ias::sandbox", "重建后 exec 仍失败，降级本地执行");
    } else if use_docker {
        tracing::warn!(target: "ias::sandbox", "容器不可用，降级本地执行");
    }

    // 降级：本地 spawn（审批已在入口完成，安全边界不受影响）
    build_local_cmd(spec, binary_path)
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .kill_on_drop(true)
        .spawn()
}