wunianxiao
e20cd19e57
- daemon: 添加 LLM 回合开始/结束标记日志，工具执行结果追踪 - daemon: 发送消费者日志增强（显示完整文本、截断用户 ID） - conversation: 移除日志截断，输出完整工具调用参数和结果 - deepseek: 代码格式化，预留请求/响应摘要日志 - 删除 systemd/ias.service
1029 lines
38 KiB
Rust
1029 lines
38 KiB
Rust
//! ## Daemon —— 常驻守护进程（核心循环）
|
||
//!
|
||
//! 这是整个项目的核心入口循环。它的职责：
|
||
//!
|
||
//! 1. **持有 WeChat 长连接** — 通过 `WeChatClient` 长轮询接收微信消息
|
||
//! 2. **持有数据库连接** — 可选 PostgreSQL，回退到文件存储
|
||
//! 3. **管理审批流程** — `ApprovalManager` 处理高风险工具的确认码审批
|
||
//! 4. **管理长期记忆** — `MemoryStore` 提供按用户隔离的记忆读写
|
||
//!
|
||
//! ### 消息处理架构（三消费者 + 公平队列）
|
||
//!
|
||
//! daemon 将消息处理拆分为 3 个独立的 tokio async task（消费者），
|
||
//! 通过 `MessageQueue` + `QueueRunner` 做公平轮转路由：
|
||
//!
|
||
//! ```text
|
||
//! receive_messages → enqueue PipelineMessage
|
||
//! │
|
||
//! ▼
|
||
//! ┌──────────────────┐
|
||
//! │ MessageQueue │ ← 公平轮转（A1, B1, A2, A3）
|
||
//! └────────┬─────────┘
|
||
//! │ dequeue
|
||
//! ▼
|
||
//! ┌──────────────────┐
|
||
//! │ QueueRunner │ ← 按 MessageKind 路由
|
||
//! └──┬──────┬──────┬─┘
|
||
//! │ │ │
|
||
//! ▼ ▼ ▼
|
||
//! LLM Tool Send
|
||
//! Cons. Cons. Cons.
|
||
//! ```
|
||
//!
|
||
//! - **LLM Consumer** — 调用 DeepSeek API 做对话 + 工具循环
|
||
//! - **Tool Consumer** — 执行内置工具，高风险工具先走审批
|
||
//! - **Send Consumer** — 将 LLM 回复通过 WeChat API 发回给用户
|
||
//!
|
||
//! ### 架构
|
||
//!
|
||
//! 采用统一的单进程 tokio task 架构。旧版 daemon+worker 分离模式（UDS IPC）已在 v2 中移除。
|
||
|
||
use crate::channel::ChannelId;
|
||
use crate::context::HistoryEntry;
|
||
use crate::context::MemoryStore;
|
||
use crate::db::Database;
|
||
use crate::llm::{
|
||
ASYNC_MARKER, ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor,
|
||
};
|
||
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
|
||
|
||
use crate::tools::approval::{ApprovalDecision, ApprovalManager};
|
||
use crate::wechat::client::WeChatClient;
|
||
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
use std::time::Instant;
|
||
use tokio::sync::Mutex;
|
||
use tracing::{error, info, warn};
|
||
use uuid::Uuid;
|
||
|
||
/// ## Daemon 上下文 —— 消费者之间共享的全部状态
|
||
///
|
||
/// 这个结构体被所有三个消费者（LLM / Tool / Send）通过 `Arc` 共享访问。
|
||
/// 每个字段都只读（或内部有锁），所以不需要额外的 `Mutex` 包装。
|
||
struct DaemonCtx {
|
||
/// 数据库句柄
|
||
db: Arc<Database>,
|
||
/// WeChat API 客户端（内部用 Arc<Mutex> 保护 token/updates_buf）
|
||
client: WeChatClient,
|
||
/// 本机器人的微信账号 ID
|
||
account_id: String,
|
||
/// 审批管理器 —— 处理高风险工具的用户确认码流程
|
||
approval: Arc<ApprovalManager>,
|
||
/// 长期记忆管理器 —— 按用户隔离的记忆读写
|
||
memory_store: Arc<MemoryStore>,
|
||
/// 注册给 LLM 的工具列表（JSON 格式，符合 function calling 规范）
|
||
tools_list: Vec<serde_json::Value>,
|
||
/// 系统提示词，覆盖默认值
|
||
system_prompt: String,
|
||
/// 使用的模型名称（如 `deepseek-v4-flash`）
|
||
model: String,
|
||
/// 待审批的上下文：`user_id → (原始文本, context_token, 工具名, 工具参数, 创建时间)`
|
||
/// 用户在聊天中发送确认码时，从这里面查找匹配项
|
||
approval_ctx: Mutex<HashMap<String, (String, Option<String>, String, String, Instant)>>,
|
||
}
|
||
|
||
/// Daemon 入口
|
||
///
|
||
/// `sock_path` 保留用于兼容旧版 daemon 命令的 CLI 接口，实际已不再使用。
|
||
pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
|
||
// 1-3. 认证 & 客户端
|
||
let auth = match load_auth(database).await {
|
||
Some(a) => a,
|
||
None => {
|
||
error!("未登录");
|
||
return;
|
||
}
|
||
};
|
||
let mut client = WeChatClient::new(Some(auth.base_url.clone()));
|
||
client
|
||
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
|
||
.await;
|
||
if let Some(buf) = crate::db::models::load_runtime(database.pool()).await {
|
||
client.set_updates_buf(&buf).await;
|
||
}
|
||
if let Err(e) = client.notify_start().await {
|
||
error!("注册监听器失败: {}", e);
|
||
return;
|
||
}
|
||
let account_id = auth.account_id.clone();
|
||
info!("Daemon 已启动");
|
||
|
||
// 4-6. 审批管理器、工具列表、模型配置
|
||
let approval_manager = Arc::new(ApprovalManager::new(Some(Arc::new(
|
||
database.pool().clone(),
|
||
))));
|
||
let tools_list = build_tools_list();
|
||
let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new());
|
||
let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
|
||
|
||
// 7. Daemon 上下文（注意 approval_ctx 为 5 元组）
|
||
let ctx = Arc::new(DaemonCtx {
|
||
db: database.clone(),
|
||
client: client.clone(),
|
||
account_id: account_id.clone(),
|
||
approval: approval_manager.clone(),
|
||
memory_store: memory_store.clone(),
|
||
tools_list: tools_list.clone(),
|
||
system_prompt,
|
||
model,
|
||
approval_ctx: Mutex::new(HashMap::new()),
|
||
});
|
||
|
||
// 8. 审批过期清理
|
||
let ctx_for_clean = ctx.clone();
|
||
let approval_clean = approval_manager.clone();
|
||
tokio::spawn(async move {
|
||
loop {
|
||
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
|
||
approval_clean.clean_expired().await;
|
||
let now = Instant::now();
|
||
ctx_for_clean
|
||
.approval_ctx
|
||
.lock()
|
||
.await
|
||
.retain(|_, (_, _, _, _, created)| now.duration_since(*created).as_secs() < 300);
|
||
}
|
||
});
|
||
|
||
// 9. 队列（容量 1024）
|
||
let (runner, channels) = QueueRunner::new(1024);
|
||
let enqueue_handle = runner.enqueue_handle();
|
||
let shutdown_handle = runner.shutdown_handle(); // 不再 prefix _
|
||
|
||
// 10-11. 消费者 + 调度器
|
||
let (consumer_handles, sessions) =
|
||
spawn_consumers(channels, ctx.clone(), enqueue_handle.clone());
|
||
|
||
let sched_db = database.pool().clone();
|
||
let sched_enqueue = enqueue_handle.clone();
|
||
tokio::spawn(async move {
|
||
let scheduler = crate::scheduler::Scheduler::new(Some(Arc::new(sched_db)));
|
||
scheduler
|
||
.run(move |to: &str, name: &str, msg: &str| {
|
||
let eq = sched_enqueue.clone();
|
||
let uid = to.to_string();
|
||
let text = msg.to_string();
|
||
let tn = name.to_string();
|
||
tokio::spawn(async move {
|
||
eq.enqueue(PipelineMessage {
|
||
channel: ChannelId::wechat(uid),
|
||
correlation_id: Uuid::new_v4(),
|
||
kind: MessageKind::ScheduledTask {
|
||
text,
|
||
task_name: tn,
|
||
},
|
||
})
|
||
.await;
|
||
});
|
||
Ok(())
|
||
})
|
||
.await;
|
||
});
|
||
info!("定时任务调度器已启动");
|
||
|
||
// 12. QueueRunner
|
||
let runner_handle = tokio::spawn(async move {
|
||
runner.run().await;
|
||
});
|
||
|
||
// 13. Session TTL 清理（5 分钟无活动）
|
||
let sessions_for_clean = sessions.clone();
|
||
tokio::spawn(async move {
|
||
loop {
|
||
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
|
||
let now = Instant::now();
|
||
sessions_for_clean
|
||
.lock()
|
||
.await
|
||
.retain(|_, (_, created)| now.duration_since(*created).as_secs() < 300);
|
||
}
|
||
});
|
||
|
||
info!("开始监听消息...");
|
||
|
||
// 13. 主循环 — 带 Ctrl+C 优雅退出
|
||
let mut running = true;
|
||
while running {
|
||
tokio::select! {
|
||
msg_result = client.receive_messages() => {
|
||
match msg_result {
|
||
Ok(resp) => {
|
||
if !resp.get_updates_buf.is_empty() {
|
||
let _ = crate::db::models::save_runtime(
|
||
ctx.db.pool(), &resp.get_updates_buf,
|
||
).await;
|
||
}
|
||
for msg in &resp.msgs {
|
||
if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) { continue; }
|
||
let from = msg.from_user_id.as_deref().unwrap_or("unknown");
|
||
let text = msg.text_content().unwrap_or("(非文本消息)");
|
||
let ctx_token = msg.context_token.as_deref();
|
||
let msg_id = msg.message_id.map(|id| id.to_string()).unwrap_or_default();
|
||
info!("收到消息 from={}: {}", from, text);
|
||
|
||
let _ = crate::db::models::insert_chat_record(
|
||
ctx.db.pool(), "inbound", from, &ctx.account_id, text, "wechat", ctx_token, &msg_id,
|
||
).await;
|
||
|
||
// 审批回复
|
||
if let Some((skill_name, decision)) = ctx.approval.handle_reply(from, text).await {
|
||
match decision {
|
||
ApprovalDecision::Approved => {
|
||
info!("审批通过: {}", skill_name);
|
||
// 5 元组: (orig_text, context_token, tool_name, tool_args, timestamp)
|
||
if let Some((orig_text, orig_ctx, _tn, _ta, _ts)) = ctx.approval_ctx.lock().await.remove(from) {
|
||
enqueue_handle.enqueue(PipelineMessage {
|
||
channel: ChannelId::wechat(from),
|
||
correlation_id: Uuid::new_v4(),
|
||
kind: MessageKind::UserMessage {
|
||
text: orig_text, message_id: String::new(),
|
||
context_token: orig_ctx, approved_tool: Some(skill_name),
|
||
},
|
||
}).await;
|
||
}
|
||
}
|
||
ApprovalDecision::Rejected | ApprovalDecision::Expired => {
|
||
ctx.approval_ctx.lock().await.remove(from);
|
||
}
|
||
}
|
||
continue;
|
||
}
|
||
|
||
enqueue_handle.enqueue(PipelineMessage {
|
||
channel: ChannelId::wechat(from),
|
||
correlation_id: Uuid::new_v4(),
|
||
kind: MessageKind::UserMessage {
|
||
text: text.to_string(), message_id: msg_id,
|
||
context_token: ctx_token.map(String::from), approved_tool: None,
|
||
},
|
||
}).await;
|
||
}
|
||
}
|
||
Err(e) => {
|
||
if !e.contains("超时") && !e.contains("timeout") {
|
||
error!("接收消息失败: {}", e);
|
||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
_ = tokio::signal::ctrl_c() => {
|
||
info!("收到 Ctrl+C，优雅关闭中...");
|
||
running = false;
|
||
}
|
||
}
|
||
}
|
||
|
||
shutdown_handle.shutdown();
|
||
info!("等待消费者处理剩余消息...");
|
||
tokio::time::sleep(std::time::Duration::from_secs(3)).await;
|
||
for h in consumer_handles {
|
||
h.abort();
|
||
}
|
||
runner_handle.abort();
|
||
info!("Daemon 已停止");
|
||
}
|
||
|
||
// ─── 消费者 ───
|
||
|
||
fn spawn_consumers(
|
||
channels: ConsumerChannels,
|
||
ctx: Arc<DaemonCtx>,
|
||
enqueue: EnqueueHandle,
|
||
) -> (
|
||
Vec<tokio::task::JoinHandle<()>>,
|
||
Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||
) {
|
||
let sessions: Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>> =
|
||
Arc::new(Mutex::new(HashMap::new()));
|
||
|
||
let mut handles = Vec::new();
|
||
// LLM Consumer
|
||
{
|
||
let mut rx = channels.llm_rx;
|
||
let ctx = ctx.clone();
|
||
let eq = enqueue.clone();
|
||
let s = sessions.clone();
|
||
handles.push(tokio::spawn(async move {
|
||
llm_consumer_loop(&mut rx, &ctx, &eq, &s).await;
|
||
}));
|
||
}
|
||
// Tool Consumer
|
||
{
|
||
let mut rx = channels.tool_rx;
|
||
let ctx = ctx.clone();
|
||
let eq = enqueue.clone();
|
||
handles.push(tokio::spawn(async move {
|
||
tool_consumer_loop(&mut rx, &ctx, &eq).await;
|
||
}));
|
||
}
|
||
// Send Consumer
|
||
{
|
||
let mut rx = channels.send_rx;
|
||
let ctx = ctx.clone();
|
||
handles.push(tokio::spawn(async move {
|
||
send_consumer_loop(&mut rx, &ctx).await;
|
||
}));
|
||
}
|
||
(handles, sessions)
|
||
}
|
||
|
||
// ─── LLM Consumer ───
|
||
|
||
async fn llm_consumer_loop(
|
||
rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>,
|
||
ctx: &DaemonCtx,
|
||
enqueue: &EnqueueHandle,
|
||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||
) {
|
||
info!("LLM Consumer 已启动");
|
||
|
||
while let Some(msg) = rx.recv().await {
|
||
let user_id = msg.channel.user_id.clone();
|
||
let correlation_id = msg.correlation_id;
|
||
|
||
match msg.kind {
|
||
MessageKind::UserMessage {
|
||
text,
|
||
message_id: _,
|
||
context_token,
|
||
approved_tool,
|
||
} => {
|
||
// 加载数据
|
||
ctx.memory_store.load_if_needed(&user_id).await;
|
||
let history = load_history(&ctx.db, &user_id).await;
|
||
|
||
let sys_prompt = if ctx.system_prompt.is_empty() {
|
||
DEFAULT_SYSTEM_PROMPT.to_string()
|
||
} else {
|
||
ctx.system_prompt.clone()
|
||
};
|
||
|
||
let cfg = ConversationConfig {
|
||
system_prompt: sys_prompt,
|
||
model: ctx.model.clone(),
|
||
tools: Some(ctx.tools_list.clone()),
|
||
..Default::default()
|
||
};
|
||
|
||
let mut conv = match Conversation::new(cfg) {
|
||
Ok(c) => c,
|
||
Err(e) => {
|
||
error!("LLM 初始化失败: {}", e);
|
||
continue;
|
||
}
|
||
};
|
||
|
||
conv.session().lock().await.db_pool = Some(Arc::new(ctx.db.pool().clone()));
|
||
|
||
{
|
||
let s_arc = conv.session();
|
||
let mut s = s_arc.lock().await;
|
||
s.load_from_history(&history);
|
||
s.current_user_id = user_id.clone();
|
||
}
|
||
|
||
if let Some(tool) = &approved_tool {
|
||
conv.session()
|
||
.lock()
|
||
.await
|
||
.add_system_note(&format!("\n[用户已批准工具: {}]", tool));
|
||
}
|
||
|
||
// 设置 ToolExecutor — 所有工具统一通过消息队列异步执行
|
||
let eq = enqueue.clone();
|
||
let uid = user_id.clone();
|
||
let corr = correlation_id;
|
||
let session = conv.session();
|
||
let approved_tool = approved_tool.clone();
|
||
|
||
let executor: ToolExecutor =
|
||
Arc::new(move |name: &str, args_json: &str, tc_id: &str| {
|
||
let eq = eq.clone();
|
||
let uid = uid.clone();
|
||
let corr = corr;
|
||
let session = session.clone();
|
||
let n = name.to_string();
|
||
let args = args_json.to_string();
|
||
let is_approved = approved_tool.as_ref() == Some(&n);
|
||
let tool_call_id = tc_id.to_string();
|
||
|
||
Box::pin(async move {
|
||
let ch = ChannelId::wechat(&uid);
|
||
let sid = { session.lock().await.session_id };
|
||
if is_approved {
|
||
eq.enqueue(PipelineMessage::approved_tool_call(
|
||
ch,
|
||
corr,
|
||
sid,
|
||
&tool_call_id,
|
||
&n,
|
||
&args,
|
||
None,
|
||
))
|
||
.await;
|
||
} else {
|
||
eq.enqueue(PipelineMessage::tool_call(
|
||
ch,
|
||
corr,
|
||
sid,
|
||
&tool_call_id,
|
||
&n,
|
||
&args,
|
||
None,
|
||
))
|
||
.await;
|
||
}
|
||
Ok(format!("{}{}", ASYNC_MARKER, n))
|
||
})
|
||
});
|
||
conv.set_tool_executor(executor);
|
||
|
||
// 缓存 Conversation（记录创建时间用于 TTL）
|
||
sessions
|
||
.lock()
|
||
.await
|
||
.insert(correlation_id, (conv, Instant::now()));
|
||
|
||
// 执行 LLM 对话
|
||
run_llm_round(
|
||
&user_id,
|
||
correlation_id,
|
||
context_token,
|
||
&text,
|
||
sessions,
|
||
enqueue,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
MessageKind::ToolResult {
|
||
session_id: _,
|
||
tool_call_id,
|
||
tool_name,
|
||
result,
|
||
context_token,
|
||
} => {
|
||
// 从缓存恢复会话，注入工具结果并恢复 LLM 循环
|
||
let should_resume = {
|
||
let mut map = sessions.lock().await;
|
||
if let Some((conv, last_access)) = map.get_mut(&correlation_id) {
|
||
conv.session().lock().await.add_tool_result(
|
||
&tool_call_id,
|
||
&tool_name,
|
||
&result,
|
||
);
|
||
*last_access = Instant::now(); // 更新访问时间，防止 TTL 误清理
|
||
conv.notify_tool_result()
|
||
} else {
|
||
warn!("会话已过期: corr={}", correlation_id);
|
||
false
|
||
}
|
||
};
|
||
|
||
if should_resume {
|
||
resume_llm_round(&user_id, correlation_id, context_token, sessions, enqueue)
|
||
.await;
|
||
}
|
||
}
|
||
|
||
MessageKind::ScheduledTask { text, task_name } => {
|
||
info!("LLM Consumer: 定时任务 user={} task={}", user_id, task_name);
|
||
let prompt = format!("[定时任务: {}]\n{}", task_name, text);
|
||
|
||
// 为保持简洁，定时任务直接作为用户消息处理（使用全新 correlation_id）
|
||
let msg = PipelineMessage {
|
||
channel: ChannelId::wechat(user_id.clone()),
|
||
correlation_id: Uuid::new_v4(),
|
||
kind: MessageKind::UserMessage {
|
||
text: prompt,
|
||
message_id: String::new(),
|
||
context_token: None,
|
||
approved_tool: None,
|
||
},
|
||
};
|
||
// 重新投递给自己 — 简单但可靠
|
||
// 注意: 这里用了一个技巧，重新入队而不是复制代码
|
||
let eq = enqueue.clone();
|
||
eq.enqueue(msg).await;
|
||
}
|
||
|
||
other => warn!(
|
||
"LLM Consumer 收到不期望的消息: {:?}",
|
||
std::mem::discriminant(&other)
|
||
),
|
||
}
|
||
}
|
||
info!("LLM Consumer 已停止");
|
||
}
|
||
|
||
/// 执行一轮 LLM 对话，仅在无异步工具等待时发送回复并清理会话。
|
||
async fn run_llm_round(
|
||
user_id: &str,
|
||
correlation_id: Uuid,
|
||
context_token: Option<String>,
|
||
text: &str,
|
||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||
enqueue: &EnqueueHandle,
|
||
) {
|
||
info!("══════════════════════════════════════════════════");
|
||
info!("🔄 LLM 回合开始 user={} text={}", user_id, text);
|
||
let result = {
|
||
let map = sessions.lock().await;
|
||
let conv = match map.get(&correlation_id) {
|
||
Some((c, _)) => c,
|
||
None => {
|
||
error!("run_llm_round: conv 不在缓存");
|
||
return;
|
||
}
|
||
};
|
||
|
||
match conv.chat_with_tools(text.to_string()).await {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
error!("LLM 对话失败: {}", e);
|
||
ChatResult {
|
||
reply: format!("处理失败: {}", e),
|
||
used_tools: false,
|
||
usage: None,
|
||
has_pending_async: false,
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
handle_chat_result(
|
||
user_id,
|
||
correlation_id,
|
||
context_token,
|
||
result,
|
||
sessions,
|
||
enqueue,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
/// 在异步工具结果返回后恢复 LLM 循环（不追加用户消息）。
|
||
async fn resume_llm_round(
|
||
user_id: &str,
|
||
correlation_id: Uuid,
|
||
context_token: Option<String>,
|
||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||
enqueue: &EnqueueHandle,
|
||
) {
|
||
info!("══════════════════════════════════════════════════");
|
||
info!("🔄 LLM 恢复 user={}", user_id);
|
||
let result = {
|
||
let map = sessions.lock().await;
|
||
let conv = match map.get(&correlation_id) {
|
||
Some((c, _)) => c,
|
||
None => {
|
||
error!("resume_llm_round: conv 不在缓存");
|
||
return;
|
||
}
|
||
};
|
||
|
||
match conv.resume_tool_loop().await {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
error!("LLM 恢复失败: {}", e);
|
||
ChatResult {
|
||
reply: format!("处理失败: {}", e),
|
||
used_tools: false,
|
||
usage: None,
|
||
has_pending_async: false,
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
handle_chat_result(
|
||
user_id,
|
||
correlation_id,
|
||
context_token,
|
||
result,
|
||
sessions,
|
||
enqueue,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
/// 统一处理 ChatResult：有异步等待则保留 session，否则发送回复并清理。
|
||
async fn handle_chat_result(
|
||
user_id: &str,
|
||
correlation_id: Uuid,
|
||
context_token: Option<String>,
|
||
result: ChatResult,
|
||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||
enqueue: &EnqueueHandle,
|
||
) {
|
||
if result.has_pending_async {
|
||
// 有异步工具等待返回 → 保留 session
|
||
if !result.reply.is_empty() {
|
||
let ch = ChannelId::wechat(user_id);
|
||
let msg = PipelineMessage::llm_reply(ch, correlation_id, &result.reply, context_token);
|
||
enqueue.enqueue(msg).await;
|
||
}
|
||
return;
|
||
}
|
||
|
||
// 对话完成：发送回复并清理 session
|
||
if !result.reply.is_empty() {
|
||
let ch = ChannelId::wechat(user_id);
|
||
let msg = PipelineMessage::llm_reply(ch, correlation_id, &result.reply, context_token);
|
||
enqueue.enqueue(msg).await;
|
||
}
|
||
sessions.lock().await.remove(&correlation_id);
|
||
info!("✅ LLM 回合结束 user={}", user_id);
|
||
}
|
||
|
||
// ─── Tool Consumer ───
|
||
|
||
async fn tool_consumer_loop(
|
||
rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>,
|
||
ctx: &DaemonCtx,
|
||
enqueue: &EnqueueHandle,
|
||
) {
|
||
info!("Tool Consumer 已启动");
|
||
|
||
while let Some(msg) = rx.recv().await {
|
||
let user_id = msg.channel.user_id.clone();
|
||
let correlation_id = msg.correlation_id;
|
||
|
||
match msg.kind {
|
||
MessageKind::ToolCall {
|
||
session_id,
|
||
tool_call_id,
|
||
tool_name,
|
||
arguments,
|
||
context_token,
|
||
approved,
|
||
} => {
|
||
info!(
|
||
"Tool Consumer: user={} tool={} args={}",
|
||
user_id, tool_name, arguments
|
||
);
|
||
|
||
// 高风险工具 → 审批（除非已被用户批准）
|
||
if !approved && crate::tools::is_high_risk(&tool_name) {
|
||
info!("高风险工具 {} 需要审批", tool_name);
|
||
ctx.approval_ctx.lock().await.insert(
|
||
user_id.clone(),
|
||
(
|
||
format!("[工具: {}]", tool_name),
|
||
context_token.clone(),
|
||
tool_name.clone(),
|
||
arguments.clone(),
|
||
Instant::now(),
|
||
),
|
||
);
|
||
match ctx.approval.create(&user_id, &tool_name).await {
|
||
Ok((code, _rx)) => {
|
||
let m = format!(
|
||
"⚠️ 操作确认\n\n技能：{}\n\n确认码：{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效)",
|
||
tool_name, code
|
||
);
|
||
// 通过消息队列发送审批提示（不直接调 send_text）
|
||
let ch = ChannelId::wechat(&user_id);
|
||
enqueue
|
||
.enqueue(PipelineMessage::llm_reply(ch, correlation_id, &m, None))
|
||
.await;
|
||
}
|
||
Err(e) => error!("创建审批失败: {}", e),
|
||
}
|
||
continue;
|
||
}
|
||
|
||
// 执行工具 — 所有工具统一在此处理
|
||
let output = execute_tool(&tool_name, &arguments, ctx, &user_id).await;
|
||
|
||
enqueue
|
||
.enqueue(PipelineMessage::tool_result(
|
||
ChannelId::wechat(&user_id),
|
||
correlation_id,
|
||
session_id,
|
||
&tool_call_id,
|
||
&tool_name,
|
||
&output,
|
||
context_token,
|
||
))
|
||
.await;
|
||
}
|
||
MessageKind::ApprovalRequest { .. } => {
|
||
warn!("Tool Consumer 收到 ApprovalRequest");
|
||
}
|
||
other => warn!(
|
||
"Tool Consumer 收到不期望: {:?}",
|
||
std::mem::discriminant(&other)
|
||
),
|
||
}
|
||
}
|
||
info!("Tool Consumer 已停止");
|
||
}
|
||
|
||
// ─── 统一工具执行 ───
|
||
|
||
/// 所有工具的统一执行入口 — 仅此一处
|
||
///
|
||
/// `call_capability` 通过递归派发到此函数，需 boxing 避免无限大小。
|
||
fn execute_tool<'a>(
|
||
name: &'a str,
|
||
arguments: &'a str,
|
||
ctx: &'a DaemonCtx,
|
||
user_id: &'a str,
|
||
) -> std::pin::Pin<std::boxed::Box<dyn std::future::Future<Output = String> + Send + 'a>> {
|
||
Box::pin(async move {
|
||
// 上下文工具
|
||
match name {
|
||
"read_memories" => {
|
||
let result = ctx.memory_store.read_for(user_id, None).await;
|
||
tracing::info!("🔧 execute_tool: read_memories → {}", result);
|
||
return result;
|
||
}
|
||
"write_memory" => {
|
||
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
|
||
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
|
||
if content.is_empty() {
|
||
return "请提供 content 参数".to_string();
|
||
}
|
||
ctx.memory_store.write_for(user_id, content).await;
|
||
let result = format!("已添加长期记忆: {}", content);
|
||
tracing::info!("🔧 execute_tool: write_memory → {}", result);
|
||
return result;
|
||
}
|
||
"delete_memory" => {
|
||
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
|
||
let memory_id = params
|
||
.get("memory_id")
|
||
.and_then(|v| v.as_i64())
|
||
.unwrap_or(-1) as i32;
|
||
if memory_id <= 0 {
|
||
return "请提供有效的 memory_id 参数".to_string();
|
||
}
|
||
let result = ctx.memory_store.delete_for(user_id, memory_id).await;
|
||
tracing::info!("🔧 execute_tool: delete_memory({}) → {}", memory_id, result);
|
||
return result;
|
||
}
|
||
"update_memory" => {
|
||
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
|
||
let memory_id = params
|
||
.get("memory_id")
|
||
.and_then(|v| v.as_i64())
|
||
.unwrap_or(-1) as i32;
|
||
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
|
||
if memory_id <= 0 {
|
||
return "请提供有效的 memory_id 参数".to_string();
|
||
}
|
||
if content.is_empty() {
|
||
return "请提供 content 参数".to_string();
|
||
}
|
||
let result = ctx
|
||
.memory_store
|
||
.update_for(user_id, memory_id, content)
|
||
.await;
|
||
tracing::info!("🔧 execute_tool: update_memory({}) → {}", memory_id, result);
|
||
return result;
|
||
}
|
||
"read_summaries" => {
|
||
let sums = load_summaries(&ctx.db, user_id).await;
|
||
let result = if sums.is_empty() {
|
||
"暂无历史摘要".to_string()
|
||
} else {
|
||
sums.iter()
|
||
.enumerate()
|
||
.map(|(i, s)| format!("{}. {}", i + 1, s))
|
||
.collect::<Vec<_>>()
|
||
.join("\n---\n")
|
||
};
|
||
tracing::info!("🔧 execute_tool: read_summaries → {}条", sums.len());
|
||
return result;
|
||
}
|
||
"query_capabilities" => {
|
||
let specs = crate::tools::specs();
|
||
let result = crate::tools::build_capability_guide(&specs);
|
||
tracing::info!(
|
||
"🔧 execute_tool: query_capabilities → {}个工具",
|
||
specs.len()
|
||
);
|
||
return result;
|
||
}
|
||
"call_capability" => {
|
||
// 解包 {name, prompt} → 提取目标工具名和参数，递归派发
|
||
let cp: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
|
||
let target_name = cp
|
||
.get("name")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("")
|
||
.to_string();
|
||
if target_name.is_empty() {
|
||
return "请提供 name 参数".to_string();
|
||
}
|
||
// 防止嵌套 call_capability
|
||
if target_name == "call_capability" {
|
||
return "不支持嵌套调用 call_capability".to_string();
|
||
}
|
||
let target_args = serde_json::to_string(&crate::tools::unpack_call_params(&cp))
|
||
.unwrap_or_default();
|
||
tracing::info!(
|
||
"🔧 execute_tool: call_capability → {} args={}",
|
||
target_name,
|
||
target_args
|
||
);
|
||
// 递归派发
|
||
return execute_tool(&target_name, &target_args, ctx, user_id).await;
|
||
}
|
||
_ => {}
|
||
}
|
||
|
||
// 内置工具注册表
|
||
match crate::tools::execute(name, arguments).await {
|
||
Some(r) => {
|
||
tracing::info!(
|
||
"🔧 execute_tool: {} → success={} output={}",
|
||
name,
|
||
r.success,
|
||
r.output
|
||
);
|
||
r.output
|
||
}
|
||
None => {
|
||
tracing::warn!("🔧 execute_tool: 未知工具 {}", name);
|
||
format!("未知工具: {}", name)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
|
||
// ─── Send Consumer ───
|
||
|
||
async fn send_consumer_loop(
|
||
rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>,
|
||
ctx: &DaemonCtx,
|
||
) {
|
||
info!("Send Consumer 已启动");
|
||
|
||
while let Some(msg) = rx.recv().await {
|
||
match msg.kind {
|
||
MessageKind::LLMReply {
|
||
text,
|
||
context_token,
|
||
} => {
|
||
let user_id = &msg.channel.user_id;
|
||
info!(
|
||
"Send Consumer: user={} text_len={} text={}",
|
||
user_id,
|
||
text.len(),
|
||
text
|
||
);
|
||
|
||
match ctx
|
||
.client
|
||
.send_text(user_id, &text, context_token.as_deref())
|
||
.await
|
||
{
|
||
Ok(sent_id) => {
|
||
info!(
|
||
"回复成功 user={} msg_id={}",
|
||
user_id.chars().take(5).collect::<String>(),
|
||
sent_id
|
||
);
|
||
let _ = crate::db::models::insert_chat_record(
|
||
ctx.db.pool(),
|
||
"outbound",
|
||
user_id,
|
||
&ctx.account_id,
|
||
&text,
|
||
"llm",
|
||
context_token.as_deref(),
|
||
&sent_id,
|
||
)
|
||
.await;
|
||
}
|
||
Err(e) => error!("发送回复失败: {}", e),
|
||
}
|
||
}
|
||
other => warn!(
|
||
"Send Consumer 收到不期望: {:?}",
|
||
std::mem::discriminant(&other)
|
||
),
|
||
}
|
||
}
|
||
info!("Send Consumer 已停止");
|
||
}
|
||
|
||
// ─── 辅助函数 ───
|
||
|
||
async fn load_auth(database: &Arc<Database>) -> Option<crate::db::models::AuthState> {
|
||
crate::db::models::load_auth(database.pool()).await
|
||
}
|
||
|
||
async fn load_history(db: &Arc<Database>, user_id: &str) -> Vec<HistoryEntry> {
|
||
match crate::db::models::list_recent_chat_records(db.pool(), user_id, 20).await {
|
||
Ok(records) => records
|
||
.into_iter()
|
||
.rev()
|
||
.map(|r| HistoryEntry {
|
||
role: if r.direction == "inbound" {
|
||
"user".into()
|
||
} else {
|
||
"assistant".into()
|
||
},
|
||
content: r.text,
|
||
})
|
||
.collect(),
|
||
Err(e) => {
|
||
warn!("加载历史失败: {}", e);
|
||
vec![]
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn load_summaries(db: &Arc<Database>, user_id: &str) -> Vec<String> {
|
||
match crate::db::models::load_summaries(db.pool(), user_id, 5).await {
|
||
Ok(s) => s,
|
||
Err(e) => {
|
||
warn!("加载摘要失败: {}", e);
|
||
vec![]
|
||
}
|
||
}
|
||
}
|
||
|
||
fn build_tools_list() -> Vec<serde_json::Value> {
|
||
let mut list = Vec::new();
|
||
list.push(serde_json::json!({
|
||
"type": "function", "function": {
|
||
"name": "query_capabilities",
|
||
"description": "列出所有可用工具的详细说明、参数格式和调用示例。",
|
||
"parameters": { "type": "object", "properties": {} }
|
||
}
|
||
}));
|
||
list.push(serde_json::json!({
|
||
"type": "function", "function": {
|
||
"name": "call_capability",
|
||
"description": "调用一个已注册的工具。name 为工具名。",
|
||
"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" } }, "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": "delete_memory",
|
||
"description": "删除指定 id 的长期记忆。id 来自 read_memories 返回的 [id] 前缀。",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"memory_id": { "type": "integer", "description": "要删除的记忆 id" }
|
||
},
|
||
"required": ["memory_id"]
|
||
}
|
||
}
|
||
}));
|
||
list.push(serde_json::json!({
|
||
"type": "function", "function": {
|
||
"name": "update_memory",
|
||
"description": "更新指定 id 的长期记忆内容。id 来自 read_memories 返回的 [id] 前缀。",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"memory_id": { "type": "integer", "description": "要更新的记忆 id" },
|
||
"content": { "type": "string", "description": "新的记忆内容" }
|
||
},
|
||
"required": ["memory_id", "content"]
|
||
}
|
||
}
|
||
}));
|
||
|
||
let specs = crate::tools::specs();
|
||
for spec in &specs {
|
||
list.push(serde_json::json!({
|
||
"type": "function",
|
||
"function": { "name": spec.name, "description": spec.description, "parameters": spec.parameters }
|
||
}));
|
||
}
|
||
list
|
||
}
|