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