Refactor tool execution pipeline

2026-06-21 14:30:40 +08:00
parent 3e5f0658b3
commit 99c1e9cd6c
81 changed files with 3587 additions and 5317 deletions
@@ -0,0 +1,297 @@
+//! # 工具执行器
+//!
+//! `execute_loop()` spawn 工具进程，解析 stdout JSON 消息：
+//! - `{"type":"http",...}` — 请求 HTTP
+//! - `{"type":"db",...}`    — 请求数据库操作
+//! - `{"type":"result","content":"..."}` — 最终结果
+//!
+//! 所有消息必须是合法 JSON。审批在 http/db 操作前处理。
+
+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,
+}
+
+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
+     ),
+) -> String {
+    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 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 ──
+        let binary = std::path::Path::new(binary_path);
+        let mut cmd = Command::new(&binary);
+        if let Some(c) = &spec.command {
+            cmd.arg("--command").arg(c);
+        }
+        // 按工具 spec 定义的 env 列表，从当前进程环境推送对应变量
+        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.stdin(std::process::Stdio::piped())
+            .stdout(std::process::Stdio::piped())
+            .stderr(std::process::Stdio::piped())
+            .kill_on_drop(true);
+
+        let mut child = match cmd.spawn() {
+            Ok(c) => c,
+            Err(e) => {
+                error!(target: "ias::err","[spawn] {e}");
+                return 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 format!("读取stdout失败: {e}");
+            }
+            Err(_) => {
+                let _ = child.start_kill();
+                return 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 format!("进程异常: {e}");
+            }
+            Err(_) => {
+                let _ = child.start_kill();
+                error!(target: "ias::err","[timeout] {}s", spec.timeout_secs);
+                return 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 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 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 content.to_string();
+            }
+            Some("http") => {
+                let method = msg["method"].as_str().unwrap_or("GET");
+                let url = msg["url"].as_str().unwrap_or("");
+                let desc = msg["desc"].as_str().unwrap_or("");
+
+                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 format!("审批失败: {e}"),
+                    };
+                    wechat_send(
+                        user_id,
+                        &format!("⚠️ {}\n确认码: {}\n回复确认码继续，回复0取消", desc, code),
+                    )
+                    .await;
+                    let d = rx.await.unwrap_or(ApprovalDecision::Expired);
+                    match d {
+                        ApprovalDecision::Approved => { /* 继续执行 */ }
+                        ApprovalDecision::Rejected => return "操作已被用户取消".to_string(),
+                        ApprovalDecision::Expired => return "审批超时，操作已取消".to_string(),
+                    }
+                }
+
+                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 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 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 desc = msg["desc"].as_str().unwrap_or(op);
+
+                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 format!("审批失败: {e}"),
+                    };
+                    wechat_send(
+                        user_id,
+                        &format!("⚠️ {}\n确认码: {}\n回复确认码继续，回复0取消", desc, code),
+                    )
+                    .await;
+                    let d = rx.await.unwrap_or(ApprovalDecision::Expired);
+                    match d {
+                        ApprovalDecision::Approved => { /* 继续执行 */ }
+                        ApprovalDecision::Rejected => return "操作已被用户取消".to_string(),
+                        ApprovalDecision::Expired => return "审批超时，操作已取消".to_string(),
+                    }
+                }
+
+                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 {
+                            let write_result = memory.write_for(user_id, c).await;
+                            write_result
+                        }
+                    }
+                    "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 format!("未知消息类型: {}", msg["type"]);
+            }
+        }
+    }
+}