From b988dca543cf41c16833fbf0d928f4c0fb1f2ffd Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Thu, 24 Jul 2025 23:01:05 +1000 Subject: [PATCH 1/2] added extra time --- .../agents/subagent_execution_tool/task_execution_tracker.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 7d6854b3e229..eb0551f4e2d6 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::{ @@ -22,6 +22,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() { @@ -309,7 +310,7 @@ impl TaskExecutionTracker { .collect(); let event = TaskExecutionNotificationEvent::tasks_complete(stats, failed_tasks); - self.try_send_notification(event, "tasks complete"); + sleep(Duration::from_millis(COMPLETION_NOTIFICATION_DELAY_MS)).await; } } From b3e4d265cf4834b3838a97dc0205fafa454677ec Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Thu, 24 Jul 2025 23:38:36 +1000 Subject: [PATCH 2/2] extra check token cancellation to a util --- crates/goose/src/agents/agent.rs | 7 ++++--- .../subagent_execution_tool/task_execution_tracker.rs | 6 +++--- crates/goose/src/utils.rs | 8 ++++++++ 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 0742ba33fe1d..f77c3fe7de53 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, tool::Tool, ToolError, ToolResult}; use regex::Regex; use rmcp::model::{Content, JsonRpcMessage, Prompt}; @@ -749,7 +750,7 @@ impl Agent { }); loop { - if cancel_token.as_ref().is_some_and(|t| t.is_cancelled()) { + if is_token_cancelled(&cancel_token) { break; } @@ -803,7 +804,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; } @@ -975,7 +976,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 eb0551f4e2d6..0c525fe4eae1 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 @@ -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; @@ -101,9 +102,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( @@ -311,6 +310,7 @@ impl TaskExecutionTracker { 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 60121f1bfe4d..1205953e83cc 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) -> bool { + cancellation_token + .as_ref() + .is_some_and(|t| t.is_cancelled()) +} + #[cfg(test)] mod tests { use super::*;