use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; use std::collections::VecDeque; use std::fs; use std::io::Write; use std::path::{Path, PathBuf}; const MINUTE_LIMIT: usize = 20; /// Estimates token count from text using standard heuristic (~4 characters per token). pub fn estimate_tokens(text: &str) -> usize { if text.is_empty() { return 0; } text.len().div_ceil(4).max(1) } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum CreditTier { LowBalance, // < $10 lifetime purchased: 50 requests/day Funded, // >= $10 lifetime purchased: 1,000 requests/day } impl CreditTier { pub fn daily_limit(&self) -> usize { match self { CreditTier::LowBalance => 50, CreditTier::Funded => 1_000, } } } #[derive(Debug, Clone, PartialEq, Eq)] pub struct BudgetStatus { pub minute_requests: usize, pub minute_limit: usize, pub daily_requests: usize, pub daily_limit: usize, pub approaching_daily_limit: bool, } #[derive(thiserror::Error, Debug, PartialEq, Eq)] pub enum BudgetExceeded { #[error("Rate limited: exceeded {0} requests in 1 minute. Please wait {1}s before retrying.")] MinuteLimit(usize, u64), #[error("Daily budget exhausted: reached {0} of {1} requests today.")] DailyQuota(usize, usize), } #[derive(Debug, Serialize, Deserialize, Default)] struct StoredBudgetState { daily_count: usize, last_day_utc: String, // YYYY-MM-DD recent_calls: VecDeque>, } pub struct BudgetTracker { storage_path: Option, credit_tier: CreditTier, daily_count: usize, last_day_utc: String, recent_calls: VecDeque>, } impl BudgetTracker { pub fn new(storage_path: Option, credit_tier: CreditTier) -> Self { let mut tracker = Self { storage_path, credit_tier, daily_count: 0, last_day_utc: Utc::now().format("%Y-%m-%d").to_string(), recent_calls: VecDeque::new(), }; if let Some(ref path) = tracker.storage_path.clone() { let _ = tracker.load_from_disk(path); } tracker.refresh_window(); tracker } fn today_utc() -> String { Utc::now().format("%Y-%m-%d").to_string() } fn refresh_window(&mut self) { let now = Utc::now(); let current_day = Self::today_utc(); // Day rollover check if self.last_day_utc != current_day { self.daily_count = 0; self.last_day_utc = current_day; } // Sliding 60-second window eviction let cutoff = now - Duration::seconds(60); while let Some(&first) = self.recent_calls.front() { if first < cutoff { self.recent_calls.pop_front(); } else { break; } } } /// Checks the current budget status against OpenRouter rate limits. /// Returns `BudgetStatus` if permitted, or `BudgetExceeded` if limit reached. pub fn check_budget(&mut self) -> Result { self.refresh_window(); let daily_limit = self.credit_tier.daily_limit(); // Check daily ceiling if self.daily_count >= daily_limit { return Err(BudgetExceeded::DailyQuota(self.daily_count, daily_limit)); } // Check 1-minute rolling ceiling if self.recent_calls.len() >= MINUTE_LIMIT { let oldest = self.recent_calls.front().copied().unwrap_or_else(Utc::now); let elapsed = (Utc::now() - oldest).num_seconds().max(0); let wait_secs = (60 - elapsed).max(1) as u64; return Err(BudgetExceeded::MinuteLimit(MINUTE_LIMIT, wait_secs)); } let approaching_daily_limit = self.daily_count >= (daily_limit * 80) / 100; Ok(BudgetStatus { minute_requests: self.recent_calls.len(), minute_limit: MINUTE_LIMIT, daily_requests: self.daily_count, daily_limit, approaching_daily_limit, }) } /// Records an outbound call and persists updated state to disk atomically. pub fn record_call(&mut self, _is_paid: bool) -> Result<(), std::io::Error> { self.refresh_window(); let now = Utc::now(); self.recent_calls.push_back(now); self.daily_count += 1; if let Some(ref path) = self.storage_path.clone() { self.save_to_disk(path)?; } Ok(()) } fn load_from_disk(&mut self, path: &Path) -> Result<(), std::io::Error> { if !path.exists() { return Ok(()); } let data = fs::read_to_string(path)?; if let Ok(state) = serde_json::from_str::(&data) { self.daily_count = state.daily_count; self.last_day_utc = state.last_day_utc; self.recent_calls = state.recent_calls; } Ok(()) } fn save_to_disk(&self, path: &Path) -> Result<(), std::io::Error> { if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; } let state = StoredBudgetState { daily_count: self.daily_count, last_day_utc: self.last_day_utc.clone(), recent_calls: self.recent_calls.clone(), }; let serialized = serde_json::to_string_pretty(&state).map_err(std::io::Error::other)?; // Atomic write via tempfile in the same parent directory let parent_dir = path.parent().unwrap_or_else(|| Path::new(".")); let mut temp_file = tempfile::NamedTempFile::new_in(parent_dir)?; temp_file.write_all(serialized.as_bytes())?; temp_file.persist(path).map_err(|e| e.error)?; Ok(()) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_budget_minute_window_and_daily_limit() { let mut tracker = BudgetTracker::new(None, CreditTier::LowBalance); // Check budget initially let status = tracker.check_budget().expect("initial check should pass"); assert_eq!(status.minute_requests, 0); assert_eq!(status.minute_limit, 20); assert_eq!(status.daily_requests, 0); assert_eq!(status.daily_limit, 50); assert!(!status.approaching_daily_limit); // Record 19 calls for _ in 0..19 { tracker.record_call(false).unwrap(); } let status = tracker.check_budget().expect("19 calls should pass"); assert_eq!(status.minute_requests, 19); // 20th call tracker.record_call(false).unwrap(); // 21st check should fail minute limit let err = tracker.check_budget().unwrap_err(); match err { BudgetExceeded::MinuteLimit(limit, _) => assert_eq!(limit, 20), _ => panic!("Expected MinuteLimit error, got {:?}", err), } } #[test] fn test_budget_daily_ceiling_and_approaching_warning() { let mut tracker = BudgetTracker::new(None, CreditTier::LowBalance); tracker.daily_count = 40; // 80% of 50 let status = tracker.check_budget().unwrap(); assert!(status.approaching_daily_limit); tracker.daily_count = 50; let err = tracker.check_budget().unwrap_err(); assert_eq!(err, BudgetExceeded::DailyQuota(50, 50)); } #[test] fn test_budget_persistence() { let temp_dir = tempfile::tempdir().unwrap(); let file_path = temp_dir.path().join("budget.json"); { let mut tracker = BudgetTracker::new(Some(file_path.clone()), CreditTier::LowBalance); tracker.record_call(false).unwrap(); tracker.record_call(false).unwrap(); } // Reload in new tracker instance let mut reloaded = BudgetTracker::new(Some(file_path), CreditTier::LowBalance); let status = reloaded.check_budget().unwrap(); assert_eq!(status.daily_requests, 2); assert_eq!(status.minute_requests, 2); } }