lean_ctx/core/
heatmap.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::path::PathBuf;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::Mutex;
6
7const HEATMAP_FLUSH_EVERY: usize = 25;
8const HEATMAP_MAX_ENTRIES: usize = 10_000;
9
10static HEATMAP_BUFFER: Mutex<Option<HeatMap>> = Mutex::new(None);
11static HEATMAP_CALLS: AtomicUsize = AtomicUsize::new(0);
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct HeatEntry {
15    pub path: String,
16    pub access_count: u32,
17    pub last_access: String,
18    pub total_tokens_saved: u64,
19    pub total_original_tokens: u64,
20    pub avg_compression_ratio: f32,
21    /// Per-agent access counts — the stigmergic pheromone field.  When multiple
22    /// agents access the same file, downstream consumers can identify shared
23    /// context (co-access patterns) and compute credit for useful preloads.
24    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
25    pub agent_accesses: HashMap<String, u32>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, Default)]
29pub struct HeatMap {
30    pub entries: HashMap<String, HeatEntry>,
31    #[serde(skip)]
32    dirty: bool,
33}
34
35impl HeatMap {
36    pub fn load() -> Self {
37        let mut guard = HEATMAP_BUFFER
38            .lock()
39            .unwrap_or_else(std::sync::PoisonError::into_inner);
40        if let Some(ref hm) = *guard {
41            return hm.clone();
42        }
43        let hm = load_from_disk();
44        *guard = Some(hm.clone());
45        hm
46    }
47
48    pub fn record_access(&mut self, file_path: &str, original_tokens: usize, saved_tokens: usize) {
49        self.record_access_with_agent(file_path, original_tokens, saved_tokens, None);
50    }
51
52    /// Record a file access with an optional agent identifier (stigmergic trace).
53    pub fn record_access_with_agent(
54        &mut self,
55        file_path: &str,
56        original_tokens: usize,
57        saved_tokens: usize,
58        agent_id: Option<&str>,
59    ) {
60        let now = chrono::Utc::now().to_rfc3339();
61        let entry = self
62            .entries
63            .entry(file_path.to_string())
64            .or_insert_with(|| HeatEntry {
65                path: file_path.to_string(),
66                access_count: 0,
67                last_access: now.clone(),
68                total_tokens_saved: 0,
69                total_original_tokens: 0,
70                avg_compression_ratio: 0.0,
71                agent_accesses: HashMap::new(),
72            });
73        entry.access_count += 1;
74        entry.last_access = now;
75        entry.total_tokens_saved += saved_tokens as u64;
76        entry.total_original_tokens += original_tokens as u64;
77        if entry.total_original_tokens > 0 {
78            entry.avg_compression_ratio = 1.0
79                - (entry.total_original_tokens - entry.total_tokens_saved) as f32
80                    / entry.total_original_tokens as f32;
81        }
82        if let Some(aid) = agent_id {
83            if !aid.is_empty() {
84                *entry.agent_accesses.entry(aid.to_string()).or_insert(0) += 1;
85            }
86        }
87        self.dirty = true;
88    }
89
90    pub fn save(&self) -> std::io::Result<()> {
91        if !self.dirty && !self.entries.is_empty() {
92            return Ok(());
93        }
94        save_to_disk(self)?;
95        let mut guard = HEATMAP_BUFFER
96            .lock()
97            .unwrap_or_else(std::sync::PoisonError::into_inner);
98        *guard = Some(self.clone());
99        Ok(())
100    }
101
102    pub fn top_files(&self, limit: usize) -> Vec<&HeatEntry> {
103        let mut sorted: Vec<&HeatEntry> = self.entries.values().collect();
104        sorted.sort_by_key(|x| std::cmp::Reverse(x.access_count));
105        sorted.truncate(limit);
106        sorted
107    }
108
109    /// Compute stigmergic context credit: which agents' file-access traces
110    /// benefited other agents? An agent A gets credit for a file F when A
111    /// accessed F before (or alongside) agent B, because A's trace effectively
112    /// pointed B to useful context. The credit for each (agent_A, file) pair is
113    /// proportional to how many *other* agents also accessed that file.
114    /// Returns `Vec<(agent_id, total_credit)>` sorted descending.
115    pub fn context_credit(&self) -> Vec<(String, f64)> {
116        let mut credit: HashMap<String, f64> = HashMap::new();
117        for entry in self.entries.values() {
118            let n_agents = entry.agent_accesses.len();
119            if n_agents < 2 {
120                continue;
121            }
122            // Shapley-inspired: each agent that accessed a shared file gets
123            // credit = (n_other_agents) / n_agents. The more agents a file
124            // served, the more each contributor is credited.
125            let share = (n_agents - 1) as f64 / n_agents as f64;
126            for agent in entry.agent_accesses.keys() {
127                *credit.entry(agent.clone()).or_insert(0.0) += share;
128            }
129        }
130        let mut sorted: Vec<(String, f64)> = credit.into_iter().collect();
131        sorted.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
132        sorted
133    }
134
135    pub fn directory_summary(&self) -> Vec<(String, u32, u64)> {
136        let mut dirs: HashMap<String, (u32, u64)> = HashMap::new();
137        for entry in self.entries.values() {
138            let dir = std::path::Path::new(&entry.path)
139                .parent()
140                .map_or_else(|| ".".to_string(), |p| p.to_string_lossy().to_string());
141            let stat = dirs.entry(dir).or_insert((0, 0));
142            stat.0 += entry.access_count;
143            stat.1 += entry.total_tokens_saved;
144        }
145        let mut result: Vec<(String, u32, u64)> = dirs
146            .into_iter()
147            .map(|(dir, (count, saved))| (dir, count, saved))
148            .collect();
149        result.sort_by_key(|x| std::cmp::Reverse(x.1));
150        result
151    }
152
153    pub fn cold_files(&self, all_files: &[String], limit: usize) -> Vec<String> {
154        let hot: std::collections::HashSet<&str> = self
155            .entries
156            .keys()
157            .map(std::string::String::as_str)
158            .collect();
159        let mut cold: Vec<String> = all_files
160            .iter()
161            .filter(|f| !hot.contains(f.as_str()))
162            .cloned()
163            .collect();
164        cold.truncate(limit);
165        cold
166    }
167
168    fn storage_path() -> PathBuf {
169        crate::core::data_dir::lean_ctx_data_dir()
170            .unwrap_or_else(|_| PathBuf::from("."))
171            .join("heatmap.json")
172    }
173}
174
175fn load_from_disk() -> HeatMap {
176    let path = HeatMap::storage_path();
177    match std::fs::read_to_string(&path) {
178        Ok(content) => serde_json::from_str(&content).unwrap_or_default(),
179        Err(_) => HeatMap::default(),
180    }
181}
182
183fn save_to_disk(hm: &HeatMap) -> std::io::Result<()> {
184    let path = HeatMap::storage_path();
185    if let Some(parent) = path.parent() {
186        std::fs::create_dir_all(parent)?;
187    }
188    let json = serde_json::to_string_pretty(hm)?;
189    let tmp = path.with_extension("json.tmp");
190    std::fs::write(&tmp, &json)?;
191    std::fs::rename(&tmp, &path)
192}
193
194pub fn record_file_access(file_path: &str, original_tokens: usize, saved_tokens: usize) {
195    // Attribute every read to the current agent identity so the per-agent
196    // pheromone field (stigmergic trace) is populated in production, not just
197    // when callers explicitly pass an id.
198    let agent = crate::core::agent_identity::current_agent_id();
199    record_file_access_with_agent(file_path, original_tokens, saved_tokens, Some(agent));
200}
201
202/// Like [`record_file_access`] but attaches an agent identifier so the heatmap
203/// builds a per-agent pheromone field (stigmergic trace for multi-agent routing).
204pub fn record_file_access_with_agent(
205    file_path: &str,
206    original_tokens: usize,
207    saved_tokens: usize,
208    agent_id: Option<&str>,
209) {
210    // Universal per-read chokepoint (CLI via tool_lifecycle, MCP via ctx_read/ctx_multi_read):
211    // also append one auditable savings event. Best-effort; never blocks/breaks the read.
212    crate::core::savings_ledger::record_read_event(original_tokens, saved_tokens);
213
214    let file_path = std::fs::canonicalize(file_path).map_or_else(
215        |_| file_path.to_string(),
216        |p| p.to_string_lossy().into_owned(),
217    );
218    let file_path = file_path.as_str();
219
220    let mut guard = HEATMAP_BUFFER
221        .lock()
222        .unwrap_or_else(std::sync::PoisonError::into_inner);
223    let hm = guard.get_or_insert_with(load_from_disk);
224    hm.record_access_with_agent(file_path, original_tokens, saved_tokens, agent_id);
225
226    // Enforce bounded retention.
227    if hm.entries.len() > HEATMAP_MAX_ENTRIES {
228        let mut items: Vec<(String, u32)> = hm
229            .entries
230            .values()
231            .map(|e| (e.path.clone(), e.access_count))
232            .collect();
233        items.sort_by_key(|x| x.1);
234        let drop_n = hm.entries.len().saturating_sub(HEATMAP_MAX_ENTRIES);
235        for (path, _) in items.into_iter().take(drop_n) {
236            hm.entries.remove(&path);
237        }
238    }
239
240    let n = HEATMAP_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
241    if n.is_multiple_of(HEATMAP_FLUSH_EVERY) && save_to_disk(hm).is_ok() {
242        hm.dirty = false;
243    }
244}
245
246pub fn flush() {
247    let guard = HEATMAP_BUFFER
248        .lock()
249        .unwrap_or_else(std::sync::PoisonError::into_inner);
250    if let Some(ref hm) = *guard {
251        if hm.dirty {
252            let _ = save_to_disk(hm);
253        }
254    }
255}
256
257/// Cheap read-only lookup against the in-process heatmap buffer:
258/// `(access_count, avg_compression_ratio)` for a file, if tracked.
259/// Paths are canonicalized the same way `record_file_access` stores them.
260pub fn entry_stats(file_path: &str) -> Option<(u32, f32)> {
261    let canonical = std::fs::canonicalize(file_path).map_or_else(
262        |_| file_path.to_string(),
263        |p| p.to_string_lossy().into_owned(),
264    );
265    let mut guard = HEATMAP_BUFFER
266        .lock()
267        .unwrap_or_else(std::sync::PoisonError::into_inner);
268    let hm = guard.get_or_insert_with(load_from_disk);
269    hm.entries
270        .get(&canonical)
271        .map(|e| (e.access_count, e.avg_compression_ratio))
272}
273
274pub fn reset() {
275    let mut guard = HEATMAP_BUFFER
276        .lock()
277        .unwrap_or_else(std::sync::PoisonError::into_inner);
278    *guard = Some(HeatMap::default());
279    if let Some(hm) = guard.as_ref() {
280        let _ = save_to_disk(hm);
281    }
282}
283
284pub fn format_heatmap_status(heatmap: &HeatMap, limit: usize) -> String {
285    let top = heatmap.top_files(limit);
286    if top.is_empty() {
287        return "No file access data recorded yet.".to_string();
288    }
289    let mut lines = vec![format!(
290        "File Access Heat Map ({} tracked files):",
291        heatmap.entries.len()
292    )];
293    lines.push(String::new());
294    for (i, entry) in top.iter().enumerate() {
295        let short = short_path(&entry.path);
296        let heat = heat_indicator(entry.access_count);
297        lines.push(format!(
298            "  {heat} #{} {} — {} accesses, {:.0}% compression, {} tok saved",
299            i + 1,
300            short,
301            entry.access_count,
302            entry.avg_compression_ratio * 100.0,
303            entry.total_tokens_saved
304        ));
305    }
306    lines.join("\n")
307}
308
309pub fn format_directory_summary(heatmap: &HeatMap) -> String {
310    let dirs = heatmap.directory_summary();
311    if dirs.is_empty() {
312        return "No directory data.".to_string();
313    }
314    let mut lines = vec!["Directory Heat Map:".to_string(), String::new()];
315    for (dir, count, saved) in dirs.iter().take(15) {
316        let heat = heat_indicator(*count);
317        lines.push(format!(
318            "  {heat} {dir}/ — {count} accesses, {saved} tok saved"
319        ));
320    }
321    lines.join("\n")
322}
323
324fn heat_indicator(count: u32) -> &'static str {
325    match count {
326        0 => "  ",
327        1..=3 => "▁▁",
328        4..=8 => "▃▃",
329        9..=15 => "▅▅",
330        16..=30 => "▇▇",
331        _ => "██",
332    }
333}
334
335fn short_path(path: &str) -> &str {
336    let parts: Vec<&str> = path.rsplitn(3, '/').collect();
337    if parts.len() >= 2 {
338        let start = path.len() - parts[0].len() - parts[1].len() - 1;
339        &path[start..]
340    } else {
341        path
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    #[test]
350    fn record_and_query() {
351        let mut hm = HeatMap::default();
352        hm.record_access("src/main.rs", 100, 80);
353        hm.record_access("src/main.rs", 100, 90);
354        hm.record_access("src/lib.rs", 200, 50);
355
356        assert_eq!(hm.entries.len(), 2);
357        assert_eq!(hm.entries["src/main.rs"].access_count, 2);
358        assert_eq!(hm.entries["src/lib.rs"].total_tokens_saved, 50);
359    }
360
361    #[test]
362    fn top_files_sorted() {
363        let mut hm = HeatMap::default();
364        hm.record_access("a.rs", 100, 50);
365        hm.record_access("b.rs", 100, 50);
366        hm.record_access("b.rs", 100, 50);
367        hm.record_access("c.rs", 100, 50);
368        hm.record_access("c.rs", 100, 50);
369        hm.record_access("c.rs", 100, 50);
370
371        let top = hm.top_files(2);
372        assert_eq!(top.len(), 2);
373        assert_eq!(top[0].path, "c.rs");
374        assert_eq!(top[1].path, "b.rs");
375    }
376
377    #[test]
378    fn directory_summary_works() {
379        let mut hm = HeatMap::default();
380        hm.record_access("src/a.rs", 100, 50);
381        hm.record_access("src/b.rs", 100, 50);
382        hm.record_access("tests/t.rs", 200, 100);
383
384        let dirs = hm.directory_summary();
385        assert!(dirs.len() >= 2);
386    }
387
388    #[test]
389    fn cold_files_detection() {
390        let mut hm = HeatMap::default();
391        hm.record_access("src/a.rs", 100, 50);
392
393        let all = vec![
394            "src/a.rs".to_string(),
395            "src/b.rs".to_string(),
396            "src/c.rs".to_string(),
397        ];
398        let cold = hm.cold_files(&all, 10);
399        assert_eq!(cold.len(), 2);
400        assert!(cold.contains(&"src/b.rs".to_string()));
401    }
402
403    #[test]
404    fn heat_indicators() {
405        assert_eq!(heat_indicator(0), "  ");
406        assert_eq!(heat_indicator(1), "▁▁");
407        assert_eq!(heat_indicator(10), "▅▅");
408        assert_eq!(heat_indicator(50), "██");
409    }
410
411    #[test]
412    fn compression_ratio() {
413        let mut hm = HeatMap::default();
414        hm.record_access("a.rs", 1000, 800);
415        let entry = &hm.entries["a.rs"];
416        assert!((entry.avg_compression_ratio - 0.8).abs() < 0.01);
417    }
418
419    #[test]
420    fn agent_scoped_access_and_context_credit() {
421        let mut hm = HeatMap::default();
422        hm.record_access_with_agent("shared.rs", 100, 50, Some("agent-a"));
423        hm.record_access_with_agent("shared.rs", 100, 60, Some("agent-b"));
424        hm.record_access_with_agent("only-a.rs", 100, 70, Some("agent-a"));
425
426        let entry = &hm.entries["shared.rs"];
427        assert_eq!(entry.agent_accesses.len(), 2);
428        assert_eq!(entry.agent_accesses["agent-a"], 1);
429        assert_eq!(entry.agent_accesses["agent-b"], 1);
430
431        let credit = hm.context_credit();
432        assert!(!credit.is_empty());
433        // Both agents get credit for the shared file; only-a.rs contributes
434        // no credit (single-agent access).
435        let a_credit = credit.iter().find(|(id, _)| id == "agent-a").unwrap().1;
436        let b_credit = credit.iter().find(|(id, _)| id == "agent-b").unwrap().1;
437        assert!(a_credit > 0.0);
438        assert!((a_credit - b_credit).abs() < 1e-9);
439    }
440}
lean_ctx/core/heatmap.rs

lean_ctx/core/
heatmap.rs
