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