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 sigs = signatures::extract_signatures(&entry.content, ext);
34            let sig_names: Vec<String> = sigs
35                .iter()
36                .take(5)
37                .map(|s| {
38                    if crp_mode.is_tdd() {
39                        s.to_tdd()
40                    } else {
41                        s.to_compact()
42                    }
43                })
44                .collect();
45            let more = if sigs.len() > 5 {
46                format!("+{}", sigs.len() - 5)
47            } else {
48                String::new()
49            };
50            sections.push(format!(
51                "{file_ref} {short} [{}L]: {}{more}",
52                entry.line_count,
53                sig_names.join(", "),
54            ));
55        } else {
56            sections.push(format!(
57                "{file_ref} {short} [{}L {}t]",
58                entry.line_count, entry.original_tokens
59            ));
60        }
61    }
62
63    let stats = cache.get_stats();
64    sections.push(String::new());
65    sections.push(format!(
66        "STATS: {} reads, {} hits ({:.0}%)",
67        stats.total_reads,
68        stats.cache_hits,
69        stats.hit_rate()
70    ));
71
72    // Cross-file codebook deduplication
73    let files_for_codebook: Vec<(String, String)> = entries
74        .iter()
75        .map(|(p, e)| (p.to_string(), e.content.clone()))
76        .collect();
77    let mut codebook = crate::core::codebook::Codebook::new();
78    codebook.build_from_files(&files_for_codebook);
79
80    let output = sections.join("\n");
81
82    let (final_output, legend) = if !codebook.is_empty() {
83        let (compressed, refs_used) = codebook.compress(&output);
84        let legend = codebook.format_legend(&refs_used);
85        if refs_used.is_empty() {
86            (output, String::new())
87        } else {
88            (compressed, format!("\n{legend}"))
89        }
90    } else {
91        (output, String::new())
92    };
93
94    // Apply filler removal to checkpoint output
95    let cleaned_output = ctx_response::handle(&final_output, crp_mode);
96
97    let compressed_tokens = count_tokens(&cleaned_output) + count_tokens(&legend);
98    let savings = protocol::format_savings(total_original, compressed_tokens);
99
100    format!("{cleaned_output}{legend}\nCOMPRESSION: {total_original} → {compressed_tokens} tok\n{savings}")
101}