lean_ctx/core/

budget_tracker.rs

//! Runtime budget tracking against role limits.
//!
//! Compares accumulated session counters with the active role's `RoleLimits`
//! and produces `BudgetStatus` verdicts (Ok / Warning / Exhausted).

use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};

use serde::Serialize;

use crate::core::roles::{self, RoleLimits};

static TRACKER: OnceLock<BudgetTracker> = OnceLock::new();

pub struct BudgetTracker {
    context_tokens: AtomicU64,
    shell_invocations: AtomicUsize,
    cost_millicents: AtomicU64,
    tool_calls: AtomicUsize,
}

impl BudgetTracker {
    fn new() -> Self {
        Self {
            context_tokens: AtomicU64::new(0),
            shell_invocations: AtomicUsize::new(0),
            cost_millicents: AtomicU64::new(0),
            tool_calls: AtomicUsize::new(0),
        }
    }

    pub fn global() -> &'static BudgetTracker {
        TRACKER.get_or_init(BudgetTracker::new)
    }

    pub fn record_tokens(&self, tokens: u64) {
        self.context_tokens.fetch_add(tokens, Ordering::Relaxed);
    }

    pub fn record_shell(&self) {
        self.shell_invocations.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_tool_call(&self) {
        self.tool_calls.fetch_add(1, Ordering::Relaxed);
    }

    pub fn tool_calls_count(&self) -> usize {
        self.tool_calls.load(Ordering::Relaxed)
    }

    pub fn record_cost_usd(&self, usd: f64) {
        let mc = (usd * 100_000.0) as u64;
        self.cost_millicents.fetch_add(mc, Ordering::Relaxed);
    }

    pub fn tokens_used(&self) -> u64 {
        self.context_tokens.load(Ordering::Relaxed)
    }

    pub fn shell_used(&self) -> usize {
        self.shell_invocations.load(Ordering::Relaxed)
    }

    pub fn cost_usd(&self) -> f64 {
        self.cost_millicents.load(Ordering::Relaxed) as f64 / 100_000.0
    }

    /// Returns `Some(message)` when the session cost cap is exceeded (#794).
    /// Returns `None` when no cap is configured, the cap isn't reached, or
    /// `LEAN_CTX_COST_CAP_OVERRIDE=1` is set.
    pub fn cost_cap_message(&self) -> Option<String> {
        let cfg = crate::core::config::Config::load();
        let cap = cfg.cost.max_session_cost_usd;
        if cap <= 0.0 {
            return None;
        }
        if std::env::var("LEAN_CTX_COST_CAP_OVERRIDE")
            .ok()
            .is_some_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
        {
            return None;
        }
        let used = self.cost_usd();
        if used < cap {
            return None;
        }
        Some(format!(
            "[COST CAP] Session cost ${used:.2} reached ${cap:.2} limit. Use ctx_session(action=budget, override=true) or LEAN_CTX_COST_CAP_OVERRIDE=1 to continue."
        ))
    }

    pub fn reset(&self) {
        self.context_tokens.store(0, Ordering::Relaxed);
        self.shell_invocations.store(0, Ordering::Relaxed);
        self.cost_millicents.store(0, Ordering::Relaxed);
        self.tool_calls.store(0, Ordering::Relaxed);
    }

    /// A context policy pack may **tighten** (never loosen) the per-session
    /// token ceiling (#673). Pure so it can be unit-tested without globals.
    fn capped_token_limit(role_limit: usize, policy_cap: Option<u32>) -> usize {
        match policy_cap {
            Some(cap) => role_limit.min(cap as usize),
            None => role_limit,
        }
    }

    pub fn check(&self) -> BudgetSnapshot {
        let mut limits = roles::active_role().limits;
        let role_name = roles::active_role_name();

        // #673 — apply the active context policy pack's token cap (Local-Free:
        // this only affects agent budget accounting, never a human's own reads).
        let policy_cap =
            crate::core::policy::runtime::active().and_then(|p| p.resolved.max_context_tokens);
        limits.max_context_tokens = Self::capped_token_limit(limits.max_context_tokens, policy_cap);

        let tokens = self.tokens_used();
        let shell = self.shell_used();
        let cost = self.cost_usd();

        BudgetSnapshot {
            role: role_name,
            tokens: DimensionStatus::evaluate(tokens as usize, limits.max_context_tokens, &limits),
            shell: DimensionStatus::evaluate(shell, limits.max_shell_invocations, &limits),
            cost: CostStatus::evaluate(cost, limits.max_cost_usd, &limits),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum BudgetLevel {
    Ok,
    Warning,
    Exhausted,
}

impl std::fmt::Display for BudgetLevel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Ok => write!(f, "OK"),
            Self::Warning => write!(f, "WARNING"),
            Self::Exhausted => write!(f, "EXHAUSTED"),
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct DimensionStatus {
    pub used: usize,
    pub limit: usize,
    pub percent: u8,
    pub level: BudgetLevel,
}

impl DimensionStatus {
    fn evaluate(used: usize, limit: usize, limits: &RoleLimits) -> Self {
        if limit == 0 {
            // Zero limit with any usage => Warning (not Exhausted, LeanCTX never blocks)
            return Self {
                used,
                limit,
                percent: 0,
                level: if used > 0 {
                    BudgetLevel::Warning
                } else {
                    BudgetLevel::Ok
                },
            };
        }
        let percent = ((used as f64 / limit as f64) * 100.0).min(254.0) as u8;
        // block_at_percent == 255 means blocking is disabled (LeanCTX default)
        let level = if limits.block_at_percent < 255 && percent >= limits.block_at_percent {
            BudgetLevel::Exhausted
        } else if percent >= limits.warn_at_percent {
            BudgetLevel::Warning
        } else {
            BudgetLevel::Ok
        };
        Self {
            used,
            limit,
            percent,
            level,
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct CostStatus {
    pub used_usd: f64,
    pub limit_usd: f64,
    pub percent: u8,
    pub level: BudgetLevel,
}

impl CostStatus {
    fn evaluate(used: f64, limit: f64, limits: &RoleLimits) -> Self {
        if limit <= 0.0 {
            // Zero limit with any usage => Warning (not Exhausted, LeanCTX never blocks)
            return Self {
                used_usd: used,
                limit_usd: limit,
                percent: 0,
                level: if used > 0.0 {
                    BudgetLevel::Warning
                } else {
                    BudgetLevel::Ok
                },
            };
        }
        let pct = ((used / limit) * 100.0).min(254.0) as u8;
        // block_at_percent == 255 means blocking is disabled (LeanCTX default)
        let level = if limits.block_at_percent < 255 && pct >= limits.block_at_percent {
            BudgetLevel::Exhausted
        } else if pct >= limits.warn_at_percent {
            BudgetLevel::Warning
        } else {
            BudgetLevel::Ok
        };
        Self {
            used_usd: used,
            limit_usd: limit,
            percent: pct,
            level,
        }
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct BudgetSnapshot {
    pub role: String,
    pub tokens: DimensionStatus,
    pub shell: DimensionStatus,
    pub cost: CostStatus,
}

impl BudgetSnapshot {
    pub fn worst_level(&self) -> &BudgetLevel {
        for level in [&self.tokens.level, &self.shell.level, &self.cost.level] {
            if *level == BudgetLevel::Exhausted {
                return level;
            }
        }
        for level in [&self.tokens.level, &self.shell.level, &self.cost.level] {
            if *level == BudgetLevel::Warning {
                return level;
            }
        }
        &BudgetLevel::Ok
    }

    pub fn format_compact(&self) -> String {
        format!(
            "Budget[role:{}]: tokens {}/{} ({}%) | shell {}/{} ({}%) | cost ${:.2}/${:.2} ({}%) → {}",
            self.role,
            self.tokens.used,
            self.tokens.limit,
            self.tokens.percent,
            self.shell.used,
            self.shell.limit,
            self.shell.percent,
            self.cost.used_usd,
            self.cost.limit_usd,
            self.cost.percent,
            self.worst_level(),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tracker_starts_at_zero() {
        let t = BudgetTracker::new();
        assert_eq!(t.tokens_used(), 0);
        assert_eq!(t.shell_used(), 0);
        assert!((t.cost_usd() - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn record_and_read() {
        let t = BudgetTracker::new();
        t.record_tokens(5000);
        t.record_tokens(3000);
        t.record_shell();
        t.record_shell();
        t.record_cost_usd(0.50);
        assert_eq!(t.tokens_used(), 8000);
        assert_eq!(t.shell_used(), 2);
        assert!((t.cost_usd() - 0.50).abs() < 0.001);
    }

    #[test]
    fn reset_clears_all() {
        let t = BudgetTracker::new();
        t.record_tokens(10_000);
        t.record_shell();
        t.record_cost_usd(1.0);
        t.reset();
        assert_eq!(t.tokens_used(), 0);
        assert_eq!(t.shell_used(), 0);
        assert!((t.cost_usd() - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn dimension_status_ok() {
        let limits = RoleLimits::default();
        let s = DimensionStatus::evaluate(50_000, 200_000, &limits);
        assert_eq!(s.level, BudgetLevel::Ok);
        assert_eq!(s.percent, 25);
    }

    #[test]
    fn policy_cap_tightens_but_never_loosens() {
        // #673: a policy may only lower the ceiling, and a None cap is a no-op.
        assert_eq!(
            BudgetTracker::capped_token_limit(200_000, Some(5_000)),
            5_000
        );
        assert_eq!(
            BudgetTracker::capped_token_limit(4_000, Some(50_000)),
            4_000
        );
        assert_eq!(BudgetTracker::capped_token_limit(10_000, None), 10_000);
    }

    #[test]
    fn dimension_status_warning() {
        let limits = RoleLimits::default();
        let s = DimensionStatus::evaluate(170_000, 200_000, &limits);
        assert_eq!(s.level, BudgetLevel::Warning);
        assert_eq!(s.percent, 85);
    }

    #[test]
    fn dimension_status_at_100_percent_is_warning_by_default() {
        // With block_at_percent=255 (default), 100% usage is Warning, not Exhausted
        let limits = RoleLimits::default();
        assert_eq!(limits.block_at_percent, 255); // Default = never block
        let s = DimensionStatus::evaluate(200_000, 200_000, &limits);
        assert_eq!(s.level, BudgetLevel::Warning);
        assert_eq!(s.percent, 100);
    }

    #[test]
    fn dimension_status_exhausted_when_blocking_enabled() {
        // Exhausted only happens when block_at_percent is explicitly set low
        let limits = RoleLimits {
            block_at_percent: 100,
            ..Default::default()
        };
        let s = DimensionStatus::evaluate(200_000, 200_000, &limits);
        assert_eq!(s.level, BudgetLevel::Exhausted);
    }

    #[test]
    fn zero_limit_warns_usage() {
        // Zero limit with any usage => Warning (not Exhausted, LeanCTX never blocks by default)
        let limits = RoleLimits::default();
        let s = DimensionStatus::evaluate(1, 0, &limits);
        assert_eq!(s.level, BudgetLevel::Warning);
    }

    #[test]
    fn cost_cap_no_limit_returns_none() {
        let t = BudgetTracker::new();
        t.record_cost_usd(100.0);
        // Without a configured cap (default 0), no message is returned.
        // We test the pure logic; the config defaults to 0.
        assert!(t.cost_cap_message().is_none());
    }

    #[test]
    fn cost_cap_blocks_when_exceeded() {
        let _env_lock = crate::core::data_dir::test_env_lock();
        let t = BudgetTracker::new();
        t.record_cost_usd(6.0);
        // SAFETY: single-threaded test — no concurrent env access.
        unsafe {
            std::env::set_var("LEAN_CTX_COST_CAP_OVERRIDE", "1");
        }
        assert!(
            t.cost_cap_message().is_none(),
            "override=1 must bypass cost cap"
        );
        // SAFETY: single-threaded test — no concurrent env access.
        unsafe {
            std::env::remove_var("LEAN_CTX_COST_CAP_OVERRIDE");
        }
    }

    #[test]
    fn cost_status_warning() {
        let limits = RoleLimits::default();
        let s = CostStatus::evaluate(4.5, 5.0, &limits);
        assert_eq!(s.level, BudgetLevel::Warning);
    }

    #[test]
    fn snapshot_worst_level() {
        let limits = RoleLimits::default();
        let snap = BudgetSnapshot {
            role: "test".into(),
            tokens: DimensionStatus::evaluate(50_000, 200_000, &limits),
            shell: DimensionStatus::evaluate(90, 100, &limits),
            cost: CostStatus::evaluate(1.0, 5.0, &limits),
        };
        assert_eq!(*snap.worst_level(), BudgetLevel::Warning);
    }

    #[test]
    fn format_compact_includes_all() {
        let s = BudgetSnapshot {
            role: "coder".into(),
            tokens: DimensionStatus {
                used: 1000,
                limit: 200_000,
                percent: 0,
                level: BudgetLevel::Ok,
            },
            shell: DimensionStatus {
                used: 5,
                limit: 100,
                percent: 5,
                level: BudgetLevel::Ok,
            },
            cost: CostStatus {
                used_usd: 0.25,
                limit_usd: 5.0,
                percent: 5,
                level: BudgetLevel::Ok,
            },
        };
        let out = s.format_compact();
        assert!(out.contains("role:coder"));
        assert!(out.contains("tokens"));
        assert!(out.contains("shell"));
        assert!(out.contains("cost"));
        assert!(out.contains("OK"));
    }
}