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