claude_code_status_line/
context.rs

1use crate::colors::SectionColors;
2use crate::config::Config;
3use crate::render::get_details_and_fg_codes;
4use crate::types::{ContextUsageInfo, ContextWindow};
5
6/// Fallback context window size (only used if Claude Code doesn't provide it)
7const DEFAULT_CONTEXT_WINDOW: u64 = 200_000;
8
9pub fn calculate_context(
10    context_window: Option<&ContextWindow>,
11    config: &Config,
12) -> Option<ContextUsageInfo> {
13    let cw = context_window?;
14    let total = cw.context_window_size.unwrap_or(DEFAULT_CONTEXT_WINDOW);
15
16    let display_remaining = config.sections.context.display_mode == "remaining";
17
18    // Try to use direct percentages from Claude Code v2.1.6+
19    let direct_percentage = if display_remaining {
20        cw.remaining_percentage
21    } else {
22        cw.used_percentage
23    };
24
25    if let Some(pct) = direct_percentage {
26        // Calculate used tokens from used_percentage (always use used, not remaining)
27        let used_pct = cw.used_percentage.unwrap_or(0.0);
28        let used_tokens = ((used_pct / 100.0) * total as f64) as u64;
29
30        return Some(ContextUsageInfo {
31            percentage: pct.clamp(0.0, 100.0),
32            current_usage_tokens: used_tokens,
33            context_window_size: total,
34            is_exact: true,
35        });
36    }
37
38    // Fall back to calculated percentage for older Claude Code versions
39    if total == 0 {
40        return None;
41    }
42    let usage = cw.current_usage.as_ref()?;
43
44    let used_tokens = usage.input_tokens.unwrap_or(0)
45        + usage.cache_creation_input_tokens.unwrap_or(0)
46        + usage.cache_read_input_tokens.unwrap_or(0);
47
48    let used_pct = (used_tokens as f64 / total as f64 * 100.0).min(100.0);
49    let percentage = if display_remaining {
50        (100.0 - used_pct).max(0.0)
51    } else {
52        used_pct
53    };
54
55    Some(ContextUsageInfo {
56        percentage,
57        current_usage_tokens: used_tokens,
58        context_window_size: total,
59        is_exact: false,
60    })
61}
62
63pub fn format_ctx_short(pct: f64, is_exact: bool, config: &Config) -> String {
64    let prefix = if is_exact { "" } else { "~" };
65    if config.sections.context.show_decimals {
66        format!("ctx: {}{:.1}%", prefix, pct)
67    } else {
68        format!("ctx: {}{}%", prefix, pct as u64)
69    }
70}
71
72pub fn format_ctx_full(info: &ContextUsageInfo, colors: &SectionColors, config: &Config) -> String {
73    let (details, fg) = get_details_and_fg_codes(colors);
74    let prefix = if info.is_exact { "" } else { "~" };
75
76    let pct_str = if config.sections.context.show_decimals {
77        format!("{}{:.1}%", prefix, info.percentage)
78    } else {
79        format!("{}{}%", prefix, info.percentage as u64)
80    };
81
82    if config.sections.context.show_token_counts && info.current_usage_tokens > 0 {
83        let used_k = info.current_usage_tokens / 1000;
84        let total_k = info.context_window_size / 1000;
85
86        format!(
87            "ctx: {} {}({}k/{}k){}",
88            pct_str, details, used_k, total_k, fg
89        )
90    } else {
91        format!("ctx: {}", pct_str)
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::types::{ContextWindow, CurrentUsage};
99
100    #[test]
101    fn test_calculate_context_direct_used_percentage() {
102        let config = Config::default();
103        let cw = ContextWindow {
104            context_window_size: Some(200_000),
105            current_usage: None,
106            used_percentage: Some(45.5),
107            remaining_percentage: Some(54.5),
108        };
109
110        let result = calculate_context(Some(&cw), &config).unwrap();
111        assert_eq!(result.percentage, 45.5);
112        assert!(result.is_exact);
113        // 45.5% of 200k = 91k
114        assert_eq!(result.current_usage_tokens, 91_000);
115    }
116
117    #[test]
118    fn test_calculate_context_direct_remaining_percentage() {
119        let mut config = Config::default();
120        config.sections.context.display_mode = "remaining".to_string();
121        let cw = ContextWindow {
122            context_window_size: Some(200_000),
123            current_usage: None,
124            used_percentage: Some(45.5),
125            remaining_percentage: Some(54.5),
126        };
127
128        let result = calculate_context(Some(&cw), &config).unwrap();
129        assert_eq!(result.percentage, 54.5);
130        assert!(result.is_exact);
131    }
132
133    #[test]
134    fn test_calculate_context_fallback_used() {
135        let config = Config::default();
136        let cw = ContextWindow {
137            context_window_size: Some(200_000),
138            current_usage: Some(CurrentUsage {
139                input_tokens: Some(50_000),
140                _output_tokens: Some(5_000),
141                cache_creation_input_tokens: Some(10_000),
142                cache_read_input_tokens: Some(5_000),
143            }),
144            used_percentage: None,
145            remaining_percentage: None,
146        };
147
148        let result = calculate_context(Some(&cw), &config).unwrap();
149        assert!(result.percentage > 0.0 && result.percentage <= 100.0);
150        assert!(!result.is_exact);
151        assert_eq!(result.context_window_size, 200_000);
152        // 50k + 10k + 5k = 65k
153        assert_eq!(result.current_usage_tokens, 65_000);
154    }
155
156    #[test]
157    fn test_calculate_context_fallback_remaining() {
158        let mut config = Config::default();
159        config.sections.context.display_mode = "remaining".to_string();
160        let cw = ContextWindow {
161            context_window_size: Some(200_000),
162            current_usage: Some(CurrentUsage {
163                input_tokens: Some(50_000),
164                _output_tokens: None,
165                cache_creation_input_tokens: Some(10_000),
166                cache_read_input_tokens: Some(5_000),
167            }),
168            used_percentage: None,
169            remaining_percentage: None,
170        };
171
172        let result = calculate_context(Some(&cw), &config).unwrap();
173        // 65k used = 32.5% used, so 67.5% remaining (allow for floating point precision)
174        assert!((result.percentage - 67.5).abs() < 0.01);
175        assert!(!result.is_exact);
176    }
177
178    #[test]
179    fn test_calculate_context_zero_window() {
180        let config = Config::default();
181        let cw = ContextWindow {
182            context_window_size: Some(0),
183            current_usage: Some(CurrentUsage {
184                input_tokens: Some(100),
185                _output_tokens: None,
186                cache_creation_input_tokens: None,
187                cache_read_input_tokens: None,
188            }),
189            used_percentage: None,
190            remaining_percentage: None,
191        };
192
193        let result = calculate_context(Some(&cw), &config);
194        assert!(result.is_none());
195    }
196
197    #[test]
198    fn test_calculate_context_missing_window() {
199        let config = Config::default();
200        let cw = ContextWindow {
201            context_window_size: None,
202            current_usage: Some(CurrentUsage {
203                input_tokens: Some(50_000),
204                _output_tokens: None,
205                cache_creation_input_tokens: None,
206                cache_read_input_tokens: None,
207            }),
208            used_percentage: None,
209            remaining_percentage: None,
210        };
211
212        let result = calculate_context(Some(&cw), &config).unwrap();
213        assert_eq!(result.context_window_size, DEFAULT_CONTEXT_WINDOW);
214    }
215
216    #[test]
217    fn test_calculate_context_exceeds_window() {
218        let config = Config::default();
219        let cw = ContextWindow {
220            context_window_size: Some(100_000),
221            current_usage: Some(CurrentUsage {
222                input_tokens: Some(95_000),
223                _output_tokens: None,
224                cache_creation_input_tokens: Some(10_000),
225                cache_read_input_tokens: None,
226            }),
227            used_percentage: None,
228            remaining_percentage: None,
229        };
230
231        let result = calculate_context(Some(&cw), &config).unwrap();
232        assert_eq!(result.percentage, 100.0);
233    }
234
235    #[test]
236    fn test_calculate_context_all_zeros() {
237        let config = Config::default();
238        let cw = ContextWindow {
239            context_window_size: Some(200_000),
240            current_usage: Some(CurrentUsage {
241                input_tokens: Some(0),
242                _output_tokens: Some(0),
243                cache_creation_input_tokens: Some(0),
244                cache_read_input_tokens: Some(0),
245            }),
246            used_percentage: None,
247            remaining_percentage: None,
248        };
249
250        let result = calculate_context(Some(&cw), &config).unwrap();
251        assert!(result.percentage < 25.0);
252        assert_eq!(result.current_usage_tokens, 0);
253    }
254
255    #[test]
256    fn test_calculate_context_missing_usage_no_direct() {
257        let config = Config::default();
258        let cw = ContextWindow {
259            context_window_size: Some(200_000),
260            current_usage: None,
261            used_percentage: None,
262            remaining_percentage: None,
263        };
264
265        let result = calculate_context(Some(&cw), &config);
266        assert!(result.is_none());
267    }
268
269    #[test]
270    fn test_calculate_context_null_tokens() {
271        let config = Config::default();
272        let cw = ContextWindow {
273            context_window_size: Some(200_000),
274            current_usage: Some(CurrentUsage {
275                input_tokens: None,
276                _output_tokens: None,
277                cache_creation_input_tokens: None,
278                cache_read_input_tokens: None,
279            }),
280            used_percentage: None,
281            remaining_percentage: None,
282        };
283
284        let result = calculate_context(Some(&cw), &config).unwrap();
285        assert_eq!(result.current_usage_tokens, 0);
286    }
287
288    #[test]
289    fn test_format_ctx_short_exact() {
290        let config = Config::default();
291        let result = format_ctx_short(42.7, true, &config);
292        assert_eq!(result, "ctx: 42%");
293    }
294
295    #[test]
296    fn test_format_ctx_short_estimated() {
297        let config = Config::default();
298        let result = format_ctx_short(42.7, false, &config);
299        assert_eq!(result, "ctx: ~42%");
300    }
301
302    #[test]
303    fn test_format_ctx_short_exact_with_decimals() {
304        let mut config = Config::default();
305        config.sections.context.show_decimals = true;
306        let result = format_ctx_short(42.7, true, &config);
307        assert_eq!(result, "ctx: 42.7%");
308    }
309
310    #[test]
311    fn test_format_ctx_short_estimated_with_decimals() {
312        let mut config = Config::default();
313        config.sections.context.show_decimals = true;
314        let result = format_ctx_short(42.7, false, &config);
315        assert_eq!(result, "ctx: ~42.7%");
316    }
317
318    #[test]
319    fn test_format_ctx_full_exact_no_token_counts() {
320        let config = Config::default();
321        let info = ContextUsageInfo {
322            percentage: 55.5,
323            current_usage_tokens: 0,
324            context_window_size: 200_000,
325            is_exact: true,
326        };
327
328        let result = format_ctx_full(&info, &config.theme.context, &config);
329        assert!(result.contains("ctx:"));
330        assert!(result.contains("55%"));
331        assert!(!result.contains("~"));
332        assert!(!result.contains("k"));
333    }
334
335    #[test]
336    fn test_format_ctx_full_exact_with_token_counts() {
337        let mut config = Config::default();
338        config.sections.context.show_token_counts = true;
339        let info = ContextUsageInfo {
340            percentage: 55.5,
341            current_usage_tokens: 111_000,
342            context_window_size: 200_000,
343            is_exact: true,
344        };
345
346        let result = format_ctx_full(&info, &config.theme.context, &config);
347        assert!(result.contains("ctx:"));
348        assert!(result.contains("55%"));
349        assert!(!result.contains("~"));
350        assert!(result.contains("111k"));
351        assert!(result.contains("200k"));
352    }
353
354    #[test]
355    fn test_format_ctx_full_estimated_with_token_counts() {
356        let mut config = Config::default();
357        config.sections.context.show_token_counts = true;
358        let info = ContextUsageInfo {
359            percentage: 55.5,
360            current_usage_tokens: 111_000,
361            context_window_size: 200_000,
362            is_exact: false,
363        };
364
365        let result = format_ctx_full(&info, &config.theme.context, &config);
366        assert!(result.contains("ctx:"));
367        assert!(result.contains("~55%"));
368        assert!(result.contains("111k"));
369        assert!(result.contains("200k"));
370    }
371
372    #[test]
373    fn test_format_ctx_full_estimated_without_token_counts() {
374        let mut config = Config::default();
375        config.sections.context.show_token_counts = false;
376        let info = ContextUsageInfo {
377            percentage: 55.5,
378            current_usage_tokens: 111_000,
379            context_window_size: 200_000,
380            is_exact: false,
381        };
382
383        let result = format_ctx_full(&info, &config.theme.context, &config);
384        assert!(result.contains("ctx:"));
385        assert!(result.contains("~55%"));
386        assert!(!result.contains("111k"));
387        assert!(!result.contains("200k"));
388    }
389}
claude_code_status_line/context.rs

claude_code_status_line/
context.rs
