1mod dispatch;
2mod execute;
3pub mod helpers;
4pub mod pipeline_stages;
5pub mod registry;
6pub mod role_guard;
7pub mod tool_trait;
8
9use rmcp::handler::server::ServerHandler;
10use rmcp::model::{
11 CallToolRequestParams, CallToolResult, Content, Implementation, InitializeRequestParams,
12 InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo,
13};
14use rmcp::service::{RequestContext, RoleServer};
15use rmcp::ErrorData;
16
17use crate::tools::{CrpMode, LeanCtxServer};
18
19impl ServerHandler for LeanCtxServer {
20 fn get_info(&self) -> ServerInfo {
21 let capabilities = ServerCapabilities::builder().enable_tools().build();
22
23 let config = crate::core::config::Config::load();
24 let level = crate::core::config::CompressionLevel::effective(&config);
25 let _ = crate::core::terse::rules_inject::inject(&level);
26
27 let instructions = crate::instructions::build_instructions(CrpMode::effective());
28
29 InitializeResult::new(capabilities)
30 .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
31 .with_instructions(instructions)
32 }
33
34 async fn initialize(
35 &self,
36 request: InitializeRequestParams,
37 _context: RequestContext<RoleServer>,
38 ) -> Result<InitializeResult, ErrorData> {
39 let name = request.client_info.name.clone();
40 tracing::info!("MCP client connected: {:?}", name);
41 *self.client_name.write().await = name.clone();
42
43 let derived_root = derive_project_root_from_cwd();
44 let cwd_str = std::env::current_dir()
45 .ok()
46 .map(|p| p.to_string_lossy().to_string())
47 .unwrap_or_default();
48 {
49 let mut session = self.session.write().await;
50 if !cwd_str.is_empty() {
51 session.shell_cwd = Some(cwd_str.clone());
52 }
53 if let Some(ref root) = derived_root {
54 session.project_root = Some(root.clone());
55 tracing::info!("Project root set to: {root}");
56 } else if let Some(ref root) = session.project_root {
57 let root_path = std::path::Path::new(root);
58 let root_has_marker = has_project_marker(root_path);
59 let root_str = root_path.to_string_lossy();
60 let root_suspicious = root_str.contains("/.claude")
61 || root_str.contains("/.codex")
62 || root_str.contains("/var/folders/")
63 || root_str.contains("/tmp/")
64 || root_str.contains("\\.claude")
65 || root_str.contains("\\.codex")
66 || root_str.contains("\\AppData\\Local\\Temp")
67 || root_str.contains("\\Temp\\");
68 if root_suspicious && !root_has_marker {
69 session.project_root = None;
70 }
71 }
72 if self.session_mode == crate::tools::SessionMode::Shared {
73 if let Some(ref root) = session.project_root {
74 if let Some(ref rt) = self.context_os {
75 rt.shared_sessions.persist_best_effort(
76 root,
77 &self.workspace_id,
78 &self.channel_id,
79 &session,
80 );
81 rt.metrics.record_session_persisted();
82 }
83 }
84 } else {
85 let _ = session.save();
86 }
87 }
88
89 let agent_name = name.clone();
90 let agent_root = derived_root.clone().unwrap_or_default();
91 let agent_id_handle = self.agent_id.clone();
92 tokio::task::spawn_blocking(move || {
93 if std::env::var("LEAN_CTX_HEADLESS").is_ok() {
94 return;
95 }
96
97 let maintenance = crate::core::startup_guard::try_acquire_lock(
101 "startup-maintenance",
102 std::time::Duration::from_secs(2),
103 std::time::Duration::from_mins(2),
104 );
105 if maintenance.is_some() {
106 if let Some(home) = dirs::home_dir() {
107 let _ = crate::rules_inject::inject_all_rules(&home);
108 }
109 crate::hooks::refresh_installed_hooks();
110 crate::core::version_check::check_background();
111 }
112 drop(maintenance);
113
114 if !agent_root.is_empty() {
115 let heuristic_role = match agent_name.to_lowercase().as_str() {
116 n if n.contains("cursor") => Some("coder"),
117 n if n.contains("claude") => Some("coder"),
118 n if n.contains("codex") => Some("coder"),
119 n if n.contains("antigravity") || n.contains("gemini") => Some("coder"),
120 n if n.contains("review") => Some("reviewer"),
121 n if n.contains("test") => Some("debugger"),
122 _ => None,
123 };
124 let env_role = std::env::var("LEAN_CTX_ROLE")
125 .or_else(|_| std::env::var("LEAN_CTX_AGENT_ROLE"))
126 .ok();
127 let effective_role = env_role.as_deref().or(heuristic_role).unwrap_or("coder");
128
129 let _ = crate::core::roles::set_active_role(effective_role);
130
131 let mut registry = crate::core::agents::AgentRegistry::load_or_create();
132 registry.cleanup_stale(24);
133 let id = registry.register("mcp", Some(effective_role), &agent_root);
134 let _ = registry.save();
135 if let Ok(mut guard) = agent_id_handle.try_write() {
136 *guard = Some(id);
137 }
138 }
139 });
140
141 let instructions =
142 crate::instructions::build_instructions_with_client(CrpMode::effective(), &name);
143 let capabilities = ServerCapabilities::builder().enable_tools().build();
144
145 Ok(InitializeResult::new(capabilities)
146 .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
147 .with_instructions(instructions))
148 }
149
150 async fn list_tools(
151 &self,
152 _request: Option<PaginatedRequestParams>,
153 _context: RequestContext<RoleServer>,
154 ) -> Result<ListToolsResult, ErrorData> {
155 let all_tools = if crate::tool_defs::is_full_mode() {
156 crate::tool_defs::granular_tool_defs()
157 } else if std::env::var("LEAN_CTX_UNIFIED").is_ok() {
158 crate::tool_defs::unified_tool_defs()
159 } else {
160 crate::tool_defs::lazy_tool_defs()
161 };
162
163 let disabled = crate::core::config::Config::load().disabled_tools_effective();
164 let tools = if disabled.is_empty() {
165 all_tools
166 } else {
167 all_tools
168 .into_iter()
169 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
170 .collect()
171 };
172
173 let tools = {
174 let active = self.workflow.read().await.clone();
175 if let Some(run) = active {
176 if let Some(state) = run.spec.state(&run.current) {
177 if let Some(allowed) = &state.allowed_tools {
178 let mut allow: std::collections::HashSet<&str> =
179 allowed.iter().map(std::string::String::as_str).collect();
180 allow.insert("ctx");
181 allow.insert("ctx_workflow");
182 return Ok(ListToolsResult {
183 tools: tools
184 .into_iter()
185 .filter(|t| allow.contains(t.name.as_ref()))
186 .collect(),
187 ..Default::default()
188 });
189 }
190 }
191 }
192 tools
193 };
194
195 let tools = {
196 let level = crate::core::config::Config::load().compression_level;
197 let mode =
198 crate::core::terse::mcp_compress::DescriptionMode::from_compression_level(&level);
199 if mode == crate::core::terse::mcp_compress::DescriptionMode::Full {
200 tools
201 } else {
202 tools
203 .into_iter()
204 .map(|mut t| {
205 let compressed = crate::core::terse::mcp_compress::compress_description(
206 t.name.as_ref(),
207 t.description.as_deref().unwrap_or(""),
208 mode,
209 );
210 t.description = Some(compressed.into());
211 t
212 })
213 .collect()
214 }
215 };
216
217 Ok(ListToolsResult {
218 tools,
219 ..Default::default()
220 })
221 }
222
223 async fn call_tool(
224 &self,
225 request: CallToolRequestParams,
226 _context: RequestContext<RoleServer>,
227 ) -> Result<CallToolResult, ErrorData> {
228 self.check_idle_expiry().await;
229
230 let original_name = request.name.as_ref().to_string();
231 let (resolved_name, resolved_args) = if original_name == "ctx" {
232 let sub = request
233 .arguments
234 .as_ref()
235 .and_then(|a| a.get("tool"))
236 .and_then(|v| v.as_str())
237 .map(std::string::ToString::to_string)
238 .ok_or_else(|| {
239 ErrorData::invalid_params("'tool' is required for ctx meta-tool", None)
240 })?;
241 let tool_name = if sub.starts_with("ctx_") {
242 sub
243 } else {
244 format!("ctx_{sub}")
245 };
246 let mut args = request.arguments.unwrap_or_default();
247 args.remove("tool");
248 (tool_name, Some(args))
249 } else {
250 (original_name, request.arguments)
251 };
252 let name = resolved_name.as_str();
253 let args = resolved_args.as_ref();
254
255 let role_check = role_guard::check_tool_access(name);
256 if let Some(denied) = role_guard::into_call_tool_result(&role_check) {
257 tracing::warn!(
258 tool = name,
259 role = %role_check.role_name,
260 "Tool blocked by role policy"
261 );
262 return Ok(denied);
263 }
264
265 if name != "ctx_workflow" {
266 let active = self.workflow.read().await.clone();
267 if let Some(run) = active {
268 if let Some(state) = run.spec.state(&run.current) {
269 if let Some(allowed) = &state.allowed_tools {
270 let allowed_ok = allowed.iter().any(|t| t == name) || name == "ctx";
271 if !allowed_ok {
272 let mut shown = allowed.clone();
273 shown.sort();
274 shown.truncate(30);
275 return Ok(CallToolResult::success(vec![Content::text(format!(
276 "Tool '{name}' blocked by workflow '{}' (state: {}). Allowed ({} shown): {}",
277 run.spec.name,
278 run.current,
279 shown.len(),
280 shown.join(", ")
281 ))]));
282 }
283 }
284 }
285 }
286 }
287
288 let auto_context = {
289 let task = {
290 let session = self.session.read().await;
291 session.task.as_ref().map(|t| t.description.clone())
292 };
293 let project_root = {
294 let session = self.session.read().await;
295 session.project_root.clone()
296 };
297 let mut cache = self.cache.write().await;
298 crate::tools::autonomy::session_lifecycle_pre_hook(
299 &self.autonomy,
300 name,
301 &mut cache,
302 task.as_deref(),
303 project_root.as_deref(),
304 CrpMode::effective(),
305 )
306 };
307
308 let throttle_result = {
309 let fp = args
310 .map(|a| {
311 crate::core::loop_detection::LoopDetector::fingerprint(
312 &serde_json::Value::Object(a.clone()),
313 )
314 })
315 .unwrap_or_default();
316 let mut detector = self.loop_detector.write().await;
317
318 let is_search = crate::core::loop_detection::LoopDetector::is_search_tool(name);
319 let is_search_shell = name == "ctx_shell" && {
320 let cmd = args
321 .as_ref()
322 .and_then(|a| a.get("command"))
323 .and_then(|v| v.as_str())
324 .unwrap_or("");
325 crate::core::loop_detection::LoopDetector::is_search_shell_command(cmd)
326 };
327
328 if is_search || is_search_shell {
329 let search_pattern = args.and_then(|a| {
330 a.get("pattern")
331 .or_else(|| a.get("query"))
332 .and_then(|v| v.as_str())
333 });
334 let shell_pattern = if is_search_shell {
335 args.and_then(|a| a.get("command"))
336 .and_then(|v| v.as_str())
337 .and_then(helpers::extract_search_pattern_from_command)
338 } else {
339 None
340 };
341 let pat = search_pattern.or(shell_pattern.as_deref());
342 detector.record_search(name, &fp, pat)
343 } else {
344 detector.record_call(name, &fp)
345 }
346 };
347
348 if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Blocked {
349 let msg = throttle_result.message.unwrap_or_default();
350 return Ok(CallToolResult::success(vec![Content::text(msg)]));
351 }
352
353 let throttle_warning =
354 if throttle_result.level == crate::core::loop_detection::ThrottleLevel::Reduced {
355 throttle_result.message.clone()
356 } else {
357 None
358 };
359
360 let config = crate::core::config::Config::load();
361 let minimal = config.minimal_overhead_effective();
362
363 {
364 use crate::core::budget_tracker::{BudgetLevel, BudgetTracker};
365 let snap = BudgetTracker::global().check();
366 if *snap.worst_level() == BudgetLevel::Exhausted
367 && name != "ctx_session"
368 && name != "ctx_cost"
369 && name != "ctx_metrics"
370 {
371 for (dim, lvl, used, limit) in [
372 (
373 "tokens",
374 &snap.tokens.level,
375 format!("{}", snap.tokens.used),
376 format!("{}", snap.tokens.limit),
377 ),
378 (
379 "shell",
380 &snap.shell.level,
381 format!("{}", snap.shell.used),
382 format!("{}", snap.shell.limit),
383 ),
384 (
385 "cost",
386 &snap.cost.level,
387 format!("${:.2}", snap.cost.used_usd),
388 format!("${:.2}", snap.cost.limit_usd),
389 ),
390 ] {
391 if *lvl == BudgetLevel::Exhausted {
392 crate::core::events::emit_budget_exhausted(&snap.role, dim, &used, &limit);
393 }
394 }
395 let msg = format!(
396 "[BUDGET EXHAUSTED] {}\n\
397 Use `ctx_session action=role` to check/switch roles, \
398 or `ctx_session action=reset` to start fresh.",
399 snap.format_compact()
400 );
401 tracing::warn!(tool = name, "{msg}");
402 return Ok(CallToolResult::success(vec![Content::text(msg)]));
403 }
404 }
405
406 if is_shell_tool_name(name) {
407 crate::core::budget_tracker::BudgetTracker::global().record_shell();
408 }
409
410 let tool_start = std::time::Instant::now();
411 let mut result_text = {
412 use futures::FutureExt;
413 use std::panic::AssertUnwindSafe;
414 match AssertUnwindSafe(self.dispatch_tool(name, args, minimal))
415 .catch_unwind()
416 .await
417 {
418 Ok(Ok(text)) => text,
419 Ok(Err(e)) => return Err(e),
420 Err(panic_payload) => {
421 let detail = if let Some(s) = panic_payload.downcast_ref::<&str>() {
422 (*s).to_string()
423 } else if let Some(s) = panic_payload.downcast_ref::<String>() {
424 s.clone()
425 } else {
426 "unknown".to_string()
427 };
428 tracing::error!(tool = name, "Tool panicked: {detail}");
429 format!("ERROR: lean-ctx internal error in tool '{name}'.\n\
430 The MCP server is still running. Please retry or use a different approach.")
431 }
432 }
433 };
434
435 let is_raw_shell = name == "ctx_shell" && {
436 let arg_raw = helpers::get_bool(args, "raw").unwrap_or(false);
437 let arg_bypass = helpers::get_bool(args, "bypass").unwrap_or(false);
438 arg_raw
439 || arg_bypass
440 || std::env::var("LEAN_CTX_DISABLED").is_ok()
441 || std::env::var("LEAN_CTX_RAW").is_ok()
442 };
443
444 let pre_terse_len = result_text.len();
445 let output_tokens = {
446 let tokens = crate::core::tokens::count_tokens(&result_text) as u64;
447 crate::core::budget_tracker::BudgetTracker::global().record_tokens(tokens);
448 tokens
449 };
450
451 crate::core::anomaly::record_metric("tokens_per_call", output_tokens as f64);
452 crate::core::anomaly::save();
454
455 let budget_warning = {
456 use crate::core::budget_tracker::{BudgetLevel, BudgetTracker};
457 let snap = BudgetTracker::global().check();
458 if *snap.worst_level() == BudgetLevel::Warning {
459 for (dim, lvl, used, limit, pct) in [
460 (
461 "tokens",
462 &snap.tokens.level,
463 format!("{}", snap.tokens.used),
464 format!("{}", snap.tokens.limit),
465 snap.tokens.percent,
466 ),
467 (
468 "shell",
469 &snap.shell.level,
470 format!("{}", snap.shell.used),
471 format!("{}", snap.shell.limit),
472 snap.shell.percent,
473 ),
474 (
475 "cost",
476 &snap.cost.level,
477 format!("${:.2}", snap.cost.used_usd),
478 format!("${:.2}", snap.cost.limit_usd),
479 snap.cost.percent,
480 ),
481 ] {
482 if *lvl == BudgetLevel::Warning {
483 crate::core::events::emit_budget_warning(
484 &snap.role, dim, &used, &limit, pct,
485 );
486 }
487 }
488 Some(format!("[BUDGET WARNING] {}", snap.format_compact()))
489 } else {
490 None
491 }
492 };
493
494 let archive_hint = if minimal || is_raw_shell {
495 None
496 } else {
497 use crate::core::archive;
498 let archivable = matches!(
499 name,
500 "ctx_shell"
501 | "ctx_read"
502 | "ctx_multi_read"
503 | "ctx_smart_read"
504 | "ctx_execute"
505 | "ctx_search"
506 | "ctx_tree"
507 );
508 if archivable && archive::should_archive(&result_text) {
509 let cmd = helpers::get_str(args, "command")
510 .or_else(|| helpers::get_str(args, "path"))
511 .unwrap_or_default();
512 let session_id = self.session.read().await.id.clone();
513 let to_store = crate::core::redaction::redact_text_if_enabled(&result_text);
514 let tokens = crate::core::tokens::count_tokens(&to_store);
515 archive::store(name, &cmd, &to_store, Some(&session_id))
516 .map(|id| archive::format_hint(&id, to_store.len(), tokens))
517 } else {
518 None
519 }
520 };
521
522 let pre_compression = result_text.clone();
523 let skip_terse = is_raw_shell
524 || (name == "ctx_shell"
525 && helpers::get_str(args, "command")
526 .is_some_and(|c| crate::shell::compress::has_structural_output(&c)));
527 let compression = crate::core::config::CompressionLevel::effective(&config);
528 if compression.is_active() && !skip_terse {
529 let terse_result =
530 crate::core::terse::pipeline::compress(&result_text, &compression, None);
531 if terse_result.quality_passed && terse_result.savings_pct >= 1.0 {
532 result_text = terse_result.output;
533 }
534 }
535
536 if !is_raw_shell {
537 let verify_cfg = crate::core::profiles::active_profile().verification;
538 let vr = crate::core::output_verification::verify_output(
539 &pre_compression,
540 &result_text,
541 &verify_cfg,
542 );
543 if !vr.warnings.is_empty() {
544 let msg = format!("[VERIFY] {}", vr.format_compact());
545 result_text = format!("{result_text}\n\n{msg}");
546 }
547 }
548
549 if let Some(hint) = archive_hint {
550 result_text = format!("{result_text}\n{hint}");
551 }
552
553 if !is_raw_shell {
554 if let Some(ctx) = auto_context {
555 result_text = format!("{ctx}\n\n{result_text}");
556 }
557 }
558
559 if let Some(warning) = throttle_warning {
560 result_text = format!("{result_text}\n\n{warning}");
561 }
562
563 if let Some(bw) = budget_warning {
564 result_text = format!("{result_text}\n\n{bw}");
565 }
566
567 if !self
568 .rules_stale_checked
569 .swap(true, std::sync::atomic::Ordering::Relaxed)
570 {
571 let client = self.client_name.read().await.clone();
572 if !client.is_empty() {
573 if let Some(stale_msg) = crate::rules_inject::check_rules_freshness(&client) {
574 result_text = format!("{result_text}\n\n{stale_msg}");
575 }
576 }
577 }
578
579 {
580 let _ = crate::core::slo::evaluate();
582 }
583
584 if name == "ctx_read" {
585 if minimal {
586 let mut cache = self.cache.write().await;
587 crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache, name);
588 } else {
589 let read_path = self
590 .resolve_path_or_passthrough(
591 &helpers::get_str(args, "path").unwrap_or_default(),
592 )
593 .await;
594 let project_root = {
595 let session = self.session.read().await;
596 session.project_root.clone()
597 };
598 let mut cache = self.cache.write().await;
599 let enrich = crate::tools::autonomy::enrich_after_read(
600 &self.autonomy,
601 &mut cache,
602 &read_path,
603 project_root.as_deref(),
604 None,
605 crate::tools::CrpMode::effective(),
606 false,
607 );
608 if let Some(hint) = enrich.related_hint {
609 result_text = format!("{result_text}\n{hint}");
610 }
611 crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache, name);
612 }
613 }
614
615 if !minimal && !is_raw_shell && name == "ctx_shell" {
616 let cmd = helpers::get_str(args, "command").unwrap_or_default();
617 let calls = self.tool_calls.read().await;
618 let last_original = calls.last().map_or(0, |c| c.original_tokens);
619 drop(calls);
620 let pre_hint_tokens = crate::core::tokens::count_tokens(&result_text);
621 if let Some(hint) = crate::tools::autonomy::shell_efficiency_hint(
622 &self.autonomy,
623 &cmd,
624 last_original,
625 pre_hint_tokens,
626 ) {
627 result_text = format!("{result_text}\n{hint}");
628 }
629 }
630
631 #[allow(clippy::cast_possible_truncation)]
632 let output_token_count = if result_text.len() == pre_terse_len {
633 output_tokens as usize
634 } else {
635 crate::core::tokens::count_tokens(&result_text)
636 };
637 let action = helpers::get_str(args, "action");
638
639 const K_STALENESS_BOUND: i64 = 10;
641 if self.session_mode == crate::tools::SessionMode::Shared {
642 if let Some(ref rt) = self.context_os {
643 let latest = rt.bus.latest_id(&self.workspace_id, &self.channel_id);
644 let cursor = self
645 .last_seen_event_id
646 .load(std::sync::atomic::Ordering::Relaxed);
647 if cursor > 0 && latest - cursor > K_STALENESS_BOUND {
648 let gap = latest - cursor;
649 result_text = format!(
650 "[CONTEXT STALE] {gap} events happened since your last read. \
651 Use ctx_session(action=\"status\") to sync.\n\n{result_text}"
652 );
653 }
654 self.last_seen_event_id
655 .store(latest, std::sync::atomic::Ordering::Relaxed);
656 }
657 }
658
659 {
660 let input = helpers::canonical_args_string(args);
661 let input_md5 = helpers::hash_fast(&input);
662 let output_md5 = helpers::hash_fast(&result_text);
663 let agent_id = self.agent_id.read().await.clone();
664 let client_name = self.client_name.read().await.clone();
665 let mut explicit_intent: Option<(
666 crate::core::intent_protocol::IntentRecord,
667 Option<String>,
668 String,
669 )> = None;
670
671 let pending_session_save = {
672 let empty_args = serde_json::Map::new();
673 let args_map = args.unwrap_or(&empty_args);
674 let mut session = self.session.write().await;
675 session.record_tool_receipt(
676 name,
677 action.as_deref(),
678 &input_md5,
679 &output_md5,
680 agent_id.as_deref(),
681 Some(&client_name),
682 );
683
684 if let Some(intent) = crate::core::intent_protocol::infer_from_tool_call(
685 name,
686 action.as_deref(),
687 args_map,
688 session.project_root.as_deref(),
689 ) {
690 let is_explicit =
691 intent.source == crate::core::intent_protocol::IntentSource::Explicit;
692 let root = session.project_root.clone();
693 let sid = session.id.clone();
694 session.record_intent(intent.clone());
695 if is_explicit {
696 explicit_intent = Some((intent, root, sid));
697 }
698 }
699 if session.should_save() {
700 session.prepare_save().ok()
701 } else {
702 None
703 }
704 };
705
706 if let Some(prepared) = pending_session_save {
707 tokio::task::spawn_blocking(move || {
708 let _ = prepared.write_to_disk();
709 });
710 }
711
712 if let Some((intent, root, session_id)) = explicit_intent {
713 let _ = crate::core::intent_protocol::apply_side_effects(
714 &intent,
715 root.as_deref(),
716 &session_id,
717 );
718 }
719
720 if self.autonomy.is_enabled() {
721 let (calls, project_root) = {
722 let session = self.session.read().await;
723 (session.stats.total_tool_calls, session.project_root.clone())
724 };
725
726 if let Some(root) = project_root {
727 if crate::tools::autonomy::should_auto_consolidate(&self.autonomy, calls) {
728 let root_clone = root.clone();
729 tokio::task::spawn_blocking(move || {
730 let _ = crate::core::consolidation_engine::consolidate_latest(
731 &root_clone,
732 crate::core::consolidation_engine::ConsolidationBudgets::default(),
733 );
734 });
735 }
736 }
737 }
738
739 let agent_key = agent_id.unwrap_or_else(|| "unknown".to_string());
740 let input_token_count = crate::core::tokens::count_tokens(&input) as u64;
741 let output_token_count_u64 = output_token_count as u64;
742 let name_owned = name.to_string();
743 tokio::task::spawn_blocking(move || {
744 let pricing = crate::core::gain::model_pricing::ModelPricing::load();
745 let quote = pricing.quote_from_env_or_agent_type(&client_name);
746 let cost_usd =
747 quote
748 .cost
749 .estimate_usd(input_token_count, output_token_count_u64, 0, 0);
750 crate::core::budget_tracker::BudgetTracker::global().record_cost_usd(cost_usd);
751
752 let mut store = crate::core::a2a::cost_attribution::CostStore::load();
753 store.record_tool_call(
754 &agent_key,
755 &client_name,
756 &name_owned,
757 input_token_count,
758 output_token_count_u64,
759 0,
760 );
761 let _ = store.save();
762 });
763 }
764
765 if self.session_mode == crate::tools::SessionMode::Shared
767 && name == "ctx_knowledge"
768 && action.as_deref() == Some("remember")
769 {
770 if let Some(ref rt) = self.context_os {
771 let my_agent = self.agent_id.read().await.clone();
772 let category = helpers::get_str(args, "category");
773 let key = helpers::get_str(args, "key");
774 if let (Some(ref cat), Some(ref k)) = (&category, &key) {
775 let recent = rt.bus.recent_by_kind(
776 &self.workspace_id,
777 &self.channel_id,
778 "knowledge_remembered",
779 20,
780 );
781 for ev in &recent {
782 let p = &ev.payload;
783 let ev_cat = p.get("category").and_then(|v| v.as_str());
784 let ev_key = p.get("key").and_then(|v| v.as_str());
785 let ev_actor = ev.actor.as_deref();
786 if ev_cat == Some(cat.as_str())
787 && ev_key == Some(k.as_str())
788 && ev_actor != my_agent.as_deref()
789 {
790 let other = ev_actor.unwrap_or("unknown");
791 result_text = format!(
792 "[CONFLICT] Agent '{other}' recently wrote to the same knowledge key \
793 '{cat}/{k}'. Review before proceeding.\n\n{result_text}"
794 );
795 break;
796 }
797 }
798 }
799 }
800 }
801
802 if self.session_mode == crate::tools::SessionMode::Shared {
804 let ws = self.workspace_id.clone();
805 let ch = self.channel_id.clone();
806 let rt = self.context_os.clone();
807 let agent = self.agent_id.read().await.clone();
808 let tool = name.to_string();
809 let tool_action = action.clone();
810 let tool_path = helpers::get_str(args, "path");
811 let tool_category = helpers::get_str(args, "category");
812 let tool_key = helpers::get_str(args, "key");
813 let session_snapshot = self.session.read().await.clone();
814 let session_task = session_snapshot.task.clone();
815 tokio::task::spawn_blocking(move || {
816 let Some(rt) = rt else {
817 return;
818 };
819 let Some(root) = session_snapshot.project_root.as_deref() else {
820 return;
821 };
822 rt.shared_sessions
823 .persist_best_effort(root, &ws, &ch, &session_snapshot);
824 rt.metrics.record_session_persisted();
825
826 let mut base_payload = serde_json::json!({
827 "tool": tool,
828 "action": tool_action,
829 });
830 if let Some(ref p) = tool_path {
831 base_payload["path"] = serde_json::Value::String(p.clone());
832 }
833 if let Some(ref c) = tool_category {
834 base_payload["category"] = serde_json::Value::String(c.clone());
835 }
836 if let Some(ref k) = tool_key {
837 base_payload["key"] = serde_json::Value::String(k.clone());
838 }
839 if let Some(ref t) = session_task {
840 base_payload["reasoning"] = serde_json::Value::String(t.description.clone());
841 }
842
843 if rt
844 .bus
845 .append(
846 &ws,
847 &ch,
848 &crate::core::context_os::ContextEventKindV1::ToolCallRecorded,
849 agent.as_deref(),
850 base_payload.clone(),
851 )
852 .is_some()
853 {
854 rt.metrics.record_event_appended();
855 rt.metrics.record_event_broadcast();
856 }
857
858 if let Some(secondary) =
859 crate::core::context_os::secondary_event_kind(&tool, tool_action.as_deref())
860 {
861 if rt
862 .bus
863 .append(&ws, &ch, &secondary, agent.as_deref(), base_payload)
864 .is_some()
865 {
866 rt.metrics.record_event_appended();
867 rt.metrics.record_event_broadcast();
868 }
869 }
870 });
871 }
872
873 let skip_checkpoint = minimal
874 || matches!(
875 name,
876 "ctx_compress"
877 | "ctx_metrics"
878 | "ctx_benchmark"
879 | "ctx_analyze"
880 | "ctx_cache"
881 | "ctx_discover"
882 | "ctx_dedup"
883 | "ctx_session"
884 | "ctx_knowledge"
885 | "ctx_agent"
886 | "ctx_share"
887 | "ctx_wrapped"
888 | "ctx_overview"
889 | "ctx_preload"
890 | "ctx_cost"
891 | "ctx_gain"
892 | "ctx_heatmap"
893 | "ctx_task"
894 | "ctx_impact"
895 | "ctx_architecture"
896 | "ctx_smells"
897 | "ctx_workflow"
898 );
899
900 if !skip_checkpoint && self.increment_and_check() {
901 if let Some(checkpoint) = self.auto_checkpoint().await {
902 let interval = LeanCtxServer::checkpoint_interval_effective();
903 let combined = format!(
904 "{result_text}\n\n--- AUTO CHECKPOINT (every {interval} calls) ---\n{checkpoint}"
905 );
906 return Ok(CallToolResult::success(vec![Content::text(combined)]));
907 }
908 }
909
910 let tool_duration_ms = tool_start.elapsed().as_millis() as u64;
911 if tool_duration_ms > 100 {
912 LeanCtxServer::append_tool_call_log(
913 name,
914 tool_duration_ms,
915 0,
916 0,
917 None,
918 &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
919 );
920 }
921
922 let current_count = self.call_count.load(std::sync::atomic::Ordering::Relaxed);
923 if current_count > 0 && current_count.is_multiple_of(100) {
924 std::thread::spawn(crate::cloud_sync::cloud_background_tasks);
925 }
926
927 Ok(CallToolResult::success(vec![Content::text(result_text)]))
928 }
929}
930
931pub fn build_instructions_for_test(crp_mode: CrpMode) -> String {
932 crate::instructions::build_instructions_for_test(crp_mode)
933}
934
935pub fn build_claude_code_instructions_for_test() -> String {
936 crate::instructions::claude_code_instructions()
937}
938
939const PROJECT_MARKERS: &[&str] = &[
940 ".git",
941 "Cargo.toml",
942 "package.json",
943 "go.mod",
944 "pyproject.toml",
945 "setup.py",
946 "pom.xml",
947 "build.gradle",
948 "Makefile",
949 ".lean-ctx.toml",
950];
951
952fn has_project_marker(dir: &std::path::Path) -> bool {
953 PROJECT_MARKERS.iter().any(|m| dir.join(m).exists())
954}
955
956fn is_home_or_agent_dir(dir: &std::path::Path) -> bool {
957 if let Some(home) = dirs::home_dir() {
958 if dir == home {
959 return true;
960 }
961 }
962 let dir_str = dir.to_string_lossy();
963 dir_str.ends_with("/.claude")
964 || dir_str.ends_with("/.codex")
965 || dir_str.contains("/.claude/")
966 || dir_str.contains("/.codex/")
967}
968
969fn git_toplevel_from(dir: &std::path::Path) -> Option<String> {
970 std::process::Command::new("git")
971 .args(["rev-parse", "--show-toplevel"])
972 .current_dir(dir)
973 .stdout(std::process::Stdio::piped())
974 .stderr(std::process::Stdio::null())
975 .output()
976 .ok()
977 .and_then(|o| {
978 if o.status.success() {
979 String::from_utf8(o.stdout)
980 .ok()
981 .map(|s| s.trim().to_string())
982 } else {
983 None
984 }
985 })
986}
987
988pub fn derive_project_root_from_cwd() -> Option<String> {
989 let cwd = std::env::current_dir().ok()?;
990 let canonical = crate::core::pathutil::safe_canonicalize_or_self(&cwd);
991
992 if is_home_or_agent_dir(&canonical) {
993 return git_toplevel_from(&canonical);
994 }
995
996 if has_project_marker(&canonical) {
997 return Some(canonical.to_string_lossy().to_string());
998 }
999
1000 if let Some(git_root) = git_toplevel_from(&canonical) {
1001 return Some(git_root);
1002 }
1003
1004 if let Some(root) = detect_multi_root_workspace(&canonical) {
1005 return Some(root);
1006 }
1007
1008 None
1009}
1010
1011fn detect_multi_root_workspace(dir: &std::path::Path) -> Option<String> {
1015 let entries = std::fs::read_dir(dir).ok()?;
1016 let mut child_projects: Vec<String> = Vec::new();
1017
1018 for entry in entries.flatten() {
1019 let path = entry.path();
1020 if path.is_dir() && has_project_marker(&path) {
1021 let canonical = crate::core::pathutil::safe_canonicalize_or_self(&path);
1022 child_projects.push(canonical.to_string_lossy().to_string());
1023 }
1024 }
1025
1026 if child_projects.len() >= 2 {
1027 let existing = std::env::var("LEAN_CTX_ALLOW_PATH").unwrap_or_default();
1028 let sep = if cfg!(windows) { ";" } else { ":" };
1029 let merged = if existing.is_empty() {
1030 child_projects.join(sep)
1031 } else {
1032 format!("{existing}{sep}{}", child_projects.join(sep))
1033 };
1034 std::env::set_var("LEAN_CTX_ALLOW_PATH", &merged);
1035 tracing::info!(
1036 "Multi-root workspace detected at {}: auto-allowing {} child projects",
1037 dir.display(),
1038 child_projects.len()
1039 );
1040 return Some(dir.to_string_lossy().to_string());
1041 }
1042
1043 None
1044}
1045
1046pub fn tool_descriptions_for_test() -> Vec<(&'static str, &'static str)> {
1047 crate::tool_defs::list_all_tool_defs()
1048 .into_iter()
1049 .map(|(name, desc, _)| (name, desc))
1050 .collect()
1051}
1052
1053pub fn tool_schemas_json_for_test() -> String {
1054 crate::tool_defs::list_all_tool_defs()
1055 .iter()
1056 .map(|(name, _, schema)| format!("{name}: {schema}"))
1057 .collect::<Vec<_>>()
1058 .join("\n")
1059}
1060
1061fn is_shell_tool_name(name: &str) -> bool {
1062 matches!(name, "ctx_shell" | "ctx_execute")
1063}
1064
1065#[cfg(test)]
1066mod tests {
1067 use super::*;
1068
1069 #[test]
1070 fn project_markers_detected() {
1071 let tmp = tempfile::tempdir().unwrap();
1072 let root = tmp.path().join("myproject");
1073 std::fs::create_dir_all(&root).unwrap();
1074 assert!(!has_project_marker(&root));
1075
1076 std::fs::create_dir(root.join(".git")).unwrap();
1077 assert!(has_project_marker(&root));
1078 }
1079
1080 #[test]
1081 fn home_dir_detected_as_agent_dir() {
1082 if let Some(home) = dirs::home_dir() {
1083 assert!(is_home_or_agent_dir(&home));
1084 }
1085 }
1086
1087 #[test]
1088 fn agent_dirs_detected() {
1089 let claude = std::path::PathBuf::from("/home/user/.claude");
1090 assert!(is_home_or_agent_dir(&claude));
1091 let codex = std::path::PathBuf::from("/home/user/.codex");
1092 assert!(is_home_or_agent_dir(&codex));
1093 let project = std::path::PathBuf::from("/home/user/projects/myapp");
1094 assert!(!is_home_or_agent_dir(&project));
1095 }
1096
1097 #[test]
1098 fn test_unified_tool_count() {
1099 let tools = crate::tool_defs::unified_tool_defs();
1100 assert_eq!(tools.len(), 5, "Expected 5 unified tools");
1101 }
1102
1103 #[test]
1104 fn test_granular_tool_count() {
1105 let tools = crate::tool_defs::granular_tool_defs();
1106 assert!(tools.len() >= 25, "Expected at least 25 granular tools");
1107 }
1108
1109 #[test]
1110 fn disabled_tools_filters_list() {
1111 let all = crate::tool_defs::granular_tool_defs();
1112 let total = all.len();
1113 let disabled = ["ctx_graph".to_string(), "ctx_agent".to_string()];
1114 let filtered: Vec<_> = all
1115 .into_iter()
1116 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
1117 .collect();
1118 assert_eq!(filtered.len(), total - 2);
1119 assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_graph"));
1120 assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_agent"));
1121 }
1122
1123 #[test]
1124 fn empty_disabled_tools_returns_all() {
1125 let all = crate::tool_defs::granular_tool_defs();
1126 let total = all.len();
1127 let disabled: Vec<String> = vec![];
1128 let filtered: Vec<_> = all
1129 .into_iter()
1130 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
1131 .collect();
1132 assert_eq!(filtered.len(), total);
1133 }
1134
1135 #[test]
1136 fn misspelled_disabled_tool_is_silently_ignored() {
1137 let all = crate::tool_defs::granular_tool_defs();
1138 let total = all.len();
1139 let disabled = ["ctx_nonexistent_tool".to_string()];
1140 let filtered: Vec<_> = all
1141 .into_iter()
1142 .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
1143 .collect();
1144 assert_eq!(filtered.len(), total);
1145 }
1146
1147 #[test]
1148 fn detect_multi_root_workspace_with_child_projects() {
1149 let tmp = tempfile::tempdir().unwrap();
1150 let workspace = tmp.path().join("workspace");
1151 std::fs::create_dir_all(&workspace).unwrap();
1152
1153 let proj_a = workspace.join("project-a");
1154 let proj_b = workspace.join("project-b");
1155 std::fs::create_dir_all(proj_a.join(".git")).unwrap();
1156 std::fs::create_dir_all(&proj_b).unwrap();
1157 std::fs::write(proj_b.join("package.json"), "{}").unwrap();
1158
1159 let result = detect_multi_root_workspace(&workspace);
1160 assert!(
1161 result.is_some(),
1162 "should detect workspace with 2 child projects"
1163 );
1164
1165 std::env::remove_var("LEAN_CTX_ALLOW_PATH");
1166 }
1167
1168 #[test]
1169 fn detect_multi_root_workspace_returns_none_for_single_project() {
1170 let tmp = tempfile::tempdir().unwrap();
1171 let workspace = tmp.path().join("workspace");
1172 std::fs::create_dir_all(&workspace).unwrap();
1173
1174 let proj_a = workspace.join("project-a");
1175 std::fs::create_dir_all(proj_a.join(".git")).unwrap();
1176
1177 let result = detect_multi_root_workspace(&workspace);
1178 assert!(
1179 result.is_none(),
1180 "should not detect workspace with only 1 child project"
1181 );
1182 }
1183}