aether_core/context/
token_tracker.rs

1use llm::{ContextUsage, TokenUsage, Tokens};
2
3/// Default threshold for triggering context compaction (85%)
4pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.85;
5
6/// One agent's context window, tracked from real LLM API usage rather than
7/// estimation. Owns the `ContextUsage` snapshot the agent publishes and
8/// evaluates the compaction policy against it.
9#[derive(Debug, Clone, Default)]
10pub struct TokenTracker {
11    snapshot: ContextUsage,
12}
13
14impl TokenTracker {
15    pub fn new(context_limit: Option<u32>) -> Self {
16        Self { snapshot: ContextUsage { context_limit: context_limit.map(Into::into), ..ContextUsage::default() } }
17    }
18
19    /// Record usage from an LLM API response.
20    pub fn record_usage(&mut self, sample: TokenUsage) {
21        self.snapshot.input_tokens = sample.input_tokens;
22        self.refresh_ratio();
23    }
24
25    /// The current window as published on `ContextEvent::UsageUpdated`.
26    pub fn snapshot(&self) -> &ContextUsage {
27        &self.snapshot
28    }
29
30    /// Current context usage as a ratio (0.0 - 1.0)
31    pub fn usage_ratio(&self) -> Option<f64> {
32        self.snapshot.usage_ratio
33    }
34
35    /// Whether current usage exceeds the given threshold
36    pub fn exceeds_threshold(&self, threshold: f64) -> bool {
37        self.usage_ratio().is_some_and(|ratio| ratio >= threshold)
38    }
39
40    /// Whether the context needs compaction
41    pub fn needs_compaction(&self, estimated_tokens: u32, threshold: f64) -> bool {
42        self.snapshot.context_limit.is_some_and(|limit| {
43            f64::from(self.snapshot.input_tokens.max(estimated_tokens.into())) >= f64::from(limit) * threshold
44        })
45    }
46
47    /// Tokens remaining before hitting limit
48    pub fn tokens_remaining(&self) -> Option<Tokens> {
49        self.snapshot.context_limit.map(|limit| limit.saturating_sub(self.snapshot.input_tokens))
50    }
51
52    /// Update the context limit (e.g. when switching models)
53    pub fn set_context_limit(&mut self, limit: Option<u32>) {
54        self.snapshot.context_limit = limit.map(Into::into);
55        self.refresh_ratio();
56    }
57
58    /// Forget the last call after context compaction so it cannot immediately
59    /// re-trigger compaction.
60    pub fn reset_current_usage(&mut self) {
61        self.snapshot.input_tokens = Tokens::ZERO;
62        self.refresh_ratio();
63    }
64
65    fn refresh_ratio(&mut self) {
66        self.snapshot.usage_ratio = self
67            .snapshot
68            .context_limit
69            .filter(|limit| !limit.is_zero())
70            .map(|limit| f64::from(self.snapshot.input_tokens) / f64::from(limit));
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_usage_tracking() {
80        let mut tracker = TokenTracker::new(Some(1000));
81
82        tracker.record_usage(TokenUsage::new(500, 100));
83        assert_eq!(tracker.usage_ratio(), Some(0.5));
84        assert!(!tracker.exceeds_threshold(0.85));
85
86        tracker.record_usage(TokenUsage::new(900, 50));
87        assert_eq!(tracker.usage_ratio(), Some(0.9));
88        assert!(tracker.exceeds_threshold(0.85));
89    }
90
91    #[test]
92    fn test_tokens_remaining() {
93        let mut tracker = TokenTracker::new(Some(1000));
94        tracker.record_usage(TokenUsage::new(700, 50));
95        assert_eq!(tracker.tokens_remaining().map(Tokens::get), Some(300));
96    }
97
98    #[test]
99    fn test_snapshot_reflects_last_call() {
100        let mut tracker = TokenTracker::new(Some(1000));
101        tracker.record_usage(TokenUsage::new(100, 50));
102        tracker.record_usage(TokenUsage::new(200, 60));
103
104        assert_eq!(tracker.snapshot().input_tokens.get(), 200);
105    }
106
107    #[test]
108    fn test_unknown_context_limit() {
109        let tracker = TokenTracker::new(None);
110        assert_eq!(tracker.usage_ratio(), None);
111        assert_eq!(tracker.tokens_remaining(), None);
112        assert!(!tracker.needs_compaction(1_000_000, 0.85));
113    }
114
115    #[test]
116    fn test_exceeds_threshold() {
117        let mut tracker = TokenTracker::new(Some(1000));
118
119        tracker.record_usage(TokenUsage::new(500, 100));
120        assert!(!tracker.exceeds_threshold(0.6));
121        assert!(tracker.exceeds_threshold(0.5));
122
123        tracker.record_usage(TokenUsage::new(850, 50));
124        assert!(tracker.exceeds_threshold(0.8));
125        assert!(tracker.exceeds_threshold(0.85));
126    }
127
128    #[test]
129    fn test_needs_compaction_from_recorded_usage() {
130        let mut tracker = TokenTracker::new(Some(10000));
131
132        tracker.record_usage(TokenUsage::new(9000, 100));
133        assert!(tracker.needs_compaction(0, 0.85));
134
135        tracker.record_usage(TokenUsage::new(7000, 100));
136        assert!(!tracker.needs_compaction(0, 0.85));
137    }
138
139    #[test]
140    fn test_needs_compaction_from_estimate_before_usage_recorded() {
141        let tracker = TokenTracker::new(Some(10000));
142
143        assert!(tracker.needs_compaction(9000, 0.85));
144        assert!(!tracker.needs_compaction(1000, 0.85));
145    }
146
147    #[test]
148    fn test_default_compaction_threshold() {
149        use super::DEFAULT_COMPACTION_THRESHOLD;
150        assert!((DEFAULT_COMPACTION_THRESHOLD - 0.85).abs() < 0.001);
151    }
152
153    #[test]
154    fn test_set_context_limit() {
155        let mut tracker = TokenTracker::new(Some(200_000));
156        assert_eq!(tracker.snapshot().context_limit.map(Tokens::get), Some(200_000));
157
158        tracker.set_context_limit(Some(128_000));
159        assert_eq!(tracker.snapshot().context_limit.map(Tokens::get), Some(128_000));
160
161        tracker.record_usage(TokenUsage::new(100_000, 50));
162        let expected_ratio = 100_000.0 / 128_000.0;
163        assert!((tracker.usage_ratio().unwrap_or_default() - expected_ratio).abs() < 0.001);
164    }
165
166    #[test]
167    fn test_reset_current_usage() {
168        let mut tracker = TokenTracker::new(Some(10000));
169        tracker.record_usage(TokenUsage::new(9000, 100));
170
171        assert!(tracker.needs_compaction(0, 0.85));
172
173        tracker.reset_current_usage();
174
175        assert!(tracker.snapshot().input_tokens.is_zero());
176        assert_eq!(tracker.usage_ratio(), Some(0.0));
177        assert!(!tracker.needs_compaction(0, 0.85));
178    }
179}
aether_core/context/token_tracker.rs

aether_core/context/
token_tracker.rs
