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