diff --git a/crates/ias-fetch/Cargo.toml b/crates/ias-fetch/Cargo.toml new file mode 100644 index 0000000..f2e3df8 --- /dev/null +++ b/crates/ias-fetch/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "ias-fetch" +version = "0.1.0" +edition = "2024" + +[dependencies] +chromiumoxide = { version = "0.7", default-features = false, features = ["tokio-runtime", "bytes"] } +futures = "0.3" +serde_json = "1.0" +tokio = { version = "1.0", features = ["full"] } diff --git a/crates/ias-fetch/src/main.rs b/crates/ias-fetch/src/main.rs new file mode 100644 index 0000000..f500959 --- /dev/null +++ b/crates/ias-fetch/src/main.rs @@ -0,0 +1,143 @@ +//! ias-fetch — 通过 headless Chrome 提取网页正文 +//! +//! 启动无头 Chrome 渲染页面 JS,取渲染后的纯文本输出到 stdout。 +//! 专供 JSON 工具 `${cmd:ias-fetch {{url}}}` 调用。 +//! +//! 用法: +//! ias-fetch # 默认:取 body.innerText +//! ias-fetch --raw # 返回完整 HTML +//! ias-fetch --timeout 15 # 自定义超时(默认 15s) + +use futures::StreamExt; +use std::time::Duration; + +use chromiumoxide::{Browser, BrowserConfig, Page}; + +const DEFAULT_TIMEOUT: u64 = 15; +const MAX_OUTPUT_CHARS: usize = 10_000; + +#[tokio::main] +async fn main() { + let args: Vec = std::env::args().collect(); + + let mut raw_mode = false; + let mut timeout_secs = DEFAULT_TIMEOUT; + let mut url: Option = None; + + let mut i = 1; + while i < args.len() { + match args[i].as_str() { + "--raw" => raw_mode = true, + "--timeout" => { + i += 1; + timeout_secs = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(DEFAULT_TIMEOUT); + } + s if s.starts_with("http://") || s.starts_with("https://") => { + url = Some(s.to_string()); + } + s if !s.starts_with('-') => url = Some(s.to_string()), + _ => {} + } + i += 1; + } + + let url = match url { + Some(u) => { + if u.starts_with("http://") || u.starts_with("https://") { + u + } else { + format!("https://{}", u) + } + } + None => { + eprintln!("用法: ias-fetch [--raw] [--timeout N] "); + std::process::exit(1); + } + }; + + let timeout = Duration::from_secs(timeout_secs); + + match tokio::time::timeout(timeout, fetch_page(&url, raw_mode)).await { + Ok(Ok(output)) => { + if output.is_empty() { + eprintln!("页面无内容"); + std::process::exit(1); + } + println!("{}", output); + } + Ok(Err(e)) => { + eprintln!("{}", e); + std::process::exit(1); + } + Err(_) => { + eprintln!("超时 ({}s): {}", timeout_secs, url); + std::process::exit(1); + } + } +} + +async fn fetch_page(url: &str, raw_mode: bool) -> Result { + let (mut browser, mut handler) = Browser::launch( + BrowserConfig::builder() + .build() + .map_err(|e| format!("Chrome 配置失败: {}", e))?, + ) + .await + .map_err(|e| format!("启动 Chrome 失败: {}\n请确保已安装 chromium 或 google-chrome", e))?; + + let handle = tokio::spawn(async move { + while let Some(_event) = handler.next().await { + // 持续消费 handler 事件 + } + }); + + // new_page 接受 URL,自动导航并等待加载 + let page = browser + .new_page(url) + .await + .map_err(|e| format!("导航到 {} 失败: {}", url, e))?; + + // 额外等待 JS 渲染 + tokio::time::sleep(Duration::from_secs(2)).await; + + // 提取内容 + let content = if raw_mode { + page.content().await.map_err(|e| format!("获取 HTML 失败: {}", e))? + } else { + get_inner_text(&page).await? + }; + + // 清理 + browser.close().await.ok(); + handle.await.ok(); + + // 截断 + if content.chars().count() > MAX_OUTPUT_CHARS { + Ok(format!( + "{}…(已截断,完整内容 {} 字符)", + content.chars().take(MAX_OUTPUT_CHARS).collect::(), + content.chars().count() + )) + } else { + Ok(content) + } +} + +/// 通过 evaluate JS 获取 body.innerText +async fn get_inner_text(page: &Page) -> Result { + let result = page + .evaluate("document.body?.innerText || document.documentElement?.innerText || ''") + .await + .map_err(|e| format!("提取正文失败: {}", e))?; + + let value: serde_json::Value = result + .into_value() + .map_err(|e| format!("解析 evaluate 结果失败: {}", e))?; + match value { + serde_json::Value::String(s) => { + let cleaned: Vec<&str> = s.lines().map(|l| l.trim()).filter(|l| !l.is_empty()).collect(); + Ok(cleaned.join("\n")) + } + _ => Ok(String::new()), + } +} diff --git a/src/bin/ias-auth.rs b/src/bin/ias-auth.rs new file mode 100644 index 0000000..d09b3f2 --- /dev/null +++ b/src/bin/ias-auth.rs @@ -0,0 +1,104 @@ +//! ias-auth — 独立 Token 生成器 +//! +//! 专供 JSON 工具定义通过 `${cmd:ias-auth gen-qweather-jwt}` 调用。 +//! 编译快、无运行时依赖,独立于主程序 `ias`。 +//! +//! 用法: +//! ias-auth gen-qweather-jwt 生成和风天气 Ed25519 JWT Token + +use std::process::exit; + +fn main() { + let args: Vec = std::env::args().collect(); + if args.len() < 2 { + eprintln!("用法: ias-auth "); + eprintln!("命令:"); + eprintln!(" gen-qweather-jwt 生成和风天气 Ed25519 JWT Token"); + exit(1); + } + match args[1].as_str() { + "gen-qweather-jwt" => gen_qweather_jwt(), + _ => { + eprintln!("未知命令: {}", args[1]); + eprintln!("支持的命令: gen-qweather-jwt"); + exit(1); + } + } +} + +fn gen_qweather_jwt() { + match generate_jwt() { + Ok(token) => println!("{}", token), + Err(e) => { + eprintln!("{}", e); + exit(1); + } + } +} + +// ─── 和风天气 Ed25519 JWT 生成 ─── +// 从 iAs 的 tools::builtins::weather 提取,保持独立 + +use base64::Engine; +use chrono::Utc; +use ed25519_dalek::pkcs8::DecodePrivateKey; +use ed25519_dalek::{Signer, SigningKey}; + +fn base64url(data: &[u8]) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data) +} + +fn generate_jwt() -> Result { + let key_file = std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE") + .unwrap_or_else(|_| "qweather/ed25519-private.pem".into()); + + let pem = std::fs::read_to_string(&key_file) + .map_err(|e| format!("读取密钥文件 {} 失败: {}", key_file, e))?; + + // 手动解析 PEM: 去掉头尾,base64 解码得到 DER + let der = pem + .lines() + .filter(|l| !l.starts_with("-----")) + .collect::>() + .join(""); + let der_bytes = base64::engine::general_purpose::STANDARD + .decode(&der) + .map_err(|e| format!("Base64 解码密钥失败: {e}"))?; + + let signing_key = SigningKey::from_pkcs8_der(&der_bytes) + .map_err(|e| format!("解析 Ed25519 密钥失败: {e}"))?; + + let key_id = std::env::var("QWEATHER_JWT_KEY_ID").unwrap_or_default(); + let project_id = std::env::var("QWEATHER_JWT_PROJECT_ID").unwrap_or_default(); + let ttl: i64 = std::env::var("QWEATHER_JWT_TTL_SECONDS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(3600); + + let now = Utc::now().timestamp(); + + // Header + let header = serde_json::json!({"alg": "EdDSA", "kid": key_id}); + let header_b64 = base64url(header.to_string().as_bytes()); + + // Payload + let payload = serde_json::json!({ + "sub": project_id, + "iat": now - 30, + "exp": now + ttl, + }); + let payload_b64 = base64url(payload.to_string().as_bytes()); + + // Sign + let message = format!("{}.{}", header_b64, payload_b64); + let signature = signing_key + .try_sign(message.as_bytes()) + .map_err(|e| format!("JWT 签名失败: {e}"))?; + + Ok(format!( + "{}.{}.{}", + header_b64, + payload_b64, + base64url(&signature.to_bytes()) + )) +} diff --git a/src/channel/mod.rs b/src/channel/mod.rs new file mode 100644 index 0000000..c7f249c --- /dev/null +++ b/src/channel/mod.rs @@ -0,0 +1,38 @@ +//! 渠道 / 通道基础类型 +//! +//! 定义消息的来源/目标渠道标识,以及从外部(微信等)收到的原始消息。 + +use serde::{Deserialize, Serialize}; +use std::fmt; + +/// 渠道标识 —— 区分消息来自哪个平台和哪个用户 +#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)] +pub struct ChannelId { + /// 平台名称,如 "wechat"、"telegram"、"console" + pub platform: String, + /// 在该平台上的用户 ID + pub user_id: String, +} + +impl ChannelId { + /// 快捷构造微信渠道 + pub fn wechat(user_id: impl Into) -> Self { + Self { platform: "wechat".into(), user_id: user_id.into() } + } +} + +impl fmt::Display for ChannelId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}:{}", self.platform, self.user_id) + } +} + +/// 从外部轮询收到的原始消息 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[allow(dead_code)] +pub struct IncomingMessage { + pub channel: ChannelId, + pub text: String, + pub message_id: String, + pub context_token: Option, +} diff --git a/src/cli.rs b/src/cli.rs index 35b3b8f..da0ce6c 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -115,7 +115,7 @@ pub enum Commands { sock: String, }, - /// Worker 进程(由 daemon 自动 spawn,也可手动测试) + /// [已废弃] Worker 进程(新架构使用 daemon 统一 tokio task) /// /// 示例: /// ias worker --sock /tmp/ias_daemon.sock diff --git a/src/context/types.rs b/src/context/types.rs index af6e630..911beba 100644 --- a/src/context/types.rs +++ b/src/context/types.rs @@ -2,6 +2,7 @@ use crate::llm::types::Message; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use std::sync::Arc; +use uuid::Uuid; /// 摘要原因 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -44,6 +45,8 @@ pub struct ChatSession { pub token_budget: u32, /// 12 小时空闲阈值(秒) pub idle_timeout_secs: i64, + /// 会话 ID(用于追踪工具调用往返) + pub session_id: Uuid, } impl Default for ChatSession { @@ -58,6 +61,7 @@ impl Default for ChatSession { last_user_at: None, token_budget: 28000, idle_timeout_secs: 12 * 3600, + session_id: Uuid::new_v4(), } } } @@ -90,6 +94,11 @@ impl ChatSession { .push(Message::tool_result(tool_call_id, name, result)); } + /// 添加一条系统说明消息(如已批准的工具) + pub fn add_system_note(&mut self, note: &str) { + self.messages.push(Message::system(note)); + } + /// 是否因空闲超过阈值需要摘要 pub fn is_idle_timeout(&self) -> bool { if let Some(last) = self.last_user_at { diff --git a/src/daemon.rs b/src/daemon.rs index 5c3661d..7f0ba07 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -1,303 +1,192 @@ //! Daemon — 常驻进程 //! -//! 持有 WeChat 长连接 + 数据库。每收到一条消息 spawn 一个 worker 进程, -//! 通过 Unix Domain Socket 通信。Worker 代码更新后编译即可,下一条消息 -//! 自动使用新版本。 +//! 持有 WeChat 长连接 + 数据库。使用消息队列架构将消息处理分解为 +//! LLM / Tool / Send 三个消费者(tokio task),通过 QueueRunner 路由。 +use crate::channel::ChannelId; use crate::context::MemoryStore; use crate::db::Database; -use crate::ipc::{HistoryEntry, OutputFrame, TaskFrame, TaskMessage, recv_output, send_frame}; +use crate::ipc::HistoryEntry; +use crate::llm::{ChatResult, Conversation, ConversationConfig, ToolExecutor, DEFAULT_SYSTEM_PROMPT, ASYNC_MARKER}; +use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner}; use crate::state::StateManager; use crate::tools::approval::{ApprovalDecision, ApprovalManager}; use crate::wechat::client::WeChatClient; -use std::collections::{HashMap, VecDeque}; +use std::collections::HashMap; use std::sync::Arc; use std::time::Instant; -use tokio::net::{UnixListener, UnixStream}; -use tokio::process::Command; use tokio::sync::Mutex; use tracing::{error, info, warn}; +use uuid::Uuid; -/// 消息队列:每个用户最多一个活跃 worker,等待队列 FIFO -struct MessageQueue { - /// 当前是否有活跃 worker - active: HashMap, - /// 每个用户的待处理消息队列 - pending: HashMap>, - /// 等待 worker 连接的用户顺序 - waiting: VecDeque, -} - -struct PendingMessage { - from: String, - text: String, +/// Daemon 上下文 —— 消费者共享的所有状态 +struct DaemonCtx { + db: Option>, + client: WeChatClient, account_id: String, - context_token: Option, - #[allow(dead_code)] - message_id: String, - approved_tool: Option, -} - -impl MessageQueue { - fn new() -> Self { - Self { - active: HashMap::new(), - pending: HashMap::new(), - waiting: VecDeque::new(), - } - } - - fn enqueue(&mut self, user_id: &str, msg: PendingMessage) { - self.pending - .entry(user_id.to_string()) - .or_default() - .push_back(msg); - if !self.active.get(user_id).unwrap_or(&false) && !self.waiting.contains(&user_id.to_string()) { - self.waiting.push_back(user_id.to_string()); - } - } - - /// 取出下一个等待处理的用户(daemon accept worker 后调用) - fn pop_waiting(&mut self) -> Option { - while let Some(uid) = self.waiting.pop_front() { - // 跳过已经活跃的 - if *self.active.get(&uid).unwrap_or(&false) { - continue; - } - // 检查是否有待处理消息 - if self.pending.get(&uid).map_or(true, |q| q.is_empty()) { - continue; - } - self.active.insert(uid.clone(), true); - return Some(uid); - } - None - } - - fn pop_message(&mut self, user_id: &str) -> Option { - self.pending.get_mut(user_id)?.pop_front() - } - - fn deactivate(&mut self, user_id: &str) { - self.active.insert(user_id.to_string(), false); - } - - fn has_pending(&self, user_id: &str) -> bool { - self.pending - .get(user_id) - .map_or(false, |q| !q.is_empty()) - } + approval: Arc, + memory_store: Arc, + tools_list: Vec, + system_prompt: String, + model: String, + /// 待审批: user_id → (原始文本, context_token, 工具名, tool_args, 创建时间) + approval_ctx: Mutex, String, String, Instant)>>, } /// Daemon 入口 pub async fn run( - sock_path: String, + _sock_path: String, database: &Option>, file_state: &StateManager, memory_store: &Arc, ) { - // 1. 加载认证 + // 1-3. 认证 & 客户端 let auth = match load_auth(database, file_state).await { Some(a) => a, - None => { - error!("未登录,请先执行 ias login"); - return; - } + None => { error!("未登录"); return; } }; - - // 2. 创建 WeChat 客户端 let mut client = WeChatClient::new(Some(auth.base_url.clone())); - client - .set_auth(&auth.token, &auth.account_id, &auth.base_url) - .await; - - // 恢复 runtime buf - if let Some(r) = file_state.load_runtime() { - client.set_updates_buf(&r.get_updates_buf).await; - } - - // 3. 注册监听器 - if let Err(e) = client.notify_start().await { - error!("注册监听器失败: {}", e); - return; - } - + client.set_auth(&auth.token, &auth.account_id, &auth.base_url).await; + if let Some(r) = file_state.load_runtime() { client.set_updates_buf(&r.get_updates_buf).await; } + if let Err(e) = client.notify_start().await { error!("注册监听器失败: {}", e); return; } let account_id = auth.account_id.clone(); + info!("Daemon 已启动 (无 worker 进程模式)"); - // 4. 绑定 Unix Domain Socket - let _ = std::fs::remove_file(&sock_path); - let listener = match UnixListener::bind(&sock_path) { - Ok(l) => l, - Err(e) => { - error!("绑定 Unix Socket 失败 ({}): {}", sock_path, e); - return; - } - }; - - info!("Daemon 已启动: sock={}", sock_path); - - // 5. 审批管理器 + // 4-6. 审批管理器、工具列表、模型配置 let approval_manager = Arc::new(ApprovalManager::new( database.as_ref().map(|db| Arc::new(db.pool().clone())), )); - - // 6. 消息队列 - let queue = Arc::new(Mutex::new(MessageQueue::new())); - - // 7. 构建工具列表(传给 worker) let tools_list = build_tools_list(); - - // 8. 环境变量映射 - let env_map = build_env_map(); - - // 9. 系统提示词 let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new()); let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string()); - // 10. 共享数据 - let shared = Arc::new(DaemonShared { + // 7. Daemon 上下文(注意 approval_ctx 为 5 元组) + let ctx = Arc::new(DaemonCtx { db: database.clone(), client: client.clone(), account_id: account_id.clone(), approval: approval_manager.clone(), memory_store: memory_store.clone(), - sock_path: sock_path.clone(), tools_list: tools_list.clone(), - env_map, system_prompt, model, approval_ctx: Mutex::new(HashMap::new()), }); - // 11. 审批过期清理 + // 8. 审批过期清理 + let ctx_for_clean = ctx.clone(); let approval_clean = approval_manager.clone(); - let shared_for_clean = shared.clone(); tokio::spawn(async move { loop { tokio::time::sleep(std::time::Duration::from_secs(60)).await; approval_clean.clean_expired().await; - // 清理过期的审批上下文(超过 5 分钟) let now = Instant::now(); - shared_for_clean.approval_ctx.lock().await - .retain(|_, (_, _, _, created)| now.duration_since(*created).as_secs() < 300); + ctx_for_clean.approval_ctx.lock().await + .retain(|_, (_, _, _, _, created)| now.duration_since(*created).as_secs() < 300); } }); - // 12. 调度器 + // 9. 队列(容量 1024) + let (runner, channels) = QueueRunner::new(1024); + let enqueue_handle = runner.enqueue_handle(); + let shutdown_handle = runner.shutdown_handle(); // 不再 prefix _ + + // 10-11. 消费者 + 调度器 + let (consumer_handles, sessions) = spawn_consumers(channels, ctx.clone(), enqueue_handle.clone()); + if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) { - let sched_client = client.clone(); + let sched_enqueue = enqueue_handle.clone(); tokio::spawn(async move { let scheduler = crate::scheduler::Scheduler::new(Some(Arc::new(sched_db))); - scheduler - .run(move |to: &str, _name: &str, msg: &str| { - let c = sched_client.clone(); - let uid = to.to_string(); - let text = msg.to_string(); - tokio::task::spawn(async move { - if let Err(e) = c.send_text(&uid, &text, None).await { - tracing::error!("发送调度结果失败: {}", e); - } - }); - Ok(()) - }) - .await; + scheduler.run(move |to: &str, name: &str, msg: &str| { + let eq = sched_enqueue.clone(); + let uid = to.to_string(); let text = msg.to_string(); let tn = name.to_string(); + tokio::spawn(async move { + eq.enqueue(PipelineMessage { + channel: ChannelId::wechat(uid), correlation_id: Uuid::new_v4(), + kind: MessageKind::ScheduledTask { text, task_name: tn }, + }).await; + }); + Ok(()) + }).await; }); info!("定时任务调度器已启动"); } + // 12. QueueRunner + let runner_handle = tokio::spawn(async move { runner.run().await; }); + + // 13. Session TTL 清理(5 分钟无活动) + let sessions_for_clean = sessions.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + let now = Instant::now(); + sessions_for_clean.lock().await + .retain(|_, (_, created)| now.duration_since(*created).as_secs() < 300); + } + }); + info!("开始监听消息..."); - // 13. 主循环 - loop { + // 13. 主循环 — 带 Ctrl+C 优雅退出 + let mut running = true; + while running { tokio::select! { - // WeChat 消息 msg_result = client.receive_messages() => { match msg_result { Ok(resp) => { - // 持久化 runtime buf if !resp.get_updates_buf.is_empty() { file_state.save_runtime(&resp.get_updates_buf); } - for msg in &resp.msgs { - if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) { - continue; - } + if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) { continue; } let from = msg.from_user_id.as_deref().unwrap_or("unknown"); let text = msg.text_content().unwrap_or("(非文本消息)"); let ctx_token = msg.context_token.as_deref(); let msg_id = msg.message_id.map(|id| id.to_string()).unwrap_or_default(); - info!("收到消息 from={}: {}", from, text); - // 入库:收到的消息 - if let Some(db) = &shared.db { - if let Err(e) = crate::db::models::insert_chat_record( - db.pool(), - "inbound", - from, - &shared.account_id, - text, - "wechat", - ctx_token, - &msg_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } + if let Some(db) = &ctx.db { + let _ = crate::db::models::insert_chat_record( + db.pool(), "inbound", from, &ctx.account_id, text, "wechat", ctx_token, &msg_id, + ).await; } - // 审批回复检查 - if let Some((skill_name, decision)) = shared.approval.handle_reply(from, text).await { + // 审批回复 + if let Some((skill_name, decision)) = ctx.approval.handle_reply(from, text).await { match decision { ApprovalDecision::Approved => { info!("审批通过: {}", skill_name); - let original = shared.approval_ctx.lock().await.remove(from); - if let Some((orig_text, orig_ctx, _tool, _ts)) = original { - let mut q = queue.lock().await; - q.enqueue(from, PendingMessage { - from: from.to_string(), - text: orig_text, - account_id: shared.account_id.clone(), - context_token: orig_ctx, - message_id: String::new(), - approved_tool: Some(skill_name), - }); - drop(q); - spawn_worker_for_user(from, &queue, &shared).await; + // 5 元组: (orig_text, context_token, tool_name, tool_args, timestamp) + if let Some((orig_text, orig_ctx, _tn, _ta, _ts)) = ctx.approval_ctx.lock().await.remove(from) { + enqueue_handle.enqueue(PipelineMessage { + channel: ChannelId::wechat(from), + correlation_id: Uuid::new_v4(), + kind: MessageKind::UserMessage { + text: orig_text, message_id: String::new(), + context_token: orig_ctx, approved_tool: Some(skill_name), + }, + }).await; } } - ApprovalDecision::Rejected => { - info!("审批被拒绝: {}", skill_name); - shared.approval_ctx.lock().await.remove(from); - } - ApprovalDecision::Expired => { - info!("审批已过期: {}", skill_name); - shared.approval_ctx.lock().await.remove(from); + ApprovalDecision::Rejected | ApprovalDecision::Expired => { + ctx.approval_ctx.lock().await.remove(from); } } continue; } - // 消息入队 - { - let mut q = queue.lock().await; - q.enqueue(from, PendingMessage { - from: from.to_string(), - text: text.to_string(), - account_id: shared.account_id.clone(), - context_token: ctx_token.map(String::from), - message_id: msg_id, - approved_tool: None, - }); - drop(q); - } - - // 尝试 spawn worker - spawn_worker_for_user(from, &queue, &shared).await; + enqueue_handle.enqueue(PipelineMessage { + channel: ChannelId::wechat(from), + correlation_id: Uuid::new_v4(), + kind: MessageKind::UserMessage { + text: text.to_string(), message_id: msg_id, + context_token: ctx_token.map(String::from), approved_tool: None, + }, + }).await; } } Err(e) => { @@ -308,318 +197,414 @@ pub async fn run( } } } - - // Worker 连接 - accept_result = listener.accept() => { - match accept_result { - Ok((stream, _)) => { - let q = queue.clone(); - let shr = shared.clone(); - tokio::spawn(async move { - handle_worker(stream, q, shr).await; - }); - } - Err(e) => { - error!("接受 Worker 连接失败: {}", e); - } - } + _ = tokio::signal::ctrl_c() => { + info!("收到 Ctrl+C,优雅关闭中..."); + running = false; } } } + + shutdown_handle.shutdown(); + info!("等待消费者处理剩余消息..."); + tokio::time::sleep(std::time::Duration::from_secs(3)).await; + for h in consumer_handles { h.abort(); } + runner_handle.abort(); + info!("Daemon 已停止"); } -/// 共享数据 -struct DaemonShared { - db: Option>, - client: WeChatClient, - account_id: String, - approval: Arc, - memory_store: Arc, - sock_path: String, - tools_list: Vec, - env_map: HashMap, - system_prompt: String, - model: String, - /// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 创建时间) - approval_ctx: Mutex, String, Instant)>>, +// ─── 消费者 ─── + +fn spawn_consumers( + channels: ConsumerChannels, + ctx: Arc, + enqueue: EnqueueHandle, +) -> (Vec>, Arc>>) { + let sessions: Arc>> = Arc::new(Mutex::new(HashMap::new())); + + let mut handles = Vec::new(); + // LLM Consumer + { + let mut rx = channels.llm_rx; + let ctx = ctx.clone(); + let eq = enqueue.clone(); + let s = sessions.clone(); + handles.push(tokio::spawn(async move { llm_consumer_loop(&mut rx, &ctx, &eq, &s).await; })); + } + // Tool Consumer + { + let mut rx = channels.tool_rx; + let ctx = ctx.clone(); + let eq = enqueue.clone(); + handles.push(tokio::spawn(async move { tool_consumer_loop(&mut rx, &ctx, &eq).await; })); + } + // Send Consumer + { + let mut rx = channels.send_rx; + let ctx = ctx.clone(); + handles.push(tokio::spawn(async move { send_consumer_loop(&mut rx, &ctx).await; })); + } + (handles, sessions) } -/// 处理一个 Worker 连接 -async fn handle_worker( - mut stream: UnixStream, - queue: Arc>, - shared: Arc, +// ─── LLM Consumer ─── + +async fn llm_consumer_loop( + rx: &mut tokio::sync::mpsc::Receiver, + ctx: &DaemonCtx, + enqueue: &EnqueueHandle, + sessions: &Arc>>, ) { - // 1. 取出等待的用户 - let user_id = { - let mut q = queue.lock().await; - q.pop_waiting() - }; + info!("LLM Consumer 已启动"); - let user_id = match user_id { - Some(uid) => uid, - None => { - warn!("Worker 连接但没有等待的用户"); - return; + while let Some(msg) = rx.recv().await { + let user_id = msg.channel.user_id.clone(); + let correlation_id = msg.correlation_id; + + match msg.kind { + MessageKind::UserMessage { text, message_id: _, context_token, approved_tool } => { + info!("LLM Consumer: user={} text={:.60}", user_id, text); + + // 加载数据 + ctx.memory_store.load(&user_id).await; + let history = load_history(&ctx.db, &user_id).await; + + let sys_prompt = if ctx.system_prompt.is_empty() { + DEFAULT_SYSTEM_PROMPT.to_string() + } else { ctx.system_prompt.clone() }; + + let cfg = ConversationConfig { + system_prompt: sys_prompt, + model: ctx.model.clone(), + tools: Some(ctx.tools_list.clone()), + ..Default::default() + }; + + let mut conv = match Conversation::new(cfg) { + Ok(c) => c, + Err(e) => { error!("LLM 初始化失败: {}", e); continue; } + }; + + if let Some(db) = &ctx.db { + conv.session().lock().await.db_pool = Some(Arc::new(db.pool().clone())); + } + + { + let s_arc = conv.session(); + let mut s = s_arc.lock().await; + s.load_from_history(&history); + s.current_user_id = user_id.clone(); + } + + if let Some(tool) = &approved_tool { + conv.session().lock().await.add_system_note(&format!("\n[用户已批准工具: {}]", tool)); + } + + // 设置 ToolExecutor — 所有工具统一通过消息队列异步执行 + let eq = enqueue.clone(); + let uid = user_id.clone(); + let corr = correlation_id; + let session = conv.session(); + + let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| { + // 所有工具统一通过消息队列异步执行 + let eq = eq.clone(); + let uid = uid.clone(); + let corr = corr; + let session = session.clone(); + let n = name.to_string(); + let args = args_json.to_string(); + + Box::pin(async move { + let ch = ChannelId::wechat(&uid); + let sid = { session.lock().await.session_id }; + eq.enqueue(PipelineMessage::tool_call( + ch, corr, sid, &Uuid::new_v4().to_string(), &n, &args, None, + )).await; + Ok(format!("{}{}", ASYNC_MARKER, n)) + }) + }); + conv.set_tool_executor(executor); + + // 缓存 Conversation(记录创建时间用于 TTL) + sessions.lock().await.insert(correlation_id, (conv, Instant::now())); + + // 执行 LLM 对话 + run_llm_round(&user_id, correlation_id, context_token, &text, sessions, enqueue).await; + } + + MessageKind::ToolResult { session_id: _, tool_call_id, tool_name, result, context_token } => { + info!("LLM Consumer: 工具结果回喂 user={} tool={}", user_id, tool_name); + + // 从缓存恢复会话,注入工具结果并恢复 LLM 循环 + let should_continue = { + let map = sessions.lock().await; + if let Some((conv, _)) = map.get(&correlation_id) { + conv.session().lock().await.add_tool_result(&tool_call_id, &tool_name, &result); + true + } else { + warn!("会话已过期: corr={}", correlation_id); + false + } + }; + + if should_continue { + resume_llm_round(&user_id, correlation_id, context_token, sessions, enqueue).await; + } + } + + MessageKind::ScheduledTask { text, task_name } => { + info!("LLM Consumer: 定时任务 user={} task={}", user_id, task_name); + let prompt = format!("[定时任务: {}]\n{}", task_name, text); + + // 为保持简洁,定时任务直接作为用户消息处理(使用全新 correlation_id) + let msg = PipelineMessage { + channel: ChannelId::wechat(user_id.clone()), + correlation_id: Uuid::new_v4(), + kind: MessageKind::UserMessage { + text: prompt, + message_id: String::new(), + context_token: None, + approved_tool: None, + }, + }; + // 重新投递给自己 — 简单但可靠 + // 注意: 这里用了一个技巧,重新入队而不是复制代码 + let eq = enqueue.clone(); + eq.enqueue(msg).await; + } + + other => warn!("LLM Consumer 收到不期望的消息: {:?}", std::mem::discriminant(&other)), + } + } + info!("LLM Consumer 已停止"); +} + +/// 执行一轮 LLM 对话,仅在无异步工具等待时发送回复并清理会话。 +async fn run_llm_round( + user_id: &str, + correlation_id: Uuid, + context_token: Option, + text: &str, + sessions: &Arc>>, + enqueue: &EnqueueHandle, +) { + let result = { + let map = sessions.lock().await; + let conv = match map.get(&correlation_id) { + Some((c, _)) => c, + None => { error!("run_llm_round: conv 不在缓存"); return; } + }; + + match conv.chat_with_tools(text.to_string()).await { + Ok(r) => r, + Err(e) => { + error!("LLM 对话失败: {}", e); + ChatResult { + reply: format!("处理失败: {}", e), + used_tools: false, + usage: None, + has_pending_async: false, + } + } } }; - // 2. 取出该用户的一条消息 - let pending_msg = { - let mut q = queue.lock().await; - q.pop_message(&user_id) - }; + handle_chat_result(user_id, correlation_id, context_token, result, sessions, enqueue).await; +} - let pending_msg = match pending_msg { - Some(m) => m, - None => { - warn!("Worker 连接但用户 {} 无待处理消息", user_id); - queue.lock().await.deactivate(&user_id); - return; +/// 在异步工具结果返回后恢复 LLM 循环(不追加用户消息)。 +async fn resume_llm_round( + user_id: &str, + correlation_id: Uuid, + context_token: Option, + sessions: &Arc>>, + enqueue: &EnqueueHandle, +) { + let result = { + let map = sessions.lock().await; + let conv = match map.get(&correlation_id) { + Some((c, _)) => c, + None => { error!("resume_llm_round: conv 不在缓存"); return; } + }; + + match conv.resume_tool_loop().await { + Ok(r) => r, + Err(e) => { + error!("LLM 恢复失败: {}", e); + ChatResult { + reply: format!("处理失败: {}", e), + used_tools: false, + usage: None, + has_pending_async: false, + } + } } }; - info!( - "Worker 处理: user={} text={:.60}", - user_id, pending_msg.text - ); + handle_chat_result(user_id, correlation_id, context_token, result, sessions, enqueue).await; +} - // 3. 加载用户记忆 - shared.memory_store.load(&user_id).await; - - // 4. 预载历史消息 - let history = load_history(&shared.db, &user_id).await; - let memories = load_memories(&shared.memory_store, &user_id).await; - let summaries = load_summaries(&shared.db, &user_id).await; - - // 5. 构建 TaskFrame - let approved_tool_info = pending_msg.approved_tool.clone(); - let task = TaskFrame { - user_id: user_id.clone(), - msg: TaskMessage { - from: pending_msg.from, - text: pending_msg.text.clone(), - account_id: pending_msg.account_id, - context_token: pending_msg.context_token.clone(), - }, - history, - memories, - summaries, - approved_tool: approved_tool_info.clone(), - env: shared.env_map.clone(), - tools: shared.tools_list.clone(), - system_prompt: shared.system_prompt.clone(), - model: shared.model.clone(), - }; - - // 6. 发送 task 帧 - if let Err(e) = send_frame(&mut stream, &task).await { - error!("发送 task 帧失败: {}", e); - queue.lock().await.deactivate(&user_id); +/// 统一处理 ChatResult:有异步等待则保留 session,否则发送回复并清理。 +async fn handle_chat_result( + user_id: &str, + correlation_id: Uuid, + context_token: Option, + result: ChatResult, + sessions: &Arc>>, + enqueue: &EnqueueHandle, +) { + if result.has_pending_async { + // 有异步工具等待返回 → 保留 session + if !result.reply.is_empty() { + let ch = ChannelId::wechat(user_id); + let msg = PipelineMessage::llm_reply(ch, correlation_id, &result.reply, context_token); + enqueue.enqueue(msg).await; + } return; } - // 7. 逐帧处理 worker 输出 - let mut reply_text = String::new(); - let mut chunk_buffer = String::new(); - let ctx_token = task.msg.context_token.clone(); + // 对话完成:发送回复并清理 session + if !result.reply.is_empty() { + let ch = ChannelId::wechat(user_id); + let msg = PipelineMessage::llm_reply(ch, correlation_id, &result.reply, context_token); + enqueue.enqueue(msg).await; + } + sessions.lock().await.remove(&correlation_id); +} - loop { - match recv_output(&mut stream).await { - Ok(frame) => match frame { - OutputFrame::Chunk { text } => { - chunk_buffer.push_str(&text); - } - OutputFrame::Reply { text } => { - reply_text = text; - } - OutputFrame::RequestApproval { tool, reason } => { - info!("Worker 请求审批: tool={} reason={}", tool, reason); - // 创建审批记录 - match shared.approval.create(&user_id, &tool).await { +// ─── Tool Consumer ─── + +async fn tool_consumer_loop( + rx: &mut tokio::sync::mpsc::Receiver, + ctx: &DaemonCtx, + enqueue: &EnqueueHandle, +) { + info!("Tool Consumer 已启动"); + + while let Some(msg) = rx.recv().await { + let user_id = msg.channel.user_id.clone(); + let correlation_id = msg.correlation_id; + + match msg.kind { + MessageKind::ToolCall { session_id, tool_call_id: _, tool_name, arguments, context_token } => { + info!("Tool Consumer: user={} tool={}", user_id, tool_name); + + // 高风险工具 → 审批 + if crate::tools::builtin::BuiltinRegistry::is_high_risk(&tool_name) { + info!("高风险工具 {} 需要审批", tool_name); + ctx.approval_ctx.lock().await.insert( + user_id.clone(), + ( + format!("[工具: {}]", tool_name), + context_token.clone(), + tool_name.clone(), + arguments.clone(), + Instant::now(), + ), + ); + match ctx.approval.create(&user_id, &tool_name).await { Ok((code, _rx)) => { - let msg = format!( + let m = format!( "⚠️ 操作确认\n\n技能:{}\n\n确认码:{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效)", - tool, code + tool_name, code ); - // 保存原始消息上下文以便审批后重新处理 - shared.approval_ctx.lock().await.insert( - user_id.clone(), - (task.msg.text.clone(), ctx_token.clone(), tool.clone(), Instant::now()), - ); - // 发送确认消息到微信 - if let Err(e) = shared.client.send_text(&user_id, &msg, ctx_token.as_deref()).await { - error!("发送审批消息失败: {}", e); - } - } - Err(e) => { - error!("创建审批失败: {}", e); - let _ = shared.client.send_text(&user_id, &format!("审批创建失败: {}", e), ctx_token.as_deref()).await; + // 通过消息队列发送审批提示(不直接调 send_text) + let ch = ChannelId::wechat(&user_id); + enqueue.enqueue(PipelineMessage::llm_reply(ch, correlation_id, &m, None)).await; } + Err(e) => error!("创建审批失败: {}", e), } + continue; } - OutputFrame::NeedApproval { tool, .. } => { - warn!("收到已废弃的 NeedApproval 帧 (tool={})", tool); - } - OutputFrame::DbWrite { table, row } => { - handle_db_write(&shared.db, &shared.memory_store, &table, &row, &user_id, &shared.account_id).await; - } - OutputFrame::UsageReport { - model, - provider, - prompt_tokens, - completion_tokens, - total_tokens: _, - cache_hit_tokens, - cache_miss_tokens, - user_id: uid, - } => { - if let Some(db) = &shared.db { - if let Err(e) = crate::db::models::insert_llm_usage( - db.pool(), - &uid, - &model, - &provider, - prompt_tokens, - completion_tokens, - cache_hit_tokens, - cache_miss_tokens, - ) - .await - { - error!("存储 LLM 用量失败: {}", e); - } - } - } - OutputFrame::Bye => break, - }, - Err(e) => { - error!("读取 Worker 输出帧失败: {}", e); - break; + + // 执行工具 — 所有工具统一在此处理 + let output = execute_tool(&tool_name, &arguments, ctx, &user_id).await; + + enqueue.enqueue(PipelineMessage::tool_result( + ChannelId::wechat(&user_id), correlation_id, session_id, "", &tool_name, &output, context_token, + )).await; } + MessageKind::ApprovalRequest { .. } => { + warn!("Tool Consumer 收到 ApprovalRequest"); + } + other => warn!("Tool Consumer 收到不期望: {:?}", std::mem::discriminant(&other)), } } + info!("Tool Consumer 已停止"); +} - // 8. 发送最终回复 - if !reply_text.is_empty() { - match shared - .client - .send_text(&user_id, &reply_text, ctx_token.as_deref()) - .await - { - Ok(sent_id) => { - info!("回复成功 user={} msg_id={}", user_id, sent_id); - // 入库:发送的回复 - if let Some(db) = &shared.db { - if let Err(e) = crate::db::models::insert_chat_record( - db.pool(), - "outbound", - &user_id, - &shared.account_id, - &reply_text, - "llm", - ctx_token.as_deref(), - &sent_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } - } - } - Err(e) => error!("发送回复失败: {}", e), +// ─── 统一工具执行 ─── + +/// 所有工具的统一执行入口 — 仅此一处 +async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &str) -> String { + // 上下文工具 + match name { + "read_memories" => { + let mems = load_memories(&ctx.memory_store, user_id).await; + if mems.is_empty() { return "暂无长期记忆".to_string(); } + let lines: Vec = mems.iter().enumerate() + .map(|(i, v)| format!("{}. {}", i + 1, v)).collect(); + return lines.join("\n"); } - } else if !chunk_buffer.is_empty() { - // 有 chunk 但没有 reply — 用 chunk 作为回复 - match shared - .client - .send_text(&user_id, &chunk_buffer, ctx_token.as_deref()) - .await - { - Ok(sent_id) => { - info!("回复成功 (chunk) user={} msg_id={}", user_id, sent_id); - if let Some(db) = &shared.db { - if let Err(e) = crate::db::models::insert_chat_record( - db.pool(), - "outbound", - &user_id, - &shared.account_id, - &chunk_buffer, - "llm", - ctx_token.as_deref(), - &sent_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } - } - } - Err(e) => error!("发送回复失败: {}", e), + "write_memory" => { + let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default(); + let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if content.is_empty() { return "请提供 content 参数".to_string(); } + ctx.memory_store.write_for(user_id, content).await; + return format!("已添加长期记忆: {}", content); } + "read_summaries" => { + let sums = load_summaries(&ctx.db, user_id).await; + if sums.is_empty() { return "暂无历史摘要".to_string(); } + return sums.iter().enumerate() + .map(|(i, s)| format!("{}. {}", i + 1, s)) + .collect::>().join("\n---\n"); + } + "query_capabilities" => { + let specs = crate::tools::builtin::BuiltinRegistry::specs(); + return crate::tools::build_capability_guide(&specs); + } + _ => {} } - // 9. 清理并继续处理该用户的下一消息 - { - let mut q = queue.lock().await; - q.deactivate(&user_id); - let has_more = q.has_pending(&user_id); - drop(q); - if has_more { - spawn_worker_for_user(&user_id, &queue, &shared).await; - } + // 内置工具注册表 + match crate::tools::builtin::BuiltinRegistry::execute(name, arguments).await { + Some(r) => r.output, + None => format!("未知工具: {}", name), } } -/// 为用户 spawn worker 进程 -async fn spawn_worker_for_user( - user_id: &str, - queue: &Arc>, - shared: &Arc, +// ─── Send Consumer ─── + +async fn send_consumer_loop( + rx: &mut tokio::sync::mpsc::Receiver, + ctx: &DaemonCtx, ) { - // 检查是否已有活跃 worker - { - let q = queue.lock().await; - if *q.active.get(user_id).unwrap_or(&false) { - return; // 已有 worker 在处理 - } - } + info!("Send Consumer 已启动"); - let sock = shared.sock_path.clone(); - let exe = match std::env::current_exe() { - Ok(p) => p, - Err(e) => { - error!("获取当前可执行文件路径失败: {}", e); - return; - } - }; + while let Some(msg) = rx.recv().await { + match msg.kind { + MessageKind::LLMReply { text, context_token } => { + let user_id = &msg.channel.user_id; + info!("Send Consumer: user={} text_len={}", user_id, text.len()); - info!("Spawning worker for user={}", user_id); - - match Command::new(&exe) - .arg("worker") - .arg("--sock") - .arg(&sock) - .spawn() - { - Ok(mut child) => { - // 后台等待 worker 退出 - tokio::spawn(async move { - match child.wait().await { - Ok(status) => { - if !status.success() { - warn!("Worker 退出码: {:?}", status.code()); + match ctx.client.send_text(user_id, &text, context_token.as_deref()).await { + Ok(sent_id) => { + info!("回复成功 user={} msg_id={}", user_id, sent_id); + if let Some(db) = &ctx.db { + let _ = crate::db::models::insert_chat_record( + db.pool(), "outbound", user_id, &ctx.account_id, + &text, "llm", context_token.as_deref(), &sent_id, + ).await; } } - Err(e) => error!("等待 Worker 退出失败: {}", e), + Err(e) => error!("发送回复失败: {}", e), } - }); - } - Err(e) => { - error!("Spawn worker 失败: {}", e); - queue.lock().await.deactivate(user_id); + } + other => warn!("Send Consumer 收到不期望: {:?}", std::mem::discriminant(&other)), } } + info!("Send Consumer 已停止"); } // ─── 辅助函数 ─── @@ -629,160 +614,54 @@ async fn load_auth( file_state: &StateManager, ) -> Option { if let Some(db) = database { - if let Some(a) = crate::db::models::load_auth(db.pool()).await { - return Some(a); - } - // 尝试从文件迁移 + if let Some(a) = crate::db::models::load_auth(db.pool()).await { return Some(a); } if let Some(fa) = file_state.load_auth() { - let a = crate::state::AuthState { - token: fa.token, - account_id: fa.account_id, - base_url: fa.base_url, - }; - if let Err(e) = crate::db::models::save_auth(db.pool(), &a).await { - error!("保存认证信息失败: {}", e); - } + let a = crate::state::AuthState { token: fa.token, account_id: fa.account_id, base_url: fa.base_url }; + if let Err(e) = crate::db::models::save_auth(db.pool(), &a).await { error!("保存 auth 失败: {}", e); } return Some(a); } - } else if let Some(a) = file_state.load_auth() { - return Some(a); - } + } else if let Some(a) = file_state.load_auth() { return Some(a); } None } async fn load_history(db: &Option>, user_id: &str) -> Vec { - let Some(db) = db else { - return vec![]; - }; + let Some(db) = db else { return vec![]; }; match crate::db::models::list_recent_chat_records(db.pool(), user_id, 20).await { - Ok(records) => records - .into_iter() - .rev() // 时间顺序 - .map(|r| HistoryEntry { - role: if r.direction == "inbound" { - "user".into() - } else { - "assistant".into() - }, - content: r.text, - }) - .collect(), - Err(e) => { - warn!("加载历史消息失败: {}", e); - vec![] - } + Ok(records) => records.into_iter().rev().map(|r| HistoryEntry { + role: if r.direction == "inbound" { "user".into() } else { "assistant".into() }, + content: r.text, + }).collect(), + Err(e) => { warn!("加载历史失败: {}", e); vec![] } } } async fn load_memories(memory_store: &MemoryStore, user_id: &str) -> Vec { let text = memory_store.read_for(user_id).await; - if text == "暂无长期记忆" { - vec![] - } else { - text.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()).collect() - } + if text == "暂无长期记忆" { vec![] } + else { text.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()).collect() } } async fn load_summaries(db: &Option>, user_id: &str) -> Vec { - let Some(db) = db else { - return vec![]; - }; + let Some(db) = db else { return vec![]; }; match crate::db::models::load_summaries(db.pool(), user_id, 5).await { Ok(s) => s, - Err(e) => { - warn!("加载摘要失败: {}", e); - vec![] - } - } -} - -async fn handle_db_write( - db: &Option>, - memory_store: &MemoryStore, - table: &str, - row: &serde_json::Value, - user_id: &str, - account_id: &str, -) { - // user_memories 不依赖 DB(memory_store.write_for 自己处理持久化) - if table == "user_memories" { - let content = row.get("content").and_then(|v| v.as_str()).unwrap_or(""); - if !content.is_empty() { - memory_store.write_for(user_id, content).await; - } - return; - } - - let Some(db) = db else { - return; - }; - match table { - "chat_records" => { - let direction = row.get("direction").and_then(|v| v.as_str()).unwrap_or("outbound"); - let text = row.get("text").and_then(|v| v.as_str()).unwrap_or(""); - let source = row.get("source").and_then(|v| v.as_str()).unwrap_or("llm"); - let ctx_token = row.get("context_token").and_then(|v| v.as_str()); - let msg_id = row.get("message_id").and_then(|v| v.as_str()).unwrap_or(""); - if let Err(e) = crate::db::models::insert_chat_record( - db.pool(), - direction, - user_id, - account_id, - text, - source, - ctx_token, - msg_id, - ) - .await - { - error!("DbWrite chat_records 失败: {}", e); - } - } - "llm_usage" => { - let model = row.get("model").and_then(|v| v.as_str()).unwrap_or(""); - let provider = row.get("provider").and_then(|v| v.as_str()).unwrap_or(""); - let prompt_tokens = row.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let completion_tokens = row.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let cache_hit = row.get("cache_hit_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let cache_miss = row.get("cache_miss_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - if let Err(e) = crate::db::models::insert_llm_usage( - db.pool(), - user_id, - model, - provider, - prompt_tokens, - completion_tokens, - cache_hit, - cache_miss, - ) - .await - { - error!("DbWrite llm_usage 失败: {}", e); - } - } - _ => warn!("未知 DbWrite table: {}", table), + Err(e) => { warn!("加载摘要失败: {}", e); vec![] } } } fn build_tools_list() -> Vec { let mut list = Vec::new(); - - // query_capabilities list.push(serde_json::json!({ - "type": "function", - "function": { + "type": "function", "function": { "name": "query_capabilities", - "description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。", + "description": "列出所有可用工具的详细说明、参数格式和调用示例。", "parameters": { "type": "object", "properties": {} } } })); - - // call_capability list.push(serde_json::json!({ - "type": "function", - "function": { + "type": "function", "function": { "name": "call_capability", - "description": "调用一个已注册的工具。name 为工具名(来自 query_capabilities),工具所需的具体参数应直接放在 JSON 中。", + "description": "调用一个已注册的工具。name 为工具名。", "parameters": { "type": "object", "properties": { @@ -793,135 +672,31 @@ fn build_tools_list() -> Vec { } } })); - - // 上下文工具 list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "read_memories", - "description": "读取用户的长期记忆(偏好、个人信息、约定)", + "type": "function", "function": { + "name": "read_memories", "description": "读取用户的长期记忆", "parameters": { "type": "object", "properties": {} } } })); list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "write_memory", - "description": "记录用户的重要信息或偏好", - "parameters": { - "type": "object", - "properties": { - "content": { "type": "string", "description": "记忆内容" } - }, - "required": ["content"] - } + "type": "function", "function": { + "name": "write_memory", "description": "记录用户的重要信息或偏好", + "parameters": { "type": "object", "properties": { "content": { "type": "string" } }, "required": ["content"] } } })); list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "read_summaries", - "description": "读取历史会话摘要", + "type": "function", "function": { + "name": "read_summaries", "description": "读取历史会话摘要", "parameters": { "type": "object", "properties": {} } } })); - // 业务工具(与 BuiltinRegistry 保持一致) - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "get_current_datetime", - "description": "获取当前日期时间(北京时间 UTC+8)", - "parameters": { "type": "object", "properties": {"format": {"type": "string"}} } - } - })); - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "manage_memos", - "description": "管理备忘录:添加/列出/删除", - "parameters": { - "type": "object", - "properties": { - "action": {"type": "string", "enum": ["add", "list", "delete"]}, - "content": {"type": "string"}, - "id": {"type": "integer"} - } - } - } - })); - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "query_weather", - "description": "查询指定城市的天气信息", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "城市名称"}, - "days": {"type": "integer", "description": "查询天数(1-7)"} - }, - "required": ["location"] - } - } - })); - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "web_search", - "description": "联网搜索最新信息。通过 Tavily API 搜索互联网,获取实时、准确的结果。适用于需要最新资讯、事实查询的场景。", - "parameters": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "搜索关键词"}, - "max_results": {"type": "integer", "description": "最大结果数(1-20,默认5)"}, - "include_answer": {"type": "boolean", "description": "是否包含 AI 生成的摘要回答"}, - "search_depth": {"type": "string", "enum": ["basic", "advanced", "fast", "ultra-fast"], "description": "搜索深度"}, - "topic": {"type": "string", "enum": ["general", "news", "finance"], "description": "搜索主题"}, - "time_range": {"type": "string", "enum": ["day", "week", "month", "year"], "description": "按发布时间过滤"}, - "start_date": {"type": "string", "description": "起始日期 YYYY-MM-DD"}, - "end_date": {"type": "string", "description": "结束日期 YYYY-MM-DD"}, - "include_domains": {"type": "array", "items": {"type": "string"}, "description": "限定域名列表"}, - "exclude_domains": {"type": "array", "items": {"type": "string"}, "description": "排除域名列表"}, - "country": {"type": "string", "description": "优先指定国家"} - }, - "required": ["query"] - } - } - })); - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "fetch_page", - "description": "抓取网页并提取正文内容。传入 URL 即可阅读文章正文。可通过 ~/.ias/site_selectors.json 配置按站点用 CSS 选择器精确定位内容区域。", - "parameters": { - "type": "object", - "properties": { - "url": {"type": "string", "description": "网页 URL"} - }, - "required": ["url"] - } - } - })); - + let specs = crate::tools::builtin::BuiltinRegistry::specs(); + for spec in &specs { + list.push(serde_json::json!({ + "type": "function", + "function": { "name": spec.name, "description": spec.description, "parameters": spec.parameters } + })); + } list } - -fn build_env_map() -> HashMap { - let mut map = HashMap::new(); - for var in &[ - "DEEPSEEK_API_KEY", - "DEEPSEEK_MODEL", - "DEEPSEEK_BASE_URL", - "QWEATHER_API_HOST", - "QWEATHER_JWT_KEY_ID", - "QWEATHER_JWT_PROJECT_ID", - "QWEATHER_JWT_PRIVATE_KEY_FILE", - "TAVILY_API_KEY", - ] { - if let Ok(val) = std::env::var(var) { - map.insert(var.to_string(), val); - } - } - map -} diff --git a/src/ipc.rs b/src/ipc.rs index 492bc28..965dc45 100644 --- a/src/ipc.rs +++ b/src/ipc.rs @@ -149,6 +149,7 @@ pub async fn recv_task(stream: &mut UnixStream) -> Result { } /// 接收并解析为 OutputFrame +#[allow(dead_code)] pub async fn recv_output(stream: &mut UnixStream) -> Result { let v = recv_frame(stream).await?; serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e)) diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index 618dc85..852943c 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -4,6 +4,7 @@ use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider}; use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage}; use std::future::Future; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::sync::Mutex; @@ -18,11 +19,26 @@ pub type Summarizer = Arc< dyn Fn(String) -> Pin> + Send>> + Send + Sync, >; +/// `chat_with_tools` 的返回结果 +#[derive(Debug)] +pub struct ChatResult { + pub reply: String, + pub used_tools: bool, + pub usage: Option, + /// 是否有异步工具调用仍在等待结果 + pub has_pending_async: bool, +} + +/// 异步工具标记 —— 工具执行器返回此前缀表示工具已异步入队 +pub const ASYNC_MARKER: &str = "__ASYNC_PENDING__"; + pub struct Conversation { config: ConversationConfig, provider: BoxedProvider, session: Arc>, tool_executor: Option, + /// 是否有异步工具调用等待结果返回 + pending_async_tools: Arc, } impl Conversation { @@ -33,12 +49,18 @@ impl Conversation { provider, session: Arc::new(Mutex::new(ChatSession::new())), tool_executor: None, + pending_async_tools: Arc::new(AtomicBool::new(false)), }) } pub fn set_tool_executor(&mut self, executor: ToolExecutor) { self.tool_executor = Some(executor); } + /// 是否有异步工具调用正在等待结果 + #[allow(dead_code)] + pub fn has_pending_async(&self) -> bool { + self.pending_async_tools.load(Ordering::SeqCst) + } pub fn session(&self) -> Arc> { Arc::clone(&self.session) } @@ -113,10 +135,7 @@ impl Conversation { } /// 对话 + 工具循环(带上下文管理) - pub async fn chat_with_tools( - &self, - user_message: String, - ) -> Result<(String, bool, Option), String> { + pub async fn chat_with_tools(&self, user_message: String) -> Result { // ── 空闲超时检查 ── { let s = self.session.lock().await; @@ -128,7 +147,17 @@ impl Conversation { } self.session.lock().await.add_user(user_message.clone()); + self.run_tool_loop().await + } + /// 在异步工具结果返回后恢复工具循环(不追加用户消息) + pub async fn resume_tool_loop(&self) -> Result { + self.pending_async_tools.store(false, Ordering::SeqCst); + self.run_tool_loop().await + } + + /// 工具循环核心:调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待 + async fn run_tool_loop(&self) -> Result { let mut used_tools = false; let mut last_usage = None; let mut turn_count = 0u32; @@ -185,6 +214,7 @@ impl Conversation { .await .add_assistant_tool_calls(calls.clone()); + let mut has_async = false; for tc in &calls { let result = match &self.tool_executor { Some(exec) => { @@ -195,6 +225,14 @@ impl Conversation { } None => "工具执行器未配置".to_string(), }; + + if result.starts_with(ASYNC_MARKER) { + has_async = true; + tracing::info!("📦 {} → 异步入队", tc.function.name); + // 不添加到 session,等待 ToolResult 异步回喂 + continue; + } + tracing::info!("📦 {} → {:.150}", tc.function.name, result); self.session.lock().await.add_tool_result( &tc.id, @@ -202,6 +240,16 @@ impl Conversation { &result, ); } + + if has_async { + self.pending_async_tools.store(true, Ordering::SeqCst); + return Ok(ChatResult { + reply: full_text, + used_tools, + usage: last_usage, + has_pending_async: true, + }); + } continue; } } @@ -224,7 +272,12 @@ impl Conversation { } } - return Ok((full_text, used_tools, last_usage)); + return Ok(ChatResult { + reply: full_text, + used_tools, + usage: last_usage, + has_pending_async: false, + }); } } } diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 4ee1972..c4abb5c 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -3,5 +3,5 @@ pub mod deepseek; pub mod provider; pub mod types; -pub use conversation::{Conversation, ToolExecutor, DEFAULT_SYSTEM_PROMPT}; +pub use conversation::{ChatResult, Conversation, ToolExecutor, ASYNC_MARKER, DEFAULT_SYSTEM_PROMPT}; pub use types::{ConversationConfig, Usage}; diff --git a/src/main.rs b/src/main.rs index 4947356..90cb96c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +mod channel; mod cli; mod context; mod daemon; @@ -5,6 +6,7 @@ mod db; mod ipc; mod llm; mod logger; +mod queue; mod scheduler; mod state; mod tools; @@ -15,11 +17,11 @@ use clap::Parser; use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand}; use context::MemoryStore; use db::Database; -use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; +use llm::{ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; use std::sync::Arc; use tools::approval::{ApprovalDecision, ApprovalManager}; use tools::types::ExecutionContext; -use tracing::{error, info}; +use tracing::{error, info, warn}; use wechat::client::WeChatClient; #[tokio::main] @@ -105,6 +107,8 @@ async fn main() { cmd_daemon(sock, &database, &file_state, &memory_store).await; } Commands::Worker { sock } => { + warn!("Worker 模式已废弃,请使用 daemon 模式"); + #[allow(deprecated)] if let Err(e) = worker::run(&sock).await { error!("Worker 失败: {}", e); } @@ -665,7 +669,7 @@ async fn generate_reply_with_tools( user_text: String, db: &Option>, ) -> Result { - let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?; + let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv.chat_with_tools(user_text).await?; if let Some(u) = &usage { info!( "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", diff --git a/src/queue/message_queue.rs b/src/queue/message_queue.rs new file mode 100644 index 0000000..e3f8d23 --- /dev/null +++ b/src/queue/message_queue.rs @@ -0,0 +1,552 @@ +//! ## 多渠道消息队列 — MessageQueue +//! +//! # 设计目标 +//! +//! 多用户场景下,保证公平性: +//! - A 发了 3 条消息、B 发了 1 条 → 处理顺序 A1, B1, A2, A3 +//! - 不会因为 A 消息多而饿死 B +//! +//! 每个消息携带 user_id + correlation_id,保证回复能追溯到正确的用户。 +//! +//! # 架构位置 +//! +//! ```text +//! receive_messages → enqueue PipelineMessage +//! ┌─────────────────────────────┐ +//! │ MessageQueue │ +//! │ [User A] [User B] │ ← 每用户独立 VecDeque +//! │ pending: [A, B, A] │ ← 公平轮转索引 +//! └─────────────────────────────┘ +//! QueueRunner::run() → dequeue → 按 kind 路由到 mpsc 通道 +//! ├─ llm_tx (UserMessage, ToolResult) +//! ├─ tool_tx (ToolCall, ApprovalRequest) +//! └─ send_tx (LLMReply) +//! ``` + +use crate::channel::ChannelId; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Notify; +use uuid::Uuid; + +// ─── 管线消息类型 ─── + +/// 消息种类 —— 区分不同消费路由 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum MessageKind { + /// 从外部收到的用户消息 → 交给 LLM Consumer + UserMessage { + text: String, + message_id: String, + context_token: Option, + /// 审批通过后携带的已批准工具名 + approved_tool: Option, + }, + /// LLM 生成的回复文本 → 发给 Send Consumer + LLMReply { + text: String, + context_token: Option, + }, + /// LLM 请求的工具调用 → 发给 Tool Consumer + ToolCall { + session_id: Uuid, + tool_call_id: String, + tool_name: String, + arguments: String, + context_token: Option, + }, + /// 工具执行结果 → 回喂 LLM Consumer + ToolResult { + session_id: Uuid, + tool_call_id: String, + tool_name: String, + result: String, + context_token: Option, + }, + /// 高风险工具 → 需用户审批 + ApprovalRequest { + session_id: Uuid, + tool_name: String, + tool_args: String, + /// 审批后要追加到用户消息中的文本 + approval_prompt: String, + }, + /// 定时任务到期 → 交给 LLM Consumer + ScheduledTask { + text: String, + task_name: String, + }, +} + +/// 管线消息 —— 携带路由元数据的统一消息类型 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PipelineMessage { + /// 来源/目标渠道(平台 + 用户),全程不变 + pub channel: ChannelId, + /// 关联 ID:追踪一次用户请求的完整链路(用户消息→LLM→工具→回喂→发送) + pub correlation_id: Uuid, + /// 消息种类 + pub kind: MessageKind, +} + +impl PipelineMessage { + /// 便捷:渠道中的 user_id + pub fn user_id(&self) -> &str { &self.channel.user_id } + /// 便捷:渠道中的 platform + #[allow(dead_code)] + pub fn platform(&self) -> &str { &self.channel.platform } + /// 快速构造一条用户消息 + #[allow(dead_code)] + pub fn user_message( + user_id: impl Into, + text: impl Into, + message_id: impl Into, + context_token: Option, + ) -> Self { + Self { + channel: ChannelId::wechat(user_id), + correlation_id: Uuid::new_v4(), + kind: MessageKind::UserMessage { + text: text.into(), + message_id: message_id.into(), + context_token, + approved_tool: None, + }, + } + } + + /// 快速构造一条 LLM 回复 + pub fn llm_reply( + channel: ChannelId, + correlation_id: Uuid, + text: impl Into, + context_token: Option, + ) -> Self { + Self { + channel, + correlation_id, + kind: MessageKind::LLMReply { + text: text.into(), + context_token, + }, + } + } + + /// 快速构造一条工具调用 + pub fn tool_call( + channel: ChannelId, + correlation_id: Uuid, + session_id: Uuid, + tool_call_id: impl Into, + tool_name: impl Into, + arguments: impl Into, + context_token: Option, + ) -> Self { + Self { + channel, + correlation_id, + kind: MessageKind::ToolCall { + session_id, + tool_call_id: tool_call_id.into(), + tool_name: tool_name.into(), + arguments: arguments.into(), + context_token, + }, + } + } + + /// 快速构造一条工具结果 + pub fn tool_result( + channel: ChannelId, + correlation_id: Uuid, + session_id: Uuid, + tool_call_id: impl Into, + tool_name: impl Into, + result: impl Into, + context_token: Option, + ) -> Self { + Self { + channel, + correlation_id, + kind: MessageKind::ToolResult { + session_id, + tool_call_id: tool_call_id.into(), + tool_name: tool_name.into(), + result: result.into(), + context_token, + }, + } + } +} + +// ─── 公平轮转队列 ─── + +/// 多渠道公平轮转消息队列 +/// +/// 按 ChannelId(platform + user_id)分队列,轮转出队保证公平。 +/// 出队后由 QueueRunner 按 MessageKind 路由到对应 mpsc 通道。 +pub struct MessageQueue { + /// 每渠道独立队列 + queues: HashMap>, + /// 有待处理消息的渠道(轮转顺序) + pending: VecDeque, + /// 空→非空时触发 Notify 唤醒消费者 + signal: Arc, + /// 定时唤醒间隔(兜底用) + poll_interval: Duration, +} + +impl MessageQueue { + pub fn new() -> Self { + Self { + queues: HashMap::new(), + pending: VecDeque::new(), + signal: Arc::new(Notify::new()), + poll_interval: Duration::from_millis(200), + } + } + + #[allow(dead_code)] + pub fn with_poll_interval(mut self, d: Duration) -> Self { + self.poll_interval = d; + self + } + + /// 入队。如果该渠道之前为空 → 加入 pending 尾部 + 通知 + pub fn enqueue(&mut self, item: PipelineMessage) { + let ch = item.channel.clone(); + let q = self.queues.entry(ch.clone()).or_default(); + let was_empty = q.is_empty(); + q.push_back(item); + if was_empty { + self.pending.push_back(ch); + self.signal.notify_one(); + } + } + + /// 公平轮转出队 + /// + /// 每次调用切换到下一个有消息的渠道: + /// - 从 pending 头部取出一个渠道 + /// - pop 其一条消息 + /// - 该渠道还有消息 → 放回 pending 尾部 + pub fn dequeue(&mut self) -> Option { + while let Some(ch) = self.pending.pop_front() { + if let Some(q) = self.queues.get_mut(&ch) { + if let Some(item) = q.pop_front() { + if !q.is_empty() { + self.pending.push_back(ch); + } else { + self.queues.remove(&ch); + } + return Some(item); + } + } + self.queues.remove(&ch); + } + None + } + + pub fn signal(&self) -> Arc { + Arc::clone(&self.signal) + } + #[allow(dead_code)] + pub fn poll_interval(&self) -> Duration { + self.poll_interval + } + + #[allow(dead_code)] + pub fn has_pending(&self, ch: &ChannelId) -> bool { + self.queues.get(ch).map_or(false, |q| !q.is_empty()) + } + + #[allow(dead_code)] + pub fn pending_count(&self) -> usize { + self.pending.len() + } + + /// 清空指定渠道的全部待处理消息 + #[allow(dead_code)] + pub fn drain_channel(&mut self, ch: &ChannelId) -> Vec { + let drained = self + .queues + .remove(ch) + .map(|mut q| q.drain(..).collect()) + .unwrap_or_default(); + self.pending.retain(|c| c != ch); + drained + } +} + +// ─── 消费者通道标识 ─── + +/// 消息路由目标 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ConsumerTarget { + /// LLM 消费者(处理 UserMessage, ToolResult) + LLM, + /// 工具消费者(处理 ToolCall, ApprovalRequest) + Tool, + /// 发送消费者(处理 LLMReply) + Send, +} + +impl PipelineMessage { + /// 判断这条消息应该路由到哪个消费者 + pub fn target(&self) -> ConsumerTarget { + match self.kind { + MessageKind::UserMessage { .. } + | MessageKind::ToolResult { .. } + | MessageKind::ScheduledTask { .. } => ConsumerTarget::LLM, + MessageKind::ToolCall { .. } | MessageKind::ApprovalRequest { .. } => { + ConsumerTarget::Tool + } + MessageKind::LLMReply { .. } => ConsumerTarget::Send, + } + } +} + +// ─── 测试 ─── + +#[cfg(test)] +mod tests { + use super::*; + use crate::queue::QueueRunner; + + fn make_user_msg(user: &str, text: &str) -> PipelineMessage { + PipelineMessage::user_message(user, text, "", None) + } + + #[test] + fn test_target_routing() { + let msg = PipelineMessage::user_message("u1", "hello", "m1", None); + assert_eq!(msg.target(), ConsumerTarget::LLM); + + let msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "hi", None); + assert_eq!(msg.target(), ConsumerTarget::Send); + + let msg = PipelineMessage::tool_call( + ChannelId::wechat("u1"), + Uuid::new_v4(), + Uuid::new_v4(), + "tc1", + "query_weather", + r#"{"location":"北京"}"#, + None, + ); + assert_eq!(msg.target(), ConsumerTarget::Tool); + } + + #[test] + fn test_fair_round_robin() { + let mut q = MessageQueue::new(); + q.enqueue(make_user_msg("A", "a1")); + q.enqueue(make_user_msg("A", "a2")); + q.enqueue(make_user_msg("A", "a3")); + q.enqueue(make_user_msg("B", "b1")); + assert_eq!(q.dequeue().unwrap().channel.user_id, "A"); + assert_eq!(q.dequeue().unwrap().channel.user_id, "B"); + assert_eq!(q.dequeue().unwrap().channel.user_id, "A"); + assert_eq!(q.dequeue().unwrap().channel.user_id, "A"); + assert!(q.dequeue().is_none()); + } + + #[test] + fn test_single_user() { + let mut q = MessageQueue::new(); + q.enqueue(make_user_msg("X", "hi")); + q.enqueue(make_user_msg("X", "hello")); + assert_eq!(q.dequeue().unwrap().channel.user_id, "X"); + assert_eq!(q.dequeue().unwrap().channel.user_id, "X"); + assert!(q.dequeue().is_none()); + } + + #[test] + fn test_empty_queue() { + let mut q = MessageQueue::new(); + assert!(q.dequeue().is_none()); + assert_eq!(q.pending_count(), 0); + } + + #[test] + fn test_user_id_preserved() { + let msg = PipelineMessage::user_message("wx_user_1", "天气?", "m1", None); + assert_eq!(msg.channel.user_id, "wx_user_1"); + + let reply = PipelineMessage::llm_reply(ChannelId::wechat("wx_user_1"), msg.correlation_id, "晴天", None); + assert_eq!(reply.channel.user_id, "wx_user_1"); + assert_eq!(reply.correlation_id, msg.correlation_id); + } + + #[test] + fn test_mixed_item_types() { + let mut q = MessageQueue::new(); + q.enqueue(PipelineMessage::user_message("U", "user msg", "1", None)); + q.enqueue(PipelineMessage { + channel: ChannelId::wechat("U"), + correlation_id: Uuid::new_v4(), + kind: MessageKind::ScheduledTask { + text: "task result".into(), + task_name: "reminder".into(), + }, + }); + assert!(matches!( + q.dequeue().unwrap().kind, + MessageKind::UserMessage { .. } + )); + assert!(matches!( + q.dequeue().unwrap().kind, + MessageKind::ScheduledTask { .. } + )); + } + + // ─── 端到端关联测试 ─── + + /// 验证 user_id 从用户消息到 LLM 回复全程不变 + #[test] + fn test_user_id_pipeline_preserved() { + let user_id = "wx_user_1"; + + // 用户消息入队 + let user_msg = PipelineMessage::user_message(user_id, "天气怎么样?", "msg1", None); + assert_eq!(user_msg.channel.user_id, user_id); + let correlation_id = user_msg.correlation_id; + + // LLM 回复(模拟 LLM Consumer 生成) + let reply = PipelineMessage::llm_reply(ChannelId::wechat(user_id), correlation_id, "今天晴天", None); + assert_eq!(reply.channel.user_id, user_id); + assert_eq!(reply.correlation_id, correlation_id); + + // 工具调用(模拟 LLM Consumer 生成) + let tool_call = PipelineMessage::tool_call( + ChannelId::wechat(user_id), + correlation_id, + Uuid::new_v4(), + "tc1", + "query_weather", + r#"{"location":"北京"}"#, + None, + ); + assert_eq!(tool_call.channel.user_id, user_id); + assert_eq!(tool_call.correlation_id, correlation_id); + + // 工具结果回喂(模拟 Tool Consumer 生成) + let tool_result = PipelineMessage::tool_result( + ChannelId::wechat(user_id), + correlation_id, + Uuid::new_v4(), + "tc1", + "query_weather", + "25°C", + None, + ); + assert_eq!(tool_result.channel.user_id, user_id); + assert_eq!(tool_result.correlation_id, correlation_id); + } + + /// 验证多个用户的消息不会串 + #[test] + fn test_multi_user_isolation() { + let mut q = MessageQueue::new(); + + q.enqueue(PipelineMessage::user_message("u1", "你好", "1", None)); + q.enqueue(PipelineMessage::user_message("u2", "hello", "2", None)); + q.enqueue(PipelineMessage::user_message("u1", "天气?", "3", None)); + + // 按公平轮转出队,验证 user_id 正确 + let msg1 = q.dequeue().unwrap(); + assert_eq!(msg1.channel.user_id, "u1"); + let msg2 = q.dequeue().unwrap(); + assert_eq!(msg2.channel.user_id, "u2"); + let msg3 = q.dequeue().unwrap(); + assert_eq!(msg3.channel.user_id, "u1"); + } + + /// 验证路由目标正确 + #[test] + fn test_consumer_target_routing() { + // UserMessage → LLM + let m = PipelineMessage::user_message("u1", "hi", "1", None); + assert_eq!(m.target(), ConsumerTarget::LLM); + + // ToolResult → LLM + let m = PipelineMessage::tool_result(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "test", "ok", None); + assert_eq!(m.target(), ConsumerTarget::LLM); + + // ScheduledTask → LLM + let m = PipelineMessage { + channel: ChannelId::wechat("u1"), + correlation_id: Uuid::new_v4(), + kind: MessageKind::ScheduledTask { text: "task".into(), task_name: "t1".into() }, + }; + assert_eq!(m.target(), ConsumerTarget::LLM); + + // ToolCall → Tool + let m = PipelineMessage::tool_call(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "weather", "{}", None); + assert_eq!(m.target(), ConsumerTarget::Tool); + + // ApprovalRequest → Tool + let m = PipelineMessage { + channel: ChannelId::wechat("u1"), + correlation_id: Uuid::new_v4(), + kind: MessageKind::ApprovalRequest { + session_id: Uuid::new_v4(), + tool_name: "high_risk_tool".into(), + tool_args: "{}".into(), + approval_prompt: "".into(), + }, + }; + assert_eq!(m.target(), ConsumerTarget::Tool); + + // LLMReply → Send + let m = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None); + assert_eq!(m.target(), ConsumerTarget::Send); + } + + /// 验证 QueueRunner 的路由逻辑:入队 → 路由到正确通道 + #[tokio::test] + async fn test_queue_runner_routing() { + let (runner, mut channels) = QueueRunner::new(16); + let enqueue = runner.enqueue_handle(); + + // 启动 runner(在后台运行) + let handle = tokio::spawn(async move { + runner.run().await; + }); + + // 入队一个 UserMessage → 应到 llm_rx + let user_msg = PipelineMessage::user_message("u1", "hello", "m1", None); + enqueue.enqueue(user_msg).await; + + // 入队一个 LLMReply → 应到 send_rx + let reply_msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None); + enqueue.enqueue(reply_msg).await; + + // 验证 llm_rx 收到 UserMessage + let llm_msg = tokio::time::timeout( + std::time::Duration::from_secs(2), + channels.llm_rx.recv(), + ) + .await + .expect("llm_rx 超时") + .expect("llm_rx 通道关闭"); + assert_eq!(llm_msg.channel.user_id, "u1"); + assert!(matches!(llm_msg.kind, MessageKind::UserMessage { .. })); + + // 验证 send_rx 收到 LLMReply + let send_msg = tokio::time::timeout( + std::time::Duration::from_secs(2), + channels.send_rx.recv(), + ) + .await + .expect("send_rx 超时") + .expect("send_rx 通道关闭"); + assert_eq!(send_msg.channel.user_id, "u1"); + assert!(matches!(send_msg.kind, MessageKind::LLMReply { .. })); + + // 关闭 runner + handle.abort(); + } +} diff --git a/src/queue/mod.rs b/src/queue/mod.rs new file mode 100644 index 0000000..19860a0 --- /dev/null +++ b/src/queue/mod.rs @@ -0,0 +1,5 @@ +pub mod message_queue; +pub mod runner; + +pub use message_queue::{ConsumerTarget, MessageKind, MessageQueue, PipelineMessage}; +pub use runner::{ConsumerChannels, EnqueueHandle, QueueRunner}; diff --git a/src/queue/runner.rs b/src/queue/runner.rs new file mode 100644 index 0000000..d952540 --- /dev/null +++ b/src/queue/runner.rs @@ -0,0 +1,212 @@ +//! QueueRunner — 消息队列运行时 +//! +//! 从公平轮转 MessageQueue 中出队消息,按 MessageKind 路由到对应的 +//! tokio mpsc 通道,由注册的消费者 task 处理。 +//! +//! # 架构 +//! +//! ```text +//! ┌──────────────┐ +//! enqueue ──────→ │ MessageQueue │ ← 公平轮转 +//! └──────┬───────┘ +//! │ dequeue +//! ▼ +//! ┌──────────────┐ +//! │ QueueRunner │ ← 路由循环 +//! └──┬───┬───┬──┘ +//! ┌─┘ │ └─┐ +//! ▼ ▼ ▼ +//! llm_tx tool_tx send_tx +//! │ │ │ +//! ┌────┘ ┌─┘ ┌──┘ +//! ▼ ▼ ▼ +//! LLMCons. ToolCons. SendCons. +//! ``` + +use crate::queue::{ConsumerTarget, MessageQueue, PipelineMessage}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::sync::{mpsc, Mutex, Notify}; +use tracing::{error, info, warn}; + +/// 消费者通道集 —— 创建 QueueRunner 时获得接收端 +pub struct ConsumerChannels { + pub llm_rx: mpsc::Receiver, + pub tool_rx: mpsc::Receiver, + pub send_rx: mpsc::Receiver, +} + +/// 队列运行时 —— 负责从公平轮转队列出队并按类型路由 +pub struct QueueRunner { + /// 公平轮转队列 + queue: Arc>, + /// LLM 消费者通道 + llm_tx: mpsc::Sender, + /// 工具消费者通道 + tool_tx: mpsc::Sender, + /// 发送消费者通道 + send_tx: mpsc::Sender, + /// 关闭信号 + shutdown: Arc, + /// 队列信号(空→非空通知) + signal: Arc, + /// 轮询间隔 + poll_interval: std::time::Duration, +} + +impl QueueRunner { + /// 创建 QueueRunner 和对应的消费者通道 + /// + /// `channel_size`: 每个 mpsc 通道的缓冲区大小 + pub fn new(channel_size: usize) -> (Self, ConsumerChannels) { + let (llm_tx, llm_rx) = mpsc::channel(channel_size); + let (tool_tx, tool_rx) = mpsc::channel(channel_size); + let (send_tx, send_rx) = mpsc::channel(channel_size); + + let queue = MessageQueue::new(); + let signal = queue.signal(); + let queue = Arc::new(Mutex::new(queue)); + + let runner = Self { + queue, + llm_tx, + tool_tx, + send_tx, + shutdown: Arc::new(AtomicBool::new(false)), + signal, + poll_interval: std::time::Duration::from_millis(200), + }; + + let channels = ConsumerChannels { + llm_rx, + tool_rx, + send_rx, + }; + + (runner, channels) + } + + /// 获取队列的写句柄(用于外部入队) + pub fn enqueue_handle(&self) -> EnqueueHandle { + EnqueueHandle { + queue: self.queue.clone(), + shutdown: self.shutdown.clone(), + } + } + + /// 获取关闭信号句柄 + pub fn shutdown_handle(&self) -> ShutdownHandle { + ShutdownHandle { + shutdown: self.shutdown.clone(), + } + } + + /// 启动路由主循环(阻塞,应在 tokio::spawn 中运行) + pub async fn run(&self) { + info!("QueueRunner 已启动"); + loop { + // 检查关闭信号 + if self.shutdown.load(Ordering::Relaxed) { + info!("QueueRunner 收到关闭信号,停止路由"); + break; + } + + // 出队 + let item = { + let mut q = self.queue.lock().await; + q.dequeue() + }; + + if let Some(msg) = item { + // 预提取字段(避免 borrow + move 冲突) + let target = msg.target(); + let user_id = msg.channel.user_id.clone(); + let correlation_id = msg.correlation_id; + + let result = match target { + ConsumerTarget::LLM => self.llm_tx.send(msg).await, + ConsumerTarget::Tool => self.tool_tx.send(msg).await, + ConsumerTarget::Send => self.send_tx.send(msg).await, + }; + + if let Err(_e) = result { + warn!( + "路由消息失败 (target={:?}, user={}, corr={}): 通道已关闭,跳过", + target, user_id, correlation_id + ); + // 不中断整个路由循环,跳过该消息继续处理 + } + } else { + // 队列为空,等待通知或超时 + tokio::select! { + _ = self.signal.notified() => { + // 有新的消息入队,继续循环 + } + _ = tokio::time::sleep(self.poll_interval) => { + // 超时兜底,继续检查 + } + } + } + } + info!("QueueRunner 已停止"); + } +} + +/// 入队句柄 —— 外部模块通过它向队列投放消息 +#[derive(Clone)] +pub struct EnqueueHandle { + queue: Arc>, + shutdown: Arc, +} + +impl EnqueueHandle { + /// 向队列投放一条消息 + pub async fn enqueue(&self, msg: PipelineMessage) { + if self.shutdown.load(Ordering::Relaxed) { + warn!( + "队列已关闭,丢弃消息: user={}", + msg.user_id() + ); + return; + } + self.queue.lock().await.enqueue(msg); + } +} + +/// 关闭句柄 +#[derive(Clone)] +pub struct ShutdownHandle { + shutdown: Arc, +} + +impl ShutdownHandle { + /// 发送关闭信号 + pub fn shutdown(&self) { + self.shutdown.store(true, Ordering::Relaxed); + } +} + +/// 等所有消费者处理完并关闭后,等待队列排空 +#[allow(dead_code)] +pub async fn drain_and_shutdown( + _enqueue: &EnqueueHandle, + runner_shutdown: &ShutdownHandle, + consumer_handles: Vec>, +) { + info!("开始关闭队列系统..."); + + // 1. 阻止新消息入队(设置关闭标志) + runner_shutdown.shutdown(); + + // 2. 等待排空(给一个合理的时间) + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + + // 3. 等待消费者 task 完成 + for handle in consumer_handles { + if let Err(e) = handle.await { + error!("消费者 task 结束异常: {:?}", e); + } + } + + info!("队列系统已完全关闭"); +} diff --git a/src/tools/definitions.rs b/src/tools/definitions.rs new file mode 100644 index 0000000..83d3f52 --- /dev/null +++ b/src/tools/definitions.rs @@ -0,0 +1,364 @@ +//! JSON 工具定义 — 像 curl + jq 一样描述一个工具 +//! +//! 三要素:URL → 取节点 → 选字段 +//! +//! ```json +//! { +//! "name": "stocks", +//! "description": "股票行情", +//! "url": "https://api.example.com/stocks", +//! "extract": "data.items", +//! "fields": ["name", "price"] +//! } +//! ``` + +use std::collections::HashMap; + +/// HTTP 方法 +#[derive(Debug, Clone, PartialEq)] +pub enum HttpMethod { + Get, + Post, + Put, + Delete, +} + +#[allow(dead_code)] +impl HttpMethod { + pub fn as_str(&self) -> &'static str { + match self { + HttpMethod::Get => "GET", + HttpMethod::Post => "POST", + HttpMethod::Put => "PUT", + HttpMethod::Delete => "DELETE", + } + } +} + +impl Default for HttpMethod { + fn default() -> Self { + HttpMethod::Get + } +} + +impl<'de> serde::Deserialize<'de> for HttpMethod { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + match s.to_uppercase().as_str() { + "GET" => Ok(HttpMethod::Get), + "POST" => Ok(HttpMethod::Post), + "PUT" => Ok(HttpMethod::Put), + "DELETE" => Ok(HttpMethod::Delete), + other => Err(serde::de::Error::custom(format!( + "不支持的 HTTP 方法: {} (支持: GET/POST/PUT/DELETE)", + other + ))), + } + } +} + +/// 一个 JSON 工具定义 +#[derive(Debug, Clone, serde::Deserialize)] +pub struct ToolDef { + /// 工具名称,用作 call_capability 的 name + pub name: String, + /// 工具描述 + pub description: String, + /// URL 模板,如 "https://api.example.com/stock/{{code}}" + /// 支持 {{var}} 和 ${param:name:type} 两种传参形式 + pub url: String, + /// HTTP 方法,默认 GET + #[serde(default)] + pub method: HttpMethod, + /// HTTP 头。值可包含 ${env:KEY}、${cmd:command}、${param:name:type} + #[serde(default, skip_serializing_if = "Option::is_none")] + pub headers: Option>, + /// POST/PUT 的 JSON body 模板。字符串值可包含 {{var}} 和 ${param:name:type} + #[serde(default, skip_serializing_if = "Option::is_none")] + pub body: Option, + /// 显式参数声明(可选) + /// 格式:{ "name": { "type": "string", "description": "..." } } + /// 未声明的参数自动从 ${param:name:type} 和 {{var}} 推导 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parameters: Option, + /// JSON 提取路径(点路径),如 "data.items" + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extract: Option, + /// 筛选字段 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub fields: Option>, + /// 超时秒数,默认 30 + #[serde(default = "default_timeout")] + pub timeout: u64, +} + +const fn default_timeout() -> u64 { + 30 +} + +impl ToolDef { + /// 从文件加载 + pub fn from_file(path: &std::path::Path) -> Result { + let text = std::fs::read_to_string(path) + .map_err(|e| format!("读取 {} 失败: {}", path.display(), e))?; + serde_json::from_str(&text) + .map_err(|e| format!("解析 {} 失败: {}", path.display(), e)) + } + + /// 生成 SkillSpec + /// + /// 参数推导优先级: + /// 1. 显式 parameters 字段(精确控制类型和描述) + /// 2. 从 URL/Body/Headers 中扫描 `${param:name:type}` 推导 + /// 3. 从 {{var}} 推导(全视为 string 类型,兼容旧格式) + pub fn to_spec(&self) -> super::types::SkillSpec { + let params = self.infer_parameters(); + super::types::SkillSpec { + name: self.name.clone(), + description: self.description.clone(), + risk_level: super::types::RiskLevel::Low, + parameters: params, + timeout_secs: self.timeout, + } + } + + /// 推导参数定义 + fn infer_parameters(&self) -> serde_json::Value { + // 1. 显式 parameters 优先 + if let Some(ref p) = self.parameters { + let p = p.clone(); + // 如果已经是完整 schema(含 type/properties),直接返回 + if p.get("type").is_some() { + return p; + } + // 否则认为是属性定义 { name: { type, description }, ... } + // 包装为完整 schema + let required: Vec = p.as_object() + .map(|o| o.keys().cloned().collect()) + .unwrap_or_default(); + return serde_json::json!({ + "type": "object", + "properties": p, + "required": required + }); + } + + // 2. 从模板中扫描 ${param:name:type} 和 {{var}} + let mut param_map: serde_json::Map = serde_json::Map::new(); + + Self::collect_params(&self.url, &mut param_map); + if let Some(ref body) = self.body { + Self::collect_params_from_json(body, &mut param_map); + } + if let Some(ref headers) = self.headers { + for v in headers.values() { + Self::collect_params(v, &mut param_map); + } + } + + let required: Vec = param_map.keys().cloned().collect(); + + serde_json::json!({ + "type": "object", + "properties": param_map, + "required": required + }) + } + + /// 从字符串中收集参数(识别 ${param:name:type} 和 {{var}}) + fn collect_params(s: &str, map: &mut serde_json::Map) { + // 匹配 ${param:name:type} 和 ${param:name:type=default} + let mut pos = 0; + while let Some(start) = s[pos..].find("${param:") { + let rel = pos + start + 8; // skip "${param:" + let rest = &s[rel..]; + if let Some(end) = rest.find('}') { + let segment = &rest[..end]; + let parts: Vec<&str> = segment.split(':').collect(); + let name = parts[0].trim().to_string(); + if name.is_empty() { + pos = rel + end + 1; + continue; + } + if !map.contains_key(&name) { + let param_type = parts.get(1).copied().unwrap_or("string"); + let description = parts.get(2).copied().unwrap_or(name.as_str()); + let entry = serde_json::json!({ + "type": param_type, + "description": description + }); + map.insert(name.clone(), entry); + } + pos = rel + end + 1; + continue; + } + pos = rel; + } + + // 兼容旧格式 {{var}} + Self::extract_vars(s, map); + } + + /// 从 JSON 值中递归收集参数 + fn collect_params_from_json( + v: &serde_json::Value, + map: &mut serde_json::Map, + ) { + match v { + serde_json::Value::String(s) => Self::collect_params(s, map), + serde_json::Value::Object(obj) => { + for val in obj.values() { + Self::collect_params_from_json(val, map); + } + } + serde_json::Value::Array(arr) => { + for item in arr { + Self::collect_params_from_json(item, map); + } + } + _ => {} + } + } + + /// 兼容旧格式 {{var}} — 全视为 string 类型 + fn extract_vars(s: &str, map: &mut serde_json::Map) { + let mut remaining = s; + while let Some(start) = remaining.find("{{") { + remaining = &remaining[start + 2..]; + if let Some(end) = remaining.find("}}") { + let name = remaining[..end].trim().to_string(); + if !name.is_empty() && !name.starts_with('#') && !map.contains_key(&name) { + map.insert( + name.clone(), + serde_json::json!({"type": "string", "description": name}), + ); + } + remaining = &remaining[end + 2..]; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_minimal() { + let json = r#"{ + "name": "ping", + "description": "Ping test", + "url": "https://httpbin.org/get" + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + assert_eq!(def.name, "ping"); + assert_eq!(def.method, HttpMethod::Get); + assert_eq!(def.timeout, 30); + } + + #[test] + fn test_parse_post() { + let json = r#"{ + "name": "translate", + "description": "翻译", + "url": "https://api.example.com/translate?text={{content}}", + "method": "POST", + "headers": { "Authorization": "${env:TR_KEY}" }, + "body": { "q": "{{content}}", "target": "zh" }, + "extract": "data.result", + "fields": ["text"], + "timeout": 60 + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + assert_eq!(def.method, HttpMethod::Post); + assert_eq!(def.timeout, 60); + assert_eq!(def.extract.as_deref(), Some("data.result")); + } + + #[test] + fn test_infer_parameters_from_url() { + let json = r#"{ + "name": "test", + "description": "test", + "url": "https://api.example.com/{{city}}/weather?days={{days}}" + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + let params = def.infer_parameters(); + let props = params["properties"].as_object().unwrap(); + assert!(props.contains_key("city")); + assert!(props.contains_key("days")); + let required = params["required"].as_array().unwrap(); + assert!(required.iter().any(|v| v == "city")); + assert!(required.iter().any(|v| v == "days")); + } + + #[test] + fn test_infer_parameters_from_body() { + let json = r#"{ + "name": "test", + "description": "test", + "url": "https://api.example.com/x", + "body": { "q": "{{text}}", "target": "{{lang}}" } + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + let params = def.infer_parameters(); + let props = params["properties"].as_object().unwrap(); + assert!(props.contains_key("text")); + assert!(props.contains_key("lang")); + } + + #[test] + fn test_method_uppercase_flexibility() { + let json = r#"{"name":"t","description":"t","url":"u","method":"post"}"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + assert_eq!(def.method, HttpMethod::Post); + } + + #[test] + fn test_infer_parameters_from_param_syntax() { + let json = r#"{ + "name": "weather", + "description": "天气查询", + "url": "https://api.example.com/weather/${param:city:string}", + "headers": { "X-Key": "${param:api_key:string}" } + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + let params = def.infer_parameters(); + let props = params["properties"].as_object().unwrap(); + + assert!(props.contains_key("city"), "应有 city 参数"); + assert_eq!(props["city"]["type"], "string", "city 类型应为 string"); + + assert!(props.contains_key("api_key"), "应有 api_key 参数"); + assert_eq!(props["api_key"]["type"], "string", "api_key 类型应为 string"); + + let required = params["required"].as_array().unwrap(); + assert!(required.iter().any(|v| v == "city")); + } + + #[test] + fn test_infer_parameters_explicit_priority() { + let json = r#"{ + "name": "my_tool", + "description": "test", + "url": "https://api.example.com/${param:city:string}", + "parameters": { + "city": { "type": "string", "description": "城市名称" }, + "days": { "type": "integer", "description": "天数" } + } + }"#; + let def: ToolDef = serde_json::from_str(json).unwrap(); + let params = def.infer_parameters(); + + // 显式 parameters 覆盖了自动推导 + let props = params["properties"].as_object().unwrap(); + assert!(props.contains_key("days"), "显式声明的 days 应在"); + assert_eq!(props["days"]["type"], "integer"); + + // 且自动 wrap 为完整 schema + assert_eq!(params["type"], "object"); + } +} diff --git a/src/tools/subprocess.rs b/src/tools/subprocess.rs new file mode 100644 index 0000000..e6eb220 --- /dev/null +++ b/src/tools/subprocess.rs @@ -0,0 +1,934 @@ +//! JSON 工具子进程执行器 +//! +//! 从 ~/.ias/tools/*.json 加载 ToolDef,运行时通过 HTTP 执行。 +//! 无需编译,放 JSON 即用。 +//! +//! 执行流程: +//! 模板替换 → HTTP 请求 → JSON 解析 +//! ├─ 失败 → 返回纯文本 +//! ├─ 提取 → json_select() +//! ├─ 筛选 → filter_fields() +//! └─ 格式化输出 + +use crate::tools::definitions::{HttpMethod, ToolDef}; +use crate::tools::types::{SkillResult, SkillSpec}; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::LazyLock; +use std::sync::Mutex; +use std::sync::RwLock; +use std::time::Instant; + +/// 全局单例子进程执行器 +static RUNNER: LazyLock = LazyLock::new(|| { + let dir = tools_dir(); + SubprocessRunner::load_all(&dir) +}); + +/// 工具目录路径 +fn tools_dir() -> std::path::PathBuf { + // 优先环境变量 + if let Ok(dir) = std::env::var("IAS_TOOLS_DIR") { + return PathBuf::from(dir); + } + // 默认 ~/.ias/tools/ + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ias") + .join("tools") +} + +/// 命令结果缓存 +/// +/// 避免 `${cmd:...}` 每次执行都重跑 shell 命令。 +/// 默认 TTL 300 秒,适合 JWT token 等有时效的值。 +struct CmdCache { + entries: Mutex>, + ttl_secs: u64, +} + +impl CmdCache { + fn new(ttl_secs: u64) -> Self { + Self { + entries: Mutex::new(HashMap::new()), + ttl_secs, + } + } + + fn get_or_exec(&self, cmd: &str) -> Result { + // 检查缓存 + { + let cache = self.entries.lock().unwrap(); + if let Some((val, time)) = cache.get(cmd) { + if time.elapsed().as_secs() < self.ttl_secs { + return Ok(val.clone()); + } + } + } + + // 执行命令 + let output = std::process::Command::new("bash") + .arg("-c") + .arg(cmd) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .output() + .map_err(|e| format!("执行命令失败: {}", e))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + return Err(format!("命令退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)); + } + + let value = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if value.is_empty() { + return Err(format!("命令输出为空: {}", cmd)); + } + + // 写入缓存 + { + let mut cache = self.entries.lock().unwrap(); + cache.insert(cmd.to_string(), (value.clone(), Instant::now())); + } + + Ok(value) + } +} + +static CMD_CACHE: LazyLock = LazyLock::new(|| CmdCache::new(300)); + +/// 获取全局 runner 引用 +pub fn runner() -> &'static SubprocessRunner { + &RUNNER +} + +/// 子进程执行器 +pub struct SubprocessRunner { + tools: RwLock>, +} + +impl SubprocessRunner { + /// 加载目录下所有 *.json 文件 + pub fn load_all(dir: &Path) -> Self { + let mut tools = HashMap::new(); + if dir.exists() { + if let Ok(entries) = std::fs::read_dir(dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().map_or(false, |e| e == "json") { + match ToolDef::from_file(&path) { + Ok(def) => { + tracing::info!("加载工具: {} ← {}", def.name, path.display()); + tools.insert(def.name.clone(), def); + } + Err(e) => { + tracing::warn!("跳过 {}: {}", path.display(), e); + } + } + } + } + } + } else { + tracing::info!("工具目录不存在,创建: {}", dir.display()); + let _ = std::fs::create_dir_all(dir); + } + Self { + tools: RwLock::new(tools), + } + } + + /// 重载全部工具(运行时热加载) + #[allow(dead_code)] + pub fn reload(&self) { + let dir = tools_dir(); + let mut new_tools = HashMap::new(); + if dir.exists() { + if let Ok(entries) = std::fs::read_dir(&dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().map_or(false, |e| e == "json") { + if let Ok(def) = ToolDef::from_file(&path) { + new_tools.insert(def.name.clone(), def); + } + } + } + } + } + if let Ok(mut w) = self.tools.write() { + *w = new_tools; + } + } + + /// 是否有该工具 + pub fn has_tool(&self, name: &str) -> bool { + self.tools.read().map_or(false, |m| m.contains_key(name)) + } + + /// 获取工具 spec(参数从 URL/Body 模板推导) + #[allow(dead_code)] + pub fn spec(&self, name: &str) -> Option { + self.tools.read().ok().and_then(|m| m.get(name).map(|d| d.to_spec())) + } + + /// 获取所有工具 specs + pub fn all_specs(&self) -> Vec { + self.tools + .read() + .map(|m| m.values().map(|d| d.to_spec()).collect()) + .unwrap_or_default() + } + + /// 执行工具 + pub async fn execute(&self, name: &str, args_json: &str) -> Option { + let def = self.tools.read().ok()?.get(name)?.clone(); + let args: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default(); + Some(self.execute_def(&def, &args).await) + } + + /// 执行一个 ToolDef + async fn execute_def(&self, def: &ToolDef, args: &serde_json::Value) -> SkillResult { + // 1. 模板替换 + let url = render_url(&def.url, args); + let headers = resolve_headers_with_args(def.headers.as_ref(), args); + let body = def.body.as_ref().map(|b| render_body(b, args)); + + // 2. HTTP 请求 + let resp_text = match http_request(def.method.clone(), &url, &headers, body.as_deref()).await + { + Ok(t) => t, + Err(e) => return SkillResult::error(format!("请求失败: {}", e)), + }; + + // 3. JSON 判定 + let json: serde_json::Value = match serde_json::from_str(&resp_text) { + Ok(j) => j, + Err(_) => { + // 纯文本,直接返回全文 + return SkillResult::ok(resp_text); + } + }; + + // 4. 提取子节点 + let extracted = match &def.extract { + Some(path) => match json_select(&json, path) { + Some(v) => v.clone(), + None => { + return SkillResult::error(format!("路径 '{}' 未找到", path)); + } + }, + None => json, + }; + + // 5. 筛选字段 + let filtered = match &def.fields { + Some(fields) => filter_fields(&extracted, fields), + None => extracted, + }; + + // 6. 格式化输出 + let output = format_output(&filtered); + + SkillResult::ok_with_data(output, filtered) + } +} + +// ═══════════════════════════════════════════════ +// 模板渲染 +// ═══════════════════════════════════════════════ + +/// 渲染 URL 模板:{{var}} / ${param:name:type} → args.name(URL 编码) +fn render_url(template: &str, args: &serde_json::Value) -> String { + // 先处理 ${param:name:type} + let after_param = replace_param_refs(template, args, true); + // 再处理 {{var}} + replace_mustache(&after_param, args) +} + +/// 渲染普通文本(非 URL 编码),支持三种引用: +/// ${param:name:type} — 参数值 +/// ${env:KEY} — 环境变量 +/// ${cmd:command} — shell 命令输出 +fn render_text(template: &str, args: &serde_json::Value) -> String { + // 先处理 ${param:name:type} + let after_param = replace_param_refs(template, args, false); + // 再处理 {{var}}(不编码) + replace_mustache_plain(&after_param, args) +} + +/// 替换 ${param:name:type} 引用 +fn replace_param_refs(template: &str, args: &serde_json::Value, url_encode: bool) -> String { + let mut result = String::new(); + let mut remaining = template; + + while let Some(start) = remaining.find("${param:") { + result.push_str(&remaining[..start]); + let rest = &remaining[start + 8..]; + if let Some(end) = rest.find('}') { + let segment = &rest[..end]; + let name = segment.split(':').next().map(|s| s.trim()).unwrap_or(""); + if !name.is_empty() { + let value = args + .get(name) + .and_then(|v| v.as_str()) + .unwrap_or(""); + if url_encode { + result.push_str(&urlencoding(value)); + } else { + result.push_str(value); + } + } + remaining = &rest[end + 1..]; + } else { + result.push_str(&remaining[start..]); + remaining = ""; + } + } + result.push_str(remaining); + result +} + +/// 替换 {{var}} — URL 编码版 +fn replace_mustache(template: &str, args: &serde_json::Value) -> String { + let mut result = String::new(); + let mut remaining = template; + while let Some(start) = remaining.find("{{") { + result.push_str(&remaining[..start]); + remaining = &remaining[start + 2..]; + if let Some(end) = remaining.find("}}") { + let key = remaining[..end].trim(); + if !key.starts_with('#') { + let value = args + .get(key) + .and_then(|v| v.as_str()) + .map(|s| urlencoding(s)) + .unwrap_or_else(|| key.to_string()); + result.push_str(&value); + } + remaining = &remaining[end + 2..]; + } else { + result.push_str(&remaining[..]); + remaining = ""; + } + } + result.push_str(remaining); + result +} + +/// 替换 {{var}} — 纯文本版(不编码) +fn replace_mustache_plain(template: &str, args: &serde_json::Value) -> String { + let mut result = String::new(); + let mut remaining = template; + while let Some(start) = remaining.find("{{") { + result.push_str(&remaining[..start]); + remaining = &remaining[start + 2..]; + if let Some(end) = remaining.find("}}") { + let key = remaining[..end].trim(); + if !key.starts_with('#') { + let value = args + .get(key) + .and_then(|v| v.as_str()) + .unwrap_or(key); + result.push_str(value); + } + remaining = &remaining[end + 2..]; + } else { + result.push_str(&remaining[..]); + remaining = ""; + } + } + result.push_str(remaining); + result +} + +/// URL 编码 +fn urlencoding(s: &str) -> String { + // 只编码需要编码的字符 + let mut result = String::with_capacity(s.len()); + for c in s.chars() { + match c { + 'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => result.push(c), + _ => { + for byte in c.to_string().as_bytes() { + result.push_str(&format!("%{:02X}", byte)); + } + } + } + } + result +} + +/// 解析模板值(无 args 上下文,主要用于测试) +#[allow(dead_code)] +fn resolve_template_value(template: &str) -> String { + if let Some(env_key) = template.strip_prefix("${env:").and_then(|s| s.strip_suffix('}')) { + return std::env::var(env_key).unwrap_or_default(); + } + if let Some(cmd) = template.strip_prefix("${cmd:").and_then(|s| s.strip_suffix('}')) { + return CMD_CACHE.get_or_exec(cmd).unwrap_or_else(|e| { + tracing::warn!("执行命令失败: {}", e); + String::new() + }); + } + // ${param:name:type} — 在 resolve_headers 中通过 args 解析, + // 单纯 resolve_template_value 没有 args 上下文的,先保留原样 + if template.contains("${param:") { + return template.to_string(); + } + template.to_string() +} + +/// 渲染 header 值(带 args 上下文,支持 ${param:name:type}) +fn resolve_header_value(template: &str, args: &serde_json::Value) -> String { + // 优先 env/cmd + if let Some(env_key) = template.strip_prefix("${env:").and_then(|s| s.strip_suffix('}')) { + return std::env::var(env_key).unwrap_or_default(); + } + if let Some(cmd) = template.strip_prefix("${cmd:").and_then(|s| s.strip_suffix('}')) { + return CMD_CACHE.get_or_exec(cmd).unwrap_or_else(|e| { + tracing::warn!("执行命令失败: {}", e); + String::new() + }); + } + // ${param:name:type} → 取 args.name + if let Some(param_rest) = template.strip_prefix("${param:") { + if let Some(end) = param_rest.find('}') { + let segment = ¶m_rest[..end]; + let name = segment.split(':').next().map(|s| s.trim()).unwrap_or(""); + if !name.is_empty() { + return args + .get(name) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + } + } + } + render_text(template, args) +} + +/// 渲染 header 列表(带 args) +fn resolve_headers_with_args( + headers: Option<&HashMap>, + args: &serde_json::Value, +) -> Vec<(String, String)> { + headers + .map(|h| { + h.iter() + .map(|(k, v)| (k.clone(), resolve_header_value(v, args))) + .collect() + }) + .unwrap_or_default() +} + +/// 渲染 header 值(无 args 上下文,主要用于测试) +#[allow(dead_code)] +fn resolve_headers( + headers: Option<&HashMap>, +) -> Vec<(String, String)> { + headers + .map(|h| { + h.iter() + .map(|(k, v)| (k.clone(), resolve_template_value(v))) + .collect() + }) + .unwrap_or_default() +} + +/// 渲染 JSON body 模板:字符串中的 {{var}} 和 ${param:name:type} 被替换 +fn render_body(template: &serde_json::Value, args: &serde_json::Value) -> String { + fn render_value(v: &serde_json::Value, args: &serde_json::Value) -> serde_json::Value { + match v { + serde_json::Value::String(s) => { + if s.contains("{{") || s.contains("${param:") { + let rendered = render_text(s, args); + // 尝试解析为 JSON + if let Ok(j) = serde_json::from_str::(&rendered) { + return j; + } + return serde_json::Value::String(rendered); + } + v.clone() + } + serde_json::Value::Object(obj) => { + let mut m = serde_json::Map::new(); + for (k, val) in obj { + m.insert(k.clone(), render_value(val, args)); + } + serde_json::Value::Object(m) + } + serde_json::Value::Array(arr) => { + let items: Vec<_> = arr.iter().map(|item| render_value(item, args)).collect(); + serde_json::Value::Array(items) + } + _ => v.clone(), + } + } + + let rendered = render_value(template, args); + serde_json::to_string(&rendered).unwrap_or_default() +} + +// ═══════════════════════════════════════════════ +// HTTP 请求 +// ═══════════════════════════════════════════════ + +async fn http_request( + method: HttpMethod, + url: &str, + headers: &[(String, String)], + body: Option<&str>, +) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .user_agent("iAs/1.0 (JSON Tool Runner)") + .build() + .map_err(|e| format!("创建 HTTP client 失败: {}", e))?; + + let mut req = match method { + HttpMethod::Get => client.get(url), + HttpMethod::Post => { + let mut r = client.post(url); + if let Some(b) = body { + r = r.header("Content-Type", "application/json").body(b.to_string()); + } + r + } + HttpMethod::Put => { + let mut r = client.put(url); + if let Some(b) = body { + r = r.header("Content-Type", "application/json").body(b.to_string()); + } + r + } + HttpMethod::Delete => client.delete(url), + }; + + for (k, v) in headers { + if !v.is_empty() { + req = req.header(k.as_str(), v.as_str()); + } + } + + let resp = req.send().await.map_err(|e| format!("HTTP 请求失败: {}", e))?; + + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + return Err(format!("HTTP {}: {}", status, text.chars().take(200).collect::())); + } + + resp.text().await.map_err(|e| format!("读取响应失败: {}", e)) +} + +// ═══════════════════════════════════════════════ +// JSON 节点选择(点路径) +// ═══════════════════════════════════════════════ + +/// 按点路径选择 JSON 节点 +/// +/// 支持格式: +/// "data" → json["data"] +/// "data.items" → json["data"]["items"] +/// "items.0.name" → json["items"][0]["name"] +/// ".data.items" → 自动忽略前导点 +fn json_select<'a>(json: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> { + let path = path.strip_prefix('.').unwrap_or(path); + if path.is_empty() { + return Some(json); + } + + let mut current = json; + for segment in path.split('.') { + if segment.is_empty() { + continue; + } + current = match current { + serde_json::Value::Object(obj) => obj.get(segment)?, + serde_json::Value::Array(arr) => { + let idx: usize = segment.parse().ok()?; + arr.get(idx)? + } + _ => return None, + }; + } + Some(current) +} + +// ═══════════════════════════════════════════════ +// 字段筛选 +// ═══════════════════════════════════════════════ + +/// 只保留指定字段 +/// +/// - object: 只保留 fields 中的 key +/// - array: 对每个元素做同样筛选 +/// - 其他: 保留原值 +fn filter_fields(value: &serde_json::Value, fields: &[String]) -> serde_json::Value { + match value { + serde_json::Value::Object(obj) => { + let mut m = serde_json::Map::new(); + for field in fields { + if let Some(v) = obj.get(field) { + m.insert(field.clone(), v.clone()); + } + } + serde_json::Value::Object(m) + } + serde_json::Value::Array(arr) => { + let items: Vec<_> = arr.iter().map(|item| filter_fields(item, fields)).collect(); + serde_json::Value::Array(items) + } + other => other.clone(), + } +} + +// ═══════════════════════════════════════════════ +// 输出格式化 +// ═══════════════════════════════════════════════ + +/// 将提取/筛选后的结果格式化为可读文本 +fn format_output(value: &serde_json::Value) -> String { + match value { + serde_json::Value::Null => "null".to_string(), + serde_json::Value::Bool(b) => b.to_string(), + serde_json::Value::Number(n) => n.to_string(), + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Array(arr) => format_array(arr), + serde_json::Value::Object(obj) => format_object(obj), + } +} + +fn format_array(arr: &[serde_json::Value]) -> String { + if arr.is_empty() { + return "(空)".to_string(); + } + + let lines: Vec = arr + .iter() + .enumerate() + .map(|(i, item)| { + let prefix = format!("{}. ", i + 1); + let body = match item { + serde_json::Value::Object(_) => { + // 对象数组:每个对象展平成 key: value 行 + let obj_lines = format_object_lines(item) + .into_iter() + .map(|l| format!(" {}", l)) + .collect::>() + .join("\n"); + obj_lines + } + serde_json::Value::Array(a) => format_array(a), + other => format_output(other), + }; + format!("{}{}", prefix, body) + }) + .collect(); + + lines.join("\n---\n") +} + +fn format_object(obj: &serde_json::Map) -> String { + let lines = format_object_lines(&serde_json::Value::Object(obj.clone())); + lines.join("\n") +} + +fn format_object_lines(value: &serde_json::Value) -> Vec { + match value { + serde_json::Value::Object(obj) => { + let mut lines = Vec::new(); + for (k, v) in obj { + let v_str = match v { + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Null => "null".to_string(), + serde_json::Value::Bool(b) => b.to_string(), + serde_json::Value::Number(n) => n.to_string(), + serde_json::Value::Array(arr) => { + if arr.len() <= 5 { + arr.iter() + .map(format_output) + .collect::>() + .join(", ") + } else { + format!("[{} 项]", arr.len()) + } + } + serde_json::Value::Object(o) => { + format!("{{ {} }}", o.keys().cloned().collect::>().join(", ")) + } + }; + lines.push(format!(" {}: {}", k, v_str)); + } + lines + } + _ => vec![format_output(value)], + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ─── 模板渲染 ─── + + #[test] + fn test_render_url_basic() { + let args = serde_json::json!({"code": "000001"}); + let url = render_url("https://api.example.com/stock/{{code}}", &args); + assert_eq!(url, "https://api.example.com/stock/000001"); + } + + #[test] + fn test_render_url_encode_chinese() { + let args = serde_json::json!({"city": "北京"}); + let url = render_url("https://api.example.com/weather?city={{city}}", &args); + assert!(url.contains("%E5%8C%97%E4%BA%AC")); + } + + #[test] + fn test_render_url_no_vars() { + let args = serde_json::json!({}); + let url = render_url("https://api.example.com/ping", &args); + assert_eq!(url, "https://api.example.com/ping"); + } + + // ─── JSON 选择 ─── + + #[test] + fn test_json_select_simple() { + let json = serde_json::json!({"data": {"price": 42.5}}); + assert_eq!(json_select(&json, "data.price").and_then(|v| v.as_f64()), Some(42.5)); + } + + #[test] + fn test_json_select_dot_prefix() { + let json = serde_json::json!({"data": {"price": 42.5}}); + assert_eq!(json_select(&json, ".data.price").and_then(|v| v.as_f64()), Some(42.5)); + } + + #[test] + fn test_json_select_array_index() { + let json = serde_json::json!({"items": [{"name": "a"}, {"name": "b"}]}); + assert_eq!( + json_select(&json, "items.1.name").and_then(|v| v.as_str()), + Some("b") + ); + } + + #[test] + fn test_json_select_not_found() { + let json = serde_json::json!({"a": 1}); + assert!(json_select(&json, "a.b.c").is_none()); + } + + // ─── 字段筛选 ─── + + #[test] + fn test_filter_fields_object() { + let obj = serde_json::json!({"name": "AAPL", "price": 150, "ignored": true}); + let fields = vec!["name".into(), "price".into()]; + let filtered = filter_fields(&obj, &fields); + assert_eq!(filtered.as_object().unwrap().len(), 2); + assert!(filtered.get("ignored").is_none()); + } + + #[test] + fn test_filter_fields_array() { + let arr = serde_json::json!([ + {"name": "A", "price": 10, "extra": true}, + {"name": "B", "price": 20, "extra": false} + ]); + let fields = vec!["name".into()]; + let filtered = filter_fields(&arr, &fields); + let items = filtered.as_array().unwrap(); + assert_eq!(items.len(), 2); + assert!(items[0].get("price").is_none()); + assert!(items[0].get("name").is_some()); + } + + // ─── 输出格式化 ─── + + #[test] + fn test_format_scalar() { + assert_eq!(format_output(&serde_json::json!(42)), "42"); + assert_eq!(format_output(&serde_json::json!("hello")), "hello"); + assert_eq!(format_output(&serde_json::json!(true)), "true"); + } + + #[test] + fn test_format_object() { + let obj = serde_json::json!({"name": "AAPL", "price": 150.5}); + let out = format_output(&obj); + assert!(out.contains("name: AAPL")); + assert!(out.contains("price: 150.5")); + } + + #[test] + fn test_format_array_of_objects() { + let arr = serde_json::json!([ + {"name": "A", "price": 10}, + {"name": "B", "price": 20} + ]); + let out = format_output(&arr); + assert!(out.starts_with("1.")); + assert!(out.contains("---")); + assert!(out.contains("2.")); + } + + // ─── 头部解析 ─── + + #[test] + fn test_resolve_headers_static() { + let mut h = HashMap::new(); + h.insert("Accept".into(), "application/json".into()); + let resolved = resolve_headers(Some(&h)); + assert_eq!(resolved[0].1, "application/json"); + } + + #[test] + fn test_resolve_cmd_inline() { + let result = resolve_template_value("${cmd:echo hello-world}"); + assert_eq!(result, "hello-world"); + } + + #[test] + fn test_resolve_cmd_cached() { + // 同一命令应缓存,第二次直接返回 + let r1 = resolve_template_value("${cmd:echo cached-test-123}"); + let r2 = resolve_template_value("${cmd:echo cached-test-123}"); + assert_eq!(r1, "cached-test-123"); + assert_eq!(r2, r1); + } + + #[test] + fn test_resolve_env() { + unsafe { std::env::set_var("TEST_API_KEY", "secret123"); } + let mut h = HashMap::new(); + h.insert("X-Key".into(), "${env:TEST_API_KEY}".into()); + let resolved = resolve_headers(Some(&h)); + assert_eq!(resolved[0].1, "secret123"); + } + + // ─── 新格式 ${param:name:type} 渲染 ─── + + #[test] + fn test_render_url_with_param() { + let args = serde_json::json!({"city": "北京", "days": "3"}); + let url = render_url( + "https://api.example.com/weather/${param:city:string}/days/${param:days:integer}", + &args, + ); + assert!(url.contains("%E5%8C%97%E4%BA%AC")); + assert!(url.contains("3")); + } + + #[test] + fn test_render_url_mixed() { + let args = serde_json::json!({"code": "000001"}); + let url = render_url( + "https://api.example.com/${param:code:string}?type={{#comment}}", + &args, + ); + // {{#comment}} 被跳过(注释语法) + assert_eq!(url, "https://api.example.com/000001?type="); + } + + #[test] + fn test_render_text_param() { + let args = serde_json::json!({"name": "hello"}); + let result = render_text("prefix_${param:name:string}_suffix", &args); + assert_eq!(result, "prefix_hello_suffix"); + } + + #[test] + fn test_resolve_header_param() { + let args = serde_json::json!({"token": "abc123"}); + let mut h = HashMap::new(); + h.insert("Authorization".into(), "Bearer ${param:token:string}".into()); + let resolved = resolve_headers_with_args(Some(&h), &args); + assert_eq!(resolved[0].1, "Bearer abc123"); + } + + #[test] + fn test_param_in_body() { + let args = serde_json::json!({"q": "你好", "lang": "zh"}); + let body = serde_json::json!({ + "query": "${param:q:string}", + "language": "${param:lang:string}" + }); + let rendered = render_body(&body, &args); + assert!(rendered.contains("你好")); + assert!(rendered.contains("zh")); + } + + #[test] + fn test_load_tools_from_disk() { + let dir = std::path::Path::new("/home/xiao/.ias/tools"); + if !dir.exists() { + eprintln!("跳过磁盘加载测试: 目录不存在"); + return; + } + let runner = SubprocessRunner::load_all(dir); + assert!(runner.has_tool("weather"), "应加载 weather 工具"); + assert!(runner.has_tool("github_repo"), "应加载 github_repo 工具"); + + let specs = runner.all_specs(); + assert!(specs.len() >= 2, "至少加载 2 个工具, 实际: {}", specs.len()); + + // 验证 weather 工具的参数从模板推导 + let weather_spec = runner.spec("weather").unwrap(); + let params = &weather_spec.parameters; + let props = params["properties"].as_object().unwrap(); + assert!(props.contains_key("city"), "weather 应有 city 参数"); + } + + // ─── 完整端到端 ─── + + #[tokio::test] + async fn test_execute_json_weather() { + let dir = std::path::Path::new("/home/xiao/.ias/tools"); + if !dir.exists() { + eprintln!("跳过: 工具目录不存在"); + return; + } + let runner = SubprocessRunner::load_all(dir); + let result = runner.execute("weather", r#"{"city":"北京"}"#).await.unwrap(); + assert!(result.success, "天气查询应成功: {}", result.output); + println!("🌤️ 天气结果: {}", result.output); + } + + #[tokio::test] + async fn test_execute_json_github() { + let dir = std::path::Path::new("/home/xiao/.ias/tools"); + if !dir.exists() { + eprintln!("跳过: 工具目录不存在"); + return; + } + let runner = SubprocessRunner::load_all(dir); + let result = runner.execute("github_repo", r#"{"owner":"rust-lang","repo":"rust"}"#).await.unwrap(); + assert!(result.success, "GitHub 查询应成功: {}", result.output); + assert!(result.output.contains("rust"), "应有仓库名: {}", result.output); + println!("🐙 GitHub: {}", result.output); + } + + #[tokio::test] + async fn test_execute_httpbin_get() { + let def = ToolDef { + name: "httpbin".into(), + description: "test".into(), + url: "https://httpbin.org/get".into(), + method: HttpMethod::Get, + headers: None, + body: None, + parameters: None, + extract: Some("url".into()), + fields: None, + timeout: 30, + }; + let args = serde_json::json!({}); + let runner = SubprocessRunner::load_all(Path::new("/nonexistent")); + let result = runner.execute_def(&def, &args).await; + assert!(result.success, "HTTP 请求应成功: {}", result.output); + assert!(result.output.contains("httpbin.org/get"), "应提取 url 字段: {}", result.output); + } +} diff --git a/src/worker.rs b/src/worker.rs index f7b4059..2ae8410 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -1,6 +1,9 @@ -//! Worker — 无状态执行进程 +#![deprecated = "Worker 进程模式已废弃,请使用 daemon 模式(统一 tokio task 架构)"] + +//! Worker — 无状态执行进程(已废弃) //! -//! 每条消息 spawn 一个 worker: +//! 此模块保留供编译参考,不再被 daemon 使用。 +//! 新架构使用 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。 //! 1. 连接 daemon 的 Unix Domain Socket //! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量) //! 3. 注入环境变量,构建 Conversation,执行 LLM 对话 + 工具调用 @@ -103,7 +106,8 @@ pub async fn run(sock_path: &str) -> Result<(), String> { conv.set_tool_executor(executor); // 7. 执行 LLM 对话(带审批检测) - let (reply, _used_tools, usage) = conv + use crate::llm::ChatResult; + let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv .chat_with_tools(task.msg.text.clone()) .await .map_err(|e| format!("LLM 对话失败: {}", e))?;