lean_ctx/core/cache/
warming.rs

1use std::collections::HashMap;
2use std::time::SystemTime;
3
4use crate::core::session::SessionState;
5use crate::core::tokens::count_tokens;
6
7use super::{SessionCache, max_cache_tokens};
8
9const MAX_WARM_FILE_BYTES: usize = 50 * 1024;
10const MAX_RECENT_FILES: usize = 50;
11
12/// A file accessed by one or more previous sessions.
13#[derive(Clone, Debug, Eq, PartialEq)]
14pub struct RecentFile {
15    pub path: String,
16    pub last_access: SystemTime,
17    pub read_count: u32,
18}
19
20trait WarmCache {
21    fn contains(&self, path: &str) -> bool;
22    fn total_tokens(&self) -> usize;
23    fn insert(&mut self, path: &str, content: &str);
24}
25
26impl WarmCache for SessionCache {
27    fn contains(&self, path: &str) -> bool {
28        self.get(path).is_some()
29    }
30
31    fn total_tokens(&self) -> usize {
32        self.total_cached_tokens()
33    }
34
35    fn insert(&mut self, path: &str, content: &str) {
36        self.store(path, content);
37    }
38}
39
40/// Pre-populates a session cache with the highest-priority recent files.
41///
42/// Warming stops when the cache reaches 80% of its configured token budget.
43/// Missing, oversized, unreadable, and already-cached files are ignored.
44pub fn warm_cache(cache: &mut SessionCache, history: &[RecentFile]) {
45    warm_cache_with(cache, history, max_cache_tokens(), |path| {
46        let metadata = std::fs::metadata(path).ok()?;
47        if !metadata.is_file() || metadata.len() > MAX_WARM_FILE_BYTES as u64 {
48            return None;
49        }
50        crate::core::io_boundary::read_file_lossy(path).ok()
51    });
52}
53
54fn warm_cache_with<C, F>(cache: &mut C, history: &[RecentFile], token_budget: usize, mut load: F)
55where
56    C: WarmCache,
57    F: FnMut(&str) -> Option<String>,
58{
59    let warming_limit = token_budget.saturating_mul(4) / 5;
60
61    for recent in history {
62        if cache.total_tokens() >= warming_limit {
63            break;
64        }
65        if cache.contains(&recent.path) {
66            continue;
67        }
68
69        let Some(content) = load(&recent.path) else {
70            continue;
71        };
72        if content.len() > MAX_WARM_FILE_BYTES {
73            continue;
74        }
75
76        let incoming_tokens = count_tokens(&content);
77        if cache.total_tokens().saturating_add(incoming_tokens) > warming_limit {
78            continue;
79        }
80        cache.insert(&recent.path, &content);
81    }
82}
83
84/// Collects, deduplicates, and ranks files touched by previous sessions.
85///
86/// Read counts are summed across sessions. A session's update time is used as
87/// the recency of its file accesses because `FileTouched` has no own timestamp.
88pub fn collect_recent_files(sessions: &[SessionState]) -> Vec<RecentFile> {
89    let mut recent_by_path: HashMap<String, RecentFile> = HashMap::new();
90
91    for session in sessions {
92        let last_access = SystemTime::from(session.updated_at);
93        for file in &session.files_touched {
94            let recent = recent_by_path
95                .entry(file.path.clone())
96                .or_insert_with(|| RecentFile {
97                    path: file.path.clone(),
98                    last_access,
99                    read_count: 0,
100                });
101            recent.read_count = recent.read_count.saturating_add(file.read_count);
102            recent.last_access = recent.last_access.max(last_access);
103        }
104    }
105
106    let mut recent: Vec<RecentFile> = recent_by_path.into_values().collect();
107    recent.sort_by(|a, b| {
108        b.read_count
109            .cmp(&a.read_count)
110            .then_with(|| b.last_access.cmp(&a.last_access))
111            .then_with(|| a.path.cmp(&b.path))
112    });
113    recent.truncate(MAX_RECENT_FILES);
114    recent
115}
116
117#[cfg(test)]
118mod tests {
119    use std::collections::HashMap;
120    use std::time::{Duration, SystemTime};
121
122    use chrono::{TimeZone, Utc};
123
124    use super::*;
125
126    #[derive(Default)]
127    struct MemoryCache {
128        entries: HashMap<String, String>,
129    }
130
131    impl WarmCache for MemoryCache {
132        fn contains(&self, path: &str) -> bool {
133            self.entries.contains_key(path)
134        }
135
136        fn total_tokens(&self) -> usize {
137            self.entries
138                .values()
139                .map(|content| count_tokens(content))
140                .sum()
141        }
142
143        fn insert(&mut self, path: &str, content: &str) {
144            self.entries.insert(path.to_string(), content.to_string());
145        }
146    }
147
148    fn recent(path: &str) -> RecentFile {
149        RecentFile {
150            path: path.to_string(),
151            last_access: SystemTime::UNIX_EPOCH,
152            read_count: 1,
153        }
154    }
155
156    fn session_at(timestamp: i64) -> SessionState {
157        let time = Utc.timestamp_opt(timestamp, 0).unwrap();
158        SessionState {
159            id: format!("session-{timestamp}"),
160            version: 0,
161            started_at: time,
162            updated_at: time,
163            project_root: None,
164            shell_cwd: None,
165            task: None,
166            findings: Vec::new(),
167            decisions: Vec::new(),
168            files_touched: Vec::new(),
169            test_results: None,
170            progress: Vec::new(),
171            next_steps: Vec::new(),
172            evidence: Vec::new(),
173            intents: Vec::new(),
174            active_structured_intent: None,
175            stats: crate::core::session::SessionStats::default(),
176            terse_mode: false,
177            compression_level: String::new(),
178            last_consolidate_ts: None,
179            last_aaak_hash: None,
180            extra_roots: Vec::new(),
181            wakeup_manifest: Vec::new(),
182            playbook: crate::core::session::Playbook::default(),
183            last_semantic_query: None,
184            last_flush: None,
185        }
186    }
187
188    #[test]
189    fn warming_skips_cached_missing_and_oversized_files() {
190        let mut cache = MemoryCache::default();
191        cache
192            .entries
193            .insert("cached.rs".to_string(), "cached".to_string());
194        let history = [
195            recent("cached.rs"),
196            recent("missing.rs"),
197            recent("large.rs"),
198            recent("good.rs"),
199        ];
200        let contents = HashMap::from([
201            ("large.rs", "x".repeat(MAX_WARM_FILE_BYTES + 1)),
202            ("good.rs", "fn good() {}".to_string()),
203        ]);
204        let mut loaded = Vec::new();
205
206        warm_cache_with(&mut cache, &history, 10_000, |path| {
207            loaded.push(path.to_string());
208            contents.get(path).cloned()
209        });
210
211        assert!(!loaded.contains(&"cached.rs".to_string()));
212        assert!(!cache.entries.contains_key("missing.rs"));
213        assert!(!cache.entries.contains_key("large.rs"));
214        assert_eq!(
215            cache.entries.get("good.rs").map(String::as_str),
216            Some("fn good() {}")
217        );
218    }
219
220    #[test]
221    fn warming_stops_at_eighty_percent_of_budget() {
222        let first_content = "first file content";
223        let first_tokens = count_tokens(first_content);
224        let budget = first_tokens.saturating_mul(5).div_ceil(4);
225        let history = [recent("first.rs"), recent("second.rs")];
226        let mut cache = MemoryCache::default();
227        let mut loaded = Vec::new();
228
229        warm_cache_with(&mut cache, &history, budget, |path| {
230            loaded.push(path.to_string());
231            Some(first_content.to_string())
232        });
233
234        assert!(cache.entries.contains_key("first.rs"));
235        assert!(!cache.entries.contains_key("second.rs"));
236        assert_eq!(loaded, vec!["first.rs"]);
237        assert!(cache.total_tokens() <= budget.saturating_mul(4) / 5);
238    }
239
240    #[test]
241    fn recent_files_are_deduplicated_and_ranked_by_frequency_then_recency() {
242        let mut older = session_at(10);
243        older.touch_file("frequent.rs", None, "full", 10);
244        older.touch_file("frequent.rs", None, "full", 10);
245        older.touch_file("recent.rs", None, "full", 10);
246        older.updated_at = Utc.timestamp_opt(10, 0).unwrap();
247
248        let mut newer = session_at(20);
249        newer.touch_file("frequent.rs", None, "full", 10);
250        newer.touch_file("recent.rs", None, "full", 10);
251        newer.touch_file("peer.rs", None, "full", 10);
252        newer.updated_at = Utc.timestamp_opt(20, 0).unwrap();
253
254        let result = collect_recent_files(&[older, newer]);
255
256        assert_eq!(
257            result
258                .iter()
259                .map(|file| file.path.as_str())
260                .collect::<Vec<_>>(),
261            vec!["frequent.rs", "recent.rs", "peer.rs"]
262        );
263        assert_eq!(result[0].read_count, 3);
264        assert_eq!(
265            result[0].last_access,
266            SystemTime::UNIX_EPOCH + Duration::from_secs(20)
267        );
268    }
269
270    #[test]
271    fn recent_files_are_capped_at_fifty() {
272        let mut session = session_at(10);
273        for index in 0..60 {
274            session.touch_file(&format!("file-{index:02}.rs"), None, "full", 1);
275        }
276
277        let result = collect_recent_files(&[session]);
278
279        assert_eq!(result.len(), MAX_RECENT_FILES);
280    }
281}
lean_ctx/core/cache/warming.rs

lean_ctx/core/cache/
warming.rs
