diff --git a/ecc2/src/config/mod.rs b/ecc2/src/config/mod.rs index 86b542dc..694e6a6d 100644 --- a/ecc2/src/config/mod.rs +++ b/ecc2/src/config/mod.rs @@ -20,6 +20,14 @@ pub struct RiskThresholds { pub block: f64, } +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct BudgetAlertThresholds { + pub advisory: f64, + pub warning: f64, + pub critical: f64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default)] pub struct Config { @@ -36,6 +44,7 @@ pub struct Config { pub auto_merge_ready_worktrees: bool, pub cost_budget_usd: f64, pub token_budget: u64, + pub budget_alert_thresholds: BudgetAlertThresholds, pub theme: Theme, pub pane_layout: PaneLayout, pub pane_navigation: PaneNavigationConfig, @@ -89,6 +98,7 @@ impl Default for Config { auto_merge_ready_worktrees: false, cost_budget_usd: 10.0, token_budget: 500_000, + budget_alert_thresholds: Self::BUDGET_ALERT_THRESHOLDS, theme: Theme::Dark, pane_layout: PaneLayout::Horizontal, pane_navigation: PaneNavigationConfig::default(), @@ -106,6 +116,12 @@ impl Config { block: 0.85, }; + pub const BUDGET_ALERT_THRESHOLDS: BudgetAlertThresholds = BudgetAlertThresholds { + advisory: 0.50, + warning: 0.75, + critical: 0.90, + }; + pub fn config_path() -> PathBuf { dirs::home_dir() .unwrap_or_else(|| PathBuf::from(".")) @@ -121,6 +137,10 @@ impl Config { .join("costs.jsonl") } + pub fn effective_budget_alert_thresholds(&self) -> BudgetAlertThresholds { + self.budget_alert_thresholds.sanitized() + } + pub fn load() -> Result { let config_path = Self::config_path(); @@ -265,9 +285,32 @@ impl Default for RiskThresholds { } } +impl Default for BudgetAlertThresholds { + fn default() -> Self { + Config::BUDGET_ALERT_THRESHOLDS + } +} + +impl BudgetAlertThresholds { + pub fn sanitized(self) -> Self { + let values = [self.advisory, self.warning, self.critical]; + let valid = values.into_iter().all(f64::is_finite) + && self.advisory > 0.0 + && self.advisory < self.warning + && self.warning < self.critical + && self.critical < 1.0; + + if valid { + self + } else { + Self::default() + } + } +} + #[cfg(test)] mod tests { - use super::{Config, PaneLayout}; + use super::{BudgetAlertThresholds, Config, PaneLayout}; use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; use uuid::Uuid; @@ -297,6 +340,10 @@ theme = "Dark" assert_eq!(config.cost_budget_usd, defaults.cost_budget_usd); assert_eq!(config.token_budget, defaults.token_budget); + assert_eq!( + config.budget_alert_thresholds, + defaults.budget_alert_thresholds + ); assert_eq!(config.pane_layout, defaults.pane_layout); assert_eq!(config.pane_navigation, defaults.pane_navigation); assert_eq!( @@ -412,6 +459,58 @@ move_right = "d" assert_eq!(Config::default().risk_thresholds, Config::RISK_THRESHOLDS); } + #[test] + fn default_budget_alert_thresholds_are_applied() { + assert_eq!( + Config::default().budget_alert_thresholds, + Config::BUDGET_ALERT_THRESHOLDS + ); + } + + #[test] + fn budget_alert_thresholds_deserialize_from_toml() { + let config: Config = toml::from_str( + r#" +[budget_alert_thresholds] +advisory = 0.40 +warning = 0.70 +critical = 0.85 +"#, + ) + .unwrap(); + + assert_eq!( + config.budget_alert_thresholds, + BudgetAlertThresholds { + advisory: 0.40, + warning: 0.70, + critical: 0.85, + } + ); + assert_eq!( + config.effective_budget_alert_thresholds(), + config.budget_alert_thresholds + ); + } + + #[test] + fn invalid_budget_alert_thresholds_fall_back_to_defaults() { + let config: Config = toml::from_str( + r#" +[budget_alert_thresholds] +advisory = 0.80 +warning = 0.70 +critical = 1.10 +"#, + ) + .unwrap(); + + assert_eq!( + config.effective_budget_alert_thresholds(), + Config::BUDGET_ALERT_THRESHOLDS + ); + } + #[test] fn save_round_trips_automation_settings() { let path = std::env::temp_dir().join(format!("ecc2-config-{}.toml", Uuid::new_v4())); @@ -420,6 +519,11 @@ move_right = "d" config.auto_dispatch_limit_per_session = 9; config.auto_create_worktrees = false; config.auto_merge_ready_worktrees = true; + config.budget_alert_thresholds = BudgetAlertThresholds { + advisory: 0.45, + warning: 0.70, + critical: 0.88, + }; config.pane_navigation.focus_metrics = "e".to_string(); config.pane_navigation.move_right = "d".to_string(); config.linear_pane_size_percent = 42; @@ -433,6 +537,14 @@ move_right = "d" assert_eq!(loaded.auto_dispatch_limit_per_session, 9); assert!(!loaded.auto_create_worktrees); assert!(loaded.auto_merge_ready_worktrees); + assert_eq!( + loaded.budget_alert_thresholds, + BudgetAlertThresholds { + advisory: 0.45, + warning: 0.70, + critical: 0.88, + } + ); assert_eq!(loaded.pane_navigation.focus_metrics, "e"); assert_eq!(loaded.pane_navigation.move_right, "d"); assert_eq!(loaded.linear_pane_size_percent, 42); diff --git a/ecc2/src/session/manager.rs b/ecc2/src/session/manager.rs index 0bdff6db..1562d796 100644 --- a/ecc2/src/session/manager.rs +++ b/ecc2/src/session/manager.rs @@ -1669,6 +1669,7 @@ mod tests { auto_merge_ready_worktrees: false, cost_budget_usd: 10.0, token_budget: 500_000, + budget_alert_thresholds: Config::BUDGET_ALERT_THRESHOLDS, theme: Theme::Dark, pane_layout: PaneLayout::Horizontal, pane_navigation: Default::default(), diff --git a/ecc2/src/tui/dashboard.rs b/ecc2/src/tui/dashboard.rs index beff4af2..5f92bf5b 100644 --- a/ecc2/src/tui/dashboard.rs +++ b/ecc2/src/tui/dashboard.rs @@ -835,11 +835,13 @@ impl Dashboard { .split(inner); let aggregate = self.aggregate_usage(); + let thresholds = self.cfg.effective_budget_alert_thresholds(); frame.render_widget( TokenMeter::tokens( "Token Budget", aggregate.total_tokens, self.cfg.token_budget, + thresholds, ), chunks[0], ); @@ -848,6 +850,7 @@ impl Dashboard { "Cost Budget", aggregate.total_cost_usd, self.cfg.cost_budget_usd, + thresholds, ), chunks[1], ); @@ -3774,6 +3777,7 @@ impl Dashboard { } fn aggregate_usage(&self) -> AggregateUsage { + let thresholds = self.cfg.effective_budget_alert_thresholds(); let total_tokens = self .sessions .iter() @@ -3784,8 +3788,12 @@ impl Dashboard { .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); + let token_state = budget_state( + total_tokens as f64, + self.cfg.token_budget as f64, + thresholds, + ); + let cost_state = budget_state(total_cost_usd, self.cfg.cost_budget_usd, thresholds); AggregateUsage { total_tokens, @@ -4072,6 +4080,7 @@ impl Dashboard { fn aggregate_cost_summary(&self) -> (String, Style) { let aggregate = self.aggregate_usage(); + let thresholds = self.cfg.effective_budget_alert_thresholds(); let mut text = if self.cfg.cost_budget_usd > 0.0 { format!( "Aggregate cost {} / {}", @@ -4085,9 +4094,9 @@ impl Dashboard { ) }; - if let Some(summary_suffix) = aggregate.overall_state.summary_suffix() { + if let Some(summary_suffix) = aggregate.overall_state.summary_suffix(thresholds) { text.push_str(" | "); - text.push_str(summary_suffix); + text.push_str(&summary_suffix); } (text, aggregate.overall_state.style()) @@ -4095,6 +4104,7 @@ impl Dashboard { fn sync_budget_alerts(&mut self) { let aggregate = self.aggregate_usage(); + let thresholds = self.cfg.effective_budget_alert_thresholds(); let current_state = aggregate.overall_state; if current_state == self.last_budget_alert_state { return; @@ -4107,7 +4117,7 @@ impl Dashboard { return; } - let Some(summary_suffix) = current_state.summary_suffix() else { + let Some(summary_suffix) = current_state.summary_suffix(thresholds) else { return; }; @@ -7098,6 +7108,26 @@ diff --git a/src/next.rs b/src/next.rs ); } + #[test] + fn aggregate_cost_summary_uses_custom_threshold_labels() { + let db = StateStore::open(Path::new(":memory:")).unwrap(); + let mut cfg = Config::default(); + cfg.cost_budget_usd = 10.0; + cfg.budget_alert_thresholds = crate::config::BudgetAlertThresholds { + advisory: 0.40, + warning: 0.70, + critical: 0.85, + }; + + let mut dashboard = Dashboard::new(db, cfg); + dashboard.sessions = vec![budget_session("sess-1", 1_000, 7.0)]; + + assert_eq!( + dashboard.aggregate_cost_summary_text(), + "Aggregate cost $7.00 / $10.00 | Budget alert 70%" + ); + } + #[test] fn aggregate_cost_summary_mentions_ninety_percent_alert() { let db = StateStore::open(Path::new(":memory:")).unwrap(); @@ -7133,6 +7163,31 @@ diff --git a/src/next.rs b/src/next.rs assert_eq!(dashboard.last_budget_alert_state, BudgetState::Alert75); } + #[test] + fn sync_budget_alerts_uses_custom_threshold_labels() { + let db = StateStore::open(Path::new(":memory:")).unwrap(); + let mut cfg = Config::default(); + cfg.token_budget = 1_000; + cfg.cost_budget_usd = 10.0; + cfg.budget_alert_thresholds = crate::config::BudgetAlertThresholds { + advisory: 0.40, + warning: 0.70, + critical: 0.85, + }; + + let mut dashboard = Dashboard::new(db, cfg); + dashboard.sessions = vec![budget_session("sess-1", 710, 2.0)]; + dashboard.last_budget_alert_state = BudgetState::Alert50; + + dashboard.sync_budget_alerts(); + + assert_eq!( + dashboard.operator_note.as_deref(), + Some("Budget alert 70% | tokens 710 / 1,000 | cost $2.00 / $10.00") + ); + assert_eq!(dashboard.last_budget_alert_state, BudgetState::Alert75); + } + #[test] fn new_session_task_uses_selected_session_context() { let dashboard = test_dashboard( @@ -9074,6 +9129,7 @@ diff --git a/src/next.rs b/src/next.rs auto_merge_ready_worktrees: false, cost_budget_usd: 10.0, token_budget: 500_000, + budget_alert_thresholds: crate::config::Config::BUDGET_ALERT_THRESHOLDS, theme: Theme::Dark, pane_layout: PaneLayout::Horizontal, pane_navigation: Default::default(), diff --git a/ecc2/src/tui/widgets.rs b/ecc2/src/tui/widgets.rs index 370011b4..1f30fcaa 100644 --- a/ecc2/src/tui/widgets.rs +++ b/ecc2/src/tui/widgets.rs @@ -1,13 +1,11 @@ +use crate::config::BudgetAlertThresholds; + use ratatui::{ prelude::*, text::{Line, Span}, widgets::{Gauge, Paragraph, Widget}, }; -pub(crate) const ALERT_THRESHOLD_50: f64 = 0.50; -pub(crate) const ALERT_THRESHOLD_75: f64 = 0.75; -pub(crate) const ALERT_THRESHOLD_90: f64 = 0.90; - #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub(crate) enum BudgetState { Unconfigured, @@ -19,23 +17,32 @@ pub(crate) enum BudgetState { } impl BudgetState { - fn badge(self) -> Option<&'static str> { + fn badge(self, thresholds: BudgetAlertThresholds) -> Option { match self { - Self::Alert50 => Some("50%"), - Self::Alert75 => Some("75%"), - Self::Alert90 => Some("90%"), - Self::OverBudget => Some("over budget"), - Self::Unconfigured => Some("no budget"), + Self::Alert50 => Some(threshold_label(thresholds.advisory)), + Self::Alert75 => Some(threshold_label(thresholds.warning)), + Self::Alert90 => Some(threshold_label(thresholds.critical)), + Self::OverBudget => Some("over budget".to_string()), + Self::Unconfigured => Some("no budget".to_string()), Self::Normal => None, } } - pub(crate) const fn summary_suffix(self) -> Option<&'static str> { + pub(crate) fn summary_suffix(self, thresholds: BudgetAlertThresholds) -> Option { match self { - Self::Alert50 => Some("Budget alert 50%"), - Self::Alert75 => Some("Budget alert 75%"), - Self::Alert90 => Some("Budget alert 90%"), - Self::OverBudget => Some("Budget exceeded"), + Self::Alert50 => Some(format!( + "Budget alert {}", + threshold_label(thresholds.advisory) + )), + Self::Alert75 => Some(format!( + "Budget alert {}", + threshold_label(thresholds.warning) + )), + Self::Alert90 => Some(format!( + "Budget alert {}", + threshold_label(thresholds.critical) + )), + Self::OverBudget => Some("Budget exceeded".to_string()), Self::Unconfigured | Self::Normal => None, } } @@ -69,30 +76,43 @@ pub(crate) struct TokenMeter<'a> { title: &'a str, used: f64, budget: f64, + thresholds: BudgetAlertThresholds, format: MeterFormat, } impl<'a> TokenMeter<'a> { - pub(crate) fn tokens(title: &'a str, used: u64, budget: u64) -> Self { + pub(crate) fn tokens( + title: &'a str, + used: u64, + budget: u64, + thresholds: BudgetAlertThresholds, + ) -> Self { Self { title, used: used as f64, budget: budget as f64, + thresholds, format: MeterFormat::Tokens, } } - pub(crate) fn currency(title: &'a str, used: f64, budget: f64) -> Self { + pub(crate) fn currency( + title: &'a str, + used: f64, + budget: f64, + thresholds: BudgetAlertThresholds, + ) -> Self { Self { title, used, budget, + thresholds, format: MeterFormat::Currency, } } pub(crate) fn state(&self) -> BudgetState { - budget_state(self.used, self.budget) + budget_state(self.used, self.budget, self.thresholds) } fn ratio(&self) -> f64 { @@ -111,7 +131,7 @@ impl<'a> TokenMeter<'a> { .add_modifier(Modifier::BOLD), )]; - if let Some(badge) = self.state().badge() { + if let Some(badge) = self.state().badge(self.thresholds) { spans.push(Span::raw(" ")); spans.push(Span::styled(format!("[{badge}]"), self.state().style())); } @@ -179,7 +199,7 @@ impl Widget for TokenMeter<'_> { .label(self.display_label()) .gauge_style( Style::default() - .fg(gradient_color(self.ratio())) + .fg(gradient_color(self.ratio(), self.thresholds)) .add_modifier(Modifier::BOLD), ) .style(Style::default().fg(Color::DarkGray)) @@ -196,39 +216,51 @@ pub(crate) fn budget_ratio(used: f64, budget: f64) -> f64 { } } -pub(crate) fn budget_state(used: f64, budget: f64) -> BudgetState { +pub(crate) fn budget_state( + used: f64, + budget: f64, + thresholds: BudgetAlertThresholds, +) -> BudgetState { if budget <= 0.0 { BudgetState::Unconfigured } else if used / budget >= 1.0 { BudgetState::OverBudget - } else if used / budget >= ALERT_THRESHOLD_90 { + } else if used / budget >= thresholds.critical { BudgetState::Alert90 - } else if used / budget >= ALERT_THRESHOLD_75 { + } else if used / budget >= thresholds.warning { BudgetState::Alert75 - } else if used / budget >= ALERT_THRESHOLD_50 { + } else if used / budget >= thresholds.advisory { BudgetState::Alert50 } else { BudgetState::Normal } } -pub(crate) fn gradient_color(ratio: f64) -> Color { +pub(crate) fn gradient_color(ratio: f64, thresholds: BudgetAlertThresholds) -> 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 <= ALERT_THRESHOLD_75 { - interpolate_rgb(GREEN, YELLOW, clamped / ALERT_THRESHOLD_75) + if clamped <= thresholds.warning { + interpolate_rgb( + GREEN, + YELLOW, + clamped / thresholds.warning.max(f64::EPSILON), + ) } else { interpolate_rgb( YELLOW, RED, - (clamped - ALERT_THRESHOLD_75) / (1.0 - ALERT_THRESHOLD_75), + (clamped - thresholds.warning) / (1.0 - thresholds.warning), ) } } +fn threshold_label(value: f64) -> String { + format!("{}%", (value * 100.0).round() as u64) +} + pub(crate) fn format_currency(value: f64) -> String { format!("${value:.2}") } @@ -264,38 +296,76 @@ fn interpolate_rgb(from: (u8, u8, u8), to: (u8, u8, u8), ratio: f64) -> Color { mod tests { use ratatui::{buffer::Buffer, layout::Rect, style::Color, widgets::Widget}; - use super::{gradient_color, BudgetState, TokenMeter}; + use crate::config::{BudgetAlertThresholds, Config}; + + use super::{gradient_color, threshold_label, BudgetState, TokenMeter}; #[test] fn budget_state_uses_alert_threshold_ladder() { assert_eq!( - TokenMeter::tokens("Token Budget", 50, 100).state(), + TokenMeter::tokens("Token Budget", 50, 100, Config::BUDGET_ALERT_THRESHOLDS).state(), BudgetState::Alert50 ); assert_eq!( - TokenMeter::tokens("Token Budget", 75, 100).state(), + TokenMeter::tokens("Token Budget", 75, 100, Config::BUDGET_ALERT_THRESHOLDS).state(), BudgetState::Alert75 ); assert_eq!( - TokenMeter::tokens("Token Budget", 90, 100).state(), + TokenMeter::tokens("Token Budget", 90, 100, Config::BUDGET_ALERT_THRESHOLDS).state(), BudgetState::Alert90 ); assert_eq!( - TokenMeter::tokens("Token Budget", 100, 100).state(), + TokenMeter::tokens("Token Budget", 100, 100, Config::BUDGET_ALERT_THRESHOLDS).state(), BudgetState::OverBudget ); } #[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.75), Color::Rgb(234, 179, 8)); - assert_eq!(gradient_color(1.0), Color::Rgb(239, 68, 68)); + assert_eq!( + gradient_color(0.0, Config::BUDGET_ALERT_THRESHOLDS), + Color::Rgb(34, 197, 94) + ); + assert_eq!( + gradient_color(0.75, Config::BUDGET_ALERT_THRESHOLDS), + Color::Rgb(234, 179, 8) + ); + assert_eq!( + gradient_color(1.0, Config::BUDGET_ALERT_THRESHOLDS), + Color::Rgb(239, 68, 68) + ); + } + + #[test] + fn token_meter_uses_custom_budget_thresholds() { + let meter = TokenMeter::tokens( + "Token Budget", + 45, + 100, + BudgetAlertThresholds { + advisory: 0.40, + warning: 0.70, + critical: 0.85, + }, + ); + + assert_eq!(meter.state(), BudgetState::Alert50); + } + + #[test] + fn threshold_label_rounds_to_percent() { + assert_eq!(threshold_label(0.4), "40%"); + assert_eq!(threshold_label(0.875), "88%"); } #[test] fn token_meter_renders_compact_usage_label() { - let meter = TokenMeter::tokens("Token Budget", 4_000, 10_000); + let meter = TokenMeter::tokens( + "Token Budget", + 4_000, + 10_000, + Config::BUDGET_ALERT_THRESHOLDS, + ); let area = Rect::new(0, 0, 48, 2); let mut buffer = Buffer::empty(area);