Merge pull request #884 from affaan-m/feat/ecc2-token-meter

feat(ecc2): token/cost meter widget

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<String> },
+    Completed {
+        summary: String,
+        files_changed: Vec<String>,
+    },
     /// Conflict detected (e.g., two agents editing the same file)
     Conflict { file: String, description: String },
 }

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

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

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)

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::<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,
+            },
+        }
+    }
 }

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