Skip to main content

lean_ctx/tools/
mod.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2use std::sync::Arc;
3use std::time::Instant;
4use tokio::sync::RwLock;
5
6use crate::core::cache::SessionCache;
7use crate::core::session::SessionState;
8
9pub mod autonomy;
10pub mod ctx_agent;
11pub mod ctx_analyze;
12pub mod ctx_benchmark;
13pub mod ctx_compress;
14pub mod ctx_context;
15pub mod ctx_dedup;
16pub mod ctx_delta;
17pub mod ctx_discover;
18pub mod ctx_edit;
19pub mod ctx_fill;
20pub mod ctx_graph;
21pub mod ctx_intent;
22pub mod ctx_knowledge;
23pub mod ctx_metrics;
24pub mod ctx_multi_read;
25pub mod ctx_overview;
26pub mod ctx_preload;
27pub mod ctx_read;
28pub mod ctx_response;
29pub mod ctx_search;
30pub mod ctx_semantic_search;
31pub mod ctx_session;
32pub mod ctx_shell;
33pub mod ctx_smart_read;
34pub mod ctx_tree;
35pub mod ctx_wrapped;
36
37const DEFAULT_CACHE_TTL_SECS: u64 = 300;
38
39struct CepComputedStats {
40    cep_score: u32,
41    cache_util: u32,
42    mode_diversity: u32,
43    compression_rate: u32,
44    total_original: u64,
45    total_compressed: u64,
46    total_saved: u64,
47    mode_counts: std::collections::HashMap<String, u64>,
48    complexity: String,
49    cache_hits: u64,
50    total_reads: u64,
51    tool_call_count: u64,
52}
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub enum CrpMode {
56    Off,
57    Compact,
58    Tdd,
59}
60
61impl CrpMode {
62    pub fn from_env() -> Self {
63        match std::env::var("LEAN_CTX_CRP_MODE")
64            .unwrap_or_default()
65            .to_lowercase()
66            .as_str()
67        {
68            "off" => Self::Off,
69            "compact" => Self::Compact,
70            _ => Self::Tdd,
71        }
72    }
73
74    pub fn is_tdd(&self) -> bool {
75        *self == Self::Tdd
76    }
77}
78
79pub type SharedCache = Arc<RwLock<SessionCache>>;
80
81#[derive(Clone)]
82pub struct LeanCtxServer {
83    pub cache: SharedCache,
84    pub session: Arc<RwLock<SessionState>>,
85    pub tool_calls: Arc<RwLock<Vec<ToolCallRecord>>>,
86    pub call_count: Arc<AtomicUsize>,
87    pub checkpoint_interval: usize,
88    pub cache_ttl_secs: u64,
89    pub last_call: Arc<RwLock<Instant>>,
90    pub crp_mode: CrpMode,
91    pub agent_id: Arc<RwLock<Option<String>>>,
92    pub client_name: Arc<RwLock<String>>,
93    pub autonomy: Arc<autonomy::AutonomyState>,
94}
95
96#[derive(Clone, Debug)]
97pub struct ToolCallRecord {
98    pub tool: String,
99    pub original_tokens: usize,
100    pub saved_tokens: usize,
101    pub mode: Option<String>,
102    pub duration_ms: u64,
103    pub timestamp: String,
104}
105
106impl Default for LeanCtxServer {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112impl LeanCtxServer {
113    pub fn new() -> Self {
114        let config = crate::core::config::Config::load();
115
116        let interval = std::env::var("LEAN_CTX_CHECKPOINT_INTERVAL")
117            .ok()
118            .and_then(|v| v.parse().ok())
119            .unwrap_or(config.checkpoint_interval as usize);
120
121        let ttl = std::env::var("LEAN_CTX_CACHE_TTL")
122            .ok()
123            .and_then(|v| v.parse().ok())
124            .unwrap_or(DEFAULT_CACHE_TTL_SECS);
125
126        let crp_mode = CrpMode::from_env();
127
128        let session = SessionState::load_latest().unwrap_or_default();
129
130        Self {
131            cache: Arc::new(RwLock::new(SessionCache::new())),
132            session: Arc::new(RwLock::new(session)),
133            tool_calls: Arc::new(RwLock::new(Vec::new())),
134            call_count: Arc::new(AtomicUsize::new(0)),
135            checkpoint_interval: interval,
136            cache_ttl_secs: ttl,
137            last_call: Arc::new(RwLock::new(Instant::now())),
138            crp_mode,
139            agent_id: Arc::new(RwLock::new(None)),
140            client_name: Arc::new(RwLock::new(String::new())),
141            autonomy: Arc::new(autonomy::AutonomyState::new()),
142        }
143    }
144
145    pub async fn check_idle_expiry(&self) {
146        if self.cache_ttl_secs == 0 {
147            return;
148        }
149        let last = *self.last_call.read().await;
150        if last.elapsed().as_secs() >= self.cache_ttl_secs {
151            {
152                let mut session = self.session.write().await;
153                let _ = session.save();
154            }
155            let mut cache = self.cache.write().await;
156            let count = cache.clear();
157            if count > 0 {
158                tracing::info!(
159                    "Cache auto-cleared after {}s idle ({count} file(s))",
160                    self.cache_ttl_secs
161                );
162            }
163        }
164        *self.last_call.write().await = Instant::now();
165    }
166
167    pub async fn record_call(
168        &self,
169        tool: &str,
170        original: usize,
171        saved: usize,
172        mode: Option<String>,
173    ) {
174        self.record_call_with_timing(tool, original, saved, mode, 0)
175            .await;
176    }
177
178    pub async fn record_call_with_timing(
179        &self,
180        tool: &str,
181        original: usize,
182        saved: usize,
183        mode: Option<String>,
184        duration_ms: u64,
185    ) {
186        let ts = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
187        let mut calls = self.tool_calls.write().await;
188        calls.push(ToolCallRecord {
189            tool: tool.to_string(),
190            original_tokens: original,
191            saved_tokens: saved,
192            mode: mode.clone(),
193            duration_ms,
194            timestamp: ts.clone(),
195        });
196
197        if duration_ms > 0 {
198            Self::append_tool_call_log(tool, duration_ms, original, saved, mode.as_deref(), &ts);
199        }
200
201        let output_tokens = original.saturating_sub(saved);
202        crate::core::stats::record(tool, original, output_tokens);
203
204        let mut session = self.session.write().await;
205        session.record_tool_call(saved as u64, original as u64);
206        if tool == "ctx_shell" {
207            session.record_command();
208        }
209        if session.should_save() {
210            let _ = session.save();
211        }
212        drop(calls);
213        drop(session);
214
215        self.write_mcp_live_stats().await;
216    }
217
218    pub async fn is_prompt_cache_stale(&self) -> bool {
219        let last = *self.last_call.read().await;
220        last.elapsed().as_secs() > 3600
221    }
222
223    pub fn upgrade_mode_if_stale(mode: &str, stale: bool) -> &str {
224        if !stale {
225            return mode;
226        }
227        match mode {
228            "full" => "full",
229            "map" => "signatures",
230            m => m,
231        }
232    }
233
234    pub fn increment_and_check(&self) -> bool {
235        let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1;
236        self.checkpoint_interval > 0 && count.is_multiple_of(self.checkpoint_interval)
237    }
238
239    pub async fn auto_checkpoint(&self) -> Option<String> {
240        let cache = self.cache.read().await;
241        if cache.get_all_entries().is_empty() {
242            return None;
243        }
244        let complexity = crate::core::adaptive::classify_from_context(&cache);
245        let checkpoint = ctx_compress::handle(&cache, true, self.crp_mode);
246        drop(cache);
247
248        let mut session = self.session.write().await;
249        let _ = session.save();
250        let session_summary = session.format_compact();
251        let has_insights = !session.findings.is_empty() || !session.decisions.is_empty();
252        let project_root = session.project_root.clone();
253        drop(session);
254
255        if has_insights {
256            if let Some(root) = project_root {
257                std::thread::spawn(move || {
258                    auto_consolidate_knowledge(&root);
259                });
260            }
261        }
262
263        self.record_call("ctx_compress", 0, 0, Some("auto".to_string()))
264            .await;
265
266        self.record_cep_snapshot().await;
267
268        Some(format!(
269            "{checkpoint}\n\n--- SESSION STATE ---\n{session_summary}\n\n{}",
270            complexity.instruction_suffix()
271        ))
272    }
273
274    pub fn append_tool_call_log(
275        tool: &str,
276        duration_ms: u64,
277        original: usize,
278        saved: usize,
279        mode: Option<&str>,
280        timestamp: &str,
281    ) {
282        const MAX_LOG_LINES: usize = 50;
283        if let Some(dir) = dirs::home_dir().map(|h| h.join(".lean-ctx")) {
284            let log_path = dir.join("tool-calls.log");
285            let mode_str = mode.unwrap_or("-");
286            let slow = if duration_ms > 5000 { " **SLOW**" } else { "" };
287            let line = format!(
288                "{timestamp}\t{tool}\t{duration_ms}ms\torig={original}\tsaved={saved}\tmode={mode_str}{slow}\n"
289            );
290
291            let mut lines: Vec<String> = std::fs::read_to_string(&log_path)
292                .unwrap_or_default()
293                .lines()
294                .map(|l| l.to_string())
295                .collect();
296
297            lines.push(line.trim_end().to_string());
298            if lines.len() > MAX_LOG_LINES {
299                lines.drain(0..lines.len() - MAX_LOG_LINES);
300            }
301
302            let _ = std::fs::write(&log_path, lines.join("\n") + "\n");
303        }
304    }
305
306    fn compute_cep_stats(
307        calls: &[ToolCallRecord],
308        stats: &crate::core::cache::CacheStats,
309        complexity: &crate::core::adaptive::TaskComplexity,
310    ) -> CepComputedStats {
311        let total_original: u64 = calls.iter().map(|c| c.original_tokens as u64).sum();
312        let total_saved: u64 = calls.iter().map(|c| c.saved_tokens as u64).sum();
313        let total_compressed = total_original.saturating_sub(total_saved);
314        let compression_rate = if total_original > 0 {
315            total_saved as f64 / total_original as f64
316        } else {
317            0.0
318        };
319
320        let modes_used: std::collections::HashSet<&str> =
321            calls.iter().filter_map(|c| c.mode.as_deref()).collect();
322        let mode_diversity = (modes_used.len() as f64 / 6.0).min(1.0);
323        let cache_util = stats.hit_rate() / 100.0;
324        let cep_score = cache_util * 0.3 + mode_diversity * 0.2 + compression_rate * 0.5;
325
326        let mut mode_counts: std::collections::HashMap<String, u64> =
327            std::collections::HashMap::new();
328        for call in calls {
329            if let Some(ref mode) = call.mode {
330                *mode_counts.entry(mode.clone()).or_insert(0) += 1;
331            }
332        }
333
334        CepComputedStats {
335            cep_score: (cep_score * 100.0).round() as u32,
336            cache_util: (cache_util * 100.0).round() as u32,
337            mode_diversity: (mode_diversity * 100.0).round() as u32,
338            compression_rate: (compression_rate * 100.0).round() as u32,
339            total_original,
340            total_compressed,
341            total_saved,
342            mode_counts,
343            complexity: format!("{:?}", complexity),
344            cache_hits: stats.cache_hits,
345            total_reads: stats.total_reads,
346            tool_call_count: calls.len() as u64,
347        }
348    }
349
350    async fn write_mcp_live_stats(&self) {
351        let cache = self.cache.read().await;
352        let calls = self.tool_calls.read().await;
353        let stats = cache.get_stats();
354        let complexity = crate::core::adaptive::classify_from_context(&cache);
355
356        let cs = Self::compute_cep_stats(&calls, stats, &complexity);
357
358        drop(cache);
359        drop(calls);
360
361        let live = serde_json::json!({
362            "cep_score": cs.cep_score,
363            "cache_utilization": cs.cache_util,
364            "mode_diversity": cs.mode_diversity,
365            "compression_rate": cs.compression_rate,
366            "task_complexity": cs.complexity,
367            "files_cached": cs.total_reads,
368            "total_reads": cs.total_reads,
369            "cache_hits": cs.cache_hits,
370            "tokens_saved": cs.total_saved,
371            "tokens_original": cs.total_original,
372            "tool_calls": cs.tool_call_count,
373            "updated_at": chrono::Local::now().to_rfc3339(),
374        });
375
376        if let Some(dir) = dirs::home_dir().map(|h| h.join(".lean-ctx")) {
377            let _ = std::fs::write(dir.join("mcp-live.json"), live.to_string());
378        }
379    }
380
381    pub async fn record_cep_snapshot(&self) {
382        let cache = self.cache.read().await;
383        let calls = self.tool_calls.read().await;
384        let stats = cache.get_stats();
385        let complexity = crate::core::adaptive::classify_from_context(&cache);
386
387        let cs = Self::compute_cep_stats(&calls, stats, &complexity);
388
389        drop(cache);
390        drop(calls);
391
392        crate::core::stats::record_cep_session(
393            cs.cep_score,
394            cs.cache_hits,
395            cs.total_reads,
396            cs.total_original,
397            cs.total_compressed,
398            &cs.mode_counts,
399            cs.tool_call_count,
400            &cs.complexity,
401        );
402    }
403}
404
405pub fn create_server() -> LeanCtxServer {
406    LeanCtxServer::new()
407}
408
409fn auto_consolidate_knowledge(project_root: &str) {
410    use crate::core::knowledge::ProjectKnowledge;
411    use crate::core::session::SessionState;
412
413    let session = match SessionState::load_latest() {
414        Some(s) => s,
415        None => return,
416    };
417
418    if session.findings.is_empty() && session.decisions.is_empty() {
419        return;
420    }
421
422    let mut knowledge = ProjectKnowledge::load_or_create(project_root);
423
424    for finding in &session.findings {
425        let key = if let Some(ref file) = finding.file {
426            if let Some(line) = finding.line {
427                format!("{file}:{line}")
428            } else {
429                file.clone()
430            }
431        } else {
432            "finding-auto".to_string()
433        };
434        knowledge.remember("finding", &key, &finding.summary, &session.id, 0.7);
435    }
436
437    for decision in &session.decisions {
438        let key = decision
439            .summary
440            .chars()
441            .take(50)
442            .collect::<String>()
443            .replace(' ', "-")
444            .to_lowercase();
445        knowledge.remember("decision", &key, &decision.summary, &session.id, 0.85);
446    }
447
448    let task_desc = session
449        .task
450        .as_ref()
451        .map(|t| t.description.clone())
452        .unwrap_or_default();
453
454    let summary = format!(
455        "Auto-consolidate session {}: {} — {} findings, {} decisions",
456        session.id,
457        task_desc,
458        session.findings.len(),
459        session.decisions.len()
460    );
461    knowledge.consolidate(&summary, vec![session.id.clone()]);
462    let _ = knowledge.save();
463}