lean_ctx/core/
grammar_usage.rs

1//! Persistent per-extension signature-backend telemetry (GH #690 Phase 2).
2//!
3//! The tiering cut ("which of the ~27 static tree-sitter grammars are
4//! actually used, and how often?") needs real usage history, but the only
5//! existing signal — `signatures::signature_backend_stats()` — is a pair of
6//! process-lifetime `AtomicU64`s with no language dimension: it resets on
7//! every restart and cannot say *which* grammar earned its binary bytes.
8//!
9//! This store records, per file extension, how often the tree-sitter path vs
10//! the regex fallback produced a signature set, persisted across sessions
11//! (mirroring `path_mode_memory`'s store shape: load-once, atomic tmp+rename
12//! writes, periodic flush via `tool_lifecycle::flush_all`).
13//!
14//! Storage: `<cache_dir>/grammar_usage.json`. Aggregate counters only — no
15//! paths, no file names, nothing project-identifying — so the store is safe
16//! to inspect and never a privacy concern.
17
18use std::collections::HashMap;
19use std::path::PathBuf;
20use std::sync::atomic::{AtomicUsize, Ordering};
21use std::sync::{Mutex, OnceLock};
22
23use serde::{Deserialize, Serialize};
24
25const STORE_FILE: &str = "grammar_usage.json";
26/// Extensions unseen for this long are dropped on load — the project mix
27/// changed and stale rows would distort a tiering decision.
28const DECAY_SECS: u64 = 180 * 24 * 3600;
29/// Hard cap; least-recently-used extensions are evicted first. Real projects
30/// see a few dozen extensions, so the cap only guards against pathological
31/// callers feeding junk "extensions".
32const MAX_EXTENSIONS: usize = 300;
33const FLUSH_EVERY: usize = 50;
34
35static STORE: OnceLock<Mutex<GrammarUsage>> = OnceLock::new();
36static RECORD_CALLS: AtomicUsize = AtomicUsize::new(0);
37
38#[derive(Debug, Clone, Serialize, Deserialize, Default)]
39pub struct ExtUsage {
40    pub tree_sitter_hits: u64,
41    pub regex_hits: u64,
42    pub last_used_unix: u64,
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize, Default)]
46pub struct GrammarUsage {
47    pub extensions: HashMap<String, ExtUsage>,
48    #[serde(skip)]
49    dirty: bool,
50}
51
52impl GrammarUsage {
53    fn load_from_disk() -> Self {
54        let Ok(raw) = std::fs::read_to_string(store_path()) else {
55            return Self::default();
56        };
57        let mut store: Self = serde_json::from_str(&raw).unwrap_or_default();
58        store.decay(now_unix());
59        store
60    }
61
62    fn decay(&mut self, now: u64) {
63        let before = self.extensions.len();
64        self.extensions
65            .retain(|_, u| now.saturating_sub(u.last_used_unix) <= DECAY_SECS);
66        if self.extensions.len() != before {
67            self.dirty = true;
68        }
69    }
70
71    fn evict_to_cap(&mut self) {
72        if self.extensions.len() <= MAX_EXTENSIONS {
73            return;
74        }
75        let mut items: Vec<(String, u64)> = self
76            .extensions
77            .iter()
78            .map(|(e, u)| (e.clone(), u.last_used_unix))
79            .collect();
80        items.sort_by_key(|(_, ts)| *ts);
81        let drop_n = self.extensions.len() - MAX_EXTENSIONS;
82        for (ext, _) in items.into_iter().take(drop_n) {
83            self.extensions.remove(&ext);
84        }
85        self.dirty = true;
86    }
87
88    pub fn record(&mut self, ext: &str, tree_sitter: bool, now: u64) {
89        let entry = self.extensions.entry(normalize_ext(ext)).or_default();
90        if tree_sitter {
91            entry.tree_sitter_hits = entry.tree_sitter_hits.saturating_add(1);
92        } else {
93            entry.regex_hits = entry.regex_hits.saturating_add(1);
94        }
95        entry.last_used_unix = now;
96        self.dirty = true;
97        self.evict_to_cap();
98    }
99
100    /// Extensions sorted by total hits (desc), then name (asc) for a stable
101    /// order — the input to any tiering decision.
102    pub fn ranked(&self) -> Vec<(String, ExtUsage)> {
103        let mut rows: Vec<(String, ExtUsage)> = self
104            .extensions
105            .iter()
106            .map(|(e, u)| (e.clone(), u.clone()))
107            .collect();
108        rows.sort_by(|(ea, ua), (eb, ub)| {
109            let ta = ua.tree_sitter_hits + ua.regex_hits;
110            let tb = ub.tree_sitter_hits + ub.regex_hits;
111            tb.cmp(&ta).then_with(|| ea.cmp(eb))
112        });
113        rows
114    }
115
116    pub fn save(&self) -> std::io::Result<()> {
117        let path = store_path();
118        if let Some(parent) = path.parent() {
119            std::fs::create_dir_all(parent)?;
120        }
121        let json = serde_json::to_string(self)?;
122        let tmp = path.with_extension("tmp");
123        std::fs::write(&tmp, json)?;
124        std::fs::rename(&tmp, &path)
125    }
126}
127
128/// Extensions arrive from callers in mixed shapes (`rs`, `.rs`, `RS`); the
129/// store keys must not fragment across them.
130fn normalize_ext(ext: &str) -> String {
131    ext.trim_start_matches('.').to_ascii_lowercase()
132}
133
134fn store_path() -> PathBuf {
135    crate::core::paths::cache_dir()
136        .unwrap_or_else(|_| PathBuf::from("."))
137        .join(STORE_FILE)
138}
139
140fn now_unix() -> u64 {
141    std::time::SystemTime::now()
142        .duration_since(std::time::UNIX_EPOCH)
143        .map_or(0, |d| d.as_secs())
144}
145
146fn global() -> &'static Mutex<GrammarUsage> {
147    STORE.get_or_init(|| Mutex::new(GrammarUsage::load_from_disk()))
148}
149
150/// Process-global: record one signature extraction for `ext` via the given
151/// backend. Called from the hot `extract_signatures` path — no I/O here;
152/// disk writes happen every `FLUSH_EVERY` records or on [`flush`].
153pub fn record(ext: &str, tree_sitter: bool) {
154    if ext.is_empty() {
155        return;
156    }
157    let Ok(mut store) = global().lock() else {
158        return;
159    };
160    store.record(ext, tree_sitter, now_unix());
161    let n = RECORD_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
162    if n.is_multiple_of(FLUSH_EVERY) && store.dirty && store.save().is_ok() {
163        store.dirty = false;
164    }
165}
166
167/// Flush pending counts (wired into `tool_lifecycle::flush_all`).
168pub fn flush() {
169    if let Ok(store) = global().lock()
170        && store.dirty
171    {
172        let _ = store.save();
173    }
174}
175
176/// Ranked per-extension usage, read straight from disk so a separate process
177/// (dashboard, `doctor`) sees what the MCP/CLI processes persisted.
178pub fn disk_ranked() -> Vec<(String, ExtUsage)> {
179    GrammarUsage::load_from_disk().ranked()
180}
181
182/// Ranked per-extension usage from the live in-process store (includes
183/// not-yet-flushed records) — what `ctx_metrics` shows.
184pub fn live_ranked() -> Vec<(String, ExtUsage)> {
185    global().lock().map(|s| s.ranked()).unwrap_or_default()
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn record_splits_backends_and_normalizes_ext() {
194        let mut u = GrammarUsage::default();
195        u.record(".RS", true, 100);
196        u.record("rs", true, 101);
197        u.record("rs", false, 102);
198        assert_eq!(u.extensions.len(), 1, "'.RS' and 'rs' must share one row");
199        let row = u.extensions.get("rs").unwrap();
200        assert_eq!(row.tree_sitter_hits, 2);
201        assert_eq!(row.regex_hits, 1);
202        assert_eq!(row.last_used_unix, 102);
203    }
204
205    #[test]
206    fn ranked_orders_by_total_then_name() {
207        let mut u = GrammarUsage::default();
208        u.record("py", true, 100);
209        u.record("py", true, 101);
210        u.record("rs", true, 102);
211        u.record("go", true, 103); // same total as rs -> alphabetical
212        let ranked = u.ranked();
213        assert_eq!(ranked[0].0, "py");
214        assert_eq!(ranked[1].0, "go");
215        assert_eq!(ranked[2].0, "rs");
216    }
217
218    #[test]
219    fn decay_drops_stale_extensions() {
220        let mut u = GrammarUsage::default();
221        u.record("old", true, 1000);
222        u.record("fresh", true, 5000);
223        u.decay(5000 + DECAY_SECS - 10);
224        assert!(!u.extensions.contains_key("old"));
225        assert!(u.extensions.contains_key("fresh"));
226    }
227
228    #[test]
229    fn eviction_keeps_most_recently_used() {
230        let mut u = GrammarUsage::default();
231        for i in 0..(MAX_EXTENSIONS + 10) {
232            u.record(&format!("e{i}"), true, 1000 + i as u64);
233        }
234        assert_eq!(u.extensions.len(), MAX_EXTENSIONS);
235        assert!(!u.extensions.contains_key("e0"), "LRU evicted");
236        let newest = format!("e{}", MAX_EXTENSIONS + 9);
237        assert!(u.extensions.contains_key(&newest));
238    }
239
240    #[test]
241    fn roundtrip_serialization() {
242        let mut u = GrammarUsage::default();
243        u.record("ts", true, 42);
244        let json = serde_json::to_string(&u).unwrap();
245        let back: GrammarUsage = serde_json::from_str(&json).unwrap();
246        assert_eq!(back.extensions.get("ts").unwrap().tree_sitter_hits, 1);
247        assert_eq!(back.extensions.get("ts").unwrap().last_used_unix, 42);
248    }
249}
lean_ctx/core/grammar_usage.rs

lean_ctx/core/
grammar_usage.rs
