From af0620aa2d9a465e59d0731fd673b3078238a1d5 Mon Sep 17 00:00:00 2001 From: wunianxiao Date: Thu, 4 Jun 2026 16:18:06 +0800 Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E6=9E=84=EF=BC=9A=E6=9E=B6=E6=9E=84?= =?UTF-8?q?=E9=87=8D=E7=BB=84=20+=206=E8=BD=AE=E7=BC=BA=E9=99=B7=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1: 工具模块统一 - 新建 tools/executor.rs — 统一工具执行器，消除 main/worker/daemon 三处重复 - tools/builtin.rs → tools/registry.rs - tools/mod.rs 新增 build_tools_list() / build_env_map()，daemon/main 共用 Phase 2: 主模块重组 context/ → core/ - core/pipeline.rs — 消息处理流水线 (从 main.rs 抽取 listen_loop) - core/session.rs — ChatSession 会话管理 - core/context.rs — 上下文构建 + token 估算 - core/summary.rs — 摘要管理 (溢出/空闲触发) - core/memory.rs — MemoryStore 长期记忆 - main.rs 精简至 ~200 行 Phase 3: listen 模式全局会话锁 + Pipeline 收尾 - PipelineContext 持有全局 conversation_lock，所有消息串行处理防串话 Phase 4: wechat/ → channel/ 重命名 缺陷修复: - daemon send_frame 失败后消息重回 waiting 队列 - daemon 审批通过后按 code_hash 精确执行，不重跑 LLM - listen 模式执行器从 session.current_user_id 动态获取用户 - listen 模式 LLM 用量统计传入实际 user_id - send_text 响应校验改用 serde_json::Value 宽松解析 - pi_subagent 工具 (JSON 事件流模式, RiskLevel::High, timeout clamp 1-300) - listen 审批超时 300s→60s，cancel_by_code 精确保护 - approval.rs 新增 cancel_by_code 方法同步 DB --- README.md | 49 ++- src/{wechat => channel}/client.rs | 19 +- src/{wechat => channel}/mod.rs | 0 src/{wechat => channel}/types.rs | 0 src/context/mod.rs | 5 - src/core/context.rs | 112 +++++ src/{context/tools.rs => core/memory.rs} | 6 +- src/core/mod.rs | 7 + src/core/pipeline.rs | 285 +++++++++++++ src/{context/types.rs => core/session.rs} | 0 src/{context/builder.rs => core/summary.rs} | 121 +----- src/daemon.rs | 218 ++-------- src/ipc.rs | 6 + src/llm/conversation.rs | 29 +- src/llm/mod.rs | 2 +- src/main.rs | 436 ++------------------ src/state.rs | 1 + src/tools/approval.rs | 26 ++ src/tools/builtins/mod.rs | 1 + src/tools/builtins/subagent.rs | 229 ++++++++++ src/tools/executor.rs | 336 +++++++++++++++ src/tools/mod.rs | 187 ++++++++- src/tools/{builtin.rs => registry.rs} | 7 + src/tools/types.rs | 2 + src/worker.rs | 211 +++------- 25 files changed, 1405 insertions(+), 890 deletions(-) rename src/{wechat => channel}/client.rs (96%) rename src/{wechat => channel}/mod.rs (100%) rename src/{wechat => channel}/types.rs (100%) delete mode 100644 src/context/mod.rs create mode 100644 src/core/context.rs rename src/{context/tools.rs => core/memory.rs} (94%) create mode 100644 src/core/mod.rs create mode 100644 src/core/pipeline.rs rename src/{context/types.rs => core/session.rs} (100%) rename src/{context/builder.rs => core/summary.rs} (59%) create mode 100644 src/tools/builtins/subagent.rs create mode 100644 src/tools/executor.rs rename src/tools/{builtin.rs => registry.rs} (87%) diff --git a/README.md b/README.md index ad76670..94bd853 100644 --- a/README.md +++ b/README.md @@ -64,15 +64,15 @@ ias service │ main.rs → clap 子命令路由 │ ├─────────────────────────────────────────────────┤ │ ┌──────────┐ ┌──────────┐ ┌───────────────┐ │ -│ │ WeChat │ │ LLM │ │ Tools │ │ -│ │ 长轮询 │ │ DeepSeek │ │ 内置工具 │ │ -│ │ HTTP API │ │ 流式对话 │ │ 审批管理 │ │ +│ │ Channel │ │ LLM │ │ Tools │ │ +│ │ 微信通道 │ │ DeepSeek │ │ 统一执行器 │ │ +│ │ 长轮询 │ │ 流式对话 │ │ 审批管理 │ │ │ └────┬─────┘ └────┬─────┘ └──────┬────────┘ │ │ │ │ │ │ │ ┌────┴──────────────┴───────────────┴────────┐ │ -│ │ Context / Memory │ │ -│ │ ChatSession · Token Budget · 摘要压缩 │ │ -│ │ MemoryStore · 长期记忆 │ │ +│ │ Core / 主模块 │ │ +│ │ Pipeline · Session · Context · Summary │ │ +│ │ Memory · 长期记忆 │ │ │ └──────────────────────┬──────────────────────┘ │ │ │ │ │ ┌──────────────────────┴──────────────────────┐ │ @@ -109,26 +109,31 @@ ias service ``` src/ -├── main.rs CLI 入口 + 命令实现 + 监听循环 + 工具执行器 +├── main.rs CLI 入口 + 命令实现 (~200行) ├── cli.rs clap 子命令定义（login/listen/send/tool 等） ├── logger.rs 日志初始化（终端 + 文件滚动） ├── state.rs 本地文件状态（auth.json / runtime.json，无DB时降级） ├── scheduler.rs 定时任务调度器 -├── context/ 上下文管理 -│ ├── types.rs ChatSession（checkpoint/messages/summaries） -│ ├── builder.rs Token budget 估算 + 上下文构建 + 摘要触发 -│ └── tools.rs MemoryStore（HashMap>） -├── db/ 数据库 -│ ├── mod.rs Database 连接池 + 迁移 -│ └── models.rs CRUD（auth/chat/usage/summary） -├── llm/ 对话系统 -│ ├── types.rs Message / Role / ToolCall / StreamChunk / Usage +├── core/ 主模块 — 消息处理流水线 +│ ├── mod.rs 模块声明 +│ ├── pipeline.rs 消息处理流水线（接收→上下文→LLM→持久化） +│ ├── session.rs 会话管理（ChatSession, checkpoint, token budget） +│ ├── context.rs 上下文构建（build_context, token 估算） +│ ├── summary.rs 摘要管理（溢出触发, 空闲触发, LLM摘要生成） +│ └── memory.rs 长期记忆（MemoryStore, read_summaries） +├── channel/ 微信通道 +│ ├── mod.rs +│ ├── types.rs iLink API 协议类型 +│ └── client.rs HTTP 客户端（login/poll/send/notify） +├── llm/ LLM 对话系统 +│ ├── types.rs Message / Role / ToolCall / StreamChunk │ ├── provider.rs LlmProvider trait + SSE 解析 │ ├── deepseek.rs DeepSeek 流式客户端 │ └── conversation.rs Conversation（chat + chat_with_tools 工具循环） -├── tools/ 工具系统 -│ ├── types.rs RiskLevel / SkillSpec / SkillResult / ExecutionContext -│ ├── builtin.rs BuiltinRegistry 路由 +├── tools/ 工具模块 +│ ├── types.rs SkillSpec / SkillResult / RiskLevel +│ ├── registry.rs BuiltinRegistry 路由器 +│ ├── executor.rs 统一工具执行器（Listen/Worker 双模式） │ ├── approval.rs ApprovalManager（确认码 + oneshot 通道） │ └── builtins/ 内置工具实现 │ ├── datetime.rs 日期时间 @@ -137,9 +142,9 @@ src/ │ ├── fetch_page.rs 网页抓取 │ ├── memos.rs 备忘录 │ └── amap.rs 高德地图（POI/地理编码/路径/旅游规划） -├── wechat/ 微信通道 -│ ├── types.rs iLink API 协议类型 -│ └── client.rs HTTP 客户端（login/poll/send/notify） +├── db/ 数据库 +│ ├── mod.rs Database 连接池 + 迁移 +│ └── models.rs CRUD（auth/chat/usage/summary） ├── daemon.rs 守护进程（持有 WeChat + DB，spawn worker） ├── worker.rs Worker 进程（无状态 LLM 执行） └── ipc.rs Unix Domain Socket 帧协议 diff --git a/src/wechat/client.rs b/src/channel/client.rs similarity index 96% rename from src/wechat/client.rs rename to src/channel/client.rs index 42c1084..24dcae0 100644 --- a/src/wechat/client.rs +++ b/src/channel/client.rs @@ -1,4 +1,4 @@ -use crate::wechat::types::*; +use crate::channel::types::*; use base64::Engine; use reqwest::Client as HttpClient; use std::sync::Arc; @@ -327,12 +327,23 @@ impl WeChatClient { let url = format!("{}/ilink/bot/sendmessage", self.base_url); let token = self.token.lock().await.clone(); - let _resp_text = self + let resp: serde_json::Value = self .post_json(&url, &body, Some(&token)) .await? - .text() + .json() .await - .map_err(|e| format!("发送消息网络错误: {}", e))?; + .map_err(|e| format!("解析发送响应失败: {}", e))?; + + // 校验业务返回码（存在 ret 字段且非 0 才算失败） + if let Some(ret) = resp.get("ret").and_then(|v| v.as_i64()) { + if ret != 0 { + let errmsg = resp.get("errmsg").and_then(|v| v.as_str()); + return Err(format!( + "发送消息业务失败 ret={} errmsg={:?}", + ret, errmsg + )); + } + } Ok(client_id) } diff --git a/src/wechat/mod.rs b/src/channel/mod.rs similarity index 100% rename from src/wechat/mod.rs rename to src/channel/mod.rs diff --git a/src/wechat/types.rs b/src/channel/types.rs similarity index 100% rename from src/wechat/types.rs rename to src/channel/types.rs diff --git a/src/context/mod.rs b/src/context/mod.rs deleted file mode 100644 index c7178ca..0000000 --- a/src/context/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod builder; -pub mod tools; -pub mod types; - -pub use tools::MemoryStore; diff --git a/src/core/context.rs b/src/core/context.rs new file mode 100644 index 0000000..7f17b7b --- /dev/null +++ b/src/core/context.rs @@ -0,0 +1,112 @@ +use crate::core::session::ChatSession; +use crate::llm::types::Message; +use std::sync::Arc; +use tokio::sync::Mutex; + +/// 估算消息的 token 数 +pub fn estimate_tokens(msg: &Message) -> u32 { + let text = msg.content.as_str(); + estimate_text_tokens(text) + 8 +} + +/// 估算纯文本的 token 数 +pub fn estimate_text_tokens(text: &str) -> u32 { + if text.is_empty() { + return 0; + } + let cjk_count = text.chars().filter(|c| is_cjk(*c)).count() as u32; + let other_count = text.chars().count() as u32 - cjk_count; + // 中文 ≈ 1.5 token/字，英文 ≈ 0.25 token/字符 + cjk_count * 15 / 10 + other_count / 4 +} + +fn is_cjk(c: char) -> bool { + ('\u{4e00}'..='\u{9fff}').contains(&c) + || ('\u{3400}'..='\u{4dbf}').contains(&c) + || ('\u{f900}'..='\u{faff}').contains(&c) +} + +/// 构建给 LLM 的消息数组（带 token budget 管理） +/// +/// 返回消息数组和是否需要摘要的信息 +pub async fn build_context( + session: &Arc>, + system_prompt: &str, +) -> Vec { + let s = session.lock().await; + + // ── 1. 空闲超时检查（消息到达前由调用方检查）── + // 这里只做构建，超时触发在上层 + + // ── 2. 构建消息 ── + let mut messages = Vec::new(); + let mut used = 0u32; + + // 系统提示词 + let sys_tokens = estimate_text_tokens(system_prompt); + messages.push(Message::system(system_prompt)); + used += sys_tokens + 8; + + // overflow 摘要（如果有） + if let Some(summary) = s.latest_overflow_summary() { + let text = format!("【历史对话摘要】\n{}", summary.text); + let tokens = estimate_text_tokens(&text); + if used + tokens < s.token_budget { + messages.push(Message::system(text)); + used += tokens + 8; + } + } + + // 近期消息（从 checkpoint 开始，倒序添加直到接近 budget） + let recent = s.recent_messages().to_vec(); + let mut included = Vec::new(); + let mut tail_used = 0u32; + + for msg in recent.iter().rev() { + let t = estimate_tokens(msg); + if used + tail_used + t > s.token_budget - 500 { + break; + } + tail_used += t; + included.push(msg.clone()); + } + included.reverse(); + + // 保护：确保最近一条用户消息始终在上下文中 + if let Some(last_user) = recent + .iter() + .rev() + .find(|m| m.role == crate::llm::types::Role::User) + { + let already_included = included + .iter() + .any(|m| m.role == last_user.role && m.content == last_user.content); + if !already_included { + let t = estimate_tokens(last_user); + if used + tail_used + t <= s.token_budget + 2000 { + included.push(last_user.clone()); + } + } + } + + messages.extend(included); + + messages +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_estimate_cjk() { + let t = estimate_text_tokens("你好世界"); + assert!(t >= 5 && t <= 10, "CJK estimate: {}", t); + } + + #[test] + fn test_estimate_ascii() { + let t = estimate_text_tokens("hello world"); + assert!(t >= 1 && t <= 6, "ASCII estimate: {}", t); + } +} diff --git a/src/context/tools.rs b/src/core/memory.rs similarity index 94% rename from src/context/tools.rs rename to src/core/memory.rs index ff01866..eb2b271 100644 --- a/src/context/tools.rs +++ b/src/core/memory.rs @@ -63,15 +63,15 @@ impl MemoryStore { /// read_summaries 工具：读取历史摘要（内存 + 数据库） pub async fn read_summaries( - session: &Arc>, + session: &Arc>, ) -> String { let s = session.lock().await; let mut lines: Vec = Vec::new(); for sum in &s.summaries { let reason = match sum.reason { - super::types::SummaryReason::Overflow => "上下文压缩", - super::types::SummaryReason::Timeout => "空闲超时", + super::session::SummaryReason::Overflow => "上下文压缩", + super::session::SummaryReason::Timeout => "空闲超时", }; lines.push(format!( "[{}] {} (原因: {})", diff --git a/src/core/mod.rs b/src/core/mod.rs new file mode 100644 index 0000000..2e4fbac --- /dev/null +++ b/src/core/mod.rs @@ -0,0 +1,7 @@ +pub mod context; +pub mod memory; +pub mod pipeline; +pub mod session; +pub mod summary; + +pub use memory::MemoryStore; diff --git a/src/core/pipeline.rs b/src/core/pipeline.rs new file mode 100644 index 0000000..684b575 --- /dev/null +++ b/src/core/pipeline.rs @@ -0,0 +1,285 @@ +//! 消息处理流水线 +//! +//! 从 main.rs 抽取消息处理核心逻辑： +//! 1. 接收微信消息 +//! 2. 用户切换 → 上下文加载 +//! 3. LLM 对话 + 工具调用 +//! 4. 结果持久化（chat_records, llm_usage） +//! +//! 供 listen 模式和 daemon 模式复用。 + +use crate::core::memory::MemoryStore; +use crate::db::Database; +use crate::llm::{Conversation, Usage}; +use crate::tools::approval::ApprovalManager; +use crate::channel::client::WeChatClient; +use crate::state; +use std::sync::Arc; +use tokio::sync::Mutex; +use tracing::{error, info}; + +/// 消息处理上下文 +pub struct PipelineContext { + pub client: WeChatClient, + pub database: Option>, + pub file_state: state::StateManager, + pub memory_store: Arc, + pub approval_manager: Arc, + pub account_id: String, + /// 全局会话锁：所有消息共享同一 Conversation，必须串行处理防止跨用户串话 + pub conversation_lock: Arc>, +} + +impl PipelineContext { + /// 监听循环：长轮询微信消息，处理每条消息 + pub async fn run_listen_loop( + &self, + enable_llm: bool, + echo: bool, + conversation: Option>, + ) { + let last_user_id = Arc::new(tokio::sync::Mutex::new(String::new())); + + loop { + match self.client.receive_messages().await { + Ok(resp) => { + // 持久化 runtime buf + if !resp.get_updates_buf.is_empty() { + self.file_state.save_runtime(&resp.get_updates_buf); + } + + for msg in &resp.msgs { + if msg.msg_type != Some(crate::channel::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) = &self.database { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "inbound", + from, + &self.account_id, + text, + "wechat", + ctx_token, + &msg_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + + // Echo 模式 + if echo { + match self.client.send_text(from, text, ctx_token).await { + Ok(sent_id) => { + info!("回显成功 msg_id={}", sent_id); + if let Some(db) = &self.database { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "outbound", + from, + &self.account_id, + text, + "echo", + ctx_token, + &sent_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + } + Err(e) => error!("回显失败: {}", e), + } + } + + // 审批回复检查 + let is_approval_reply = self + .approval_manager + .handle_reply(from, text) + .await + .is_some(); + + if is_approval_reply { + info!("消息匹配审批确认码，不传给 LLM"); + continue; + } + + // 用户切换 + LLM 回复（全局会话锁保证串行） + if enable_llm { + if let Some(conv) = &conversation { + let conv = conv.clone(); + let client = self.client.clone(); + let database = self.database.clone(); + let from_owned = from.to_string(); + let text_owned = text.to_string(); + let ctx_owned = ctx_token.map(String::from); + let aid = self.account_id.clone(); + let memory_store = self.memory_store.clone(); + let conv_lock = Arc::clone(&self.conversation_lock); + let last_user = Arc::clone(&last_user_id); + + tokio::spawn(async move { + let _guard = conv_lock.lock().await; + + // ── 用户切换 ── + { + let mut last = last_user.lock().await; + if *last != from_owned { + info!("用户切换: {} → {}", *last, from_owned); + let session = conv.session(); + let mut s = session.lock().await; + s.messages.clear(); + s.checkpoint = 0; + s.summaries.clear(); + s.last_user_at = None; + s.current_user_id = from_owned.clone(); + s.load_recent_messages(20).await; + drop(s); + memory_store.load(&from_owned).await; + } + *last = from_owned.clone(); + } + + // ── LLM 回复 ── + let reply_result = if conv.spec().tools.is_some() { + generate_reply_with_tools( + &conv, + text_owned, + &from_owned, + &database, + ) + .await + } else { + generate_reply( + &conv, + text_owned, + &from_owned, + &database, + ) + .await + }; + let reply = match reply_result { + Ok(r) if !r.trim().is_empty() => r, + Ok(_) => "抱歉，生成回复为空，请重试。".to_string(), + Err(e) => format!("处理消息时出错：{:.200}", e), + }; + match client + .send_text(&from_owned, &reply, ctx_owned.as_deref()) + .await + { + Ok(sent_id) => { + info!("LLM 回复成功 msg_id={}", sent_id); + if let Some(db) = database { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "outbound", + &from_owned, + &aid, + &reply, + "llm", + ctx_owned.as_deref(), + &sent_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + } + Err(e) => error!("发送 LLM 回复失败: {}", e), + } + }); + } + } + } + } + Err(e) => { + if !e.contains("超时") && !e.contains("timeout") { + error!("接收消息失败: {}", e); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + } + } + } + } + } +} + +async fn generate_reply( + conv: &Conversation, + user_text: String, + user_id: &str, + db: &Option>, +) -> Result { + let mut handle = conv.chat(user_text).await?; + let reply = handle.consume().await?; + if let Some(u) = handle.get_usage() { + info!( + "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", + u.total_tokens, + u.prompt_tokens, + u.completion_tokens, + u.prompt_cache_hit_tokens, + u.prompt_cache_miss_tokens + ); + store_usage(db, user_id, conv.model(), conv.provider_name(), u).await; + } + Ok(reply) +} + +async fn generate_reply_with_tools( + conv: &Conversation, + user_text: String, + user_id: &str, + db: &Option>, +) -> Result { + let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?; + if let Some(u) = &usage { + info!( + "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", + u.total_tokens, + u.prompt_tokens, + u.completion_tokens, + u.prompt_cache_hit_tokens, + u.prompt_cache_miss_tokens + ); + store_usage(db, user_id, conv.model(), conv.provider_name(), u).await; + } + Ok(reply) +} + +async fn store_usage( + db: &Option>, + user_id: &str, + model: &str, + provider: &str, + u: &Usage, +) { + if let Some(db) = db { + if let Err(e) = crate::db::models::insert_llm_usage( + db.pool(), + user_id, + model, + provider, + u.prompt_tokens, + u.completion_tokens, + u.prompt_cache_hit_tokens, + u.prompt_cache_miss_tokens, + ) + .await + { + tracing::error!("存储 LLM 用量失败: {}", e); + } + } +} diff --git a/src/context/types.rs b/src/core/session.rs similarity index 100% rename from src/context/types.rs rename to src/core/session.rs diff --git a/src/context/builder.rs b/src/core/summary.rs similarity index 59% rename from src/context/builder.rs rename to src/core/summary.rs index 2b4fee8..ec58cb0 100644 --- a/src/context/builder.rs +++ b/src/core/summary.rs @@ -1,97 +1,9 @@ -use crate::context::types::ChatSession; +use crate::core::session::{ChatSession, SummaryEntry, SummaryReason}; use crate::llm::conversation::Summarizer; use crate::llm::types::Message; use std::sync::Arc; use tokio::sync::Mutex; -/// 估算消息的 token 数 -pub fn estimate_tokens(msg: &Message) -> u32 { - let text = msg.content.as_str(); - estimate_text_tokens(text) + 8 -} - -/// 估算纯文本的 token 数 -pub fn estimate_text_tokens(text: &str) -> u32 { - if text.is_empty() { - return 0; - } - let cjk_count = text.chars().filter(|c| is_cjk(*c)).count() as u32; - let other_count = text.chars().count() as u32 - cjk_count; - // 中文 ≈ 1.5 token/字，英文 ≈ 0.25 token/字符 - cjk_count * 15 / 10 + other_count / 4 -} - -fn is_cjk(c: char) -> bool { - ('\u{4e00}'..='\u{9fff}').contains(&c) - || ('\u{3400}'..='\u{4dbf}').contains(&c) - || ('\u{f900}'..='\u{faff}').contains(&c) -} - -/// 构建给 LLM 的消息数组（带 token budget 管理） -/// -/// 返回消息数组和是否需要摘要的信息 -pub async fn build_context(session: &Arc>, system_prompt: &str) -> Vec { - let s = session.lock().await; - - // ── 1. 空闲超时检查（消息到达前由调用方检查）── - // 这里只做构建，超时触发在上层 - - // ── 2. 构建消息 ── - let mut messages = Vec::new(); - let mut used = 0u32; - - // 系统提示词 - let sys_tokens = estimate_text_tokens(system_prompt); - messages.push(Message::system(system_prompt)); - used += sys_tokens + 8; - - // overflow 摘要（如果有） - if let Some(summary) = s.latest_overflow_summary() { - let text = format!("【历史对话摘要】\n{}", summary.text); - let tokens = estimate_text_tokens(&text); - if used + tokens < s.token_budget { - messages.push(Message::system(text)); - used += tokens + 8; - } - } - - // 近期消息（从 checkpoint 开始，倒序添加直到接近 budget） - let recent = s.recent_messages().to_vec(); - let mut included = Vec::new(); - let mut tail_used = 0u32; - - for msg in recent.iter().rev() { - let t = estimate_tokens(msg); - if used + tail_used + t > s.token_budget - 500 { - break; - } - tail_used += t; - included.push(msg.clone()); - } - included.reverse(); - - // 保护：确保最近一条用户消息始终在上下文中 - if let Some(last_user) = recent - .iter() - .rev() - .find(|m| m.role == crate::llm::types::Role::User) - { - let already_included = included - .iter() - .any(|m| m.role == last_user.role && m.content == last_user.content); - if !already_included { - let t = estimate_tokens(last_user); - if used + tail_used + t <= s.token_budget + 2000 { - included.push(last_user.clone()); - } - } - } - - messages.extend(included); - - messages -} - /// 触发溢出摘要：压缩 checkpoint 到当前位置之间的新消息 pub async fn trigger_overflow_summary( session: &Arc>, @@ -123,10 +35,10 @@ pub async fn trigger_overflow_summary( let summary_text = generate_summary(&to_summarize, summarizer).await; - s.summaries.push(super::types::SummaryEntry { + s.summaries.push(SummaryEntry { checkpoint: prev_checkpoint, text: summary_text.clone(), - reason: super::types::SummaryReason::Overflow, + reason: SummaryReason::Overflow, created_at: chrono::Utc::now(), }); @@ -166,10 +78,10 @@ pub async fn trigger_idle_summary( let summary_text = generate_summary(&s.messages, summarizer).await; let cp = s.messages.len(); - s.summaries.push(super::types::SummaryEntry { + s.summaries.push(SummaryEntry { checkpoint: cp, text: summary_text.clone(), - reason: super::types::SummaryReason::Timeout, + reason: SummaryReason::Timeout, created_at: chrono::Utc::now(), }); @@ -214,7 +126,8 @@ fn build_summary_prompt(messages: &[Message]) -> String { let convo: String = messages .iter() .filter(|m| { - m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant + m.role == crate::llm::types::Role::User + || m.role == crate::llm::types::Role::Assistant }) .map(|m| { let role = if m.role == crate::llm::types::Role::User { @@ -240,7 +153,8 @@ fn summarize_messages(messages: &[Message]) -> String { let lines: Vec = messages .iter() .filter(|m| { - m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant + m.role == crate::llm::types::Role::User + || m.role == crate::llm::types::Role::Assistant }) .map(|m| { let role = match m.role { @@ -271,20 +185,3 @@ fn summarize_messages(messages: &[Message]) -> String { } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_estimate_cjk() { - let t = estimate_text_tokens("你好世界"); - assert!(t >= 5 && t <= 10, "CJK estimate: {}", t); - } - - #[test] - fn test_estimate_ascii() { - let t = estimate_text_tokens("hello world"); - assert!(t >= 1 && t <= 6, "ASCII estimate: {}", t); - } -} diff --git a/src/daemon.rs b/src/daemon.rs index 5c3661d..672a2c1 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -4,12 +4,12 @@ //! 通过 Unix Domain Socket 通信。Worker 代码更新后编译即可，下一条消息 //! 自动使用新版本。 -use crate::context::MemoryStore; +use crate::core::MemoryStore; use crate::db::Database; use crate::ipc::{HistoryEntry, OutputFrame, TaskFrame, TaskMessage, recv_output, send_frame}; use crate::state::StateManager; use crate::tools::approval::{ApprovalDecision, ApprovalManager}; -use crate::wechat::client::WeChatClient; +use crate::channel::client::WeChatClient; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use std::time::Instant; @@ -28,6 +28,7 @@ struct MessageQueue { waiting: VecDeque, } +#[derive(Clone)] struct PendingMessage { from: String, text: String, @@ -36,6 +37,7 @@ struct PendingMessage { #[allow(dead_code)] message_id: String, approved_tool: Option, + approved_tool_args: Option, } impl MessageQueue { @@ -78,6 +80,20 @@ impl MessageQueue { self.pending.get_mut(user_id)?.pop_front() } + /// 将消息放回队首（用于 send_frame 失败恢复），并确保用户回到 waiting + fn push_front(&mut self, user_id: &str, msg: PendingMessage) { + self.pending + .entry(user_id.to_string()) + .or_default() + .push_front(msg); + // 如果用户未活跃且不在 waiting 中，重新加入 waiting + if !self.active.get(user_id).unwrap_or(&false) + && !self.waiting.contains(&user_id.to_string()) + { + self.waiting.push_back(user_id.to_string()); + } + } + fn deactivate(&mut self, user_id: &str) { self.active.insert(user_id.to_string(), false); } @@ -144,11 +160,11 @@ pub async fn run( // 6. 消息队列 let queue = Arc::new(Mutex::new(MessageQueue::new())); - // 7. 构建工具列表（传给 worker） - let tools_list = build_tools_list(); + // 7. 构建工具列表（统一使用 tools 模块） + let tools_list = crate::tools::build_tools_list(); - // 8. 环境变量映射 - let env_map = build_env_map(); + // 8. 环境变量映射（统一使用 tools 模块） + let env_map = crate::tools::build_env_map(); // 9. 系统提示词 let system_prompt = @@ -181,7 +197,7 @@ pub async fn run( // 清理过期的审批上下文（超过 5 分钟） let now = Instant::now(); shared_for_clean.approval_ctx.lock().await - .retain(|_, (_, _, _, created)| now.duration_since(*created).as_secs() < 300); + .retain(|_, (_, _, _, _, created)| now.duration_since(*created).as_secs() < 300); } }); @@ -222,7 +238,7 @@ pub async fn run( } for msg in &resp.msgs { - if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) { + if msg.msg_type != Some(crate::channel::types::WeixinMessage::TYPE_USER) { continue; } let from = msg.from_user_id.as_deref().unwrap_or("unknown"); @@ -256,7 +272,7 @@ pub async fn run( ApprovalDecision::Approved => { info!("审批通过: {}", skill_name); let original = shared.approval_ctx.lock().await.remove(from); - if let Some((orig_text, orig_ctx, _tool, _ts)) = original { + if let Some((orig_text, orig_ctx, _tool, tool_args, _ts)) = original { let mut q = queue.lock().await; q.enqueue(from, PendingMessage { from: from.to_string(), @@ -264,7 +280,8 @@ pub async fn run( account_id: shared.account_id.clone(), context_token: orig_ctx, message_id: String::new(), - approved_tool: Some(skill_name), + approved_tool: Some(skill_name.clone()), + approved_tool_args: if tool_args.is_empty() { None } else { Some(tool_args) }, }); drop(q); spawn_worker_for_user(from, &queue, &shared).await; @@ -292,6 +309,7 @@ pub async fn run( context_token: ctx_token.map(String::from), message_id: msg_id, approved_tool: None, + approved_tool_args: None, }); drop(q); } @@ -340,8 +358,8 @@ struct DaemonShared { env_map: HashMap, system_prompt: String, model: String, - /// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 创建时间) - approval_ctx: Mutex, String, Instant)>>, + /// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 工具参数, 创建时间) + approval_ctx: Mutex, String, String, Instant)>>, } /// 处理一个 Worker 连接 @@ -394,6 +412,8 @@ async fn handle_worker( // 5. 构建 TaskFrame let approved_tool_info = pending_msg.approved_tool.clone(); + let approved_tool_args_info = pending_msg.approved_tool_args.clone(); + let pending_clone = pending_msg.clone(); // 保留副本用于失败恢复 let task = TaskFrame { user_id: user_id.clone(), msg: TaskMessage { @@ -406,6 +426,7 @@ async fn handle_worker( memories, summaries, approved_tool: approved_tool_info.clone(), + approved_tool_args: approved_tool_args_info.clone(), env: shared.env_map.clone(), tools: shared.tools_list.clone(), system_prompt: shared.system_prompt.clone(), @@ -415,7 +436,12 @@ async fn handle_worker( // 6. 发送 task 帧 if let Err(e) = send_frame(&mut stream, &task).await { error!("发送 task 帧失败: {}", e); - queue.lock().await.deactivate(&user_id); + let mut q = queue.lock().await; + q.deactivate(&user_id); // 先取消 active，push_front 才能加入 waiting + q.push_front(&user_id, pending_clone); + drop(q); + // 尝试重新 spawn + spawn_worker_for_user(&user_id, &queue, &shared).await; return; } @@ -433,7 +459,7 @@ async fn handle_worker( OutputFrame::Reply { text } => { reply_text = text; } - OutputFrame::RequestApproval { tool, reason } => { + OutputFrame::RequestApproval { tool, reason, tool_args } => { info!("Worker 请求审批: tool={} reason={}", tool, reason); // 创建审批记录 match shared.approval.create(&user_id, &tool).await { @@ -445,7 +471,7 @@ async fn handle_worker( // 保存原始消息上下文以便审批后重新处理 shared.approval_ctx.lock().await.insert( user_id.clone(), - (task.msg.text.clone(), ctx_token.clone(), tool.clone(), Instant::now()), + (task.msg.text.clone(), ctx_token.clone(), tool.clone(), tool_args.clone(), Instant::now()), ); // 发送确认消息到微信 if let Err(e) = shared.client.send_text(&user_id, &msg, ctx_token.as_deref()).await { @@ -763,165 +789,3 @@ async fn handle_db_write( _ => warn!("未知 DbWrite table: {}", table), } } - -fn build_tools_list() -> Vec { - let mut list = Vec::new(); - - // query_capabilities - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "query_capabilities", - "description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。", - "parameters": { "type": "object", "properties": {} } - } - })); - - // call_capability - list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "call_capability", - "description": "调用一个已注册的工具。name 为工具名（来自 query_capabilities），工具所需的具体参数应直接放在 JSON 中。", - "parameters": { - "type": "object", - "properties": { - "name": { "type": "string", "description": "工具名称" }, - "prompt": { "type": "string", "description": "json格式参数" } - }, - "required": ["name"] - } - } - })); - - // 上下文工具 - list.push(serde_json::json!({ - "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"] - } - } - })); - list.push(serde_json::json!({ - "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"] - } - } - })); - - 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..fb2f309 100644 --- a/src/ipc.rs +++ b/src/ipc.rs @@ -22,6 +22,9 @@ pub struct TaskFrame { pub summaries: Vec, #[serde(default)] pub approved_tool: Option, + /// 审批通过后的工具调用参数 + #[serde(default)] + pub approved_tool_args: Option, pub env: std::collections::HashMap, pub tools: Vec, pub system_prompt: String, @@ -65,6 +68,9 @@ pub enum OutputFrame { RequestApproval { tool: String, reason: String, + /// 工具调用参数（审批通过后直接执行） + #[serde(default)] + tool_args: String, }, /// 数据库写入 #[serde(rename = "db_write")] diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index 618dc85..35a06ac 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -1,5 +1,6 @@ -use crate::context::builder; -use crate::context::types::ChatSession; +use crate::core::context as ctx_builder; +use crate::core::session::ChatSession; +use crate::core::summary as summary_builder; use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider}; use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage}; use std::future::Future; @@ -95,12 +96,13 @@ impl Conversation { let s = self.session.lock().await; if s.is_idle_timeout() { drop(s); - builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await; + summary_builder::trigger_idle_summary(&self.session, Some(&self.summarizer())) + .await; } } self.session.lock().await.add_user(user_message.clone()); - let messages = builder::build_context(&self.session, &self.config.system_prompt).await; + let messages = ctx_builder::build_context(&self.session, &self.config.system_prompt).await; let rx = self.provider.chat_stream(&self.config, &messages).await?; Ok(ChatHandle { @@ -122,7 +124,8 @@ impl Conversation { let s = self.session.lock().await; if s.is_idle_timeout() { drop(s); - builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await; + summary_builder::trigger_idle_summary(&self.session, Some(&self.summarizer())) + .await; tracing::info!("⏰ 检测到 12h 空闲，已生成摘要"); } } @@ -135,11 +138,12 @@ impl Conversation { loop { turn_count += 1; - if turn_count > 5 { - return Err("工具调用次数过多（最多5轮）".to_string()); + if turn_count > 30 { + return Err("工具调用次数过多（最多30轮）".to_string()); } - let messages = builder::build_context(&self.session, &self.config.system_prompt).await; + let messages = + ctx_builder::build_context(&self.session, &self.config.system_prompt).await; let rx = self.provider.chat_stream(&self.config, &messages).await?; let mut full_text = String::new(); @@ -215,11 +219,14 @@ impl Conversation { { let s = self.session.lock().await; let recent = s.recent_messages(); - let estimated: u32 = recent.iter().map(|m| builder::estimate_tokens(m)).sum(); + let estimated: u32 = recent.iter().map(|m| ctx_builder::estimate_tokens(m)).sum(); if estimated > s.token_budget { drop(s); - builder::trigger_overflow_summary(&self.session, Some(&self.summarizer())) - .await; + summary_builder::trigger_overflow_summary( + &self.session, + Some(&self.summarizer()), + ) + .await; tracing::info!("📦 上下文超预算，已触发溢出摘要"); } } diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 4ee1972..7b73af8 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::{Conversation, DEFAULT_SYSTEM_PROMPT}; pub use types::{ConversationConfig, Usage}; diff --git a/src/main.rs b/src/main.rs index 4947356..08d7db6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,5 @@ mod cli; -mod context; +mod core; mod daemon; mod db; mod ipc; @@ -8,19 +8,21 @@ mod logger; mod scheduler; mod state; mod tools; -mod wechat; +mod channel; mod worker; use clap::Parser; use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand}; -use context::MemoryStore; +use core::MemoryStore; +use core::pipeline::PipelineContext; use db::Database; -use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; +use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT}; use std::sync::Arc; -use tools::approval::{ApprovalDecision, ApprovalManager}; -use tools::types::ExecutionContext; +use tokio::sync::Mutex; +use tools::approval::ApprovalManager; +use tools::executor::ToolExecutorBuilder; use tracing::{error, info}; -use wechat::client::WeChatClient; +use channel::client::WeChatClient; #[tokio::main] async fn main() { @@ -240,66 +242,7 @@ async fn cmd_listen( // 各项能力通过这两个元工具实现：查询能力 + 调用能力 // 具体工具描述通过 query_capabilities 返回 - let mut tools_list: Vec = Vec::new(); - - // query_capabilities：列出所有可用工具及其说明 - tools_list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "query_capabilities", - "description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。", - "parameters": { "type": "object", "properties": {} } - } - })); - - // call_capability：按名称调用任意工具 - tools_list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "call_capability", - "description": "调用一个已注册的工具。name 为工具名（来自 query_capabilities），工具所需的具体参数（如 location、days）应直接放在 JSON 中。", - "parameters": { - "type": "object", - "properties": { - "name": { "type": "string", "description": "工具名称" }, - "prompt": { "type": "string", "description": "json格式参数，来自 query_capabilities 中的对应工具的描述。" } - }, - "required": ["name"] - } - } - })); - - // 上下文工具（仍然直接暴露给 LLM） - tools_list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "read_memories", - "description": "读取用户的长期记忆（偏好、个人信息、约定）", - "parameters": { "type": "object", "properties": {} } - } - })); - tools_list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "write_memory", - "description": "记录用户的重要信息或偏好", - "parameters": { - "type": "object", - "properties": { - "content": { "type": "string", "description": "记忆内容" } - }, - "required": ["content"] - } - } - })); - tools_list.push(serde_json::json!({ - "type": "function", - "function": { - "name": "read_summaries", - "description": "读取历史会话摘要", - "parameters": { "type": "object", "properties": {} } - } - })); + let tools_list = tools::build_tools_list(); cfg.tools = Some(tools_list); @@ -321,10 +264,7 @@ async fn cmd_listen( } }; - // 构建工具执行器 - let session = conv.session(); - let approval = approval_manager.clone(); - let memory_store = memory_store.clone(); + // 构建工具执行器（统一使用 ToolExecutorBuilder） let send_wechat: tools::types::WechatSender = { let client = client.clone(); Arc::new(move |user_id: &str, text: &str| { @@ -341,77 +281,12 @@ async fn cmd_listen( }) }; - let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| { - let approval = approval.clone(); - let send = send_wechat.clone(); - let session = session.clone(); - let memory_store = memory_store.clone(); - let n = name.to_string(); - let args = args_json.to_string(); - - Box::pin(async move { - let user_id = session.lock().await.current_user_id.clone(); - - // 上下文工具直接处理 - if n == "read_memories" { - return Ok(memory_store.read_for(&user_id).await); - } - if n == "write_memory" { - let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); - let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); - if content.is_empty() { - return Ok("请提供 content 参数".to_string()); - } - return Ok(memory_store.write_for(&user_id, content).await); - } - if n == "read_summaries" { - return Ok(context::tools::read_summaries(&session).await); - } - - // query_capabilities: 列出所有内置工具 - if n == "query_capabilities" { - let specs = tools::builtin::BuiltinRegistry::specs(); - return Ok(tools::build_capability_guide(&specs)); - } - - // 确定目标工具名和参数 - // call_capability: 从 {name, prompt} 中提取 name，从 prompt 中解包真实参数 - let (target_name, target_args) = if n == "call_capability" { - let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); - let name = cp - .get("name") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - .unwrap_or_default(); - let unpacked = tools::unpack_call_params(&cp); - (name, serde_json::to_string(&unpacked).unwrap_or_default()) - } else { - (n.clone(), args.clone()) - }; - - // 内置工具 - if tools::builtin::BuiltinRegistry::is_builtin(&target_name) { - if tools::builtin::BuiltinRegistry::is_high_risk(&target_name) { - let mut ctx = ExecutionContext::new(&user_id); - ctx = ctx.with_approval(approval.clone()); - ctx = ctx.with_wechat(send.clone()); - let approved = builtin_approve(&ctx, &target_name).await; - match approved { - Ok(true) => {} - Ok(false) => return Ok(SkillResult::rejected(&target_name).output), - Err(e) => return Ok(e), - } - } - if let Some(result) = - tools::builtin::BuiltinRegistry::execute(&target_name, &target_args).await - { - return Ok(result.output); - } - } - - Ok(format!("未知工具: {}", target_name)) - }) - }); + let executor = ToolExecutorBuilder::new("") + .with_session(conv.session()) + .with_approval(approval_manager.clone()) + .with_wechat(send_wechat) + .with_memory(memory_store.clone()) + .build(); conv.set_tool_executor(executor); @@ -455,6 +330,17 @@ async fn cmd_listen( info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo); println!("已开始监听消息，按 Ctrl+C 退出"); + // 构建流水线 + let pipeline = PipelineContext { + client: client.clone(), + database: database.clone(), + file_state: file_state.clone(), + memory_store: memory_store.clone(), + approval_manager: approval_manager.clone(), + account_id, + conversation_lock: Arc::new(Mutex::new(())), + }; + let shutdown = async { tokio::signal::ctrl_c().await.expect("Ctrl+C 失败") }; tokio::select! { @@ -464,244 +350,7 @@ async fn cmd_listen( file_state.save_runtime(&client.get_updates_buf().await); info!("已退出"); } - _ = listen_loop(&client, enable_llm, echo, &account_id, database, file_state, conversation, &approval_manager, memory_store) => {} - } -} - -async fn listen_loop( - client: &WeChatClient, - enable_llm: bool, - echo: bool, - account_id: &str, - database: &Option>, - file_state: &state::StateManager, - conversation: Option>, - approval_manager: &ApprovalManager, - memory_store: &Arc, -) { - let last_user_id = Arc::new(tokio::sync::Mutex::new(String::new())); - loop { - match client.receive_messages().await { - 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(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) = database { - if let Err(e) = db::models::insert_chat_record( - db.pool(), - "inbound", - from, - account_id, - text, - "wechat", - ctx_token, - &msg_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } - } - - // Echo 模式 - if echo { - match client.send_text(from, text, ctx_token).await { - Ok(sent_id) => { - info!("回显成功 msg_id={}", sent_id); - // === 入库：发送的回显 === - if let Some(db) = database { - if let Err(e) = db::models::insert_chat_record( - db.pool(), - "outbound", - from, - account_id, - text, - "echo", - ctx_token, - &sent_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } - } - } - Err(e) => error!("回显失败: {}", e), - } - } - - // 审批回复检查 - let is_approval_reply = - approval_manager.handle_reply(from, text).await.is_some(); - - if is_approval_reply { - info!("消息匹配审批确认码，不传给 LLM"); - continue; - } - - // 用户切换：清空上下文，加载新用户历史 - { - let mut last_user = last_user_id.lock().await; - if *last_user != *from { - info!("用户切换: {} → {}", *last_user, from); - if let Some(conv) = &conversation { - let session = conv.session(); - let mut s = session.lock().await; - s.messages.clear(); - s.checkpoint = 0; - s.summaries.clear(); - s.last_user_at = None; - s.current_user_id = from.to_string(); - // 加载该用户最近 20 条历史 - s.load_recent_messages(20).await; - drop(s); - memory_store.load(from).await; - } - } - *last_user = from.to_string(); - } - - // LLM 回复（异步执行，避免审批阻塞主循环） - if enable_llm { - if let Some(conv) = &conversation { - let conv = conv.clone(); - let client = client.clone(); - let database = database.clone(); - let from_owned = from.to_string(); - let text_owned = text.to_string(); - let ctx_owned = ctx_token.map(String::from); - let aid = account_id.to_string(); - - tokio::spawn(async move { - let reply_result = if conv.spec().tools.is_some() { - generate_reply_with_tools(&conv, text_owned, &database).await - } else { - generate_reply(&conv, text_owned, &database).await - }; - let reply = match reply_result { - Ok(r) if !r.trim().is_empty() => r, - Ok(_) => "抱歉，生成回复为空，请重试。".to_string(), - Err(e) => format!("处理消息时出错：{:.200}", e), - }; - match client - .send_text(&from_owned, &reply, ctx_owned.as_deref()) - .await - { - Ok(sent_id) => { - info!("LLM 回复成功 msg_id={}", sent_id); - if let Some(db) = database { - if let Err(e) = db::models::insert_chat_record( - db.pool(), - "outbound", - &from_owned, - &aid, - &reply, - "llm", - ctx_owned.as_deref(), - &sent_id, - ) - .await - { - error!("保存聊天记录失败: {}", e); - } - } - } - Err(e) => error!("发送 LLM 回复失败: {}", e), - } - }); - } - } - } - } - Err(e) => { - if !e.contains("超时") && !e.contains("timeout") { - error!("接收消息失败: {}", e); - tokio::time::sleep(std::time::Duration::from_secs(2)).await; - } - } - } - } -} - -async fn generate_reply( - conv: &Conversation, - user_text: String, - db: &Option>, -) -> Result { - let mut handle = conv.chat(user_text).await?; - // 先消费流（usage 在 Done chunk 中才可用） - let reply = handle.consume().await?; - if let Some(u) = handle.get_usage() { - info!( - "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", - u.total_tokens, - u.prompt_tokens, - u.completion_tokens, - u.prompt_cache_hit_tokens, - u.prompt_cache_miss_tokens - ); - store_usage(db, "", conv.model(), conv.provider_name(), u).await; - } - Ok(reply) -} - -async fn generate_reply_with_tools( - conv: &Conversation, - user_text: String, - db: &Option>, -) -> Result { - let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?; - if let Some(u) = &usage { - info!( - "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", - u.total_tokens, - u.prompt_tokens, - u.completion_tokens, - u.prompt_cache_hit_tokens, - u.prompt_cache_miss_tokens - ); - store_usage(db, "", conv.model(), conv.provider_name(), u).await; - } - Ok(reply) -} - -async fn store_usage( - db: &Option>, - user_id: &str, - model: &str, - provider: &str, - u: &Usage, -) { - if let Some(db) = db { - if let Err(e) = db::models::insert_llm_usage( - db.pool(), - user_id, - model, - provider, - u.prompt_tokens, - u.completion_tokens, - u.prompt_cache_hit_tokens, - u.prompt_cache_miss_tokens, - ) - .await - { - tracing::error!("存储 LLM 用量失败: {}", e); - } + _ = pipeline.run_listen_loop(enable_llm, echo, conversation) => {} } } @@ -849,33 +498,6 @@ async fn cmd_send( } } -use tools::types::SkillResult; - -async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result { - let manager = match ctx.approval_manager.as_ref() { - Some(m) => m.clone(), - None => return Ok(true), - }; - let (code, rx) = manager.create(&ctx.user_id, name).await?; - - let msg = format!( - "⚠️ 操作确认\n\n技能：{}\n\n确认码：{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效，最多3次尝试)", - name, code - ); - if let Some(ref send) = ctx.send_wechat { - if let Err(e) = send(&ctx.user_id, &msg).await { - return Err(e); - } - } - - match tokio::time::timeout(std::time::Duration::from_secs(300), rx).await { - Ok(Ok(ApprovalDecision::Approved)) => Ok(true), - _ => Ok(false), - } -} - -// ─── 工具 CLI 命令 ─── - async fn cmd_tool(cmd: ToolCommand) { match cmd { ToolCommand::Datetime => { diff --git a/src/state.rs b/src/state.rs index 2403c15..f0bf2ea 100644 --- a/src/state.rs +++ b/src/state.rs @@ -16,6 +16,7 @@ pub struct RuntimeState { } /// 状态管理器 +#[derive(Clone)] pub struct StateManager { state_dir: PathBuf, } diff --git a/src/tools/approval.rs b/src/tools/approval.rs index 4b49816..a1598a4 100644 --- a/src/tools/approval.rs +++ b/src/tools/approval.rs @@ -147,6 +147,32 @@ impl ApprovalManager { Some((name, result)) } + /// 按确认码哈希精确取消一条待审批记录（不触发决策通知） + /// 改为 cancelled 状态同步 DB，用于 listen 模式审批超时时清理 + pub async fn cancel_by_code(&self, user_id: &str, code_hash: &str) { + let mut map = self.pending.lock().await; + if let Some(entries) = map.get_mut(user_id) { + if let Some(idx) = entries.iter().position(|e| e.code_hash == code_hash) { + let entry = entries.remove(idx); + // sender 被 drop，rx 端收到 Cancelled 错误 + tracing::debug!("取消审批: user={} skill={}", user_id, entry.skill_name); + } + if entries.is_empty() { + map.remove(user_id); + } + } + + // 同步 DB 状态 + if let Some(ref pool) = self.pool { + let _ = sqlx::query( + "UPDATE pending_approvals SET status = 'cancelled', consumed_at = NOW() WHERE code_hash = $1 AND status = 'pending'", + ) + .bind(code_hash) + .execute(pool.as_ref()) + .await; + } + } + /// 清理过期审批（按时间判定，通知等待方 + 更新 DB） pub async fn clean_expired(&self) { let now = Instant::now(); diff --git a/src/tools/builtins/mod.rs b/src/tools/builtins/mod.rs index 4b86d71..829592c 100644 --- a/src/tools/builtins/mod.rs +++ b/src/tools/builtins/mod.rs @@ -2,5 +2,6 @@ pub mod amap; pub mod datetime; pub mod fetch_page; pub mod memos; +pub mod subagent; pub mod weather; pub mod web_search; diff --git a/src/tools/builtins/subagent.rs b/src/tools/builtins/subagent.rs new file mode 100644 index 0000000..2994c19 --- /dev/null +++ b/src/tools/builtins/subagent.rs @@ -0,0 +1,229 @@ +//! Pi 子代理工具 — 调用独立 pi 实例执行任务 +//! +//! 通过 JSON 事件流模式启动一个独立的 pi 进程，传入 prompt 并收集流式回复。 +//! 用于将复杂子任务委派给独立 agent，隔离上下文。 +//! +//! 协议：`pi --mode json --no-session "prompt"` → stdout JSON 事件流 + +use super::super::types::{RiskLevel, SkillResult, SkillSpec}; +use serde::Deserialize; +use std::time::Duration; + +/// 默认超时（秒） +const DEFAULT_TIMEOUT: u64 = 120; +/// 最大回复长度 +const MAX_REPLY_CHARS: usize = 8000; + +// ─── Spec ─── + +pub fn spec() -> SkillSpec { + SkillSpec { + name: "pi_subagent".into(), + description: "启动独立的 pi agent 子进程处理复杂子任务。传入 prompt，返回 agent 完整回复。适合将复杂分析、代码生成等任务委派给独立上下文执行。".into(), + risk_level: RiskLevel::High, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "发送给子 agent 的完整 prompt" + }, + "timeout_secs": { + "type": "integer", + "description": "超时秒数（1-300，默认 120）" + } + }, + "required": ["prompt"] + }), + timeout_secs: DEFAULT_TIMEOUT, + } +} + +// ─── JSON 事件流类型 ─── + +/// 事件顶层结构：根据 type 字段区分 +#[derive(Deserialize, Debug)] +#[allow(dead_code)] +struct JsonEvent { + #[serde(rename = "type")] + event_type: String, + #[serde(default)] + message: Option, + #[serde(default)] + #[serde(rename = "assistantMessageEvent")] + assistant_message_event: Option, +} + +#[derive(Deserialize, Debug)] +#[allow(dead_code)] +struct EventMessage { + #[serde(default)] + content: Option, +} + +/// message_update 中的嵌套增量事件 +#[derive(Deserialize, Debug)] +struct AssistantMessageEvent { + #[serde(rename = "type")] + delta_type: Option, + #[serde(default)] + delta: Option, + #[serde(default)] + content: Option, +} + +// ─── 执行 ─── + +pub async fn execute(params: serde_json::Value) -> SkillResult { + let prompt = params + .get("prompt") + .and_then(|v| v.as_str()) + .unwrap_or(""); + + let timeout_secs = params + .get("timeout_secs") + .and_then(|v| v.as_u64()) + .unwrap_or(DEFAULT_TIMEOUT) + .clamp(1, 300); // 限制在 1~300 秒 + + if prompt.is_empty() { + return SkillResult::error("缺少 prompt 参数"); + } + + match run_subagent(prompt, timeout_secs).await { + Ok(text) => SkillResult::ok(text), + Err(e) => SkillResult::error(format!("pi 子代理执行失败: {}", e)), + } +} + +async fn run_subagent(prompt: &str, timeout_secs: u64) -> Result { + let pi_exe = find_pi_exe()?; + + let mut cmd = tokio::process::Command::new(&pi_exe); + cmd.args(["--mode", "json", "--no-session", prompt]) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::null()) // 不读 stderr，避免管道满阻塞 + .kill_on_drop(true); + + let mut child = cmd + .spawn() + .map_err(|e| format!("启动 pi 进程失败: {}", e))?; + + let stdout = child + .stdout + .take() + .ok_or("无法获取 pi stdout")?; + + use tokio::io::{AsyncBufReadExt, BufReader}; + + let reader = BufReader::new(stdout); + let mut lines = reader.lines(); + let mut reply = String::new(); + let mut done = false; + + let timeout = Duration::from_secs(timeout_secs); + let start = std::time::Instant::now(); + + while !done { + // 检查超时 + if start.elapsed() > timeout { + child.kill().await.ok(); + return Err("子代理执行超时".to_string()); + } + + let line = match tokio::time::timeout(Duration::from_secs(10), lines.next_line()).await { + Ok(Ok(Some(l))) => l, + Ok(Ok(None)) => break, // EOF + Ok(Err(e)) => return Err(format!("读取 pi 输出失败: {}", e)), + Err(_) => continue, // 读超时，继续循环 + }; + + if line.trim().is_empty() { + continue; + } + + let event: JsonEvent = match serde_json::from_str(&line) { + Ok(e) => e, + Err(_) => continue, // 跳过无法解析的行 + }; + + match event.event_type.as_str() { + "session" => { + // 首行：会话头部，跳过 + } + "message_update" => { + if let Some(delta_ev) = event.assistant_message_event { + match delta_ev.delta_type.as_deref() { + Some("text_delta") => { + if let Some(ref delta) = delta_ev.delta { + reply.push_str(delta); + } + } + Some("text_end") => { + if let Some(ref content) = delta_ev.content { + reply = content.clone(); // 完整内容覆盖增量 + } + } + _ => {} + } + } + } + "agent_end" => { + done = true; + } + _ => {} + } + } + + // 确保子进程退出 + let _ = child.wait().await; + + if reply.trim().is_empty() { + return Err("pi 子代理返回空回复".to_string()); + } + + // 截断过长回复 + if reply.chars().count() > MAX_REPLY_CHARS { + reply = format!( + "{}…（已截断，完整输出 {} 字符）", + reply.chars().take(MAX_REPLY_CHARS).collect::(), + reply.chars().count() + ); + } + + Ok(reply) +} + +/// 查找 pi 可执行文件路径 +fn find_pi_exe() -> Result { + // 1. 检查环境变量 PI_EXE + if let Ok(path) = std::env::var("PI_EXE") { + if std::path::Path::new(&path).exists() { + return Ok(path); + } + } + + // 2. 在 PATH 中搜索 + if let Ok(path_var) = std::env::var("PATH") { + for dir in path_var.split(':') { + let candidate = std::path::PathBuf::from(dir).join("pi"); + if candidate.exists() { + return Ok(candidate.to_string_lossy().to_string()); + } + } + } + + // 3. 检查常见全局路径 + let candidates = [ + "/home/xiao/.npm/bin/pi", + "/usr/local/bin/pi", + "/usr/bin/pi", + ]; + for c in &candidates { + if std::path::Path::new(c).exists() { + return Ok(c.to_string()); + } + } + + Err("未找到 pi 可执行文件。请设置 PI_EXE 环境变量或确保 pi 在 PATH 中。".to_string()) +} diff --git a/src/tools/executor.rs b/src/tools/executor.rs new file mode 100644 index 0000000..8f09676 --- /dev/null +++ b/src/tools/executor.rs @@ -0,0 +1,336 @@ +//! 统一工具执行器构建器 +//! +//! 消除 main.rs / worker.rs / daemon.rs 中三处重复的 ToolExecutor 构建逻辑。 +//! 支持两种模式： +//! - Listen 模式：有 ApprovalManager + WechatSender + DB +//! - Worker 模式：无 DB 连接，使用本地缓存（memories/summaries） + +use crate::core::memory::MemoryStore; +use crate::core::session::ChatSession; +use crate::llm::conversation::ToolExecutor; +use crate::tools::approval::ApprovalManager; +use crate::tools::types::{SkillResult, WechatSender}; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::Mutex; + +/// 执行器构建器 +pub struct ToolExecutorBuilder { + /// 审批管理器（Listen 模式） + approval: Option>, + /// 微信发送器（Listen 模式） + send_wechat: Option, + /// 会话引用 + session: Option>>, + /// 记忆管理器 + memory_store: Option>, + /// Worker 模式：已批准的工具名 + approved_tool: Option, + /// Worker 模式：预载记忆 + memories_cache: HashMap>, + /// Worker 模式：预载摘要 + summaries_cache: Vec, + /// Worker 模式：待发送的审批请求 + pending_approval: Option>>>, + /// Worker 模式：新写入的记忆 + new_memories: Option>>>, + /// 当前用户 ID + user_id: String, +} + +impl ToolExecutorBuilder { + pub fn new(user_id: impl Into) -> Self { + Self { + approval: None, + send_wechat: None, + session: None, + memory_store: None, + approved_tool: None, + memories_cache: HashMap::new(), + summaries_cache: Vec::new(), + pending_approval: None, + new_memories: None, + user_id: user_id.into(), + } + } + + pub fn with_approval(mut self, mgr: Arc) -> Self { + self.approval = Some(mgr); + self + } + + pub fn with_wechat(mut self, sender: WechatSender) -> Self { + self.send_wechat = Some(sender); + self + } + + pub fn with_session(mut self, session: Arc>) -> Self { + self.session = Some(session); + self + } + + pub fn with_memory(mut self, store: Arc) -> Self { + self.memory_store = Some(store); + self + } + + /// Worker 模式：注入预载上下文 + pub fn with_worker_context( + mut self, + approved_tool: Option, + memories: HashMap>, + summaries: Vec, + pending_approval: Arc>>, + new_memories: Arc>>, + ) -> Self { + self.approved_tool = approved_tool; + self.memories_cache = memories; + self.summaries_cache = summaries; + self.pending_approval = Some(pending_approval); + self.new_memories = Some(new_memories); + self + } + + /// 构建最终 ToolExecutor + pub fn build(self) -> ToolExecutor { + let approval = self.approval; + let send_wechat = self.send_wechat; + let session = self.session; + let memory_store = self.memory_store; + let approved_tool = self.approved_tool; + let memories_cache = self.memories_cache; + let summaries_cache = self.summaries_cache; + let pending_approval = self.pending_approval; + let new_memories = self.new_memories; + let user_id = self.user_id; + let is_worker = pending_approval.is_some(); + + Arc::new(move |name: &str, args_json: &str| { + let approval = approval.clone(); + let send_wechat = send_wechat.clone(); + let session = session.clone(); + let memory_store = memory_store.clone(); + let approved_tool = approved_tool.clone(); + let memories_cache = memories_cache.clone(); + let summaries_cache = summaries_cache.clone(); + let pending_approval = pending_approval.clone(); + let new_memories = new_memories.clone(); + let user_id = user_id.clone(); + let is_worker_flag = is_worker; + + let n = name.to_string(); + let args = args_json.to_string(); + + Box::pin(async move { + // listen 模式下从 session 动态获取当前用户 + let effective_user = if is_worker_flag { + user_id.clone() + } else if let Some(ref s) = session { + s.lock().await.current_user_id.clone() + } else { + user_id.clone() + }; + + // ── 上下文工具：read_memories ── + if n == "read_memories" { + return if is_worker_flag { + read_worker_memories(&effective_user, &memories_cache).await + } else if let Some(store) = &memory_store { + Ok(store.read_for(&effective_user).await) + } else { + Ok("暂无长期记忆".to_string()) + }; + } + + // ── 上下文工具：write_memory ── + if n == "write_memory" { + let params: serde_json::Value = + serde_json::from_str(&args).unwrap_or_default(); + let content = + params.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if content.is_empty() { + return Ok("请提供 content 参数".to_string()); + } + if let Some(ref new_mems) = new_memories { + new_mems.lock().await.push(content.to_string()); + } + if let Some(store) = &memory_store { + return Ok(store.write_for(&effective_user, content).await); + } + return Ok(format!("已添加长期记忆: {}", content)); + } + + // ── 上下文工具：read_summaries ── + if n == "read_summaries" { + if is_worker_flag { + return read_worker_summaries(&summaries_cache).await; + } + if let Some(ref s) = session { + return Ok(crate::core::memory::read_summaries(s).await); + } + return Ok("暂无历史摘要".to_string()); + } + + // ── query_capabilities ── + if n == "query_capabilities" { + let specs = crate::tools::registry::BuiltinRegistry::specs(); + return Ok(crate::tools::build_capability_guide(&specs)); + } + + // ── 解析目标工具名和参数 ── + let (target_name, target_args) = if n == "call_capability" { + let cp: serde_json::Value = + serde_json::from_str(&args).unwrap_or_default(); + let name = cp + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let unpacked = crate::tools::unpack_call_params(&cp); + (name, serde_json::to_string(&unpacked).unwrap_or_default()) + } else { + (n.clone(), args.clone()) + }; + + if target_name.is_empty() { + return Ok(format!("未知工具: {}", n)); + } + + // ── 高风险工具审批 ── + if crate::tools::registry::BuiltinRegistry::is_high_risk(&target_name) { + if is_worker_flag { + // Worker 模式：已批准则放行一次，否则请求审批 + let is_approved = approved_tool.as_deref() == Some(&target_name); + if !is_approved { + if let Some(ref pa) = pending_approval { + let reason = format!( + "LLM 请求执行高风险工具: {}", + target_name + ); + *pa.lock().await = + Some((target_name.clone(), target_args.clone(), reason)); + } + return Ok(SkillResult::rejected(&target_name).output); + } + } else if let Some(ref mgr) = approval { + // Listen 模式：通过 ApprovalManager 处理 + let approved = listen_approve( + mgr, + &effective_user, + &target_name, + send_wechat.as_ref(), + ) + .await; + match approved { + Ok(true) => {} + Ok(false) => { + return Ok(SkillResult::rejected(&target_name).output) + } + Err(e) => return Ok(e), + } + } + } + + // ── 执行内置工具 ── + if let Some(result) = + crate::tools::registry::BuiltinRegistry::execute( + &target_name, + &target_args, + ) + .await + { + // manage_memos add 同步到记忆缓存 + if target_name == "manage_memos" { + let cp: serde_json::Value = + serde_json::from_str(&target_args).unwrap_or_default(); + if cp.get("action").and_then(|v| v.as_str()) == Some("add") { + if let Some(content) = + cp.get("content").and_then(|v| v.as_str()) + { + if let Some(ref new_mems) = new_memories { + new_mems.lock().await.push(content.to_string()); + } + if let Some(store) = &memory_store { + store.write_for(&effective_user, content).await; + } + } + } + } + return Ok(result.output); + } + + Ok(format!("未知工具: {}", target_name)) + }) + }) + } +} + +/// Worker 模式：读取记忆缓存 +async fn read_worker_memories( + user_id: &str, + cache: &HashMap>, +) -> Result { + match cache.get(user_id) { + Some(m) if !m.is_empty() => { + let lines: Vec = m + .iter() + .enumerate() + .map(|(i, v)| format!("{}. {}", i + 1, v)) + .collect(); + Ok(lines.join("\n")) + } + _ => Ok("暂无长期记忆".to_string()), + } +} + +/// Worker 模式：读取摘要缓存 +async fn read_worker_summaries(summaries: &[String]) -> Result { + if summaries.is_empty() { + Ok("暂无历史摘要".to_string()) + } else { + Ok(summaries + .iter() + .enumerate() + .map(|(i, s)| format!("{}. {}", i + 1, s)) + .collect::>() + .join("\n---\n")) + } +} + +/// Listen 模式：通过 ApprovalManager 处理高风险工具审批 +/// +/// 注意：listen 模式共享全局会话锁，审批期间所有用户消息都会排队。 +/// 生产环境建议使用 daemon 模式（`ias daemon`），审批不阻塞其他用户。 +async fn listen_approve( + mgr: &Arc, + user_id: &str, + tool_name: &str, + send_wechat: Option<&WechatSender>, +) -> Result { + use crate::tools::approval::ApprovalDecision; + use sha2::{Digest, Sha256}; + + let (code, rx) = mgr.create(user_id, tool_name).await?; + let code_hash = format!("{:x}", Sha256::digest(code.as_bytes())); + + let msg = format!( + "⚠️ 操作确认\n\n技能：{}\n\n确认码：{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(60秒内有效，最多3次尝试)", + tool_name, code + ); + if let Some(ref send) = send_wechat { + if let Err(e) = send(user_id, &msg).await { + return Err(e); + } + } + + // listen 模式使用较短超时（60s），避免长时间持有全局会话锁 + match tokio::time::timeout(std::time::Duration::from_secs(60), rx).await { + Ok(Ok(ApprovalDecision::Approved)) => Ok(true), + Ok(_) => Ok(false), // Rejected 或 sender dropped + Err(_elapsed) => { + // 超时：精确取消该条审批，防止迟到确认码被误吞 + 同步 DB + mgr.cancel_by_code(user_id, &code_hash).await; + Ok(false) + } + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index a57e61f..73817a6 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -1,6 +1,7 @@ pub mod approval; -pub mod builtin; pub mod builtins; +pub mod executor; +pub mod registry; pub mod types; /// 从 call_capability 的 `{name, prompt}` JSON 中提取目标工具的真实参数。 @@ -75,6 +76,7 @@ mod tests { } } +use std::collections::HashMap; use types::SkillSpec; /// 构建工具调用指南（含优先级说明） @@ -124,6 +126,189 @@ pub fn build_capability_guide(specs: &[SkillSpec]) -> String { lines.join("\n") } +/// 构建所有 LLM 可见工具的 JSON 定义列表 +/// 供 main.rs cmd_listen 和 daemon.rs 共用 +pub fn build_tools_list() -> Vec { + let mut list = Vec::new(); + + // query_capabilities + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "query_capabilities", + "description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。", + "parameters": { "type": "object", "properties": {} } + } + })); + + // call_capability + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "call_capability", + "description": "调用一个已注册的工具。name 为工具名（来自 query_capabilities），工具所需的具体参数应直接放在 JSON 中。", + "parameters": { + "type": "object", + "properties": { + "name": { "type": "string", "description": "工具名称" }, + "prompt": { "type": "string", "description": "json格式参数" } + }, + "required": ["name"] + } + } + })); + + // 上下文工具 + list.push(serde_json::json!({ + "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"] + } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "read_summaries", + "description": "读取历史会话摘要", + "parameters": { "type": "object", "properties": {} } + } + })); + + // 业务工具 + 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"] + } + } + })); + + // pi 子代理 + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "pi_subagent", + "description": "启动独立的 pi agent 子进程处理复杂子任务。传入 prompt 让子 agent 独立执行，返回完整回复。适合将复杂分析、代码生成、多步推理等任务委派给独立 agent。", + "parameters": { + "type": "object", + "properties": { + "prompt": {"type": "string", "description": "发送给子 agent 的完整 prompt"}, + "timeout_secs": {"type": "integer", "description": "超时秒数（1-300，默认 120）"} + }, + "required": ["prompt"] + } + } + })); + + list +} + +/// 构建传递给 Worker 的环境变量映射 +/// 供 daemon.rs 使用 +pub 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 +} + fn format_spec(lines: &mut Vec, spec: &SkillSpec) { let risk = if spec.risk_level == crate::tools::types::RiskLevel::High { " ⚠️需确认" diff --git a/src/tools/builtin.rs b/src/tools/registry.rs similarity index 87% rename from src/tools/builtin.rs rename to src/tools/registry.rs index 6ffb97d..a7aead9 100644 --- a/src/tools/builtin.rs +++ b/src/tools/registry.rs @@ -23,6 +23,11 @@ impl BuiltinRegistry { serde_json::from_str(args_json).unwrap_or_default(); Some(super::builtins::fetch_page::execute(params).await) } + "pi_subagent" => { + let params: serde_json::Value = + serde_json::from_str(args_json).unwrap_or_default(); + Some(super::builtins::subagent::execute(params).await) + } n if n.starts_with("amap_") => { let params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default(); @@ -40,6 +45,7 @@ impl BuiltinRegistry { super::builtins::weather::spec(), super::builtins::web_search::spec(), super::builtins::fetch_page::spec(), + super::builtins::subagent::spec(), ]; specs.extend(super::builtins::amap::specs()); specs @@ -51,6 +57,7 @@ impl BuiltinRegistry { .any(|s| s.name == name && s.risk_level == RiskLevel::High) } + #[allow(dead_code)] pub fn is_builtin(name: &str) -> bool { Self::specs().iter().any(|s| s.name == name) } diff --git a/src/tools/types.rs b/src/tools/types.rs index c0bb9a8..af1248e 100644 --- a/src/tools/types.rs +++ b/src/tools/types.rs @@ -95,12 +95,14 @@ pub type WechatSender = Arc< dyn Fn(&str, &str) -> Pin> + Send>> + Send + Sync, >; +#[allow(dead_code)] pub struct ExecutionContext { pub user_id: String, pub approval_manager: Option>, pub send_wechat: Option, } +#[allow(dead_code)] impl ExecutionContext { pub fn new(user_id: impl Into) -> Self { Self { diff --git a/src/worker.rs b/src/worker.rs index f7b4059..764764f 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -9,7 +9,7 @@ use crate::ipc::{OutputFrame, TaskFrame, recv_task, send_frame}; use crate::llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT}; -use crate::tools::types::SkillResult; +use crate::tools::executor::ToolExecutorBuilder; use std::collections::HashMap; use std::sync::Arc; use tokio::net::UnixStream; @@ -17,18 +17,10 @@ use tracing::info; /// Worker 与工具执行器之间的共享状态 struct WorkerShared { - /// 待发送的审批请求 (tool, reason) - pending_approval: tokio::sync::Mutex>, - /// 本次对话中已由用户批准的工具名（避免重复审批） - approved_tool: tokio::sync::Mutex>, - /// 本地记忆缓存 - memories: tokio::sync::Mutex>>, + /// 待发送的审批请求 (tool, args, reason) + pending_approval: Arc>>, /// 本次写入的新记忆（退出前发送 DbWrite 帧同步到 daemon） - new_memories: tokio::sync::Mutex>, - /// 预载的摘要 - summaries: Vec, - /// 用户 ID - user_id: String, + new_memories: Arc>>, } /// Worker 入口: 连接 daemon，处理一条消息后退出 @@ -87,21 +79,60 @@ pub async fn run(sock_path: &str) -> Result<(), String> { // 5. 共享状态 let approved_tool = task.approved_tool.clone().filter(|t| !t.is_empty()); let shared = Arc::new(WorkerShared { - pending_approval: tokio::sync::Mutex::new(None), - approved_tool: tokio::sync::Mutex::new(approved_tool), - memories: tokio::sync::Mutex::new(HashMap::from([( - task.user_id.clone(), - task.memories.clone(), - )])), - new_memories: tokio::sync::Mutex::new(Vec::new()), - summaries: task.summaries.clone(), - user_id: task.user_id.clone(), + pending_approval: Arc::new(tokio::sync::Mutex::new(None)), + new_memories: Arc::new(tokio::sync::Mutex::new(Vec::new())), }); - // 6. 构建工具执行器 - let executor = build_worker_executor(conv.session(), shared.clone()); + // 6. 构建工具执行器（统一使用 ToolExecutorBuilder） + // 注意：若审批已通过，approved_tool 不传给 executor（已在 6b 直接执行） + let pending_approval = Arc::clone(&shared.pending_approval); + let new_memories = Arc::clone(&shared.new_memories); + let has_approved = task.approved_tool.as_deref().filter(|t| !t.is_empty()).is_some() + && task.approved_tool_args.as_deref().filter(|a| !a.is_empty()).is_some(); + let executor = ToolExecutorBuilder::new(task.user_id.clone()) + .with_worker_context( + if has_approved { None } else { approved_tool }, + HashMap::from([(task.user_id.clone(), task.memories.clone())]), + task.summaries.clone(), + pending_approval, + new_memories, + ) + .build(); conv.set_tool_executor(executor); + // 6b. 若审批已通过且带有参数，直接执行工具并注入结果 + if let (Some(tool_name), Some(tool_args)) = + (&task.approved_tool, &task.approved_tool_args) + { + if !tool_name.is_empty() && !tool_args.is_empty() { + info!("审批通过，直接执行工具: {} args={:.80}", tool_name, tool_args); + let tool_id = format!("approved_{}", uuid::Uuid::new_v4()); + // 手动执行工具 + let result_text: String = match crate::tools::registry::BuiltinRegistry::execute( + tool_name, tool_args, + ) + .await + { + Some(r) => r.output, + None => format!("未知工具: {}", tool_name), + }; + // 注入到 session：先加 assistant tool_call，再加 tool_result + { + let session = conv.session(); + let mut s = session.lock().await; + s.add_assistant_tool_calls(vec![crate::llm::types::ToolCall { + id: tool_id.clone(), + call_type: "function".to_string(), + function: crate::llm::types::ToolFunctionCall { + name: tool_name.clone(), + arguments: tool_args.clone(), + }, + }]); + s.add_tool_result(&tool_id, tool_name, &result_text); + } + } + } + // 7. 执行 LLM 对话（带审批检测） let (reply, _used_tools, usage) = conv .chat_with_tools(task.msg.text.clone()) @@ -110,9 +141,17 @@ pub async fn run(sock_path: &str) -> Result<(), String> { // 8. 检查是否有审批请求（由工具执行器设置） let approval_req = shared.pending_approval.lock().await.take(); - if let Some((tool, reason)) = approval_req { + if let Some((tool, tool_args, reason)) = approval_req { info!("Worker 发送审批请求: tool={}", tool); - let _ = send_frame(&mut stream, &OutputFrame::RequestApproval { tool, reason }).await; + let _ = send_frame( + &mut stream, + &OutputFrame::RequestApproval { + tool, + reason, + tool_args, + }, + ) + .await; let _ = send_frame(&mut stream, &OutputFrame::Bye).await; return Ok(()); } @@ -167,125 +206,3 @@ pub async fn run(sock_path: &str) -> Result<(), String> { Ok(()) } -/// 构建 Worker 侧的工具执行器（无 DB 连接） -fn build_worker_executor( - _session: Arc>, - shared: Arc, -) -> crate::llm::conversation::ToolExecutor { - Arc::new(move |name: &str, args_json: &str| { - let shared = shared.clone(); - let n = name.to_string(); - let args = args_json.to_string(); - - Box::pin(async move { - match n.as_str() { - "read_memories" => { - let cache = shared.memories.lock().await; - let mems = cache.get(&shared.user_id); - match mems { - Some(m) if !m.is_empty() => { - let lines: Vec = m - .iter() - .enumerate() - .map(|(i, v)| format!("{}. {}", i + 1, v)) - .collect(); - Ok(lines.join("\n")) - } - _ => Ok("暂无长期记忆".to_string()), - } - } - "write_memory" => { - let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); - let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); - if content.is_empty() { - return Ok("请提供 content 参数".to_string()); - } - shared - .memories - .lock() - .await - .entry(shared.user_id.clone()) - .or_default() - .push(content.to_string()); - shared.new_memories.lock().await.push(content.to_string()); - Ok(format!("已添加长期记忆: {}", content)) - } - "read_summaries" => { - if shared.summaries.is_empty() { - Ok("暂无历史摘要".to_string()) - } else { - Ok(shared - .summaries - .iter() - .enumerate() - .map(|(i, s)| format!("{}. {}", i + 1, s)) - .collect::>() - .join("\n---\n")) - } - } - "query_capabilities" => { - let specs = crate::tools::builtin::BuiltinRegistry::specs(); - Ok(crate::tools::build_capability_guide(&specs)) - } - // call_capability 或直接调用 - _ => { - let (target_name, target_args) = if n == "call_capability" { - let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); - let name = cp - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let unpacked = crate::tools::unpack_call_params(&cp); - (name, serde_json::to_string(&unpacked).unwrap_or_default()) - } else { - (n.clone(), args.clone()) - }; - - if target_name.is_empty() { - return Ok(format!("未知工具: {}", n)); - } - - // 高风险工具 → 请求审批（除非已被用户批准） - if crate::tools::builtin::BuiltinRegistry::is_high_risk(&target_name) { - let is_approved = shared.approved_tool.lock().await.as_deref() == Some(&target_name); - if !is_approved { - let reason = format!("LLM 请求执行高风险工具: {}", target_name); - *shared.pending_approval.lock().await = Some((target_name.clone(), reason)); - return Ok(SkillResult::rejected(&target_name).output); - } - // 批准过 → 执行一次后清除（同一对话内同一工具不会再被拦截） - *shared.approved_tool.lock().await = None; - } - - // 执行内置工具 - if let Some(result) = - crate::tools::builtin::BuiltinRegistry::execute(&target_name, &target_args) - .await - { - if target_name == "manage_memos" { - let cp: serde_json::Value = - serde_json::from_str(&target_args).unwrap_or_default(); - if cp.get("action").and_then(|v| v.as_str()) == Some("add") { - if let Some(content) = cp.get("content").and_then(|v| v.as_str()) { - shared - .memories - .lock() - .await - .entry(shared.user_id.clone()) - .or_default() - .push(content.to_string()); - shared.new_memories.lock().await.push(content.to_string()); - } - } - } - return Ok(result.output); - } - - Ok(format!("未知工具: {}", target_name)) - } - } - }) - }) -} -