Skip to main content

lean_ctx/server/
mod.rs

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            // Avoid startup stampedes when multiple agent sessions initialize at once.
98            // These are best-effort maintenance tasks; it's fine to skip if another
99            // lean-ctx instance is already doing them.
100            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 pre_terse_len = result_text.len();
436        let output_tokens = {
437            let tokens = crate::core::tokens::count_tokens(&result_text) as u64;
438            crate::core::budget_tracker::BudgetTracker::global().record_tokens(tokens);
439            tokens
440        };
441
442        crate::core::anomaly::record_metric("tokens_per_call", output_tokens as f64);
443        // Persist anomaly detector so dashboard state survives restarts.
444        crate::core::anomaly::save();
445
446        let budget_warning = {
447            use crate::core::budget_tracker::{BudgetLevel, BudgetTracker};
448            let snap = BudgetTracker::global().check();
449            if *snap.worst_level() == BudgetLevel::Warning {
450                for (dim, lvl, used, limit, pct) in [
451                    (
452                        "tokens",
453                        &snap.tokens.level,
454                        format!("{}", snap.tokens.used),
455                        format!("{}", snap.tokens.limit),
456                        snap.tokens.percent,
457                    ),
458                    (
459                        "shell",
460                        &snap.shell.level,
461                        format!("{}", snap.shell.used),
462                        format!("{}", snap.shell.limit),
463                        snap.shell.percent,
464                    ),
465                    (
466                        "cost",
467                        &snap.cost.level,
468                        format!("${:.2}", snap.cost.used_usd),
469                        format!("${:.2}", snap.cost.limit_usd),
470                        snap.cost.percent,
471                    ),
472                ] {
473                    if *lvl == BudgetLevel::Warning {
474                        crate::core::events::emit_budget_warning(
475                            &snap.role, dim, &used, &limit, pct,
476                        );
477                    }
478                }
479                Some(format!("[BUDGET WARNING] {}", snap.format_compact()))
480            } else {
481                None
482            }
483        };
484
485        let archive_hint = if minimal {
486            None
487        } else {
488            use crate::core::archive;
489            let archivable = matches!(
490                name,
491                "ctx_shell"
492                    | "ctx_read"
493                    | "ctx_multi_read"
494                    | "ctx_smart_read"
495                    | "ctx_execute"
496                    | "ctx_search"
497                    | "ctx_tree"
498            );
499            if archivable && archive::should_archive(&result_text) {
500                let cmd = helpers::get_str(args, "command")
501                    .or_else(|| helpers::get_str(args, "path"))
502                    .unwrap_or_default();
503                let session_id = self.session.read().await.id.clone();
504                let to_store = crate::core::redaction::redact_text_if_enabled(&result_text);
505                let tokens = crate::core::tokens::count_tokens(&to_store);
506                archive::store(name, &cmd, &to_store, Some(&session_id))
507                    .map(|id| archive::format_hint(&id, to_store.len(), tokens))
508            } else {
509                None
510            }
511        };
512
513        let pre_compression = result_text.clone();
514        let skip_terse = name == "ctx_shell"
515            && helpers::get_str(args, "command")
516                .is_some_and(|c| crate::shell::compress::has_structural_output(&c));
517        let compression = crate::core::config::CompressionLevel::effective(&config);
518        if compression.is_active() && !skip_terse {
519            let terse_result =
520                crate::core::terse::pipeline::compress(&result_text, &compression, None);
521            if terse_result.quality_passed && terse_result.savings_pct >= 1.0 {
522                result_text = terse_result.output;
523            }
524        }
525
526        {
527            let verify_cfg = crate::core::profiles::active_profile().verification;
528            let vr = crate::core::output_verification::verify_output(
529                &pre_compression,
530                &result_text,
531                &verify_cfg,
532            );
533            if !vr.warnings.is_empty() {
534                let msg = format!("[VERIFY] {}", vr.format_compact());
535                result_text = format!("{result_text}\n\n{msg}");
536            }
537        }
538
539        if let Some(hint) = archive_hint {
540            result_text = format!("{result_text}\n{hint}");
541        }
542
543        if let Some(ctx) = auto_context {
544            result_text = format!("{ctx}\n\n{result_text}");
545        }
546
547        if let Some(warning) = throttle_warning {
548            result_text = format!("{result_text}\n\n{warning}");
549        }
550
551        if let Some(bw) = budget_warning {
552            result_text = format!("{result_text}\n\n{bw}");
553        }
554
555        if !self
556            .rules_stale_checked
557            .swap(true, std::sync::atomic::Ordering::Relaxed)
558        {
559            let client = self.client_name.read().await.clone();
560            if !client.is_empty() {
561                if let Some(stale_msg) = crate::rules_inject::check_rules_freshness(&client) {
562                    result_text = format!("{result_text}\n\n{stale_msg}");
563                }
564            }
565        }
566
567        {
568            // Evaluate SLOs for observability (watch/dashboard), but keep tool outputs clean.
569            let _ = crate::core::slo::evaluate();
570        }
571
572        if name == "ctx_read" {
573            if minimal {
574                let mut cache = self.cache.write().await;
575                crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache, name);
576            } else {
577                let read_path = self
578                    .resolve_path_or_passthrough(
579                        &helpers::get_str(args, "path").unwrap_or_default(),
580                    )
581                    .await;
582                let project_root = {
583                    let session = self.session.read().await;
584                    session.project_root.clone()
585                };
586                let mut cache = self.cache.write().await;
587                let enrich = crate::tools::autonomy::enrich_after_read(
588                    &self.autonomy,
589                    &mut cache,
590                    &read_path,
591                    project_root.as_deref(),
592                    None,
593                    crate::tools::CrpMode::effective(),
594                    false,
595                );
596                if let Some(hint) = enrich.related_hint {
597                    result_text = format!("{result_text}\n{hint}");
598                }
599                crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache, name);
600            }
601        }
602
603        if !minimal && name == "ctx_shell" {
604            let cmd = helpers::get_str(args, "command").unwrap_or_default();
605            let calls = self.tool_calls.read().await;
606            let last_original = calls.last().map_or(0, |c| c.original_tokens);
607            drop(calls);
608            let pre_hint_tokens = crate::core::tokens::count_tokens(&result_text);
609            if let Some(hint) = crate::tools::autonomy::shell_efficiency_hint(
610                &self.autonomy,
611                &cmd,
612                last_original,
613                pre_hint_tokens,
614            ) {
615                result_text = format!("{result_text}\n{hint}");
616            }
617        }
618
619        #[allow(clippy::cast_possible_truncation)]
620        let output_token_count = if result_text.len() == pre_terse_len {
621            output_tokens as usize
622        } else {
623            crate::core::tokens::count_tokens(&result_text)
624        };
625        let action = helpers::get_str(args, "action");
626
627        // K-bounded staleness guard: warn if shared context has diverged.
628        const K_STALENESS_BOUND: i64 = 10;
629        if self.session_mode == crate::tools::SessionMode::Shared {
630            if let Some(ref rt) = self.context_os {
631                let latest = rt.bus.latest_id(&self.workspace_id, &self.channel_id);
632                let cursor = self
633                    .last_seen_event_id
634                    .load(std::sync::atomic::Ordering::Relaxed);
635                if cursor > 0 && latest - cursor > K_STALENESS_BOUND {
636                    let gap = latest - cursor;
637                    result_text = format!(
638                        "[CONTEXT STALE] {gap} events happened since your last read. \
639                         Use ctx_session(action=\"status\") to sync.\n\n{result_text}"
640                    );
641                }
642                self.last_seen_event_id
643                    .store(latest, std::sync::atomic::Ordering::Relaxed);
644            }
645        }
646
647        {
648            let input = helpers::canonical_args_string(args);
649            let input_md5 = helpers::hash_fast(&input);
650            let output_md5 = helpers::hash_fast(&result_text);
651            let agent_id = self.agent_id.read().await.clone();
652            let client_name = self.client_name.read().await.clone();
653            let mut explicit_intent: Option<(
654                crate::core::intent_protocol::IntentRecord,
655                Option<String>,
656                String,
657            )> = None;
658
659            let pending_session_save = {
660                let empty_args = serde_json::Map::new();
661                let args_map = args.unwrap_or(&empty_args);
662                let mut session = self.session.write().await;
663                session.record_tool_receipt(
664                    name,
665                    action.as_deref(),
666                    &input_md5,
667                    &output_md5,
668                    agent_id.as_deref(),
669                    Some(&client_name),
670                );
671
672                if let Some(intent) = crate::core::intent_protocol::infer_from_tool_call(
673                    name,
674                    action.as_deref(),
675                    args_map,
676                    session.project_root.as_deref(),
677                ) {
678                    let is_explicit =
679                        intent.source == crate::core::intent_protocol::IntentSource::Explicit;
680                    let root = session.project_root.clone();
681                    let sid = session.id.clone();
682                    session.record_intent(intent.clone());
683                    if is_explicit {
684                        explicit_intent = Some((intent, root, sid));
685                    }
686                }
687                if session.should_save() {
688                    session.prepare_save().ok()
689                } else {
690                    None
691                }
692            };
693
694            if let Some(prepared) = pending_session_save {
695                tokio::task::spawn_blocking(move || {
696                    let _ = prepared.write_to_disk();
697                });
698            }
699
700            if let Some((intent, root, session_id)) = explicit_intent {
701                let _ = crate::core::intent_protocol::apply_side_effects(
702                    &intent,
703                    root.as_deref(),
704                    &session_id,
705                );
706            }
707
708            if self.autonomy.is_enabled() {
709                let (calls, project_root) = {
710                    let session = self.session.read().await;
711                    (session.stats.total_tool_calls, session.project_root.clone())
712                };
713
714                if let Some(root) = project_root {
715                    if crate::tools::autonomy::should_auto_consolidate(&self.autonomy, calls) {
716                        let root_clone = root.clone();
717                        tokio::task::spawn_blocking(move || {
718                            let _ = crate::core::consolidation_engine::consolidate_latest(
719                                &root_clone,
720                                crate::core::consolidation_engine::ConsolidationBudgets::default(),
721                            );
722                        });
723                    }
724                }
725            }
726
727            let agent_key = agent_id.unwrap_or_else(|| "unknown".to_string());
728            let input_token_count = crate::core::tokens::count_tokens(&input) as u64;
729            let output_token_count_u64 = output_token_count as u64;
730            let name_owned = name.to_string();
731            tokio::task::spawn_blocking(move || {
732                let pricing = crate::core::gain::model_pricing::ModelPricing::load();
733                let quote = pricing.quote_from_env_or_agent_type(&client_name);
734                let cost_usd =
735                    quote
736                        .cost
737                        .estimate_usd(input_token_count, output_token_count_u64, 0, 0);
738                crate::core::budget_tracker::BudgetTracker::global().record_cost_usd(cost_usd);
739
740                let mut store = crate::core::a2a::cost_attribution::CostStore::load();
741                store.record_tool_call(
742                    &agent_key,
743                    &client_name,
744                    &name_owned,
745                    input_token_count,
746                    output_token_count_u64,
747                    0,
748                );
749                let _ = store.save();
750            });
751        }
752
753        // Context Bus: conflict detection for knowledge writes in shared mode.
754        if self.session_mode == crate::tools::SessionMode::Shared
755            && name == "ctx_knowledge"
756            && action.as_deref() == Some("remember")
757        {
758            if let Some(ref rt) = self.context_os {
759                let my_agent = self.agent_id.read().await.clone();
760                let category = helpers::get_str(args, "category");
761                let key = helpers::get_str(args, "key");
762                if let (Some(ref cat), Some(ref k)) = (&category, &key) {
763                    let recent = rt.bus.recent_by_kind(
764                        &self.workspace_id,
765                        &self.channel_id,
766                        "knowledge_remembered",
767                        20,
768                    );
769                    for ev in &recent {
770                        let p = &ev.payload;
771                        let ev_cat = p.get("category").and_then(|v| v.as_str());
772                        let ev_key = p.get("key").and_then(|v| v.as_str());
773                        let ev_actor = ev.actor.as_deref();
774                        if ev_cat == Some(cat.as_str())
775                            && ev_key == Some(k.as_str())
776                            && ev_actor != my_agent.as_deref()
777                        {
778                            let other = ev_actor.unwrap_or("unknown");
779                            result_text = format!(
780                                "[CONFLICT] Agent '{other}' recently wrote to the same knowledge key \
781                                 '{cat}/{k}'. Review before proceeding.\n\n{result_text}"
782                            );
783                            break;
784                        }
785                    }
786                }
787            }
788        }
789
790        // Context OS: persist shared session + publish events.
791        if self.session_mode == crate::tools::SessionMode::Shared {
792            let ws = self.workspace_id.clone();
793            let ch = self.channel_id.clone();
794            let rt = self.context_os.clone();
795            let agent = self.agent_id.read().await.clone();
796            let tool = name.to_string();
797            let tool_action = action.clone();
798            let tool_path = helpers::get_str(args, "path");
799            let tool_category = helpers::get_str(args, "category");
800            let tool_key = helpers::get_str(args, "key");
801            let session_snapshot = self.session.read().await.clone();
802            let session_task = session_snapshot.task.clone();
803            tokio::task::spawn_blocking(move || {
804                let Some(rt) = rt else {
805                    return;
806                };
807                let Some(root) = session_snapshot.project_root.as_deref() else {
808                    return;
809                };
810                rt.shared_sessions
811                    .persist_best_effort(root, &ws, &ch, &session_snapshot);
812                rt.metrics.record_session_persisted();
813
814                let mut base_payload = serde_json::json!({
815                    "tool": tool,
816                    "action": tool_action,
817                });
818                if let Some(ref p) = tool_path {
819                    base_payload["path"] = serde_json::Value::String(p.clone());
820                }
821                if let Some(ref c) = tool_category {
822                    base_payload["category"] = serde_json::Value::String(c.clone());
823                }
824                if let Some(ref k) = tool_key {
825                    base_payload["key"] = serde_json::Value::String(k.clone());
826                }
827                if let Some(ref t) = session_task {
828                    base_payload["reasoning"] = serde_json::Value::String(t.description.clone());
829                }
830
831                if rt
832                    .bus
833                    .append(
834                        &ws,
835                        &ch,
836                        &crate::core::context_os::ContextEventKindV1::ToolCallRecorded,
837                        agent.as_deref(),
838                        base_payload.clone(),
839                    )
840                    .is_some()
841                {
842                    rt.metrics.record_event_appended();
843                    rt.metrics.record_event_broadcast();
844                }
845
846                if let Some(secondary) =
847                    crate::core::context_os::secondary_event_kind(&tool, tool_action.as_deref())
848                {
849                    if rt
850                        .bus
851                        .append(&ws, &ch, &secondary, agent.as_deref(), base_payload)
852                        .is_some()
853                    {
854                        rt.metrics.record_event_appended();
855                        rt.metrics.record_event_broadcast();
856                    }
857                }
858            });
859        }
860
861        let skip_checkpoint = minimal
862            || matches!(
863                name,
864                "ctx_compress"
865                    | "ctx_metrics"
866                    | "ctx_benchmark"
867                    | "ctx_analyze"
868                    | "ctx_cache"
869                    | "ctx_discover"
870                    | "ctx_dedup"
871                    | "ctx_session"
872                    | "ctx_knowledge"
873                    | "ctx_agent"
874                    | "ctx_share"
875                    | "ctx_wrapped"
876                    | "ctx_overview"
877                    | "ctx_preload"
878                    | "ctx_cost"
879                    | "ctx_gain"
880                    | "ctx_heatmap"
881                    | "ctx_task"
882                    | "ctx_impact"
883                    | "ctx_architecture"
884                    | "ctx_smells"
885                    | "ctx_workflow"
886            );
887
888        if !skip_checkpoint && self.increment_and_check() {
889            if let Some(checkpoint) = self.auto_checkpoint().await {
890                let interval = LeanCtxServer::checkpoint_interval_effective();
891                let combined = format!(
892                    "{result_text}\n\n--- AUTO CHECKPOINT (every {interval} calls) ---\n{checkpoint}"
893                );
894                return Ok(CallToolResult::success(vec![Content::text(combined)]));
895            }
896        }
897
898        let tool_duration_ms = tool_start.elapsed().as_millis() as u64;
899        if tool_duration_ms > 100 {
900            LeanCtxServer::append_tool_call_log(
901                name,
902                tool_duration_ms,
903                0,
904                0,
905                None,
906                &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
907            );
908        }
909
910        let current_count = self.call_count.load(std::sync::atomic::Ordering::Relaxed);
911        if current_count > 0 && current_count.is_multiple_of(100) {
912            std::thread::spawn(crate::cloud_sync::cloud_background_tasks);
913        }
914
915        Ok(CallToolResult::success(vec![Content::text(result_text)]))
916    }
917}
918
919pub fn build_instructions_for_test(crp_mode: CrpMode) -> String {
920    crate::instructions::build_instructions_for_test(crp_mode)
921}
922
923pub fn build_claude_code_instructions_for_test() -> String {
924    crate::instructions::claude_code_instructions()
925}
926
927const PROJECT_MARKERS: &[&str] = &[
928    ".git",
929    "Cargo.toml",
930    "package.json",
931    "go.mod",
932    "pyproject.toml",
933    "setup.py",
934    "pom.xml",
935    "build.gradle",
936    "Makefile",
937    ".lean-ctx.toml",
938];
939
940fn has_project_marker(dir: &std::path::Path) -> bool {
941    PROJECT_MARKERS.iter().any(|m| dir.join(m).exists())
942}
943
944fn is_home_or_agent_dir(dir: &std::path::Path) -> bool {
945    if let Some(home) = dirs::home_dir() {
946        if dir == home {
947            return true;
948        }
949    }
950    let dir_str = dir.to_string_lossy();
951    dir_str.ends_with("/.claude")
952        || dir_str.ends_with("/.codex")
953        || dir_str.contains("/.claude/")
954        || dir_str.contains("/.codex/")
955}
956
957fn git_toplevel_from(dir: &std::path::Path) -> Option<String> {
958    std::process::Command::new("git")
959        .args(["rev-parse", "--show-toplevel"])
960        .current_dir(dir)
961        .stdout(std::process::Stdio::piped())
962        .stderr(std::process::Stdio::null())
963        .output()
964        .ok()
965        .and_then(|o| {
966            if o.status.success() {
967                String::from_utf8(o.stdout)
968                    .ok()
969                    .map(|s| s.trim().to_string())
970            } else {
971                None
972            }
973        })
974}
975
976pub fn derive_project_root_from_cwd() -> Option<String> {
977    let cwd = std::env::current_dir().ok()?;
978    let canonical = crate::core::pathutil::safe_canonicalize_or_self(&cwd);
979
980    if is_home_or_agent_dir(&canonical) {
981        return git_toplevel_from(&canonical);
982    }
983
984    if has_project_marker(&canonical) {
985        return Some(canonical.to_string_lossy().to_string());
986    }
987
988    if let Some(git_root) = git_toplevel_from(&canonical) {
989        return Some(git_root);
990    }
991
992    if let Some(root) = detect_multi_root_workspace(&canonical) {
993        return Some(root);
994    }
995
996    None
997}
998
999/// Detect a multi-root workspace: a directory that has no project markers
1000/// itself, but contains child directories that do. In this case, use the
1001/// parent as jail root and auto-allow all child projects via LEAN_CTX_ALLOW_PATH.
1002fn detect_multi_root_workspace(dir: &std::path::Path) -> Option<String> {
1003    let entries = std::fs::read_dir(dir).ok()?;
1004    let mut child_projects: Vec<String> = Vec::new();
1005
1006    for entry in entries.flatten() {
1007        let path = entry.path();
1008        if path.is_dir() && has_project_marker(&path) {
1009            let canonical = crate::core::pathutil::safe_canonicalize_or_self(&path);
1010            child_projects.push(canonical.to_string_lossy().to_string());
1011        }
1012    }
1013
1014    if child_projects.len() >= 2 {
1015        let existing = std::env::var("LEAN_CTX_ALLOW_PATH").unwrap_or_default();
1016        let sep = if cfg!(windows) { ";" } else { ":" };
1017        let merged = if existing.is_empty() {
1018            child_projects.join(sep)
1019        } else {
1020            format!("{existing}{sep}{}", child_projects.join(sep))
1021        };
1022        std::env::set_var("LEAN_CTX_ALLOW_PATH", &merged);
1023        tracing::info!(
1024            "Multi-root workspace detected at {}: auto-allowing {} child projects",
1025            dir.display(),
1026            child_projects.len()
1027        );
1028        return Some(dir.to_string_lossy().to_string());
1029    }
1030
1031    None
1032}
1033
1034pub fn tool_descriptions_for_test() -> Vec<(&'static str, &'static str)> {
1035    crate::tool_defs::list_all_tool_defs()
1036        .into_iter()
1037        .map(|(name, desc, _)| (name, desc))
1038        .collect()
1039}
1040
1041pub fn tool_schemas_json_for_test() -> String {
1042    crate::tool_defs::list_all_tool_defs()
1043        .iter()
1044        .map(|(name, _, schema)| format!("{name}: {schema}"))
1045        .collect::<Vec<_>>()
1046        .join("\n")
1047}
1048
1049fn is_shell_tool_name(name: &str) -> bool {
1050    matches!(name, "ctx_shell" | "ctx_execute")
1051}
1052
1053#[cfg(test)]
1054mod tests {
1055    use super::*;
1056
1057    #[test]
1058    fn project_markers_detected() {
1059        let tmp = tempfile::tempdir().unwrap();
1060        let root = tmp.path().join("myproject");
1061        std::fs::create_dir_all(&root).unwrap();
1062        assert!(!has_project_marker(&root));
1063
1064        std::fs::create_dir(root.join(".git")).unwrap();
1065        assert!(has_project_marker(&root));
1066    }
1067
1068    #[test]
1069    fn home_dir_detected_as_agent_dir() {
1070        if let Some(home) = dirs::home_dir() {
1071            assert!(is_home_or_agent_dir(&home));
1072        }
1073    }
1074
1075    #[test]
1076    fn agent_dirs_detected() {
1077        let claude = std::path::PathBuf::from("/home/user/.claude");
1078        assert!(is_home_or_agent_dir(&claude));
1079        let codex = std::path::PathBuf::from("/home/user/.codex");
1080        assert!(is_home_or_agent_dir(&codex));
1081        let project = std::path::PathBuf::from("/home/user/projects/myapp");
1082        assert!(!is_home_or_agent_dir(&project));
1083    }
1084
1085    #[test]
1086    fn test_unified_tool_count() {
1087        let tools = crate::tool_defs::unified_tool_defs();
1088        assert_eq!(tools.len(), 5, "Expected 5 unified tools");
1089    }
1090
1091    #[test]
1092    fn test_granular_tool_count() {
1093        let tools = crate::tool_defs::granular_tool_defs();
1094        assert!(tools.len() >= 25, "Expected at least 25 granular tools");
1095    }
1096
1097    #[test]
1098    fn disabled_tools_filters_list() {
1099        let all = crate::tool_defs::granular_tool_defs();
1100        let total = all.len();
1101        let disabled = ["ctx_graph".to_string(), "ctx_agent".to_string()];
1102        let filtered: Vec<_> = all
1103            .into_iter()
1104            .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
1105            .collect();
1106        assert_eq!(filtered.len(), total - 2);
1107        assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_graph"));
1108        assert!(!filtered.iter().any(|t| t.name.as_ref() == "ctx_agent"));
1109    }
1110
1111    #[test]
1112    fn empty_disabled_tools_returns_all() {
1113        let all = crate::tool_defs::granular_tool_defs();
1114        let total = all.len();
1115        let disabled: Vec<String> = vec![];
1116        let filtered: Vec<_> = all
1117            .into_iter()
1118            .filter(|t| !disabled.iter().any(|d| t.name.as_ref() == d.as_str()))
1119            .collect();
1120        assert_eq!(filtered.len(), total);
1121    }
1122
1123    #[test]
1124    fn misspelled_disabled_tool_is_silently_ignored() {
1125        let all = crate::tool_defs::granular_tool_defs();
1126        let total = all.len();
1127        let disabled = ["ctx_nonexistent_tool".to_string()];
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 detect_multi_root_workspace_with_child_projects() {
1137        let tmp = tempfile::tempdir().unwrap();
1138        let workspace = tmp.path().join("workspace");
1139        std::fs::create_dir_all(&workspace).unwrap();
1140
1141        let proj_a = workspace.join("project-a");
1142        let proj_b = workspace.join("project-b");
1143        std::fs::create_dir_all(proj_a.join(".git")).unwrap();
1144        std::fs::create_dir_all(&proj_b).unwrap();
1145        std::fs::write(proj_b.join("package.json"), "{}").unwrap();
1146
1147        let result = detect_multi_root_workspace(&workspace);
1148        assert!(
1149            result.is_some(),
1150            "should detect workspace with 2 child projects"
1151        );
1152
1153        std::env::remove_var("LEAN_CTX_ALLOW_PATH");
1154    }
1155
1156    #[test]
1157    fn detect_multi_root_workspace_returns_none_for_single_project() {
1158        let tmp = tempfile::tempdir().unwrap();
1159        let workspace = tmp.path().join("workspace");
1160        std::fs::create_dir_all(&workspace).unwrap();
1161
1162        let proj_a = workspace.join("project-a");
1163        std::fs::create_dir_all(proj_a.join(".git")).unwrap();
1164
1165        let result = detect_multi_root_workspace(&workspace);
1166        assert!(
1167            result.is_none(),
1168            "should not detect workspace with only 1 child project"
1169        );
1170    }
1171}