diff --git a/migrations/20260601000015_dummy_noop.sql b/migrations/20260601000015_dummy_noop.sql new file mode 100644 index 0000000..9eba19c --- /dev/null +++ b/migrations/20260601000015_dummy_noop.sql @@ -0,0 +1,4 @@ +-- 空迁移：用于对齐 _sqlx_migrations 表中已记录但文件丢失的迁移 +-- 该迁移已被之前的某次运行应用过，但对应文件被删除/重命名 +-- 此文件仅用于让 sqlx 迁移校验通过，不执行任何 DDL +SELECT 1; diff --git a/src/cli.rs b/src/cli.rs index 63cc105..51e7562 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -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, diff --git a/src/context/builder.rs b/src/context/builder.rs index 120a9cb..4811e57 100644 --- a/src/context/builder.rs +++ b/src/context/builder.rs @@ -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 { + 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, 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 = 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>, 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 = 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); + } } diff --git a/src/daemon.rs b/src/daemon.rs index 976c6f8..10c47a4 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -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>, + /// 数据库句柄 + db: Arc, /// WeChat API 客户端（内部用 Arc 保护 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>, - file_state: &StateManager, - memory_store: &Arc, -) { +pub async fn run(database: &Arc, memory_store: &Arc) { // 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); } }); @@ -140,32 +148,43 @@ pub async fn run( // 9. 队列（容量 1024） let (runner, channels) = QueueRunner::new(1024); let enqueue_handle = runner.enqueue_handle(); - let shutdown_handle = runner.shutdown_handle(); // 不再 prefix _ + 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_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; + 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; }); + 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, + 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 { @@ -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, enqueue: EnqueueHandle, -) -> (Vec>, Arc>>) { - let sessions: Arc>> = Arc::new(Mutex::new(HashMap::new())); +) -> ( + Vec>, + Arc>>, +) { + let sessions: Arc>> = + 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::>().join("\n---\n"); + .collect::>() + .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; - } + 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)), + 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>, - file_state: &StateManager, -) -> Option { - 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) -> Option { + crate::db::models::load_auth(database.pool()).await } -async fn load_history(db: &Option>, user_id: &str) -> Vec { - let Some(db) = db else { return vec![]; }; +async fn load_history(db: &Arc, user_id: &str) -> 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![] } + 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: &Option>, user_id: &str) -> Vec { - let Some(db) = db else { return vec![]; }; +async fn load_summaries(db: &Arc, user_id: &str) -> Vec { 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![] + } } } diff --git a/src/db/mod.rs b/src/db/mod.rs index 80afc73..248c867 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -31,8 +31,8 @@ pub mod models; -use sqlx::postgres::PgPoolOptions; use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; use std::time::Duration; /// ## 数据库管理器 diff --git a/src/db/models.rs b/src/db/models.rs index 83f3914..8022bf9 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -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 { + 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)] diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index 0aedb8d..9d62844 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -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> + Send>> + dyn Fn(&str, &str, &str) -> Pin> + Send>> + Send + Sync, >; @@ -105,6 +105,8 @@ pub struct Conversation { tool_executor: Option, /// 是否有异步工具调用等待结果返回 pending_async_tools: Arc, + /// 待处理的异步工具数量（全部返回后才能 resume） + pending_async_count: Arc, } 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 { 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, diff --git a/src/main.rs b/src/main.rs index 9175158..db157f7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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>, - file_state: &state::StateManager, - memory_store: &Arc, -) { - daemon::run(sock_path, database, file_state, memory_store).await; +async fn cmd_daemon(database: &Arc, memory_store: &Arc) { + daemon::run(database, memory_store).await; } -async fn cmd_login( - timeout_secs: u64, - database: &Option>, - file_state: &state::StateManager, -) { +async fn cmd_login(timeout_secs: u64, database: &Arc) { 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 { - error!("保存 auth 到数据库失败: {}", e); - file_state.save_auth(&auth); - } else { - info!("认证信息已保存到数据库"); - // 清理文件状态（迁移完成） - file_state.clear_auth(); - } + // 存数据库 + if let Err(e) = db::models::save_auth(database.pool(), &auth).await { + error!("保存 auth 到数据库失败: {}", e); } else { - file_state.save_auth(&auth); + info!("认证信息已保存到数据库"); } println!("\n✅ 登录成功！"); @@ -189,46 +176,27 @@ async fn cmd_login( } } - - - - - -async fn cmd_whoami(database: &Option>, file_state: &state::StateManager) { - let auth: Option = if let Some(db) = database { - db::models::load_auth(db.pool()).await - } else { - file_state.load_auth() - }; +async fn cmd_whoami(database: &Arc) { + let auth: Option = 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()); - } + println!("存储: PostgreSQL"); } None => println!("未登录，请先执行 login 命令"), } } async fn cmd_usage( - database: &Option>, + database: &Arc, since: Option, until: Option, model: Option, ) { - 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>, - file_state: &state::StateManager, + database: &Arc, ) { - let auth: Option = if let Some(db) = database { - db::models::load_auth(db.pool()).await - } else { - file_state.load_auth() - }; + let auth: Option = 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, @@ -329,65 +291,57 @@ async fn cmd_send( &msg_id, ) .await - { - error!("保存聊天记录失败: {}", e); - } + { + error!("保存聊天记录失败: {}", e); + } } Err(e) => error!("发送失败: {}", e), } } - // ─── 定时任务 CLI 命令 ─── -async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option>) { - let db = match database { - Some(d) => d, - None => { - println!("❌ 数据库未配置，无法管理定时任务"); - return; - } - }; +async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Arc) { + let db = database; match action { - ScheduledTaskAction::List => { - match db::models::list_scheduled_tasks(db.pool()).await { - Ok(tasks) => { - if tasks.is_empty() { - println!("📋 没有定时任务"); - return; - } - println!("📋 定时任务列表:"); - println!("{:=<80}", ""); - for task in &tasks { - let status = if task.enabled { "🟢" } else { "🔴" }; - 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"); - let interval = if task.interval_seconds >= 86400 { - format!("{}天", task.interval_seconds / 86400) - } else if task.interval_seconds >= 3600 { - format!("{}小时", task.interval_seconds / 3600) - } else { - format!("{}秒", task.interval_seconds) - }; - println!("{} {}", status, task.name); - println!(" ID: {}", task.id); - println!(" 用户: {}", task.user_id.as_deref().unwrap_or("-")); - println!(" 命令: {}", task.command); - println!(" 间隔: {}", interval); - println!(" 上次执行: {}", last); - println!(" 下次执行: {}", next); - if let Some(ref s) = task.last_status { - println!(" 上次状态: {:.100}", s); - } - println!("{:=<80}", ""); - } + ScheduledTaskAction::List => match db::models::list_scheduled_tasks(db.pool()).await { + Ok(tasks) => { + if tasks.is_empty() { + println!("📋 没有定时任务"); + return; + } + println!("📋 定时任务列表:"); + println!("{:=<80}", ""); + for task in &tasks { + let status = if task.enabled { "🟢" } else { "🔴" }; + 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"); + let interval = if task.interval_seconds >= 86400 { + format!("{}天", task.interval_seconds / 86400) + } else if task.interval_seconds >= 3600 { + format!("{}小时", task.interval_seconds / 3600) + } else { + format!("{}秒", task.interval_seconds) + }; + println!("{} {}", status, task.name); + println!(" ID: {}", task.id); + println!(" 用户: {}", task.user_id.as_deref().unwrap_or("-")); + println!(" 命令: {}", task.command); + println!(" 间隔: {}", interval); + println!(" 上次执行: {}", last); + println!(" 下次执行: {}", next); + if let Some(ref s) = task.last_status { + println!(" 上次状态: {:.100}", s); + } + println!("{:=<80}", ""); } - Err(e) => println!("❌ 查询失败: {}", e), } - } + Err(e) => println!("❌ 查询失败: {}", e), + }, ScheduledTaskAction::Add { name, user, @@ -479,7 +433,11 @@ async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option { - 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 diff --git a/src/state.rs b/src/state.rs deleted file mode 100644 index 994fdf1..0000000 --- a/src/state.rs +++ /dev/null @@ -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) -> 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 { - 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 { - 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(); - } - } -}