From d7bcc92007f63cf38fa6b428a625c099d8530fba Mon Sep 17 00:00:00 2001 From: Affaan Mustafa Date: Mon, 23 Mar 2026 03:46:25 -0700 Subject: [PATCH] feat(ecc2): add token/cost meter widget (#775) - TokenMeter widget using ratatui Gauge with color gradient (green->yellow->red) - Budget fields (cost_budget_usd, token_budget) in Config - Aggregate cost display in status bar - Warning state at 80%+ budget consumption - Tests for gradient, config fallback, and meter rendering --- ecc2/src/comms/mod.rs | 5 +- ecc2/src/config/mod.rs | 38 +++++ ecc2/src/main.rs | 15 +- ecc2/src/session/store.rs | 1 - ecc2/src/tui/dashboard.rs | 273 ++++++++++++++++++++++++++++++++---- ecc2/src/tui/widgets.rs | 287 +++++++++++++++++++++++++++++++++++++- 6 files changed, 577 insertions(+), 42 deletions(-) diff --git a/ecc2/src/comms/mod.rs b/ecc2/src/comms/mod.rs index be176e96..8be89f2b 100644 --- a/ecc2/src/comms/mod.rs +++ b/ecc2/src/comms/mod.rs @@ -13,7 +13,10 @@ pub enum MessageType { /// Response to a query Response { answer: String }, /// Notification of completion - Completed { summary: String, files_changed: Vec }, + Completed { + summary: String, + files_changed: Vec, + }, /// Conflict detected (e.g., two agents editing the same file) Conflict { file: String, description: String }, } diff --git a/ecc2/src/config/mod.rs b/ecc2/src/config/mod.rs index 1e7eeab7..c6fe807d 100644 --- a/ecc2/src/config/mod.rs +++ b/ecc2/src/config/mod.rs @@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize}; use std::path::PathBuf; #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] pub struct Config { pub db_path: PathBuf, pub worktree_root: PathBuf, @@ -11,6 +12,8 @@ pub struct Config { pub session_timeout_secs: u64, pub heartbeat_interval_secs: u64, pub default_agent: String, + pub cost_budget_usd: f64, + pub token_budget: u64, pub theme: Theme, } @@ -31,6 +34,8 @@ impl Default for Config { session_timeout_secs: 3600, heartbeat_interval_secs: 30, default_agent: "claude".to_string(), + cost_budget_usd: 10.0, + token_budget: 500_000, theme: Theme::Dark, } } @@ -52,3 +57,36 @@ impl Config { } } } + +#[cfg(test)] +mod tests { + use super::Config; + + #[test] + fn default_includes_positive_budget_thresholds() { + let config = Config::default(); + + assert!(config.cost_budget_usd > 0.0); + assert!(config.token_budget > 0); + } + + #[test] + fn missing_budget_fields_fall_back_to_defaults() { + let legacy_config = r#" +db_path = "/tmp/ecc2.db" +worktree_root = "/tmp/ecc-worktrees" +max_parallel_sessions = 8 +max_parallel_worktrees = 6 +session_timeout_secs = 3600 +heartbeat_interval_secs = 30 +default_agent = "claude" +theme = "Dark" +"#; + + let config: Config = toml::from_str(legacy_config).unwrap(); + let defaults = Config::default(); + + assert_eq!(config.cost_budget_usd, defaults.cost_budget_usd); + assert_eq!(config.token_budget, defaults.token_budget); + } +} diff --git a/ecc2/src/main.rs b/ecc2/src/main.rs index 850b7b49..afa50a2f 100644 --- a/ecc2/src/main.rs +++ b/ecc2/src/main.rs @@ -1,9 +1,9 @@ +mod comms; mod config; +mod observability; mod session; mod tui; mod worktree; -mod observability; -mod comms; use anyhow::Result; use clap::Parser; @@ -63,10 +63,13 @@ async fn main() -> Result<()> { Some(Commands::Dashboard) | None => { tui::app::run(db, cfg).await?; } - Some(Commands::Start { task, agent, worktree: use_worktree }) => { - let session_id = session::manager::create_session( - &db, &cfg, &task, &agent, use_worktree, - ).await?; + Some(Commands::Start { + task, + agent, + worktree: use_worktree, + }) => { + let session_id = + session::manager::create_session(&db, &cfg, &task, &agent, use_worktree).await?; println!("Session started: {session_id}"); } Some(Commands::Sessions) => { diff --git a/ecc2/src/session/store.rs b/ecc2/src/session/store.rs index 499141dd..60d2a5b2 100644 --- a/ecc2/src/session/store.rs +++ b/ecc2/src/session/store.rs @@ -236,7 +236,6 @@ impl StateStore { .into_iter() .find(|s| s.id == id || s.id.starts_with(id))) } - pub fn send_message(&self, from: &str, to: &str, content: &str, msg_type: &str) -> Result<()> { self.conn.execute( "INSERT INTO messages (from_session, to_session, content, msg_type, timestamp) diff --git a/ecc2/src/tui/dashboard.rs b/ecc2/src/tui/dashboard.rs index aca1e995..42b41b84 100644 --- a/ecc2/src/tui/dashboard.rs +++ b/ecc2/src/tui/dashboard.rs @@ -1,11 +1,12 @@ use ratatui::{ prelude::*, - widgets::{Block, Borders, List, ListItem, Paragraph, Tabs}, + widgets::{Block, Borders, List, ListItem, Paragraph, Tabs, Wrap}, }; +use super::widgets::{budget_state, format_currency, format_token_count, BudgetState, TokenMeter}; use crate::config::Config; -use crate::session::{Session, SessionState}; use crate::session::store::StateStore; +use crate::session::{Session, SessionState}; pub struct Dashboard { db: StateStore, @@ -24,6 +25,15 @@ enum Pane { Metrics, } +#[derive(Debug, Clone, Copy)] +struct AggregateUsage { + total_tokens: u64, + total_cost_usd: f64, + token_state: BudgetState, + cost_state: BudgetState, + overall_state: BudgetState, +} + impl Dashboard { pub fn new(db: StateStore, cfg: Config) -> Self { let sessions = db.list_sessions().unwrap_or_default(); @@ -42,7 +52,7 @@ impl Dashboard { let chunks = Layout::default() .direction(Direction::Vertical) .constraints([ - Constraint::Length(3), // Header + Constraint::Length(3), // Header Constraint::Min(10), // Main content Constraint::Length(3), // Status bar ]) @@ -79,7 +89,11 @@ impl Dashboard { } fn render_header(&self, frame: &mut Frame, area: Rect) { - let running = self.sessions.iter().filter(|s| s.state == SessionState::Running).count(); + let running = self + .sessions + .iter() + .filter(|s| s.state == SessionState::Running) + .count(); let total = self.sessions.len(); let title = format!(" ECC 2.0 | {running} running / {total} total "); @@ -90,7 +104,11 @@ impl Dashboard { Pane::Output => 1, Pane::Metrics => 2, }) - .highlight_style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD)); + .highlight_style( + Style::default() + .fg(Color::Cyan) + .add_modifier(Modifier::BOLD), + ); frame.render_widget(tabs, area); } @@ -110,11 +128,18 @@ impl Dashboard { SessionState::Pending => "◌", }; let style = if i == self.selected_session { - Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD) + Style::default() + .fg(Color::Cyan) + .add_modifier(Modifier::BOLD) } else { Style::default() }; - let text = format!("{state_icon} {} [{}] {}", &s.id[..8.min(s.id.len())], s.agent_type, s.task); + let text = format!( + "{state_icon} {} [{}] {}", + &s.id[..8.min(s.id.len())], + s.agent_type, + s.task + ); ListItem::new(text).style(style) }) .collect(); @@ -136,7 +161,10 @@ impl Dashboard { fn render_output(&self, frame: &mut Frame, area: Rect) { let content = if let Some(session) = self.sessions.get(self.selected_session) { - format!("Agent output for session {}...\n\n(Live streaming coming soon)", session.id) + format!( + "Agent output for session {}...\n\n(Live streaming coming soon)", + session.id + ) } else { "No sessions. Press 'n' to start one.".to_string() }; @@ -157,37 +185,87 @@ impl Dashboard { } fn render_metrics(&self, frame: &mut Frame, area: Rect) { - let content = if let Some(session) = self.sessions.get(self.selected_session) { - let m = &session.metrics; - format!( - "Tokens: {} | Tools: {} | Files: {} | Cost: ${:.4} | Duration: {}s", - m.tokens_used, m.tool_calls, m.files_changed, m.cost_usd, m.duration_secs - ) - } else { - "No metrics available".to_string() - }; - let border_style = if self.selected_pane == Pane::Metrics { Style::default().fg(Color::Cyan) } else { Style::default() }; - let paragraph = Paragraph::new(content).block( - Block::default() - .borders(Borders::ALL) - .title(" Metrics ") - .border_style(border_style), + let block = Block::default() + .borders(Borders::ALL) + .title(" Metrics ") + .border_style(border_style); + let inner = block.inner(area); + frame.render_widget(block, area); + + if inner.is_empty() { + return; + } + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Length(2), + Constraint::Length(2), + Constraint::Min(1), + ]) + .split(inner); + + let aggregate = self.aggregate_usage(); + frame.render_widget( + TokenMeter::tokens( + "Token Budget", + aggregate.total_tokens, + self.cfg.token_budget, + ), + chunks[0], + ); + frame.render_widget( + TokenMeter::currency( + "Cost Budget", + aggregate.total_cost_usd, + self.cfg.cost_budget_usd, + ), + chunks[1], + ); + frame.render_widget( + Paragraph::new(self.selected_session_metrics_text()).wrap(Wrap { trim: true }), + chunks[2], ); - frame.render_widget(paragraph, area); } fn render_status_bar(&self, frame: &mut Frame, area: Rect) { let text = " [n]ew session [s]top [Tab] switch pane [j/k] scroll [?] help [q]uit "; - let paragraph = Paragraph::new(text) - .style(Style::default().fg(Color::DarkGray)) - .block(Block::default().borders(Borders::ALL)); - frame.render_widget(paragraph, area); + let aggregate = self.aggregate_usage(); + let (summary_text, summary_style) = self.aggregate_cost_summary(); + let block = Block::default() + .borders(Borders::ALL) + .border_style(aggregate.overall_state.style()); + let inner = block.inner(area); + frame.render_widget(block, area); + + if inner.is_empty() { + return; + } + + let summary_width = summary_text + .len() + .min(inner.width.saturating_sub(1) as usize) as u16; + let chunks = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Min(1), Constraint::Length(summary_width)]) + .split(inner); + + frame.render_widget( + Paragraph::new(text).style(Style::default().fg(Color::DarkGray)), + chunks[0], + ); + frame.render_widget( + Paragraph::new(summary_text) + .style(summary_style) + .alignment(Alignment::Right), + chunks[1], + ); } fn render_help(&self, frame: &mut Frame, area: Rect) { @@ -270,4 +348,143 @@ impl Dashboard { // Periodic refresh every few ticks self.sessions = self.db.list_sessions().unwrap_or_default(); } + + fn aggregate_usage(&self) -> AggregateUsage { + let total_tokens = self + .sessions + .iter() + .map(|session| session.metrics.tokens_used) + .sum(); + let total_cost_usd = self + .sessions + .iter() + .map(|session| session.metrics.cost_usd) + .sum::(); + let token_state = budget_state(total_tokens as f64, self.cfg.token_budget as f64); + let cost_state = budget_state(total_cost_usd, self.cfg.cost_budget_usd); + + AggregateUsage { + total_tokens, + total_cost_usd, + token_state, + cost_state, + overall_state: token_state.max(cost_state), + } + } + + fn selected_session_metrics_text(&self) -> String { + if let Some(session) = self.sessions.get(self.selected_session) { + let metrics = &session.metrics; + format!( + "Selected {} [{}]\nTokens {} | Tools {} | Files {}\nCost ${:.4} | Duration {}s", + &session.id[..8.min(session.id.len())], + session.state, + format_token_count(metrics.tokens_used), + metrics.tool_calls, + metrics.files_changed, + metrics.cost_usd, + metrics.duration_secs + ) + } else { + "No metrics available".to_string() + } + } + + fn aggregate_cost_summary(&self) -> (String, Style) { + let aggregate = self.aggregate_usage(); + let mut text = if self.cfg.cost_budget_usd > 0.0 { + format!( + "Aggregate cost {} / {}", + format_currency(aggregate.total_cost_usd), + format_currency(self.cfg.cost_budget_usd), + ) + } else { + format!( + "Aggregate cost {} (no budget)", + format_currency(aggregate.total_cost_usd) + ) + }; + + match aggregate.overall_state { + BudgetState::Warning => text.push_str(" | Budget warning"), + BudgetState::OverBudget => text.push_str(" | Budget exceeded"), + _ => {} + } + + (text, aggregate.overall_state.style()) + } + + fn aggregate_cost_summary_text(&self) -> String { + self.aggregate_cost_summary().0 + } +} + +#[cfg(test)] +mod tests { + use std::path::Path; + + use chrono::Utc; + + use super::Dashboard; + use crate::config::Config; + use crate::session::store::StateStore; + use crate::session::{Session, SessionMetrics, SessionState}; + use crate::tui::widgets::BudgetState; + + #[test] + fn aggregate_usage_sums_tokens_and_cost_with_warning_state() { + let db = StateStore::open(Path::new(":memory:")).unwrap(); + let mut cfg = Config::default(); + cfg.token_budget = 10_000; + cfg.cost_budget_usd = 10.0; + + let mut dashboard = Dashboard::new(db, cfg); + dashboard.sessions = vec![ + session("sess-1", 4_000, 3.50), + session("sess-2", 4_500, 4.80), + ]; + + let aggregate = dashboard.aggregate_usage(); + + assert_eq!(aggregate.total_tokens, 8_500); + assert!((aggregate.total_cost_usd - 8.30).abs() < 1e-9); + assert_eq!(aggregate.token_state, BudgetState::Warning); + assert_eq!(aggregate.cost_state, BudgetState::Warning); + assert_eq!(aggregate.overall_state, BudgetState::Warning); + } + + #[test] + fn aggregate_cost_summary_mentions_total_cost() { + let db = StateStore::open(Path::new(":memory:")).unwrap(); + let mut cfg = Config::default(); + cfg.cost_budget_usd = 10.0; + + let mut dashboard = Dashboard::new(db, cfg); + dashboard.sessions = vec![session("sess-1", 3_500, 8.25)]; + + assert_eq!( + dashboard.aggregate_cost_summary_text(), + "Aggregate cost $8.25 / $10.00 | Budget warning" + ); + } + + fn session(id: &str, tokens_used: u64, cost_usd: f64) -> Session { + let now = Utc::now(); + Session { + id: id.to_string(), + task: "Budget tracking".to_string(), + agent_type: "claude".to_string(), + state: SessionState::Running, + worktree: None, + created_at: now, + updated_at: now, + metrics: SessionMetrics { + tokens_used, + tool_calls: 0, + files_changed: 0, + duration_secs: 0, + cost_usd, + }, + } + } } diff --git a/ecc2/src/tui/widgets.rs b/ecc2/src/tui/widgets.rs index 604d6a0e..784e4b50 100644 --- a/ecc2/src/tui/widgets.rs +++ b/ecc2/src/tui/widgets.rs @@ -1,6 +1,281 @@ -// Custom TUI widgets for ECC 2.0 -// TODO: Implement custom widgets: -// - TokenMeter: visual token usage bar with budget threshold -// - DiffViewer: side-by-side syntax-highlighted diff display -// - ProgressTimeline: session timeline with tool call markers -// - AgentTree: hierarchical view of parent/child agent sessions +use ratatui::{ + prelude::*, + text::{Line, Span}, + widgets::{Gauge, Paragraph, Widget}, +}; + +pub(crate) const WARNING_THRESHOLD: f64 = 0.8; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum BudgetState { + Unconfigured, + Normal, + Warning, + OverBudget, +} + +impl BudgetState { + pub(crate) const fn is_warning(self) -> bool { + matches!(self, Self::Warning | Self::OverBudget) + } + + fn badge(self) -> Option<&'static str> { + match self { + Self::Warning => Some("warning"), + Self::OverBudget => Some("over budget"), + Self::Unconfigured => Some("no budget"), + Self::Normal => None, + } + } + + pub(crate) fn style(self) -> Style { + let base = Style::default().fg(match self { + Self::Unconfigured => Color::DarkGray, + Self::Normal => Color::DarkGray, + Self::Warning => Color::Yellow, + Self::OverBudget => Color::Red, + }); + + if self.is_warning() { + base.add_modifier(Modifier::BOLD) + } else { + base + } + } +} + +#[derive(Debug, Clone, Copy)] +enum MeterFormat { + Tokens, + Currency, +} + +#[derive(Debug, Clone)] +pub(crate) struct TokenMeter<'a> { + title: &'a str, + used: f64, + budget: f64, + format: MeterFormat, +} + +impl<'a> TokenMeter<'a> { + pub(crate) fn tokens(title: &'a str, used: u64, budget: u64) -> Self { + Self { + title, + used: used as f64, + budget: budget as f64, + format: MeterFormat::Tokens, + } + } + + pub(crate) fn currency(title: &'a str, used: f64, budget: f64) -> Self { + Self { + title, + used, + budget, + format: MeterFormat::Currency, + } + } + + pub(crate) fn state(&self) -> BudgetState { + budget_state(self.used, self.budget) + } + + fn ratio(&self) -> f64 { + budget_ratio(self.used, self.budget) + } + + fn clamped_ratio(&self) -> f64 { + self.ratio().clamp(0.0, 1.0) + } + + fn title_line(&self) -> Line<'static> { + let mut spans = vec![Span::styled( + self.title.to_string(), + Style::default() + .fg(Color::Gray) + .add_modifier(Modifier::BOLD), + )]; + + if let Some(badge) = self.state().badge() { + spans.push(Span::raw(" ")); + spans.push(Span::styled(format!("[{badge}]"), self.state().style())); + } + + Line::from(spans) + } + + fn display_label(&self) -> String { + if self.budget <= 0.0 { + return match self.format { + MeterFormat::Tokens => format!("{} tok used | no budget", self.used_label()), + MeterFormat::Currency => format!("{} spent | no budget", self.used_label()), + }; + } + + format!( + "{} / {}{} ({}%)", + self.used_label(), + self.budget_label(), + self.unit_suffix(), + (self.ratio() * 100.0).round() as u64 + ) + } + + fn used_label(&self) -> String { + match self.format { + MeterFormat::Tokens => format_token_count(self.used.max(0.0).round() as u64), + MeterFormat::Currency => format_currency(self.used.max(0.0)), + } + } + + fn budget_label(&self) -> String { + match self.format { + MeterFormat::Tokens => format_token_count(self.budget.max(0.0).round() as u64), + MeterFormat::Currency => format_currency(self.budget.max(0.0)), + } + } + + fn unit_suffix(&self) -> &'static str { + match self.format { + MeterFormat::Tokens => " tok", + MeterFormat::Currency => "", + } + } +} + +impl Widget for TokenMeter<'_> { + fn render(self, area: Rect, buf: &mut Buffer) { + if area.is_empty() { + return; + } + + let mut gauge_area = area; + if area.height > 1 { + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Length(1), Constraint::Min(1)]) + .split(area); + Paragraph::new(self.title_line()).render(chunks[0], buf); + gauge_area = chunks[1]; + } + + Gauge::default() + .ratio(self.clamped_ratio()) + .label(self.display_label()) + .gauge_style( + Style::default() + .fg(gradient_color(self.ratio())) + .add_modifier(Modifier::BOLD), + ) + .style(Style::default().fg(Color::DarkGray)) + .use_unicode(true) + .render(gauge_area, buf); + } +} + +pub(crate) fn budget_ratio(used: f64, budget: f64) -> f64 { + if budget <= 0.0 { + 0.0 + } else { + used / budget + } +} + +pub(crate) fn budget_state(used: f64, budget: f64) -> BudgetState { + if budget <= 0.0 { + BudgetState::Unconfigured + } else if used / budget >= 1.0 { + BudgetState::OverBudget + } else if used / budget >= WARNING_THRESHOLD { + BudgetState::Warning + } else { + BudgetState::Normal + } +} + +pub(crate) fn gradient_color(ratio: f64) -> Color { + const GREEN: (u8, u8, u8) = (34, 197, 94); + const YELLOW: (u8, u8, u8) = (234, 179, 8); + const RED: (u8, u8, u8) = (239, 68, 68); + + let clamped = ratio.clamp(0.0, 1.0); + if clamped <= WARNING_THRESHOLD { + interpolate_rgb(GREEN, YELLOW, clamped / WARNING_THRESHOLD) + } else { + interpolate_rgb( + YELLOW, + RED, + (clamped - WARNING_THRESHOLD) / (1.0 - WARNING_THRESHOLD), + ) + } +} + +pub(crate) fn format_currency(value: f64) -> String { + format!("${value:.2}") +} + +pub(crate) fn format_token_count(value: u64) -> String { + let digits = value.to_string(); + let mut formatted = String::with_capacity(digits.len() + digits.len() / 3); + + for (index, ch) in digits.chars().rev().enumerate() { + if index != 0 && index % 3 == 0 { + formatted.push(','); + } + formatted.push(ch); + } + + formatted.chars().rev().collect() +} + +fn interpolate_rgb(from: (u8, u8, u8), to: (u8, u8, u8), ratio: f64) -> Color { + let ratio = ratio.clamp(0.0, 1.0); + let channel = |start: u8, end: u8| -> u8 { + (f64::from(start) + (f64::from(end) - f64::from(start)) * ratio).round() as u8 + }; + + Color::Rgb( + channel(from.0, to.0), + channel(from.1, to.1), + channel(from.2, to.2), + ) +} + +#[cfg(test)] +mod tests { + use ratatui::{buffer::Buffer, layout::Rect, style::Color, widgets::Widget}; + + use super::{gradient_color, BudgetState, TokenMeter}; + + #[test] + fn warning_state_starts_at_eighty_percent() { + let meter = TokenMeter::tokens("Token Budget", 80, 100); + + assert_eq!(meter.state(), BudgetState::Warning); + } + + #[test] + fn gradient_runs_from_green_to_yellow_to_red() { + assert_eq!(gradient_color(0.0), Color::Rgb(34, 197, 94)); + assert_eq!(gradient_color(0.8), Color::Rgb(234, 179, 8)); + assert_eq!(gradient_color(1.0), Color::Rgb(239, 68, 68)); + } + + #[test] + fn token_meter_renders_compact_usage_label() { + let meter = TokenMeter::tokens("Token Budget", 4_000, 10_000); + let area = Rect::new(0, 0, 48, 2); + let mut buffer = Buffer::empty(area); + + meter.render(area, &mut buffer); + + let rendered = buffer + .content() + .chunks(area.width as usize) + .flat_map(|row| row.iter().map(|cell| cell.symbol())) + .collect::(); + + assert!(rendered.contains("4,000 / 10,000 tok (40%)")); + } +}