Skip to content

fix: ensure execution task result is shown #3629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lifeizhou-ap merged 3 commits into from
Jul 25, 2025
Merged

fix: ensure execution task result is shown #3629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b988dca
added extra time
lifeizhou-ap Jul 24, 2025
b3e4d26
extra check token cancellation to a util
lifeizhou-ap Jul 24, 2025
59f2517
Merge branch 'main' into lifei/fix-ensure-execution-result-shown
lifeizhou-ap Jul 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions crates/goose/src/agents/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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 {
Expand Down
11 changes: 6 additions & 5 deletions crates/goose/src/agents/subagent_execution_tool/task_execution_tracker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand All @@ -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;

Expand All @@ -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() {
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Copy link
Collaborator

@DOsinga DOsinga Jul 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good, but are you saying the message doesn't get there if we complete here?

Copy link
Collaborator Author

@lifeizhou-ap lifeizhou-ap Jul 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The extra sleep is to handle the race condition.

The notification flow:

  1. TaskExecutionTracker → try_send → mpsc channel queue
  2. ReceiverStream reads from queue → agent's event stream
  3. Session loop processes from agent stream → CLI display

When the last task completes, notifications is sent async and the sub agent execution tool completes. The extra sleep is to ensure the asynchronous notification processing completes before the tool returns its result to the agent.

Without sleep:
T=1ms: Notifications queued in channel
T=2ms: Tool returns immediately (no sleep)
T=15ms: Agent stream terminates
T=15ms: Session loop exits
T=16ms: Channel dropped with notifications inside

Copy link
Collaborator

@michaelneale michaelneale Jul 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this could be a future source of bugs - really no other way?

}
}
8 changes: 8 additions & 0 deletions crates/goose/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.)
Expand All @@ -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::*;
Expand Down
Loading