claude_agent/tokens/
window.rs

1use super::tier::{DEFAULT_CRITICAL_THRESHOLD, DEFAULT_WARNING_THRESHOLD};
2use crate::models::{Capabilities, ModelSpec};
3
4#[derive(Debug, Clone, Copy, PartialEq)]
5pub enum WindowStatus {
6    Ok { utilization: f64, remaining: u64 },
7    Warning { utilization: f64, remaining: u64 },
8    Critical { utilization: f64, remaining: u64 },
9    Exceeded { overage: u64 },
10}
11
12impl WindowStatus {
13    pub fn should_proceed(&self) -> bool {
14        !matches!(self, Self::Exceeded { .. })
15    }
16
17    pub fn utilization(&self) -> Option<f64> {
18        match self {
19            Self::Ok { utilization, .. }
20            | Self::Warning { utilization, .. }
21            | Self::Critical { utilization, .. } => Some(*utilization),
22            Self::Exceeded { .. } => None,
23        }
24    }
25}
26
27#[derive(Debug, Clone)]
28pub struct ContextWindow {
29    capabilities: Capabilities,
30    extended_enabled: bool,
31    current_usage: u64,
32    peak_usage: u64,
33    warning_threshold: f64,
34    critical_threshold: f64,
35}
36
37impl ContextWindow {
38    pub fn new(spec: &ModelSpec, extended_enabled: bool) -> Self {
39        Self {
40            capabilities: spec.capabilities,
41            extended_enabled,
42            current_usage: 0,
43            peak_usage: 0,
44            warning_threshold: DEFAULT_WARNING_THRESHOLD,
45            critical_threshold: DEFAULT_CRITICAL_THRESHOLD,
46        }
47    }
48
49    pub fn limit(&self) -> u64 {
50        self.capabilities.effective_context(self.extended_enabled)
51    }
52
53    pub fn usage(&self) -> u64 {
54        self.current_usage
55    }
56
57    pub fn remaining(&self) -> u64 {
58        self.limit().saturating_sub(self.current_usage)
59    }
60
61    pub fn utilization(&self) -> f64 {
62        let limit = self.limit();
63        if limit == 0 {
64            return 0.0;
65        }
66        self.current_usage as f64 / limit as f64
67    }
68
69    pub fn status(&self) -> WindowStatus {
70        let limit = self.limit();
71        let utilization = self.utilization();
72
73        if self.current_usage > limit {
74            WindowStatus::Exceeded {
75                overage: self.current_usage - limit,
76            }
77        } else if utilization >= self.critical_threshold {
78            WindowStatus::Critical {
79                utilization,
80                remaining: self.remaining(),
81            }
82        } else if utilization >= self.warning_threshold {
83            WindowStatus::Warning {
84                utilization,
85                remaining: self.remaining(),
86            }
87        } else {
88            WindowStatus::Ok {
89                utilization,
90                remaining: self.remaining(),
91            }
92        }
93    }
94
95    pub fn can_fit(&self, additional: u64) -> bool {
96        self.current_usage + additional <= self.limit()
97    }
98
99    pub fn update(&mut self, new_usage: u64) {
100        self.current_usage = new_usage;
101        if new_usage > self.peak_usage {
102            self.peak_usage = new_usage;
103        }
104    }
105
106    pub fn add(&mut self, tokens: u64) {
107        self.update(self.current_usage.saturating_add(tokens));
108    }
109
110    pub fn reset(&mut self, new_usage: u64) {
111        self.current_usage = new_usage;
112    }
113
114    pub fn peak(&self) -> u64 {
115        self.peak_usage
116    }
117
118    pub fn warning_threshold(&self) -> f64 {
119        self.warning_threshold
120    }
121
122    pub fn with_thresholds(mut self, warning: f64, critical: f64) -> Self {
123        self.warning_threshold = warning;
124        self.critical_threshold = critical;
125        self
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use crate::models::read_registry;
133
134    #[test]
135    fn test_context_window_status() {
136        let reg = read_registry();
137        let spec = reg.resolve("sonnet").unwrap();
138        let mut window = ContextWindow::new(spec, false);
139
140        window.update(100_000);
141        assert!(matches!(window.status(), WindowStatus::Ok { .. }));
142
143        window.update(180_000);
144        assert!(matches!(window.status(), WindowStatus::Warning { .. }));
145
146        window.update(195_000);
147        assert!(matches!(window.status(), WindowStatus::Critical { .. }));
148
149        window.update(250_000);
150        assert!(matches!(window.status(), WindowStatus::Exceeded { .. }));
151    }
152
153    #[test]
154    fn test_extended_context() {
155        let reg = read_registry();
156        let spec = reg.resolve("sonnet").unwrap();
157
158        let standard = ContextWindow::new(spec, false);
159        assert_eq!(standard.limit(), 200_000);
160
161        let extended = ContextWindow::new(spec, true);
162        assert_eq!(extended.limit(), 1_000_000);
163    }
164}
claude_agent/tokens/window.rs

claude_agent/tokens/
window.rs
