tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: streaming shell output while commands run (#10808)

Browse Source
This commit is contained in:
Lifei Zhou
2026-07-30 21:06:35 +10:00
committed by GitHub
parent 2694fff7e3
commit cf312f1d1d
14 changed files with 1030 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+67
View File
@@ -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(&params, 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::<Vec<_>>();
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::<Vec<_>>();
lines.reverse();
lines
}
/// Format a logging notification from MCP, returns (formatted_message, subagent_id, notification_type)
fn format_logging_notification(
data: &Value,
crates/goose-cli/src/session/output.rs
+78 -6
View File
@@ -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<String, ProgressBar>,
log_spinner: Option<ProgressBar>,
shell_output_lines: VecDeque<String>,
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<String>, 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<f64>, 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<String>,
lines: impl IntoIterator<Item = String>,
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::<Vec<_>>()
.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!(
crates/goose/src/acp/server/tool_notifications.rs
+52 -5
View File
@@ -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(&notification) {
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(
crates/goose/src/agents/extension_manager.rs
+193 -7
View File
@@ -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<dyn McpClientTrait>;
const TOOL_CALL_NOTIFICATION_CHANNEL_CAPACITY: usize = 32;
struct ActionRequiredStream {
inner: ReceiverStream<crate::conversation::message::Message>,
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<dyn Stream<Item = ServerNotification> + 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<String>,
cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
MockClient {}
.list_tools(session_id, next_cursor, cancellation_token)
.await
}
async fn call_tool(
&self,
ctx: &ToolCallContext,
_name: &str,
_arguments: Option<JsonObject>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
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<ServerNotification> {
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<String> {
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::<Vec<_>>()
.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::<Vec<_>>()
.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;
crates/goose/src/agents/platform_extensions/developer/mod.rs
+8 -1
View File
@@ -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::<ShellParams>(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)),
},
crates/goose/src/agents/platform_extensions/developer/shell.rs
+92 -4
View File
@@ -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<ToolCallNotificationEmitter>,
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<ToolCallNotificationEmitter>,
cancellation_token: CancellationToken,
) -> Result<ExecutionOutput, String> {
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<ToolCallNotificationEmitter>,
) -> 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() {
crates/goose/src/agents/platform_extensions/developer/shell_output_streaming.rs
+226
View File
@@ -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<ShellOutputNotificationChunk>,
pub truncated: bool,
}
pub fn parse_shell_output_notification(
notification: &CustomNotification,
) -> Option<ShellOutputNotificationParams> {
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<ShellOutputNotificationChunk>,
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<ShellOutputNotificationChunk>, 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<ServerNotification>,
) -> 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(&notification)
.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());
}
}
crates/goose/src/agents/tool_execution.rs
+31
View File
@@ -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<ServerNotification>,
}
impl ToolCallNotificationEmitter {
pub(crate) fn new(sender: mpsc::Sender<ServerNotification>) -> 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<PathBuf>,
pub tool_call_request_id: Option<String>,
notification_emitter: Option<ToolCallNotificationEmitter>,
}
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
crates/goose/tests/acp_common_tests/mod.rs
+8 -1
View File
@@ -1447,8 +1447,15 @@ pub async fn run_shell_terminal_false<C: Connection>() {
.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(),
&notifications,
&[
Notification::ToolCall,
Notification::ToolCallContent("content".into()),
ui/desktop/src/acp/__tests__/toolNotifications.test.ts
+71
View File
@@ -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();
});
});
ui/desktop/src/acp/adapter/toolNotifications.ts
+42 -1
View File
@@ -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<string, unknown>;
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<AcpChatStateChange, { type: 'notification' }> | 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';
}
}
ui/desktop/src/components/ToolCallWithResponse.test.tsx
+103
View File
@@ -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(
<ToolCallWithResponse
isCancelledMessage={false}
toolRequest={toolRequest}
toolResponse={response}
notifications={[liveOutputNotification]}
isStreamingMessage={!response}
isPendingApproval={false}
/>,
{ 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(
<ToolCallWithResponse
isCancelledMessage={false}
toolRequest={toolRequest}
toolResponse={toolResponse}
notifications={[liveOutputNotification]}
isStreamingMessage={false}
isPendingApproval={false}
/>
);
expect(screen.queryByText(/starting/)).not.toBeInTheDocument();
expect(await screen.findByText('final result')).toBeInTheDocument();
});
});
ui/desktop/src/components/ToolCallWithResponse.tsx
+48 -3
View File
@@ -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 (
<ToolCallExpandable
isStartExpanded={isRenderingProgress || isExpandToolDetails}
isStartExpanded={isRenderingActivity || isExpandToolDetails}
isForceExpand={false}
label={
extensionTooltip ? (
@@ -833,6 +848,12 @@ function ToolCallView({
</div>
)}
{liveOutput && (
<div className="border-t border-border-primary">
<LiveOutputView output={liveOutput} />
</div>
)}
{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<HTMLDivElement>(null);
useEffect(() => {
if (outputRef.current) {
outputRef.current.scrollTop = outputRef.current.scrollHeight;
}
}, [output]);
return (
<ToolCallExpandable
label={<span className="pl-4 py-1 font-sans text-sm">{intl.formatMessage(i18n.output)}</span>}
isStartExpanded={true}
>
<div ref={outputRef} className="max-h-[20rem] overflow-y-auto px-4 py-3">
<pre className="font-mono text-xs text-textSubtle whitespace-pre-wrap break-words">
{output}
</pre>
</div>
</ToolCallExpandable>
);
}
function ToolResultView({ result, isStartExpanded }: ToolResultViewProps) {
const intl = useIntl();
const hasText = (c: ContentBlock): c is ContentBlock & { text: string } =>
ui/desktop/src/types/message.ts
+11
View File
@@ -258,6 +258,17 @@ export type ToolConfirmationRequestContent = ToolConfirmationRequest & {
};
export type NotificationEvent = Extract<MessageEvent, { type: 'Notification' }>;
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;