diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 0547282c67..eb024d28cd 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -3,6 +3,7 @@ mod completion; mod export; mod input; mod output; +mod output_timing_tests; mod prompt; mod task_execution_display; mod thinking; @@ -44,7 +45,10 @@ use serde_json::Value; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex; use std::time::Instant; +use console; + use tokio; use tokio_util::sync::CancellationToken; @@ -64,6 +68,7 @@ pub struct CliSession { max_turns: Option, edit_mode: Option, retry_config: Option, + tool_call_timings: Arc>>, } // Cache structure for completion data @@ -151,6 +156,7 @@ impl CliSession { max_turns, edit_mode, retry_config, + tool_call_timings: Arc::new(Mutex::new(HashMap::new())), } } @@ -1023,6 +1029,10 @@ impl CliSession { tool_name = %tool_call.name, "Tool call started" ); + // Record the start time for this tool call + if let Ok(mut timings) = self.tool_call_timings.lock() { + timings.insert(tool_request.id.clone(), Instant::now()); + } } } if let MessageContent::ToolResponse(tool_response) = content { @@ -1062,7 +1072,32 @@ impl CliSession { if interactive {output::hide_thinking()}; let _ = progress_bars.hide(); - output::render_message(&message, self.debug); + + // Collect timing information for this message and render with timing + let mut tool_timings = HashMap::new(); + + // First pass: collect all tool response IDs and their timings + let tool_response_ids: Vec = message.content.iter() + .filter_map(|content| { + if let MessageContent::ToolResponse(resp) = content { + Some(resp.id.clone()) + } else { + None + } + }) + .collect(); + + // Get timing data for all tool responses in this message + if let Ok(mut timings) = self.tool_call_timings.lock() { + for tool_id in tool_response_ids { + if let Some(start_time) = timings.remove(&tool_id) { + let duration = start_time.elapsed(); + tool_timings.insert(tool_id.clone(), duration); + } + } + } + + output::render_message_with_timing(&message, self.debug, &tool_timings); } } Some(Ok(AgentEvent::McpNotification((_id, message)))) => { diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 639d7da961..af3cf86d9d 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -161,6 +161,44 @@ pub fn set_thinking_message(s: &String) { } } +pub fn render_message_with_timing(message: &Message, debug: bool, tool_timings: &std::collections::HashMap) { + let theme = get_theme(); + + for content in &message.content { + match content { + MessageContent::Text(text) => print_markdown(&text.text, theme), + MessageContent::ToolRequest(req) => render_tool_request(req, theme, debug), + MessageContent::ToolResponse(resp) => { + let timing = tool_timings.get(&resp.id).copied(); + render_tool_response(resp, theme, debug, timing); + } + MessageContent::Image(image) => { + println!("Image: [data: {}, type: {}]", image.data, image.mime_type); + } + MessageContent::Thinking(thinking) => { + if std::env::var("GOOSE_CLI_SHOW_THINKING").is_ok() + && std::io::stdout().is_terminal() + { + println!("\n{}", console::style("Thinking:").dim().italic()); + print_markdown(&thinking.thinking, theme); + } + } + MessageContent::RedactedThinking(_) => { + println!("\n{}", console::style("Thinking:").dim().italic()); + print_markdown("Thinking was redacted", theme); + } + MessageContent::SummarizationRequested(summarization) => { + println!("\n{}", console::style(&summarization.msg).yellow()); + } + _ => { + println!("WARNING: Message content type could not be rendered"); + } + } + } + + let _ = std::io::stdout().flush(); +} + pub fn render_message(message: &Message, debug: bool) { let theme = get_theme(); @@ -168,7 +206,7 @@ pub fn render_message(message: &Message, debug: bool) { match content { MessageContent::Text(text) => print_markdown(&text.text, theme), MessageContent::ToolRequest(req) => render_tool_request(req, theme, debug), - MessageContent::ToolResponse(resp) => render_tool_response(resp, theme, debug), + MessageContent::ToolResponse(resp) => render_tool_response(resp, theme, debug, None), MessageContent::Image(image) => { println!("Image: [data: {}, type: {}]", image.data, image.mime_type); } @@ -245,7 +283,7 @@ pub fn goose_mode_message(text: &str) { println!("\n{}", style(text).yellow(),); } -fn render_tool_request(req: &ToolRequest, theme: Theme, debug: bool) { +pub fn render_tool_request(req: &ToolRequest, theme: Theme, debug: bool) { match &req.tool_call { Ok(call) => match call.name.to_string().as_str() { "developer__text_editor" => render_text_editor_request(call, debug), @@ -258,7 +296,7 @@ fn render_tool_request(req: &ToolRequest, theme: Theme, debug: bool) { } } -fn render_tool_response(resp: &ToolResponse, theme: Theme, debug: bool) { +pub fn render_tool_response(resp: &ToolResponse, theme: Theme, debug: bool, timing: Option) { let config = Config::global(); match &resp.tool_result { @@ -292,6 +330,13 @@ fn render_tool_response(resp: &ToolResponse, theme: Theme, debug: bool) { } Err(e) => print_markdown(&e.to_string(), theme), } + + // Display timing information if available + if let Some(duration) = timing { + let elapsed_str = format_elapsed_time(duration); + println!("{}", console::style(format!("⚡️ {}", elapsed_str)).dim()); + println!(); // Add extra newline for spacing + } } pub fn render_error(message: &str) { @@ -539,7 +584,7 @@ pub fn env_no_color() -> bool { std::env::var_os("NO_COLOR").is_none() } -fn print_markdown(content: &str, theme: Theme) { +pub fn print_markdown(content: &str, theme: Theme) { if std::io::stdout().is_terminal() { bat::PrettyPrinter::new() .input(bat::Input::from_bytes(content.as_bytes())) @@ -941,6 +986,19 @@ impl McpSpinners { } } +/// Format elapsed time duration +/// Shows seconds if less than 60, otherwise shows minutes:seconds +pub fn format_elapsed_time(duration: std::time::Duration) -> String { + let total_secs = duration.as_secs(); + if total_secs < 60 { + format!("{:.2}s", duration.as_secs_f64()) + } else { + let minutes = total_secs / 60; + let seconds = total_secs % 60; + format!("{}m {:02}s", minutes, seconds) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose-cli/src/session/output_timing_tests.rs b/crates/goose-cli/src/session/output_timing_tests.rs new file mode 100644 index 0000000000..c48ead37a9 --- /dev/null +++ b/crates/goose-cli/src/session/output_timing_tests.rs @@ -0,0 +1,40 @@ +#[cfg(test)] +mod tests { + use std::time::Duration; + use std::collections::HashMap; + use goose::conversation::message::{Message, MessageContent, ToolResponse}; + use rmcp::model::Content; + use crate::session::output::{render_message_with_timing, format_elapsed_time}; + + #[test] + fn test_render_tool_response_with_timing() { + // Create a mock tool response + let tool_response = ToolResponse { + id: "test-tool-123".to_string(), + tool_result: Ok(vec![Content::text("Test output")]), + }; + + // Create timing information + let mut tool_timings = HashMap::new(); + tool_timings.insert("test-tool-123".to_string(), Duration::from_millis(1500)); + + // Create a message with the tool response + let message = Message::assistant().with_content(MessageContent::ToolResponse(tool_response)); + + // This test mainly verifies the function doesn't panic and compiles correctly + // In a real test environment, we'd capture stdout to verify the timing display + render_message_with_timing(&message, false, &tool_timings); + } + + #[test] + fn test_format_elapsed_time_function_exists() { + // Test that our format_elapsed_time function works correctly + let duration = Duration::from_millis(1500); + let formatted = format_elapsed_time(duration); + assert_eq!(formatted, "1.50s"); + + let duration = Duration::from_secs(75); + let formatted = format_elapsed_time(duration); + assert_eq!(formatted, "1m 15s"); + } +} \ No newline at end of file