Skip to main content

lean_ctx/core/
cli_cache.rs

1use md5::{Digest, Md5};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::PathBuf;
5use std::time::{SystemTime, UNIX_EPOCH};
6
7const CACHE_TTL_SECS: u64 = 300;
8const MAX_ENTRIES: usize = 200;
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct CliCacheEntry {
12    pub path: String,
13    pub hash: String,
14    pub line_count: usize,
15    pub original_tokens: usize,
16    pub timestamp: u64,
17    pub read_count: u32,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize, Default)]
21pub struct CliCacheStore {
22    pub entries: HashMap<String, CliCacheEntry>,
23    pub total_hits: u64,
24    pub total_reads: u64,
25}
26
27pub enum CacheResult {
28    Hit {
29        entry: CliCacheEntry,
30        file_ref: String,
31    },
32    Miss {
33        content: String,
34    },
35}
36
37fn cache_dir() -> Option<PathBuf> {
38    crate::core::data_dir::lean_ctx_data_dir()
39        .ok()
40        .map(|d| d.join("cli-cache"))
41}
42
43fn cache_file() -> Option<PathBuf> {
44    cache_dir().map(|d| d.join("cache.json"))
45}
46
47fn now_secs() -> u64 {
48    SystemTime::now()
49        .duration_since(UNIX_EPOCH)
50        .unwrap_or_default()
51        .as_secs()
52}
53
54fn compute_md5(content: &str) -> String {
55    let mut hasher = Md5::new();
56    hasher.update(content.as_bytes());
57    format!("{:x}", hasher.finalize())
58}
59
60fn normalize_key(path: &str) -> String {
61    crate::hooks::normalize_tool_path(path)
62}
63
64fn load_store() -> CliCacheStore {
65    let path = match cache_file() {
66        Some(p) => p,
67        None => return CliCacheStore::default(),
68    };
69    match std::fs::read_to_string(&path) {
70        Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
71        Err(_) => CliCacheStore::default(),
72    }
73}
74
75fn save_store(store: &CliCacheStore) {
76    let dir = match cache_dir() {
77        Some(d) => d,
78        None => return,
79    };
80    let _ = std::fs::create_dir_all(&dir);
81    let path = dir.join("cache.json");
82    if let Ok(data) = serde_json::to_string(store) {
83        let _ = std::fs::write(path, data);
84    }
85}
86
87fn file_ref(key: &str, store: &CliCacheStore) -> String {
88    let keys: Vec<&String> = store.entries.keys().collect();
89    let idx = keys
90        .iter()
91        .position(|k| k.as_str() == key)
92        .unwrap_or(store.entries.len());
93    format!("F{}", idx + 1)
94}
95
96pub fn check_and_read(path: &str) -> CacheResult {
97    let content = match crate::tools::ctx_read::read_file_lossy(path) {
98        Ok(c) => c,
99        Err(_) => {
100            return CacheResult::Miss {
101                content: String::new(),
102            }
103        }
104    };
105
106    let key = normalize_key(path);
107    let hash = compute_md5(&content);
108    let now = now_secs();
109    let mut store = load_store();
110
111    store.total_reads += 1;
112
113    if let Some(entry) = store.entries.get_mut(&key) {
114        if entry.hash == hash && (now - entry.timestamp) < CACHE_TTL_SECS {
115            entry.read_count += 1;
116            entry.timestamp = now;
117            store.total_hits += 1;
118            let result = CacheResult::Hit {
119                entry: entry.clone(),
120                file_ref: file_ref(&key, &store),
121            };
122            save_store(&store);
123            return result;
124        }
125    }
126
127    let line_count = content.lines().count();
128    let original_tokens = crate::core::tokens::count_tokens(&content);
129
130    let entry = CliCacheEntry {
131        path: key.clone(),
132        hash,
133        line_count,
134        original_tokens,
135        timestamp: now,
136        read_count: 1,
137    };
138    store.entries.insert(key, entry);
139
140    evict_stale(&mut store, now);
141
142    save_store(&store);
143    CacheResult::Miss { content }
144}
145
146pub fn invalidate(path: &str) {
147    let key = normalize_key(path);
148    let mut store = load_store();
149    store.entries.remove(&key);
150    save_store(&store);
151}
152
153pub fn clear() -> usize {
154    let mut store = load_store();
155    let count = store.entries.len();
156    store.entries.clear();
157    save_store(&store);
158    count
159}
160
161pub fn clear_project(project_root: &str) -> usize {
162    let mut store = load_store();
163    let prefix = normalize_key(project_root);
164    let before = store.entries.len();
165    store
166        .entries
167        .retain(|key, entry| !key.starts_with(&prefix) && !entry.path.starts_with(&prefix));
168    let removed = before - store.entries.len();
169    save_store(&store);
170    removed
171}
172
173pub fn stats() -> (u64, u64, usize) {
174    let store = load_store();
175    (store.total_hits, store.total_reads, store.entries.len())
176}
177
178fn evict_stale(store: &mut CliCacheStore, now: u64) {
179    store
180        .entries
181        .retain(|_, e| (now - e.timestamp) < CACHE_TTL_SECS);
182
183    if store.entries.len() > MAX_ENTRIES {
184        let mut entries: Vec<(String, u64)> = store
185            .entries
186            .iter()
187            .map(|(k, e)| (k.clone(), e.timestamp))
188            .collect();
189        entries.sort_by_key(|(_, ts)| *ts);
190        let to_remove = store.entries.len() - MAX_ENTRIES;
191        for (key, _) in entries.into_iter().take(to_remove) {
192            store.entries.remove(&key);
193        }
194    }
195}
196
197pub fn format_hit(entry: &CliCacheEntry, file_ref: &str, short_path: &str) -> String {
198    format!(
199        "{file_ref} cached {short_path} [{}L {}t] (read #{})",
200        entry.line_count, entry.original_tokens, entry.read_count
201    )
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn compute_md5_deterministic() {
210        let h1 = compute_md5("test content");
211        let h2 = compute_md5("test content");
212        assert_eq!(h1, h2);
213        assert_ne!(h1, compute_md5("different"));
214    }
215
216    #[test]
217    fn evict_stale_removes_old_entries() {
218        let mut store = CliCacheStore::default();
219        store.entries.insert(
220            "/old.rs".to_string(),
221            CliCacheEntry {
222                path: "/old.rs".to_string(),
223                hash: "h1".into(),
224                line_count: 10,
225                original_tokens: 50,
226                timestamp: 1000,
227                read_count: 1,
228            },
229        );
230        store.entries.insert(
231            "/new.rs".to_string(),
232            CliCacheEntry {
233                path: "/new.rs".to_string(),
234                hash: "h2".into(),
235                line_count: 20,
236                original_tokens: 100,
237                timestamp: now_secs(),
238                read_count: 1,
239            },
240        );
241
242        evict_stale(&mut store, now_secs());
243        assert!(!store.entries.contains_key("/old.rs"));
244        assert!(store.entries.contains_key("/new.rs"));
245    }
246
247    #[test]
248    fn evict_respects_max_entries() {
249        let mut store = CliCacheStore::default();
250        let now = now_secs();
251        for i in 0..MAX_ENTRIES + 10 {
252            store.entries.insert(
253                format!("/file_{i}.rs"),
254                CliCacheEntry {
255                    path: format!("/file_{i}.rs"),
256                    hash: format!("h{i}"),
257                    line_count: 1,
258                    original_tokens: 10,
259                    timestamp: now - i as u64,
260                    read_count: 1,
261                },
262            );
263        }
264        evict_stale(&mut store, now);
265        assert!(store.entries.len() <= MAX_ENTRIES);
266    }
267
268    #[test]
269    fn format_hit_output() {
270        let entry = CliCacheEntry {
271            path: "/test.rs".into(),
272            hash: "abc".into(),
273            line_count: 42,
274            original_tokens: 500,
275            timestamp: now_secs(),
276            read_count: 3,
277        };
278        let output = format_hit(&entry, "F1", "test.rs");
279        assert!(output.contains("F1 cached"));
280        assert!(output.contains("42L"));
281        assert!(output.contains("500t"));
282        assert!(output.contains("read #3"));
283    }
284
285    #[test]
286    fn stats_returns_defaults_on_empty() {
287        let s = CliCacheStore::default();
288        assert_eq!(s.total_hits, 0);
289        assert_eq!(s.total_reads, 0);
290        assert!(s.entries.is_empty());
291    }
292
293    #[test]
294    fn cache_result_integration() {
295        let _lock = crate::core::data_dir::test_env_lock();
296
297        let nanos = std::time::SystemTime::now()
298            .duration_since(std::time::UNIX_EPOCH)
299            .unwrap()
300            .as_nanos();
301        let test_data_dir = std::env::temp_dir().join(format!("lean_ctx_cache_iso_{nanos}"));
302        std::fs::create_dir_all(&test_data_dir).unwrap();
303        std::env::set_var("LEAN_CTX_DATA_DIR", &test_data_dir);
304
305        let tmp = test_data_dir.join("test_file.txt");
306        std::fs::write(&tmp, "fn main() {}\n").unwrap();
307        let path_str = tmp.to_str().unwrap();
308
309        invalidate(path_str);
310
311        let result = check_and_read(path_str);
312        assert!(matches!(result, CacheResult::Miss { .. }));
313
314        let result2 = check_and_read(path_str);
315        assert!(matches!(result2, CacheResult::Hit { .. }));
316        if let CacheResult::Hit { entry, .. } = result2 {
317            assert_eq!(entry.line_count, 1);
318            assert!(entry.read_count >= 2);
319        }
320
321        invalidate(path_str);
322        let result3 = check_and_read(path_str);
323        assert!(matches!(result3, CacheResult::Miss { .. }));
324
325        std::env::remove_var("LEAN_CTX_DATA_DIR");
326        let _ = std::fs::remove_dir_all(&test_data_dir);
327    }
328}