Skip to main content

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
75}
76
77pub fn should_archive(content: &str) -> bool {
78    is_enabled() && content.len() >= threshold_chars()
79}
80
81const MAX_ARCHIVE_SIZE: usize = 10 * 1024 * 1024; // 10 MB
82
83pub fn store(tool: &str, command: &str, content: &str, session_id: Option<&str>) -> Option<String> {
84    if !is_enabled() || content.is_empty() {
85        return None;
86    }
87
88    let content = if content.len() > MAX_ARCHIVE_SIZE {
89        &content[..content.floor_char_boundary(MAX_ARCHIVE_SIZE)]
90    } else {
91        content
92    };
93
94    let id = compute_id(content);
95    let c_path = content_path(&id);
96
97    // Fast path: content already archived (idempotent, no race)
98    if c_path.exists() {
99        return Some(id);
100    }
101
102    let dir = entry_dir(&id);
103    if std::fs::create_dir_all(&dir).is_err() {
104        return None;
105    }
106
107    // Atomic write: PID-unique tmp file prevents race between parallel writers.
108    // rename() is atomic on POSIX; on Windows it replaces atomically too.
109    // If two processes race past the exists() check, both write their own tmp
110    // file and both rename to the same target — last writer wins, content is
111    // identical (same hash), so the result is correct either way.
112    let pid = std::process::id();
113    let tmp_path = c_path.with_extension(format!("tmp.{pid}"));
114    if std::fs::write(&tmp_path, content).is_err() {
115        return None;
116    }
117    if std::fs::rename(&tmp_path, &c_path).is_err() {
118        let _ = std::fs::remove_file(&tmp_path);
119        // Another process may have won the race — check if content is there now
120        if c_path.exists() {
121            return Some(id);
122        }
123        return None;
124    }
125    #[cfg(unix)]
126    set_private_file_perms(&c_path);
127
128    let tokens = super::tokens::count_tokens(content);
129    let entry = ArchiveEntry {
130        id: id.clone(),
131        tool: tool.to_string(),
132        command: command.to_string(),
133        size_chars: content.len(),
134        size_tokens: tokens,
135        created_at: Utc::now(),
136        session_id: session_id.map(std::string::ToString::to_string),
137    };
138
139    if let Ok(json) = serde_json::to_string_pretty(&entry) {
140        let meta_tmp = meta_path(&id).with_extension(format!("tmp.{pid}"));
141        if std::fs::write(&meta_tmp, &json).is_ok() {
142            let meta_final = meta_path(&id);
143            let _ = std::fs::rename(&meta_tmp, &meta_final);
144            #[cfg(unix)]
145            set_private_file_perms(&meta_final);
146        }
147    }
148
149    super::archive_fts::index_entry(&id, tool, command, content);
150
151    Some(id)
152}
153
154pub fn retrieve(id: &str) -> Option<String> {
155    let path = content_path(id);
156    std::fs::read_to_string(path).ok()
157}
158
159pub fn retrieve_with_range(id: &str, start: usize, end: usize) -> Option<String> {
160    let content = retrieve(id)?;
161    let lines: Vec<&str> = content.lines().collect();
162    let start = start.saturating_sub(1).min(lines.len());
163    let end = end.min(lines.len());
164    if start >= end {
165        return Some(String::new());
166    }
167    Some(
168        lines[start..end]
169            .iter()
170            .enumerate()
171            .map(|(i, line)| format!("{:>6}|{line}", start + i + 1))
172            .collect::<Vec<_>>()
173            .join("\n"),
174    )
175}
176
177pub fn retrieve_with_search(id: &str, pattern: &str) -> Option<String> {
178    let content = retrieve(id)?;
179    let pattern_lower = pattern.to_lowercase();
180    let matches: Vec<String> = content
181        .lines()
182        .enumerate()
183        .filter(|(_, line)| line.to_lowercase().contains(&pattern_lower))
184        .map(|(i, line)| format!("{:>6}|{line}", i + 1))
185        .collect();
186
187    if matches.is_empty() {
188        Some(format!("No matches for \"{pattern}\" in archive {id}"))
189    } else {
190        Some(format!(
191            "{} match(es) for \"{}\":\n{}",
192            matches.len(),
193            pattern,
194            matches.join("\n")
195        ))
196    }
197}
198
199/// Retrieve the first `n` lines of an archived entry, with a line-number gutter.
200pub fn retrieve_head(id: &str, n: usize) -> Option<String> {
201    retrieve_with_range(id, 1, n)
202}
203
204/// Retrieve the last `n` lines of an archived entry, with a line-number gutter.
205pub fn retrieve_tail(id: &str, n: usize) -> Option<String> {
206    let content = retrieve(id)?;
207    let total = content.lines().count();
208    let start = if total > n { total - n + 1 } else { 1 };
209    retrieve_with_range(id, start, total)
210}
211
212/// Describe the JSON structure of an archived entry: top-level keys with type hints,
213/// array lengths + element types, etc. An optional dot/slash `path` (e.g. `data.items.0`)
214/// navigates into the structure first. Returns `None` when the archive is missing or its
215/// content is not valid JSON, so callers can fall back to a raw retrieval hint.
216pub fn retrieve_json_keys(id: &str, path: Option<&str>) -> Option<String> {
217    let content = retrieve(id)?;
218    let root: serde_json::Value = serde_json::from_str(content.trim()).ok()?;
219    let mut cur = &root;
220    let mut walked = String::from("$");
221    if let Some(p) = path {
222        for seg in p.split(['.', '/']).filter(|s| !s.is_empty()) {
223            let next = if let Ok(idx) = seg.parse::<usize>() {
224                cur.get(idx)
225            } else {
226                cur.get(seg)
227            };
228            match next {
229                Some(v) => {
230                    cur = v;
231                    walked.push('.');
232                    walked.push_str(seg);
233                }
234                None => {
235                    return Some(format!(
236                        "Path '{p}' not found at '{walked}' in archive {id}"
237                    ));
238                }
239            }
240        }
241    }
242    Some(format!("{walked} => {}", describe_json(cur)))
243}
244
245fn json_type_hint(v: &serde_json::Value) -> String {
246    use serde_json::Value;
247    match v {
248        Value::Object(m) => format!("object({})", m.len()),
249        Value::Array(a) => format!("array({})", a.len()),
250        Value::String(s) => {
251            let preview: String = s.chars().take(40).collect();
252            if s.chars().count() > 40 {
253                format!("string \"{preview}…\"")
254            } else {
255                format!("string \"{preview}\"")
256            }
257        }
258        Value::Number(n) => format!("number {n}"),
259        Value::Bool(b) => format!("bool {b}"),
260        Value::Null => "null".to_string(),
261    }
262}
263
264fn describe_json(v: &serde_json::Value) -> String {
265    use serde_json::Value;
266    match v {
267        Value::Object(map) => {
268            let mut keys: Vec<&String> = map.keys().collect();
269            keys.sort();
270            let rendered: Vec<String> = keys
271                .iter()
272                .map(|k| format!("  {k}: {}", json_type_hint(&map[*k])))
273                .collect();
274            format!("object ({} keys)\n{}", map.len(), rendered.join("\n"))
275        }
276        Value::Array(arr) => {
277            let elem = arr.first().map_or("empty", |e| match e {
278                Value::Object(_) => "object",
279                Value::Array(_) => "array",
280                Value::String(_) => "string",
281                Value::Number(_) => "number",
282                Value::Bool(_) => "bool",
283                Value::Null => "null",
284            });
285            let mut out = format!("array ({} items of {elem})", arr.len());
286            if let Some(Value::Object(map)) = arr.first() {
287                let mut keys: Vec<&String> = map.keys().collect();
288                keys.sort();
289                out.push_str(&format!(
290                    "\n  [0] keys: {}",
291                    keys.iter()
292                        .map(|s| s.as_str())
293                        .collect::<Vec<_>>()
294                        .join(", ")
295                ));
296            }
297            out
298        }
299        Value::String(s) => format!("string ({} chars)", s.len()),
300        Value::Number(n) => format!("number ({n})"),
301        Value::Bool(b) => format!("bool ({b})"),
302        Value::Null => "null".to_string(),
303    }
304}
305
306pub fn list_entries(session_id: Option<&str>) -> Vec<ArchiveEntry> {
307    let base = archive_base_dir();
308    if !base.exists() {
309        return Vec::new();
310    }
311    let mut entries = Vec::new();
312    if let Ok(dirs) = std::fs::read_dir(&base) {
313        for dir_entry in dirs.flatten() {
314            if !dir_entry.path().is_dir() {
315                continue;
316            }
317            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
318                for file in files.flatten() {
319                    let path = file.path();
320                    if path.extension().and_then(|e| e.to_str()) != Some("json") {
321                        continue;
322                    }
323                    if let Ok(data) = std::fs::read_to_string(&path) {
324                        if let Ok(entry) = serde_json::from_str::<ArchiveEntry>(&data) {
325                            if let Some(sid) = session_id {
326                                if entry.session_id.as_deref() != Some(sid) {
327                                    continue;
328                                }
329                            }
330                            entries.push(entry);
331                        }
332                    }
333                }
334            }
335        }
336    }
337    entries.sort_by_key(|e| std::cmp::Reverse(e.created_at));
338    entries
339}
340
341pub fn cleanup() -> u32 {
342    let max_hours = max_age_hours();
343    let cutoff = Utc::now() - chrono::Duration::hours(max_hours as i64);
344    let base = archive_base_dir();
345    if !base.exists() {
346        return 0;
347    }
348    let mut removed = 0u32;
349    if let Ok(dirs) = std::fs::read_dir(&base) {
350        for dir_entry in dirs.flatten() {
351            if !dir_entry.path().is_dir() {
352                continue;
353            }
354            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
355                for file in files.flatten() {
356                    let path = file.path();
357                    if path.extension().and_then(|e| e.to_str()) != Some("json") {
358                        continue;
359                    }
360                    if let Ok(data) = std::fs::read_to_string(&path) {
361                        if let Ok(entry) = serde_json::from_str::<ArchiveEntry>(&data) {
362                            if entry.created_at < cutoff {
363                                let c = content_path(&entry.id);
364                                let _ = std::fs::remove_file(&c);
365                                let _ = std::fs::remove_file(&path);
366                                super::archive_fts::remove_entry(&entry.id);
367                                removed += 1;
368                            }
369                        }
370                    }
371                }
372            }
373        }
374    }
375    removed
376}
377
378pub fn disk_usage_bytes() -> u64 {
379    let base = archive_base_dir();
380    if !base.exists() {
381        return 0;
382    }
383    let mut total = 0u64;
384    if let Ok(dirs) = std::fs::read_dir(&base) {
385        for dir_entry in dirs.flatten() {
386            if let Ok(files) = std::fs::read_dir(dir_entry.path()) {
387                for file in files.flatten() {
388                    total += file.metadata().map_or(0, |m| m.len());
389                }
390            }
391        }
392    }
393    total
394}
395
396pub fn format_hint(id: &str, size_chars: usize, size_tokens: usize) -> String {
397    format!("[Archived: {size_chars} chars ({size_tokens} tok). Retrieve: ctx_expand(id=\"{id}\")]")
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403
404    #[test]
405    fn compute_id_deterministic() {
406        let id1 = compute_id("test content");
407        let id2 = compute_id("test content");
408        assert_eq!(id1, id2);
409        let id3 = compute_id("different content");
410        assert_ne!(id1, id3);
411    }
412
413    #[test]
414    fn nonexistent_id_returns_none() {
415        assert!(retrieve("nonexistent_archive_id_xyz").is_none());
416    }
417
418    #[test]
419    fn format_hint_readable() {
420        let hint = format_hint("abc123", 5000, 1200);
421        assert!(hint.contains("5000 chars"));
422        assert!(hint.contains("1200 tok"));
423        assert!(hint.contains("ctx_expand"));
424        assert!(hint.contains("abc123"));
425    }
426}