// src/agent/tools/todo.rs — 任务规划工具 (TodoWrite) // // P1 改进:任务状态持久化到 SQLite agent_tasks 表, // 支持 blockedBy 依赖关系和跨 turn 状态恢复。 use async_trait::async_trait; use serde_json::json; use sqlx::SqlitePool; use std::str::FromStr; use tracing::{info, warn}; use super::{AgentTool, ToolContext, ToolOutput}; /// 任务规划工具:让 LLM 在开始复杂研究前先制定计划,执行中更新进度。 /// 任务状态持久化到 SQLite,支持 DAG 依赖。 pub struct TodoWriteTool; #[async_trait] impl AgentTool for TodoWriteTool { fn name(&self) -> &str { "todo_write" } fn display_name(&self) -> &str { "任务管理" } fn is_internal(&self) -> bool { true } fn description(&self) -> &str { "任务规划工具。在开始复杂研究前列出待办事项,执行中标记进度(每项状态:pending/in_progress/completed)。\ 一次只能有一个 in_progress 任务。支持任务依赖(blockedBy:依赖的其他任务ID列表)。\ 适用范围:任何需要多步工具调用的研究任务。" } fn parameters(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "todos": { "type": "array", "description": "任务列表,每项包含 id(唯一标识)、content(任务描述)、status(pending/in_progress/completed)、blockedBy(可选,依赖的其他任务ID列表)", "items": { "type": "object", "properties": { "id": { "type": "string", "description": "任务唯一标识" }, "content": { "type": "string", "description": "任务描述" }, "status": { "type": "string", "enum": ["pending", "in_progress", "completed"], "description": "任务状态" }, "blockedBy": { "type": "array", "items": { "type": "string" }, "description": "该任务依赖的其他任务ID列表(这些任务必须先完成)" } }, "required": ["id", "content", "status"] } } }, "required": ["todos"] }) } /// 纯格式化输出,无副作用(持久化由 runtime 层处理),并发安全 fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { true } async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { let todos = match args.get("todos").and_then(|t| t.as_array()) { Some(t) => t, None => return ToolOutput::error("缺少必需参数 'todos'"), }; // 从 ToolContext 中无法直接获取 session_id // TodoWrite 工具生成格式化的输出,实际持久化在 runtime 层完成 // 这里只做验证和格式化 let mut formatted = String::from("📋 当前任务计划:\n\n"); let mut in_progress_count = 0; for todo in todos { let id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("?"); let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or("?"); let status = todo .get("status") .and_then(|v| v.as_str()) .unwrap_or("pending"); let task_status = crate::agent::enums::TaskStatus::from_str(status) .unwrap_or(crate::agent::enums::TaskStatus::Pending); let icon = task_status.icon(); if task_status == crate::agent::enums::TaskStatus::InProgress { in_progress_count += 1; } // 显示依赖关系 let blocked_by: Vec = todo .get("blockedBy") .and_then(|b| b.as_array()) .map(|arr| { arr.iter() .filter_map(|v| v.as_str().map(String::from)) .collect() }) .unwrap_or_default(); let mut line = format!("{} [{}] {}", icon, id, content); if !blocked_by.is_empty() { line.push_str(&format!(" (依赖: {})", blocked_by.join(", "))); } formatted.push_str(&line); formatted.push('\n'); } // 验证约束 if in_progress_count > 1 { formatted.push_str( "\n⚠️ 提醒:你当前有多个 in_progress 任务。请先完成当前任务再开始下一个。", ); } else if in_progress_count == 0 && todos .iter() .any(|t| t.get("status").and_then(|v| v.as_str()) == Some("pending")) { formatted .push_str("\n💡 提示:还有待处理任务,请选择一个设为 in_progress 并开始执行。"); } ToolOutput::success(formatted, json!({ "task_count": todos.len() })) } } /// 将 TodoWrite 的任务列表持久化到 agent_tasks 表。 /// /// 使用 INSERT OR REPLACE 实现 upsert(基于 session_id + task_id 唯一约束)。 /// `owner` 参数指定任务的归属 agent(默认 "lead")。 pub async fn persist_tasks( db: &SqlitePool, session_id: &str, todos: &[serde_json::Value], owner: &str, ) -> anyhow::Result<()> { for todo in todos { let task_id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("unknown"); let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or(""); let status = todo .get("status") .and_then(|v| v.as_str()) .unwrap_or("pending"); // 收集 blockedBy 数组 let blocked_by: Vec = todo .get("blockedBy") .and_then(|b| b.as_array()) .map(|arr| { arr.iter() .filter_map(|v| v.as_str().map(String::from)) .collect() }) .unwrap_or_default(); let blocked_by_json = serde_json::to_string(&blocked_by).unwrap_or_else(|_| "[]".to_string()); // 简单 DAG 验证:不能依赖自身 if blocked_by.contains(&task_id.to_string()) { warn!( "[TodoWrite] 任务 {} 依赖自身,已跳过 blockedBy 中的自引用", task_id ); } sqlx::query( "INSERT INTO agent_tasks (session_id, task_id, content, status, blocked_by, owner) \ VALUES (?, ?, ?, ?, ?, ?) \ ON CONFLICT(session_id, task_id) DO UPDATE SET \ content=excluded.content, \ status=excluded.status, \ blocked_by=excluded.blocked_by, \ owner=excluded.owner, \ updated_at=CURRENT_TIMESTAMP", ) .bind(session_id) .bind(task_id) .bind(content) .bind(status) .bind(&blocked_by_json) .bind(owner) .execute(db) .await?; } info!( "[TodoWrite] 已持久化 {} 个任务到会话 {}", todos.len(), session_id ); Ok(()) } #[cfg(test)] mod tests { use super::*; use crate::agent::runtime::file_cache::FileStateCache; use crate::api::AppState; use crate::clients::ads::AdsClient; use crate::clients::arxiv::ArxivClient; use crate::clients::cds::vizier::VizierClient; use crate::clients::llm::{EmbeddingClient, LlmClient}; use crate::clients::qiniu::QiniuClient; use crate::services::batch::asset::AssetBatchStatus; use crate::services::batch::meta::MetaSyncStatus; use crate::services::download::Downloader; use crate::services::translation::Dictionary; use crate::Config; use std::path::PathBuf; use std::sync::Arc; /// 构建一个最小化的 ToolContext 用于测试(不依赖真实外部服务) async fn make_test_tool_context() -> ToolContext { let pool = SqlitePool::connect("sqlite::memory:") .await .expect("in-memory db"); sqlx::migrate!("./migrations") .run(&pool) .await .expect("migrations"); let config = Config::from_env(); let llm = LlmClient::new( "test_key".into(), "http://localhost".into(), "test-model".into(), ) .unwrap(); let medium_llm = llm.clone(); let fast_llm = llm.clone(); let embedding = EmbeddingClient::new( "test_key".into(), "http://localhost".into(), "test-embed".into(), ) .unwrap(); let ads = AdsClient::new("test_token".into()).unwrap(); let arxiv = ArxivClient::new().unwrap(); let vizier = VizierClient::new("https://tapvizier.cds.unistra.fr/TAPVizieR/tap", 60).unwrap(); let qiniu = QiniuClient::new( "test_ak".into(), "test_sk".into(), "test_bucket".into(), "http://localhost".into(), ); let downloader = Downloader::new().expect("failed to create downloader"); let dict = Dictionary::default(); let app_state = Arc::new(AppState { config, db: pool, dict, qiniu, downloader, http_client: reqwest::Client::new(), harvest_status: Arc::new(tokio::sync::Mutex::new(MetaSyncStatus::default())), batch_status: Arc::new(tokio::sync::Mutex::new(AssetBatchStatus::default())), active_bibcode: Arc::new(tokio::sync::Mutex::new(None)), skill_registry: Arc::new(tokio::sync::RwLock::new( crate::agent::skills::SkillRegistry::new(PathBuf::from("/tmp/test_skills")), )), sse_broadcast: None, memory_manager: Arc::new(tokio::sync::Mutex::new( crate::agent::memory::MemoryManager::new(PathBuf::from("/tmp/test_memory")), )), login_rate_limiter: Arc::new(dashmap::DashMap::new()), upload_rate_limiter: Arc::new(dashmap::DashMap::new()), llm: crate::api::LlmState { primary: llm, medium: medium_llm, fast: fast_llm, vision: None, embedding, }, sources: crate::api::DataSourceState { ads, arxiv, vizier, lamost: crate::clients::lamost::LamostClient::new("https://www.lamost.org", 60) .unwrap(), gaia: crate::clients::gaia::GaiaClient::new( "https://gea.esac.esa.int/tap-server/tap", "https://gea.esac.esa.int/data-server", 90, ) .unwrap(), sdss: crate::clients::sdss::SdssClient::new( "https://datalab.noirlab.edu/tap/sync", 90, ) .unwrap(), desi: crate::clients::desi::DesiClient::new( "https://datalab.noirlab.edu/tap/sync", 120, ) .unwrap(), irsa: crate::clients::irsa::IrsaClient::new("https://irsa.ipac.caltech.edu", 60) .unwrap(), mast: crate::clients::mast::MastClient::new("https://mast.stsci.edu", 90).unwrap(), observation_registry: std::sync::Arc::new( crate::services::observation::ObservationRegistry::default(), ), }, session: crate::api::SessionState { sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())), session_last_active: Arc::new(dashmap::DashMap::new()), cancelled_runs: Arc::new(dashmap::DashMap::new()), session_permission_checkers: Arc::new(dashmap::DashMap::new()), pending_questions: Arc::new(tokio::sync::Mutex::new( std::collections::HashMap::new(), )), pending_permissions: Arc::new(tokio::sync::Mutex::new( std::collections::HashMap::new(), )), }, }); ToolContext { app_state, session_id: "test_session".into(), silent: true, read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())), sse_tx: None, enable_thinking: false, additional_allowed_dirs: vec![], tool_call_id: None, max_output_chars: 4000, } } #[test] fn test_todo_output_format() { let tool = TodoWriteTool; // 验证参数 schema 包含必需字段 let params = tool.parameters(); assert!(params["required"] .as_array() .unwrap() .contains(&json!("todos"))); } #[tokio::test] async fn test_todo_multiple_in_progress_warning() { let tool = TodoWriteTool; let ctx = make_test_tool_context().await; let args = json!({ "todos": [ {"id": "1", "content": "任务A", "status": "in_progress"}, {"id": "2", "content": "任务B", "status": "in_progress"} ] }); let output = tool.execute(args, &ctx).await; assert!(!output.is_error, "should succeed"); assert!( output.content.contains("多个 in_progress 任务"), "should warn about multiple in_progress tasks, got: {}", output.content ); assert!( output.content.contains("🔄"), "should show in_progress icons" ); } #[tokio::test] async fn test_todo_with_blocked_by() { let tool = TodoWriteTool; let ctx = make_test_tool_context().await; let args = json!({ "todos": [ {"id": "1", "content": "文献检索", "status": "completed"}, {"id": "2", "content": "文献分析", "status": "pending", "blockedBy": ["1"]} ] }); let output = tool.execute(args, &ctx).await; assert!(!output.is_error, "should succeed"); assert!( output.content.contains("依赖: 1"), "should display dependency, got: {}", output.content ); assert!(output.content.contains("✅"), "should show completed icon"); assert!(output.content.contains("⏳"), "should show pending icon"); } #[tokio::test] async fn test_todo_single_in_progress_no_warning() { let tool = TodoWriteTool; let ctx = make_test_tool_context().await; let args = json!({ "todos": [ {"id": "1", "content": "任务A", "status": "in_progress"}, {"id": "2", "content": "任务B", "status": "pending"} ] }); let output = tool.execute(args, &ctx).await; assert!(!output.is_error, "should succeed"); assert!( !output.content.contains("多个 in_progress"), "should NOT warn when only one in_progress, got: {}", output.content ); // 只有一个 in_progress 时不显示"待处理任务"提示(该提示仅在 in_progress_count==0 时触发) } #[tokio::test] async fn test_todo_no_in_progress_with_pending_hint() { let tool = TodoWriteTool; let ctx = make_test_tool_context().await; let args = json!({ "todos": [ {"id": "1", "content": "任务A", "status": "completed"}, {"id": "2", "content": "任务B", "status": "pending"} ] }); let output = tool.execute(args, &ctx).await; assert!(!output.is_error, "should succeed"); assert!( output.content.contains("待处理任务"), "should hint about pending tasks when no in_progress, got: {}", output.content ); } #[tokio::test] async fn test_todo_missing_todos_field() { let tool = TodoWriteTool; let ctx = make_test_tool_context().await; let args = json!({}); let output = tool.execute(args, &ctx).await; assert!(output.is_error, "should return error for missing todos"); assert!( output.content.contains("缺少必需参数"), "should mention missing param, got: {}", output.content ); } }