use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenBudget {
pub total: u32,
pub used: u32,
pub reserved: u32,
}
impl TokenBudget {
pub fn new(total: u32) -> Self {
Self {
total,
used: 0,
reserved: 0,
}
}
pub fn remaining(&self) -> u32 {
self.total
.saturating_sub(self.used)
.saturating_sub(self.reserved)
}
pub fn available(&self) -> u32 {
self.total.saturating_sub(self.used)
}
pub fn consume(&mut self, amount: u32) {
self.used = self.used.saturating_add(amount);
}
pub fn reserve(&mut self, amount: u32) {
self.reserved = self.reserved.saturating_add(amount);
}
pub fn release(&mut self, amount: u32) {
self.reserved = self.reserved.saturating_sub(amount);
}
pub fn reset(&mut self) {
self.used = 0;
self.reserved = 0;
}
pub fn percent_used(&self) -> f64 {
if self.total == 0 {
return 0.0;
}
(self.used as f64 / self.total as f64) * 100.0
}
pub fn is_exhausted(&self) -> bool {
self.remaining() == 0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_token_budget() {
let mut budget = TokenBudget::new(1000);
assert_eq!(budget.remaining(), 1000);
budget.consume(300);
assert_eq!(budget.used, 300);
assert_eq!(budget.remaining(), 700);
assert!((budget.percent_used() - 30.0).abs() < 0.01);
}
}