diff --git a/src/gateway/tool_registry_factory.rs b/src/gateway/tool_registry_factory.rs index cffaccc..9554f64 100644 --- a/src/gateway/tool_registry_factory.rs +++ b/src/gateway/tool_registry_factory.rs @@ -124,7 +124,7 @@ impl ToolRegistryFactory { } if self.is_enabled("todo_write") { if let Some(ref state) = self.todo_state { - registry.register(TodoWriteTool::new(state.clone())); + registry.register(TodoWriteTool::new(state.clone(), self.todo_repository.clone())); registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone())); } } @@ -232,7 +232,7 @@ impl ToolRegistryFactory { // Todo 追踪工具 if self.is_enabled("todo_write") { if let Some(ref state) = self.todo_state { - registry.register(TodoWriteTool::new(state.clone())); + registry.register(TodoWriteTool::new(state.clone(), self.todo_repository.clone())); registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone())); } } diff --git a/src/tools/todo_write.rs b/src/tools/todo_write.rs index d70a912..a9d7d49 100644 --- a/src/tools/todo_write.rs +++ b/src/tools/todo_write.rs @@ -6,6 +6,7 @@ use serde::Serialize; use serde_json::json; use tokio::sync::RwLock; +use crate::storage::TodoRepository; use crate::tools::traits::{Tool, ToolContext, ToolResult}; // ── 数据模型 ────────────────────────────────────────────── @@ -61,11 +62,16 @@ pub struct TodoWriteTool { /// 内存状态:scope_key → Vec /// scope_key = topic_id.unwrap_or(session_id) state: Arc>>>, + /// 持久化仓库:内存为空时从 DB 回填,保证 merge 模式不丢失旧项 + repository: Arc, } impl TodoWriteTool { - pub(crate) fn new(state: Arc>>>) -> Self { - Self { state } + pub(crate) fn new( + state: Arc>>>, + repository: Arc, + ) -> Self { + Self { state, repository } } } @@ -158,10 +164,39 @@ impl Tool for TodoWriteTool { .and_then(|v| v.as_bool()) .unwrap_or(true); - // 3. 读锁获取旧状态 + // 3. 读锁获取旧状态;内存为空时从 DB 回填(与 TodoReadTool 一致,保证 merge 模式不丢失旧项) let old_items = { let guard = self.state.read().await; - guard.get(&scope_key).cloned().unwrap_or_default() + match guard.get(&scope_key).cloned() { + Some(items) if !items.is_empty() => items, + _ => { + drop(guard); + let db_items = match self.repository.list_todos(&scope_key) { + Ok(records) if !records.is_empty() => { + records + .into_iter() + .map(|r| TodoItem { + id: r.id, + content: r.content, + status: r.status, + created_by_message_id: r.created_by_message_id, + }) + .collect::>() + } + _ => Vec::new(), + }; + if !db_items.is_empty() { + let mut write_guard = self.state.write().await; + write_guard.insert(scope_key.clone(), db_items.clone()); + tracing::info!( + scope_key = %scope_key, + todo_count = db_items.len(), + "TodoWriteTool: backfilled memory from SQLite before merge" + ); + } + db_items + } + } }; // 构建 id → TodoItem 的旧状态映射 @@ -444,9 +479,59 @@ mod tests { Arc::new(RwLock::new(HashMap::new())) } + struct MockTodoRepository { + records: Vec, + } + + impl TodoRepository for MockTodoRepository { + fn replace_todos( + &self, + _scope_key: &str, + _items: &[crate::storage::TodoRecord], + ) -> Result, crate::storage::StorageError> { + Ok(vec![]) + } + + fn list_todos( + &self, + scope_key: &str, + ) -> Result, crate::storage::StorageError> { + Ok(self + .records + .iter() + .filter(|r| r.scope_key == scope_key) + .cloned() + .collect()) + } + } + + fn mock_repo() -> Arc { + Arc::new(MockTodoRepository { records: vec![] }) + } + + fn mock_record( + scope_key: &str, + id: &str, + content: &str, + status: &str, + ) -> crate::storage::TodoRecord { + crate::storage::TodoRecord { + id: id.to_string(), + scope_key: scope_key.to_string(), + session_id: "cli:chat-1".to_string(), + topic_id: None, + content: content.to_string(), + status: status.to_string(), + priority: "medium".to_string(), + created_at: 1000, + updated_at: 1000, + created_by_message_id: None, + } + } + #[tokio::test] async fn test_create_initial_todos() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); let result = tool @@ -473,7 +558,7 @@ mod tests { #[tokio::test] async fn test_single_in_progress_constraint() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -510,7 +595,7 @@ mod tests { #[tokio::test] async fn test_state_transition_in_progress_to_completed() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -560,7 +645,7 @@ mod tests { #[tokio::test] async fn test_completed_can_revert_to_in_progress() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); // 创建并完成一个任务 @@ -609,7 +694,7 @@ mod tests { #[tokio::test] async fn test_cancelled_can_revert_to_pending() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -657,7 +742,7 @@ mod tests { #[tokio::test] async fn test_in_progress_cannot_revert_to_pending() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -703,7 +788,7 @@ mod tests { #[tokio::test] async fn test_new_item_can_be_any_status() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); // 新项直接 completed — 应该允许(id 必填后不再限制初始状态) @@ -724,7 +809,7 @@ mod tests { #[tokio::test] async fn test_new_item_can_start_as_in_progress() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); let result = tool @@ -751,7 +836,7 @@ mod tests { #[tokio::test] async fn test_remove_items_by_omission() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -791,7 +876,7 @@ mod tests { #[tokio::test] async fn test_topic_isolation() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let main_context = ToolContext { session_id: Some("cli:chat-1".to_string()), @@ -841,7 +926,7 @@ mod tests { #[tokio::test] async fn test_empty_list() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); let result = tool @@ -859,7 +944,7 @@ mod tests { #[tokio::test] async fn test_missing_todos_param() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); let result = tool @@ -873,7 +958,7 @@ mod tests { #[tokio::test] async fn test_no_context() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let result = tool.execute(json!({})).await.unwrap(); @@ -884,7 +969,7 @@ mod tests { #[tokio::test] async fn test_subagent_isolation() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let parent_ctx = ToolContext { session_id: Some("cli:chat-1".to_string()), @@ -935,7 +1020,7 @@ mod tests { #[tokio::test] async fn test_merge_mode_preserves_unreferenced_items() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -977,7 +1062,7 @@ mod tests { #[tokio::test] async fn test_merge_mode_add_new_item() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -1015,7 +1100,7 @@ mod tests { #[tokio::test] async fn test_merge_mode_never_removes() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -1052,7 +1137,7 @@ mod tests { #[tokio::test] async fn test_non_merge_still_removes_by_omission() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -1091,7 +1176,7 @@ mod tests { #[tokio::test] async fn test_merge_match_by_content_fallback() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); let _ = tool @@ -1137,7 +1222,7 @@ mod tests { #[tokio::test] async fn test_missing_id_validation_error() { - let tool = TodoWriteTool::new(test_state()); + let tool = TodoWriteTool::new(test_state(), mock_repo()); let context = test_context(); let result = tool @@ -1159,7 +1244,7 @@ mod tests { #[tokio::test] async fn test_default_is_merge_mode() { let state = test_state(); - let tool = TodoWriteTool::new(state.clone()); + let tool = TodoWriteTool::new(state.clone(), mock_repo()); let context = test_context(); // 先创建 2 个 todo @@ -1196,4 +1281,59 @@ mod tests { let task_a = todos.iter().find(|t| t["id"] == "x1").unwrap(); assert_eq!(task_a["status"], "in_progress"); } + + #[tokio::test] + async fn test_merge_backfills_from_db_when_memory_empty() { + // 场景:进程重启后内存清空,DB 保留旧项 [a(completed), b(pending), c(in_progress)] + // 子智能体用 merge 模式只传 [{a, completed}],应保留 b、c,不丢失 + let state = test_state(); // 空 HashMap + let scope_key = "task:sub-1".to_string(); + let repo = Arc::new(MockTodoRepository { + records: vec![ + mock_record(&scope_key, "a", "任务A", "completed"), + mock_record(&scope_key, "b", "任务B", "pending"), + mock_record(&scope_key, "c", "任务C", "in_progress"), + ], + }); + let tool = TodoWriteTool::new(state.clone(), repo); + + let ctx = ToolContext { + session_id: Some("cli:chat-1".to_string()), + task_id: Some(scope_key.clone()), + nesting_depth: 1, // 模拟子智能体 + ..ToolContext::default() + }; + + // merge=true 只传 a(completed),b/c 应从 DB 回填并保留 + let args = json!({ + "merge": true, + "todos": [ + { "id": "a", "content": "任务A", "status": "completed" } + ] + }); + + let result = tool.execute_with_context(&ctx, args).await.unwrap(); + assert!(result.success, "merge should succeed: {:?}", result.error); + + let output: serde_json::Value = serde_json::from_str(&result.output).unwrap(); + let todos = output["current_todos"].as_array().unwrap(); + assert_eq!( + todos.len(), + 3, + "should preserve all 3 items after merge with DB backfill" + ); + + let ids: Vec<&str> = todos.iter().map(|t| t["id"].as_str().unwrap()).collect(); + assert!(ids.contains(&"a")); + assert!(ids.contains(&"b"), "pending item b must be preserved"); + assert!( + ids.contains(&"c"), + "in_progress item c must be preserved" + ); + + // 验证内存已被回填 + let guard = state.read().await; + let memory_items = guard.get(&scope_key).unwrap(); + assert_eq!(memory_items.len(), 3); + } }