From dfd037acace27258cd1417669e7f281148958db9 Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Sat, 26 Sep 2026 13:46:27 +0530 Subject: [PATCH] feat: implement account-wide budget tracker and rate limiter --- src/budget.rs | 241 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 + 2 files changed, 243 insertions(+) create mode 100644 src/budget.rs diff --git a/src/budget.rs b/src/budget.rs new file mode 100644 index 0000000..3083c11 --- /dev/null +++ b/src/budget.rs @@ -0,0 +1,241 @@ +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; + +#[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(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?; + + // 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); + } +} diff --git a/src/lib.rs b/src/lib.rs index f161cf8..59f87f3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,7 @@ //! BrokeCrew: A Rust-Native Multi-Model AI Team Library +pub mod budget; pub mod persona; +pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier}; pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};