diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index b31ad53bd..1c76c2f0f 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -23,6 +23,9 @@ use tokio_util::task::AbortOnDropHandle; pub use self::export::{message_to_markdown, user_projected_message_to_markdown}; pub use builder::{build_session, SessionBuilderConfig}; use console::Color; +use goose::agents::platform_extensions::developer::shell::{ + parse_shell_output_notification, ShellOutputNotificationParams, ShellOutputStream, +}; use goose::agents::AgentEvent; use goose::agents::SUBAGENT_TOOL_REQUEST_TYPE; use goose::permission::permission_confirmation::PrincipalType; @@ -63,6 +66,9 @@ use tokio_util::sync::CancellationToken; use tracing::warn; const GOOSE_PLANNER_CONTEXT_LIMIT: &str = "GOOSE_PLANNER_CONTEXT_LIMIT"; +const SHELL_STATUS_FALLBACK_WIDTH: usize = 120; +const SHELL_STATUS_MAX_LINES: usize = 3; +const SHELL_STATUS_RESERVED_WIDTH: usize = 2; fn planner_provider_messages(plan_messages: &Conversation) -> Conversation { let projected_messages = plan_messages.agent_visible_messages(); @@ -2265,10 +2271,71 @@ fn handle_mcp_notification( ); } } + ServerNotification::CustomNotification(notification) => { + if let Some(params) = parse_shell_output_notification(notification) { + if is_stream_json_mode + || is_json_mode + || !interactive + || !std::io::stdout().is_terminal() + { + return; + } + display_shell_output_notification(params, progress_bars); + } + } _ => (), } } +fn display_shell_output_notification( + params: ShellOutputNotificationParams, + progress_bars: &mut output::McpSpinners, +) { + if params.truncated { + return; + } + + let max_width = console::Term::stdout() + .size_checked() + .map(|(_, width)| usize::from(width).saturating_sub(SHELL_STATUS_RESERVED_WIDTH)) + .unwrap_or(SHELL_STATUS_FALLBACK_WIDTH); + let lines = latest_shell_output_lines(¶ms, max_width) + .into_iter() + .map(|(stream, line)| match stream { + ShellOutputStream::Stdout => console::style(line).dim().to_string(), + ShellOutputStream::Stderr => console::style(line).yellow().dim().to_string(), + }) + .collect::>(); + if !lines.is_empty() { + progress_bars.log_shell_output(lines, SHELL_STATUS_MAX_LINES); + } +} + +fn latest_shell_output_lines( + params: &ShellOutputNotificationParams, + max_width: usize, +) -> Vec<(ShellOutputStream, String)> { + let mut lines = params + .chunks + .iter() + .rev() + .flat_map(|chunk| { + chunk + .output + .lines() + .rev() + .map(move |line| (chunk.stream, line)) + }) + .take(SHELL_STATUS_MAX_LINES) + .map(|(stream, line)| { + let line = output::sanitize_terminal_line(line); + (stream, safe_truncate(&line, max_width)) + }) + .collect::>(); + lines.reverse(); + lines +} + /// Format a logging notification from MCP, returns (formatted_message, subagent_id, notification_type) fn format_logging_notification( data: &Value, diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 7f04796af..e76040fb3 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -1,4 +1,4 @@ -use anstream::println; +use anstream::{adapter::strip_str, println}; use bat::WrappingMode; use console::{measure_text_width, style, Color, StyledObject, Term}; use goose::config::Config; @@ -15,7 +15,7 @@ use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use rmcp::model::{CallToolRequestParams, JsonObject, PromptArgument}; use serde_json::Value; use std::cell::RefCell; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; use std::fmt::Display; use std::io::{Error, IsTerminal, Write}; use std::path::Path; @@ -559,6 +559,17 @@ fn render_tool_response(resp: &ToolResponse, debug: bool) { } } +pub(super) fn sanitize_terminal_line(line: &str) -> String { + strip_str(line) + .flat_map(str::chars) + .filter(|character| *character == '\t' || !character.is_control()) + .collect() +} + +fn print_tool_output_line(line: &str) { + println!(" {}", style(sanitize_terminal_line(line)).dim()); +} + fn print_tool_output(text: &str) { if text.is_empty() { return; @@ -575,13 +586,13 @@ fn print_tool_output(text: &str) { let lines: Vec<&str> = text.lines().collect(); if lines.len() <= max_lines { for line in &lines { - println!(" {}", style(line).dim()); + print_tool_output_line(line); } } else { let head = max_lines / 2; let tail = max_lines - head; for line in &lines[..head] { - println!(" {}", style(line).dim()); + print_tool_output_line(line); } println!( " {}", @@ -593,7 +604,7 @@ fn print_tool_output(text: &str) { .italic() ); for line in &lines[lines.len() - tail..] { - println!(" {}", style(line).dim()); + print_tool_output_line(line); } } } @@ -1463,7 +1474,7 @@ pub fn display_cost_usage(provider: &str, model: &str, usage: &Usage) { pub struct McpSpinners { bars: HashMap, log_spinner: Option, - + shell_output_lines: VecDeque, multi_bar: MultiProgress, } @@ -1472,6 +1483,7 @@ impl McpSpinners { McpSpinners { bars: HashMap::new(), log_spinner: None, + shell_output_lines: VecDeque::new(), multi_bar: MultiProgress::new(), } } @@ -1494,6 +1506,13 @@ impl McpSpinners { spinner.set_message(message.to_string()); } + pub fn log_shell_output(&mut self, lines: Vec, max_lines: usize) { + let message = update_recent_lines(&mut self.shell_output_lines, lines, max_lines); + if !message.is_empty() { + self.log(&message); + } + } + pub fn update(&mut self, token: &str, value: f64, total: Option, message: Option<&str>) { let bar = self.bars.entry(token.to_string()).or_insert_with(|| { if let Some(total) = total { @@ -1520,16 +1539,69 @@ impl McpSpinners { if let Some(spinner) = self.log_spinner.as_mut() { spinner.disable_steady_tick(); } + self.shell_output_lines.clear(); self.multi_bar.clear() } } +fn update_recent_lines( + recent_lines: &mut VecDeque, + lines: impl IntoIterator, + max_lines: usize, +) -> String { + recent_lines.extend(lines); + while recent_lines.len() > max_lines { + recent_lines.pop_front(); + } + recent_lines + .iter() + .map(String::as_str) + .collect::>() + .join("\n ") +} + #[cfg(test)] mod tests { use super::*; use serde_json::json; use std::env; + #[test] + fn recent_lines_accumulate_across_updates() { + let mut recent_lines = VecDeque::new(); + let mut rendered = String::new(); + + for line in ["one", "two", "three", "four"] { + rendered = update_recent_lines(&mut recent_lines, [line.to_string()], 3); + } + + assert_eq!(rendered, "two\n three\n four"); + } + + #[test] + fn terminal_line_sanitizer_removes_escape_sequences_and_controls() { + assert_eq!( + sanitize_terminal_line( + "\x1b[31mred\x1b[0m \x1b[2J\x1b[H\ + \x1b]0;spoofed title\x07\ + \x1b]52;c;Y2xpcGJvYXJk\x1b\\safe" + ), + "red safe" + ); + assert_eq!( + sanitize_terminal_line("before\x08after\x07\r\tvisible"), + "beforeafter\tvisible" + ); + } + + #[test] + fn terminal_line_sanitizer_preserves_plain_unicode_text() { + assert_eq!( + sanitize_terminal_line("goose 🪿\t日本語"), + "goose 🪿\t日本語" + ); + } + #[test] fn formats_subagent_tool_call_names() { assert_eq!( diff --git a/crates/goose/src/acp/server/tool_notifications.rs b/crates/goose/src/acp/server/tool_notifications.rs index 1fab42053..391a99528 100644 --- a/crates/goose/src/acp/server/tool_notifications.rs +++ b/crates/goose/src/acp/server/tool_notifications.rs @@ -6,6 +6,10 @@ use rmcp::model::LoggingMessageNotificationParam; use rmcp::model::{ProgressNotificationParam, ServerNotification}; use serde::Serialize; +use crate::agents::platform_extensions::developer::shell::{ + parse_shell_output_notification, ShellOutputNotificationParams, +}; + #[expect(deprecated)] #[derive(Debug, Serialize)] #[serde(tag = "type", rename_all = "snake_case")] @@ -19,6 +23,9 @@ enum ToolNotification { PlatformEvent { params: serde_json::Value, }, + LiveOutput { + params: ShellOutputNotificationParams, + }, } pub(super) fn tool_notification_update( @@ -32,11 +39,15 @@ pub(super) fn tool_notification_update( ServerNotification::ProgressNotification(notification) => ToolNotification::Progress { params: notification.params, }, - ServerNotification::CustomNotification(notification) - if notification.method == "platform_event" => - { - ToolNotification::PlatformEvent { - params: notification.params.unwrap_or(serde_json::Value::Null), + ServerNotification::CustomNotification(notification) => { + if let Some(params) = parse_shell_output_notification(¬ification) { + ToolNotification::LiveOutput { params } + } else if notification.method == "platform_event" { + ToolNotification::PlatformEvent { + params: notification.params.unwrap_or(serde_json::Value::Null), + } + } else { + return None; } } _ => return None, @@ -60,6 +71,7 @@ pub(super) fn tool_notification_update( #[cfg(test)] mod tests { use super::tool_notification_update; + use crate::agents::platform_extensions::developer::shell::DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD; use agent_client_protocol::schema::v1::SessionUpdate; use rmcp::model::{ CancelledNotificationParam, CustomNotification, Notification, NumberOrString, @@ -188,6 +200,41 @@ mod tests { ); } + #[test] + fn maps_shell_output_custom_notification_to_live_output_meta() { + let notification = ServerNotification::CustomNotification(CustomNotification::new( + DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD, + Some(json!({ + "sequence": 2, + "chunks": [{ + "stream": "stdout", + "output": "ready\n" + }], + "truncated": false + })), + )); + + let update = tool_notification_update("tool_1", notification).expect("expected update"); + let value = serde_json::to_value(SessionUpdate::ToolCallUpdate(update)) + .expect("update should serialize"); + + assert_eq!(value["sessionUpdate"], "tool_call_update"); + assert_eq!(value["toolCallId"], "tool_1"); + assert_eq!(value["status"], "in_progress"); + assert_eq!(value["_meta"]["toolNotification"]["type"], "live_output"); + assert_eq!( + value["_meta"]["toolNotification"]["params"], + json!({ + "sequence": 2, + "chunks": [{ + "stream": "stdout", + "output": "ready\n" + }], + "truncated": false + }) + ); + } + #[test] fn ignores_non_platform_event_custom_notifications() { let notification = ServerNotification::CustomNotification(CustomNotification::new( diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index a89387699..ef3fe8281 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -1,7 +1,7 @@ use anyhow::Result; use axum::http::{HeaderMap, HeaderName, HeaderValue}; use chrono::{DateTime, Utc}; -use futures::stream::{FuturesUnordered, StreamExt}; +use futures::stream::{self, FuturesUnordered, StreamExt}; use futures::Stream; use futures::{future, FutureExt}; use once_cell::sync::Lazy; @@ -23,7 +23,7 @@ use std::time::Duration; use tempfile::{tempdir, TempDir}; use tokio::io::AsyncReadExt; use tokio::process::Command; -use tokio::sync::Mutex; +use tokio::sync::{mpsc, Mutex}; use tokio_stream::wrappers::ReceiverStream; use tokio_util::sync::CancellationToken; use tracing::{error, warn}; @@ -33,7 +33,7 @@ use super::extension::{ ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, PlatformExtensionContext, ToolInfo, PLATFORM_EXTENSIONS, }; -use super::tool_execution::{ToolCallContext, ToolCallResult}; +use super::tool_execution::{ToolCallContext, ToolCallNotificationEmitter, ToolCallResult}; use super::types::SharedProvider; use crate::action_required_manager::ActionRequiredManager; use crate::agents::extension::{Envs, ProcessExit}; @@ -50,7 +50,7 @@ use crate::prompt_template; use crate::subprocess::configure_subprocess; use rmcp::model::{ CallToolRequestParams, CallToolResult, ContentBlock, ErrorCode, ErrorData, GetPromptResult, - MetaObject, Prompt, Resource, ResourceContents, ServerInfo, Tool, + MetaObject, Prompt, Resource, ResourceContents, ServerInfo, ServerNotification, Tool, }; use rmcp::transport::auth::{AuthClient, CredentialStore}; use schemars::_private::NoSerialize; @@ -58,6 +58,8 @@ use serde_json::Value; type McpClientBox = Arc; +const TOOL_CALL_NOTIFICATION_CHANNEL_CAPACITY: usize = 32; + struct ActionRequiredStream { inner: ReceiverStream, session_id: String, @@ -1819,7 +1821,7 @@ impl ExtensionManager { let arguments = tool_call.arguments.clone(); let client = resolved.client.clone(); let hydration_client = client.clone(); - let notifications_receiver = client.subscribe().await; + let client_notifications_receiver = client.subscribe().await; let session_id = ctx.session_id.clone(); let action_required_tool_call_request_id = ctx.tool_call_request_id.clone(); let action_required_receiver = @@ -1852,6 +1854,32 @@ impl ExtensionManager { ctx.working_dir.clone(), ctx.tool_call_request_id.clone(), ); + let (owned_ctx, tool_call_notifications_receiver) = + if let Some(notification_emitter) = ctx.notification_emitter().cloned() { + ( + owned_ctx.with_notification_emitter(notification_emitter), + None, + ) + } else if owned_ctx.tool_call_request_id.is_some() { + let (tool_call_notifications_sender, tool_call_notifications_receiver) = + mpsc::channel(TOOL_CALL_NOTIFICATION_CHANNEL_CAPACITY); + ( + owned_ctx.with_notification_emitter(ToolCallNotificationEmitter::new( + tool_call_notifications_sender, + )), + Some(tool_call_notifications_receiver), + ) + } else { + (owned_ctx, None) + }; + let notification_stream: Box + Send + Unpin> = + match tool_call_notifications_receiver { + Some(tool_call_notifications_receiver) => Box::new(stream::select( + ReceiverStream::new(client_notifications_receiver), + ReceiverStream::new(tool_call_notifications_receiver), + )), + None => Box::new(ReceiverStream::new(client_notifications_receiver)), + }; let fut = async move { tracing::debug!( @@ -1892,7 +1920,7 @@ impl ExtensionManager { Ok(ToolCallResult { result: Box::new(fut.boxed()), - notification_stream: Some(Box::new(ReceiverStream::new(notifications_receiver))), + notification_stream: Some(notification_stream), action_required_stream: action_required_receiver.map( |(rx, session_id, tool_call_request_id)| { Box::new(ActionRequiredStream::new( @@ -2115,7 +2143,7 @@ impl ExtensionManager { mod tests { use super::*; use rmcp::model::CallToolResult; - use rmcp::model::{InitializeResult, JsonObject}; + use rmcp::model::{CustomNotification, InitializeResult, JsonObject}; use rmcp::{object, ServiceError as Error}; use rmcp::model::ListPromptsResult; @@ -2267,6 +2295,164 @@ mod tests { } } + struct ContextNotificationClient; + + #[async_trait::async_trait] + impl McpClientTrait for ContextNotificationClient { + fn get_info(&self) -> Option<&InitializeResult> { + None + } + + async fn list_tools( + &self, + session_id: &str, + next_cursor: Option, + cancellation_token: CancellationToken, + ) -> Result { + MockClient {} + .list_tools(session_id, next_cursor, cancellation_token) + .await + } + + async fn call_tool( + &self, + ctx: &ToolCallContext, + _name: &str, + _arguments: Option, + _cancellation_token: CancellationToken, + ) -> Result { + if let Some(emitter) = ctx.notification_emitter() { + let request_id = ctx + .tool_call_request_id + .as_deref() + .expect("an emitter requires a request ID"); + emitter.emit_best_effort(ServerNotification::CustomNotification( + CustomNotification::new(format!("scoped/{request_id}"), None), + )); + } + Ok(CallToolResult::success(vec![])) + } + + async fn subscribe(&self) -> mpsc::Receiver { + let (sender, receiver) = mpsc::channel(1); + sender + .try_send(ServerNotification::CustomNotification( + CustomNotification::new("client/subscription", None), + )) + .expect("test notification should fit"); + receiver + } + } + + async fn dispatch_notification_methods(ctx: ToolCallContext) -> Vec { + let temp_dir = tempfile::tempdir().unwrap(); + let extension_manager = + ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + extension_manager + .add_mock_extension( + "notifications".to_string(), + Arc::new(ContextNotificationClient), + ) + .await; + + let tool_call = CallToolRequestParams::new("notifications__tool".to_string()) + .with_arguments(object!({})); + let dispatched = extension_manager + .dispatch_tool_call(&ctx, tool_call, CancellationToken::default()) + .await + .expect("tool call should dispatch"); + + assert!(dispatched.result.await.is_ok()); + + let mut methods = dispatched + .notification_stream + .expect("notification stream should exist") + .filter_map(|notification| async move { + match notification { + ServerNotification::CustomNotification(notification) => { + Some(notification.method) + } + _ => None, + } + }) + .collect::>() + .await; + methods.sort(); + methods + } + + #[tokio::test] + async fn dispatch_merges_request_scoped_and_client_notifications() { + let methods = dispatch_notification_methods(ToolCallContext::new( + "session".to_string(), + None, + Some("request".to_string()), + )) + .await; + + assert_eq!(methods, vec!["client/subscription", "scoped/request"]); + } + + #[tokio::test] + async fn dispatch_reuses_existing_notification_emitter() { + let temp_dir = tempfile::tempdir().unwrap(); + let extension_manager = + ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + extension_manager + .add_mock_extension( + "notifications".to_string(), + Arc::new(ContextNotificationClient), + ) + .await; + let (sender, mut receiver) = mpsc::channel(1); + let ctx = ToolCallContext::new( + "nested-session".to_string(), + None, + Some("nested-request".to_string()), + ) + .with_notification_emitter(ToolCallNotificationEmitter::new(sender)); + let tool_call = CallToolRequestParams::new("notifications__tool".to_string()) + .with_arguments(object!({})); + + let dispatched = extension_manager + .dispatch_tool_call(&ctx, tool_call, CancellationToken::default()) + .await + .expect("tool call should dispatch"); + assert!(dispatched.result.await.is_ok()); + + let notification = receiver + .try_recv() + .expect("parent emitter should receive nested notification"); + let ServerNotification::CustomNotification(notification) = notification else { + panic!("expected a custom notification"); + }; + assert_eq!(notification.method, "scoped/nested-request"); + + let methods = dispatched + .notification_stream + .expect("client notification stream should exist") + .filter_map(|notification| async move { + match notification { + ServerNotification::CustomNotification(notification) => { + Some(notification.method) + } + _ => None, + } + }) + .collect::>() + .await; + assert_eq!(methods, vec!["client/subscription"]); + } + + #[tokio::test] + async fn dispatch_without_request_id_uses_only_client_notifications() { + let methods = + dispatch_notification_methods(ToolCallContext::new("session".to_string(), None, None)) + .await; + + assert_eq!(methods, vec!["client/subscription"]); + } + #[tokio::test] async fn test_dispatch_tool_call() { use super::super::tool_execution::ToolCallContext; diff --git a/crates/goose/src/agents/platform_extensions/developer/mod.rs b/crates/goose/src/agents/platform_extensions/developer/mod.rs index 2d69ad216..ce1eb7b28 100644 --- a/crates/goose/src/agents/platform_extensions/developer/mod.rs +++ b/crates/goose/src/agents/platform_extensions/developer/mod.rs @@ -1,6 +1,7 @@ pub mod edit; pub mod image; pub mod shell; +mod shell_output_streaming; pub mod tree; use crate::agents::extension::PlatformExtensionContext; @@ -206,7 +207,13 @@ impl McpClientTrait for DeveloperClient { "shell" => match Self::parse_args::(arguments) { Ok(params) => Ok(self .shell_tool - .shell_with_cwd(params, working_dir, Some(&ctx.session_id), cancel_token) + .shell_with_cwd_and_emitter( + params, + working_dir, + Some(&ctx.session_id), + ctx.notification_emitter().cloned(), + cancel_token, + ) .await), Err(error) => Ok(ShellTool::error_result(&format!("Error: {error}"), None)), }, diff --git a/crates/goose/src/agents/platform_extensions/developer/shell.rs b/crates/goose/src/agents/platform_extensions/developer/shell.rs index ab82b5f26..193215832 100644 --- a/crates/goose/src/agents/platform_extensions/developer/shell.rs +++ b/crates/goose/src/agents/platform_extensions/developer/shell.rs @@ -20,8 +20,15 @@ use tokio::task::JoinHandle; use tokio_stream::{wrappers::SplitStream, StreamExt}; use tokio_util::sync::CancellationToken; +use crate::agents::tool_execution::ToolCallNotificationEmitter; use crate::subprocess::SubprocessExt; +pub use super::shell_output_streaming::{ + parse_shell_output_notification, ShellOutputNotificationChunk, ShellOutputNotificationParams, + ShellOutputStream, DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD, +}; +use super::shell_output_streaming::{ShellOutputBatcher, SHELL_LIVE_OUTPUT_FLUSH_INTERVAL}; + /// Check if the current process is running inside a Flatpak sandbox. /// /// When inside Flatpak, shell commands must be wrapped with `flatpak-spawn --host` @@ -364,6 +371,18 @@ impl ShellTool { working_dir: Option<&std::path::Path>, session_id: Option<&str>, cancellation_token: CancellationToken, + ) -> CallToolResult { + self.shell_with_cwd_and_emitter(params, working_dir, session_id, None, cancellation_token) + .await + } + + pub(crate) async fn shell_with_cwd_and_emitter( + &self, + params: ShellParams, + working_dir: Option<&std::path::Path>, + session_id: Option<&str>, + notification_emitter: Option, + cancellation_token: CancellationToken, ) -> CallToolResult { if params.command.trim().is_empty() { return Self::error_result("Command cannot be empty.", None); @@ -382,6 +401,7 @@ impl ShellTool { working_dir, login_path_ref, session_id, + notification_emitter, cancellation_token, ) .await @@ -529,6 +549,7 @@ async fn run_command( working_dir: Option<&std::path::Path>, login_path: Option<&str>, session_id: Option<&str>, + notification_emitter: Option, cancellation_token: CancellationToken, ) -> Result { let timeout_secs = Some(resolve_shell_timeout(timeout_secs)); @@ -553,7 +574,12 @@ async fn run_command( .ok_or_else(|| "Failed to capture stderr".to_string())?; let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); - let output_task = tokio::spawn(collect_tagged_lines(child_stdout, child_stderr, tx)); + let output_task = tokio::spawn(collect_tagged_lines( + child_stdout, + child_stderr, + tx, + notification_emitter, + )); let abort_handle = output_task.abort_handle(); let mut timed_out = false; @@ -750,14 +776,51 @@ async fn collect_tagged_lines( stdout: tokio::process::ChildStdout, stderr: tokio::process::ChildStderr, tx: tokio::sync::mpsc::UnboundedSender<(bool, String)>, + notification_emitter: Option, ) -> Result<(), std::io::Error> { let stdout_lines = SplitStream::new(BufReader::new(stdout).split(b'\n')).map(|l| (false, l)); let stderr_lines = SplitStream::new(BufReader::new(stderr).split(b'\n')).map(|l| (true, l)); let mut merged = stdout_lines.merge(stderr_lines); + let mut output_batcher = notification_emitter.map(ShellOutputBatcher::new); + let mut flush_interval = tokio::time::interval_at( + tokio::time::Instant::now() + SHELL_LIVE_OUTPUT_FLUSH_INTERVAL, + SHELL_LIVE_OUTPUT_FLUSH_INTERVAL, + ); + flush_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - while let Some((is_stderr, line)) = merged.next().await { - let line = line?; - let _ = tx.send((is_stderr, String::from_utf8_lossy(&line).into_owned())); + loop { + tokio::select! { + tagged_line = merged.next() => { + let Some((is_stderr, line)) = tagged_line else { + break; + }; + let line = match line { + Ok(line) => String::from_utf8_lossy(&line).into_owned(), + Err(error) => { + if let Some(output_batcher) = output_batcher.as_mut() { + output_batcher.flush(); + } + return Err(error); + } + }; + + if let Some(output_batcher) = output_batcher.as_mut() { + if output_batcher.push_line(is_stderr, &line) { + flush_interval.reset(); + } + } + let _ = tx.send((is_stderr, line)); + } + _ = flush_interval.tick(), if output_batcher.is_some() => { + if let Some(output_batcher) = output_batcher.as_mut() { + output_batcher.flush(); + } + } + } + } + + if let Some(output_batcher) = output_batcher.as_mut() { + output_batcher.flush(); } Ok(()) } @@ -879,6 +942,31 @@ mod tests { assert!(extract_text(&result).contains("hello")); } + #[cfg(not(windows))] + #[tokio::test] + async fn full_live_notification_channel_does_not_change_final_output() { + let tool = ShellTool::new_for_test().unwrap(); + let (sender, _receiver) = tokio::sync::mpsc::channel(1); + let result = tool + .shell_with_cwd_and_emitter( + ShellParams { + command: "printf 'out-1\\nout-2\\n'; printf 'err-1\\nerr-2\\n' >&2".to_string(), + timeout_secs: None, + }, + None, + None, + Some(ToolCallNotificationEmitter::new(sender)), + CancellationToken::new(), + ) + .await; + + let shell_output = extract_shell_output(&result); + assert_eq!(result.is_error, Some(false)); + assert_eq!(shell_output.stdout, "out-1\nout-2"); + assert_eq!(shell_output.stderr, "err-1\nerr-2"); + assert_eq!(shell_output.exit_code, Some(0)); + } + #[cfg(not(windows))] #[tokio::test] async fn shell_returns_error_for_non_zero_exit() { diff --git a/crates/goose/src/agents/platform_extensions/developer/shell_output_streaming.rs b/crates/goose/src/agents/platform_extensions/developer/shell_output_streaming.rs new file mode 100644 index 000000000..9a24cf451 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/developer/shell_output_streaming.rs @@ -0,0 +1,226 @@ +use std::time::Duration; + +use rmcp::model::{CustomNotification, ServerNotification}; +use serde::{Deserialize, Serialize}; + +use crate::agents::tool_execution::ToolCallNotificationEmitter; + +pub(super) const SHELL_LIVE_OUTPUT_FLUSH_INTERVAL: Duration = Duration::from_millis(150); +const SHELL_LIVE_OUTPUT_BATCH_BYTES: usize = 16 * 1024; +const SHELL_LIVE_OUTPUT_LIMIT_BYTES: usize = 256 * 1024; + +pub const DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD: &str = "goose/developer_shell_output"; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum ShellOutputStream { + Stdout, + Stderr, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct ShellOutputNotificationChunk { + pub stream: ShellOutputStream, + pub output: String, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct ShellOutputNotificationParams { + pub sequence: u64, + pub chunks: Vec, + pub truncated: bool, +} + +pub fn parse_shell_output_notification( + notification: &CustomNotification, +) -> Option { + if notification.method != DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD { + return None; + } + serde_json::from_value(notification.params.clone()?).ok() +} + +pub(super) struct ShellOutputBatcher { + emitter: ToolCallNotificationEmitter, + sequence: u64, + chunks: Vec, + buffered_bytes: usize, + live_output_bytes: usize, + emitted_first_line: bool, + truncated: bool, +} + +impl ShellOutputBatcher { + pub(super) fn new(emitter: ToolCallNotificationEmitter) -> Self { + Self { + emitter, + sequence: 0, + chunks: Vec::new(), + buffered_bytes: 0, + live_output_bytes: 0, + emitted_first_line: false, + truncated: false, + } + } + + pub(super) fn push_line(&mut self, is_stderr: bool, line: &str) -> bool { + if self.truncated { + return false; + } + + let line = line.strip_suffix('\r').unwrap_or(line); + let output_bytes = line.len() + 1; + if self.live_output_bytes + output_bytes > SHELL_LIVE_OUTPUT_LIMIT_BYTES { + self.flush(); + self.truncated = true; + self.emit_notification(Vec::new(), true); + return true; + } + + self.live_output_bytes += output_bytes; + self.buffered_bytes += output_bytes; + + let stream = if is_stderr { + ShellOutputStream::Stderr + } else { + ShellOutputStream::Stdout + }; + if let Some(chunk) = self + .chunks + .last_mut() + .filter(|chunk| chunk.stream == stream) + { + chunk.output.push_str(line); + chunk.output.push('\n'); + } else { + self.chunks.push(ShellOutputNotificationChunk { + stream, + output: format!("{line}\n"), + }); + } + + if !self.emitted_first_line || self.buffered_bytes >= SHELL_LIVE_OUTPUT_BATCH_BYTES { + self.emitted_first_line = true; + return self.flush(); + } + + false + } + + pub(super) fn flush(&mut self) -> bool { + if self.chunks.is_empty() { + return false; + } + + self.buffered_bytes = 0; + let chunks = std::mem::take(&mut self.chunks); + self.emit_notification(chunks, false); + true + } + + fn emit_notification(&mut self, chunks: Vec, truncated: bool) { + self.sequence += 1; + let params = ShellOutputNotificationParams { + sequence: self.sequence, + chunks, + truncated, + }; + if let Ok(params) = serde_json::to_value(params) { + self.emitter + .emit_best_effort(ServerNotification::CustomNotification( + CustomNotification::new( + DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD, + Some(params), + ), + )); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::mpsc; + + fn receive_params( + receiver: &mut mpsc::Receiver, + ) -> ShellOutputNotificationParams { + let notification = receiver + .try_recv() + .expect("expected a shell output notification"); + let ServerNotification::CustomNotification(notification) = notification else { + panic!("expected a custom notification"); + }; + parse_shell_output_notification(¬ification) + .expect("expected valid shell output notification params") + } + + #[test] + fn first_line_emits_immediately_and_normalizes_crlf() { + let (sender, mut receiver) = mpsc::channel(4); + let mut batcher = ShellOutputBatcher::new(ToolCallNotificationEmitter::new(sender)); + + assert!(batcher.push_line(false, "hello\r")); + + let params = receive_params(&mut receiver); + assert_eq!(params.sequence, 1); + assert!(!params.truncated); + assert_eq!(params.chunks.len(), 1); + assert_eq!(params.chunks[0].stream, ShellOutputStream::Stdout); + assert_eq!(params.chunks[0].output, "hello\n"); + } + + #[test] + fn flush_coalesces_consecutive_lines_and_advances_sequence() { + let (sender, mut receiver) = mpsc::channel(4); + let mut batcher = ShellOutputBatcher::new(ToolCallNotificationEmitter::new(sender)); + batcher.push_line(false, "first"); + receive_params(&mut receiver); + + assert!(!batcher.push_line(false, "second")); + assert!(!batcher.push_line(false, "third")); + assert!(!batcher.push_line(true, "warning")); + assert!(batcher.flush()); + + let params = receive_params(&mut receiver); + assert_eq!(params.sequence, 2); + assert!(!params.truncated); + assert_eq!(params.chunks.len(), 2); + assert_eq!(params.chunks[0].stream, ShellOutputStream::Stdout); + assert_eq!(params.chunks[0].output, "second\nthird\n"); + assert_eq!(params.chunks[1].stream, ShellOutputStream::Stderr); + assert_eq!(params.chunks[1].output, "warning\n"); + } + + #[test] + fn live_output_limit_emits_one_truncation_notification() { + let (sender, mut receiver) = mpsc::channel(4); + let mut batcher = ShellOutputBatcher::new(ToolCallNotificationEmitter::new(sender)); + let output_at_limit = "x".repeat(SHELL_LIVE_OUTPUT_LIMIT_BYTES - 1); + + assert!(batcher.push_line(false, &output_at_limit)); + receive_params(&mut receiver); + assert!(batcher.push_line(false, "omitted")); + + let params = receive_params(&mut receiver); + assert_eq!(params.sequence, 2); + assert!(params.truncated); + assert!(params.chunks.is_empty()); + + assert!(!batcher.push_line(false, "also omitted")); + assert!(!batcher.flush()); + assert!(receiver.try_recv().is_err()); + } + + #[test] + fn parser_ignores_unrelated_or_malformed_custom_notifications() { + let unrelated = CustomNotification::new("goose/other", Some(serde_json::json!({}))); + let malformed = CustomNotification::new( + DEVELOPER_SHELL_OUTPUT_NOTIFICATION_METHOD, + Some(serde_json::json!({ "sequence": "not-a-number" })), + ); + + assert!(parse_shell_output_notification(&unrelated).is_none()); + assert!(parse_shell_output_notification(&malformed).is_none()); + } +} diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index 94f05c2d4..3a579587c 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -4,6 +4,7 @@ use futures::{Stream, StreamExt}; use rmcp::model::CallToolResult; use std::collections::HashMap; use std::future::Future; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use std::path::PathBuf; @@ -14,12 +15,29 @@ use crate::mcp_utils::ToolResult; use crate::permission::Permission; use rmcp::model::{ContentBlock, ServerNotification}; +#[derive(Clone)] +pub(crate) struct ToolCallNotificationEmitter { + sender: mpsc::Sender, +} + +impl ToolCallNotificationEmitter { + pub(crate) fn new(sender: mpsc::Sender) -> Self { + Self { sender } + } + + pub(crate) fn emit_best_effort(&self, notification: ServerNotification) { + // Do not let a slow notification consumer delay tool execution. + let _ = self.sender.try_send(notification); + } +} + /// Context passed through the tool call dispatch chain. #[derive(Clone)] pub struct ToolCallContext { pub session_id: String, pub working_dir: Option, pub tool_call_request_id: Option, + notification_emitter: Option, } impl ToolCallContext { @@ -32,12 +50,25 @@ impl ToolCallContext { session_id, working_dir, tool_call_request_id, + notification_emitter: None, } } pub fn working_dir_str(&self) -> Option<&str> { self.working_dir.as_ref().and_then(|p| p.to_str()) } + + pub(crate) fn with_notification_emitter( + mut self, + notification_emitter: ToolCallNotificationEmitter, + ) -> Self { + self.notification_emitter = Some(notification_emitter); + self + } + + pub(crate) fn notification_emitter(&self) -> Option<&ToolCallNotificationEmitter> { + self.notification_emitter.as_ref() + } } // ToolCallResult combines the result of a tool call with an optional notification stream that diff --git a/crates/goose/tests/acp_common_tests/mod.rs b/crates/goose/tests/acp_common_tests/mod.rs index 3e57e2190..08bc1a840 100644 --- a/crates/goose/tests/acp_common_tests/mod.rs +++ b/crates/goose/tests/acp_common_tests/mod.rs @@ -1447,8 +1447,15 @@ pub async fn run_shell_terminal_false() { .await .unwrap(); assert!(!output.text.is_empty()); + let mut notifications = session.notifications(); + notifications.retain(|notification| { + !matches!( + notification, + Notification::ToolCallStatus(ToolCallStatus::InProgress) + ) + }); assert_notifications( - &session.notifications(), + ¬ifications, &[ Notification::ToolCall, Notification::ToolCallContent("content".into()), diff --git a/ui/desktop/src/acp/__tests__/toolNotifications.test.ts b/ui/desktop/src/acp/__tests__/toolNotifications.test.ts new file mode 100644 index 000000000..a642031d3 --- /dev/null +++ b/ui/desktop/src/acp/__tests__/toolNotifications.test.ts @@ -0,0 +1,71 @@ +import type { ToolCallUpdate } from '@agentclientprotocol/sdk'; +import { describe, expect, it } from 'vitest'; +import { toolNotificationEvent } from '../adapter/toolNotifications'; + +function liveOutputUpdate(params: unknown): ToolCallUpdate { + return { + toolCallId: 'tool-1', + status: 'in_progress', + _meta: { + toolNotification: { + type: 'live_output', + params, + }, + }, + }; +} + +describe('toolNotificationEvent', () => { + it('maps live output metadata to a tool-correlated notification', () => { + const event = toolNotificationEvent( + liveOutputUpdate({ + sequence: 2, + chunks: [ + { + stream: 'stdout', + output: 'ready\n', + }, + { + stream: 'stderr', + output: 'warning\n', + }, + ], + truncated: false, + }) + ); + + expect(event).toEqual({ + type: 'Notification', + request_id: 'tool-1', + message: { + method: 'goose/live_output', + params: { + sequence: 2, + chunks: [ + { + stream: 'stdout', + output: 'ready\n', + }, + { + stream: 'stderr', + output: 'warning\n', + }, + ], + truncated: false, + }, + }, + }); + }); + + it('ignores malformed live output metadata', () => { + expect( + toolNotificationEvent( + liveOutputUpdate({ + sequence: 'two', + chunks: [], + truncated: false, + }) + ) + ).toBeUndefined(); + }); +}); diff --git a/ui/desktop/src/acp/adapter/toolNotifications.ts b/ui/desktop/src/acp/adapter/toolNotifications.ts index f23ab7042..fccd3a940 100644 --- a/ui/desktop/src/acp/adapter/toolNotifications.ts +++ b/ui/desktop/src/acp/adapter/toolNotifications.ts @@ -1,5 +1,9 @@ import type { ToolCallUpdate } from '@agentclientprotocol/sdk'; -import type { NotificationEvent } from '../../types/message'; +import type { + LiveOutputNotificationChunk, + LiveOutputNotificationParams, + NotificationEvent, +} from '../../types/message'; import type { AcpChatStateChange } from './shared'; import { isRecord } from './shared'; @@ -15,6 +19,10 @@ type ToolNotification = | { type: 'platform_event'; params: PlatformEventParams; + } + | { + type: 'live_output'; + params: LiveOutputNotificationParams; }; type LoggingMessageNotificationParams = { @@ -32,6 +40,14 @@ type ProgressNotificationParams = { type PlatformEventParams = Record; +function isLiveOutputChunk(value: unknown): value is LiveOutputNotificationChunk { + return ( + isRecord(value) && + (value.stream === 'stdout' || value.stream === 'stderr') && + typeof value.output === 'string' + ); +} + export function toolNotificationChange( update: ToolCallUpdate ): Extract | undefined { @@ -80,6 +96,11 @@ function parseToolNotification(meta: unknown): ToolNotification | undefined { return params ? { type: 'platform_event', params } : undefined; } + if (toolNotification.type === 'live_output') { + const params = parseLiveOutputParams(toolNotification.params); + return params ? { type: 'live_output', params } : undefined; + } + return undefined; } @@ -116,6 +137,24 @@ function parsePlatformEventParams(value: unknown): PlatformEventParams | undefin return isRecord(value) ? value : undefined; } +function parseLiveOutputParams(value: unknown): LiveOutputNotificationParams | undefined { + if ( + !isRecord(value) || + typeof value.sequence !== 'number' || + !Array.isArray(value.chunks) || + !value.chunks.every(isLiveOutputChunk) || + typeof value.truncated !== 'boolean' + ) { + return undefined; + } + + return { + sequence: value.sequence, + chunks: value.chunks, + truncated: value.truncated, + }; +} + function toNotificationEvent( toolCallId: string, toolNotification: ToolNotification @@ -138,5 +177,7 @@ function notificationMethod(toolNotification: ToolNotification): string { return 'notifications/progress'; case 'platform_event': return 'platform_event'; + case 'live_output': + return 'goose/live_output'; } } diff --git a/ui/desktop/src/components/ToolCallWithResponse.test.tsx b/ui/desktop/src/components/ToolCallWithResponse.test.tsx new file mode 100644 index 000000000..42b7e786c --- /dev/null +++ b/ui/desktop/src/components/ToolCallWithResponse.test.tsx @@ -0,0 +1,103 @@ +import { render, screen } from '@testing-library/react'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { IntlTestWrapper } from '../i18n/test-utils'; +import type { + NotificationEvent, + ToolRequestMessageContent, + ToolResponseMessageContent, +} from '../types/message'; +import ToolCallWithResponse from './ToolCallWithResponse'; + +const toolRequest: ToolRequestMessageContent = { + type: 'toolRequest', + id: 'tool-1', + toolCall: { + status: 'success', + value: { + name: 'developer__shell', + arguments: { + command: 'build', + }, + }, + }, +}; + +const liveOutputNotification: NotificationEvent = { + type: 'Notification', + request_id: 'tool-1', + message: { + method: 'goose/live_output', + params: { + sequence: 1, + chunks: [ + { + stream: 'stdout', + output: 'starting\n', + }, + { + stream: 'stderr', + output: 'checking\n', + }, + ], + truncated: false, + }, + }, +}; + +const toolResponse: ToolResponseMessageContent = { + type: 'toolResponse', + id: 'tool-1', + toolResult: { + status: 'success', + value: { + content: [ + { + type: 'text', + text: 'final result', + }, + ], + isError: false, + }, + }, +}; + +function renderToolCall(response?: ToolResponseMessageContent) { + return render( + , + { wrapper: IntlTestWrapper } + ); +} + +describe('ToolCallWithResponse live output', () => { + beforeEach(() => { + vi.mocked(window.electron.getSetting).mockResolvedValue('detailed'); + }); + + it('renders raw live output while running and replaces it with the final result', async () => { + const { rerender } = renderToolCall(); + + expect(screen.getByText(/starting/)).toHaveTextContent('starting checking'); + expect(screen.queryByText(/stdout|stderr/)).not.toBeInTheDocument(); + + rerender( + + ); + + expect(screen.queryByText(/starting/)).not.toBeInTheDocument(); + expect(await screen.findByText('final result')).toBeInTheDocument(); + }); +}); diff --git a/ui/desktop/src/components/ToolCallWithResponse.tsx b/ui/desktop/src/components/ToolCallWithResponse.tsx index a3fab0fc1..f47e4d9f5 100644 --- a/ui/desktop/src/components/ToolCallWithResponse.tsx +++ b/ui/desktop/src/components/ToolCallWithResponse.tsx @@ -9,6 +9,7 @@ import { ToolRequestMessageContent, ToolResponseMessageContent, NotificationEvent, + LiveOutputNotificationParams, ToolConfirmationData, } from '../types/message'; import { cn, snakeToTitleCase } from '../utils'; @@ -451,6 +452,18 @@ const notificationToProgress = (notification: NotificationEvent): Progress => { return message.params as Progress; }; +const liveOutputToString = (notifications: NotificationEvent[] | undefined): string => + notifications + ?.filter((notification) => { + const message = notification.message as { method?: string }; + return message.method === 'goose/live_output'; + }) + .flatMap((notification) => { + const message = notification.message as { params?: LiveOutputNotificationParams }; + return message.params?.chunks.map((chunk) => chunk.output) ?? []; + }) + .join('') ?? ''; + // Helper function to extract toolcall name const getToolName = (toolCallName: string): string => { const lastIndex = toolCallName.lastIndexOf('__'); @@ -534,6 +547,7 @@ function ToolCallView({ loadingStatus === 'success' && toolResponse?.toolResult ? getToolResultContent(toolResponse.toolResult) : []; + const liveOutput = toolResponse ? '' : liveOutputToString(notifications); const logs = notifications ?.filter((notification) => { @@ -561,8 +575,9 @@ function ToolCallView({ (entries) => entries.sort((a, b) => b.progress - a.progress)[0] ); - const isRenderingProgress = - loadingStatus === 'loading' && (progressEntries.length > 0 || (logs || []).length > 0); + const isRenderingActivity = + loadingStatus === 'loading' && + (progressEntries.length > 0 || (logs || []).length > 0 || liveOutput.length > 0); // Function to create a descriptive representation of what the tool is doing const getToolDescription = (): string | null => { @@ -783,7 +798,7 @@ function ToolCallView({ ); return ( )} + {liveOutput && ( +
+ +
+ )} + {toolResults.length === 0 && progressEntries.length > 0 && progressEntries.map((entry, index) => ( @@ -958,6 +979,30 @@ interface ToolResultViewProps { isStartExpanded: boolean; } +function LiveOutputView({ output }: { output: string }) { + const intl = useIntl(); + const outputRef = useRef(null); + + useEffect(() => { + if (outputRef.current) { + outputRef.current.scrollTop = outputRef.current.scrollHeight; + } + }, [output]); + + return ( + {intl.formatMessage(i18n.output)}} + isStartExpanded={true} + > +
+ 
+          {output}
+
+
+ + ); +} + function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) { const intl = useIntl(); const hasText = (c: ContentBlock): c is ContentBlock & { text: string } => diff --git a/ui/desktop/src/types/message.ts b/ui/desktop/src/types/message.ts index efc6f1608..ec48d134a 100644 --- a/ui/desktop/src/types/message.ts +++ b/ui/desktop/src/types/message.ts @@ -258,6 +258,17 @@ export type ToolConfirmationRequestContent = ToolConfirmationRequest & { }; export type NotificationEvent = Extract; +export type LiveOutputNotificationParams = { + sequence: number; + chunks: LiveOutputNotificationChunk[]; + truncated: boolean; +}; + +export type LiveOutputNotificationChunk = { + stream: 'stdout' | 'stderr'; + output: string; +}; + export interface ImageData { data: string; // base64 encoded image data mimeType: string;