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