1mod dispatch;
2mod execute;
3pub mod helpers;
4
5use rmcp::handler::server::ServerHandler;
6use rmcp::model::*;
7use rmcp::service::{RequestContext, RoleServer};
8use rmcp::ErrorData;
9
10use crate::tools::{CrpMode, LeanCtxServer};
11
12use helpers::{canonical_args_string, extract_search_pattern_from_command, get_str, md5_hex};
13
14impl ServerHandler for LeanCtxServer {
15 fn get_info(&self) -> ServerInfo {
16 let capabilities = ServerCapabilities::builder().enable_tools().build();
17
18 let instructions = crate::instructions::build_instructions(self.crp_mode);
19
20 InitializeResult::new(capabilities)
21 .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
22 .with_instructions(instructions)
23 }
24
25 async fn initialize(
26 &self,
27 request: InitializeRequestParams,
28 _context: RequestContext<RoleServer>,
29 ) -> Result<InitializeResult, ErrorData> {
30 let name = request.client_info.name.clone();
31 tracing::info!("MCP client connected: {:?}", name);
32 *self.client_name.write().await = name.clone();
33
34 tokio::task::spawn_blocking(|| {
35 if let Some(home) = dirs::home_dir() {
36 let _ = crate::rules_inject::inject_all_rules(&home);
37 }
38 crate::hooks::refresh_installed_hooks();
39 crate::core::version_check::check_background();
40 });
41
42 let instructions =
43 crate::instructions::build_instructions_with_client(self.crp_mode, &name);
44 let capabilities = ServerCapabilities::builder().enable_tools().build();
45
46 Ok(InitializeResult::new(capabilities)
47 .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
48 .with_instructions(instructions))
49 }
50
51 async fn list_tools(
52 &self,
53 _request: Option<PaginatedRequestParams>,
54 _context: RequestContext<RoleServer>,
55 ) -> Result<ListToolsResult, ErrorData> {
56 let all_tools = if crate::tool_defs::is_lazy_mode() {
57 crate::tool_defs::lazy_tool_defs()
58 } else if std::env::var("LEAN_CTX_UNIFIED").is_ok()
59 && std::env::var("LEAN_CTX_FULL_TOOLS").is_err()
60 {
61 crate::tool_defs::unified_tool_defs()
62 } else {
63 crate::tool_defs::granular_tool_defs()
64 };
65
66 let disabled = crate::core::config::Config::load().disabled_tools_effective();
67 let tools = if disabled.is_empty() {
68 all_tools
69 } else {
70 all_tools
71 .into_iter()
72 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
73 .collect()
74 };
75
76 let tools = {
77 let active = self.workflow.read().await.clone();
78 if let Some(run) = active {
79 if let Some(state) = run.spec.state(&run.current) {
80 if let Some(allowed) = &state.allowed_tools {
81 let mut allow: std::collections::HashSet<&str> =
82 allowed.iter().map(|s| s.as_str()).collect();
83 allow.insert("ctx");
84 allow.insert("ctx_workflow");
85 return Ok(ListToolsResult {
86 tools: tools
87 .into_iter()
88 .filter(|t| allow.contains(t.name.as_ref()))
89 .collect(),
90 ..Default::default()
91 });
92 }
93 }
94 }
95 tools
96 };
97
98 Ok(ListToolsResult {
99 tools,
100 ..Default::default()
101 })
102 }
103
104 async fn call_tool(
105 &self,
106 request: CallToolRequestParams,
107 _context: RequestContext<RoleServer>,
108 ) -> Result<CallToolResult, ErrorData> {
109 self.check_idle_expiry().await;
110
111 let original_name = request.name.as_ref().to_string();
112 let (resolved_name, resolved_args) = if original_name == "ctx" {
113 let sub = request
114 .arguments
115 .as_ref()
116 .and_then(|a| a.get("tool"))
117 .and_then(|v| v.as_str())
118 .map(|s| s.to_string())
119 .ok_or_else(|| {
120 ErrorData::invalid_params("'tool' is required for ctx meta-tool", None)
121 })?;
122 let tool_name = if sub.starts_with("ctx_") {
123 sub
124 } else {
125 format!("ctx_{sub}")
126 };
127 let mut args = request.arguments.unwrap_or_default();
128 args.remove("tool");
129 (tool_name, Some(args))
130 } else {
131 (original_name, request.arguments)
132 };
133 let name = resolved_name.as_str();
134 let args = &resolved_args;
135
136 if name != "ctx_workflow" {
137 let active = self.workflow.read().await.clone();
138 if let Some(run) = active {
139 if let Some(state) = run.spec.state(&run.current) {
140 if let Some(allowed) = &state.allowed_tools {
141 let allowed_ok = allowed.iter().any(|t| t == name) || name == "ctx";
142 if !allowed_ok {
143 let mut shown = allowed.clone();
144 shown.sort();
145 shown.truncate(30);
146 return Ok(CallToolResult::success(vec![Content::text(format!(
147 "Tool '{name}' blocked by workflow '{}' (state: {}). Allowed ({} shown): {}",
148 run.spec.name,
149 run.current,
150 shown.len(),
151 shown.join(", ")
152 ))]));
153 }
154 }
155 }
156 }
157 }
158
159 let auto_context = {
160 let task = {
161 let session = self.session.read().await;
162 session.task.as_ref().map(|t| t.description.clone())
163 };
164 let project_root = {
165 let session = self.session.read().await;
166 session.project_root.clone()
167 };
168 let mut cache = self.cache.write().await;
169 crate::tools::autonomy::session_lifecycle_pre_hook(
170 &self.autonomy,
171 name,
172 &mut cache,
173 task.as_deref(),
174 project_root.as_deref(),
175 self.crp_mode,
176 )
177 };
178
179 let throttle_result = {
180 let fp = args
181 .as_ref()
182 .map(|a| {
183 crate::core::loop_detection::LoopDetector::fingerprint(
184 &serde_json::Value::Object(a.clone()),
185 )
186 })
187 .unwrap_or_default();
188 let mut detector = self.loop_detector.write().await;
189
190 let is_search = crate::core::loop_detection::LoopDetector::is_search_tool(name);
191 let is_search_shell = name == "ctx_shell" && {
192 let cmd = args
193 .as_ref()
194 .and_then(|a| a.get("command"))
195 .and_then(|v| v.as_str())
196 .unwrap_or("");
197 crate::core::loop_detection::LoopDetector::is_search_shell_command(cmd)
198 };
199
200 if is_search || is_search_shell {
201 let search_pattern = args.as_ref().and_then(|a| {
202 a.get("pattern")
203 .or_else(|| a.get("query"))
204 .and_then(|v| v.as_str())
205 });
206 let shell_pattern = if is_search_shell {
207 args.as_ref()
208 .and_then(|a| a.get("command"))
209 .and_then(|v| v.as_str())
210 .and_then(extract_search_pattern_from_command)
211 } else {
212 None
213 };
214 let pat = search_pattern.or(shell_pattern.as_deref());
215 detector.record_search(name, &fp, pat)
216 } else {
217 detector.record_call(name, &fp)
218 }
219 };
220
221 if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Blocked {
222 let msg = throttle_result.message.unwrap_or_default();
223 return Ok(CallToolResult::success(vec![Content::text(msg)]));
224 }
225
226 let throttle_warning =
227 if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Reduced {
228 throttle_result.message.clone()
229 } else {
230 None
231 };
232
233 let tool_start = std::time::Instant::now();
234 let result_text = self.dispatch_tool(name, args).await?;
235
236 let mut result_text = result_text;
237
238 {
239 let config = crate::core::config::Config::load();
240 let density = crate::core::config::OutputDensity::effective(&config.output_density);
241 result_text = crate::core::protocol::compress_output(&result_text, &density);
242 }
243
244 if let Some(ctx) = auto_context {
245 result_text = format!("{ctx}\n\n{result_text}");
246 }
247
248 if let Some(warning) = throttle_warning {
249 result_text = format!("{result_text}\n\n{warning}");
250 }
251
252 if name == "ctx_read" {
253 let read_path = self
254 .resolve_path_or_passthrough(&get_str(args, "path").unwrap_or_default())
255 .await;
256 let project_root = {
257 let session = self.session.read().await;
258 session.project_root.clone()
259 };
260 let mut cache = self.cache.write().await;
261 let enrich = crate::tools::autonomy::enrich_after_read(
262 &self.autonomy,
263 &mut cache,
264 &read_path,
265 project_root.as_deref(),
266 );
267 if let Some(hint) = enrich.related_hint {
268 result_text = format!("{result_text}\n{hint}");
269 }
270
271 crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache);
272 }
273
274 if name == "ctx_shell" {
275 let cmd = get_str(args, "command").unwrap_or_default();
276 let output_tokens = crate::core::tokens::count_tokens(&result_text);
277 let calls = self.tool_calls.read().await;
278 let last_original = calls.last().map(|c| c.original_tokens).unwrap_or(0);
279 drop(calls);
280 if let Some(hint) = crate::tools::autonomy::shell_efficiency_hint(
281 &self.autonomy,
282 &cmd,
283 last_original,
284 output_tokens,
285 ) {
286 result_text = format!("{result_text}\n{hint}");
287 }
288 }
289
290 {
291 let input = canonical_args_string(args);
292 let input_md5 = md5_hex(&input);
293 let output_md5 = md5_hex(&result_text);
294 let action = get_str(args, "action");
295 let agent_id = self.agent_id.read().await.clone();
296 let client_name = self.client_name.read().await.clone();
297 let mut explicit_intent: Option<(
298 crate::core::intent_protocol::IntentRecord,
299 Option<String>,
300 String,
301 )> = None;
302
303 {
304 let empty_args = serde_json::Map::new();
305 let args_map = args.as_ref().unwrap_or(&empty_args);
306 let mut session = self.session.write().await;
307 session.record_tool_receipt(
308 name,
309 action.as_deref(),
310 &input_md5,
311 &output_md5,
312 agent_id.as_deref(),
313 Some(&client_name),
314 );
315
316 if let Some(intent) = crate::core::intent_protocol::infer_from_tool_call(
317 name,
318 action.as_deref(),
319 args_map,
320 session.project_root.as_deref(),
321 ) {
322 let is_explicit =
323 intent.source == crate::core::intent_protocol::IntentSource::Explicit;
324 let root = session.project_root.clone();
325 let sid = session.id.clone();
326 session.record_intent(intent.clone());
327 if is_explicit {
328 explicit_intent = Some((intent, root, sid));
329 }
330 }
331 if session.should_save() {
332 let _ = session.save();
333 }
334 }
335
336 if let Some((intent, root, session_id)) = explicit_intent {
337 crate::core::intent_protocol::apply_side_effects(
338 &intent,
339 root.as_deref(),
340 &session_id,
341 );
342 }
343
344 if self.autonomy.is_enabled() {
345 let (calls, project_root) = {
346 let session = self.session.read().await;
347 (session.stats.total_tool_calls, session.project_root.clone())
348 };
349
350 if let Some(root) = project_root {
351 if crate::tools::autonomy::should_auto_consolidate(&self.autonomy, calls) {
352 let root_clone = root.clone();
353 tokio::task::spawn_blocking(move || {
354 let _ = crate::core::consolidation_engine::consolidate_latest(
355 &root_clone,
356 crate::core::consolidation_engine::ConsolidationBudgets::default(),
357 );
358 });
359 }
360 }
361 }
362
363 let agent_key = agent_id.unwrap_or_else(|| "unknown".to_string());
364 let input_tokens = crate::core::tokens::count_tokens(&input) as u64;
365 let output_tokens = crate::core::tokens::count_tokens(&result_text) as u64;
366 let mut store = crate::core::a2a::cost_attribution::CostStore::load();
367 store.record_tool_call(&agent_key, &client_name, name, input_tokens, output_tokens);
368 let _ = store.save();
369 }
370
371 let skip_checkpoint = matches!(
372 name,
373 "ctx_compress"
374 | "ctx_metrics"
375 | "ctx_benchmark"
376 | "ctx_analyze"
377 | "ctx_cache"
378 | "ctx_discover"
379 | "ctx_dedup"
380 | "ctx_session"
381 | "ctx_knowledge"
382 | "ctx_agent"
383 | "ctx_share"
384 | "ctx_wrapped"
385 | "ctx_overview"
386 | "ctx_preload"
387 | "ctx_cost"
388 | "ctx_gain"
389 | "ctx_heatmap"
390 | "ctx_task"
391 | "ctx_impact"
392 | "ctx_architecture"
393 | "ctx_workflow"
394 );
395
396 if !skip_checkpoint && self.increment_and_check() {
397 if let Some(checkpoint) = self.auto_checkpoint().await {
398 let combined = format!(
399 "{result_text}\n\n--- AUTO CHECKPOINT (every {} calls) ---\n{checkpoint}",
400 self.checkpoint_interval
401 );
402 return Ok(CallToolResult::success(vec![Content::text(combined)]));
403 }
404 }
405
406 let tool_duration_ms = tool_start.elapsed().as_millis() as u64;
407 if tool_duration_ms > 100 {
408 LeanCtxServer::append_tool_call_log(
409 name,
410 tool_duration_ms,
411 0,
412 0,
413 None,
414 &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
415 );
416 }
417
418 let current_count = self.call_count.load(std::sync::atomic::Ordering::Relaxed);
419 if current_count > 0 && current_count.is_multiple_of(100) {
420 std::thread::spawn(crate::cloud_sync::cloud_background_tasks);
421 }
422
423 Ok(CallToolResult::success(vec![Content::text(result_text)]))
424 }
425}
426
427pub fn build_instructions_for_test(crp_mode: CrpMode) -> String {
428 crate::instructions::build_instructions(crp_mode)
429}
430
431pub fn build_claude_code_instructions_for_test() -> String {
432 crate::instructions::claude_code_instructions()
433}
434
435pub fn tool_descriptions_for_test() -> Vec<(&'static str, &'static str)> {
436 crate::tool_defs::list_all_tool_defs()
437 .into_iter()
438 .map(|(name, desc, _)| (name, desc))
439 .collect()
440}
441
442pub fn tool_schemas_json_for_test() -> String {
443 crate::tool_defs::list_all_tool_defs()
444 .iter()
445 .map(|(name, _, schema)| format!("{}: {}", name, schema))
446 .collect::<Vec<_>>()
447 .join("\n")
448}
449
450#[cfg(test)]
451mod tests {
452 #[test]
453 fn test_unified_tool_count() {
454 let tools = crate::tool_defs::unified_tool_defs();
455 assert_eq!(tools.len(), 5, "Expected 5 unified tools");
456 }
457
458 #[test]
459 fn test_granular_tool_count() {
460 let tools = crate::tool_defs::granular_tool_defs();
461 assert!(tools.len() >= 25, "Expected at least 25 granular tools");
462 }
463
464 #[test]
465 fn disabled_tools_filters_list() {
466 let all = crate::tool_defs::granular_tool_defs();
467 let total = all.len();
468 let disabled = ["ctx_graph".to_string(), "ctx_agent".to_string()];
469 let filtered: Vec<_> = all
470 .into_iter()
471 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
472 .collect();
473 assert_eq!(filtered.len(), total - 2);
474 assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_graph"));
475 assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_agent"));
476 }
477
478 #[test]
479 fn empty_disabled_tools_returns_all() {
480 let all = crate::tool_defs::granular_tool_defs();
481 let total = all.len();
482 let disabled: Vec<String> = vec![];
483 let filtered: Vec<_> = all
484 .into_iter()
485 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
486 .collect();
487 assert_eq!(filtered.len(), total);
488 }
489
490 #[test]
491 fn misspelled_disabled_tool_is_silently_ignored() {
492 let all = crate::tool_defs::granular_tool_defs();
493 let total = all.len();
494 let disabled = ["ctx_nonexistent_tool".to_string()];
495 let filtered: Vec<_> = all
496 .into_iter()
497 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
498 .collect();
499 assert_eq!(filtered.len(), total);
500 }
501}