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 { 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 { // 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 = 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 = 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, }) } }