aether_core/context/
token_tracker.rs

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

aether_core/context/
token_tracker.rs
