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 fn increment_and_check(&self) -> bool {
181 let count = self.call_count.fetch_add(1, Ordering::Relaxed) + 1;
182 self.checkpoint_interval > 0 && count.is_multiple_of(self.checkpoint_interval)
183 }
184
185 pub async fn auto_checkpoint(&self) -> Option<String> {
186 let cache = self.cache.read().await;
187 if cache.get_all_entries().is_empty() {
188 return None;
189 }
190 let complexity = crate::core::adaptive::classify_from_context(&cache);
191 let checkpoint = ctx_compress::handle(&cache, true, self.crp_mode);
192 drop(cache);
193
194 let mut session = self.session.write().await;
195 let _ = session.save();
196 let session_summary = session.format_compact();
197 drop(session);
198
199 self.record_call("ctx_compress", 0, 0, Some("auto".to_string()))
200 .await;
201
202 self.write_mcp_live_stats().await;
203
204 Some(format!(
205 "{checkpoint}\n\n--- SESSION STATE ---\n{session_summary}\n\n{}",
206 complexity.instruction_suffix()
207 ))
208 }
209
210 async fn write_mcp_live_stats(&self) {
211 let cache = self.cache.read().await;
212 let calls = self.tool_calls.read().await;
213 let stats = cache.get_stats();
214 let complexity = crate::core::adaptive::classify_from_context(&cache);
215
216 let total_original: u64 = calls.iter().map(|c| c.original_tokens as u64).sum();
217 let total_saved: u64 = calls.iter().map(|c| c.saved_tokens as u64).sum();
218 let total_compressed = total_original.saturating_sub(total_saved);
219 let compression_rate = if total_original > 0 {
220 total_saved as f64 / total_original as f64
221 } else {
222 0.0
223 };
224
225 let modes_used: std::collections::HashSet<&str> =
226 calls.iter().filter_map(|c| c.mode.as_deref()).collect();
227 let mode_diversity = (modes_used.len() as f64 / 6.0).min(1.0);
228 let cache_util = stats.hit_rate() / 100.0;
229 let cep_score = cache_util * 0.3 + mode_diversity * 0.2 + compression_rate * 0.5;
230 let cep_score_u32 = (cep_score * 100.0).round() as u32;
231
232 let live = serde_json::json!({
233 "cep_score": cep_score_u32,
234 "cache_utilization": (cache_util * 100.0).round() as u32,
235 "mode_diversity": (mode_diversity * 100.0).round() as u32,
236 "compression_rate": (compression_rate * 100.0).round() as u32,
237 "task_complexity": format!("{:?}", complexity),
238 "files_cached": stats.files_tracked,
239 "total_reads": stats.total_reads,
240 "cache_hits": stats.cache_hits,
241 "tokens_saved": total_saved,
242 "tokens_original": total_original,
243 "tool_calls": calls.len(),
244 "updated_at": chrono::Local::now().to_rfc3339(),
245 });
246
247 let mut mode_counts: std::collections::HashMap<String, u64> =
248 std::collections::HashMap::new();
249 for call in calls.iter() {
250 if let Some(ref mode) = call.mode {
251 *mode_counts.entry(mode.clone()).or_insert(0) += 1;
252 }
253 }
254
255 let tool_call_count = calls.len() as u64;
256 let complexity_str = format!("{:?}", complexity);
257 let cache_hits = stats.cache_hits;
258 let total_reads = stats.total_reads;
259
260 drop(cache);
261 drop(calls);
262
263 if let Some(dir) = dirs::home_dir().map(|h| h.join(".lean-ctx")) {
264 let _ = std::fs::write(dir.join("mcp-live.json"), live.to_string());
265 }
266
267 crate::core::stats::record_cep_session(
268 cep_score_u32,
269 cache_hits,
270 total_reads,
271 total_original,
272 total_compressed,
273 &mode_counts,
274 tool_call_count,
275 &complexity_str,
276 );
277 }
278}
279
280pub fn create_server() -> LeanCtxServer {
281 LeanCtxServer::new()
282}