PicoBot/src/tools/wait_tool.rs

use std::time::Duration;

use async_trait::async_trait;
use serde_json::json;

use crate::tools::{Tool, ToolContext, ToolResult, WaitEvent};

/// wait_for_subagents 工具 — 等待异步子代理完成或用户消息到达。
///
/// 调用后释放 serial_lock，进入 select! 等待：
/// - 子代理完成 → 返回结果 + 未完成列表
/// - 用户消息到达 → 返回 "有新用户消息"（消息已注入 history）
/// - 超时 → 返回超时 + 未完成列表
///
/// 等待结束后重新获取 serial_lock，保证后续工具调用串行。
pub struct WaitForSubagentsTool {
    /// 默认超时（秒），LLM 未指定时使用
    default_timeout_secs: u64,
}

impl WaitForSubagentsTool {
    pub const TOOL_NAME: &'static str = "wait_for_subagents";

    pub fn new(default_timeout_secs: u64) -> Self {
        Self {
            default_timeout_secs,
        }
    }
}

#[async_trait]
impl Tool for WaitForSubagentsTool {
    fn name(&self) -> &str {
        Self::TOOL_NAME
    }

    fn description(&self) -> &str {
        "Wait for asynchronous subagents to complete, or for new user messages to arrive. \
         Use this after launching subagents via the task tool to receive their results. \
         Returns the first completed subagent's result and a list of still-pending task IDs. \
         If pending_task_ids is non-empty, call this tool again to wait for the next one."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "timeout_secs": {
                    "type": "integer",
                    "description": "Maximum seconds to wait. Default 60. The tool returns early if a subagent completes or a user message arrives.",
                    "default": 60
                }
            },
            "required": []
        })
    }

    fn read_only(&self) -> bool {
        false
    }

    fn exclusive(&self) -> bool {
        // wait 工具释放/重获取 serial_lock，不应与其他工具并发
        true
    }

    async fn execute(&self, _args: serde_json::Value) -> anyhow::Result<ToolResult> {
        Ok(ToolResult {
            success: false,
            output: String::new(),
            error: Some(
                "wait_for_subagents requires tool context with wait_coordinator".to_string(),
            ),
        })
    }

    async fn execute_with_context(
        &self,
        context: &ToolContext,
        args: serde_json::Value,
    ) -> anyhow::Result<ToolResult> {
        // 1. 获取 wait_coordinator（仅主 agent 有值）
        let coordinator = match &context.wait_coordinator {
            Some(c) => c.clone(),
            None => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(
                        "wait_for_subagents is not available in this context (no wait_coordinator)"
                            .to_string(),
                    ),
                });
            }
        };

        // 2. 解析超时参数
        let timeout_secs = args
            .get("timeout_secs")
            .and_then(|v| v.as_u64())
            .unwrap_or(self.default_timeout_secs);
        let timeout = Duration::from_secs(timeout_secs);

        // 3. 先尝试排空队列中已缓冲的结果
        //    场景：子代理已完成 → 结果 send 到队列 + DB 更新为 completed，
        //    但 LLM 上轮未调用 wait（或 wait 超时未消费）→ 结果缓冲在队列中。
        //    若不排空，query_pending_task_ids 返回空（DB 已非 running）→
        //    返回 "No pending" → 缓冲结果永远丢失。
        let drained = coordinator.try_drain_queued_results().await;
        if !drained.is_empty() {
            let pending_after = coordinator.query_pending_task_ids();
            let pending_str = if pending_after.is_empty() {
                "none".to_string()
            } else {
                pending_after.join(", ")
            };
            let formatted: Vec<String> = drained
                .iter()
                .map(|r| {
                    format!(
                        "Subagent {} completed (status: {:?}). Output: {}",
                        r.task_id, r.status, r.output
                    )
                })
                .collect();
            return Ok(ToolResult {
                success: true,
                output: format!(
                    "Retrieved {} buffered subagent result(s):\n{}\nStill pending: [{}]",
                    drained.len(),
                    formatted.join("\n"),
                    pending_str
                ),
                error: None,
            });
        }

        // 4. 查询 pending 子代理
        let pending = coordinator.query_pending_task_ids();
        if pending.is_empty() {
            return Ok(ToolResult {
                success: true,
                output: "No pending subagents to wait for.".to_string(),
                error: None,
            });
        }

        tracing::info!(
            topic_id = ?context.topic_id,
            pending_count = pending.len(),
            timeout_secs,
            "wait_for_subagents: entering wait"
        );

        // 4. 进入等待（coordinator 内部：释放锁 → select! → 重获取锁）
        //    传入 cancel_rx 使 /stop 命令能立即中断等待。
        //    coordinator 在 select! 中以 biased 优先级处理：
        //    子代理结果 > 用户消息 > 取消信号 > 超时
        let event = coordinator.wait(timeout, context.cancel_rx.clone()).await;

        // 5. 格式化返回结果
        let output = match event {
            WaitEvent::SubagentResult(result) => {
                let pending_str = if result.pending_task_ids.is_empty() {
                    "none".to_string()
                } else {
                    result.pending_task_ids.join(", ")
                };
                format!(
                    "Subagent {} completed (status: {:?}). Output: {}\nStill pending: [{}]",
                    result.task_id, result.status, result.output, pending_str
                )
            }
            WaitEvent::UserMessage(messages) => {
                let pending = coordinator.query_pending_task_ids();
                let pending_str = if pending.is_empty() {
                    "none".to_string()
                } else {
                    pending.join(", ")
                };
                if messages.is_empty() {
                    format!(
                        "A new user message arrived while waiting (content could not be retrieved). \
                         Still pending subagents: [{}]",
                        pending_str
                    )
                } else {
                    let formatted_msgs: Vec<String> = messages
                        .iter()
                        .enumerate()
                        .map(|(i, msg)| format!("  [{}] {}", i + 1, msg))
                        .collect();
                    format!(
                        "New user message(s) arrived while waiting:\n{}\n\
                         These messages have been added to the conversation history. \
                         Still pending subagents: [{}]",
                        formatted_msgs.join("\n"),
                        pending_str
                    )
                }
            }
            WaitEvent::Timeout => {
                let pending = coordinator.query_pending_task_ids();
                format!(
                    "Wait timed out after {}s. Still pending subagents: [{}]",
                    timeout_secs,
                    if pending.is_empty() {
                        "none".to_string()
                    } else {
                        pending.join(", ")
                    }
                )
            }
            WaitEvent::Cancelled => {
                // /stop 命令中断了等待。coordinator 已完成全部状态清理
                //（重获取 serial_lock、回填 guard、清除 is_waiting、归还 receiver）。
                // 返回提示性输出，Agent 下一轮迭代会检测到 cancel 并退出。
                let pending = coordinator.query_pending_task_ids();
                tracing::info!(
                    topic_id = ?context.topic_id,
                    pending_count = pending.len(),
                    "wait_for_subagents: cancelled by /stop"
                );
                format!(
                    "Wait was cancelled by /stop command. \
                     Pending subagents (if any) have been cancelled separately. \
                     Still pending in DB: [{}]",
                    if pending.is_empty() {
                        "none".to_string()
                    } else {
                        pending.join(", ")
                    }
                )
            }
        };

        tracing::info!(
            topic_id = ?context.topic_id,
            "wait_for_subagents: returning result"
        );

        Ok(ToolResult {
            success: true,
            output,
            error: None,
        })
    }
}