lean_ctx/core/
archive.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::path::PathBuf;
4
5use super::data_dir::lean_ctx_data_dir;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ArchiveEntry {
9    pub id: String,
10    pub tool: String,
11    pub command: String,
12    pub size_chars: usize,
13    pub size_tokens: usize,
14    pub created_at: DateTime<Utc>,
15    pub session_id: Option<String>,
16}
17
18fn archive_base_dir() -> PathBuf {
19    lean_ctx_data_dir()
20        .unwrap_or_else(|_| PathBuf::from(".lean-ctx"))
21        .join("archives")
22}
23
24fn entry_dir(id: &str) -> PathBuf {
25    let prefix = if id.len() >= 2 { &id[..2] } else { id };
26    archive_base_dir().join(prefix)
27}
28
29fn content_path(id: &str) -> PathBuf {
30    entry_dir(id).join(format!("{id}.txt"))
31}
32
33fn meta_path(id: &str) -> PathBuf {
34    entry_dir(id).join(format!("{id}.meta.json"))
35}
36
37#[cfg(unix)]
38fn set_private_file_perms(path: &PathBuf) {
39    use std::os::unix::fs::PermissionsExt;
40    let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600));
41}
42
43fn compute_id(content: &str) -> String {
44    use std::collections::hash_map::DefaultHasher;
45    use std::hash::{Hash, Hasher};
46    let mut hasher = DefaultHasher::new();
47    content.hash(&mut hasher);
48    let hash = hasher.finish();
49    format!("{hash:016x}")
50}
51
52pub fn is_enabled() -> bool {
53    if let Ok(v) = std::env::var("LEAN_CTX_ARCHIVE") {
54        return !matches!(v.as_str(), "0" | "false" | "off");
55    }
56    super::config::Config::load().archive.enabled
57}
58
59fn threshold_chars() -> usize {
60    if let Ok(v) = std::env::var("LEAN_CTX_ARCHIVE_THRESHOLD") {
61        if let Ok(n) = v.parse::<usize>() {
62            return n;
63        }
64    }
65    super::config::Config::load().archive.threshold_chars
66}
67
68fn max_age_hours() -> u64 {
69    if let Ok(v) = std::env::var("LEAN_CTX_ARCHIVE_TTL") {
70        if let Ok(n) = v.parse::<u64>() {
71            return n;
72        }
73    }
74    super::config::Config::load().archive_max_age_hours_effective()
75}
76
77/// Effective on-disk byte budget for archived `.txt`/`.meta.json` content,
78/// derived from `[archive] max_disk_mb` (or the simplified global `max_disk_mb`)
79/// so it matches what `doctor` reports. `0` disables the size cap; the TTL still
80/// applies. This is what bounds the archive store on disk (#417).
81fn max_disk_bytes() -> u64 {
82    super::config::Config::load()
83        .archive_max_disk_mb_effective()
84        .saturating_mul(1024 * 1024)
85}
86
87pub fn should_archive(content: &str) -> bool {
88    is_enabled() && content.len() >= threshold_chars()
89}
90
91const MAX_ARCHIVE_SIZE: usize = 10 * 1024 * 1024; // 10 MB
92
93pub fn store(tool: &str, command: &str, content: &str, session_id: Option<&str>) -> Option<String> {
94    if !is_enabled() || content.is_empty() {
95        return None;
96    }
97
98    let content = if content.len() > MAX_ARCHIVE_SIZE {
99        &content[..content.floor_char_boundary(MAX_ARCHIVE_SIZE)]
100    } else {
101        content
102    };
103
104    let id = compute_id(content);
105    let c_path = content_path(&id);
106
107    // Fast path: content already archived (idempotent, no race)
108    if c_path.exists() {
109        return Some(id);
110    }
111
112    let dir = entry_dir(&id);
113    if std::fs::create_dir_all(&dir).is_err() {
114        return None;
115    }
116
117    // Atomic write: PID-unique tmp file prevents race between parallel writers.
118    // rename() is atomic on POSIX; on Windows it replaces atomically too.
119    // If two processes race past the exists() check, both write their own tmp
120    // file and both rename to the same target — last writer wins, content is
121    // identical (same hash), so the result is correct either way.
122    let pid = std::process::id();
123    let tmp_path = c_path.with_extension(format!("tmp.{pid}"));
124    if std::fs::write(&tmp_path, content).is_err() {
125        return None;
126    }
127    if std::fs::rename(&tmp_path, &c_path).is_err() {
128        let _ = std::fs::remove_file(&tmp_path);
129        // Another process may have won the race — check if content is there now
130        if c_path.exists() {
131            return Some(id);
132        }
133        return None;
134    }
135    #[cfg(unix)]
136    set_private_file_perms(&c_path);
137
138    let tokens = super::tokens::count_tokens(content);
139    let entry = ArchiveEntry {
140        id: id.clone(),
141        tool: tool.to_string(),
142        command: command.to_string(),
143        size_chars: content.len(),
144        size_tokens: tokens,
145        created_at: Utc::now(),
146        session_id: session_id.map(std::string::ToString::to_string),
147    };
148
149    if let Ok(json) = serde_json::to_string_pretty(&entry) {
150        let meta_tmp = meta_path(&id).with_extension(format!("tmp.{pid}"));
151        if std::fs::write(&meta_tmp, &json).is_ok() {
152            let meta_final = meta_path(&id);
153            let _ = std::fs::rename(&meta_tmp, &meta_final);
154            #[cfg(unix)]
155            set_private_file_perms(&meta_final);
156        }
157    }
158
159    super::archive_fts::index_entry(&id, tool, command, content);
160
161    Some(id)
162}
163
164pub fn retrieve(id: &str) -> Option<String> {
165    let path = content_path(id);
166    std::fs::read_to_string(path).ok()
167}
168
169pub fn retrieve_with_range(id: &str, start: usize, end: usize) -> Option<String> {
170    let content = retrieve(id)?;
171    let lines: Vec<&str> = content.lines().collect();
172    let start = start.saturating_sub(1).min(lines.len());
173    let end = end.min(lines.len());
174    if start >= end {
175        return Some(String::new());
176    }
177    Some(
178        lines[start..end]
179            .iter()
180            .enumerate()
181            .map(|(i, line)| format!("{:>6}|{line}", start + i + 1))
182            .collect::<Vec<_>>()
183            .join("\n"),
184    )
185}
186
187pub fn retrieve_with_search(id: &str, pattern: &str) -> Option<String> {
188    let content = retrieve(id)?;
189    let pattern_lower = pattern.to_lowercase();
190    let matches: Vec<String> = content
191        .lines()
192        .enumerate()
193        .filter(|(_, line)| line.to_lowercase().contains(&pattern_lower))
194        .map(|(i, line)| format!("{:>6}|{line}", i + 1))
195        .collect();
196
197    if matches.is_empty() {
198        Some(format!("No matches for \"{pattern}\" in archive {id}"))
199    } else {
200        Some(format!(
201            "{} match(es) for \"{}\":\n{}",
202            matches.len(),
203            pattern,
204            matches.join("\n")
205        ))
206    }
207}
208
209/// Retrieve the first `n` lines of an archived entry, with a line-number gutter.
210pub fn retrieve_head(id: &str, n: usize) -> Option<String> {
211    retrieve_with_range(id, 1, n)
212}
213
214/// Retrieve the last `n` lines of an archived entry, with a line-number gutter.
215pub fn retrieve_tail(id: &str, n: usize) -> Option<String> {
216    let content = retrieve(id)?;
217    let total = content.lines().count();
218    let start = if total > n { total - n + 1 } else { 1 };
219    retrieve_with_range(id, start, total)
220}
221
222/// Describe the JSON structure of an archived entry: top-level keys with type hints,
223/// array lengths + element types, etc. An optional dot/slash `path` (e.g. `data.items.0`)
224/// navigates into the structure first. Returns `None` when the archive is missing or its
225/// content is not valid JSON, so callers can fall back to a raw retrieval hint.
226pub fn retrieve_json_keys(id: &str, path: Option<&str>) -> Option<String> {
227    let content = retrieve(id)?;
228    let root: serde_json::Value = serde_json::from_str(content.trim()).ok()?;
229    let mut cur = &root;
230    let mut walked = String::from("$");
231    if let Some(p) = path {
232        for seg in p.split(['.', '/']).filter(|s| !s.is_empty()) {
233            let next = if let Ok(idx) = seg.parse::<usize>() {
234                cur.get(idx)
235            } else {
236                cur.get(seg)
237            };
238            match next {
239                Some(v) => {
240                    cur = v;
241                    walked.push('.');
242                    walked.push_str(seg);
243                }
244                None => {
245                    return Some(format!(
246                        "Path '{p}' not found at '{walked}' in archive {id}"
247                    ));
248                }
249            }
250        }
251    }
252    Some(format!("{walked} => {}", describe_json(cur)))
253}
254
255fn json_type_hint(v: &serde_json::Value) -> String {
256    use serde_json::Value;
257    match v {
258        Value::Object(m) => format!("object({})", m.len()),
259        Value::Array(a) => format!("array({})", a.len()),
260        Value::String(s) => {
261            let preview: String = s.chars().take(40).collect();
262            if s.chars().count() > 40 {
263                format!("string \"{preview}…\"")
264            } else {
265                format!("string \"{preview}\"")
266            }
267        }
268        Value::Number(n) => format!("number {n}"),
269        Value::Bool(b) => format!("bool {b}"),
270        Value::Null => "null".to_string(),
271    }
272}
273
274fn describe_json(v: &serde_json::Value) -> String {
275    use serde_json::Value;
276    match v {
277        Value::Object(map) => {
278            let mut keys: Vec<&String> = map.keys().collect();
279            keys.sort();
280            let rendered: Vec<String> = keys
281                .iter()
282                .map(|k| format!("  {k}: {}", json_type_hint(&map[*k])))
283                .collect();
284            format!("object ({} keys)\n{}", map.len(), rendered.join("\n"))
285        }
286        Value::Array(arr) => {
287            let elem = arr.first().map_or("empty", |e| match e {
288                Value::Object(_) => "object",
289                Value::Array(_) => "array",
290                Value::String(_) => "string",
291                Value::Number(_) => "number",
292                Value::Bool(_) => "bool",
293                Value::Null => "null",
294            });
295            let mut out = format!("array ({} items of {elem})", arr.len());
296            if let Some(Value::Object(map)) = arr.first() {
297                let mut keys: Vec<&String> = map.keys().collect();
298                keys.sort();
299                out.push_str(&format!(
300                    "\n  [0] keys: {}",
301                    keys.iter()
302                        .map(|s| s.as_str())
303                        .collect::<Vec<_>>()
304                        .join(", ")
305                ));
306            }
307            out
308        }
309        Value::String(s) => format!("string ({} chars)", s.len()),
310        Value::Number(n) => format!("number ({n})"),
311        Value::Bool(b) => format!("bool ({b})"),
312        Value::Null => "null".to_string(),
313    }
314}
315
316pub fn list_entries(session_id: Option<&str>) -> Vec<ArchiveEntry> {
317    let base = archive_base_dir();
318    if !base.exists() {
319        return Vec::new();
320    }
321    let mut entries = Vec::new();
322    if let Ok(dirs) = std::fs::read_dir(&base) {
323        for dir_entry in dirs.flatten() {
324            if !dir_entry.path().is_dir() {
325                continue;
326            }
327            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
328                for file in files.flatten() {
329                    let path = file.path();
330                    if path.extension().and_then(|e| e.to_str()) != Some("json") {
331                        continue;
332                    }
333                    if let Ok(data) = std::fs::read_to_string(&path) {
334                        if let Ok(entry) = serde_json::from_str::<ArchiveEntry>(&data) {
335                            if let Some(sid) = session_id {
336                                if entry.session_id.as_deref() != Some(sid) {
337                                    continue;
338                                }
339                            }
340                            entries.push(entry);
341                        }
342                    }
343                }
344            }
345        }
346    }
347    entries.sort_by_key(|e| std::cmp::Reverse(e.created_at));
348    entries
349}
350
351/// Remove only the on-disk content + metadata files for an archive id, leaving
352/// the FTS index untouched. Used by the FTS cap-enforcer so the `.txt`/`.meta.json`
353/// blobs of rows it evicts can't outlive their index entry as orphans (#417).
354pub fn remove_files(id: &str) {
355    let _ = std::fs::remove_file(content_path(id));
356    let _ = std::fs::remove_file(meta_path(id));
357}
358
359/// Prune archived entries that exceed the age TTL (`max_age_hours`) or that push
360/// the on-disk store past its size budget (`max_disk_mb`). The content file,
361/// metadata, and FTS index are removed together so the two stores stay in sync.
362/// Returns the number of entries removed.
363///
364/// Wired into MCP-start + periodic maintenance ([`super::storage_maintenance`])
365/// and `lean-ctx cache prune`; without an enforcer the archive grew unbounded on
366/// disk and starved the host of RAM via the page cache (#417).
367pub fn cleanup() -> u32 {
368    let cutoff = Utc::now() - chrono::Duration::hours(max_age_hours() as i64);
369    cleanup_with(cutoff, max_disk_bytes())
370}
371
372/// Core of [`cleanup`], parameterized for testing: drop entries older than
373/// `cutoff`, then evict the oldest survivors until the total on-disk footprint is
374/// at or below `budget_bytes` (`0` = no size cap).
375fn cleanup_with(cutoff: DateTime<Utc>, budget_bytes: u64) -> u32 {
376    let base = archive_base_dir();
377    if !base.exists() {
378        return 0;
379    }
380
381    struct Scanned {
382        id: String,
383        created_at: DateTime<Utc>,
384        bytes: u64,
385    }
386
387    let mut entries: Vec<Scanned> = Vec::new();
388    if let Ok(dirs) = std::fs::read_dir(&base) {
389        for dir_entry in dirs.flatten() {
390            if !dir_entry.path().is_dir() {
391                continue;
392            }
393            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
394                for file in files.flatten() {
395                    let path = file.path();
396                    if path.extension().and_then(|e| e.to_str()) != Some("json") {
397                        continue;
398                    }
399                    let Ok(data) = std::fs::read_to_string(&path) else {
400                        continue;
401                    };
402                    let Ok(entry) = serde_json::from_str::<ArchiveEntry>(&data) else {
403                        continue;
404                    };
405                    let content_bytes =
406                        std::fs::metadata(content_path(&entry.id)).map_or(0, |m| m.len());
407                    let meta_bytes = file.metadata().map_or(0, |m| m.len());
408                    entries.push(Scanned {
409                        id: entry.id,
410                        created_at: entry.created_at,
411                        bytes: content_bytes + meta_bytes,
412                    });
413                }
414            }
415        }
416    }
417
418    // Oldest first: TTL victims drop first, then the oldest survivors are evicted
419    // until the store is back under budget. Sorted order lets us stop early.
420    entries.sort_by_key(|e| e.created_at);
421    let mut live_bytes: u64 = entries.iter().map(|e| e.bytes).sum();
422
423    let mut removed = 0u32;
424    for e in &entries {
425        let expired = e.created_at < cutoff;
426        let over_budget = budget_bytes > 0 && live_bytes > budget_bytes;
427        if !expired && !over_budget {
428            break;
429        }
430        remove_files(&e.id);
431        super::archive_fts::remove_entry(&e.id);
432        live_bytes = live_bytes.saturating_sub(e.bytes);
433        removed += 1;
434    }
435    removed
436}
437
438pub fn disk_usage_bytes() -> u64 {
439    let base = archive_base_dir();
440    if !base.exists() {
441        return 0;
442    }
443    let mut total = 0u64;
444    if let Ok(dirs) = std::fs::read_dir(&base) {
445        for dir_entry in dirs.flatten() {
446            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
447                for file in files.flatten() {
448                    total += file.metadata().map_or(0, |m| m.len());
449                }
450            }
451        }
452    }
453    total
454}
455
456pub fn format_hint(id: &str, size_chars: usize, size_tokens: usize) -> String {
457    format!("[Archived: {size_chars} chars ({size_tokens} tok). Retrieve: ctx_expand(id=\"{id}\")]")
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463
464    #[test]
465    fn compute_id_deterministic() {
466        let id1 = compute_id("test content");
467        let id2 = compute_id("test content");
468        assert_eq!(id1, id2);
469        let id3 = compute_id("different content");
470        assert_ne!(id1, id3);
471    }
472
473    #[test]
474    fn nonexistent_id_returns_none() {
475        assert!(retrieve("nonexistent_archive_id_xyz").is_none());
476    }
477
478    #[test]
479    fn format_hint_readable() {
480        let hint = format_hint("abc123", 5000, 1200);
481        assert!(hint.contains("5000 chars"));
482        assert!(hint.contains("1200 tok"));
483        assert!(hint.contains("ctx_expand"));
484        assert!(hint.contains("abc123"));
485    }
486
487    fn write_test_entry(id: &str, created_at: DateTime<Utc>, content_bytes: usize) {
488        std::fs::create_dir_all(entry_dir(id)).unwrap();
489        std::fs::write(content_path(id), "x".repeat(content_bytes)).unwrap();
490        let entry = ArchiveEntry {
491            id: id.to_string(),
492            tool: "ctx_shell".to_string(),
493            command: "test".to_string(),
494            size_chars: content_bytes,
495            size_tokens: content_bytes / 4,
496            created_at,
497            session_id: None,
498        };
499        std::fs::write(meta_path(id), serde_json::to_string(&entry).unwrap()).unwrap();
500    }
501
502    #[test]
503    fn cleanup_removes_expired_keeps_fresh() {
504        let _lock = crate::core::data_dir::test_env_lock();
505        let tmp = tempfile::tempdir().unwrap();
506        std::env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
507
508        let now = Utc::now();
509        write_test_entry("aa_old", now - chrono::Duration::hours(100), 100);
510        write_test_entry("bb_new", now - chrono::Duration::hours(1), 100);
511
512        // Cutoff = 48h ago; budget effectively unlimited so only the TTL applies.
513        let removed = cleanup_with(now - chrono::Duration::hours(48), u64::MAX);
514        assert_eq!(removed, 1);
515        assert!(!content_path("aa_old").exists());
516        assert!(!meta_path("aa_old").exists());
517        assert!(content_path("bb_new").exists());
518
519        std::env::remove_var("LEAN_CTX_DATA_DIR");
520    }
521
522    #[test]
523    fn cleanup_enforces_disk_budget_oldest_first() {
524        let _lock = crate::core::data_dir::test_env_lock();
525        let tmp = tempfile::tempdir().unwrap();
526        std::env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
527
528        let now = Utc::now();
529        write_test_entry("c1_oldest", now - chrono::Duration::minutes(30), 10_000);
530        write_test_entry("c2_middle", now - chrono::Duration::minutes(20), 10_000);
531        write_test_entry("c3_newest", now - chrono::Duration::minutes(10), 10_000);
532
533        // Nothing expired (cutoff far in the past). Budget 25 KB holds the two
534        // newest (~20 KB content + meta); the single oldest entry is evicted.
535        let removed = cleanup_with(now - chrono::Duration::days(365), 25_000);
536        assert_eq!(removed, 1, "only the oldest over-budget entry is evicted");
537        assert!(!content_path("c1_oldest").exists());
538        assert!(content_path("c2_middle").exists());
539        assert!(content_path("c3_newest").exists());
540
541        std::env::remove_var("LEAN_CTX_DATA_DIR");
542    }
543}
lean_ctx/core/archive.rs

lean_ctx/core/
archive.rs
