From f57bd21ee44676f7d406442e9cd4181cf52941fe Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Fri, 14 Aug 2026 16:53:37 -0600 Subject: [PATCH] fix: Improved subagent escalation handling --- src/function/mod.rs | 157 ++++++++++++++++++++++++++++++------- src/function/supervisor.rs | 110 +++++++++++++++++++++++++- 2 files changed, 239 insertions(+), 28 deletions(-) diff --git a/src/function/mod.rs b/src/function/mod.rs index b7742fa..12e432d 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -186,42 +186,29 @@ pub async fn eval_tool_calls( }) .collect(); for (idx, call, result) in future::join_all(futs).await { - indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result?)))); + let value = match result { + Ok(v) => normalize_tool_result(v), + Err(e) => json!({"tool_call_error": format!("{e}")}), + }; + indexed_results.push((idx, ToolResult::new(call, value))); } } for (idx, call) in sequential_calls { - let result = call.eval(ctx).await?; - indexed_results.push((idx, ToolResult::new(call, normalize_tool_result(result)))); + let value = match call.eval(ctx).await { + Ok(v) => normalize_tool_result(v), + Err(e) => json!({ + "tool_call_error": format!( + "{e}. This tool is not available or the call failed; use only tools listed in your catalog." + ) + }), + }; + indexed_results.push((idx, ToolResult::new(call, value))); } indexed_results.sort_unstable_by_key(|(idx, _)| *idx); output = indexed_results.into_iter().map(|(_, r)| r).collect(); - if !output.is_empty() { - let (has_escalations, summary) = if ctx.current_depth == 0 - && let Some(queue) = ctx.root_escalation_queue() - && queue.has_pending() - { - (true, queue.pending_summary()) - } else { - (false, vec![]) - }; - - if has_escalations { - let notification = json!({ - "pending_escalations": summary, - "instruction": "Child agents are BLOCKED waiting for your reply. Call agent__reply_escalation for each pending escalation to unblock them." - }); - let synthetic_call = ToolCall::new( - "__escalation_notification".to_string(), - json!({}), - Some("escalation_check".to_string()), - ); - output.push(ToolResult::new(synthetic_call, notification)); - } - } - { let max_chars = ctx .agent @@ -236,6 +223,14 @@ pub async fn eval_tool_calls( } } + if ctx.current_depth == 0 + && let Some(queue) = ctx.root_escalation_queue() + && queue.has_pending() + && let Some(last) = output.last_mut() + { + inject_escalation_notification(last, queue.pending_summary()); + } + Ok(output) } @@ -252,6 +247,24 @@ fn normalize_tool_result(result: Value) -> Value { } } +fn inject_escalation_notification(last: &mut ToolResult, summary: Vec) { + let instruction = "Child agents are BLOCKED waiting for your reply. \ + Call agent__reply_escalation for each pending escalation to unblock them."; + match &mut last.output { + Value::Object(map) => { + map.insert("pending_escalations".into(), json!(summary)); + map.insert("escalation_instruction".into(), json!(instruction)); + } + other => { + *other = json!({ + "output": other.take(), + "pending_escalations": summary, + "escalation_instruction": instruction, + }); + } + } +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct ToolResult { pub call: ToolCall, @@ -1741,7 +1754,10 @@ fn format_json_colored_keys(value: &serde_json::Value) -> String { #[cfg(test)] mod tests { use super::*; + use crate::config::{AppState, WorkingMode}; + use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use serde_json::json; + use std::sync::Arc; fn call(name: &str, id: Option<&str>) -> ToolCall { ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string())) @@ -1751,11 +1767,98 @@ mod tests { ToolCall::new(name.to_string(), args, Some("id1".to_string())) } + fn run_async(f: F) -> F::Output { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(f) + } + + fn submit_escalation(queue: &EscalationQueue, id: &str) { + let (tx, _rx) = tokio::sync::oneshot::channel(); + queue.submit(EscalationRequest { + id: id.to_string(), + from_agent_id: "a1".into(), + from_agent_name: "explore".into(), + question: "What do?".into(), + options: None, + reply_tx: tx, + }); + } + #[test] fn normalize_tool_result_substitutes_done_for_null() { assert_eq!(normalize_tool_result(Value::Null), json!("DONE")); } + #[test] + fn inject_escalation_notification_extends_object_output() { + let mut result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"})); + inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_1"})]); + assert_eq!(result.output["status"], "ok"); + assert_eq!( + result.output["pending_escalations"], + json!([{"escalation_id": "esc_1"}]) + ); + assert!( + result.output["escalation_instruction"] + .as_str() + .unwrap() + .contains("agent__reply_escalation") + ); + assert!(result.text.is_none()); + } + + #[test] + fn inject_escalation_notification_wraps_non_object_output() { + let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE")); + inject_escalation_notification(&mut result, vec![json!({"escalation_id": "esc_2"})]); + assert_eq!(result.output["output"], json!("DONE")); + assert_eq!( + result.output["pending_escalations"], + json!([{"escalation_id": "esc_2"}]) + ); + assert!(result.output["escalation_instruction"].is_string()); + } + + #[test] + fn eval_tool_calls_soft_fails_unknown_tool() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![call("__escalation_notification", Some("id-1"))]; + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + assert_eq!(results.len(), 1); + let err = results[0].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call")); + assert!(err.contains("use only tools listed in your catalog")); + } + + #[test] + fn eval_tool_calls_injects_escalations_into_last_result() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let queue = ctx.ensure_root_escalation_queue(); + submit_escalation(&queue, "esc_1"); + + let calls = vec![call("unknown_tool", Some("id-1"))]; + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 1); + assert!( + results + .iter() + .all(|r| r.call.name != "__escalation_notification") + ); + let out = &results[0].output; + assert!(out["tool_call_error"].is_string()); + assert_eq!(out["pending_escalations"][0]["escalation_id"], "esc_1"); + assert!( + out["escalation_instruction"] + .as_str() + .unwrap() + .contains("agent__reply_escalation") + ); + } + #[test] fn normalize_tool_result_preserves_non_null_values() { assert_eq!( diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index d392821..05cdf65 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -85,7 +85,17 @@ pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailActi } ctx.pending_agents_guardrail_count += 1; - GuardrailAction::Inject(build_pending_agents_guardrail_prompt(&pending)) + let mut prompt = build_pending_agents_guardrail_prompt(&pending); + if let Some(queue) = ctx.root_escalation_queue() + && queue.has_pending() + { + let summary = serde_json::to_string(&queue.pending_summary()).unwrap_or_default(); + prompt.push_str(&format!( + "\n\nAdditionally, child agents have pending escalations blocking them. Reply to each \ + via `agent__reply_escalation` first:\n{summary}" + )); + } + GuardrailAction::Inject(prompt) } pub fn escalation_function_declarations() -> Vec { @@ -1996,4 +2006,102 @@ mod tests { let q2 = ctx.ensure_root_escalation_queue(); assert!(Arc::ptr_eq(&q1, &q2)); } + + #[test] + fn guardrail_prompt_mentions_pending_escalations() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let join_handle = tokio::spawn(async { + time::sleep(Duration::from_secs(60)).await; + Ok(AgentResult { + id: "slow".into(), + agent_name: "test".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: "slow".into(), + agent_name: "test".into(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let queue = ctx.ensure_root_escalation_queue(); + let (tx, _rx) = tokio::sync::oneshot::channel(); + queue.submit(EscalationRequest { + id: "esc_9".into(), + from_agent_id: "a1".into(), + from_agent_name: "explore".into(), + question: "Which option?".into(), + options: None, + reply_tx: tx, + }); + + match check_pending_agents_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("agent__reply_escalation")); + assert!(prompt.contains("esc_9")); + } + _ => panic!("expected Inject action"), + } + }); + } + + #[test] + fn guardrail_prompt_omits_escalations_when_none_pending() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let join_handle = tokio::spawn(async { + time::sleep(Duration::from_secs(60)).await; + Ok(AgentResult { + id: "slow".into(), + agent_name: "test".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: "slow".into(), + agent_name: "test".into(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + match check_pending_agents_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(!prompt.contains("agent__reply_escalation")); + } + _ => panic!("expected Inject action"), + } + }); + } }