From 1bccd2613fcb4de97362ca48e19b0ad0c946cd4d Mon Sep 17 00:00:00 2001
From: Lifei Zhou <lifei@squareup.com>
Date: Fri, 25 Jul 2025 12:30:22 +1000
Subject: [PATCH] fix: ensure execution task result is shown (#3629)

---
 crates/goose/src/agents/agent.rs                      |  7 ++++---
 .../subagent_execution_tool/task_execution_tracker.rs | 11 ++++++-----
 crates/goose/src/utils.rs                             |  8 ++++++++
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs
index 270fd580..c6369249 100644
--- a/crates/goose/src/agents/agent.rs
+++ b/crates/goose/src/agents/agent.rs
@@ -42,6 +42,7 @@ use crate::providers::errors::ProviderError;
 use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
 use crate::scheduler_trait::SchedulerTrait;
 use crate::tool_monitor::{ToolCall, ToolMonitor};
+use crate::utils::is_token_cancelled;
 use mcp_core::{protocol::GetPromptResult, ToolError, ToolResult};
 use regex::Regex;
 use rmcp::model::Tool;
@@ -742,7 +743,7 @@ impl Agent {
                 });
 
             loop {
-                if cancel_token.as_ref().is_some_and(|t| t.is_cancelled()) {
+                if is_token_cancelled(&cancel_token) {
                     break;
                 }
 
@@ -778,7 +779,7 @@ impl Agent {
                 let mut tools_updated = false;
 
                 while let Some(next) = stream.next().await {
-                    if cancel_token.as_ref().is_some_and(|t| t.is_cancelled()) {
+                    if is_token_cancelled(&cancel_token) {
                         break;
                     }
 
@@ -950,7 +951,7 @@ impl Agent {
                                     let mut all_install_successful = true;
 
                                     while let Some((request_id, item)) = combined.next().await {
-                                        if cancel_token.as_ref().is_some_and(|t| t.is_cancelled()) {
+                                        if is_token_cancelled(&cancel_token) {
                                             break;
                                         }
                                         match item {
diff --git a/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs b/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs
index 3610ea78..908bb465 100644
--- a/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs
+++ b/crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs
@@ -3,7 +3,7 @@ use rmcp::object;
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::sync::{mpsc, RwLock};
-use tokio::time::{Duration, Instant};
+use tokio::time::{sleep, Duration, Instant};
 use tokio_util::sync::CancellationToken;
 
 use crate::agents::subagent_execution_tool::notification_events::{
@@ -12,6 +12,7 @@ use crate::agents::subagent_execution_tool::notification_events::{
 };
 use crate::agents::subagent_execution_tool::task_types::{Task, TaskInfo, TaskResult, TaskStatus};
 use crate::agents::subagent_execution_tool::utils::{count_by_status, get_task_name};
+use crate::utils::is_token_cancelled;
 use serde_json::Value;
 use tokio::sync::mpsc::Sender;
 
@@ -22,6 +23,7 @@ pub enum DisplayMode {
 }
 
 const THROTTLE_INTERVAL_MS: u64 = 250;
+const COMPLETION_NOTIFICATION_DELAY_MS: u64 = 500;
 
 fn format_task_metadata(task_info: &TaskInfo) -> String {
     if let Some(params) = task_info.task.get_command_parameters() {
@@ -90,9 +92,7 @@ impl TaskExecutionTracker {
     }
 
     fn is_cancelled(&self) -> bool {
-        self.cancellation_token
-            .as_ref()
-            .is_some_and(|t| t.is_cancelled())
+        is_token_cancelled(&self.cancellation_token)
     }
 
     fn log_notification_error(
@@ -299,7 +299,8 @@ impl TaskExecutionTracker {
             .collect();
 
         let event = TaskExecutionNotificationEvent::tasks_complete(stats, failed_tasks);
-
         self.try_send_notification(event, "tasks complete");
+        // Wait for the notification to be recieved and displayed before clearing the tasks
+        sleep(Duration::from_millis(COMPLETION_NOTIFICATION_DELAY_MS)).await;
     }
 }
diff --git a/crates/goose/src/utils.rs b/crates/goose/src/utils.rs
index 60121f1b..1205953e 100644
--- a/crates/goose/src/utils.rs
+++ b/crates/goose/src/utils.rs
@@ -1,3 +1,5 @@
+use tokio_util::sync::CancellationToken;
+
 /// Safely truncate a string at character boundaries, not byte boundaries
 ///
 /// This function ensures that multi-byte UTF-8 characters (like Japanese, emoji, etc.)
@@ -18,6 +20,12 @@ pub fn safe_truncate(s: &str, max_chars: usize) -> String {
     }
 }
 
+pub fn is_token_cancelled(cancellation_token: &Option<CancellationToken>) -> bool {
+    cancellation_token
+        .as_ref()
+        .is_some_and(|t| t.is_cancelled())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;