diff --git a/crates/goose/src/agents/platform_extensions/summon.rs b/crates/goose/src/agents/platform_extensions/summon.rs index c0836493c..a8f78d0fa 100644 --- a/crates/goose/src/agents/platform_extensions/summon.rs +++ b/crates/goose/src/agents/platform_extensions/summon.rs @@ -2,7 +2,7 @@ use crate::agents::extension::PlatformExtensionContext; use crate::agents::mcp_client::{Error, McpClientTrait}; use crate::agents::subagent_handler::{run_subagent_task, OnMessageCallback, SubagentRunParams}; use crate::agents::subagent_task_config::{TaskConfig, DEFAULT_SUBAGENT_MAX_TURNS}; -use crate::agents::tool_execution::ToolCallContext; +use crate::agents::tool_execution::{ToolCallContext, ToolCallNotificationEmitter}; use crate::agents::AgentConfig; use crate::config::paths::Paths; use crate::config::{Config, GooseMode}; @@ -23,11 +23,12 @@ use rmcp::model::{ }; use serde::Deserialize; use std::collections::HashMap; +use std::future::Future; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::Mutex; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -72,7 +73,7 @@ pub struct BackgroundTask { pub last_activity: Arc, pub handle: JoinHandle>, pub cancellation_token: CancellationToken, - pub notification_buffer: Arc>>, + notification_sink: SharedNotificationSink, } pub struct CompletedTask { @@ -82,6 +83,59 @@ pub struct CompletedTask { pub turns_taken: u32, pub duration: Duration, pub completed_at: Instant, + notification_sink: SharedNotificationSink, +} + +enum NotificationSink { + Buffer(Vec), + Emitter(ToolCallNotificationEmitter), +} + +type SharedNotificationSink = Arc>; + +async fn yield_to_outer_tool_stream() { + // The outer select may have polled its receiver before this future queues a + // notification. Keep the result pending for the following select pass so + // the now-ready receiver is observed before the terminal result. + tokio::task::yield_now().await; + tokio::task::yield_now().await; +} + +impl NotificationSink { + fn route(&mut self, notification: ServerNotification) { + match self { + Self::Buffer(buffer) => buffer.push(notification), + Self::Emitter(emitter) => emitter.emit_best_effort(notification), + } + } + + async fn attach(&mut self, emitter: Option) { + let Some(emitter) = emitter else { + return; + }; + while let Self::Buffer(buffered) = self { + let Some(notification) = buffered.first().cloned() else { + break; + }; + emitter.emit_best_effort(notification); + yield_to_outer_tool_stream().await; + buffered.remove(0); + } + *self = Self::Emitter(emitter); + } + + fn detach(&mut self) { + if matches!(self, Self::Emitter(_)) { + *self = Self::Buffer(Vec::new()); + } + } + + fn buffered_len(&self) -> usize { + match self { + Self::Buffer(buffer) => buffer.len(), + Self::Emitter(_) => 0, + } + } } fn merge_subrecipe_parameters( @@ -471,7 +525,6 @@ pub struct SummonClient { source_cache: Mutex)>>, background_tasks: Mutex>, completed_tasks: Mutex>, - notification_subscribers: Arc>>>, } impl Drop for SummonClient { @@ -496,7 +549,6 @@ impl SummonClient { source_cache: Mutex::new(None), background_tasks: Mutex::new(HashMap::new()), completed_tasks: Mutex::new(HashMap::new()), - notification_subscribers: Arc::new(Mutex::new(Vec::new())), }) } @@ -530,26 +582,49 @@ impl SummonClient { Ok(session) } - fn spawn_notification_bridge( - mut notif_rx: tokio::sync::mpsc::UnboundedReceiver, - subscribers: Arc>>>, - buffer: Arc>>, + fn notification_sink(emitter: Option) -> SharedNotificationSink { + Arc::new(Mutex::new(match emitter { + Some(emitter) => NotificationSink::Emitter(emitter), + None => NotificationSink::Buffer(Vec::new()), + })) + } + + async fn attach_notification_emitter( + sink: &SharedNotificationSink, + emitter: Option, ) { - tokio::spawn(async move { - while let Some(notification) = notif_rx.recv().await { - let mut subs = subscribers.lock().await; - if subs.is_empty() { - drop(subs); - buffer.lock().await.push(notification); - } else { - subs.retain(|tx| match tx.try_send(notification.clone()) { - Ok(()) => true, - Err(mpsc::error::TrySendError::Full(_)) => true, - Err(mpsc::error::TrySendError::Closed(_)) => false, - }); + sink.lock().await.attach(emitter).await; + } + + async fn run_subagent_with_notifications( + sink: SharedNotificationSink, + run_subagent: Run, + ) -> Result + where + Run: FnOnce(tokio::sync::mpsc::UnboundedSender) -> RunFuture, + RunFuture: Future>, + { + let (notification_tx, mut notification_rx) = tokio::sync::mpsc::unbounded_channel(); + let run = run_subagent(notification_tx); + tokio::pin!(run); + + loop { + tokio::select! { + biased; + result = &mut run => { + while let Ok(notification) = notification_rx.try_recv() { + sink.lock().await.route(notification); + yield_to_outer_tool_stream().await; + } + yield_to_outer_tool_stream().await; + return result; + } + Some(notification) = notification_rx.recv() => { + sink.lock().await.route(notification); + yield_to_outer_tool_stream().await; } } - }); + } } fn create_load_tool(&self) -> Tool { @@ -849,6 +924,7 @@ impl SummonClient { &self, session_id: &str, arguments: Option, + notification_emitter: Option, ) -> Result { self.cleanup_completed_tasks().await; @@ -881,7 +957,9 @@ impl SummonClient { let name = source_name.unwrap(); if is_session_id(name) { - let task_result = self.handle_load_task_result(name, cancel, peek).await?; + let task_result = self + .handle_load_task_result(name, cancel, peek, notification_emitter) + .await?; let mut meta = MetaObject::new(); meta.0.insert( "subagent_session_id".to_string(), @@ -916,30 +994,27 @@ impl SummonClient { task_id: &str, cancel: bool, peek: bool, + notification_emitter: Option, ) -> Result { let mut completed = self.completed_tasks.lock().await; - let completed_entry = if peek { - completed.get(task_id).map(|task| { - ( - task.result.clone(), - task.description.clone(), - task.duration, - task.turns_taken, - ) - }) - } else { - completed.remove(task_id).map(|task| { - ( - task.result, - task.description, - task.duration, - task.turns_taken, - ) - }) - }; + let completed_entry = completed.get(task_id).map(|task| { + ( + task.result.clone(), + task.description.clone(), + task.duration, + task.turns_taken, + Arc::clone(&task.notification_sink), + ) + }); - if let Some((result, description, duration, turns_taken)) = completed_entry { + if let Some((result, description, duration, turns_taken, notification_sink)) = + completed_entry + { + if !peek { + Self::attach_notification_emitter(¬ification_sink, notification_emitter).await; + completed.remove(task_id); + } let status_key = match &result { Ok(_) => "completed", Err(e) if e.starts_with("Task panicked:") => "panicked", @@ -986,7 +1061,7 @@ impl SummonClient { let idle_ms = now.saturating_sub(task.last_activity.load(Ordering::Relaxed)); let description = task.description.clone(); - let buffered_count = task.notification_buffer.lock().await.len(); + let buffered_count = task.notification_sink.lock().await.buffered_len(); drop(running); @@ -1013,9 +1088,11 @@ impl SummonClient { } if cancel { + let notification_sink = + Arc::clone(&running.get(task_id).unwrap().notification_sink); + Self::attach_notification_emitter(¬ification_sink, notification_emitter).await; let task = running.remove(task_id).unwrap(); drop(running); - task.cancellation_token.cancel(); let duration = task.started_at.elapsed(); @@ -1057,22 +1134,11 @@ impl SummonClient { // Wait for the running task to complete, keeping the tool call // alive so notifications (subagent tool calls) stream in real time. + let notification_sink = Arc::clone(&running.get(task_id).unwrap().notification_sink); + Self::attach_notification_emitter(¬ification_sink, notification_emitter).await; let mut task = running.remove(task_id).unwrap(); drop(running); - let buffered = { - let mut buf = task.notification_buffer.lock().await; - std::mem::take(&mut *buf) - }; - if !buffered.is_empty() { - let subs = self.notification_subscribers.lock().await; - for notif in buffered { - for tx in subs.iter() { - let _ = tx.try_send(notif.clone()); - } - } - } - tokio::select! { result = &mut task.handle => { let (output, status_key) = match result { @@ -1108,6 +1174,7 @@ impl SummonClient { }); } _ = tokio::time::sleep(Duration::from_secs(300)) => { + task.notification_sink.lock().await.detach(); self.background_tasks.lock().await.insert(task_id.to_string(), task); return Err(format!( @@ -1245,6 +1312,7 @@ impl SummonClient { session_id: &str, arguments: Option, cancellation_token: CancellationToken, + notification_emitter: Option, ) -> Result { self.cleanup_completed_tasks().await; @@ -1305,16 +1373,9 @@ impl SummonClient { .create_subagent_session(&task_config, "Delegated task".to_string()) .await?; - let (notif_tx, notif_rx) = tokio::sync::mpsc::unbounded_channel::(); - Self::spawn_notification_bridge( - notif_rx, - Arc::clone(&self.notification_subscribers), - Arc::new(Mutex::new(Vec::new())), - ); - let subagent_session_id = subagent_session.id.clone(); - let result = run_subagent_task(SubagentRunParams { + let params = SubagentRunParams { config: agent_config, recipe, task_config, @@ -1322,8 +1383,16 @@ impl SummonClient { session_id: subagent_session.id, cancellation_token: Some(cancellation_token), on_message: None, - notification_tx: Some(notif_tx), - }) + notification_tx: None, + }; + let result = Self::run_subagent_with_notifications( + Self::notification_sink(notification_emitter), + move |notification_tx| { + let mut params = params; + params.notification_tx = Some(notification_tx); + run_subagent_task(params) + }, + ) .await; let mut meta = MetaObject::new(); @@ -1794,6 +1863,7 @@ impl SummonClient { turns_taken, duration, completed_at: Instant::now(), + notification_sink: task.notification_sink, }, ); } @@ -1878,17 +1948,11 @@ impl SummonClient { let task_token = CancellationToken::new(); let task_token_clone = task_token.clone(); - let notification_buffer = Arc::new(Mutex::new(Vec::new())); - - let (notif_tx, notif_rx) = tokio::sync::mpsc::unbounded_channel::(); - Self::spawn_notification_bridge( - notif_rx, - Arc::clone(&self.notification_subscribers), - Arc::clone(¬ification_buffer), - ); + let notification_sink = Self::notification_sink(None); + let task_notification_sink = Arc::clone(¬ification_sink); let handle = tokio::spawn(async move { - run_subagent_task(SubagentRunParams { + let params = SubagentRunParams { config: agent_config, recipe, task_config, @@ -1896,7 +1960,12 @@ impl SummonClient { session_id: subagent_session.id, cancellation_token: Some(task_token_clone), on_message: Some(on_message), - notification_tx: Some(notif_tx), + notification_tx: None, + }; + Self::run_subagent_with_notifications(task_notification_sink, move |notification_tx| { + let mut params = params; + params.notification_tx = Some(notification_tx); + run_subagent_task(params) }) .await }); @@ -1909,7 +1978,7 @@ impl SummonClient { last_activity, handle, cancellation_token: task_token, - notification_buffer, + notification_sink, }; self.background_tasks @@ -1967,7 +2036,10 @@ impl McpClientTrait for SummonClient { ) -> Result { let session_id = &ctx.session_id; match name { - "load" => match self.handle_load(session_id, arguments).await { + "load" => match self + .handle_load(session_id, arguments, ctx.notification_emitter().cloned()) + .await + { Ok(result) => Ok(result), Err(error) => Ok(CallToolResult::error(vec![ContentBlock::text(format!( "Error: {}", @@ -1976,7 +2048,12 @@ impl McpClientTrait for SummonClient { }, "delegate" => { match self - .handle_delegate(session_id, arguments, cancellation_token) + .handle_delegate( + session_id, + arguments, + cancellation_token, + ctx.notification_emitter().cloned(), + ) .await { Ok(result) => Ok(result), @@ -2006,12 +2083,6 @@ impl McpClientTrait for SummonClient { } } - async fn subscribe(&self) -> mpsc::Receiver { - let (tx, rx) = mpsc::channel(16); - self.notification_subscribers.lock().await.push(tx); - rx - } - async fn get_moim(&self, _session_id: &str) -> Option { self.cleanup_completed_tasks().await; @@ -2104,6 +2175,7 @@ fn resolve_working_dir(parent_dir: &Path, requested: &str) -> Result ServerNotification { + use crate::agents::subagent_handler::create_tool_notification; + use crate::conversation::message::MessageContent; + use rmcp::model::CallToolRequestParams; + + let tool_call = CallToolRequestParams::new("developer__shell").with_arguments( + serde_json::json!({"command": request_id}) + .as_object() + .unwrap() + .clone(), + ); + let content = MessageContent::tool_request(request_id, Ok(tool_call)); + create_tool_notification(&content, subagent_id).unwrap() + } + + fn notification_subagent_id(notification: &ServerNotification) -> Option { + let ServerNotification::LoggingMessageNotification(log) = notification else { + return None; + }; + serde_json::to_value(&log.params) + .ok()? + .get("data")? + .get("subagent_id")? + .as_str() + .map(str::to_string) + } + + fn notification_command(notification: &ServerNotification) -> Option { + let ServerNotification::LoggingMessageNotification(log) = notification else { + return None; + }; + serde_json::to_value(&log.params) + .ok()? + .get("data")? + .get("tool_call")? + .get("arguments")? + .get("command")? + .as_str() + .map(str::to_string) + } + + fn notification_channel() -> ( + ToolCallNotificationEmitter, + tokio::sync::mpsc::Receiver, + ) { + let (sender, receiver) = tokio::sync::mpsc::channel(32); + (ToolCallNotificationEmitter::new(sender), receiver) + } + + fn buffered_notification_sink( + notifications: Vec, + ) -> SharedNotificationSink { + Arc::new(Mutex::new(NotificationSink::Buffer(notifications))) + } + #[test] fn test_is_session_id() { assert!(is_session_id("20260204_1")); @@ -2896,32 +3023,272 @@ You review code."#; assert!(!is_session_id("20260204")); } + #[tokio::test] + async fn test_notification_sinks_isolate_concurrent_delegate_calls() { + let (emitter_a, mut notifications_a) = notification_channel(); + let (emitter_b, mut notifications_b) = notification_channel(); + let sink_a = SummonClient::notification_sink(Some(emitter_a)); + let sink_b = SummonClient::notification_sink(Some(emitter_b)); + + let (result_a, result_b) = tokio::join!( + SummonClient::run_subagent_with_notifications(sink_a, |notification_tx| async move { + notification_tx + .send(test_tool_notification("inner-a", "subagent-a")) + .unwrap(); + tokio::task::yield_now().await; + Ok("delegate-a".to_string()) + }), + SummonClient::run_subagent_with_notifications(sink_b, |notification_tx| async move { + notification_tx + .send(test_tool_notification("inner-b", "subagent-b")) + .unwrap(); + tokio::task::yield_now().await; + Ok("delegate-b".to_string()) + }) + ); + assert_eq!(result_a.unwrap(), "delegate-a"); + assert_eq!(result_b.unwrap(), "delegate-b"); + + let notification_a = notifications_a.recv().await.unwrap(); + let notification_b = notifications_b.recv().await.unwrap(); + assert_eq!( + notification_subagent_id(¬ification_a).as_deref(), + Some("subagent-a") + ); + assert_eq!( + notification_subagent_id(¬ification_b).as_deref(), + Some("subagent-b") + ); + assert!(notifications_a.try_recv().is_err()); + assert!(notifications_b.try_recv().is_err()); + } + + #[tokio::test] + async fn test_live_notifications_precede_delegate_result() { + use crate::agents::tool_execution::{tool_stream, ToolStreamItem}; + use tokio_stream::wrappers::ReceiverStream; + + for _ in 0..32 { + let (emitter, notifications) = notification_channel(); + let sink = SummonClient::notification_sink(Some(emitter)); + let mut output = tool_stream( + ReceiverStream::new(notifications), + futures::stream::empty(), + async move { + let result = SummonClient::run_subagent_with_notifications( + sink, + |notification_tx| async move { + for command in ["inner-live-0", "inner-live-1", "inner-live-2"] { + notification_tx + .send(test_tool_notification(command, "subagent-live")) + .unwrap(); + } + Ok("delegate-result".to_string()) + }, + ) + .await + .unwrap(); + Ok::<_, rmcp::model::ErrorData>(CallToolResult::success(vec![ + ContentBlock::text(result), + ])) + }, + ); + + let mut commands = Vec::new(); + let result = loop { + match output.next().await.unwrap() { + ToolStreamItem::Message(notification) => { + assert_eq!( + notification_subagent_id(¬ification).as_deref(), + Some("subagent-live") + ); + commands.push(notification_command(¬ification).unwrap()); + } + ToolStreamItem::Result(result) => break result, + ToolStreamItem::ActionRequired(_) => { + panic!("delegate must not request an action") + } + } + }; + + assert_eq!(commands, ["inner-live-0", "inner-live-1", "inner-live-2"]); + assert!(result.is_ok()); + assert!(output.next().await.is_none()); + } + } + + #[tokio::test] + async fn test_async_completion_before_load_replays_notifications() { + use crate::agents::tool_execution::{tool_stream, ToolStreamItem}; + use tokio_stream::wrappers::ReceiverStream; + + let client = Arc::new(SummonClient::new(create_test_context()).unwrap()); + let task_id = "20260204_1"; + let buffered = vec![test_tool_notification("inner-completed", task_id)]; + client.completed_tasks.lock().await.insert( + task_id.to_string(), + CompletedTask { + id: task_id.to_string(), + description: "Completed task".to_string(), + result: Ok("done".to_string()), + turns_taken: 1, + duration: Duration::from_secs(1), + completed_at: Instant::now(), + notification_sink: buffered_notification_sink(buffered), + }, + ); + let (emitter, notifications) = notification_channel(); + let load_client = Arc::clone(&client); + let mut output = tool_stream( + ReceiverStream::new(notifications), + futures::stream::empty(), + async move { + let result = load_client + .handle_load_task_result(task_id, false, false, Some(emitter)) + .await + .unwrap(); + Ok::<_, rmcp::model::ErrorData>(CallToolResult::success(result.content)) + }, + ); + + let ToolStreamItem::Message(notification) = output.next().await.unwrap() else { + panic!("buffered notification must be emitted before the load result"); + }; + assert_eq!( + notification_subagent_id(¬ification).as_deref(), + Some(task_id) + ); + assert_eq!( + notification_command(¬ification).as_deref(), + Some("inner-completed") + ); + let ToolStreamItem::Result(result) = output.next().await.unwrap() else { + panic!("load result must follow buffered notifications"); + }; + assert!(result.is_ok()); + assert!(output.next().await.is_none()); + assert!(!client.completed_tasks.lock().await.contains_key(task_id)); + } + + #[tokio::test] + async fn test_cancelled_completed_load_remains_retrievable() { + let client = Arc::new(SummonClient::new(create_test_context()).unwrap()); + let task_id = "20260204_1"; + client.completed_tasks.lock().await.insert( + task_id.to_string(), + CompletedTask { + id: task_id.to_string(), + description: "Completed task".to_string(), + result: Ok("done".to_string()), + turns_taken: 1, + duration: Duration::from_secs(1), + completed_at: Instant::now(), + notification_sink: buffered_notification_sink(vec![ + test_tool_notification("inner-0", task_id), + test_tool_notification("inner-1", task_id), + ]), + }, + ); + let (emitter, mut notifications) = notification_channel(); + let load_client = Arc::clone(&client); + let load = tokio::spawn(async move { + load_client + .handle_load_task_result(task_id, false, false, Some(emitter)) + .await + }); + + let first = notifications.recv().await.unwrap(); + assert_eq!(notification_command(&first).as_deref(), Some("inner-0")); + load.abort(); + assert!(load.await.unwrap_err().is_cancelled()); + assert!(client.completed_tasks.lock().await.contains_key(task_id)); + + let (retry_emitter, mut retry_notifications) = notification_channel(); + let result = client + .handle_load_task_result(task_id, false, false, Some(retry_emitter)) + .await + .unwrap(); + + assert_eq!(result.status, "completed"); + for command in ["inner-0", "inner-1"] { + let notification = retry_notifications.try_recv().unwrap(); + assert_eq!( + notification_command(¬ification).as_deref(), + Some(command) + ); + } + assert!(retry_notifications.try_recv().is_err()); + assert!(!client.completed_tasks.lock().await.contains_key(task_id)); + } + + #[tokio::test] + async fn test_buffered_replay_preserves_order_and_emitter_capacity() { + let sink = buffered_notification_sink( + (0..33) + .map(|index| test_tool_notification(&format!("inner-{index}"), "subagent")) + .collect(), + ); + let (emitter, mut notifications) = notification_channel(); + + SummonClient::attach_notification_emitter(&sink, Some(emitter)).await; + + for index in 0..32 { + let notification = notifications.try_recv().unwrap(); + assert_eq!( + notification_command(¬ification), + Some(format!("inner-{index}")) + ); + } + assert!(notifications.try_recv().is_err()); + } + + #[tokio::test] + async fn test_load_completes_when_caller_does_not_consume_notifications() { + let client = SummonClient::new(create_test_context()).unwrap(); + let task_id = "20260204_1"; + client.completed_tasks.lock().await.insert( + task_id.to_string(), + CompletedTask { + id: task_id.to_string(), + description: "Completed task".to_string(), + result: Ok("done".to_string()), + turns_taken: 1, + duration: Duration::from_secs(1), + completed_at: Instant::now(), + notification_sink: buffered_notification_sink( + (0..64) + .map(|index| test_tool_notification(&format!("inner-{index}"), task_id)) + .collect(), + ), + }, + ); + let (emitter, _notifications) = notification_channel(); + + let result = tokio::time::timeout( + Duration::from_secs(1), + client.handle_load_task_result(task_id, false, false, Some(emitter)), + ) + .await + .expect("load must not wait for a notification consumer") + .unwrap(); + + assert_eq!(result.status, "completed"); + } + #[tokio::test] async fn test_async_task_result_lifecycle() { let client = SummonClient::new(create_test_context()).unwrap(); let temp_dir = TempDir::new().unwrap(); let result = client - .handle_load_task_result("20260204_999", false, false) + .handle_load_task_result("20260204_999", false, false, None) .await; assert!(result.is_err()); assert!(result.unwrap_err().contains("not found")); { - use crate::agents::subagent_handler::create_tool_notification; - use crate::conversation::message::MessageContent; - use rmcp::model::CallToolRequestParams; - - let tool_call = CallToolRequestParams::new("developer__shell").with_arguments( - serde_json::json!({"command": "ls"}) - .as_object() - .unwrap() - .clone(), - ); - let content = MessageContent::tool_request("req1", Ok(tool_call)); - let notif = create_tool_notification(&content, "20260204_1").unwrap(); - - let buffer = Arc::new(Mutex::new(vec![notif])); + let notification_sink = + buffered_notification_sink(vec![test_tool_notification("req1", "20260204_1")]); let mut running = client.background_tasks.lock().await; running.insert( @@ -2937,24 +3304,22 @@ You review code."#; Ok("done".to_string()) }), cancellation_token: CancellationToken::new(), - notification_buffer: buffer, + notification_sink, }, ); } - let mut subscriber = client.subscribe().await; - - let result = client - .handle_load_task_result("20260204_1", false, false) - .await - .expect("load should wait and return result"); + let (emitter, mut notifications) = notification_channel(); + let (result, notification) = tokio::join!( + client.handle_load_task_result("20260204_1", false, false, Some(emitter)), + notifications.recv() + ); + let result = result.expect("load should wait and return result"); let text = extract_text(&result.content[0]); assert!(text.contains("Completed")); assert!(text.contains("done")); - let notif = subscriber - .try_recv() - .expect("subscriber should receive buffered notification"); + let notif = notification.expect("load emitter should receive buffered notification"); if let ServerNotification::LoggingMessageNotification(log) = notif { let params = serde_json::to_value(&log.params).unwrap(); let data = params.get("data").and_then(|v| v.as_object()).unwrap(); @@ -2977,6 +3342,7 @@ You review code."#; turns_taken: 5, duration: Duration::from_secs(60), completed_at: Instant::now(), + notification_sink: buffered_notification_sink(Vec::new()), }, ); completed.insert( @@ -2988,6 +3354,7 @@ You review code."#; turns_taken: 3, duration: Duration::from_secs(30), completed_at: Instant::now(), + notification_sink: buffered_notification_sink(Vec::new()), }, ); } @@ -3008,7 +3375,7 @@ You review code."#; assert!(discovery_text.contains("20260204_3")); let result = client - .handle_load_task_result("20260204_2", false, false) + .handle_load_task_result("20260204_2", false, false, None) .await .unwrap(); let text = extract_text(&result.content[0]); @@ -3028,7 +3395,7 @@ You review code."#; .contains_key("20260204_2")); let result = client - .handle_load_task_result("20260204_3", false, false) + .handle_load_task_result("20260204_3", false, false, None) .await .unwrap(); let text = extract_text(&result.content[0]); @@ -3037,7 +3404,7 @@ You review code."#; assert_eq!(result.status, "failed"); let result = client - .handle_load_task_result("20260204_3", false, false) + .handle_load_task_result("20260204_3", false, false, None) .await; assert!(result.is_err()); assert!(result.unwrap_err().contains("not found")); @@ -3050,43 +3417,174 @@ You review code."#; async fn test_cancel_running_task() { let client = SummonClient::new(create_test_context()).unwrap(); let token = CancellationToken::new(); + let task_id = "20260204_1"; + let notification_sink = buffered_notification_sink(Vec::new()); + let task_notification_sink = Arc::clone(¬ification_sink); + let task_token = token.clone(); { let mut running = client.background_tasks.lock().await; running.insert( - "20260204_1".to_string(), + task_id.to_string(), BackgroundTask { - id: "20260204_1".to_string(), + id: task_id.to_string(), description: "Cancellable task".to_string(), started_at: Instant::now(), turns: Arc::new(AtomicU32::new(3)), last_activity: Arc::new(AtomicU64::new(current_epoch_millis())), - handle: tokio::spawn(async { - tokio::time::sleep(Duration::from_secs(1000)).await; - Ok("should not see this".to_string()) + handle: tokio::spawn(async move { + task_token.cancelled().await; + task_notification_sink + .lock() + .await + .route(test_tool_notification("cancel", task_id)); + Ok("cancelled gracefully".to_string()) }), cancellation_token: token.clone(), - notification_buffer: Arc::new(Mutex::new(Vec::new())), + notification_sink, }, ); } - let result = client - .handle_load_task_result("20260204_1", true, false) - .await - .unwrap(); + let (emitter, mut notifications) = notification_channel(); + let (result, notification) = tokio::join!( + client.handle_load_task_result(task_id, true, false, Some(emitter)), + notifications.recv() + ); + let result = result.unwrap(); let text = extract_text(&result.content[0]); assert!(text.contains("Cancelled")); - assert!(text.contains("20260204_1")); + assert!(text.contains(task_id)); assert!(text.contains("Cancellable task")); + assert!(text.contains("cancelled gracefully")); assert_eq!(result.status, "cancelled"); assert_eq!(result.turns, Some(3)); + assert_eq!( + notification_subagent_id(¬ification.unwrap()).as_deref(), + Some(task_id) + ); assert!(token.is_cancelled()); - assert!(!client - .background_tasks - .lock() + assert!(!client.background_tasks.lock().await.contains_key(task_id)); + } + + #[tokio::test] + async fn test_cancelled_running_load_remains_retrievable() { + let client = Arc::new(SummonClient::new(create_test_context()).unwrap()); + let token = CancellationToken::new(); + let task_id = "20260204_1"; + let task_token = token.clone(); + + client.background_tasks.lock().await.insert( + task_id.to_string(), + BackgroundTask { + id: task_id.to_string(), + description: "Cancellable task".to_string(), + started_at: Instant::now(), + turns: Arc::new(AtomicU32::new(1)), + last_activity: Arc::new(AtomicU64::new(current_epoch_millis())), + handle: tokio::spawn(async move { + task_token.cancelled().await; + Ok("cancelled gracefully".to_string()) + }), + cancellation_token: token.clone(), + notification_sink: buffered_notification_sink(vec![ + test_tool_notification("inner-0", task_id), + test_tool_notification("inner-1", task_id), + ]), + }, + ); + + let (emitter, mut notifications) = notification_channel(); + let load_client = Arc::clone(&client); + let load = tokio::spawn(async move { + load_client + .handle_load_task_result(task_id, true, false, Some(emitter)) + .await + }); + + let first = notifications.recv().await.unwrap(); + assert_eq!(notification_command(&first).as_deref(), Some("inner-0")); + load.abort(); + assert!(load.await.unwrap_err().is_cancelled()); + assert!(client.background_tasks.lock().await.contains_key(task_id)); + assert!(!token.is_cancelled()); + + let (retry_emitter, mut retry_notifications) = notification_channel(); + let result = client + .handle_load_task_result(task_id, true, false, Some(retry_emitter)) .await - .contains_key("20260204_1")); + .unwrap(); + + assert_eq!(result.status, "cancelled"); + assert!(token.is_cancelled()); + assert!(!client.background_tasks.lock().await.contains_key(task_id)); + + let commands: Vec = std::iter::from_fn(|| retry_notifications.try_recv().ok()) + .filter_map(|notification| notification_command(¬ification)) + .collect(); + assert!( + commands == ["inner-0", "inner-1"] || commands == ["inner-1"], + "retry must replay the remaining notifications, with at-least-once delivery allowed" + ); + } + + #[tokio::test] + async fn test_cancelled_waiting_load_remains_retrievable() { + let client = Arc::new(SummonClient::new(create_test_context()).unwrap()); + let task_id = "20260204_1"; + let (finish_tx, finish_rx) = tokio::sync::oneshot::channel(); + + client.background_tasks.lock().await.insert( + task_id.to_string(), + BackgroundTask { + id: task_id.to_string(), + description: "Running task".to_string(), + started_at: Instant::now(), + turns: Arc::new(AtomicU32::new(1)), + last_activity: Arc::new(AtomicU64::new(current_epoch_millis())), + handle: tokio::spawn(async move { + finish_rx.await.unwrap(); + Ok("done".to_string()) + }), + cancellation_token: CancellationToken::new(), + notification_sink: buffered_notification_sink(vec![ + test_tool_notification("inner-0", task_id), + test_tool_notification("inner-1", task_id), + ]), + }, + ); + + let (emitter, mut notifications) = notification_channel(); + let load_client = Arc::clone(&client); + let load = tokio::spawn(async move { + load_client + .handle_load_task_result(task_id, false, false, Some(emitter)) + .await + }); + + let first = notifications.recv().await.unwrap(); + assert_eq!(notification_command(&first).as_deref(), Some("inner-0")); + load.abort(); + assert!(load.await.unwrap_err().is_cancelled()); + assert!(client.background_tasks.lock().await.contains_key(task_id)); + + finish_tx.send(()).unwrap(); + let (retry_emitter, mut retry_notifications) = notification_channel(); + let result = client + .handle_load_task_result(task_id, false, false, Some(retry_emitter)) + .await + .unwrap(); + + assert_eq!(result.status, "completed"); + assert!(!client.background_tasks.lock().await.contains_key(task_id)); + + let commands: Vec = std::iter::from_fn(|| retry_notifications.try_recv().ok()) + .filter_map(|notification| notification_command(¬ification)) + .collect(); + assert!( + commands == ["inner-0", "inner-1"] || commands == ["inner-1"], + "retry must replay the remaining notifications, with at-least-once delivery allowed" + ); } #[tokio::test] @@ -3108,14 +3606,14 @@ You review code."#; Ok("eventual result".to_string()) }), cancellation_token: CancellationToken::new(), - notification_buffer: Arc::new(Mutex::new(Vec::new())), + notification_sink: buffered_notification_sink(Vec::new()), }, ); } // Peek should return status without removing the task let result = client - .handle_load_task_result("20260204_1", false, true) + .handle_load_task_result("20260204_1", false, true, None) .await .unwrap(); let text = extract_text(&result.content[0]); @@ -3136,7 +3634,7 @@ You review code."#; let client = SummonClient::new(create_test_context()).unwrap(); let result = client - .handle_load_task_result("20260204_999", false, true) + .handle_load_task_result("20260204_999", false, true, None) .await; assert!(result.is_err()); assert!(result.unwrap_err().contains("not found")); @@ -3157,13 +3655,14 @@ You review code."#; turns_taken: 4, duration: Duration::from_secs(30), completed_at: Instant::now(), + notification_sink: buffered_notification_sink(Vec::new()), }, ); } // Peek on a completed task should return the full result (same as non-peek) let result = client - .handle_load_task_result("20260204_1", false, true) + .handle_load_task_result("20260204_1", false, true, None) .await .unwrap(); let text = extract_text(&result.content[0]); @@ -3177,7 +3676,7 @@ You review code."#; .await .contains_key("20260204_1")); let result = client - .handle_load_task_result("20260204_1", false, false) + .handle_load_task_result("20260204_1", false, false, None) .await .unwrap(); assert!(extract_text(&result.content[0]).contains("final output"));