llm_kernel/tokens/
budget.rs

1//! Token budget tracking for LLM context windows.
2//!
3//! Provides a thread-safe token accounting type for managing context window
4//! budgets without external dependencies.
5
6use core::sync::atomic::{AtomicU32, Ordering};
7
8/// Thread-safe token budget tracker for LLM context windows.
9///
10/// Tracks total capacity, current usage, and provides atomic reserve/release
11/// operations for concurrent token accounting.
12pub struct TokenBudget {
13    /// Total token capacity.
14    total: u32,
15    /// Currently used tokens (atomic for thread safety).
16    used: AtomicU32,
17}
18
19impl TokenBudget {
20    /// Create a new budget with the given total token capacity.
21    ///
22    /// Initial usage is zero.
23    pub fn new(total: u32) -> Self {
24        Self {
25            total,
26            used: AtomicU32::new(0),
27        }
28    }
29
30    /// Returns the total token capacity.
31    pub fn total(&self) -> u32 {
32        self.total
33    }
34
35    /// Returns the number of remaining tokens.
36    pub fn remaining(&self) -> u32 {
37        self.total.saturating_sub(self.used.load(Ordering::Relaxed))
38    }
39
40    /// Returns the number of currently used tokens.
41    pub fn used(&self) -> u32 {
42        self.used.load(Ordering::Relaxed)
43    }
44
45    /// Try to reserve `n` tokens from the budget.
46    ///
47    /// Returns `true` if the reservation succeeded (enough capacity remaining),
48    /// `false` if there were insufficient tokens.
49    pub fn try_reserve(&self, n: u32) -> bool {
50        loop {
51            let current = self.used.load(Ordering::Relaxed);
52            let new_used = current.saturating_add(n);
53            if new_used > self.total {
54                return false;
55            }
56            match self.used.compare_exchange_weak(
57                current,
58                new_used,
59                Ordering::Relaxed,
60                Ordering::Relaxed,
61            ) {
62                Ok(_) => return true,
63                Err(_) => continue,
64            }
65        }
66    }
67
68    /// Release `n` tokens back to the budget.
69    ///
70    /// Clamps to zero — cannot release more than currently used.
71    pub fn release(&self, n: u32) {
72        loop {
73            let current = self.used.load(Ordering::Relaxed);
74            let new_used = current.saturating_sub(n);
75            match self.used.compare_exchange_weak(
76                current,
77                new_used,
78                Ordering::Relaxed,
79                Ordering::Relaxed,
80            ) {
81                Ok(_) => return,
82                Err(_) => continue,
83            }
84        }
85    }
86
87    /// Reset the budget, setting used tokens back to zero.
88    pub fn reset(&self) {
89        self.used.store(0, Ordering::Relaxed);
90    }
91}
92
93impl std::fmt::Debug for TokenBudget {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        f.debug_struct("TokenBudget")
96            .field("total", &self.total)
97            .field("used", &self.used())
98            .field("remaining", &self.remaining())
99            .finish()
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn new_budget_starts_empty() {
109        let b = TokenBudget::new(1000);
110        assert_eq!(b.total(), 1000);
111        assert_eq!(b.used(), 0);
112        assert_eq!(b.remaining(), 1000);
113    }
114
115    #[test]
116    fn reserve_success() {
117        let b = TokenBudget::new(1000);
118        assert!(b.try_reserve(600));
119        assert_eq!(b.remaining(), 400);
120        assert_eq!(b.used(), 600);
121    }
122
123    #[test]
124    fn reserve_insufficient() {
125        let b = TokenBudget::new(1000);
126        assert!(b.try_reserve(600));
127        assert!(!b.try_reserve(500));
128        assert_eq!(b.remaining(), 400);
129    }
130
131    #[test]
132    fn release_tokens() {
133        let b = TokenBudget::new(1000);
134        b.try_reserve(600);
135        b.release(200);
136        assert_eq!(b.remaining(), 600);
137        assert_eq!(b.used(), 400);
138    }
139
140    #[test]
141    fn release_clamps_to_zero() {
142        let b = TokenBudget::new(1000);
143        b.try_reserve(100);
144        b.release(200);
145        assert_eq!(b.used(), 0);
146        assert_eq!(b.remaining(), 1000);
147    }
148
149    #[test]
150    fn reset_clears_usage() {
151        let b = TokenBudget::new(1000);
152        b.try_reserve(500);
153        b.reset();
154        assert_eq!(b.used(), 0);
155        assert_eq!(b.remaining(), 1000);
156    }
157
158    #[test]
159    fn debug_format() {
160        let b = TokenBudget::new(1000);
161        b.try_reserve(300);
162        let debug = format!("{:?}", b);
163        assert!(debug.contains("total: 1000"));
164        assert!(debug.contains("used: 300"));
165        assert!(debug.contains("remaining: 700"));
166    }
167
168    #[test]
169    fn reserve_exact_total() {
170        let b = TokenBudget::new(100);
171        assert!(b.try_reserve(100));
172        assert_eq!(b.remaining(), 0);
173    }
174
175    #[test]
176    fn reserve_zero_is_noop() {
177        let b = TokenBudget::new(100);
178        assert!(b.try_reserve(0));
179        assert_eq!(b.used(), 0);
180    }
181}
llm_kernel/tokens/budget.rs

llm_kernel/tokens/
budget.rs
