use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::RwLock;
use serde::{Deserialize, Serialize};
use tracing::debug;
use crate::pipeline::traits::{BudgetResult, CostTrackable};
#[derive(Debug, Clone, Serialize, Deserialize)]
struct UserSpend {
daily_spent: f64,
monthly_spent: f64,
daily_reserved: f64,
monthly_reserved: f64,
last_daily_reset: chrono::DateTime<chrono::Utc>,
last_monthly_reset: chrono::DateTime<chrono::Utc>,
}
impl UserSpend {
fn new() -> Self {
let now = chrono::Utc::now();
Self {
daily_spent: 0.0,
monthly_spent: 0.0,
daily_reserved: 0.0,
monthly_reserved: 0.0,
last_daily_reset: now,
last_monthly_reset: now,
}
}
fn daily_effective(&self) -> f64 {
self.daily_spent + self.daily_reserved
}
fn monthly_effective(&self) -> f64 {
self.monthly_spent + self.monthly_reserved
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CostSnapshot {
spends: HashMap<String, UserSpend>,
reset_hour_utc: u8,
}
pub struct CostTracker {
spends: RwLock<HashMap<String, UserSpend>>,
reset_hour_utc: u8,
persistence_enabled: bool,
persistence_path: Option<PathBuf>,
}
impl CostTracker {
pub fn new(reset_hour_utc: u8) -> Self {
Self {
spends: RwLock::new(HashMap::new()),
reset_hour_utc: if reset_hour_utc < 24 {
reset_hour_utc
} else {
0
},
persistence_enabled: false,
persistence_path: None,
}
}
pub fn with_persistence(mut self, path: PathBuf) -> Self {
self.persistence_enabled = true;
self.persistence_path = Some(path);
self
}
pub fn reserve_budget(
&self,
sender_id: &str,
estimated_cost_usd: f64,
user_daily_limit: f64,
user_monthly_limit: f64,
) -> BudgetResult {
if estimated_cost_usd <= 0.0 {
return BudgetResult::Approved;
}
let mut spends = self.spends.write().expect("cost tracker lock poisoned");
let entry = spends
.entry(sender_id.to_string())
.or_insert_with(UserSpend::new);
self.maybe_reset_entry(entry);
if user_daily_limit > 0.0 {
let effective = entry.daily_effective();
if effective + estimated_cost_usd > user_daily_limit {
return BudgetResult::DailyLimitExceeded {
spent: effective,
limit: user_daily_limit,
};
}
}
if user_monthly_limit > 0.0 {
let effective = entry.monthly_effective();
if effective + estimated_cost_usd > user_monthly_limit {
return BudgetResult::MonthlyLimitExceeded {
spent: effective,
limit: user_monthly_limit,
};
}
}
entry.daily_reserved += estimated_cost_usd;
entry.monthly_reserved += estimated_cost_usd;
BudgetResult::Approved
}
pub fn reconcile_actual(
&self,
sender_id: &str,
estimated_cost_usd: f64,
actual_cost_usd: f64,
) {
let mut spends = self.spends.write().expect("cost tracker lock poisoned");
let entry = spends
.entry(sender_id.to_string())
.or_insert_with(UserSpend::new);
entry.daily_reserved = (entry.daily_reserved - estimated_cost_usd).max(0.0);
entry.monthly_reserved = (entry.monthly_reserved - estimated_cost_usd).max(0.0);
entry.daily_spent = (entry.daily_spent + actual_cost_usd).max(0.0);
entry.monthly_spent = (entry.monthly_spent + actual_cost_usd).max(0.0);
}
pub fn check_budget(
&self,
sender_id: &str,
estimated_cost_usd: f64,
user_daily_limit: f64,
user_monthly_limit: f64,
) -> BudgetResult {
if estimated_cost_usd <= 0.0 {
return BudgetResult::Approved;
}
let mut spends = self.spends.write().expect("cost tracker lock poisoned");
let entry = spends
.entry(sender_id.to_string())
.or_insert_with(UserSpend::new);
self.maybe_reset_entry(entry);
if user_daily_limit > 0.0 {
let effective = entry.daily_effective();
if effective + estimated_cost_usd > user_daily_limit {
return BudgetResult::DailyLimitExceeded {
spent: effective,
limit: user_daily_limit,
};
}
}
if user_monthly_limit > 0.0 {
let effective = entry.monthly_effective();
if effective + estimated_cost_usd > user_monthly_limit {
return BudgetResult::MonthlyLimitExceeded {
spent: effective,
limit: user_monthly_limit,
};
}
}
BudgetResult::Approved
}
pub fn record_estimated(&self, sender_id: &str, estimated_cost: f64) {
if estimated_cost <= 0.0 {
return;
}
let mut spends = self.spends.write().expect("cost tracker lock poisoned");
let entry = spends
.entry(sender_id.to_string())
.or_insert_with(UserSpend::new);
self.maybe_reset_entry(entry);
entry.daily_spent += estimated_cost;
entry.monthly_spent += estimated_cost;
}
pub fn record_actual(&self, sender_id: &str, estimated_cost: f64, actual_cost: f64) {
let mut spends = self.spends.write().expect("cost tracker lock poisoned");
let entry = spends
.entry(sender_id.to_string())
.or_insert_with(UserSpend::new);
let delta = actual_cost - estimated_cost;
entry.daily_spent = (entry.daily_spent + delta).max(0.0);
entry.monthly_spent = (entry.monthly_spent + delta).max(0.0);
}
pub fn estimate_cost(cost_per_1k_tokens: f64, estimated_tokens: usize) -> f64 {
cost_per_1k_tokens * (estimated_tokens as f64) / 1000.0
}
pub fn get_spend(&self, sender_id: &str) -> (f64, f64) {
let spends = self.spends.read().expect("cost tracker lock poisoned");
match spends.get(sender_id) {
Some(entry) => (entry.daily_effective(), entry.monthly_effective()),
None => (0.0, 0.0),
}
}
pub fn persist(&self) -> std::io::Result<()> {
let Some(ref path) = self.persistence_path else {
return Ok(());
};
if !self.persistence_enabled {
return Ok(());
}
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let snapshot = {
let spends = self.spends.read().expect("cost tracker lock poisoned");
CostSnapshot {
spends: spends.clone(),
reset_hour_utc: self.reset_hour_utc,
}
};
let json = serde_json::to_string_pretty(&snapshot)
.map_err(std::io::Error::other)?;
let tmp_path = path.with_extension("json.tmp");
std::fs::write(&tmp_path, &json)?;
std::fs::rename(&tmp_path, path)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = std::fs::Permissions::from_mode(0o600);
std::fs::set_permissions(path, perms)?;
}
debug!("saved cost tracking state to {}", path.display());
Ok(())
}
pub fn load(path: &std::path::Path) -> std::io::Result<Self> {
let data = std::fs::read_to_string(path)?;
let snapshot: CostSnapshot = serde_json::from_str(&data)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
let tracker = Self {
spends: RwLock::new(snapshot.spends),
reset_hour_utc: snapshot.reset_hour_utc,
persistence_enabled: true,
persistence_path: Some(path.to_path_buf()),
};
debug!("loaded cost tracking state from {}", path.display());
Ok(tracker)
}
fn maybe_reset_entry(&self, entry: &mut UserSpend) {
let now = chrono::Utc::now();
if self.should_reset_monthly(now, entry.last_monthly_reset) {
debug!("resetting monthly cost counters");
entry.daily_spent = 0.0;
entry.daily_reserved = 0.0;
entry.monthly_spent = 0.0;
entry.monthly_reserved = 0.0;
entry.last_daily_reset = now;
entry.last_monthly_reset = now;
return;
}
if self.should_reset_daily(now, entry.last_daily_reset) {
debug!("resetting daily cost counters");
entry.daily_spent = 0.0;
entry.daily_reserved = 0.0;
entry.last_daily_reset = now;
}
}
fn should_reset_daily(
&self,
now: chrono::DateTime<chrono::Utc>,
last_reset: chrono::DateTime<chrono::Utc>,
) -> bool {
let reset_offset = chrono::Duration::hours(self.reset_hour_utc as i64);
let adjusted_now = now - reset_offset;
let adjusted_last = last_reset - reset_offset;
let day_now = adjusted_now.date_naive();
let day_last = adjusted_last.date_naive();
day_now > day_last
}
fn should_reset_monthly(
&self,
now: chrono::DateTime<chrono::Utc>,
last_reset: chrono::DateTime<chrono::Utc>,
) -> bool {
use chrono::Datelike;
(now.year(), now.month()) != (last_reset.year(), last_reset.month())
}
}
impl std::fmt::Debug for CostTracker {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let user_count = self
.spends
.read()
.map(|s| s.len())
.unwrap_or(0);
f.debug_struct("CostTracker")
.field("users", &user_count)
.field("reset_hour_utc", &self.reset_hour_utc)
.field("persistence_enabled", &self.persistence_enabled)
.field("persistence_path", &self.persistence_path)
.finish()
}
}
impl CostTrackable for CostTracker {
fn check_budget(
&self,
sender_id: &str,
estimated_cost: f64,
daily_limit: f64,
monthly_limit: f64,
) -> BudgetResult {
CostTracker::check_budget(self, sender_id, estimated_cost, daily_limit, monthly_limit)
}
fn record_estimated(&self, sender_id: &str, estimated_cost: f64) {
CostTracker::record_estimated(self, sender_id, estimated_cost);
}
fn record_actual(&self, sender_id: &str, estimated_cost: f64, actual_cost: f64) {
CostTracker::record_actual(self, sender_id, estimated_cost, actual_cost);
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
fn tracker() -> CostTracker {
CostTracker::new(0)
}
#[test]
fn reserve_budget_approved_within_limits() {
let t = tracker();
let result = t.reserve_budget("alice", 1.0, 5.0, 100.0);
assert_eq!(result, BudgetResult::Approved);
}
#[test]
fn check_budget_approved_within_limits() {
let t = tracker();
let result = t.check_budget("alice", 1.0, 5.0, 100.0);
assert_eq!(result, BudgetResult::Approved);
}
#[test]
fn reserve_budget_daily_limit_exceeded() {
let t = tracker();
assert_eq!(t.reserve_budget("alice", 4.5, 5.0, 100.0), BudgetResult::Approved);
let result = t.reserve_budget("alice", 1.0, 5.0, 100.0);
assert!(matches!(result, BudgetResult::DailyLimitExceeded { .. }));
}
#[test]
fn check_budget_daily_limit_exceeded_after_record() {
let t = tracker();
t.record_estimated("alice", 4.5);
let result = t.check_budget("alice", 1.0, 5.0, 100.0);
assert!(matches!(result, BudgetResult::DailyLimitExceeded { .. }));
}
#[test]
fn reserve_budget_monthly_limit_exceeded() {
let t = tracker();
assert_eq!(t.reserve_budget("alice", 9.5, 20.0, 10.0), BudgetResult::Approved);
let result = t.reserve_budget("alice", 1.0, 20.0, 10.0);
assert!(matches!(result, BudgetResult::MonthlyLimitExceeded { .. }));
}
#[test]
fn check_budget_monthly_limit_exceeded_after_record() {
let t = tracker();
t.record_estimated("alice", 99.5);
let result = t.check_budget("alice", 1.0, 200.0, 100.0);
assert!(matches!(result, BudgetResult::MonthlyLimitExceeded { .. }));
}
#[test]
fn reconcile_actual_lower_cost() {
let t = tracker();
assert_eq!(t.reserve_budget("alice", 3.00, 10.0, 100.0), BudgetResult::Approved);
t.reconcile_actual("alice", 3.00, 2.00);
let (daily, monthly) = t.get_spend("alice");
assert!((daily - 2.0).abs() < 1e-10);
assert!((monthly - 2.0).abs() < 1e-10);
}
#[test]
fn reconcile_actual_higher_cost() {
let t = tracker();
assert_eq!(t.reserve_budget("alice", 2.00, 10.0, 100.0), BudgetResult::Approved);
t.reconcile_actual("alice", 2.00, 3.50);
let (daily, monthly) = t.get_spend("alice");
assert!((daily - 3.5).abs() < 1e-10);
assert!((monthly - 3.5).abs() < 1e-10);
}
#[test]
fn independent_user_tracking() {
let t = tracker();
t.record_estimated("alice", 3.0);
t.record_estimated("bob", 7.0);
t.record_estimated("alice", 2.0);
let (alice_d, alice_m) = t.get_spend("alice");
let (bob_d, bob_m) = t.get_spend("bob");
assert!((alice_d - 5.0).abs() < 1e-10);
assert!((alice_m - 5.0).abs() < 1e-10);
assert!((bob_d - 7.0).abs() < 1e-10);
assert!((bob_m - 7.0).abs() < 1e-10);
}
#[test]
fn daily_reset_boundary_detection() {
let t = CostTracker::new(0);
let now = chrono::DateTime::parse_from_rfc3339("2025-06-15T12:00:00Z")
.unwrap()
.with_timezone(&chrono::Utc);
let yesterday = now - chrono::Duration::hours(25);
assert!(!t.should_reset_daily(now, now));
assert!(!t.should_reset_daily(now, now - chrono::Duration::hours(1)));
assert!(t.should_reset_daily(now, yesterday));
}
#[test]
fn daily_reset_respects_reset_hour() {
let t = CostTracker::new(6); let now = chrono::Utc::now();
let long_ago = now - chrono::Duration::hours(25);
assert!(t.should_reset_daily(now, long_ago));
}
#[test]
fn monthly_reset_boundary_detection() {
let t = tracker();
let now = chrono::Utc::now();
assert!(!t.should_reset_monthly(now, now));
let different_month = now - chrono::Duration::days(35);
assert!(t.should_reset_monthly(now, different_month));
}
#[test]
fn estimate_cost_calculation() {
let cost = CostTracker::estimate_cost(0.001, 1000);
assert!((cost - 0.001).abs() < 1e-10);
let cost = CostTracker::estimate_cost(0.0, 2000);
assert!(cost.abs() < 1e-10);
let cost = CostTracker::estimate_cost(0.05, 10000);
assert!((cost - 0.50).abs() < 1e-10);
let cost = CostTracker::estimate_cost(0.01, 6096);
assert!((cost - 0.06096).abs() < 1e-10);
}
#[test]
fn zero_limit_means_unlimited() {
let t = tracker();
t.record_estimated("alice", 999.0);
let result = t.check_budget("alice", 100.0, 0.0, 0.0);
assert_eq!(result, BudgetResult::Approved);
}
#[test]
fn concurrent_record_no_panic() {
let t = Arc::new(CostTracker::new(0));
let mut handles = vec![];
for i in 0..10 {
let tracker = Arc::clone(&t);
handles.push(std::thread::spawn(move || {
for _ in 0..100 {
let user = format!("user_{i}");
tracker.record_estimated(&user, 0.01);
let _ = tracker.check_budget(&user, 0.01, 50.0, 500.0);
}
}));
}
for h in handles {
h.join().expect("thread panicked");
}
let total: f64 = (0..10)
.map(|i| t.get_spend(&format!("user_{i}")).0)
.sum();
assert!((total - 10.0).abs() < 0.01);
}
#[test]
fn concurrent_reserve_budget_same_user() {
let t = Arc::new(CostTracker::new(0));
let mut handles = vec![];
for _ in 0..20 {
let tracker = Arc::clone(&t);
handles.push(std::thread::spawn(move || {
tracker.reserve_budget("alice", 0.50, 8.0, 100.0)
}));
}
let results: Vec<BudgetResult> = handles
.into_iter()
.map(|h| h.join().expect("thread panicked"))
.collect();
let ok_count = results.iter().filter(|r| r.is_approved()).count();
let over_count = results.iter().filter(|r| !r.is_approved()).count();
assert!(ok_count <= 16);
assert_eq!(ok_count + over_count, 20);
let (daily, _) = t.get_spend("alice");
assert!(daily <= 8.0 + 1e-10);
}
#[test]
fn persistence_roundtrip() {
let dir = std::env::temp_dir().join("clawft_cost_tracker_test_roundtrip");
let _ = std::fs::create_dir_all(&dir);
let path = dir.join("cost_tracking.json");
let _ = std::fs::remove_file(&path);
{
let t = CostTracker::new(0).with_persistence(path.clone());
t.record_estimated("alice", 3.50);
t.record_estimated("bob", 1.25);
t.persist().expect("persist failed");
}
{
let t = CostTracker::load(&path).expect("load failed");
let (alice_d, alice_m) = t.get_spend("alice");
let (bob_d, bob_m) = t.get_spend("bob");
assert!((alice_d - 3.50).abs() < 1e-10);
assert!((alice_m - 3.50).abs() < 1e-10);
assert!((bob_d - 1.25).abs() < 1e-10);
assert!((bob_m - 1.25).abs() < 1e-10);
}
let _ = std::fs::remove_file(&path);
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn empty_sender_id_works() {
let t = tracker();
let result = t.reserve_budget("", 1.0, 5.0, 100.0);
assert_eq!(result, BudgetResult::Approved);
let (daily, monthly) = t.get_spend("");
assert!((daily - 1.0).abs() < 1e-10);
assert!((monthly - 1.0).abs() < 1e-10);
}
#[test]
fn get_spend_unknown_user_returns_zero() {
let t = tracker();
let (daily, monthly) = t.get_spend("nonexistent");
assert!(daily.abs() < 1e-10);
assert!(monthly.abs() < 1e-10);
}
#[test]
fn get_spend_reflects_committed_and_reserved() {
let t = tracker();
t.reserve_budget("alice", 2.0, 10.0, 100.0);
let (daily, monthly) = t.get_spend("alice");
assert!((daily - 2.0).abs() < 1e-10);
assert!((monthly - 2.0).abs() < 1e-10);
t.reconcile_actual("alice", 2.0, 1.5);
let (daily, monthly) = t.get_spend("alice");
assert!((daily - 1.5).abs() < 1e-10);
assert!((monthly - 1.5).abs() < 1e-10);
}
#[test]
fn budget_result_is_approved() {
assert!(BudgetResult::Approved.is_approved());
assert!(!BudgetResult::DailyLimitExceeded { spent: 5.0, limit: 5.0 }.is_approved());
assert!(!BudgetResult::MonthlyLimitExceeded { spent: 100.0, limit: 100.0 }.is_approved());
}
#[test]
fn reserve_budget_no_record_on_failure() {
let t = tracker();
assert_eq!(t.reserve_budget("alice", 4.0, 5.0, 100.0), BudgetResult::Approved);
let result = t.reserve_budget("alice", 2.0, 5.0, 100.0);
assert!(matches!(result, BudgetResult::DailyLimitExceeded { .. }));
let (daily, _) = t.get_spend("alice");
assert!((daily - 4.0).abs() < 1e-10);
}
#[test]
fn reserve_budget_zero_cost_approved() {
let t = tracker();
let result = t.reserve_budget("alice", 0.0, 5.0, 100.0);
assert_eq!(result, BudgetResult::Approved);
let (daily, _) = t.get_spend("alice");
assert!(daily.abs() < 1e-10);
}
#[test]
fn reserve_budget_negative_cost_approved() {
let t = tracker();
let result = t.reserve_budget("alice", -1.0, 5.0, 100.0);
assert_eq!(result, BudgetResult::Approved);
let (daily, _) = t.get_spend("alice");
assert!(daily.abs() < 1e-10);
}
#[test]
fn reconcile_actual_clamps_to_zero() {
let t = tracker();
t.reserve_budget("alice", 1.0, 10.0, 100.0);
t.reconcile_actual("alice", 1.0, 0.0);
let (daily, monthly) = t.get_spend("alice");
assert!(daily >= 0.0);
assert!(monthly >= 0.0);
}
#[test]
fn record_estimated_zero_is_noop() {
let t = tracker();
t.record_estimated("alice", 0.0);
let (daily, _) = t.get_spend("alice");
assert!(daily.abs() < 1e-10);
}
#[test]
fn record_actual_delegates_correctly() {
let t = tracker();
t.record_estimated("alice", 2.0);
t.record_actual("alice", 2.0, 1.5);
let (daily, monthly) = t.get_spend("alice");
assert!((daily - 1.5).abs() < 1e-10);
assert!((monthly - 1.5).abs() < 1e-10);
}
#[cfg(unix)]
#[test]
fn persistence_file_permissions() {
use std::os::unix::fs::PermissionsExt;
let dir = std::env::temp_dir().join("clawft_cost_tracker_test_perms");
let _ = std::fs::create_dir_all(&dir);
let path = dir.join("cost_tracking.json");
let _ = std::fs::remove_file(&path);
let t = CostTracker::new(0).with_persistence(path.clone());
t.record_estimated("alice", 1.0);
t.persist().expect("persist failed");
let metadata = std::fs::metadata(&path).expect("metadata failed");
let mode = metadata.permissions().mode() & 0o777;
assert_eq!(mode, 0o600, "persistence file should have 0600 permissions");
let _ = std::fs::remove_file(&path);
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn debug_impl_does_not_panic() {
let t = tracker();
t.record_estimated("alice", 1.0);
let debug_str = format!("{:?}", t);
assert!(debug_str.contains("CostTracker"));
assert!(debug_str.contains("users"));
}
#[test]
fn daily_checked_before_monthly() {
let t = tracker();
t.record_estimated("alice", 9.5);
let result = t.check_budget("alice", 1.0, 5.0, 10.0);
assert!(matches!(result, BudgetResult::DailyLimitExceeded { .. }));
}
}