Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 29 additions & 6 deletions crates/goose-cli/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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() {
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
32 changes: 32 additions & 0 deletions crates/goose-cli/src/session/output.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, ProgressBar>,
log_spinner: Option<ProgressBar>,
Expand Down
Loading