AstroResearch/src/agent/coordinator/tools.rs

// src/agent/coordinator/tools.rs
//
// Coordinator 元工具：delegate_task、check_task、task_stop、synthesize。
//
// 每个工具实现 AgentTool trait。工具间通过 Arc<WorkerPool> 共享任务板。

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

use super::worker::WorkerPool;
use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};

// ── delegate_task ──

pub struct DelegateTaskTool {
    pool: Arc<WorkerPool>,
    enable_thinking: bool,
}

impl DelegateTaskTool {
    pub fn new(pool: Arc<WorkerPool>, enable_thinking: bool) -> Self {
        DelegateTaskTool {
            pool,
            enable_thinking,
        }
    }
}

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

    fn description(&self) -> &str {
        "将子任务委托给 Worker 代理执行。Worker 拥有完整工具访问权限。\
         使用此工具将独立子任务分发给 Worker，然后使用 check_task 检查进度。\
         可以先委托所有任务再统一检查，也可以逐个委托立即检查。\
         适用于：多篇文献综述、并行数据收集、独立计算任务。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "task_description": {
                    "type": "string",
                    "description": "要委托的任务完整描述（Worker 会据此开始工作）"
                },
                "context": {
                    "type": "string",
                    "description": "可选：附加上下文信息帮助 Worker 理解任务背景"
                }
            },
            "required": ["task_description"]
        })
    }

    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Cancel
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let task_desc = match args.get("task_description").and_then(|v| v.as_str()) {
            Some(s) => s.to_string(),
            None => return ToolOutput::error("缺少 task_description 参数"),
        };
        let context = args.get("context").and_then(|v| v.as_str()).unwrap_or("");
        let full_desc = if context.is_empty() {
            task_desc
        } else {
            format!("{}\n\n附加上下文：{}", task_desc, context)
        };

        let task_id = self
            .pool
            .delegate(&full_desc, &ctx.session_id, self.enable_thinking)
            .await;

        info!("[Coordinator] 委托任务: id={}", task_id);
        ToolOutput::success(
            format!(
                "任务已委托给 Worker。task_id: {}\n使用 check_task(task_id=\"{}\") 查看进度。",
                task_id, task_id
            ),
            json!({"task_id": task_id}),
        )
    }
}

// ── check_task ──

pub struct CheckTaskTool {
    pool: Arc<WorkerPool>,
}

impl CheckTaskTool {
    pub fn new(pool: Arc<WorkerPool>) -> Self {
        CheckTaskTool { pool }
    }
}

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

    fn description(&self) -> &str {
        "查询已委托任务的当前状态和结果。返回 Pending/Running/Completed/Failed/TimedOut。\
         对于已完成的任务，结果包含在返回中。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "task_id": {
                    "type": "string",
                    "description": "要查询的任务 ID（delegate_task 返回的）"
                }
            },
            "required": ["task_id"]
        })
    }

    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Cancel
    }

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

        match self.pool.check(task_id).await {
            Some(task) => {
                let status_str = match &task.status {
                    super::CoordinatorTaskStatus::Pending => "Pending",
                    super::CoordinatorTaskStatus::Running => "Running",
                    super::CoordinatorTaskStatus::Completed => "Completed",
                    super::CoordinatorTaskStatus::Failed { .. } => "Failed",
                    super::CoordinatorTaskStatus::TimedOut => "TimedOut",
                };
                ToolOutput::success(
                    format!(
                        "任务 {}: {}\n描述: {}\n结果: {}",
                        task.id,
                        status_str,
                        task.description,
                        task.result.as_deref().unwrap_or("(尚无结果)")
                    ),
                    json!({
                        "task": {
                            "id": task.id,
                            "status": status_str,
                            "description": task.description,
                            "result": task.result,
                        }
                    }),
                )
            }
            None => ToolOutput::error(format!("未找到 task_id: {}", task_id)),
        }
    }
}

// ── task_stop ──

pub struct TaskStopTool {
    pool: Arc<WorkerPool>,
}

impl TaskStopTool {
    pub fn new(pool: Arc<WorkerPool>) -> Self {
        TaskStopTool { pool }
    }
}

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

    fn description(&self) -> &str {
        "停止一个正在运行的任务（标记为失败）。如果任务已完成，此操作无效果。\
         适用于发现 Worker 方向错误时提前终止。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "task_id": {
                    "type": "string",
                    "description": "要停止的任务 ID"
                }
            },
            "required": ["task_id"]
        })
    }

    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Cancel
    }

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

        match self.pool.check(task_id).await {
            Some(_task) => {
                // 软取消 —— 标记为 Failed。
                // 完整的取消需要 CancellationToken 传入 Worker，这留待 P3 增强。
                let mut tasks = self.pool.tasks.write().await;
                if let Some(t) = tasks.iter_mut().find(|t| t.id == task_id) {
                    if matches!(
                        t.status,
                        super::CoordinatorTaskStatus::Running
                            | super::CoordinatorTaskStatus::Pending
                    ) {
                        t.status = super::CoordinatorTaskStatus::Failed {
                            reason: "协调者手动取消".into(),
                        };
                        t.completed_at = Some(chrono::Utc::now());
                    }
                }
                ToolOutput::success(
                    format!("任务 {} 已标记为取消。", task_id),
                    json!({"status": "cancelled", "task_id": task_id}),
                )
            }
            None => ToolOutput::error(format!("未找到 task_id: {}", task_id)),
        }
    }
}

// ── synthesize ──

pub struct SynthesizeTool {
    pool: Arc<WorkerPool>,
}

impl SynthesizeTool {
    pub fn new(pool: Arc<WorkerPool>) -> Self {
        SynthesizeTool { pool }
    }
}

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

    fn description(&self) -> &str {
        "收集所有已完成 Worker 任务的结果并合成最终答案。\
         在所有委托的任务完成后调用此工具，将各 Worker 结果合并为连贯的最终输出。\
         如果还有未完成的任务，工具会明确提示。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {},
            "required": []
        })
    }

    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Cancel
    }

    async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput {
        let all = self.pool.all_tasks().await;
        let completed: Vec<_> = self.pool.completed_results().await;
        let pending: Vec<_> = all
            .iter()
            .filter(|t| {
                matches!(
                    t.status,
                    super::CoordinatorTaskStatus::Pending | super::CoordinatorTaskStatus::Running
                )
            })
            .collect();

        let mut output = String::new();

        if !pending.is_empty() {
            output.push_str(&format!(
                "⚠ 还有 {} 个任务未完成: {}\n\n",
                pending.len(),
                pending
                    .iter()
                    .map(|t| t.id.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            ));
        }

        if completed.is_empty() {
            output.push_str("尚无已完成的任务结果可供合成。");
        } else {
            output.push_str(&format!("已完成 {} 个任务的结果合成:\n\n", completed.len()));
            for task in &completed {
                output.push_str(&format!(
                    "## {}\n{}\n\n---\n",
                    task.description.chars().take(120).collect::<String>(),
                    task.result.as_deref().unwrap_or("(无结果)")
                ));
            }
        }

        // 收集结构化结果
        let combined_results: Vec<serde_json::Value> = completed
            .iter()
            .map(|t| {
                json!({
                    "task_id": t.id,
                    "description": t.description,
                    "result": t.result,
                })
            })
            .collect();

        ToolOutput::success(
            format!(
                "[协调者合成结果]\n\n{} 个任务已完成。请基于以上结果向用户提供最终答案。",
                completed.len()
            ),
            json!({"synthesized": output, "completed_tasks": combined_results}),
        )
    }
}