agentic_forge_core/query/
budget.rs

1//! Token budget enforcement.
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct TokenBudget {
7    pub max_tokens: u64,
8    pub used_tokens: u64,
9}
10
11impl TokenBudget {
12    pub fn new(max_tokens: u64) -> Self {
13        Self {
14            max_tokens,
15            used_tokens: 0,
16        }
17    }
18
19    pub fn unlimited() -> Self {
20        Self {
21            max_tokens: u64::MAX,
22            used_tokens: 0,
23        }
24    }
25
26    pub fn remaining(&self) -> u64 {
27        self.max_tokens.saturating_sub(self.used_tokens)
28    }
29
30    pub fn is_exhausted(&self) -> bool {
31        self.used_tokens >= self.max_tokens
32    }
33
34    pub fn can_afford(&self, cost: u64) -> bool {
35        self.remaining() >= cost
36    }
37
38    pub fn spend(&mut self, tokens: u64) -> bool {
39        if self.can_afford(tokens) {
40            self.used_tokens += tokens;
41            true
42        } else {
43            false
44        }
45    }
46
47    pub fn force_spend(&mut self, tokens: u64) {
48        self.used_tokens += tokens;
49    }
50
51    pub fn utilization(&self) -> f64 {
52        if self.max_tokens == 0 || self.max_tokens == u64::MAX {
53            return 0.0;
54        }
55        self.used_tokens as f64 / self.max_tokens as f64
56    }
57}
58
59impl Default for TokenBudget {
60    fn default() -> Self {
61        Self::unlimited()
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_budget_new() {
71        let b = TokenBudget::new(1000);
72        assert_eq!(b.remaining(), 1000);
73        assert!(!b.is_exhausted());
74    }
75
76    #[test]
77    fn test_budget_spend() {
78        let mut b = TokenBudget::new(100);
79        assert!(b.spend(50));
80        assert_eq!(b.remaining(), 50);
81        assert!(b.spend(50));
82        assert!(b.is_exhausted());
83        assert!(!b.spend(1));
84    }
85
86    #[test]
87    fn test_budget_can_afford() {
88        let b = TokenBudget::new(100);
89        assert!(b.can_afford(100));
90        assert!(!b.can_afford(101));
91    }
92
93    #[test]
94    fn test_budget_unlimited() {
95        let b = TokenBudget::unlimited();
96        assert!(b.can_afford(u64::MAX - 1));
97        assert!(!b.is_exhausted());
98    }
99
100    #[test]
101    fn test_budget_utilization() {
102        let mut b = TokenBudget::new(200);
103        assert_eq!(b.utilization(), 0.0);
104        b.spend(100);
105        assert!((b.utilization() - 0.5).abs() < f64::EPSILON);
106    }
107
108    #[test]
109    fn test_budget_force_spend() {
110        let mut b = TokenBudget::new(10);
111        b.force_spend(100); // Over budget
112        assert!(b.is_exhausted());
113        assert_eq!(b.used_tokens, 100);
114    }
115}
agentic_forge_core/query/budget.rs

agentic_forge_core/query/
budget.rs
