AstroResearch/src/agent/tools/background.rs

// src/agent/tools/background.rs — 后台任务工具 (bg_task_run / bg_task_check)
//
// 参考 Claude Code s08 Background Tasks 设计。
// 慢速操作可在后台异步执行，结果在下一轮 LLM 调用前注入上下文。

use async_trait::async_trait;
use serde_json::json;
use std::sync::Arc;
use tracing::info;

use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};
use crate::agent::background::{self, BgNotificationQueue};

/// 支持后台执行的工具列表
const BG_SUPPORTED_TOOLS: &[&str] = &["process_paper"];

/// 后台任务启动工具
pub struct BgTaskRunTool {
    queue: Arc<BgNotificationQueue>,
}

impl BgTaskRunTool {
    pub fn new(queue: Arc<BgNotificationQueue>) -> Self {
        BgTaskRunTool { queue }
    }
}

#[async_trait]
impl AgentTool for BgTaskRunTool {
    fn name(&self) -> &str {
        "bg_task_run"
    }

    fn display_name(&self) -> &str {
        "后台任务"
    }

    fn description(&self) -> &str {
        "在后台异步执行慢速工具（process_paper）。\
         返回任务ID后立即让 LLM 继续思考，后台完成的结果会在下一轮对话中自动通知。\
         适用于：文献下载、解析、向量化等耗时操作。使用 bg_task_check 查询任务状态。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "tool_name": {
                    "type": "string",
                    "description": "要在后台执行的工具名称",
                    "enum": ["process_paper"]
                },
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符（ADS bibcode）"
                },
                "tasks": {
                    "type": "array",
                    "items": {
                        "type": "string",
                        "enum": ["download", "parse", "embed"]
                    },
                    "description": "需执行的任务列表。不指定时默认执行 ['download', 'parse']"
                }
            },
            "required": ["tool_name", "bibcode"]
        })
    }

    /// 后台任务启动有副作用（spawn tokio task），中断时应阻塞以完成
    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Block
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let tool_name = match args.get("tool_name").and_then(|v| v.as_str()) {
            Some(s) => s.to_string(),
            None => return ToolOutput::error("缺少必需参数 'tool_name'"),
        };

        let bibcode = match args.get("bibcode").and_then(|v| v.as_str()) {
            Some(s) => s.to_string(),
            None => return ToolOutput::error("缺少必需参数 'bibcode'"),
        };

        if !BG_SUPPORTED_TOOLS.contains(&tool_name.as_str()) {
            return ToolOutput::error(format!(
                "工具 '{}' 不支持后台执行。支持的工具: {}",
                tool_name,
                BG_SUPPORTED_TOOLS.join(", ")
            ));
        }

        // 提取 tasks 参数（用于 process_paper）
        let tasks: Vec<String> = args
            .get("tasks")
            .and_then(|t| t.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_else(|| vec!["download".into(), "parse".into()]);

        info!(
            "[BgTaskRun] 启动后台任务: tool={}, bibcode={}, tasks={:?}",
            tool_name, bibcode, tasks
        );

        let handle = background::spawn_background_task(
            ctx.app_state.clone(),
            self.queue.clone(),
            tool_name.clone(),
            bibcode.clone(),
            tasks,
        )
        .await;

        ToolOutput::success(
            format!(
                "✅ 后台任务已启动。\n\
                 任务ID: {}\n\
                 工具: {}\n\
                 文献: {}\n\
                 状态: 运行中\n\n\
                 使用 bg_task_check 查询任务状态。完成后结果会自动通知。",
                handle.task_id, handle.tool_name, handle.bibcode,
            ),
            json!({
                "task_id": handle.task_id,
                "tool_name": handle.tool_name,
                "bibcode": handle.bibcode,
                "status": "running"
            }),
        )
    }

    fn defer_loading(&self) -> bool {
        true
    }
}

/// 后台任务查询工具
pub struct BgTaskCheckTool {
    queue: Arc<BgNotificationQueue>,
}

impl BgTaskCheckTool {
    pub fn new(queue: Arc<BgNotificationQueue>) -> Self {
        BgTaskCheckTool { queue }
    }
}

#[async_trait]
impl AgentTool for BgTaskCheckTool {
    fn name(&self) -> &str {
        "bg_task_check"
    }

    fn display_name(&self) -> &str {
        "检查后台"
    }

    fn is_internal(&self) -> bool {
        true
    }

    fn description(&self) -> &str {
        "查询后台任务状态。不指定 task_id 时返回所有任务。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "task_id": {
                    "type": "string",
                    "description": "可选：要查询的任务ID。不指定则返回所有任务。"
                }
            },
            "required": []
        })
    }

    /// 纯读取内存队列状态，并发安全
    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
        let task_id = args.get("task_id").and_then(|v| v.as_str());

        match task_id {
            Some(tid) => match self.queue.get_task(tid).await {
                Some(task) => {
                    let status_icon = match task.status {
                        background::BgTaskStatus::Running => "🔄",
                        background::BgTaskStatus::Completed => "✅",
                        background::BgTaskStatus::Failed => "❌",
                    };
                    ToolOutput::success(
                        format!(
                            "{} 任务 {}: {} ({})\n文献: {}",
                            status_icon,
                            task.task_id,
                            task.tool_name,
                            match task.status {
                                background::BgTaskStatus::Running => "运行中",
                                background::BgTaskStatus::Completed => "已完成",
                                background::BgTaskStatus::Failed => "失败",
                            },
                            task.bibcode,
                        ),
                        serde_json::to_value(&task).unwrap_or_default(),
                    )
                }
                None => ToolOutput::error(format!("任务 '{}' 未找到", tid)),
            },
            None => {
                let tasks = self.queue.get_all_tasks().await;
                if tasks.is_empty() {
                    return ToolOutput::success("当前没有后台任务。", json!({ "tasks": [] }));
                }

                let mut lines = vec!["📊 后台任务状态：\n".to_string()];
                for task in &tasks {
                    let icon = match task.status {
                        background::BgTaskStatus::Running => "🔄",
                        background::BgTaskStatus::Completed => "✅",
                        background::BgTaskStatus::Failed => "❌",
                    };
                    lines.push(format!(
                        "{} [{}] {} — {}",
                        icon, task.task_id, task.tool_name, task.bibcode
                    ));
                }

                ToolOutput::success(lines.join("\n"), json!({ "tasks": tasks }))
            }
        }
    }

    fn defer_loading(&self) -> bool {
        true
    }
}