1use std::cell::RefCell;
2use std::sync::atomic::{AtomicUsize, Ordering};
3
4static SESSION_ORIGINAL: AtomicUsize = AtomicUsize::new(0);
5static SESSION_SAVED: AtomicUsize = AtomicUsize::new(0);
6static SESSION_CALL_COUNT: AtomicUsize = AtomicUsize::new(0);
7
8const SESSION_TOTAL_INTERVAL: usize = 10;
9
10thread_local! {
11 static CURRENT_MODE: RefCell<Option<String>> = const { RefCell::new(None) };
12 static CURRENT_DETAIL: RefCell<Option<String>> = const { RefCell::new(None) };
13}
14
15pub struct SavingsInfo<'a> {
16 pub original: usize,
17 pub compressed: usize,
18 pub mode: Option<&'a str>,
19 pub detail: Option<&'a str>,
20}
21
22pub struct ModeGuard;
23
24impl ModeGuard {
25 pub fn new(mode: &str) -> Self {
26 CURRENT_MODE.with(|m| *m.borrow_mut() = Some(mode.to_string()));
27 Self
28 }
29
30 pub fn with_detail(mode: &str, detail: &str) -> Self {
31 CURRENT_MODE.with(|m| *m.borrow_mut() = Some(mode.to_string()));
32 CURRENT_DETAIL.with(|d| *d.borrow_mut() = Some(detail.to_string()));
33 Self
34 }
35}
36
37impl Drop for ModeGuard {
38 fn drop(&mut self) {
39 CURRENT_MODE.with(|m| *m.borrow_mut() = None);
40 CURRENT_DETAIL.with(|d| *d.borrow_mut() = None);
41 }
42}
43
44fn current_mode() -> Option<String> {
45 CURRENT_MODE.with(|m| m.borrow().clone())
46}
47
48fn current_detail() -> Option<String> {
49 CURRENT_DETAIL.with(|d| d.borrow().clone())
50}
51
52pub fn record_savings(original: usize, saved: usize) {
53 SESSION_ORIGINAL.fetch_add(original, Ordering::Relaxed);
54 SESSION_SAVED.fetch_add(saved, Ordering::Relaxed);
55 SESSION_CALL_COUNT.fetch_add(1, Ordering::Relaxed);
56}
57
58pub fn session_totals() -> (usize, usize, usize) {
59 (
60 SESSION_ORIGINAL.load(Ordering::Relaxed),
61 SESSION_SAVED.load(Ordering::Relaxed),
62 SESSION_CALL_COUNT.load(Ordering::Relaxed),
63 )
64}
65
66pub fn reset_session() {
67 SESSION_ORIGINAL.store(0, Ordering::Relaxed);
68 SESSION_SAVED.store(0, Ordering::Relaxed);
69 SESSION_CALL_COUNT.store(0, Ordering::Relaxed);
70}
71
72fn format_number(n: usize) -> String {
73 if n >= 1_000_000 {
74 let m = n as f64 / 1_000_000.0;
75 format!("{m:.1}M")
76 } else if n >= 10_000 {
77 let k = n as f64 / 1_000.0;
78 format!("{k:.1}k")
79 } else if n >= 1_000 {
80 let whole = n / 1_000;
81 format!("{whole},{:03}", n % 1_000)
82 } else {
83 n.to_string()
84 }
85}
86
87fn is_explicitly_enabled() -> bool {
88 matches!(std::env::var("LEAN_CTX_SHOW_SAVINGS"), Ok(v) if v.trim() == "1")
89}
90
91fn is_ultra_suppressed() -> bool {
92 if is_explicitly_enabled() {
93 return false;
94 }
95 let level = super::config::CompressionLevel::effective(&super::config::Config::load());
96 matches!(level, super::config::CompressionLevel::Max)
97}
98
99pub fn format_footer(info: &SavingsInfo<'_>) -> String {
100 if !super::protocol::savings_footer_visible() {
101 return String::new();
102 }
103 if is_ultra_suppressed() {
104 return String::new();
105 }
106 format_footer_inner(info)
107}
108
109fn format_footer_inner(info: &SavingsInfo<'_>) -> String {
110 if info.original == 0 {
111 return String::new();
112 }
113 let saved = info.original.saturating_sub(info.compressed);
114 if saved == 0 {
115 return String::new();
116 }
117 let pct = (saved as f64 / info.original as f64 * 100.0).round() as usize;
118
119 let orig_str = format_number(info.original);
120 let comp_str = format_number(info.compressed);
121
122 let mut parts = vec![format!(
123 "{orig_str} \u{2192} {comp_str} tok (\u{2193}{pct}%)"
124 )];
125
126 if let Some(mode) = info.mode {
127 parts.push(format!("mode: {mode}"));
128 }
129 if let Some(detail) = info.detail {
130 parts.push(detail.to_string());
131 }
132
133 record_savings(info.original, saved);
134
135 let call_count = SESSION_CALL_COUNT.load(Ordering::Relaxed);
136 if call_count > 0 && call_count.is_multiple_of(SESSION_TOTAL_INTERVAL) {
137 let (_, total_saved, _) = session_totals();
138 let total_str = format_number(total_saved);
139 parts.push(format!("session: {total_str} saved"));
140 }
141
142 let body = parts.join(" | ");
143 format!("\u{2500}\u{2500}\u{2500} {body} \u{2500}\u{2500}\u{2500}")
144}
145
146pub fn format_footer_basic(original: usize, compressed: usize) -> String {
147 let mode = current_mode();
148 let detail = current_detail();
149 format_footer(&SavingsInfo {
150 original,
151 compressed,
152 mode: mode.as_deref(),
153 detail: detail.as_deref(),
154 })
155}
156
157pub fn append_footer(output: &str, info: &SavingsInfo<'_>) -> String {
158 let footer = format_footer(info);
159 if footer.is_empty() {
160 output.to_string()
161 } else {
162 format!("{output}\n{footer}")
163 }
164}
165
166pub fn append_footer_basic(output: &str, original: usize, compressed: usize) -> String {
167 let footer = format_footer_basic(original, compressed);
168 if footer.is_empty() {
169 output.to_string()
170 } else {
171 format!("{output}\n{footer}")
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 #[test]
180 fn format_number_small() {
181 assert_eq!(format_number(42), "42");
182 assert_eq!(format_number(999), "999");
183 }
184
185 #[test]
186 fn format_number_thousands() {
187 assert_eq!(format_number(1_000), "1,000");
188 assert_eq!(format_number(4_200), "4,200");
189 assert_eq!(format_number(9_999), "9,999");
190 }
191
192 #[test]
193 fn format_number_large() {
194 assert_eq!(format_number(12_300), "12.3k");
195 assert_eq!(format_number(45_200), "45.2k");
196 }
197
198 #[test]
199 fn format_number_millions() {
200 assert_eq!(format_number(1_500_000), "1.5M");
201 }
202
203 #[test]
204 fn basic_footer_format() {
205 let info = SavingsInfo {
206 original: 4200,
207 compressed: 840,
208 mode: Some("map"),
209 detail: None,
210 };
211 let result = format_footer_inner(&info);
212 assert!(
213 result.starts_with("\u{2500}\u{2500}\u{2500} "),
214 "should start with box-drawing: {result}"
215 );
216 assert!(
217 result.ends_with(" \u{2500}\u{2500}\u{2500}"),
218 "should end with box-drawing: {result}"
219 );
220 assert!(
221 result.contains("4,200"),
222 "should contain formatted original: {result}"
223 );
224 assert!(
225 result.contains("840"),
226 "should contain compressed: {result}"
227 );
228 assert!(
229 result.contains("\u{2193}80%"),
230 "should contain percentage: {result}"
231 );
232 assert!(
233 result.contains("mode: map"),
234 "should contain mode: {result}"
235 );
236 }
237
238 #[test]
239 fn footer_with_detail() {
240 let info = SavingsInfo {
241 original: 12300,
242 compressed: 620,
243 mode: None,
244 detail: Some("3 patterns matched"),
245 };
246 let result = format_footer_inner(&info);
247 assert!(
248 result.contains("3 patterns matched"),
249 "detail missing: {result}"
250 );
251 assert!(
252 result.contains("12.3k"),
253 "should format large numbers: {result}"
254 );
255 }
256
257 #[test]
258 fn footer_returns_empty_when_no_savings() {
259 let result = format_footer_inner(&SavingsInfo {
260 original: 100,
261 compressed: 100,
262 mode: None,
263 detail: None,
264 });
265 assert!(
266 result.is_empty(),
267 "should be empty with 0 savings: {result}"
268 );
269 }
270
271 #[test]
272 fn footer_returns_empty_when_zero_original() {
273 let result = format_footer_inner(&SavingsInfo {
274 original: 0,
275 compressed: 0,
276 mode: None,
277 detail: None,
278 });
279 assert!(
280 result.is_empty(),
281 "should be empty with 0 original: {result}"
282 );
283 }
284
285 #[test]
286 fn visibility_gated_tests() {
287 let _lock = crate::core::data_dir::test_env_lock();
288
289 std::env::set_var("LEAN_CTX_SHOW_SAVINGS", "0");
290 std::env::set_var("LEAN_CTX_SAVINGS_FOOTER", "never");
291 let result = format_footer_basic(100, 50);
292 assert!(
293 result.is_empty(),
294 "should be empty with never mode: {result}"
295 );
296
297 let result = append_footer_basic("hello", 100, 50);
298 assert_eq!(result, "hello");
299
300 std::env::set_var("LEAN_CTX_SHOW_SAVINGS", "1");
301 std::env::set_var("LEAN_CTX_SAVINGS_FOOTER", "always");
302 std::env::remove_var("LEAN_CTX_QUIET");
303 super::super::protocol::set_mcp_context(false);
304
305 let result = append_footer_basic("hello", 100, 50);
306 assert!(
307 result.starts_with("hello\n"),
308 "should start with original: {result}"
309 );
310 assert!(
311 result.contains("\u{2500}\u{2500}\u{2500}"),
312 "should contain box-drawing: {result}"
313 );
314
315 std::env::remove_var("LEAN_CTX_SHOW_SAVINGS");
316 }
317
318 #[test]
319 fn session_accumulator_tracks() {
320 reset_session();
321 record_savings(100, 50);
322 record_savings(200, 80);
323 let (orig, saved, calls) = session_totals();
324 assert_eq!(orig, 300);
325 assert_eq!(saved, 130);
326 assert_eq!(calls, 2);
327 reset_session();
328 }
329
330 #[test]
331 fn session_total_shown_at_interval() {
332 reset_session();
333 for _ in 0..(SESSION_TOTAL_INTERVAL - 1) {
334 record_savings(100, 50);
335 }
336 let info = SavingsInfo {
337 original: 100,
338 compressed: 50,
339 mode: None,
340 detail: None,
341 };
342 let result = format_footer_inner(&info);
343 assert!(
344 result.contains("session:"),
345 "should contain session total at interval: {result}"
346 );
347 reset_session();
348 }
349
350 #[test]
351 fn mode_guard_sets_and_clears() {
352 assert!(current_mode().is_none());
353 {
354 let _guard = ModeGuard::new("map");
355 assert_eq!(current_mode().as_deref(), Some("map"));
356 }
357 assert!(current_mode().is_none());
358 }
359
360 #[test]
361 fn mode_guard_with_detail() {
362 {
363 let _guard = ModeGuard::with_detail("shell", "3 patterns");
364 assert_eq!(current_mode().as_deref(), Some("shell"));
365 assert_eq!(current_detail().as_deref(), Some("3 patterns"));
366 }
367 assert!(current_mode().is_none());
368 assert!(current_detail().is_none());
369 }
370}