重构:架构重组 + 6轮缺陷修复

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
This commit is contained in:
2026-06-04 16:18:06 +08:00
parent 8dc232ebab
commit af0620aa2d
25 changed files with 1405 additions and 890 deletions
+27 -22
View File
@@ -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 ChatSessioncheckpoint/messages/summaries
│ ├── builder.rs Token budget 估算 + 上下文构建 + 摘要触发
── tools.rs MemoryStoreHashMap<user_id, Vec<mem>>
├── db/ 数据库
│ ├── mod.rs Database 连接池 + 迁移
│ └── models.rs CRUDauth/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 Conversationchat + 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 CRUDauth/chat/usage/summary
├── daemon.rs 守护进程(持有 WeChat + DBspawn worker
├── worker.rs Worker 进程(无状态 LLM 执行)
└── ipc.rs Unix Domain Socket 帧协议
+15 -4
View File
@@ -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)
}
-5
View File
@@ -1,5 +0,0 @@
pub mod builder;
pub mod tools;
pub mod types;
pub use tools::MemoryStore;
+112
View File
@@ -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<Mutex<ChatSession>>,
system_prompt: &str,
) -> Vec<Message> {
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);
}
}
+3 -3
View File
@@ -63,15 +63,15 @@ impl MemoryStore {
/// read_summaries 工具:读取历史摘要(内存 + 数据库)
pub async fn read_summaries(
session: &Arc<Mutex<super::types::ChatSession>>,
session: &Arc<Mutex<super::session::ChatSession>>,
) -> String {
let s = session.lock().await;
let mut lines: Vec<String> = 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!(
"[{}] {} (原因: {})",
+7
View File
@@ -0,0 +1,7 @@
pub mod context;
pub mod memory;
pub mod pipeline;
pub mod session;
pub mod summary;
pub use memory::MemoryStore;
+285
View File
@@ -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<Arc<Database>>,
pub file_state: state::StateManager,
pub memory_store: Arc<MemoryStore>,
pub approval_manager: Arc<ApprovalManager>,
pub account_id: String,
/// 全局会话锁:所有消息共享同一 Conversation,必须串行处理防止跨用户串话
pub conversation_lock: Arc<Mutex<()>>,
}
impl PipelineContext {
/// 监听循环:长轮询微信消息,处理每条消息
pub async fn run_listen_loop(
&self,
enable_llm: bool,
echo: bool,
conversation: Option<Arc<Conversation>>,
) {
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<Arc<Database>>,
) -> Result<String, String> {
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<Arc<Database>>,
) -> Result<String, String> {
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<Arc<Database>>,
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);
}
}
}
+9 -112
View File
@@ -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<Mutex<ChatSession>>, system_prompt: &str) -> Vec<Message> {
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<Mutex<ChatSession>>,
@@ -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<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 = 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);
}
}
+41 -177
View File
@@ -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<String>,
}
#[derive(Clone)]
struct PendingMessage {
from: String,
text: String,
@@ -36,6 +37,7 @@ struct PendingMessage {
#[allow(dead_code)]
message_id: String,
approved_tool: Option<String>,
approved_tool_args: Option<String>,
}
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<String, String>,
system_prompt: String,
model: String,
/// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 创建时间)
approval_ctx: Mutex<HashMap<String, (String, Option<String>, String, Instant)>>,
/// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 工具参数, 创建时间)
approval_ctx: Mutex<HashMap<String, (String, Option<String>, 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); // 先取消 activepush_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<serde_json::Value> {
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<String, String> {
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
}
+6
View File
@@ -22,6 +22,9 @@ pub struct TaskFrame {
pub summaries: Vec<String>,
#[serde(default)]
pub approved_tool: Option<String>,
/// 审批通过后的工具调用参数
#[serde(default)]
pub approved_tool_args: Option<String>,
pub env: std::collections::HashMap<String, String>,
pub tools: Vec<serde_json::Value>,
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")]
+18 -11
View File
@@ -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!("📦 上下文超预算,已触发溢出摘要");
}
}
+1 -1
View File
@@ -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};
+29 -407
View File
@@ -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<serde_json::Value> = 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<Arc<Database>>,
file_state: &state::StateManager,
conversation: Option<Arc<Conversation>>,
approval_manager: &ApprovalManager,
memory_store: &Arc<MemoryStore>,
) {
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<Arc<Database>>,
) -> Result<String, String> {
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<Arc<Database>>,
) -> Result<String, String> {
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<Arc<Database>>,
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<bool, String> {
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 => {
+1
View File
@@ -16,6 +16,7 @@ pub struct RuntimeState {
}
/// 状态管理器
#[derive(Clone)]
pub struct StateManager {
state_dir: PathBuf,
}
+26
View File
@@ -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 被 droprx 端收到 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();
+1
View File
@@ -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;
+229
View File
@@ -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<EventMessage>,
#[serde(default)]
#[serde(rename = "assistantMessageEvent")]
assistant_message_event: Option<AssistantMessageEvent>,
}
#[derive(Deserialize, Debug)]
#[allow(dead_code)]
struct EventMessage {
#[serde(default)]
content: Option<serde_json::Value>,
}
/// message_update 中的嵌套增量事件
#[derive(Deserialize, Debug)]
struct AssistantMessageEvent {
#[serde(rename = "type")]
delta_type: Option<String>,
#[serde(default)]
delta: Option<String>,
#[serde(default)]
content: Option<String>,
}
// ─── 执行 ───
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<String, String> {
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::<String>(),
reply.chars().count()
);
}
Ok(reply)
}
/// 查找 pi 可执行文件路径
fn find_pi_exe() -> Result<String, String> {
// 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())
}
+336
View File
@@ -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<Arc<ApprovalManager>>,
/// 微信发送器(Listen 模式)
send_wechat: Option<WechatSender>,
/// 会话引用
session: Option<Arc<Mutex<ChatSession>>>,
/// 记忆管理器
memory_store: Option<Arc<MemoryStore>>,
/// Worker 模式:已批准的工具名
approved_tool: Option<String>,
/// Worker 模式:预载记忆
memories_cache: HashMap<String, Vec<String>>,
/// Worker 模式:预载摘要
summaries_cache: Vec<String>,
/// Worker 模式:待发送的审批请求
pending_approval: Option<Arc<Mutex<Option<(String, String, String)>>>>,
/// Worker 模式:新写入的记忆
new_memories: Option<Arc<Mutex<Vec<String>>>>,
/// 当前用户 ID
user_id: String,
}
impl ToolExecutorBuilder {
pub fn new(user_id: impl Into<String>) -> 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<ApprovalManager>) -> 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<Mutex<ChatSession>>) -> Self {
self.session = Some(session);
self
}
pub fn with_memory(mut self, store: Arc<MemoryStore>) -> Self {
self.memory_store = Some(store);
self
}
/// Worker 模式:注入预载上下文
pub fn with_worker_context(
mut self,
approved_tool: Option<String>,
memories: HashMap<String, Vec<String>>,
summaries: Vec<String>,
pending_approval: Arc<Mutex<Option<(String, String, String)>>>,
new_memories: Arc<Mutex<Vec<String>>>,
) -> 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<String, Vec<String>>,
) -> Result<String, String> {
match cache.get(user_id) {
Some(m) if !m.is_empty() => {
let lines: Vec<String> = 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<String, String> {
if summaries.is_empty() {
Ok("暂无历史摘要".to_string())
} else {
Ok(summaries
.iter()
.enumerate()
.map(|(i, s)| format!("{}. {}", i + 1, s))
.collect::<Vec<_>>()
.join("\n---\n"))
}
}
/// Listen 模式:通过 ApprovalManager 处理高风险工具审批
///
/// 注意:listen 模式共享全局会话锁,审批期间所有用户消息都会排队。
/// 生产环境建议使用 daemon 模式(`ias daemon`),审批不阻塞其他用户。
async fn listen_approve(
mgr: &Arc<ApprovalManager>,
user_id: &str,
tool_name: &str,
send_wechat: Option<&WechatSender>,
) -> Result<bool, String> {
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)
}
}
}
+186 -1
View File
@@ -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<serde_json::Value> {
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<String, String> {
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<String>, spec: &SkillSpec) {
let risk = if spec.risk_level == crate::tools::types::RiskLevel::High {
" ⚠️需确认"
@@ -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)
}
+2
View File
@@ -95,12 +95,14 @@ pub type WechatSender = Arc<
dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync,
>;
#[allow(dead_code)]
pub struct ExecutionContext {
pub user_id: String,
pub approval_manager: Option<Arc<crate::tools::approval::ApprovalManager>>,
pub send_wechat: Option<WechatSender>,
}
#[allow(dead_code)]
impl ExecutionContext {
pub fn new(user_id: impl Into<String>) -> Self {
Self {
+64 -147
View File
@@ -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<Option<(String, String)>>,
/// 本次对话中已由用户批准的工具名(避免重复审批)
approved_tool: tokio::sync::Mutex<Option<String>>,
/// 本地记忆缓存
memories: tokio::sync::Mutex<HashMap<String, Vec<String>>>,
/// 待发送的审批请求 (tool, args, reason)
pending_approval: Arc<tokio::sync::Mutex<Option<(String, String, String)>>>,
/// 本次写入的新记忆(退出前发送 DbWrite 帧同步到 daemon
new_memories: tokio::sync::Mutex<Vec<String>>,
/// 预载的摘要
summaries: Vec<String>,
/// 用户 ID
user_id: String,
new_memories: Arc<tokio::sync::Mutex<Vec<String>>>,
}
/// 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<tokio::sync::Mutex<crate::context::types::ChatSession>>,
shared: Arc<WorkerShared>,
) -> 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<String> = 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::<Vec<_>>()
.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))
}
}
})
})
}