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
This commit is contained in:
Affaan Mustafa
2026-03-23 03:46:25 -07:00
parent 67306c22cd
commit 63410afcad
8 changed files with 587 additions and 51 deletions

View File

@@ -13,7 +13,10 @@ pub enum MessageType {
/// Response to a query
Response { answer: String },
/// Notification of completion
Completed { summary: String, files_changed: Vec<String> },
Completed {
summary: String,
files_changed: Vec<String>,
},
/// Conflict detected (e.g., two agents editing the same file)
Conflict { file: String, description: String },
}

View File

@@ -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);
}
}

View File

@@ -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) => {

View File

@@ -1,8 +1,8 @@
use anyhow::Result;
use std::fmt;
use super::{Session, SessionMetrics, SessionState};
use super::store::StateStore;
use super::{Session, SessionMetrics, SessionState};
use crate::config::Config;
use crate::worktree;

View File

@@ -170,16 +170,12 @@ impl StateStore {
pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
let sessions = self.list_sessions()?;
Ok(sessions.into_iter().find(|s| s.id == id || s.id.starts_with(id)))
Ok(sessions
.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<()> {
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)
VALUES (?1, ?2, ?3, ?4, ?5)",

View File

@@ -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::<f64>();
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,
},
}
}
}

View File

@@ -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::<String>();
assert!(rendered.contains("4,000 / 10,000 tok (40%)"));
}
}

View File

@@ -28,7 +28,11 @@ pub fn create_for_session(session_id: &str, cfg: &Config) -> Result<WorktreeInfo
anyhow::bail!("git worktree add failed: {stderr}");
}
tracing::info!("Created worktree at {} on branch {}", path.display(), branch);
tracing::info!(
"Created worktree at {} on branch {}",
path.display(),
branch
);
Ok(WorktreeInfo {
path,