AstroResearch/src/agent/tools/todo.rs

// 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<String> = 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<String> = 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
        );
    }
}