feat: 高德地图工具 + 多轮 bug 修复
新增: - 高德地图 6 个内置工具: amap_poi_search, amap_geocode, amap_reverse_geocode, amap_route_plan, amap_travel_plan, amap_map_link - CLI 子命令: ias tool amap (poi-search/geocode/route-plan/travel-plan/map-link) - 工具优先级指南: build_capability_guide() 含 amap > weather > web_search 优先级 - DEFAULT_SYSTEM_PROMPT 内置工具优先级说明（不再依赖环境变量） Bug 修复: 1. (HIGH) call_capability 参数契约修复 — 新增 unpack_call_params() 统一解包 prompt 字段 2. (MEDIUM) 天气 hours=0 不再调用逐时 API 3. (MEDIUM) 长期记忆 daemon 双写缓存+DB，无 DB 不丢失 4. (HIGH) 审批流程状态语义修复 — handle_reply 返回 ApprovalDecision，拒绝/过期不再当作通过 5. (HIGH) 审批有效期基于 expires_at 时间戳而非 receiver drop 6. (MEDIUM) 摘要保存改为 current_user_id 与读取一致 7. (MEDIUM) worker 审批通过后删除孤立 tool result，仅保留 approved_tool 状态放行 架构: - daemon + worker IPC 架构 (Unix Domain Socket) - build_capability_guide 共享函数消除 main.rs/worker.rs 重复
This commit is contained in:
+927
@@ -0,0 +1,927 @@
|
||||
//! Daemon — 常驻进程
|
||||
//!
|
||||
//! 持有 WeChat 长连接 + 数据库。每收到一条消息 spawn 一个 worker 进程，
|
||||
//! 通过 Unix Domain Socket 通信。Worker 代码更新后编译即可，下一条消息
|
||||
//! 自动使用新版本。
|
||||
|
||||
use crate::context::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 std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
/// 消息队列：每个用户最多一个活跃 worker，等待队列 FIFO
|
||||
struct MessageQueue {
|
||||
/// 当前是否有活跃 worker
|
||||
active: HashMap<String, bool>,
|
||||
/// 每个用户的待处理消息队列
|
||||
pending: HashMap<String, VecDeque<PendingMessage>>,
|
||||
/// 等待 worker 连接的用户顺序
|
||||
waiting: VecDeque<String>,
|
||||
}
|
||||
|
||||
struct PendingMessage {
|
||||
from: String,
|
||||
text: String,
|
||||
account_id: String,
|
||||
context_token: Option<String>,
|
||||
#[allow(dead_code)]
|
||||
message_id: String,
|
||||
approved_tool: Option<String>,
|
||||
}
|
||||
|
||||
impl MessageQueue {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
active: HashMap::new(),
|
||||
pending: HashMap::new(),
|
||||
waiting: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn enqueue(&mut self, user_id: &str, msg: PendingMessage) {
|
||||
self.pending
|
||||
.entry(user_id.to_string())
|
||||
.or_default()
|
||||
.push_back(msg);
|
||||
if !self.active.get(user_id).unwrap_or(&false) && !self.waiting.contains(&user_id.to_string()) {
|
||||
self.waiting.push_back(user_id.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// 取出下一个等待处理的用户（daemon accept worker 后调用）
|
||||
fn pop_waiting(&mut self) -> Option<String> {
|
||||
while let Some(uid) = self.waiting.pop_front() {
|
||||
// 跳过已经活跃的
|
||||
if *self.active.get(&uid).unwrap_or(&false) {
|
||||
continue;
|
||||
}
|
||||
// 检查是否有待处理消息
|
||||
if self.pending.get(&uid).map_or(true, |q| q.is_empty()) {
|
||||
continue;
|
||||
}
|
||||
self.active.insert(uid.clone(), true);
|
||||
return Some(uid);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn pop_message(&mut self, user_id: &str) -> Option<PendingMessage> {
|
||||
self.pending.get_mut(user_id)?.pop_front()
|
||||
}
|
||||
|
||||
fn deactivate(&mut self, user_id: &str) {
|
||||
self.active.insert(user_id.to_string(), false);
|
||||
}
|
||||
|
||||
fn has_pending(&self, user_id: &str) -> bool {
|
||||
self.pending
|
||||
.get(user_id)
|
||||
.map_or(false, |q| !q.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
/// Daemon 入口
|
||||
pub async fn run(
|
||||
sock_path: String,
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &StateManager,
|
||||
memory_store: &Arc<MemoryStore>,
|
||||
) {
|
||||
// 1. 加载认证
|
||||
let auth = match load_auth(database, file_state).await {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
error!("未登录，请先执行 ias login");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 创建 WeChat 客户端
|
||||
let mut client = WeChatClient::new(Some(auth.base_url.clone()));
|
||||
client
|
||||
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
|
||||
.await;
|
||||
|
||||
// 恢复 runtime buf
|
||||
if let Some(r) = file_state.load_runtime() {
|
||||
client.set_updates_buf(&r.get_updates_buf).await;
|
||||
}
|
||||
|
||||
// 3. 注册监听器
|
||||
if let Err(e) = client.notify_start().await {
|
||||
error!("注册监听器失败: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
let account_id = auth.account_id.clone();
|
||||
|
||||
// 4. 绑定 Unix Domain Socket
|
||||
let _ = std::fs::remove_file(&sock_path);
|
||||
let listener = match UnixListener::bind(&sock_path) {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
error!("绑定 Unix Socket 失败 ({}): {}", sock_path, e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
info!("Daemon 已启动: sock={}", sock_path);
|
||||
|
||||
// 5. 审批管理器
|
||||
let approval_manager = Arc::new(ApprovalManager::new(
|
||||
database.as_ref().map(|db| Arc::new(db.pool().clone())),
|
||||
));
|
||||
|
||||
// 6. 消息队列
|
||||
let queue = Arc::new(Mutex::new(MessageQueue::new()));
|
||||
|
||||
// 7. 构建工具列表（传给 worker）
|
||||
let tools_list = build_tools_list();
|
||||
|
||||
// 8. 环境变量映射
|
||||
let env_map = build_env_map();
|
||||
|
||||
// 9. 系统提示词
|
||||
let system_prompt =
|
||||
std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new());
|
||||
let model =
|
||||
std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
|
||||
|
||||
// 10. 共享数据
|
||||
let shared = Arc::new(DaemonShared {
|
||||
db: database.clone(),
|
||||
client: client.clone(),
|
||||
account_id: account_id.clone(),
|
||||
approval: approval_manager.clone(),
|
||||
memory_store: memory_store.clone(),
|
||||
sock_path: sock_path.clone(),
|
||||
tools_list: tools_list.clone(),
|
||||
env_map,
|
||||
system_prompt,
|
||||
model,
|
||||
approval_ctx: Mutex::new(HashMap::new()),
|
||||
});
|
||||
|
||||
// 11. 审批过期清理
|
||||
let approval_clean = approval_manager.clone();
|
||||
let shared_for_clean = shared.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
|
||||
approval_clean.clean_expired().await;
|
||||
// 清理过期的审批上下文（超过 5 分钟）
|
||||
let now = Instant::now();
|
||||
shared_for_clean.approval_ctx.lock().await
|
||||
.retain(|_, (_, _, _, created)| now.duration_since(*created).as_secs() < 300);
|
||||
}
|
||||
});
|
||||
|
||||
// 12. 调度器
|
||||
if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) {
|
||||
let sched_client = client.clone();
|
||||
tokio::spawn(async move {
|
||||
let scheduler = crate::scheduler::Scheduler::new(Some(Arc::new(sched_db)));
|
||||
scheduler
|
||||
.run(move |to: &str, _name: &str, msg: &str| {
|
||||
let c = sched_client.clone();
|
||||
let uid = to.to_string();
|
||||
let text = msg.to_string();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(e) = c.send_text(&uid, &text, None).await {
|
||||
tracing::error!("发送调度结果失败: {}", e);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
.await;
|
||||
});
|
||||
info!("定时任务调度器已启动");
|
||||
}
|
||||
|
||||
info!("开始监听消息...");
|
||||
|
||||
// 13. 主循环
|
||||
loop {
|
||||
tokio::select! {
|
||||
// WeChat 消息
|
||||
msg_result = client.receive_messages() => {
|
||||
match msg_result {
|
||||
Ok(resp) => {
|
||||
// 持久化 runtime buf
|
||||
if !resp.get_updates_buf.is_empty() {
|
||||
file_state.save_runtime(&resp.get_updates_buf);
|
||||
}
|
||||
|
||||
for msg in &resp.msgs {
|
||||
if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) {
|
||||
continue;
|
||||
}
|
||||
let from = msg.from_user_id.as_deref().unwrap_or("unknown");
|
||||
let text = msg.text_content().unwrap_or("(非文本消息)");
|
||||
let ctx_token = msg.context_token.as_deref();
|
||||
let msg_id = msg.message_id.map(|id| id.to_string()).unwrap_or_default();
|
||||
|
||||
info!("收到消息 from={}: {}", from, text);
|
||||
|
||||
// 入库：收到的消息
|
||||
if let Some(db) = &shared.db {
|
||||
if let Err(e) = crate::db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"inbound",
|
||||
from,
|
||||
&shared.account_id,
|
||||
text,
|
||||
"wechat",
|
||||
ctx_token,
|
||||
&msg_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// 审批回复检查
|
||||
if let Some((skill_name, decision)) = shared.approval.handle_reply(from, text).await {
|
||||
match decision {
|
||||
ApprovalDecision::Approved => {
|
||||
info!("审批通过: {}", skill_name);
|
||||
let original = shared.approval_ctx.lock().await.remove(from);
|
||||
if let Some((orig_text, orig_ctx, _tool, _ts)) = original {
|
||||
let mut q = queue.lock().await;
|
||||
q.enqueue(from, PendingMessage {
|
||||
from: from.to_string(),
|
||||
text: orig_text,
|
||||
account_id: shared.account_id.clone(),
|
||||
context_token: orig_ctx,
|
||||
message_id: String::new(),
|
||||
approved_tool: Some(skill_name),
|
||||
});
|
||||
drop(q);
|
||||
spawn_worker_for_user(from, &queue, &shared).await;
|
||||
}
|
||||
}
|
||||
ApprovalDecision::Rejected => {
|
||||
info!("审批被拒绝: {}", skill_name);
|
||||
shared.approval_ctx.lock().await.remove(from);
|
||||
}
|
||||
ApprovalDecision::Expired => {
|
||||
info!("审批已过期: {}", skill_name);
|
||||
shared.approval_ctx.lock().await.remove(from);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// 消息入队
|
||||
{
|
||||
let mut q = queue.lock().await;
|
||||
q.enqueue(from, PendingMessage {
|
||||
from: from.to_string(),
|
||||
text: text.to_string(),
|
||||
account_id: shared.account_id.clone(),
|
||||
context_token: ctx_token.map(String::from),
|
||||
message_id: msg_id,
|
||||
approved_tool: None,
|
||||
});
|
||||
drop(q);
|
||||
}
|
||||
|
||||
// 尝试 spawn worker
|
||||
spawn_worker_for_user(from, &queue, &shared).await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if !e.contains("超时") && !e.contains("timeout") {
|
||||
error!("接收消息失败: {}", e);
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Worker 连接
|
||||
accept_result = listener.accept() => {
|
||||
match accept_result {
|
||||
Ok((stream, _)) => {
|
||||
let q = queue.clone();
|
||||
let shr = shared.clone();
|
||||
tokio::spawn(async move {
|
||||
handle_worker(stream, q, shr).await;
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
error!("接受 Worker 连接失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 共享数据
|
||||
struct DaemonShared {
|
||||
db: Option<Arc<Database>>,
|
||||
client: WeChatClient,
|
||||
account_id: String,
|
||||
approval: Arc<ApprovalManager>,
|
||||
memory_store: Arc<MemoryStore>,
|
||||
sock_path: String,
|
||||
tools_list: Vec<serde_json::Value>,
|
||||
env_map: HashMap<String, String>,
|
||||
system_prompt: String,
|
||||
model: String,
|
||||
/// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 创建时间)
|
||||
approval_ctx: Mutex<HashMap<String, (String, Option<String>, String, Instant)>>,
|
||||
}
|
||||
|
||||
/// 处理一个 Worker 连接
|
||||
async fn handle_worker(
|
||||
mut stream: UnixStream,
|
||||
queue: Arc<Mutex<MessageQueue>>,
|
||||
shared: Arc<DaemonShared>,
|
||||
) {
|
||||
// 1. 取出等待的用户
|
||||
let user_id = {
|
||||
let mut q = queue.lock().await;
|
||||
q.pop_waiting()
|
||||
};
|
||||
|
||||
let user_id = match user_id {
|
||||
Some(uid) => uid,
|
||||
None => {
|
||||
warn!("Worker 连接但没有等待的用户");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 2. 取出该用户的一条消息
|
||||
let pending_msg = {
|
||||
let mut q = queue.lock().await;
|
||||
q.pop_message(&user_id)
|
||||
};
|
||||
|
||||
let pending_msg = match pending_msg {
|
||||
Some(m) => m,
|
||||
None => {
|
||||
warn!("Worker 连接但用户 {} 无待处理消息", user_id);
|
||||
queue.lock().await.deactivate(&user_id);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
"Worker 处理: user={} text={:.60}",
|
||||
user_id, pending_msg.text
|
||||
);
|
||||
|
||||
// 3. 加载用户记忆
|
||||
shared.memory_store.load(&user_id).await;
|
||||
|
||||
// 4. 预载历史消息
|
||||
let history = load_history(&shared.db, &user_id).await;
|
||||
let memories = load_memories(&shared.memory_store, &user_id).await;
|
||||
let summaries = load_summaries(&shared.db, &user_id).await;
|
||||
|
||||
// 5. 构建 TaskFrame
|
||||
let approved_tool_info = pending_msg.approved_tool.clone();
|
||||
let task = TaskFrame {
|
||||
user_id: user_id.clone(),
|
||||
msg: TaskMessage {
|
||||
from: pending_msg.from,
|
||||
text: pending_msg.text.clone(),
|
||||
account_id: pending_msg.account_id,
|
||||
context_token: pending_msg.context_token.clone(),
|
||||
},
|
||||
history,
|
||||
memories,
|
||||
summaries,
|
||||
approved_tool: approved_tool_info.clone(),
|
||||
env: shared.env_map.clone(),
|
||||
tools: shared.tools_list.clone(),
|
||||
system_prompt: shared.system_prompt.clone(),
|
||||
model: shared.model.clone(),
|
||||
};
|
||||
|
||||
// 6. 发送 task 帧
|
||||
if let Err(e) = send_frame(&mut stream, &task).await {
|
||||
error!("发送 task 帧失败: {}", e);
|
||||
queue.lock().await.deactivate(&user_id);
|
||||
return;
|
||||
}
|
||||
|
||||
// 7. 逐帧处理 worker 输出
|
||||
let mut reply_text = String::new();
|
||||
let mut chunk_buffer = String::new();
|
||||
let ctx_token = task.msg.context_token.clone();
|
||||
|
||||
loop {
|
||||
match recv_output(&mut stream).await {
|
||||
Ok(frame) => match frame {
|
||||
OutputFrame::Chunk { text } => {
|
||||
chunk_buffer.push_str(&text);
|
||||
}
|
||||
OutputFrame::Reply { text } => {
|
||||
reply_text = text;
|
||||
}
|
||||
OutputFrame::RequestApproval { tool, reason } => {
|
||||
info!("Worker 请求审批: tool={} reason={}", tool, reason);
|
||||
// 创建审批记录
|
||||
match shared.approval.create(&user_id, &tool).await {
|
||||
Ok((code, _rx)) => {
|
||||
let msg = format!(
|
||||
"⚠️ 操作确认\n\n技能：{}\n\n确认码：{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效)",
|
||||
tool, code
|
||||
);
|
||||
// 保存原始消息上下文以便审批后重新处理
|
||||
shared.approval_ctx.lock().await.insert(
|
||||
user_id.clone(),
|
||||
(task.msg.text.clone(), ctx_token.clone(), tool.clone(), Instant::now()),
|
||||
);
|
||||
// 发送确认消息到微信
|
||||
if let Err(e) = shared.client.send_text(&user_id, &msg, ctx_token.as_deref()).await {
|
||||
error!("发送审批消息失败: {}", e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("创建审批失败: {}", e);
|
||||
let _ = shared.client.send_text(&user_id, &format!("审批创建失败: {}", e), ctx_token.as_deref()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
OutputFrame::NeedApproval { tool, .. } => {
|
||||
warn!("收到已废弃的 NeedApproval 帧 (tool={})", tool);
|
||||
}
|
||||
OutputFrame::DbWrite { table, row } => {
|
||||
handle_db_write(&shared.db, &shared.memory_store, &table, &row, &user_id, &shared.account_id).await;
|
||||
}
|
||||
OutputFrame::UsageReport {
|
||||
model,
|
||||
provider,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens: _,
|
||||
cache_hit_tokens,
|
||||
cache_miss_tokens,
|
||||
user_id: uid,
|
||||
} => {
|
||||
if let Some(db) = &shared.db {
|
||||
if let Err(e) = crate::db::models::insert_llm_usage(
|
||||
db.pool(),
|
||||
&uid,
|
||||
&model,
|
||||
&provider,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cache_hit_tokens,
|
||||
cache_miss_tokens,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("存储 LLM 用量失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
OutputFrame::Bye => break,
|
||||
},
|
||||
Err(e) => {
|
||||
error!("读取 Worker 输出帧失败: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 8. 发送最终回复
|
||||
if !reply_text.is_empty() {
|
||||
match shared
|
||||
.client
|
||||
.send_text(&user_id, &reply_text, ctx_token.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(sent_id) => {
|
||||
info!("回复成功 user={} msg_id={}", user_id, sent_id);
|
||||
// 入库：发送的回复
|
||||
if let Some(db) = &shared.db {
|
||||
if let Err(e) = crate::db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
&user_id,
|
||||
&shared.account_id,
|
||||
&reply_text,
|
||||
"llm",
|
||||
ctx_token.as_deref(),
|
||||
&sent_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送回复失败: {}", e),
|
||||
}
|
||||
} else if !chunk_buffer.is_empty() {
|
||||
// 有 chunk 但没有 reply — 用 chunk 作为回复
|
||||
match shared
|
||||
.client
|
||||
.send_text(&user_id, &chunk_buffer, ctx_token.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(sent_id) => {
|
||||
info!("回复成功 (chunk) user={} msg_id={}", user_id, sent_id);
|
||||
if let Some(db) = &shared.db {
|
||||
if let Err(e) = crate::db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
&user_id,
|
||||
&shared.account_id,
|
||||
&chunk_buffer,
|
||||
"llm",
|
||||
ctx_token.as_deref(),
|
||||
&sent_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送回复失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// 9. 清理并继续处理该用户的下一消息
|
||||
{
|
||||
let mut q = queue.lock().await;
|
||||
q.deactivate(&user_id);
|
||||
let has_more = q.has_pending(&user_id);
|
||||
drop(q);
|
||||
if has_more {
|
||||
spawn_worker_for_user(&user_id, &queue, &shared).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 为用户 spawn worker 进程
|
||||
async fn spawn_worker_for_user(
|
||||
user_id: &str,
|
||||
queue: &Arc<Mutex<MessageQueue>>,
|
||||
shared: &Arc<DaemonShared>,
|
||||
) {
|
||||
// 检查是否已有活跃 worker
|
||||
{
|
||||
let q = queue.lock().await;
|
||||
if *q.active.get(user_id).unwrap_or(&false) {
|
||||
return; // 已有 worker 在处理
|
||||
}
|
||||
}
|
||||
|
||||
let sock = shared.sock_path.clone();
|
||||
let exe = match std::env::current_exe() {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
error!("获取当前可执行文件路径失败: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
info!("Spawning worker for user={}", user_id);
|
||||
|
||||
match Command::new(&exe)
|
||||
.arg("worker")
|
||||
.arg("--sock")
|
||||
.arg(&sock)
|
||||
.spawn()
|
||||
{
|
||||
Ok(mut child) => {
|
||||
// 后台等待 worker 退出
|
||||
tokio::spawn(async move {
|
||||
match child.wait().await {
|
||||
Ok(status) => {
|
||||
if !status.success() {
|
||||
warn!("Worker 退出码: {:?}", status.code());
|
||||
}
|
||||
}
|
||||
Err(e) => error!("等待 Worker 退出失败: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Spawn worker 失败: {}", e);
|
||||
queue.lock().await.deactivate(user_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 辅助函数 ───
|
||||
|
||||
async fn load_auth(
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &StateManager,
|
||||
) -> Option<crate::state::AuthState> {
|
||||
if let Some(db) = database {
|
||||
if let Some(a) = crate::db::models::load_auth(db.pool()).await {
|
||||
return Some(a);
|
||||
}
|
||||
// 尝试从文件迁移
|
||||
if let Some(fa) = file_state.load_auth() {
|
||||
let a = crate::state::AuthState {
|
||||
token: fa.token,
|
||||
account_id: fa.account_id,
|
||||
base_url: fa.base_url,
|
||||
};
|
||||
if let Err(e) = crate::db::models::save_auth(db.pool(), &a).await {
|
||||
error!("保存认证信息失败: {}", e);
|
||||
}
|
||||
return Some(a);
|
||||
}
|
||||
} else if let Some(a) = file_state.load_auth() {
|
||||
return Some(a);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn load_history(db: &Option<Arc<Database>>, user_id: &str) -> Vec<HistoryEntry> {
|
||||
let Some(db) = db else {
|
||||
return vec![];
|
||||
};
|
||||
match crate::db::models::list_recent_chat_records(db.pool(), user_id, 20).await {
|
||||
Ok(records) => records
|
||||
.into_iter()
|
||||
.rev() // 时间顺序
|
||||
.map(|r| HistoryEntry {
|
||||
role: if r.direction == "inbound" {
|
||||
"user".into()
|
||||
} else {
|
||||
"assistant".into()
|
||||
},
|
||||
content: r.text,
|
||||
})
|
||||
.collect(),
|
||||
Err(e) => {
|
||||
warn!("加载历史消息失败: {}", e);
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_memories(memory_store: &MemoryStore, user_id: &str) -> Vec<String> {
|
||||
let text = memory_store.read_for(user_id).await;
|
||||
if text == "暂无长期记忆" {
|
||||
vec![]
|
||||
} else {
|
||||
text.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_summaries(db: &Option<Arc<Database>>, user_id: &str) -> Vec<String> {
|
||||
let Some(db) = db else {
|
||||
return vec![];
|
||||
};
|
||||
match crate::db::models::load_summaries(db.pool(), user_id, 5).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
warn!("加载摘要失败: {}", e);
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_db_write(
|
||||
db: &Option<Arc<Database>>,
|
||||
memory_store: &MemoryStore,
|
||||
table: &str,
|
||||
row: &serde_json::Value,
|
||||
user_id: &str,
|
||||
account_id: &str,
|
||||
) {
|
||||
// user_memories 不依赖 DB（memory_store.write_for 自己处理持久化）
|
||||
if table == "user_memories" {
|
||||
let content = row.get("content").and_then(|v| v.as_str()).unwrap_or("");
|
||||
if !content.is_empty() {
|
||||
memory_store.write_for(user_id, content).await;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(db) = db else {
|
||||
return;
|
||||
};
|
||||
match table {
|
||||
"chat_records" => {
|
||||
let direction = row.get("direction").and_then(|v| v.as_str()).unwrap_or("outbound");
|
||||
let text = row.get("text").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let source = row.get("source").and_then(|v| v.as_str()).unwrap_or("llm");
|
||||
let ctx_token = row.get("context_token").and_then(|v| v.as_str());
|
||||
let msg_id = row.get("message_id").and_then(|v| v.as_str()).unwrap_or("");
|
||||
if let Err(e) = crate::db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
direction,
|
||||
user_id,
|
||||
account_id,
|
||||
text,
|
||||
source,
|
||||
ctx_token,
|
||||
msg_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("DbWrite chat_records 失败: {}", e);
|
||||
}
|
||||
}
|
||||
"llm_usage" => {
|
||||
let model = row.get("model").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let provider = row.get("provider").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let prompt_tokens = row.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
||||
let completion_tokens = row.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
||||
let cache_hit = row.get("cache_hit_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
||||
let cache_miss = row.get("cache_miss_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
||||
if let Err(e) = crate::db::models::insert_llm_usage(
|
||||
db.pool(),
|
||||
user_id,
|
||||
model,
|
||||
provider,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
cache_hit,
|
||||
cache_miss,
|
||||
)
|
||||
.await
|
||||
{
|
||||
error!("DbWrite llm_usage 失败: {}", e);
|
||||
}
|
||||
}
|
||||
_ => warn!("未知 DbWrite table: {}", table),
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user