Skip to main content

lean_ctx/tools/
ctx_compress.rs

1use crate::core::cache::SessionCache;
2use crate::core::protocol;
3use crate::core::signatures;
4use crate::core::tokens::count_tokens;
5use crate::tools::ctx_response;
6use crate::tools::CrpMode;
7
8pub fn handle(cache: &SessionCache, include_signatures: bool, crp_mode: CrpMode) -> String {
9    let entries = cache.get_all_entries();
10    let file_count = entries.len();
11
12    if file_count == 0 {
13        return "CTX CHECKPOINT (0 files)\nNo files cached yet.".to_string();
14    }
15
16    let mut sections = Vec::new();
17    sections.push(format!("CTX CHECKPOINT ({file_count} files)"));
18    sections.push(String::new());
19
20    let mut total_original = 0usize;
21    let refs = cache.file_ref_map();
22
23    for (path, entry) in &entries {
24        total_original += entry.original_tokens;
25        let file_ref = refs.get(*path).map_or("F?", |s| s.as_str());
26        let short = protocol::shorten_path(path);
27
28        if include_signatures {
29            let ext = std::path::Path::new(path)
30                .extension()
31                .and_then(|e| e.to_str())
32                .unwrap_or("");
33            let Some(content) = entry.content() else {
34                continue;
35            };
36            let sigs = signatures::extract_signatures(&content, ext);
37            let sig_names: Vec<String> = sigs
38                .iter()
39                .take(5)
40                .map(|s| {
41                    if crp_mode.is_tdd() {
42                        s.to_tdd()
43                    } else {
44                        s.to_compact()
45                    }
46                })
47                .collect();
48            let more = if sigs.len() > 5 {
49                format!("+{}", sigs.len() - 5)
50            } else {
51                String::new()
52            };
53            sections.push(format!(
54                "{file_ref} {short} [{}L]: {}{more}",
55                entry.line_count,
56                sig_names.join(", "),
57            ));
58        } else {
59            sections.push(format!(
60                "{file_ref} {short} [{}L {}t]",
61                entry.line_count, entry.original_tokens
62            ));
63        }
64    }
65
66    let stats = cache.get_stats();
67    sections.push(String::new());
68    sections.push(format!(
69        "STATS: {} reads, {} hits ({:.0}%)",
70        stats.total_reads,
71        stats.cache_hits,
72        stats.hit_rate()
73    ));
74
75    let contents: Vec<(String, String)> = entries
76        .iter()
77        .filter_map(|(p, e)| Some(((*p).clone(), e.content()?)))
78        .collect();
79    let files_for_codebook: Vec<(&str, &str)> = contents
80        .iter()
81        .map(|(p, c)| (p.as_str(), c.as_str()))
82        .collect();
83    let mut codebook = crate::core::codebook::Codebook::new();
84    codebook.build_from_files(&files_for_codebook);
85
86    let output = sections.join("\n");
87
88    let (final_output, legend) = if codebook.is_empty() {
89        (output, String::new())
90    } else {
91        let (compressed, refs_used) = codebook.compress(&output);
92        let legend = codebook.format_legend(&refs_used);
93        if refs_used.is_empty() {
94            (output, String::new())
95        } else {
96            (compressed, format!("\n{legend}"))
97        }
98    };
99
100    // Apply filler removal to checkpoint output
101    let cleaned_output = ctx_response::handle(&final_output, crp_mode);
102
103    let compressed_tokens = count_tokens(&cleaned_output) + count_tokens(&legend);
104    let savings = protocol::format_savings(total_original, compressed_tokens);
105
106    format!("{cleaned_output}{legend}\nCOMPRESSION: {total_original} → {compressed_tokens} tok\n{savings}")
107}