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    if total == 0 {
16        return None;
17    }
18    let usage = cw.current_usage.as_ref()?;
19
20    // Calculate content tokens using official method
21    // current_tokens = input_tokens + cache_creation_input_tokens + cache_read_input_tokens
22    let content_tokens = usage.input_tokens.unwrap_or(0)
23        + usage.cache_creation_input_tokens.unwrap_or(0)
24        + usage.cache_read_input_tokens.unwrap_or(0);
25
26    // Add autocompact buffer
27    let total_used = content_tokens + config.sections.context.autocompact_buffer_size;
28
29    // Raw percentage (matches /context display)
30    let percentage = (total_used as f64 / total as f64 * 100.0).min(100.0);
31
32    Some(ContextUsageInfo {
33        percentage,
34        current_usage_tokens: total_used,
35        context_window_size: total,
36    })
37}
38
39pub fn format_ctx_short(pct: f64, config: &Config) -> String {
40    if config.sections.context.show_decimals {
41        format!("ctx: ~{:.1}%", pct)
42    } else {
43        format!("ctx: ~{}%", pct as u64)
44    }
45}
46
47pub fn format_ctx_full(info: &ContextUsageInfo, colors: &SectionColors, config: &Config) -> String {
48    let (details, fg) = get_details_and_fg_codes(colors, config);
49
50    // Format tokens in k (thousands)
51    let used_k = info.current_usage_tokens / 1000;
52    let total_k = info.context_window_size / 1000;
53
54    let pct_str = if config.sections.context.show_decimals {
55        format!("~{:.1}%", info.percentage)
56    } else {
57        format!("~{}%", info.percentage as u64)
58    };
59
60    format!(
61        "ctx: {} {}({}k/{}k){}",
62        pct_str, details, used_k, total_k, fg
63    )
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use crate::types::{ContextWindow, CurrentUsage};
70
71    #[test]
72    fn test_calculate_context_normal() {
73        let config = Config::default();
74        let cw = ContextWindow {
75            context_window_size: Some(200_000),
76            current_usage: Some(CurrentUsage {
77                input_tokens: Some(50_000),
78                _output_tokens: Some(5_000),
79                cache_creation_input_tokens: Some(10_000),
80                cache_read_input_tokens: Some(5_000),
81            }),
82        };
83
84        let result = calculate_context(Some(&cw), &config).unwrap();
85        assert!(result.percentage > 0.0 && result.percentage <= 100.0);
86        assert_eq!(result.context_window_size, 200_000);
87        // 50k + 10k + 5k + 45k buffer = 110k
88        assert_eq!(result.current_usage_tokens, 110_000);
89    }
90
91    #[test]
92    fn test_calculate_context_zero_window() {
93        let config = Config::default();
94        let cw = ContextWindow {
95            context_window_size: Some(0),
96            current_usage: Some(CurrentUsage {
97                input_tokens: Some(100),
98                _output_tokens: None,
99                cache_creation_input_tokens: None,
100                cache_read_input_tokens: None,
101            }),
102        };
103
104        let result = calculate_context(Some(&cw), &config);
105        assert!(result.is_none()); // Should handle zero window
106    }
107
108    #[test]
109    fn test_calculate_context_missing_window() {
110        let config = Config::default();
111        let cw = ContextWindow {
112            context_window_size: None,
113            current_usage: Some(CurrentUsage {
114                input_tokens: Some(50_000),
115                _output_tokens: None,
116                cache_creation_input_tokens: None,
117                cache_read_input_tokens: None,
118            }),
119        };
120
121        let result = calculate_context(Some(&cw), &config).unwrap();
122        // Should use DEFAULT_CONTEXT_WINDOW
123        assert_eq!(result.context_window_size, DEFAULT_CONTEXT_WINDOW);
124    }
125
126    #[test]
127    fn test_calculate_context_exceeds_window() {
128        let config = Config::default();
129        let cw = ContextWindow {
130            context_window_size: Some(100_000),
131            current_usage: Some(CurrentUsage {
132                input_tokens: Some(95_000),
133                _output_tokens: None,
134                cache_creation_input_tokens: Some(10_000),
135                cache_read_input_tokens: None,
136            }),
137        };
138
139        let result = calculate_context(Some(&cw), &config).unwrap();
140        // Should cap at 100%
141        assert_eq!(result.percentage, 100.0);
142    }
143
144    #[test]
145    fn test_calculate_context_all_zeros() {
146        let config = Config::default();
147        let cw = ContextWindow {
148            context_window_size: Some(200_000),
149            current_usage: Some(CurrentUsage {
150                input_tokens: Some(0),
151                _output_tokens: Some(0),
152                cache_creation_input_tokens: Some(0),
153                cache_read_input_tokens: Some(0),
154            }),
155        };
156
157        let result = calculate_context(Some(&cw), &config).unwrap();
158        // Just buffer, should be low percentage
159        assert!(result.percentage < 25.0);
160        assert_eq!(result.current_usage_tokens, 45_000); // Just buffer
161    }
162
163    #[test]
164    fn test_calculate_context_missing_usage() {
165        let config = Config::default();
166        let cw = ContextWindow {
167            context_window_size: Some(200_000),
168            current_usage: None,
169        };
170
171        let result = calculate_context(Some(&cw), &config);
172        assert!(result.is_none());
173    }
174
175    #[test]
176    fn test_calculate_context_null_tokens() {
177        let config = Config::default();
178        let cw = ContextWindow {
179            context_window_size: Some(200_000),
180            current_usage: Some(CurrentUsage {
181                input_tokens: None,
182                _output_tokens: None,
183                cache_creation_input_tokens: None,
184                cache_read_input_tokens: None,
185            }),
186        };
187
188        let result = calculate_context(Some(&cw), &config).unwrap();
189        // All None should be treated as 0
190        assert_eq!(result.current_usage_tokens, 45_000); // Just buffer
191    }
192
193    #[test]
194    fn test_format_ctx_short_with_decimals() {
195        let mut config = Config::default();
196        config.sections.context.show_decimals = true;
197
198        let result = format_ctx_short(42.7, &config);
199        assert_eq!(result, "ctx: ~42.7%");
200    }
201
202    #[test]
203    fn test_format_ctx_short_without_decimals() {
204        let mut config = Config::default();
205        config.sections.context.show_decimals = false;
206
207        let result = format_ctx_short(42.7, &config);
208        assert_eq!(result, "ctx: ~42%");
209    }
210
211    #[test]
212    fn test_format_ctx_full() {
213        let config = Config::default();
214        let info = ContextUsageInfo {
215            percentage: 55.5,
216            current_usage_tokens: 111_000,
217            context_window_size: 200_000,
218        };
219
220        let result = format_ctx_full(&info, &config.theme.context, &config);
221        assert!(result.contains("ctx:"));
222        assert!(result.contains("111k"));
223        assert!(result.contains("200k"));
224    }
225}
claude_code_status_line/context.rs

claude_code_status_line/
context.rs
