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