lean_ctx/core/
mode_predictor.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Mutex};
3use std::time::Instant;
4
5const STATS_FILE: &str = "mode_stats.json";
6const PREDICTOR_FLUSH_SECS: u64 = 10;
7
8static PREDICTOR_BUFFER: Mutex<Option<(Arc<ModePredictor>, Instant)>> = Mutex::new(None);
9
10/// Observed outcome of a read mode: tokens in/out and information density.
11#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
12pub struct ModeOutcome {
13    pub mode: String,
14    pub tokens_in: usize,
15    pub tokens_out: usize,
16    pub density: f64,
17}
18
19impl ModeOutcome {
20    /// Computes an efficiency score: density / compression ratio.
21    pub fn efficiency(&self) -> f64 {
22        if self.tokens_out == 0 {
23            return 0.0;
24        }
25        self.density / (self.tokens_out as f64 / self.tokens_in.max(1) as f64)
26    }
27}
28
29/// File identity for mode prediction: extension + token-count size bucket.
30#[derive(Clone, Debug, Hash, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
31pub struct FileSignature {
32    pub ext: String,
33    pub size_bucket: u8,
34}
35
36impl FileSignature {
37    /// Creates a file signature from its path and token count.
38    pub fn from_path(path: &str, token_count: usize) -> Self {
39        let ext = std::path::Path::new(path)
40            .extension()
41            .and_then(|e| e.to_str())
42            .unwrap_or("")
43            .to_string();
44        let size_bucket = match token_count {
45            0..=500 => 0,
46            501..=2000 => 1,
47            2001..=5000 => 2,
48            5001..=20000 => 3,
49            _ => 4,
50        };
51        Self { ext, size_bucket }
52    }
53}
54
55/// Learns the best read mode per file signature from historical outcomes.
56#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
57pub struct ModePredictor {
58    history: HashMap<FileSignature, Vec<ModeOutcome>>,
59    project_root: Option<String>,
60}
61
62impl ModePredictor {
63    /// Loads or creates the predictor, using an in-memory buffer for caching.
64    pub fn new() -> Self {
65        let mut guard = PREDICTOR_BUFFER
66            .lock()
67            .unwrap_or_else(std::sync::PoisonError::into_inner);
68        if let Some((ref predictor, _)) = *guard {
69            return Self {
70                history: predictor.history.clone(),
71                project_root: predictor.project_root.clone(),
72            };
73        }
74        let mut loaded = Self::load_from_disk().unwrap_or_default();
75        if loaded.project_root.is_none() {
76            loaded.project_root = std::env::current_dir()
77                .ok()
78                .map(|p| p.to_string_lossy().to_string());
79        }
80        *guard = Some((Arc::new(loaded.clone()), Instant::now()));
81        loaded
82    }
83
84    pub fn with_project_root(mut self, root: &str) -> Self {
85        self.project_root = Some(root.to_string());
86        self
87    }
88
89    pub fn set_project_root(&mut self, root: &str) {
90        self.project_root = Some(root.to_string());
91    }
92
93    /// Records a mode outcome for a file signature (capped at 100 entries).
94    pub fn record(&mut self, sig: FileSignature, outcome: ModeOutcome) {
95        let entries = self.history.entry(sig).or_default();
96        entries.push(outcome);
97        if entries.len() > 100 {
98            entries.drain(0..50);
99        }
100    }
101
102    /// Returns the best mode based on historical efficiency.
103    /// Chain: local history -> cloud adaptive models -> built-in defaults.
104    pub fn predict_best_mode(&self, sig: &FileSignature) -> Option<String> {
105        let default_mode = Self::predict_from_defaults(sig);
106
107        let allow_override = |candidate: &str| -> bool {
108            let Some(def) = default_mode.as_deref() else {
109                return true;
110            };
111            if candidate == "full" {
112                return false;
113            }
114            // For code-structured defaults, never override to lossy modes.
115            if (def == "map" || def == "signatures")
116                && (candidate == "aggressive" || candidate == "entropy")
117            {
118                return false;
119            }
120            true
121        };
122
123        if let Some(local) = self.predict_from_local(sig)
124            && allow_override(&local)
125        {
126            return Some(local);
127        }
128        if let Some(bandit) = self.predict_from_bandit(sig)
129            && allow_override(&bandit)
130        {
131            return Some(bandit);
132        }
133        if let Some(cloud) = self.predict_from_cloud(sig)
134            && allow_override(&cloud)
135        {
136            return Some(cloud);
137        }
138        default_mode
139    }
140
141    fn predict_from_bandit(&self, sig: &FileSignature) -> Option<String> {
142        let key = format!("{}_feedback", sig.ext);
143        let store =
144            crate::core::bandit::BanditStore::load(self.project_root.as_deref().unwrap_or("."));
145        let bandit = store.bandits.get(&key)?;
146        if bandit.total_pulls < 5 {
147            return None;
148        }
149        let best_arm = bandit.arms.iter().max_by(|a, b| {
150            a.mean()
151                .partial_cmp(&b.mean())
152                .unwrap_or(std::cmp::Ordering::Equal)
153        })?;
154        let mode = match best_arm.name.as_str() {
155            "conservative" => "full",
156            "balanced" => "signatures",
157            "aggressive" => "aggressive",
158            _ => return None,
159        };
160        Some(mode.to_string())
161    }
162
163    fn predict_from_local(&self, sig: &FileSignature) -> Option<String> {
164        let entries = self.history.get(sig)?;
165        if entries.len() < 3 {
166            return None;
167        }
168
169        let mut mode_scores: HashMap<&str, (f64, usize)> = HashMap::new();
170        for entry in entries {
171            let (sum, count) = mode_scores.entry(&entry.mode).or_insert((0.0, 0));
172            *sum += entry.efficiency();
173            *count += 1;
174        }
175
176        mode_scores
177            .into_iter()
178            .max_by(|a, b| {
179                let avg_a = a.1.0 / a.1.1 as f64;
180                let avg_b = b.1.0 / b.1.1 as f64;
181                avg_a
182                    .partial_cmp(&avg_b)
183                    .unwrap_or(std::cmp::Ordering::Equal)
184            })
185            .map(|(mode, _)| mode.to_string())
186    }
187
188    /// Loads cloud adaptive models (synced from LeanCTX Cloud).
189    /// Models are cached locally and auto-updated for cloud users.
190    #[allow(clippy::unused_self)]
191    fn predict_from_cloud(&self, sig: &FileSignature) -> Option<String> {
192        let data = crate::cloud_client::load_cloud_models()?;
193        let models = data["models"].as_array()?;
194
195        let ext_with_dot = format!(".{}", sig.ext);
196        let bucket_name = match sig.size_bucket {
197            0 => "0-500",
198            1 => "500-2k",
199            2 => "2k-10k",
200            _ => "10k+",
201        };
202
203        let mut best: Option<(&str, f64)> = None;
204
205        for model in models {
206            let m_ext = model["file_ext"].as_str().unwrap_or("");
207            let m_bucket = model["size_bucket"].as_str().unwrap_or("");
208            let confidence = model["confidence"].as_f64().unwrap_or(0.0);
209
210            if m_ext == ext_with_dot
211                && m_bucket == bucket_name
212                && confidence > 0.5
213                && let Some(mode) = model["recommended_mode"].as_str()
214                && best.is_none_or(|(_, c)| confidence > c)
215            {
216                best = Some((mode, confidence));
217            }
218        }
219
220        if let Some((mode, _)) = best {
221            return Some(mode.to_string());
222        }
223
224        for model in models {
225            let m_ext = model["file_ext"].as_str().unwrap_or("");
226            let confidence = model["confidence"].as_f64().unwrap_or(0.0);
227            if m_ext == ext_with_dot && confidence > 0.5 {
228                return model["recommended_mode"]
229                    .as_str()
230                    .map(std::string::ToString::to_string);
231            }
232        }
233
234        None
235    }
236
237    /// Built-in defaults for common file types and sizes.
238    /// Ensures reasonable compression even without local history or cloud models.
239    /// Respects Kolmogorov-Gate: files with K>0.7 skip aggressive modes.
240    fn predict_from_defaults(sig: &FileSignature) -> Option<String> {
241        if sig.size_bucket == 0 {
242            return None;
243        }
244        if matches!(sig.ext.as_str(), "md" | "mdx" | "txt" | "rst") {
245            return None;
246        }
247
248        let mode = match (sig.ext.as_str(), sig.size_bucket) {
249            // Large code files: signatures only
250            (
251                "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "rb"
252                | "swift" | "kt" | "cs" | "vue" | "svelte" | "gd",
253                4..,
254            ) => "signatures",
255
256            // Code 2k-10k, SQL, lock, config/data: structured map
257            ("lock" | "json" | "yaml" | "yml" | "toml", _)
258            | (
259                "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "rb"
260                | "swift" | "kt" | "cs" | "vue" | "svelte" | "gd",
261                2 | 3,
262            )
263            | ("sql", 2..) => "map",
264
265            // CSS, XML/CSV, and large unknown files: aggressive
266            ("xml" | "csv", _) | ("css" | "scss" | "less" | "sass", 2..) | (_, 3..) => "aggressive",
267
268            _ => return None,
269        };
270        Some(mode.to_string())
271    }
272
273    /// Saves to the in-memory buffer and flushes to disk if the interval elapsed.
274    pub fn save(&self) {
275        let mut guard = PREDICTOR_BUFFER
276            .lock()
277            .unwrap_or_else(std::sync::PoisonError::into_inner);
278        let should_flush = match *guard {
279            Some((_, ref last_flush)) => last_flush.elapsed().as_secs() >= PREDICTOR_FLUSH_SECS,
280            None => true,
281        };
282        *guard = Some((Arc::new(self.clone()), Instant::now()));
283        if should_flush {
284            self.save_to_disk();
285        }
286    }
287
288    fn save_to_disk(&self) {
289        let Ok(dir) = crate::core::data_dir::lean_ctx_data_dir() else {
290            return;
291        };
292        let _ = std::fs::create_dir_all(&dir);
293        let path = dir.join(STATS_FILE);
294        if let Ok(json) = serde_json::to_string_pretty(self) {
295            let tmp = dir.join(".mode_stats.tmp");
296            if std::fs::write(&tmp, &json).is_ok() {
297                let _ = std::fs::rename(&tmp, &path);
298            }
299        }
300    }
301
302    /// Forces an immediate write of the buffered predictor state to disk.
303    pub fn flush() {
304        let guard = PREDICTOR_BUFFER
305            .lock()
306            .unwrap_or_else(std::sync::PoisonError::into_inner);
307        if let Some((ref predictor, _)) = *guard {
308            predictor.save_to_disk();
309        }
310    }
311
312    fn load_from_disk() -> Option<Self> {
313        let path = crate::core::data_dir::lean_ctx_data_dir()
314            .ok()?
315            .join(STATS_FILE);
316        let data = std::fs::read_to_string(path).ok()?;
317        serde_json::from_str(&data).ok()
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn file_signature_buckets() {
327        assert_eq!(FileSignature::from_path("main.rs", 100).size_bucket, 0);
328        assert_eq!(FileSignature::from_path("main.rs", 1000).size_bucket, 1);
329        assert_eq!(FileSignature::from_path("main.rs", 3000).size_bucket, 2);
330        assert_eq!(FileSignature::from_path("main.rs", 10000).size_bucket, 3);
331        assert_eq!(FileSignature::from_path("main.rs", 50000).size_bucket, 4);
332    }
333
334    #[test]
335    fn predict_returns_none_without_history() {
336        let predictor = ModePredictor::default();
337        let sig = FileSignature::from_path("test.zzz", 500);
338        assert!(predictor.predict_from_local(&sig).is_none());
339    }
340
341    #[test]
342    fn predict_returns_none_with_too_few_entries() {
343        let mut predictor = ModePredictor::default();
344        let sig = FileSignature::from_path("test.zzz", 500);
345        predictor.record(
346            sig.clone(),
347            ModeOutcome {
348                mode: "full".to_string(),
349                tokens_in: 100,
350                tokens_out: 100,
351                density: 0.5,
352            },
353        );
354        assert!(predictor.predict_from_local(&sig).is_none());
355    }
356
357    #[test]
358    fn predict_learns_best_mode() {
359        let mut predictor = ModePredictor::default();
360        let sig = FileSignature::from_path("big.rs", 5000);
361        for _ in 0..5 {
362            predictor.record(
363                sig.clone(),
364                ModeOutcome {
365                    mode: "full".to_string(),
366                    tokens_in: 5000,
367                    tokens_out: 5000,
368                    density: 0.3,
369                },
370            );
371            predictor.record(
372                sig.clone(),
373                ModeOutcome {
374                    mode: "map".to_string(),
375                    tokens_in: 5000,
376                    tokens_out: 800,
377                    density: 0.6,
378                },
379            );
380        }
381        let best = predictor.predict_best_mode(&sig);
382        assert_eq!(best, Some("map".to_string()));
383    }
384
385    #[test]
386    fn history_caps_at_100() {
387        let mut predictor = ModePredictor::default();
388        let sig = FileSignature::from_path("test.rs", 100);
389        for _ in 0..120 {
390            predictor.record(
391                sig.clone(),
392                ModeOutcome {
393                    mode: "full".to_string(),
394                    tokens_in: 100,
395                    tokens_out: 100,
396                    density: 0.5,
397                },
398            );
399        }
400        assert!(predictor.history.get(&sig).unwrap().len() <= 100);
401    }
402
403    #[test]
404    fn defaults_return_none_for_small_files() {
405        let sig = FileSignature::from_path("small.rs", 200);
406        assert!(ModePredictor::predict_from_defaults(&sig).is_none());
407    }
408
409    #[test]
410    fn defaults_recommend_map_for_medium_code() {
411        let sig = FileSignature::from_path("medium.rs", 3000);
412        assert_eq!(
413            ModePredictor::predict_from_defaults(&sig),
414            Some("map".to_string())
415        );
416    }
417
418    #[test]
419    fn defaults_recommend_map_for_json() {
420        let sig = FileSignature::from_path("config.json", 1000);
421        assert_eq!(
422            ModePredictor::predict_from_defaults(&sig),
423            Some("map".to_string())
424        );
425    }
426
427    #[test]
428    fn defaults_recommend_signatures_for_huge_code() {
429        let sig = FileSignature::from_path("huge.ts", 25000);
430        assert_eq!(
431            ModePredictor::predict_from_defaults(&sig),
432            Some("signatures".to_string())
433        );
434    }
435
436    #[test]
437    fn defaults_recommend_aggressive_for_large_unknown() {
438        let sig = FileSignature::from_path("data.xyz", 8000);
439        assert_eq!(
440            ModePredictor::predict_from_defaults(&sig),
441            Some("aggressive".to_string())
442        );
443    }
444
445    #[test]
446    fn defaults_never_compress_markdown() {
447        for tokens in [600, 3000, 8000, 25000] {
448            let sig = FileSignature::from_path("SKILL.md", tokens);
449            assert!(
450                ModePredictor::predict_from_defaults(&sig).is_none(),
451                "SKILL.md at {tokens} tokens should get full (None), not compressed"
452            );
453        }
454        let sig = FileSignature::from_path("AGENTS.md", 5000);
455        assert!(ModePredictor::predict_from_defaults(&sig).is_none());
456        let sig = FileSignature::from_path("README.md", 12000);
457        assert!(ModePredictor::predict_from_defaults(&sig).is_none());
458    }
459
460    #[test]
461    fn mode_outcome_efficiency() {
462        let o = ModeOutcome {
463            mode: "map".to_string(),
464            tokens_in: 1000,
465            tokens_out: 200,
466            density: 0.6,
467        };
468        assert!(o.efficiency() > 0.0);
469    }
470}
lean_ctx/core/mode_predictor.rs

lean_ctx/core/
mode_predictor.rs
