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