diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 9eab36b1..d980ba7c 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -16,8 +16,7 @@ pub use goose::session::Identifier; use anyhow::{Context, Result}; use completion::GooseCompleter; -use etcetera::choose_app_strategy; -use etcetera::AppStrategy; +use etcetera::{choose_app_strategy, AppStrategy}; use goose::agents::extension::{Envs, ExtensionConfig}; use goose::agents::{Agent, SessionConfig}; use goose::config::Config; @@ -26,9 +25,9 @@ use goose::session; use input::InputResult; use mcp_core::handler::ToolError; use mcp_core::prompt::PromptMessage; - use mcp_core::protocol::JsonRpcMessage; use mcp_core::protocol::JsonRpcNotification; + use rand::{distributions::Alphanumeric, Rng}; use serde_json::Value; use std::collections::HashMap; @@ -354,9 +353,10 @@ impl Session { // Create and use a global history file in ~/.config/goose directory // This allows command history to persist across different chat sessions // instead of being tied to each individual session's messages - let history_file = choose_app_strategy(crate::APP_STRATEGY.clone()) - .expect("goose requires a home dir") - .in_config_dir("history.txt"); + let strategy = + choose_app_strategy(crate::APP_STRATEGY.clone()).expect("goose requires a home dir"); + let config_dir = strategy.config_dir(); + let history_file = config_dir.join("history.txt"); // Ensure config directory exists if let Some(parent) = history_file.parent() { @@ -382,6 +382,9 @@ impl Session { output::display_greeting(); loop { + // Display context usage before each prompt + self.display_context_usage().await?; + match input::get_input(&mut editor)? { input::InputResult::Message(content) => { match self.run_mode { @@ -1118,6 +1121,26 @@ impl Session { Ok(metadata.total_tokens) } + /// Display enhanced context usage with session totals + pub async fn display_context_usage(&self) -> Result<()> { + let provider = self.agent.provider().await?; + let model_config = provider.get_model_config(); + let context_limit = model_config.context_limit.unwrap_or(32000); + + match self.get_metadata() { + Ok(metadata) => { + let total_tokens = metadata.total_tokens.unwrap_or(0) as usize; + + output::display_context_usage(total_tokens, context_limit); + } + Err(_) => { + output::display_context_usage(0, context_limit); + } + } + + Ok(()) + } + /// Handle prompt command execution async fn handle_prompt_command(&mut self, opts: input::PromptCommandOptions) -> Result<()> { // name is required diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 78b4eb76..3bda532b 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -574,6 +574,38 @@ pub fn display_greeting() { println!("\nGoose is running! Enter your instructions, or try asking what goose can do.\n"); } +/// Display context window usage with both current and session totals +pub fn display_context_usage(total_tokens: usize, context_limit: usize) { + use console::style; + + // Calculate percentage used + let percentage = (total_tokens as f64 / context_limit as f64 * 100.0).round() as usize; + + // Create dot visualization + let dot_count = 10; + let filled_dots = ((percentage as f64 / 100.0) * dot_count as f64).round() as usize; + let empty_dots = dot_count - filled_dots; + + let filled = "●".repeat(filled_dots); + let empty = "○".repeat(empty_dots); + + // Combine dots and apply color + let dots = format!("{}{}", filled, empty); + let colored_dots = if percentage < 50 { + style(dots).green() + } else if percentage < 85 { + style(dots).yellow() + } else { + style(dots).red() + }; + + // Print the status line + println!( + "Context: {} {}% ({}/{} tokens)", + colored_dots, percentage, total_tokens, context_limit + ); +} + pub struct McpSpinners { bars: HashMap, log_spinner: Option,