lean_ctx/tools/
ctx_compress_memory.rs1use std::path::Path;
2
3use crate::core::tokens::count_tokens;
4
5const SOURCE_EXTENSIONS: &[&str] = &[
6 ".rs", ".ts", ".tsx", ".js", ".jsx", ".py", ".go", ".java", ".c", ".cpp", ".h", ".hpp", ".cs",
7 ".rb", ".swift", ".kt", ".scala", ".zig", ".toml", ".yaml", ".yml", ".json", ".lock", ".sum",
8 ".mod",
9];
10
11pub fn handle(path: &str) -> String {
12 if SOURCE_EXTENSIONS.iter().any(|ext| path.ends_with(ext)) {
13 return format!(
14 "ERROR: refusing to compress source file {path} — \
15 ctx_compress_memory is for memory-bank/doc files only (#1323)"
16 );
17 }
18
19 if let Err(e) = crate::core::pathjail::enforce_writable(Path::new(path)) {
20 return format!("ERROR: {e}");
21 }
22
23 let content = match std::fs::read_to_string(path) {
24 Ok(c) => c,
25 Err(e) => return format!("ERROR: Cannot read {path}: {e}"),
26 };
27
28 let original_tokens = count_tokens(&content);
29
30 let backup_path = build_backup_path(path);
31 if !Path::new(&backup_path).exists()
32 && let Err(e) = std::fs::write(&backup_path, &content)
33 {
34 return format!("ERROR: Cannot create backup {backup_path}: {e}");
35 }
36
37 let compressed = compress_memory_file(&content);
38 let compressed_tokens = count_tokens(&compressed);
39
40 if let Err(e) = std::fs::write(path, &compressed) {
41 return format!("ERROR: Cannot write compressed file: {e}");
42 }
43
44 let saved = original_tokens.saturating_sub(compressed_tokens);
45 let pct = if original_tokens > 0 {
46 (saved as f64 / original_tokens as f64 * 100.0).round() as usize
47 } else {
48 0
49 };
50
51 format!(
52 "Compressed {}: {} → {} tokens ({saved} saved, {pct}%)\n\
53 Backup: {backup_path}",
54 Path::new(path)
55 .file_name()
56 .and_then(|f| f.to_str())
57 .unwrap_or(path),
58 original_tokens,
59 compressed_tokens,
60 )
61}
62
63fn build_backup_path(path: &str) -> String {
64 let p = Path::new(path);
65 let stem = p.file_stem().and_then(|s| s.to_str()).unwrap_or("file");
66 let parent = p.parent().unwrap_or_else(|| Path::new("."));
67 parent
68 .join(format!("{stem}.original.md"))
69 .to_string_lossy()
70 .to_string()
71}
72
73fn compress_memory_file(content: &str) -> String {
74 let mut output = Vec::new();
75 let mut in_code_block = false;
76 let mut code_fence = String::new();
77
78 for line in content.lines() {
79 let trimmed = line.trim();
80
81 if !in_code_block && is_code_fence_start(trimmed) {
82 in_code_block = true;
83 code_fence = trimmed
84 .chars()
85 .take_while(|c| *c == '`' || *c == '~')
86 .collect();
87 output.push(line.to_string());
88 continue;
89 }
90
91 if in_code_block {
92 output.push(line.to_string());
93 if trimmed.starts_with(&code_fence) && trimmed.len() <= code_fence.len() + 1 {
94 in_code_block = false;
95 code_fence.clear();
96 }
97 continue;
98 }
99
100 if is_protected_line(trimmed) {
101 output.push(line.to_string());
102 continue;
103 }
104
105 if trimmed.is_empty() {
106 if output.last().is_some_and(|l| l.trim().is_empty()) {
107 continue;
108 }
109 output.push(String::new());
110 continue;
111 }
112
113 let compressed = compress_prose_line(line);
114 if !compressed.trim().is_empty() {
115 output.push(compressed);
116 }
117 }
118
119 output.join("\n")
120}
121
122fn is_code_fence_start(line: &str) -> bool {
123 line.starts_with("```") || line.starts_with("~~~")
124}
125
126fn is_protected_line(line: &str) -> bool {
127 if line.starts_with('#') {
128 return true;
129 }
130 if line.starts_with("- ") || line.starts_with("* ") || line.starts_with("> ") {
131 return true;
132 }
133 if line.starts_with('|') {
134 return true;
135 }
136 if contains_url_or_path(line) && line.split_whitespace().count() <= 3 {
137 return true;
138 }
139 false
140}
141
142fn contains_url_or_path(line: &str) -> bool {
143 line.contains("http://")
144 || line.contains("https://")
145 || line.contains("ftp://")
146 || (line.contains('/') && line.contains('.') && !line.contains(' '))
147}
148
149fn compress_prose_line(line: &str) -> String {
150 let leading_ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
151 let trimmed = line.trim();
152
153 let mut words: Vec<&str> = trimmed.split_whitespace().collect();
154
155 words.retain(|w| !is_filler_word(w));
156
157 let mut result: Vec<String> = Vec::new();
158 let mut i = 0;
159 while i < words.len() {
160 if let Some((replacement, skip)) = try_shorten_phrase(&words, i) {
161 result.push(replacement.to_string());
162 i += skip;
163 } else {
164 result.push(words[i].to_string());
165 i += 1;
166 }
167 }
168
169 format!("{}{}", leading_ws, result.join(" "))
170}
171
172fn is_filler_word(word: &str) -> bool {
173 let w = word.to_lowercase();
174 let w = w.trim_matches(|c: char| c.is_ascii_punctuation());
175 matches!(
176 w,
177 "just" | "really" | "basically" | "actually" | "simply" | "please" | "very" | "quite"
178 )
179}
180
181fn try_shorten_phrase(words: &[&str], pos: usize) -> Option<(&'static str, usize)> {
182 if pos + 2 < words.len() {
183 let three = format!(
184 "{} {} {}",
185 words[pos].to_lowercase(),
186 words[pos + 1].to_lowercase(),
187 words[pos + 2].to_lowercase()
188 );
189 match three.as_str() {
190 "in order to" => return Some(("to", 3)),
191 "as well as" => return Some(("and", 3)),
192 "due to the" => return Some(("because", 3)),
193 "make sure to" => return Some(("ensure", 3)),
194 "a lot of" => return Some(("many", 3)),
195 "on top of" => return Some(("besides", 3)),
196 _ => {}
197 }
198 }
199
200 if pos + 1 < words.len() {
201 let two = format!(
202 "{} {}",
203 words[pos].to_lowercase(),
204 words[pos + 1].to_lowercase()
205 );
206 match two.as_str() {
207 "make sure" => return Some(("ensure", 2)),
208 "a lot" => return Some(("many", 2)),
209 "as well" => return Some(("also", 2)),
210 "in order" => return Some(("to", 2)),
211 "prior to" => return Some(("before", 2)),
212 "due to" => return Some(("because", 2)),
213 _ => {}
214 }
215 }
216
217 None
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn preserves_code_blocks() {
226 let input = "Some text just really here.\n\n```rust\nfn main() {\n println!(\"hello\");\n}\n```\n\nMore text.";
227 let result = compress_memory_file(input);
228 assert!(result.contains("fn main()"));
229 assert!(result.contains("println!"));
230 }
231
232 #[test]
233 fn preserves_headings() {
234 let input = "# Main Heading\n\nJust some filler text here.\n\n## Sub Heading";
235 let result = compress_memory_file(input);
236 assert!(result.contains("# Main Heading"));
237 assert!(result.contains("## Sub Heading"));
238 }
239
240 #[test]
241 fn preserves_urls() {
242 let input = "Visit https://example.com for details.\nJust some really basic text.";
243 let result = compress_memory_file(input);
244 assert!(result.contains("https://example.com"));
245 }
246
247 #[test]
248 fn removes_filler_words() {
249 let input = "You should just really basically make sure to check this.";
250 let result = compress_prose_line(input);
251 assert!(!result.contains("just"));
252 assert!(!result.contains("really"));
253 assert!(!result.contains("basically"));
254 assert!(result.contains("ensure"));
255 }
256
257 #[test]
258 fn shortens_phrases() {
259 let input = "In order to fix this, make sure to check the config.";
260 let result = compress_prose_line(input);
261 assert!(!result.contains("In order to"));
262 assert!(result.contains("to"));
263 assert!(result.contains("ensure"));
264 }
265
266 #[test]
267 fn collapses_blank_lines() {
268 let input = "Line 1\n\n\n\nLine 2\n\n\nLine 3";
269 let result = compress_memory_file(input);
270 assert!(!result.contains("\n\n\n"));
271 }
272
273 #[test]
274 fn preserves_tables() {
275 let input = "| Col A | Col B |\n|-------|-------|\n| val1 | val2 |";
276 let result = compress_memory_file(input);
277 assert!(result.contains("| Col A | Col B |"));
278 }
279
280 #[test]
281 fn backup_path_computed_correctly() {
282 assert_eq!(
283 Path::new(&build_backup_path("/home/user/.cursorrules")),
284 Path::new("/home/user")
285 .join(".cursorrules.original.md")
286 .as_path()
287 );
288 assert_eq!(
289 Path::new(&build_backup_path("/project/CLAUDE.md")),
290 Path::new("/project").join("CLAUDE.original.md").as_path()
291 );
292 }
293
294 #[cfg(not(feature = "no-jail"))]
297 #[test]
298 fn handle_denies_write_into_read_only_root() {
299 let _iso = crate::core::data_dir::isolated_data_dir();
300 let dir = tempfile::tempdir().unwrap();
301 let ro = dir.path().join("refrepo");
302 std::fs::create_dir_all(&ro).unwrap();
303 let file = ro.join("CLAUDE.md");
304 let original = "# Title\n\nJust some really verbose prose here.\n";
305 std::fs::write(&file, original).unwrap();
306
307 let ro_canon = crate::core::pathjail::canonicalize_or_self(&ro);
308 crate::test_env::set_var(
309 "LEAN_CTX_READ_ONLY_ROOTS",
310 ro_canon.to_string_lossy().as_ref(),
311 );
312 let out = handle(file.to_string_lossy().as_ref());
313 crate::test_env::remove_var("LEAN_CTX_READ_ONLY_ROOTS");
314
315 assert!(
316 out.starts_with("ERROR") && out.contains("read-only"),
317 "compaction into a read-only root must be refused: {out}"
318 );
319 assert_eq!(
320 std::fs::read_to_string(&file).unwrap(),
321 original,
322 "the file must be untouched"
323 );
324 assert!(
325 !ro.join("CLAUDE.original.md").exists(),
326 "no backup may be written into a read-only root"
327 );
328 }
329}