agentic_evolve_core/query/
budget.rs

1//! Token budget — track and enforce a per-request token spending limit.
2
3use serde::{Deserialize, Serialize};
4
5/// A token budget that tracks spending against a maximum allocation.
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct TokenBudget {
8    max_tokens: u64,
9    used_tokens: u64,
10}
11
12impl TokenBudget {
13    /// Create a new budget with the given maximum.
14    pub fn new(max_tokens: u64) -> Self {
15        Self {
16            max_tokens,
17            used_tokens: 0,
18        }
19    }
20
21    /// Maximum tokens allowed.
22    pub fn max_tokens(&self) -> u64 {
23        self.max_tokens
24    }
25
26    /// Tokens consumed so far.
27    pub fn used_tokens(&self) -> u64 {
28        self.used_tokens
29    }
30
31    /// Spend tokens from the budget. Returns `true` if the budget was
32    /// sufficient, `false` if it would exceed the limit (in which case
33    /// the spend is still applied to track overruns).
34    pub fn spend(&mut self, tokens: u64) -> bool {
35        self.used_tokens += tokens;
36        self.used_tokens <= self.max_tokens
37    }
38
39    /// Remaining tokens before the budget is exhausted.
40    pub fn remaining(&self) -> u64 {
41        self.max_tokens.saturating_sub(self.used_tokens)
42    }
43
44    /// Whether the budget is fully consumed.
45    pub fn is_exhausted(&self) -> bool {
46        self.used_tokens >= self.max_tokens
47    }
48
49    /// Check if a given cost can be afforded without exceeding the budget.
50    pub fn can_afford(&self, cost: u64) -> bool {
51        self.used_tokens + cost <= self.max_tokens
52    }
53
54    /// Utilization ratio in `[0.0, 1.0+]` (can exceed 1.0 if overspent).
55    pub fn utilization(&self) -> f64 {
56        if self.max_tokens == 0 {
57            return if self.used_tokens == 0 {
58                0.0
59            } else {
60                f64::INFINITY
61            };
62        }
63        self.used_tokens as f64 / self.max_tokens as f64
64    }
65
66    /// Reset the budget to zero usage.
67    pub fn reset(&mut self) {
68        self.used_tokens = 0;
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn new_budget_is_empty() {
78        let b = TokenBudget::new(100);
79        assert_eq!(b.used_tokens(), 0);
80        assert_eq!(b.remaining(), 100);
81        assert!(!b.is_exhausted());
82    }
83
84    #[test]
85    fn spend_reduces_remaining() {
86        let mut b = TokenBudget::new(100);
87        assert!(b.spend(40));
88        assert_eq!(b.remaining(), 60);
89        assert_eq!(b.used_tokens(), 40);
90    }
91
92    #[test]
93    fn spend_returns_false_when_exceeding() {
94        let mut b = TokenBudget::new(10);
95        assert!(!b.spend(20));
96        assert!(b.is_exhausted());
97    }
98
99    #[test]
100    fn can_afford_check() {
101        let mut b = TokenBudget::new(100);
102        assert!(b.can_afford(100));
103        assert!(!b.can_afford(101));
104        b.spend(50);
105        assert!(b.can_afford(50));
106        assert!(!b.can_afford(51));
107    }
108
109    #[test]
110    fn utilization_tracks_ratio() {
111        let mut b = TokenBudget::new(100);
112        assert_eq!(b.utilization(), 0.0);
113        b.spend(50);
114        assert!((b.utilization() - 0.5).abs() < f64::EPSILON);
115        b.spend(50);
116        assert!((b.utilization() - 1.0).abs() < f64::EPSILON);
117    }
118
119    #[test]
120    fn reset_clears_usage() {
121        let mut b = TokenBudget::new(100);
122        b.spend(80);
123        b.reset();
124        assert_eq!(b.used_tokens(), 0);
125        assert_eq!(b.remaining(), 100);
126    }
127
128    #[test]
129    fn zero_budget_exhausted() {
130        let b = TokenBudget::new(0);
131        assert!(b.is_exhausted());
132        assert!(!b.can_afford(1));
133    }
134}
agentic_evolve_core/query/budget.rs

agentic_evolve_core/query/
budget.rs
