aether_core/context/
token_tracker.rs

1use llm::TokenUsage;
2
3/// Default threshold for triggering context compaction (85%)
4pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.85;
5
6/// Tracks token usage from LLM API responses.
7/// Uses real usage data from API, not estimation.
8///
9/// Cumulative totals are stored for the dimensions consumers care about today
10/// (input/output, cache read/creation, reasoning). The `last_usage` field
11/// preserves the full `TokenUsage` from the most recent API call, so audio /
12/// video / prediction dimensions are still accessible without growing
13/// dedicated accumulators until a consumer asks for them.
14#[derive(Debug, Clone, Default)]
15pub struct TokenTracker {
16    total_input_tokens: u64,
17    total_output_tokens: u64,
18    total_cache_read_tokens: u64,
19    total_cache_creation_tokens: u64,
20    total_reasoning_tokens: u64,
21    last_usage: TokenUsage,
22    context_limit: Option<u32>,
23}
24
25impl TokenTracker {
26    pub fn new(context_limit: Option<u32>) -> Self {
27        Self { context_limit, ..Self::default() }
28    }
29
30    /// Record usage from an LLM API response.
31    pub fn record_usage(&mut self, sample: TokenUsage) {
32        self.total_input_tokens += u64::from(sample.input_tokens);
33        self.total_output_tokens += u64::from(sample.output_tokens);
34        self.total_cache_read_tokens += u64::from(sample.cache_read_tokens.unwrap_or(0));
35        self.total_cache_creation_tokens += u64::from(sample.cache_creation_tokens.unwrap_or(0));
36        self.total_reasoning_tokens += u64::from(sample.reasoning_tokens.unwrap_or(0));
37        self.last_usage = sample;
38    }
39
40    /// Current context usage as a ratio (0.0 - 1.0)
41    pub fn usage_ratio(&self) -> Option<f64> {
42        let context_limit = self.context_limit?;
43        if context_limit == 0 {
44            return None;
45        }
46        Some(f64::from(self.last_usage.input_tokens) / f64::from(context_limit))
47    }
48
49    /// Whether current usage exceeds the given threshold
50    pub fn exceeds_threshold(&self, threshold: f64) -> bool {
51        self.usage_ratio().is_some_and(|ratio| ratio >= threshold)
52    }
53
54    /// Whether the context needs compaction
55    pub fn needs_compaction(&self, estimated_tokens: u32, threshold: f64) -> bool {
56        self.context_limit.is_some_and(|limit| {
57            f64::from(self.last_usage.input_tokens.max(estimated_tokens)) >= f64::from(limit) * threshold
58        })
59    }
60
61    /// Tokens remaining before hitting limit
62    pub fn tokens_remaining(&self) -> Option<u32> {
63        self.context_limit.map(|context_limit| context_limit.saturating_sub(self.last_usage.input_tokens))
64    }
65
66    /// Update the context limit (e.g. when switching models)
67    pub fn set_context_limit(&mut self, limit: Option<u32>) {
68        self.context_limit = limit;
69    }
70
71    /// Get the context limit
72    pub fn context_limit(&self) -> Option<u32> {
73        self.context_limit
74    }
75
76    /// Get last recorded input tokens (current context size)
77    pub fn last_input_tokens(&self) -> u32 {
78        self.last_usage.input_tokens
79    }
80
81    /// Get the full `TokenUsage` from the most recent API call. Returns the
82    /// default (all zeros / `None`) before any call has been recorded.
83    pub fn last_usage(&self) -> &TokenUsage {
84        &self.last_usage
85    }
86
87    /// Get total input tokens across all calls
88    pub fn total_input_tokens(&self) -> u64 {
89        self.total_input_tokens
90    }
91
92    /// Get total output tokens across all calls
93    pub fn total_output_tokens(&self) -> u64 {
94        self.total_output_tokens
95    }
96
97    /// Get total cache-read tokens across all calls
98    pub fn total_cache_read_tokens(&self) -> u64 {
99        self.total_cache_read_tokens
100    }
101
102    /// Get total cache-creation tokens across all calls
103    pub fn total_cache_creation_tokens(&self) -> u64 {
104        self.total_cache_creation_tokens
105    }
106
107    /// Get total reasoning tokens across all calls
108    pub fn total_reasoning_tokens(&self) -> u64 {
109        self.total_reasoning_tokens
110    }
111
112    /// Reset current usage tracking after context compaction.
113    /// Preserves cumulative totals for metrics while clearing `last_usage` to
114    /// prevent immediate re-triggering of compaction.
115    pub fn reset_current_usage(&mut self) {
116        self.last_usage = TokenUsage::default();
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_usage_tracking() {
126        let mut tracker = TokenTracker::new(Some(1000));
127
128        tracker.record_usage(TokenUsage::new(500, 100));
129        assert_eq!(tracker.usage_ratio(), Some(0.5));
130        assert!(!tracker.exceeds_threshold(0.85));
131
132        tracker.record_usage(TokenUsage::new(900, 50));
133        assert_eq!(tracker.usage_ratio(), Some(0.9));
134        assert!(tracker.exceeds_threshold(0.85));
135    }
136
137    #[test]
138    fn test_tokens_remaining() {
139        let mut tracker = TokenTracker::new(Some(1000));
140        tracker.record_usage(TokenUsage::new(700, 50));
141        assert_eq!(tracker.tokens_remaining(), Some(300));
142    }
143
144    #[test]
145    fn test_cumulative_totals() {
146        let mut tracker = TokenTracker::new(Some(1000));
147        tracker.record_usage(TokenUsage::new(100, 50));
148        tracker.record_usage(TokenUsage::new(200, 60));
149
150        assert_eq!(tracker.total_input_tokens(), 300);
151        assert_eq!(tracker.total_output_tokens(), 110);
152        assert_eq!(tracker.last_input_tokens(), 200);
153    }
154
155    #[test]
156    fn test_unknown_context_limit() {
157        let tracker = TokenTracker::new(None);
158        assert_eq!(tracker.usage_ratio(), None);
159        assert_eq!(tracker.tokens_remaining(), None);
160        assert!(!tracker.needs_compaction(1_000_000, 0.85));
161    }
162
163    #[test]
164    fn test_exceeds_threshold() {
165        let mut tracker = TokenTracker::new(Some(1000));
166
167        tracker.record_usage(TokenUsage::new(500, 100));
168        assert!(!tracker.exceeds_threshold(0.6));
169        assert!(tracker.exceeds_threshold(0.5));
170
171        tracker.record_usage(TokenUsage::new(850, 50));
172        assert!(tracker.exceeds_threshold(0.8));
173        assert!(tracker.exceeds_threshold(0.85));
174    }
175
176    #[test]
177    fn test_needs_compaction_from_recorded_usage() {
178        let mut tracker = TokenTracker::new(Some(10000));
179
180        tracker.record_usage(TokenUsage::new(9000, 100));
181        assert!(tracker.needs_compaction(0, 0.85));
182
183        tracker.record_usage(TokenUsage::new(7000, 100));
184        assert!(!tracker.needs_compaction(0, 0.85));
185    }
186
187    #[test]
188    fn test_needs_compaction_from_estimate_before_usage_recorded() {
189        let tracker = TokenTracker::new(Some(10000));
190
191        assert!(tracker.needs_compaction(9000, 0.85));
192        assert!(!tracker.needs_compaction(1000, 0.85));
193    }
194
195    #[test]
196    fn test_default_compaction_threshold() {
197        use super::DEFAULT_COMPACTION_THRESHOLD;
198        assert!((DEFAULT_COMPACTION_THRESHOLD - 0.85).abs() < 0.001);
199    }
200
201    #[test]
202    fn test_set_context_limit() {
203        let mut tracker = TokenTracker::new(Some(200_000));
204        assert_eq!(tracker.context_limit(), Some(200_000));
205
206        tracker.set_context_limit(Some(128_000));
207        assert_eq!(tracker.context_limit(), Some(128_000));
208
209        // Verify usage ratio recalculates against new limit
210        tracker.record_usage(TokenUsage::new(100_000, 50));
211        let expected_ratio = 100_000.0 / 128_000.0;
212        assert!((tracker.usage_ratio().unwrap_or_default() - expected_ratio).abs() < 0.001);
213    }
214
215    #[test]
216    fn test_reset_current_usage() {
217        let mut tracker = TokenTracker::new(Some(10000));
218        tracker.record_usage(TokenUsage::new(9000, 100));
219
220        assert!(tracker.needs_compaction(0, 0.85));
221
222        tracker.reset_current_usage();
223
224        assert_eq!(tracker.last_input_tokens(), 0);
225        assert!(!tracker.needs_compaction(0, 0.85));
226        assert_eq!(tracker.total_input_tokens(), 9000);
227        assert_eq!(tracker.total_output_tokens(), 100);
228    }
229
230    #[test]
231    fn test_cache_and_reasoning_totals_accumulate() {
232        let mut tracker = TokenTracker::new(Some(10000));
233
234        tracker.record_usage(TokenUsage {
235            input_tokens: 500,
236            output_tokens: 100,
237            cache_read_tokens: Some(200),
238            cache_creation_tokens: Some(50),
239            reasoning_tokens: Some(30),
240            ..TokenUsage::default()
241        });
242        tracker.record_usage(TokenUsage {
243            input_tokens: 600,
244            output_tokens: 80,
245            cache_read_tokens: Some(300),
246            cache_creation_tokens: None,
247            reasoning_tokens: Some(20),
248            ..TokenUsage::default()
249        });
250
251        assert_eq!(tracker.total_cache_read_tokens(), 500);
252        assert_eq!(tracker.total_cache_creation_tokens(), 50);
253        assert_eq!(tracker.total_reasoning_tokens(), 50);
254    }
255
256    #[test]
257    fn test_last_usage_exposes_full_token_usage() {
258        let mut tracker = TokenTracker::new(Some(10000));
259        let sample = TokenUsage {
260            input_tokens: 500,
261            output_tokens: 100,
262            cache_read_tokens: Some(200),
263            cache_creation_tokens: Some(50),
264            reasoning_tokens: Some(30),
265            input_audio_tokens: Some(5),
266            ..TokenUsage::default()
267        };
268
269        tracker.record_usage(sample);
270
271        assert_eq!(*tracker.last_usage(), sample);
272    }
273
274    #[test]
275    fn test_reset_clears_last_usage_but_keeps_cache_totals() {
276        let mut tracker = TokenTracker::new(Some(10000));
277        tracker.record_usage(TokenUsage {
278            input_tokens: 500,
279            output_tokens: 100,
280            cache_read_tokens: Some(200),
281            cache_creation_tokens: Some(50),
282            reasoning_tokens: Some(30),
283            ..TokenUsage::default()
284        });
285
286        tracker.reset_current_usage();
287
288        assert_eq!(*tracker.last_usage(), TokenUsage::default());
289        assert_eq!(tracker.total_cache_read_tokens(), 200);
290        assert_eq!(tracker.total_cache_creation_tokens(), 50);
291        assert_eq!(tracker.total_reasoning_tokens(), 30);
292    }
293}
aether_core/context/token_tracker.rs

aether_core/context/
token_tracker.rs
