Timing information for toolcalls and turns in cli

This commit is contained in:
Dhanji Prasanna
2025-10-13 15:34:02 +11:00
parent 224d486c6f
commit c7f38061ff
3 changed files with 138 additions and 5 deletions
+36 -1
View File
@@ -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<u32>,
edit_mode: Option<EditMode>,
retry_config: Option<RetryConfig>,
tool_call_timings: Arc<Mutex<HashMap<String, Instant>>>,
}
// 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<String> = 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)))) => {
+62 -4
View File
@@ -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<String, std::time::Duration>) {
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<std::time::Duration>) {
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::*;
@@ -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");
}
}