use chrono::{Datelike, Local, NaiveDate, Weekday}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; /// Represents a single day's cost data #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct DailyCost { pub date: String, // ISO date string (YYYY-MM-DD) pub input_tokens: u64, pub output_tokens: u64, pub cost_usd: f64, pub messages_sent: u64, pub sessions_count: u64, } /// Historical cost tracking data #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct CostHistory { /// Daily costs indexed by date string (YYYY-MM-DD) pub daily_costs: HashMap, /// Cost alert thresholds pub daily_alert_threshold: Option, pub weekly_alert_threshold: Option, pub monthly_alert_threshold: Option, /// Whether alerts have been triggered today pub daily_alert_triggered: bool, pub weekly_alert_triggered: bool, pub monthly_alert_triggered: bool, pub last_alert_reset_date: Option, } impl CostHistory { pub fn new() -> Self { Self::default() } /// Get today's date as a string fn today_str() -> String { Local::now().format("%Y-%m-%d").to_string() } /// Get the start of the current week (Monday) fn week_start() -> NaiveDate { let today = Local::now().date_naive(); let days_since_monday = today.weekday().num_days_from_monday(); today - chrono::Duration::days(days_since_monday as i64) } /// Get the start of the current month fn month_start() -> NaiveDate { let today = Local::now().date_naive(); NaiveDate::from_ymd_opt(today.year(), today.month(), 1).unwrap_or(today) } /// Add cost for today pub fn add_cost(&mut self, input_tokens: u64, output_tokens: u64, cost_usd: f64) { let today = Self::today_str(); // Reset alert flags if it's a new day if self.last_alert_reset_date.as_ref() != Some(&today) { self.daily_alert_triggered = false; // Reset weekly on Monday if Local::now().weekday() == Weekday::Mon { self.weekly_alert_triggered = false; } // Reset monthly on the 1st if Local::now().day() == 1 { self.monthly_alert_triggered = false; } self.last_alert_reset_date = Some(today.clone()); } let daily = self.daily_costs.entry(today).or_default(); daily.input_tokens += input_tokens; daily.output_tokens += output_tokens; daily.cost_usd += cost_usd; daily.messages_sent += 1; } /// Increment session count for today pub fn increment_sessions(&mut self) { let today = Self::today_str(); let daily = self.daily_costs.entry(today.clone()).or_insert_with(|| DailyCost { date: today, ..Default::default() }); daily.sessions_count += 1; } /// Get today's cost pub fn get_today_cost(&self) -> f64 { self.daily_costs .get(&Self::today_str()) .map(|d| d.cost_usd) .unwrap_or(0.0) } /// Get this week's cost (Monday to Sunday) pub fn get_week_cost(&self) -> f64 { let week_start = Self::week_start(); self.daily_costs .values() .filter(|d| { NaiveDate::parse_from_str(&d.date, "%Y-%m-%d") .map(|date| date >= week_start) .unwrap_or(false) }) .map(|d| d.cost_usd) .sum() } /// Get this month's cost pub fn get_month_cost(&self) -> f64 { let month_start = Self::month_start(); self.daily_costs .values() .filter(|d| { NaiveDate::parse_from_str(&d.date, "%Y-%m-%d") .map(|date| date >= month_start) .unwrap_or(false) }) .map(|d| d.cost_usd) .sum() } /// Get cost summary for a date range pub fn get_summary(&self, days: u32) -> CostSummary { let today = Local::now().date_naive(); let start_date = today - chrono::Duration::days(days as i64 - 1); let mut total_input_tokens = 0u64; let mut total_output_tokens = 0u64; let mut total_cost = 0.0f64; let mut total_messages = 0u64; let mut total_sessions = 0u64; let mut daily_breakdown = Vec::new(); for i in 0..days { let date = start_date + chrono::Duration::days(i as i64); let date_str = date.format("%Y-%m-%d").to_string(); if let Some(daily) = self.daily_costs.get(&date_str) { total_input_tokens += daily.input_tokens; total_output_tokens += daily.output_tokens; total_cost += daily.cost_usd; total_messages += daily.messages_sent; total_sessions += daily.sessions_count; daily_breakdown.push(daily.clone()); } else { daily_breakdown.push(DailyCost { date: date_str, ..Default::default() }); } } CostSummary { period_days: days, total_input_tokens, total_output_tokens, total_cost, total_messages, total_sessions, average_daily_cost: if days > 0 { total_cost / days as f64 } else { 0.0 }, daily_breakdown, } } /// Check if any alert thresholds are exceeded and return which ones pub fn check_alerts(&mut self) -> Vec { let mut alerts = Vec::new(); if let Some(threshold) = self.daily_alert_threshold { let today_cost = self.get_today_cost(); if today_cost >= threshold && !self.daily_alert_triggered { self.daily_alert_triggered = true; alerts.push(CostAlert { alert_type: AlertType::Daily, threshold, current_cost: today_cost, }); } } if let Some(threshold) = self.weekly_alert_threshold { let week_cost = self.get_week_cost(); if week_cost >= threshold && !self.weekly_alert_triggered { self.weekly_alert_triggered = true; alerts.push(CostAlert { alert_type: AlertType::Weekly, threshold, current_cost: week_cost, }); } } if let Some(threshold) = self.monthly_alert_threshold { let month_cost = self.get_month_cost(); if month_cost >= threshold && !self.monthly_alert_triggered { self.monthly_alert_triggered = true; alerts.push(CostAlert { alert_type: AlertType::Monthly, threshold, current_cost: month_cost, }); } } alerts } /// Set alert thresholds pub fn set_alert_thresholds( &mut self, daily: Option, weekly: Option, monthly: Option, ) { self.daily_alert_threshold = daily; self.weekly_alert_threshold = weekly; self.monthly_alert_threshold = monthly; } /// Clean up old data (keep last N days) #[allow(dead_code)] pub fn cleanup_old_data(&mut self, keep_days: u32) { let cutoff = Local::now().date_naive() - chrono::Duration::days(keep_days as i64); self.daily_costs.retain(|date_str, _| { NaiveDate::parse_from_str(date_str, "%Y-%m-%d") .map(|date| date >= cutoff) .unwrap_or(false) }); } /// Export to CSV format pub fn export_csv(&self, days: u32) -> String { let summary = self.get_summary(days); let mut csv = String::from("Date,Input Tokens,Output Tokens,Cost (USD),Messages,Sessions\n"); for daily in &summary.daily_breakdown { csv.push_str(&format!( "{},{},{},{:.4},{},{}\n", daily.date, daily.input_tokens, daily.output_tokens, daily.cost_usd, daily.messages_sent, daily.sessions_count )); } // Add totals row csv.push_str(&format!( "TOTAL,{},{},{:.4},{},{}\n", summary.total_input_tokens, summary.total_output_tokens, summary.total_cost, summary.total_messages, summary.total_sessions )); csv } } /// Cost summary for a period #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CostSummary { pub period_days: u32, pub total_input_tokens: u64, pub total_output_tokens: u64, pub total_cost: f64, pub total_messages: u64, pub total_sessions: u64, pub average_daily_cost: f64, pub daily_breakdown: Vec, } /// Alert types #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub enum AlertType { Daily, Weekly, Monthly, } /// Cost alert notification #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CostAlert { pub alert_type: AlertType, pub threshold: f64, pub current_cost: f64, } #[cfg(test)] mod tests { use super::*; #[test] fn test_add_cost() { let mut history = CostHistory::new(); history.add_cost(1000, 500, 0.05); let today_cost = history.get_today_cost(); assert!((today_cost - 0.05).abs() < 0.0001); } #[test] fn test_accumulate_daily_cost() { let mut history = CostHistory::new(); history.add_cost(1000, 500, 0.05); history.add_cost(2000, 1000, 0.10); let today_cost = history.get_today_cost(); assert!((today_cost - 0.15).abs() < 0.0001); } #[test] fn test_summary() { let mut history = CostHistory::new(); history.add_cost(1000, 500, 0.05); let summary = history.get_summary(7); assert_eq!(summary.period_days, 7); assert!((summary.total_cost - 0.05).abs() < 0.0001); } #[test] fn test_daily_alert() { let mut history = CostHistory::new(); history.set_alert_thresholds(Some(0.10), None, None); history.add_cost(1000, 500, 0.05); let alerts = history.check_alerts(); assert!(alerts.is_empty()); history.add_cost(1000, 500, 0.06); let alerts = history.check_alerts(); assert_eq!(alerts.len(), 1); assert_eq!(alerts[0].alert_type, AlertType::Daily); } #[test] fn test_alert_only_triggers_once() { let mut history = CostHistory::new(); history.set_alert_thresholds(Some(0.10), None, None); history.add_cost(1000, 500, 0.15); let alerts = history.check_alerts(); assert_eq!(alerts.len(), 1); // Second check should not trigger again let alerts = history.check_alerts(); assert!(alerts.is_empty()); } #[test] fn test_export_csv() { let mut history = CostHistory::new(); history.add_cost(1000, 500, 0.05); let csv = history.export_csv(1); assert!(csv.contains("Date,Input Tokens")); assert!(csv.contains("TOTAL")); } #[test] fn test_increment_sessions() { let mut history = CostHistory::new(); history.increment_sessions(); history.increment_sessions(); let summary = history.get_summary(1); assert_eq!(summary.total_sessions, 2); } }