feat: add ecc2 configurable budget thresholds

This commit is contained in:
Affaan Mustafa
2026-04-09 06:36:22 -07:00
parent 95c33d3c04
commit 67d06687a0
4 changed files with 282 additions and 43 deletions

View File

@@ -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<String> {
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<String> {
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);