fix: TodoWriteTool merge 模式内存为空时从 DB 回填,防止丢失旧项
todo_write.rs: 内存状态为空时从 repository.list_todos 回填,保证 merge 模式不丢失旧项(进程重启后内存清空场景)。新增 MockTodoRepository 和回填测试。 tool_registry_factory.rs: TodoWriteTool::new 传入 todo_repository。
This commit is contained in:
parent
7652bb16e2
commit
4da7b5f505
@ -124,7 +124,7 @@ impl ToolRegistryFactory {
|
|||||||
}
|
}
|
||||||
if self.is_enabled("todo_write") {
|
if self.is_enabled("todo_write") {
|
||||||
if let Some(ref state) = self.todo_state {
|
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()));
|
registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone()));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -232,7 +232,7 @@ impl ToolRegistryFactory {
|
|||||||
// Todo 追踪工具
|
// Todo 追踪工具
|
||||||
if self.is_enabled("todo_write") {
|
if self.is_enabled("todo_write") {
|
||||||
if let Some(ref state) = self.todo_state {
|
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()));
|
registry.register(TodoReadTool::new(state.clone(), self.todo_repository.clone()));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@ -6,6 +6,7 @@ use serde::Serialize;
|
|||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use tokio::sync::RwLock;
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
|
use crate::storage::TodoRepository;
|
||||||
use crate::tools::traits::{Tool, ToolContext, ToolResult};
|
use crate::tools::traits::{Tool, ToolContext, ToolResult};
|
||||||
|
|
||||||
// ── 数据模型 ──────────────────────────────────────────────
|
// ── 数据模型 ──────────────────────────────────────────────
|
||||||
@ -61,11 +62,16 @@ pub struct TodoWriteTool {
|
|||||||
/// 内存状态:scope_key → Vec<TodoItem>
|
/// 内存状态:scope_key → Vec<TodoItem>
|
||||||
/// scope_key = topic_id.unwrap_or(session_id)
|
/// scope_key = topic_id.unwrap_or(session_id)
|
||||||
state: Arc<RwLock<HashMap<String, Vec<TodoItem>>>>,
|
state: Arc<RwLock<HashMap<String, Vec<TodoItem>>>>,
|
||||||
|
/// 持久化仓库:内存为空时从 DB 回填,保证 merge 模式不丢失旧项
|
||||||
|
repository: Arc<dyn TodoRepository>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TodoWriteTool {
|
impl TodoWriteTool {
|
||||||
pub(crate) fn new(state: Arc<RwLock<HashMap<String, Vec<TodoItem>>>>) -> Self {
|
pub(crate) fn new(
|
||||||
Self { state }
|
state: Arc<RwLock<HashMap<String, Vec<TodoItem>>>>,
|
||||||
|
repository: Arc<dyn TodoRepository>,
|
||||||
|
) -> Self {
|
||||||
|
Self { state, repository }
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -158,10 +164,39 @@ impl Tool for TodoWriteTool {
|
|||||||
.and_then(|v| v.as_bool())
|
.and_then(|v| v.as_bool())
|
||||||
.unwrap_or(true);
|
.unwrap_or(true);
|
||||||
|
|
||||||
// 3. 读锁获取旧状态
|
// 3. 读锁获取旧状态;内存为空时从 DB 回填(与 TodoReadTool 一致,保证 merge 模式不丢失旧项)
|
||||||
let old_items = {
|
let old_items = {
|
||||||
let guard = self.state.read().await;
|
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<_>>()
|
||||||
|
}
|
||||||
|
_ => 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 的旧状态映射
|
// 构建 id → TodoItem 的旧状态映射
|
||||||
@ -444,9 +479,59 @@ mod tests {
|
|||||||
Arc::new(RwLock::new(HashMap::new()))
|
Arc::new(RwLock::new(HashMap::new()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct MockTodoRepository {
|
||||||
|
records: Vec<crate::storage::TodoRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TodoRepository for MockTodoRepository {
|
||||||
|
fn replace_todos(
|
||||||
|
&self,
|
||||||
|
_scope_key: &str,
|
||||||
|
_items: &[crate::storage::TodoRecord],
|
||||||
|
) -> Result<Vec<crate::storage::TodoRecord>, crate::storage::StorageError> {
|
||||||
|
Ok(vec![])
|
||||||
|
}
|
||||||
|
|
||||||
|
fn list_todos(
|
||||||
|
&self,
|
||||||
|
scope_key: &str,
|
||||||
|
) -> Result<Vec<crate::storage::TodoRecord>, crate::storage::StorageError> {
|
||||||
|
Ok(self
|
||||||
|
.records
|
||||||
|
.iter()
|
||||||
|
.filter(|r| r.scope_key == scope_key)
|
||||||
|
.cloned()
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mock_repo() -> Arc<MockTodoRepository> {
|
||||||
|
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]
|
#[tokio::test]
|
||||||
async fn test_create_initial_todos() {
|
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 context = test_context();
|
||||||
|
|
||||||
let result = tool
|
let result = tool
|
||||||
@ -473,7 +558,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_single_in_progress_constraint() {
|
async fn test_single_in_progress_constraint() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -510,7 +595,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_state_transition_in_progress_to_completed() {
|
async fn test_state_transition_in_progress_to_completed() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -560,7 +645,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_completed_can_revert_to_in_progress() {
|
async fn test_completed_can_revert_to_in_progress() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
// 创建并完成一个任务
|
// 创建并完成一个任务
|
||||||
@ -609,7 +694,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_cancelled_can_revert_to_pending() {
|
async fn test_cancelled_can_revert_to_pending() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -657,7 +742,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_in_progress_cannot_revert_to_pending() {
|
async fn test_in_progress_cannot_revert_to_pending() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -703,7 +788,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_new_item_can_be_any_status() {
|
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();
|
let context = test_context();
|
||||||
|
|
||||||
// 新项直接 completed — 应该允许(id 必填后不再限制初始状态)
|
// 新项直接 completed — 应该允许(id 必填后不再限制初始状态)
|
||||||
@ -724,7 +809,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_new_item_can_start_as_in_progress() {
|
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 context = test_context();
|
||||||
|
|
||||||
let result = tool
|
let result = tool
|
||||||
@ -751,7 +836,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_remove_items_by_omission() {
|
async fn test_remove_items_by_omission() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -791,7 +876,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_topic_isolation() {
|
async fn test_topic_isolation() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
|
|
||||||
let main_context = ToolContext {
|
let main_context = ToolContext {
|
||||||
session_id: Some("cli:chat-1".to_string()),
|
session_id: Some("cli:chat-1".to_string()),
|
||||||
@ -841,7 +926,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_empty_list() {
|
async fn test_empty_list() {
|
||||||
let tool = TodoWriteTool::new(test_state());
|
let tool = TodoWriteTool::new(test_state(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let result = tool
|
let result = tool
|
||||||
@ -859,7 +944,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_missing_todos_param() {
|
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 context = test_context();
|
||||||
|
|
||||||
let result = tool
|
let result = tool
|
||||||
@ -873,7 +958,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_no_context() {
|
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();
|
let result = tool.execute(json!({})).await.unwrap();
|
||||||
|
|
||||||
@ -884,7 +969,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_subagent_isolation() {
|
async fn test_subagent_isolation() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
|
|
||||||
let parent_ctx = ToolContext {
|
let parent_ctx = ToolContext {
|
||||||
session_id: Some("cli:chat-1".to_string()),
|
session_id: Some("cli:chat-1".to_string()),
|
||||||
@ -935,7 +1020,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_merge_mode_preserves_unreferenced_items() {
|
async fn test_merge_mode_preserves_unreferenced_items() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -977,7 +1062,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_merge_mode_add_new_item() {
|
async fn test_merge_mode_add_new_item() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -1015,7 +1100,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_merge_mode_never_removes() {
|
async fn test_merge_mode_never_removes() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -1052,7 +1137,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_non_merge_still_removes_by_omission() {
|
async fn test_non_merge_still_removes_by_omission() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -1091,7 +1176,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_merge_match_by_content_fallback() {
|
async fn test_merge_match_by_content_fallback() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
let _ = tool
|
let _ = tool
|
||||||
@ -1137,7 +1222,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_missing_id_validation_error() {
|
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 context = test_context();
|
||||||
|
|
||||||
let result = tool
|
let result = tool
|
||||||
@ -1159,7 +1244,7 @@ mod tests {
|
|||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_default_is_merge_mode() {
|
async fn test_default_is_merge_mode() {
|
||||||
let state = test_state();
|
let state = test_state();
|
||||||
let tool = TodoWriteTool::new(state.clone());
|
let tool = TodoWriteTool::new(state.clone(), mock_repo());
|
||||||
let context = test_context();
|
let context = test_context();
|
||||||
|
|
||||||
// 先创建 2 个 todo
|
// 先创建 2 个 todo
|
||||||
@ -1196,4 +1281,59 @@ mod tests {
|
|||||||
let task_a = todos.iter().find(|t| t["id"] == "x1").unwrap();
|
let task_a = todos.iter().find(|t| t["id"] == "x1").unwrap();
|
||||||
assert_eq!(task_a["status"], "in_progress");
|
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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user