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