refactor: 移除文件状态并修复工具上下文关联

This commit is contained in:
2026-06-16 20:18:34 +08:00
parent 1d78567102
commit f608c27f8a
9 changed files with 552 additions and 402 deletions
+4
View File
@@ -0,0 +1,4 @@
-- 空迁移:用于对齐 _sqlx_migrations 表中已记录但文件丢失的迁移
-- 该迁移已被之前的某次运行应用过,但对应文件被删除/重命名
-- 此文件仅用于让 sqlx 迁移校验通过,不执行任何 DDL
SELECT 1;
-5
View File
@@ -54,13 +54,8 @@ pub enum Commands {
///
/// 示例:
/// ias daemon
/// ias daemon --sock /tmp/ias.sock
/// ias daemon --log-file # 启用文件日志(配合 systemd)
Daemon {
/// Unix Domain Socket 路径
#[arg(long, default_value = "/tmp/ias_daemon.sock")]
sock: String,
/// 启用文件日志(默认输出到 stderr)
#[arg(long)]
log_file: bool,
+95 -15
View File
@@ -13,7 +13,8 @@
use crate::context::types::ChatSession;
use crate::llm::conversation::Summarizer;
use crate::llm::types::Message;
use crate::llm::types::{Message, Role};
use std::collections::BTreeSet;
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -51,6 +52,49 @@ fn is_cjk(c: char) -> bool {
|| ('\u{f900}'..='\u{faff}').contains(&c)
}
fn matching_tool_call_anchor(recent: &[Message], tool_call_id: &str) -> Option<usize> {
recent.iter().enumerate().rev().find_map(|(idx, msg)| {
if msg.role != Role::Assistant {
return None;
}
let has_match = msg.tool_calls.as_ref().is_some_and(|calls| {
calls.iter().any(|call| call.id == tool_call_id)
});
has_match.then_some(idx)
})
}
fn collect_context_block(recent: &[Message], start_idx: usize) -> (Vec<usize>, u32) {
let mut indexes = BTreeSet::from([start_idx]);
let msg = &recent[start_idx];
if msg.role == Role::Tool
&& let Some(tool_call_id) = msg.tool_call_id.as_deref()
&& let Some(anchor_idx) = matching_tool_call_anchor(recent, tool_call_id)
{
indexes.insert(anchor_idx);
}
if msg.role == Role::Assistant
&& let Some(tool_calls) = &msg.tool_calls
{
for (idx, next_msg) in recent.iter().enumerate().skip(start_idx + 1) {
if next_msg.role != Role::Tool {
continue;
}
if let Some(tool_call_id) = next_msg.tool_call_id.as_deref()
&& tool_calls.iter().any(|call| call.id == tool_call_id)
{
indexes.insert(idx);
}
}
}
let ordered: Vec<usize> = indexes.into_iter().collect();
let tokens = ordered.iter().map(|&idx| estimate_tokens(&recent[idx])).sum();
(ordered, tokens)
}
/// ## 构建给 LLM 的消息数组(带 token budget 管理)
///
/// 这是 LLM 调用前的关键步骤 —— 从 `ChatSession` 中提取合适的消息,
@@ -89,36 +133,41 @@ pub async fn build_context(session: &Arc<Mutex<ChatSession>>, system_prompt: &st
// 近期消息(从 checkpoint 开始,倒序添加直到接近 budget)
let recent = s.recent_messages().to_vec();
let mut included = Vec::new();
let mut included_indexes = BTreeSet::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 {
for start_idx in (0..recent.len()).rev() {
if included_indexes.contains(&start_idx) {
continue;
}
let (block_indexes, block_tokens) = collect_context_block(&recent, start_idx);
if used + tail_used + block_tokens > s.token_budget - 500 {
break;
}
tail_used += t;
included.push(msg.clone());
tail_used += block_tokens;
included_indexes.extend(block_indexes);
}
included.reverse();
// 保护:确保最近一条用户消息始终在上下文中
if let Some(last_user) = recent
if let Some((last_user_idx, last_user)) = recent
.iter()
.enumerate()
.rev()
.find(|m| m.role == crate::llm::types::Role::User)
.find(|(_, m)| m.role == Role::User)
{
let already_included = included
.iter()
.any(|m| m.role == last_user.role && m.content == last_user.content);
if !already_included {
if !included_indexes.contains(&last_user_idx) {
let t = estimate_tokens(last_user);
if used + tail_used + t <= s.token_budget + 2000 {
included.push(last_user.clone());
included_indexes.insert(last_user_idx);
}
}
}
let included: Vec<Message> = included_indexes
.into_iter()
.map(|idx| recent[idx].clone())
.collect();
messages.extend(included);
messages
@@ -323,6 +372,8 @@ fn summarize_messages(messages: &[Message]) -> String {
#[cfg(test)]
mod tests {
use super::*;
use crate::context::types::ChatSession;
use crate::llm::types::{ToolCall, ToolFunctionCall};
#[test]
fn test_estimate_cjk() {
@@ -335,4 +386,33 @@ mod tests {
let t = estimate_text_tokens("hello world");
assert!(t >= 1 && t <= 6, "ASCII estimate: {}", t);
}
#[tokio::test]
async fn keeps_tool_result_with_matching_assistant_tool_call() {
let mut session = ChatSession::new();
session.token_budget = 80;
session.add_user("请重新查看工具列表".to_string());
session.messages.push(Message::assistant_with_tool_calls(vec![ToolCall {
id: "tc_1".to_string(),
call_type: "function".to_string(),
function: ToolFunctionCall {
name: "query_capabilities".to_string(),
arguments: "{}".to_string(),
},
}]));
session.add_tool_result("tc_1", "query_capabilities", "工具列表如下");
let session = Arc::new(Mutex::new(session));
let messages = build_context(&session, "sys").await;
let has_tool = messages.iter().any(|m| m.role == Role::Tool);
let has_anchor = messages.iter().any(|m| {
m.tool_calls
.as_ref()
.is_some_and(|calls| calls.iter().any(|c| c.id == "tc_1"))
});
assert!(has_tool);
assert!(has_anchor);
}
}
+287 -117
View File
@@ -39,12 +39,14 @@
//! 采用统一的单进程 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::context::HistoryEntry;
use crate::llm::{ChatResult, Conversation, ConversationConfig, ToolExecutor, DEFAULT_SYSTEM_PROMPT, ASYNC_MARKER};
use crate::llm::{
ASYNC_MARKER, ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor,
};
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
use crate::state::StateManager;
use crate::tools::approval::{ApprovalDecision, ApprovalManager};
use crate::wechat::client::WeChatClient;
use std::collections::HashMap;
@@ -59,8 +61,8 @@ use uuid::Uuid;
/// 这个结构体被所有三个消费者(LLM / Tool / Send)通过 `Arc` 共享访问。
/// 每个字段都只读(或内部有锁),所以不需要额外的 `Mutex` 包装。
struct DaemonCtx {
/// 数据库句柄(可选,没有时走文件存储)
db: Option<Arc<Database>>,
/// 数据库句柄
db: Arc<Database>,
/// WeChat API 客户端(内部用 Arc<Mutex> 保护 token/updates_buf
client: WeChatClient,
/// 本机器人的微信账号 ID
@@ -83,33 +85,36 @@ struct DaemonCtx {
/// Daemon 入口
///
/// `sock_path` 保留用于兼容旧版 daemon 命令的 CLI 接口,实际已不再使用。
pub async fn run(
_sock_path: String,
database: &Option<Arc<Database>>,
file_state: &StateManager,
memory_store: &Arc<MemoryStore>,
) {
pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
// 1-3. 认证 & 客户端
let auth = match load_auth(database, file_state).await {
let auth = match load_auth(database).await {
Some(a) => a,
None => { error!("未登录"); return; }
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(r) = file_state.load_runtime() { client.set_updates_buf(&r.get_updates_buf).await; }
if let Err(e) = client.notify_start().await { error!("注册监听器失败: {}", e); return; }
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(
database.as_ref().map(|db| Arc::new(db.pool().clone())),
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());
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 {
@@ -132,7 +137,10 @@ pub async fn run(
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
ctx_for_clean
.approval_ctx
.lock()
.await
.retain(|_, (_, _, _, _, created)| now.duration_since(*created).as_secs() < 300);
}
});
@@ -143,29 +151,40 @@ pub async fn run(
let shutdown_handle = runner.shutdown_handle(); // 不再 prefix _
// 10-11. 消费者 + 调度器
let (consumer_handles, sessions) = spawn_consumers(channels, ctx.clone(), enqueue_handle.clone());
let (consumer_handles, sessions) =
spawn_consumers(channels, ctx.clone(), enqueue_handle.clone());
if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) {
let sched_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| {
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();
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;
channel: ChannelId::wechat(uid),
correlation_id: Uuid::new_v4(),
kind: MessageKind::ScheduledTask {
text,
task_name: tn,
},
})
.await;
});
Ok(())
}).await;
})
.await;
});
info!("定时任务调度器已启动");
}
// 12. QueueRunner
let runner_handle = tokio::spawn(async move { runner.run().await; });
let runner_handle = tokio::spawn(async move {
runner.run().await;
});
// 13. Session TTL 清理(5 分钟无活动)
let sessions_for_clean = sessions.clone();
@@ -173,7 +192,9 @@ pub async fn run(
loop {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
let now = Instant::now();
sessions_for_clean.lock().await
sessions_for_clean
.lock()
.await
.retain(|_, (_, created)| now.duration_since(*created).as_secs() < 300);
}
});
@@ -188,7 +209,9 @@ pub async fn run(
match msg_result {
Ok(resp) => {
if !resp.get_updates_buf.is_empty() {
file_state.save_runtime(&resp.get_updates_buf);
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; }
@@ -198,11 +221,9 @@ pub async fn run(
let msg_id = msg.message_id.map(|id| id.to_string()).unwrap_or_default();
info!("收到消息 from={}: {}", from, text);
if let Some(db) = &ctx.db {
let _ = crate::db::models::insert_chat_record(
db.pool(), "inbound", from, &ctx.account_id, text, "wechat", ctx_token, &msg_id,
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 {
@@ -256,7 +277,9 @@ pub async fn run(
shutdown_handle.shutdown();
info!("等待消费者处理剩余消息...");
tokio::time::sleep(std::time::Duration::from_secs(3)).await;
for h in consumer_handles { h.abort(); }
for h in consumer_handles {
h.abort();
}
runner_handle.abort();
info!("Daemon 已停止");
}
@@ -267,8 +290,12 @@ 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()));
) -> (
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
@@ -277,20 +304,26 @@ fn spawn_consumers(
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; }));
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; }));
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.push(tokio::spawn(async move {
send_consumer_loop(&mut rx, &ctx).await;
}));
}
(handles, sessions)
}
@@ -310,7 +343,12 @@ async fn llm_consumer_loop(
let correlation_id = msg.correlation_id;
match msg.kind {
MessageKind::UserMessage { text, message_id: _, context_token, approved_tool } => {
MessageKind::UserMessage {
text,
message_id: _,
context_token,
approved_tool,
} => {
info!("LLM Consumer: user={} text={:.60}", user_id, text);
// 加载数据
@@ -319,7 +357,9 @@ async fn llm_consumer_loop(
let sys_prompt = if ctx.system_prompt.is_empty() {
DEFAULT_SYSTEM_PROMPT.to_string()
} else { ctx.system_prompt.clone() };
} else {
ctx.system_prompt.clone()
};
let cfg = ConversationConfig {
system_prompt: sys_prompt,
@@ -330,12 +370,13 @@ async fn llm_consumer_loop(
let mut conv = match Conversation::new(cfg) {
Ok(c) => c,
Err(e) => { error!("LLM 初始化失败: {}", e); continue; }
Err(e) => {
error!("LLM 初始化失败: {}", e);
continue;
}
};
if let Some(db) = &ctx.db {
conv.session().lock().await.db_pool = Some(Arc::new(db.pool().clone()));
}
conv.session().lock().await.db_pool = Some(Arc::new(ctx.db.pool().clone()));
{
let s_arc = conv.session();
@@ -345,7 +386,10 @@ async fn llm_consumer_loop(
}
if let Some(tool) = &approved_tool {
conv.session().lock().await.add_system_note(&format!("\n[用户已批准工具: {}]", tool));
conv.session()
.lock()
.await
.add_system_note(&format!("\n[用户已批准工具: {}]", tool));
}
// 设置 ToolExecutor — 所有工具统一通过消息队列异步执行
@@ -355,7 +399,7 @@ async fn llm_consumer_loop(
let session = conv.session();
let approved_tool = approved_tool.clone();
let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| {
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;
@@ -363,19 +407,33 @@ async fn llm_consumer_loop(
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 };
let tool_call_id = Uuid::new_v4().to_string();
if is_approved {
eq.enqueue(PipelineMessage::approved_tool_call(
ch, corr, sid, &tool_call_id, &n, &args, None,
)).await;
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;
ch,
corr,
sid,
&tool_call_id,
&n,
&args,
None,
))
.await;
}
Ok(format!("{}{}", ASYNC_MARKER, n))
})
@@ -383,30 +441,55 @@ async fn llm_consumer_loop(
conv.set_tool_executor(executor);
// 缓存 Conversation(记录创建时间用于 TTL
sessions.lock().await.insert(correlation_id, (conv, Instant::now()));
sessions
.lock()
.await
.insert(correlation_id, (conv, Instant::now()));
// 执行 LLM 对话
run_llm_round(&user_id, correlation_id, context_token, &text, sessions, enqueue).await;
run_llm_round(
&user_id,
correlation_id,
context_token,
&text,
sessions,
enqueue,
)
.await;
}
MessageKind::ToolResult { session_id: _, tool_call_id, tool_name, result, context_token } => {
info!("LLM Consumer: 工具结果回喂 user={} tool={}", user_id, tool_name);
MessageKind::ToolResult {
session_id: _,
tool_call_id,
tool_name,
result,
context_token,
} => {
info!(
"LLM Consumer: 工具结果回喂 user={} tool={}",
user_id, tool_name
);
// 从缓存恢复会话,注入工具结果并恢复 LLM 循环
let should_continue = {
let 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);
conv.session().lock().await.add_tool_result(
&tool_call_id,
&tool_name,
&result,
);
*last_access = Instant::now(); // 更新访问时间,防止 TTL 误清理
true
conv.notify_tool_result()
} else {
warn!("会话已过期: corr={}", correlation_id);
false
}
};
if should_continue {
resume_llm_round(&user_id, correlation_id, context_token, sessions, enqueue).await;
if should_resume {
resume_llm_round(&user_id, correlation_id, context_token, sessions, enqueue)
.await;
}
}
@@ -431,7 +514,10 @@ async fn llm_consumer_loop(
eq.enqueue(msg).await;
}
other => warn!("LLM Consumer 收到不期望的消息: {:?}", std::mem::discriminant(&other)),
other => warn!(
"LLM Consumer 收到不期望的消息: {:?}",
std::mem::discriminant(&other)
),
}
}
info!("LLM Consumer 已停止");
@@ -450,7 +536,10 @@ async fn run_llm_round(
let map = sessions.lock().await;
let conv = match map.get(&correlation_id) {
Some((c, _)) => c,
None => { error!("run_llm_round: conv 不在缓存"); return; }
None => {
error!("run_llm_round: conv 不在缓存");
return;
}
};
match conv.chat_with_tools(text.to_string()).await {
@@ -467,7 +556,15 @@ async fn run_llm_round(
}
};
handle_chat_result(user_id, correlation_id, context_token, result, sessions, enqueue).await;
handle_chat_result(
user_id,
correlation_id,
context_token,
result,
sessions,
enqueue,
)
.await;
}
/// 在异步工具结果返回后恢复 LLM 循环(不追加用户消息)。
@@ -482,7 +579,10 @@ async fn resume_llm_round(
let map = sessions.lock().await;
let conv = match map.get(&correlation_id) {
Some((c, _)) => c,
None => { error!("resume_llm_round: conv 不在缓存"); return; }
None => {
error!("resume_llm_round: conv 不在缓存");
return;
}
};
match conv.resume_tool_loop().await {
@@ -499,7 +599,15 @@ async fn resume_llm_round(
}
};
handle_chat_result(user_id, correlation_id, context_token, result, sessions, enqueue).await;
handle_chat_result(
user_id,
correlation_id,
context_token,
result,
sessions,
enqueue,
)
.await;
}
/// 统一处理 ChatResult:有异步等待则保留 session,否则发送回复并清理。
@@ -544,7 +652,14 @@ async fn tool_consumer_loop(
let correlation_id = msg.correlation_id;
match msg.kind {
MessageKind::ToolCall { session_id, tool_call_id, tool_name, arguments, context_token, approved } => {
MessageKind::ToolCall {
session_id,
tool_call_id,
tool_name,
arguments,
context_token,
approved,
} => {
info!("Tool Consumer: user={} tool={}", user_id, tool_name);
// 高风险工具 → 审批(除非已被用户批准)
@@ -568,7 +683,9 @@ async fn tool_consumer_loop(
);
// 通过消息队列发送审批提示(不直接调 send_text)
let ch = ChannelId::wechat(&user_id);
enqueue.enqueue(PipelineMessage::llm_reply(ch, correlation_id, &m, None)).await;
enqueue
.enqueue(PipelineMessage::llm_reply(ch, correlation_id, &m, None))
.await;
}
Err(e) => error!("创建审批失败: {}", e),
}
@@ -578,14 +695,25 @@ async fn tool_consumer_loop(
// 执行工具 — 所有工具统一在此处理
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;
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)),
other => warn!(
"Tool Consumer 收到不期望: {:?}",
std::mem::discriminant(&other)
),
}
}
info!("Tool Consumer 已停止");
@@ -611,30 +739,52 @@ fn execute_tool<'a>(
"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(); }
if content.is_empty() {
return "请提供 content 参数".to_string();
}
ctx.memory_store.write_for(user_id, content).await;
return format!("已添加长期记忆: {}", content);
}
"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 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();
}
return ctx.memory_store.delete_for(user_id, memory_id).await;
}
"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 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(); }
return ctx.memory_store.update_for(user_id, memory_id, content).await;
if memory_id <= 0 {
return "请提供有效的 memory_id 参数".to_string();
}
if content.is_empty() {
return "请提供 content 参数".to_string();
}
return ctx
.memory_store
.update_for(user_id, memory_id, content)
.await;
}
"read_summaries" => {
let sums = load_summaries(&ctx.db, user_id).await;
if sums.is_empty() { return "暂无历史摘要".to_string(); }
return sums.iter().enumerate()
if sums.is_empty() {
return "暂无历史摘要".to_string();
}
return sums
.iter()
.enumerate()
.map(|(i, s)| format!("{}. {}", i + 1, s))
.collect::<Vec<_>>().join("\n---\n");
.collect::<Vec<_>>()
.join("\n---\n");
}
"query_capabilities" => {
let specs = crate::tools::specs();
@@ -643,7 +793,11 @@ fn execute_tool<'a>(
"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();
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();
}
@@ -677,24 +831,39 @@ async fn send_consumer_loop(
while let Some(msg) = rx.recv().await {
match msg.kind {
MessageKind::LLMReply { text, context_token } => {
MessageKind::LLMReply {
text,
context_token,
} => {
let user_id = &msg.channel.user_id;
info!("Send Consumer: user={} text_len={}", user_id, text.len());
match ctx.client.send_text(user_id, &text, context_token.as_deref()).await {
match ctx
.client
.send_text(user_id, &text, context_token.as_deref())
.await
{
Ok(sent_id) => {
info!("回复成功 user={} msg_id={}", user_id, sent_id);
if let Some(db) = &ctx.db {
let _ = crate::db::models::insert_chat_record(
db.pool(), "outbound", user_id, &ctx.account_id,
&text, "llm", context_token.as_deref(), &sent_id,
).await;
}
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)),
other => warn!(
"Send Consumer 收到不期望: {:?}",
std::mem::discriminant(&other)
),
}
}
info!("Send Consumer 已停止");
@@ -702,37 +871,38 @@ async fn send_consumer_loop(
// ─── 辅助函数 ───
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!("保存 auth 失败: {}", e); }
return Some(a);
}
} else if let Some(a) = file_state.load_auth() { return Some(a); }
None
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: &Option<Arc<Database>>, user_id: &str) -> Vec<HistoryEntry> {
let Some(db) = db else { return vec![]; };
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() },
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![] }
})
.collect(),
Err(e) => {
warn!("加载历史失败: {}", e);
vec![]
}
}
}
async fn load_summaries(db: &Option<Arc<Database>>, user_id: &str) -> Vec<String> {
let Some(db) = db else { return 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![] }
Err(e) => {
warn!("加载摘要失败: {}", e);
vec![]
}
}
}
+1 -1
View File
@@ -31,8 +31,8 @@
pub mod models;
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use sqlx::postgres::PgPoolOptions;
use std::time::Duration;
/// ## 数据库管理器
+39 -1
View File
@@ -14,7 +14,13 @@ use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use crate::state::AuthState;
/// 认证状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthState {
pub token: String,
pub account_id: String,
pub base_url: String,
}
// ─── 认证状态 ───
@@ -50,6 +56,38 @@ pub async fn save_auth(pool: &PgPool, auth: &AuthState) -> Result<(), String> {
Ok(())
}
// ─── 运行时状态 (get_updates_buf) ───
/// 从数据库加载运行时状态 (get_updates_buf)
pub async fn load_runtime(pool: &PgPool) -> Option<String> {
let row = sqlx::query_as::<_, (serde_json::Value,)>(&
"SELECT value FROM app_state WHERE key = 'runtime'",
)
.fetch_optional(pool)
.await
.ok()??;
row.0.as_str().map(|s| s.to_string())
}
/// 保存运行时状态 (get_updates_buf) 到数据库
pub async fn save_runtime(pool: &PgPool, buf: &str) -> Result<(), String> {
let value = serde_json::Value::String(buf.to_string());
sqlx::query(
r#"
INSERT INTO app_state (key, value, updated_at)
VALUES ('runtime', $1, NOW())
ON CONFLICT (key) DO UPDATE
SET value = EXCLUDED.value, updated_at = NOW()
"#,
)
.bind(&value)
.execute(pool)
.await
.map_err(|e| format!("保存 runtime 到数据库失败: {}", e))?;
Ok(())
}
// ─── 聊天记录 ───
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
+14 -3
View File
@@ -26,7 +26,7 @@ use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -41,7 +41,7 @@ use tokio::sync::Mutex;
/// 2. **异步模式** — executor 将工具调用入队到消息队列,返回 `ASYNC_MARKER + 工具名`
/// 表示"工具已异步入队,结果稍后通过 ToolResult 消息回喂"
pub type ToolExecutor = Arc<
dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>
dyn Fn(&str, &str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>
+ Send
+ Sync,
>;
@@ -105,6 +105,8 @@ pub struct Conversation {
tool_executor: Option<ToolExecutor>,
/// 是否有异步工具调用等待结果返回
pending_async_tools: Arc<AtomicBool>,
/// 待处理的异步工具数量(全部返回后才能 resume)
pending_async_count: Arc<AtomicU32>,
}
impl Conversation {
@@ -121,6 +123,7 @@ impl Conversation {
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
pending_async_tools: Arc::new(AtomicBool::new(false)),
pending_async_count: Arc::new(AtomicU32::new(0)),
})
}
@@ -257,6 +260,11 @@ impl Conversation {
self.run_tool_loop().await
}
/// 通知一个异步工具结果已返回,返回 true 表示所有异步工具均已返回
pub fn notify_tool_result(&self) -> bool {
self.pending_async_count.fetch_sub(1, Ordering::SeqCst) == 1
}
/// 工具循环核心:调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待
async fn run_tool_loop(&self) -> Result<ChatResult, String> {
let mut used_tools = false;
@@ -316,10 +324,11 @@ impl Conversation {
.add_assistant_tool_calls(calls.clone());
let mut has_async = false;
let mut async_count: u32 = 0;
for tc in &calls {
let result = match &self.tool_executor {
Some(exec) => {
match exec(&tc.function.name, &tc.function.arguments).await {
match exec(&tc.function.name, &tc.function.arguments, &tc.id).await {
Ok(r) => r,
Err(e) => format!("执行失败: {}", e),
}
@@ -329,6 +338,7 @@ impl Conversation {
if result.starts_with(ASYNC_MARKER) {
has_async = true;
async_count += 1;
tracing::info!("📦 {} → 异步入队", tc.function.name);
// 不添加到 session,等待 ToolResult 异步回喂
continue;
@@ -344,6 +354,7 @@ impl Conversation {
if has_async {
self.pending_async_tools.store(true, Ordering::SeqCst);
self.pending_async_count.store(async_count, Ordering::SeqCst);
return Ok(ChatResult {
reply: full_text,
used_tools,
+63 -95
View File
@@ -37,7 +37,7 @@ mod llm;
mod logger;
mod queue;
mod scheduler;
mod state;
mod tools;
mod wechat;
@@ -52,14 +52,6 @@ use wechat::client::WeChatClient;
#[tokio::main]
async fn main() {
let cli = Cli::parse();
let log_file = matches!(&cli.command, Commands::Daemon { log_file: true, .. });
let log_dir = dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ias")
.join("logs");
let log_dir_str = log_dir.to_string_lossy().to_string();
logger::init_logger(Some(&log_dir_str), log_file);
// .env 加载:先当前目录,再 ~/.ias/.env
if std::path::Path::new(".env").exists() {
dotenvy::dotenv().ok();
@@ -71,7 +63,19 @@ async fn main() {
}
}
let file_state = state::StateManager::new(".data/weixin-ilink");
let log_file = matches!(&cli.command, Commands::Daemon { log_file: true, .. });
let log_dir = std::env::var_os("LOG_FILE_PATH")
.map(|os_str| os_str.to_string_lossy().into_owned())
.unwrap_or_else(|| {
dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ias")
.join("logs")
.to_string_lossy()
.into_owned()
});
logger::init_logger(Some(&log_dir), log_file);
// 工具命令无需登录/数据库,提前处理
let tool_cmd = match &cli.command {
@@ -86,33 +90,32 @@ async fn main() {
let database = match Database::connect().await {
Ok(db) => {
info!("✅ 数据库连接成功");
Some(Arc::new(db))
Arc::new(db)
}
Err(e) => {
info!("数据库不可用,使用文件存储: {}", e);
None
error!("数据库连接失败: {}", e);
eprintln!("错误: 数据库连接失败 - {}", e);
std::process::exit(1);
}
};
// 长期记忆存储
let memory_store = Arc::new(MemoryStore::new(
database.as_ref().map(|db| Arc::new(db.pool().clone())),
));
let memory_store = Arc::new(MemoryStore::new(Some(Arc::new(database.pool().clone()))));
memory_store.load("").await;
match cli.command {
Commands::Login { timeout } => {
cmd_login(timeout, &database, &file_state).await;
cmd_login(timeout, &database).await;
}
Commands::Send {
to,
text,
context_token,
} => {
cmd_send(&to, &text, context_token.as_deref(), &database, &file_state).await;
cmd_send(&to, &text, context_token.as_deref(), &database).await;
}
Commands::Whoami => {
cmd_whoami(&database, &file_state).await;
cmd_whoami(&database).await;
}
Commands::Usage {
since,
@@ -121,8 +124,8 @@ async fn main() {
} => {
cmd_usage(&database, since, until, model).await;
}
Commands::Daemon { sock, .. } => {
cmd_daemon(sock, &database, &file_state, &memory_store).await;
Commands::Daemon { .. } => {
cmd_daemon(&database, &memory_store).await;
}
Commands::ScheduledTask(action) => {
cmd_scheduled_task(action, &database).await;
@@ -136,20 +139,11 @@ async fn main() {
// ─── 命令实现 ───
async fn cmd_daemon(
sock_path: String,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
memory_store: &Arc<MemoryStore>,
) {
daemon::run(sock_path, database, file_state, memory_store).await;
async fn cmd_daemon(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
daemon::run(database, memory_store).await;
}
async fn cmd_login(
timeout_secs: u64,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
) {
async fn cmd_login(timeout_secs: u64, database: &Arc<Database>) {
let client = WeChatClient::new(None);
match client
@@ -157,24 +151,17 @@ async fn cmd_login(
.await
{
Ok(result) => {
let auth = state::AuthState {
let auth = db::models::AuthState {
token: result.token.clone(),
account_id: result.account_id.clone(),
base_url: result.base_url.clone(),
};
// 优先存数据库
if let Some(db) = database {
if let Err(e) = db::models::save_auth(db.pool(), &auth).await {
// 存数据库
if let Err(e) = db::models::save_auth(database.pool(), &auth).await {
error!("保存 auth 到数据库失败: {}", e);
file_state.save_auth(&auth);
} else {
info!("认证信息已保存到数据库");
// 清理文件状态(迁移完成)
file_state.clear_auth();
}
} else {
file_state.save_auth(&auth);
}
println!("\n✅ 登录成功!");
@@ -189,46 +176,27 @@ async fn cmd_login(
}
}
async fn cmd_whoami(database: &Option<Arc<Database>>, file_state: &state::StateManager) {
let auth: Option<state::AuthState> = if let Some(db) = database {
db::models::load_auth(db.pool()).await
} else {
file_state.load_auth()
};
async fn cmd_whoami(database: &Arc<Database>) {
let auth: Option<db::models::AuthState> = db::models::load_auth(database.pool()).await;
match auth {
Some(a) => {
println!("账号: {}", a.account_id);
println!("Token: {}...", &a.token[..std::cmp::min(16, a.token.len())]);
println!("API: {}", a.base_url);
if database.is_some() {
println!("存储: PostgreSQL");
} else {
println!("存储: 文件 ({})", file_state.state_dir().display());
}
}
None => println!("未登录,请先执行 login 命令"),
}
}
async fn cmd_usage(
database: &Option<Arc<Database>>,
database: &Arc<Database>,
since: Option<String>,
until: Option<String>,
model: Option<String>,
) {
let db = match database {
Some(d) => d,
None => {
println!("数据库未配置,无法查询用量");
return;
}
};
let db = database;
let since_dt = since.and_then(|s| {
chrono::DateTime::parse_from_rfc3339(&s)
@@ -288,14 +256,9 @@ async fn cmd_send(
to: &str,
text: &str,
context_token: Option<&str>,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
database: &Arc<Database>,
) {
let auth: Option<state::AuthState> = if let Some(db) = database {
db::models::load_auth(db.pool()).await
} else {
file_state.load_auth()
};
let auth: Option<db::models::AuthState> = db::models::load_auth(database.pool()).await;
let auth = match auth {
Some(a) => a,
@@ -317,9 +280,8 @@ async fn cmd_send(
println!("✅ 消息已发送, msg_id={}", msg_id);
// 入库:发送的消息
if let Some(db) = database
&& let Err(e) = db::models::insert_chat_record(
db.pool(),
if let Err(e) = db::models::insert_chat_record(
database.pool(),
"outbound",
to,
&account_id,
@@ -337,21 +299,13 @@ async fn cmd_send(
}
}
// ─── 定时任务 CLI 命令 ───
async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<Database>>) {
let db = match database {
Some(d) => d,
None => {
println!("❌ 数据库未配置,无法管理定时任务");
return;
}
};
async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Arc<Database>) {
let db = database;
match action {
ScheduledTaskAction::List => {
match db::models::list_scheduled_tasks(db.pool()).await {
ScheduledTaskAction::List => match db::models::list_scheduled_tasks(db.pool()).await {
Ok(tasks) => {
if tasks.is_empty() {
println!("📋 没有定时任务");
@@ -361,7 +315,8 @@ async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<D
println!("{:=<80}", "");
for task in &tasks {
let status = if task.enabled { "🟢" } else { "🔴" };
let last = task.last_run_at
let last = task
.last_run_at
.map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
.unwrap_or_else(|| "-".to_string());
let next = task.next_run_at.format("%Y-%m-%d %H:%M:%S");
@@ -386,8 +341,7 @@ async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<D
}
}
Err(e) => println!("❌ 查询失败: {}", e),
}
}
},
ScheduledTaskAction::Add {
name,
user,
@@ -479,7 +433,11 @@ async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<D
};
match db::models::toggle_scheduled_task(db.pool(), &uuid).await {
Ok((true, enabled)) => {
let status = if enabled { "🟢 已启用" } else { "🔴 已禁用" };
let status = if enabled {
"🟢 已启用"
} else {
"🔴 已禁用"
};
println!("✅ 定时任务状态已切换: {}", status);
}
Ok((false, _)) => println!("⚠️ 未找到该任务"),
@@ -512,12 +470,21 @@ async fn cmd_tools() {
println!(" 超时: {}", spec.timeout_secs);
// 显示参数
if let Some(props) = spec.parameters.get("properties").and_then(|v| v.as_object()) {
if let Some(props) = spec
.parameters
.get("properties")
.and_then(|v| v.as_object())
{
if !props.is_empty() {
println!(" 参数:");
for (name, info) in props {
let desc = info.get("description").and_then(|v| v.as_str()).unwrap_or("");
let required = spec.parameters.get("required")
let desc = info
.get("description")
.and_then(|v| v.as_str())
.unwrap_or("");
let required = spec
.parameters
.get("required")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().any(|r| r.as_str() == Some(name)))
.unwrap_or(false);
@@ -544,7 +511,8 @@ async fn cmd_tool(cmd: ToolCommand) {
tools::builtins::memos::list_execute(serde_json::json!({})).await
}
MemosAction::Add { content } => {
tools::builtins::memos::add_execute(serde_json::json!({"content": content})).await
tools::builtins::memos::add_execute(serde_json::json!({"content": content}))
.await
}
MemosAction::Delete { id } => {
tools::builtins::memos::delete_execute(serde_json::json!({"id": id})).await
-116
View File
@@ -1,116 +0,0 @@
//! ## 文件状态管理(PostgreSQL 不可用时的回退方案)
//!
//! ### 设计意图
//! 在数据库不可用时,使用本地 JSON 文件持久化认证和运行时状态。
//! 当数据库可用时,数据优先存数据库,并从文件迁移到数据库后清理文件。
//!
//! ### 存储位置
//! 默认 `~/.ias/data/weixin-ilink/`
//! * `auth.json` — 认证信息(token, account_id, base_url
//! * `runtime.json` — 运行时状态(get_updates_buf
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
/// 认证状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthState {
pub token: String,
pub account_id: String,
pub base_url: String,
}
/// 运行时状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuntimeState {
pub get_updates_buf: String,
}
/// ## 状态管理器 —— 基于文件的状态持久化
///
/// 在数据库不可用时代替 DB 存储认证和运行时状态。
///
/// ### 文件格式
/// * `auth.json` — `AuthState { token, account_id, base_url }`
/// * `runtime.json` — `RuntimeState { get_updates_buf }`
///
/// ### 安全
/// Unix 系统上 auth.json 的权限被设置为 600(仅所有者可读写)
pub struct StateManager {
state_dir: PathBuf,
}
impl StateManager {
pub fn new(state_dir: impl Into<PathBuf>) -> Self {
Self {
state_dir: state_dir.into(),
}
}
fn auth_file(&self) -> PathBuf {
self.state_dir.join("auth.json")
}
fn runtime_file(&self) -> PathBuf {
self.state_dir.join("runtime.json")
}
fn ensure_dir(&self) {
std::fs::create_dir_all(&self.state_dir).ok();
}
pub fn state_dir(&self) -> &std::path::Path {
&self.state_dir
}
pub fn save_auth(&self, auth: &AuthState) {
self.ensure_dir();
let json = serde_json::to_string_pretty(auth).expect("序列化 auth 失败");
if let Err(e) = std::fs::write(self.auth_file(), &json) {
tracing::error!("保存 auth 状态失败: {}", e);
}
// 权限
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(self.auth_file(), std::fs::Permissions::from_mode(0o600)).ok();
}
}
pub fn load_auth(&self) -> Option<AuthState> {
let path = self.auth_file();
if !path.exists() {
return None;
}
let json = std::fs::read_to_string(&path).ok()?;
serde_json::from_str(&json).ok()
}
pub fn save_runtime(&self, buf: &str) {
self.ensure_dir();
let state = RuntimeState {
get_updates_buf: buf.to_string(),
};
let json = serde_json::to_string_pretty(&state).expect("序列化 runtime 失败");
if let Err(e) = std::fs::write(self.runtime_file(), &json) {
tracing::error!("保存 runtime 状态失败: {}", e);
}
}
pub fn load_runtime(&self) -> Option<RuntimeState> {
let path = self.runtime_file();
if !path.exists() {
return None;
}
let json = std::fs::read_to_string(&path).ok()?;
serde_json::from_str(&json).ok()
}
/// 删除 auth 文件(迁移到数据库后清理)
pub fn clear_auth(&self) {
let path = self.auth_file();
if path.exists() {
std::fs::remove_file(&path).ok();
}
}
}