Skip to main content

lean_ctx/core/
context_radar.rs

1use std::path::{Path, PathBuf};
2
3use serde::{Deserialize, Serialize};
4
5/// ContextRadar aggregates all context sources into a single budget model.
6/// Data flows in from: hooks (JSONL), proxy introspector, rules scanner, session cache.
7pub struct ContextRadar {
8    pub events: Vec<RadarEvent>,
9    pub rules_tokens: RulesTokens,
10    pub window_size: usize,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct RadarEvent {
15    pub ts: u64,
16    pub event_type: String,
17    pub tokens: usize,
18    #[serde(default)]
19    pub tool_name: Option<String>,
20    #[serde(default)]
21    pub detail: Option<String>,
22    #[serde(default)]
23    pub content: Option<String>,
24    #[serde(default)]
25    pub model: Option<String>,
26    #[serde(default)]
27    pub conversation_id: Option<String>,
28}
29
30#[derive(Debug, Default, Clone)]
31pub struct RulesTokens {
32    pub files: Vec<(String, usize)>,
33    pub total: usize,
34}
35
36#[derive(Debug, Serialize)]
37pub struct BudgetBreakdown {
38    pub window_size: usize,
39    pub system_prompt_tokens: usize,
40    pub user_message_tokens: usize,
41    pub agent_response_tokens: usize,
42    pub lean_ctx_tool_tokens: usize,
43    pub other_mcp_tokens: usize,
44    pub native_read_tokens: usize,
45    pub shell_tokens: usize,
46    pub thinking_tokens: usize,
47    pub tracked_total: usize,
48    pub available: usize,
49    pub compaction_count: usize,
50    pub session_total_tokens: usize,
51    pub session_user_tokens: usize,
52    pub session_agent_tokens: usize,
53    pub session_lctx_tokens: usize,
54    pub session_mcp_tokens: usize,
55    pub session_native_tokens: usize,
56    pub session_shell_tokens: usize,
57    pub session_thinking_tokens: usize,
58    pub source: String,
59}
60
61impl ContextRadar {
62    pub fn new(window_size: usize) -> Self {
63        Self {
64            events: Vec::new(),
65            rules_tokens: RulesTokens::default(),
66            window_size,
67        }
68    }
69
70    pub fn load(data_dir: &Path, window_size: usize) -> Self {
71        let mut radar = Self::new(window_size);
72        radar.load_events(data_dir);
73        radar.scan_rules();
74        radar
75    }
76
77    fn load_events(&mut self, data_dir: &Path) {
78        let mut all: Vec<RadarEvent> = Vec::new();
79
80        let prev_path = data_dir.join("context_radar.prev.jsonl");
81        if let Ok(content) = std::fs::read_to_string(&prev_path) {
82            for line in content.lines() {
83                if let Ok(ev) = serde_json::from_str::<RadarEvent>(line) {
84                    all.push(ev);
85                }
86            }
87        }
88
89        let radar_path = data_dir.join("context_radar.jsonl");
90        if let Ok(content) = std::fs::read_to_string(&radar_path) {
91            for line in content.lines() {
92                if let Ok(ev) = serde_json::from_str::<RadarEvent>(line) {
93                    all.push(ev);
94                }
95            }
96        }
97
98        const MAX_EVENTS: usize = 50_000;
99        if all.len() > MAX_EVENTS {
100            self.events = all[all.len() - MAX_EVENTS..].to_vec();
101        } else {
102            self.events = all;
103        }
104    }
105
106    pub fn scan_rules(&mut self) {
107        let Some(home) = crate::core::home::resolve_home_dir() else {
108            return;
109        };
110
111        let cwd = std::env::current_dir().unwrap_or_default();
112        let mut files: Vec<(String, usize)> = Vec::new();
113
114        let paths_to_scan: Vec<PathBuf> = vec![
115            cwd.join(".cursorrules"),
116            cwd.join("AGENTS.md"),
117            cwd.join("CLAUDE.md"),
118            cwd.join("LEAN-CTX.md"),
119            home.join(".cursor").join("rules"),
120            home.join(".cursorrules"),
121            cwd.join(".cursor").join("rules"),
122        ];
123
124        for path in &paths_to_scan {
125            if path.is_file() {
126                if Self::is_rules_file(path) {
127                    if let Ok(content) = std::fs::read_to_string(path) {
128                        let tokens = content.len() / 4;
129                        if tokens > 0 {
130                            files.push((path.display().to_string(), tokens));
131                        }
132                    }
133                }
134            } else if path.is_dir() {
135                if let Ok(entries) = std::fs::read_dir(path) {
136                    for entry in entries.flatten() {
137                        let p = entry.path();
138                        if p.is_file() && Self::is_rules_file(&p) {
139                            if let Ok(content) = std::fs::read_to_string(&p) {
140                                let tokens = content.len() / 4;
141                                if tokens > 0 {
142                                    files.push((p.display().to_string(), tokens));
143                                }
144                            }
145                        }
146                    }
147                }
148            }
149        }
150
151        let total = files.iter().map(|(_, t)| *t).sum();
152        self.rules_tokens = RulesTokens { files, total };
153    }
154
155    fn is_rules_file(path: &Path) -> bool {
156        let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
157            return false;
158        };
159        if name.starts_with('.') && name != ".cursorrules" {
160            return false;
161        }
162        if name.contains(".bak") || name.contains(".tmp") || name.contains(".swp") {
163            return false;
164        }
165        if name == ".cursorrules"
166            || name == "AGENTS.md"
167            || name == "CLAUDE.md"
168            || name == "LEAN-CTX.md"
169        {
170            return true;
171        }
172        let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
173        matches!(ext, "md" | "mdc" | "markdown" | "txt")
174    }
175
176    pub fn budget_breakdown(&self) -> BudgetBreakdown {
177        let mut compaction_count = 0;
178        let mut last_compaction_idx: Option<usize> = None;
179
180        for (i, event) in self.events.iter().enumerate() {
181            if event.event_type == "compaction" {
182                compaction_count += 1;
183                last_compaction_idx = Some(i);
184            }
185        }
186
187        let current_window_start = last_compaction_idx.map_or(0, |i| i + 1);
188        let current_events = &self.events[current_window_start..];
189
190        let (mut user_cur, mut agent_cur, mut lctx_cur, mut mcp_cur) = (0, 0, 0, 0);
191        let (mut native_cur, mut shell_cur, mut thinking_cur) = (0, 0, 0);
192        let (mut user_all, mut agent_all, mut lctx_all, mut mcp_all) = (0, 0, 0, 0);
193        let (mut native_all, mut shell_all, mut thinking_all) = (0, 0, 0);
194
195        for event in &self.events {
196            Self::classify_event(
197                event,
198                &mut user_all,
199                &mut agent_all,
200                &mut lctx_all,
201                &mut mcp_all,
202                &mut native_all,
203                &mut shell_all,
204                &mut thinking_all,
205            );
206        }
207        for event in current_events {
208            Self::classify_event(
209                event,
210                &mut user_cur,
211                &mut agent_cur,
212                &mut lctx_cur,
213                &mut mcp_cur,
214                &mut native_cur,
215                &mut shell_cur,
216                &mut thinking_cur,
217            );
218        }
219
220        let system_prompt_tokens = self.rules_tokens.total;
221        let tracked_total = system_prompt_tokens
222            + user_cur
223            + agent_cur
224            + lctx_cur
225            + mcp_cur
226            + native_cur
227            + shell_cur;
228        let available = self.window_size.saturating_sub(tracked_total);
229
230        let session_total = system_prompt_tokens
231            + user_all
232            + agent_all
233            + lctx_all
234            + mcp_all
235            + native_all
236            + shell_all;
237
238        BudgetBreakdown {
239            window_size: self.window_size,
240            system_prompt_tokens,
241            user_message_tokens: user_cur,
242            agent_response_tokens: agent_cur,
243            lean_ctx_tool_tokens: lctx_cur,
244            other_mcp_tokens: mcp_cur,
245            native_read_tokens: native_cur,
246            shell_tokens: shell_cur,
247            thinking_tokens: thinking_cur,
248            tracked_total,
249            available,
250            compaction_count,
251            session_total_tokens: session_total,
252            session_user_tokens: user_all,
253            session_agent_tokens: agent_all,
254            session_lctx_tokens: lctx_all,
255            session_mcp_tokens: mcp_all,
256            session_native_tokens: native_all,
257            session_shell_tokens: shell_all,
258            session_thinking_tokens: thinking_all,
259            source: "hooks + rules-scan".to_string(),
260        }
261    }
262
263    fn classify_event(
264        event: &RadarEvent,
265        user: &mut usize,
266        agent: &mut usize,
267        lctx: &mut usize,
268        mcp: &mut usize,
269        native: &mut usize,
270        shell: &mut usize,
271        thinking: &mut usize,
272    ) {
273        match event.event_type.as_str() {
274            "user_message" => *user += event.tokens,
275            "agent_response" => *agent += event.tokens,
276            "mcp_call" => {
277                let is_leanctx = event
278                    .detail
279                    .as_deref()
280                    .is_some_and(|d| d.contains("lean-ctx"))
281                    || event
282                        .tool_name
283                        .as_deref()
284                        .is_some_and(|t| t.starts_with("ctx_"));
285                if is_leanctx {
286                    *lctx += event.tokens;
287                } else {
288                    *mcp += event.tokens;
289                }
290            }
291            "native_tool" | "file_read" => *native += event.tokens,
292            "shell" => *shell += event.tokens,
293            "thinking" => *thinking += event.tokens,
294            _ => {}
295        }
296    }
297
298    pub fn format_display(&self) -> String {
299        let b = self.budget_breakdown();
300        let pct = |tokens: usize| -> f64 {
301            if b.window_size == 0 {
302                0.0
303            } else {
304                (tokens as f64 / b.window_size as f64 * 100.0).min(100.0)
305            }
306        };
307        let bar = |tokens: usize| -> String {
308            let width = (pct(tokens) / 2.0).min(40.0) as usize;
309            "█".repeat(width)
310        };
311
312        let mut out = String::new();
313        out.push_str(&format!(
314            "CONTEXT RADAR — Current Window ({:.0}k)\n",
315            b.window_size as f64 / 1000.0
316        ));
317        if b.compaction_count > 0 {
318            out.push_str(&format!(
319                "  (after {} compaction(s) — showing current window only)\n",
320                b.compaction_count
321            ));
322        }
323        out.push_str(&format!(
324            "  System Prompt (est.): {:>8} tok {:>5.1}%  {}\n",
325            fmt_num(b.system_prompt_tokens),
326            pct(b.system_prompt_tokens),
327            bar(b.system_prompt_tokens)
328        ));
329        out.push_str(&format!(
330            "  User Messages:        {:>8} tok {:>5.1}%  {}\n",
331            fmt_num(b.user_message_tokens),
332            pct(b.user_message_tokens),
333            bar(b.user_message_tokens)
334        ));
335        out.push_str(&format!(
336            "  Agent Responses:      {:>8} tok {:>5.1}%  {}\n",
337            fmt_num(b.agent_response_tokens),
338            pct(b.agent_response_tokens),
339            bar(b.agent_response_tokens)
340        ));
341        out.push_str(&format!(
342            "  lean-ctx Tools:       {:>8} tok {:>5.1}%  {}\n",
343            fmt_num(b.lean_ctx_tool_tokens),
344            pct(b.lean_ctx_tool_tokens),
345            bar(b.lean_ctx_tool_tokens)
346        ));
347        out.push_str(&format!(
348            "  Other MCP:            {:>8} tok {:>5.1}%  {}\n",
349            fmt_num(b.other_mcp_tokens),
350            pct(b.other_mcp_tokens),
351            bar(b.other_mcp_tokens)
352        ));
353        out.push_str(&format!(
354            "  Native Reads:         {:>8} tok {:>5.1}%  {}\n",
355            fmt_num(b.native_read_tokens),
356            pct(b.native_read_tokens),
357            bar(b.native_read_tokens)
358        ));
359        out.push_str(&format!(
360            "  Shell Output:         {:>8} tok {:>5.1}%  {}\n",
361            fmt_num(b.shell_tokens),
362            pct(b.shell_tokens),
363            bar(b.shell_tokens)
364        ));
365        out.push_str("  ──────────────────────────────────────────\n");
366        out.push_str(&format!(
367            "  TRACKED:              {:>8} tok {:>5.1}%\n",
368            fmt_num(b.tracked_total),
369            pct(b.tracked_total)
370        ));
371        out.push_str(&format!(
372            "  Available:            {:>8} tok {:>5.1}%\n",
373            fmt_num(b.available),
374            pct(b.available)
375        ));
376        if b.thinking_tokens > 0 {
377            out.push_str(&format!(
378                "  Thinking (not in window): {:>5} tok\n",
379                fmt_num(b.thinking_tokens)
380            ));
381        }
382        if b.session_total_tokens > b.tracked_total {
383            out.push_str(&format!(
384                "\n  SESSION TOTAL:        {:>8} tok (across {} compaction(s))\n",
385                fmt_num(b.session_total_tokens),
386                b.compaction_count
387            ));
388        }
389        out.push_str(&format!("  Source: {}\n", b.source));
390        out
391    }
392}
393
394fn fmt_num(n: usize) -> String {
395    if n >= 1000 {
396        format!("{},{:03}", n / 1000, n % 1000)
397    } else {
398        n.to_string()
399    }
400}
401
402/// Default context window size based on client name.
403pub fn default_window_for_client(client: &str) -> usize {
404    if let Some((_model, window)) = crate::hook_handlers::load_detected_model() {
405        return window;
406    }
407    match client.to_lowercase().as_str() {
408        "gemini" => 1_000_000,
409        "windsurf" | "zed" | "copilot" => 128_000,
410        _ => 200_000,
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    #[test]
419    fn budget_breakdown_empty() {
420        let radar = ContextRadar::new(200_000);
421        let b = radar.budget_breakdown();
422        assert_eq!(b.window_size, 200_000);
423        assert_eq!(b.tracked_total, 0);
424        assert_eq!(b.available, 200_000);
425    }
426
427    fn ev(
428        ts: u64,
429        event_type: &str,
430        tokens: usize,
431        tool_name: Option<&str>,
432        detail: Option<&str>,
433    ) -> RadarEvent {
434        RadarEvent {
435            ts,
436            event_type: event_type.to_string(),
437            tokens,
438            tool_name: tool_name.map(String::from),
439            detail: detail.map(String::from),
440            content: None,
441            model: None,
442            conversation_id: None,
443        }
444    }
445
446    #[test]
447    fn budget_breakdown_with_events() {
448        let mut radar = ContextRadar::new(200_000);
449        radar.events.push(ev(1000, "user_message", 500, None, None));
450        radar
451            .events
452            .push(ev(1001, "agent_response", 2000, None, None));
453        radar
454            .events
455            .push(ev(1002, "shell", 300, None, Some("git status")));
456        let b = radar.budget_breakdown();
457        assert_eq!(b.user_message_tokens, 500);
458        assert_eq!(b.agent_response_tokens, 2000);
459        assert_eq!(b.shell_tokens, 300);
460        assert_eq!(b.tracked_total, 2800);
461        assert_eq!(b.available, 200_000 - 2800);
462    }
463
464    #[test]
465    fn budget_breakdown_resets_after_compaction() {
466        let mut radar = ContextRadar::new(100_000);
467        radar.events.push(ev(1, "user_message", 50_000, None, None));
468        radar.events.push(ev(2, "compaction", 0, None, None));
469        radar.events.push(ev(3, "user_message", 10_000, None, None));
470        let b = radar.budget_breakdown();
471        assert_eq!(
472            b.user_message_tokens, 10_000,
473            "only counts since compaction"
474        );
475        assert_eq!(b.available, 90_000);
476        assert_eq!(b.compaction_count, 1);
477        assert_eq!(b.session_user_tokens, 60_000, "session total includes all");
478    }
479
480    #[test]
481    fn format_display_not_empty() {
482        let radar = ContextRadar::new(200_000);
483        let display = radar.format_display();
484        assert!(display.contains("CONTEXT RADAR"));
485        assert!(display.contains("200k"));
486    }
487
488    #[test]
489    fn default_window_sizes() {
490        // If a detected model file exists on the system, default_window_for_client
491        // returns that model's window for all clients. Skip client-specific asserts
492        // in that case and only verify the function returns a reasonable value.
493        if crate::hook_handlers::load_detected_model().is_some() {
494            let w = default_window_for_client("cursor");
495            assert!(
496                (128_000..=2_000_000).contains(&w),
497                "window {w} out of range"
498            );
499        } else {
500            assert_eq!(default_window_for_client("cursor"), 200_000);
501            assert_eq!(default_window_for_client("gemini"), 1_000_000);
502            assert_eq!(default_window_for_client("windsurf"), 128_000);
503        }
504    }
505}