Skip to main content

lean_ctx/
server.rs

1use rmcp::handler::server::ServerHandler;
2use rmcp::model::*;
3use rmcp::service::{RequestContext, RoleServer};
4use rmcp::ErrorData;
5use serde_json::Value;
6
7use crate::tools::{CrpMode, LeanCtxServer};
8
9impl ServerHandler for LeanCtxServer {
10    fn get_info(&self) -> ServerInfo {
11        let capabilities = ServerCapabilities::builder().enable_tools().build();
12
13        let instructions = crate::instructions::build_instructions(self.crp_mode);
14
15        InitializeResult::new(capabilities)
16            .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
17            .with_instructions(instructions)
18    }
19
20    async fn initialize(
21        &self,
22        request: InitializeRequestParams,
23        _context: RequestContext<RoleServer>,
24    ) -> Result<InitializeResult, ErrorData> {
25        let name = request.client_info.name.clone();
26        tracing::info!("MCP client connected: {:?}", name);
27        *self.client_name.write().await = name.clone();
28
29        tokio::task::spawn_blocking(|| {
30            if let Some(home) = dirs::home_dir() {
31                let _ = crate::rules_inject::inject_all_rules(&home);
32            }
33            crate::hooks::refresh_installed_hooks();
34            crate::core::version_check::check_background();
35        });
36
37        let instructions =
38            crate::instructions::build_instructions_with_client(self.crp_mode, &name);
39        let capabilities = ServerCapabilities::builder().enable_tools().build();
40
41        Ok(InitializeResult::new(capabilities)
42            .with_server_info(Implementation::new("lean-ctx", env!("CARGO_PKG_VERSION")))
43            .with_instructions(instructions))
44    }
45
46    async fn list_tools(
47        &self,
48        _request: Option<PaginatedRequestParams>,
49        _context: RequestContext<RoleServer>,
50    ) -> Result<ListToolsResult, ErrorData> {
51        if std::env::var("LEAN_CTX_UNIFIED").is_ok()
52            && std::env::var("LEAN_CTX_FULL_TOOLS").is_err()
53        {
54            return Ok(ListToolsResult {
55                tools: crate::tool_defs::unified_tool_defs(),
56                ..Default::default()
57            });
58        }
59
60        Ok(ListToolsResult {
61            tools: crate::tool_defs::granular_tool_defs(),
62            ..Default::default()
63        })
64    }
65
66    async fn call_tool(
67        &self,
68        request: CallToolRequestParams,
69        _context: RequestContext<RoleServer>,
70    ) -> Result<CallToolResult, ErrorData> {
71        self.check_idle_expiry().await;
72
73        let original_name = request.name.as_ref().to_string();
74        let (resolved_name, resolved_args) = if original_name == "ctx" {
75            let sub = request
76                .arguments
77                .as_ref()
78                .and_then(|a| a.get("tool"))
79                .and_then(|v| v.as_str())
80                .map(|s| s.to_string())
81                .ok_or_else(|| {
82                    ErrorData::invalid_params("'tool' is required for ctx meta-tool", None)
83                })?;
84            let tool_name = if sub.starts_with("ctx_") {
85                sub
86            } else {
87                format!("ctx_{sub}")
88            };
89            let mut args = request.arguments.unwrap_or_default();
90            args.remove("tool");
91            (tool_name, Some(args))
92        } else {
93            (original_name, request.arguments)
94        };
95        let name = resolved_name.as_str();
96        let args = &resolved_args;
97
98        let auto_context = {
99            let task = {
100                let session = self.session.read().await;
101                session.task.as_ref().map(|t| t.description.clone())
102            };
103            let project_root = {
104                let session = self.session.read().await;
105                session.project_root.clone()
106            };
107            let mut cache = self.cache.write().await;
108            crate::tools::autonomy::session_lifecycle_pre_hook(
109                &self.autonomy,
110                name,
111                &mut cache,
112                task.as_deref(),
113                project_root.as_deref(),
114                self.crp_mode,
115            )
116        };
117
118        let tool_start = std::time::Instant::now();
119        let result_text = match name {
120            "ctx_read" => {
121                let path = get_str(args, "path")
122                    .map(|p| crate::hooks::normalize_tool_path(&p))
123                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
124                let current_task = {
125                    let session = self.session.read().await;
126                    session.task.as_ref().map(|t| t.description.clone())
127                };
128                let task_ref = current_task.as_deref();
129                let mut mode = match get_str(args, "mode") {
130                    Some(m) => m,
131                    None => {
132                        let cache = self.cache.read().await;
133                        crate::tools::ctx_smart_read::select_mode_with_task(&cache, &path, task_ref)
134                    }
135                };
136                let fresh = get_bool(args, "fresh").unwrap_or(false);
137                let start_line = get_int(args, "start_line");
138                if let Some(sl) = start_line {
139                    let sl = sl.max(1_i64);
140                    mode = format!("lines:{sl}-999999");
141                }
142                let stale = self.is_prompt_cache_stale().await;
143                let effective_mode = LeanCtxServer::upgrade_mode_if_stale(&mode, stale).to_string();
144                let mut cache = self.cache.write().await;
145                let output = if fresh {
146                    crate::tools::ctx_read::handle_fresh_with_task(
147                        &mut cache,
148                        &path,
149                        &effective_mode,
150                        self.crp_mode,
151                        task_ref,
152                    )
153                } else {
154                    crate::tools::ctx_read::handle_with_task(
155                        &mut cache,
156                        &path,
157                        &effective_mode,
158                        self.crp_mode,
159                        task_ref,
160                    )
161                };
162                let stale_note = if effective_mode != mode {
163                    format!("[cache stale, {mode}→{effective_mode}]\n")
164                } else {
165                    String::new()
166                };
167                let original = cache.get(&path).map_or(0, |e| e.original_tokens);
168                let output_tokens = crate::core::tokens::count_tokens(&output);
169                let saved = original.saturating_sub(output_tokens);
170                let is_cache_hit = output.contains(" cached ");
171                let output = format!("{stale_note}{output}");
172                let file_ref = cache.file_ref_map().get(&path).cloned();
173                drop(cache);
174                {
175                    let mut session = self.session.write().await;
176                    session.touch_file(&path, file_ref.as_deref(), &effective_mode, original);
177                    if is_cache_hit {
178                        session.record_cache_hit();
179                    }
180                    if session.project_root.is_none() {
181                        if let Some(root) = crate::core::protocol::detect_project_root(&path) {
182                            session.project_root = Some(root.clone());
183                            let mut current = self.agent_id.write().await;
184                            if current.is_none() {
185                                let mut registry =
186                                    crate::core::agents::AgentRegistry::load_or_create();
187                                registry.cleanup_stale(24);
188                                let role = std::env::var("LEAN_CTX_AGENT_ROLE").ok();
189                                let id = registry.register("mcp", role.as_deref(), &root);
190                                let _ = registry.save();
191                                *current = Some(id);
192                            }
193                        }
194                    }
195                }
196                self.record_call("ctx_read", original, saved, Some(mode.clone()))
197                    .await;
198                {
199                    let sig =
200                        crate::core::mode_predictor::FileSignature::from_path(&path, original);
201                    let density = if output_tokens > 0 {
202                        original as f64 / output_tokens as f64
203                    } else {
204                        1.0
205                    };
206                    let outcome = crate::core::mode_predictor::ModeOutcome {
207                        mode: mode.clone(),
208                        tokens_in: original,
209                        tokens_out: output_tokens,
210                        density: density.min(1.0),
211                    };
212                    let mut predictor = crate::core::mode_predictor::ModePredictor::new();
213                    predictor.record(sig, outcome);
214                    predictor.save();
215
216                    let ext = std::path::Path::new(&path)
217                        .extension()
218                        .and_then(|e| e.to_str())
219                        .unwrap_or("")
220                        .to_string();
221                    let thresholds = crate::core::adaptive_thresholds::thresholds_for_path(&path);
222                    let cache = self.cache.read().await;
223                    let stats = cache.get_stats();
224                    let feedback_outcome = crate::core::feedback::CompressionOutcome {
225                        session_id: format!("{}", std::process::id()),
226                        language: ext,
227                        entropy_threshold: thresholds.bpe_entropy,
228                        jaccard_threshold: thresholds.jaccard,
229                        total_turns: stats.total_reads as u32,
230                        tokens_saved: saved as u64,
231                        tokens_original: original as u64,
232                        cache_hits: stats.cache_hits as u32,
233                        total_reads: stats.total_reads as u32,
234                        task_completed: true,
235                        timestamp: chrono::Local::now().to_rfc3339(),
236                    };
237                    drop(cache);
238                    let mut store = crate::core::feedback::FeedbackStore::load();
239                    store.record_outcome(feedback_outcome);
240                }
241                output
242            }
243            "ctx_multi_read" => {
244                let paths = get_str_array(args, "paths")
245                    .ok_or_else(|| ErrorData::invalid_params("paths array is required", None))?
246                    .into_iter()
247                    .map(|p| crate::hooks::normalize_tool_path(&p))
248                    .collect::<Vec<_>>();
249                let mode = get_str(args, "mode").unwrap_or_else(|| "full".to_string());
250                let current_task = {
251                    let session = self.session.read().await;
252                    session.task.as_ref().map(|t| t.description.clone())
253                };
254                let mut cache = self.cache.write().await;
255                let output = crate::tools::ctx_multi_read::handle_with_task(
256                    &mut cache,
257                    &paths,
258                    &mode,
259                    self.crp_mode,
260                    current_task.as_deref(),
261                );
262                let mut total_original: usize = 0;
263                for path in &paths {
264                    total_original = total_original
265                        .saturating_add(cache.get(path).map(|e| e.original_tokens).unwrap_or(0));
266                }
267                let tokens = crate::core::tokens::count_tokens(&output);
268                drop(cache);
269                self.record_call(
270                    "ctx_multi_read",
271                    total_original,
272                    total_original.saturating_sub(tokens),
273                    Some(mode),
274                )
275                .await;
276                output
277            }
278            "ctx_tree" => {
279                let path = crate::hooks::normalize_tool_path(
280                    &get_str(args, "path").unwrap_or_else(|| ".".to_string()),
281                );
282                let depth = get_int(args, "depth").unwrap_or(3) as usize;
283                let show_hidden = get_bool(args, "show_hidden").unwrap_or(false);
284                let (result, original) = crate::tools::ctx_tree::handle(&path, depth, show_hidden);
285                let sent = crate::core::tokens::count_tokens(&result);
286                let saved = original.saturating_sub(sent);
287                self.record_call("ctx_tree", original, saved, None).await;
288                let savings_note = if saved > 0 {
289                    format!("\n[saved {saved} tokens vs native ls]")
290                } else {
291                    String::new()
292                };
293                format!("{result}{savings_note}")
294            }
295            "ctx_shell" => {
296                let command = get_str(args, "command")
297                    .ok_or_else(|| ErrorData::invalid_params("command is required", None))?;
298
299                if let Some(rejection) = crate::tools::ctx_shell::validate_command(&command) {
300                    self.record_call("ctx_shell", 0, 0, None).await;
301                    return Ok(CallToolResult::success(vec![Content::text(rejection)]));
302                }
303
304                let raw = get_bool(args, "raw").unwrap_or(false)
305                    || std::env::var("LEAN_CTX_DISABLED").is_ok();
306                let cmd_clone = command.clone();
307                let (output, real_exit_code) =
308                    tokio::task::spawn_blocking(move || execute_command(&cmd_clone))
309                        .await
310                        .unwrap_or_else(|e| (format!("ERROR: shell task failed: {e}"), 1));
311
312                if raw {
313                    let original = crate::core::tokens::count_tokens(&output);
314                    self.record_call("ctx_shell", original, 0, None).await;
315                    output
316                } else {
317                    let result = crate::tools::ctx_shell::handle(&command, &output, self.crp_mode);
318                    let original = crate::core::tokens::count_tokens(&output);
319                    let sent = crate::core::tokens::count_tokens(&result);
320                    let saved = original.saturating_sub(sent);
321                    self.record_call("ctx_shell", original, saved, None).await;
322
323                    let cfg = crate::core::config::Config::load();
324                    let tee_hint = match cfg.tee_mode {
325                        crate::core::config::TeeMode::Always => {
326                            crate::shell::save_tee(&command, &output)
327                                .map(|p| format!("\n[full output: {p}]"))
328                                .unwrap_or_default()
329                        }
330                        crate::core::config::TeeMode::Failures
331                            if !output.trim().is_empty() && output.contains("error")
332                                || output.contains("Error")
333                                || output.contains("ERROR") =>
334                        {
335                            crate::shell::save_tee(&command, &output)
336                                .map(|p| format!("\n[full output: {p}]"))
337                                .unwrap_or_default()
338                        }
339                        _ => String::new(),
340                    };
341
342                    let savings_note = if saved > 0 {
343                        format!("\n[saved {saved} tokens vs native Shell]")
344                    } else {
345                        String::new()
346                    };
347
348                    // Bug Memory: detect errors / resolve pending
349                    {
350                        let sess = self.session.read().await;
351                        let root = sess.project_root.clone();
352                        let sid = sess.id.clone();
353                        let files: Vec<String> = sess
354                            .files_touched
355                            .iter()
356                            .map(|ft| ft.path.clone())
357                            .collect();
358                        drop(sess);
359
360                        if let Some(ref root) = root {
361                            let mut store = crate::core::gotcha_tracker::GotchaStore::load(root);
362
363                            if real_exit_code != 0 {
364                                store.detect_error(&output, &command, real_exit_code, &files, &sid);
365                            } else {
366                                // Success: check if any injected gotchas prevented a repeat
367                                let relevant = store.top_relevant(&files, 7);
368                                let relevant_ids: Vec<String> =
369                                    relevant.iter().map(|g| g.id.clone()).collect();
370                                for gid in &relevant_ids {
371                                    store.mark_prevented(gid);
372                                }
373
374                                if store.try_resolve_pending(&command, &files, &sid).is_some() {
375                                    store.cross_session_boost();
376                                }
377
378                                // Promote mature gotchas to ProjectKnowledge
379                                let promotions = store.check_promotions();
380                                if !promotions.is_empty() {
381                                    let mut knowledge =
382                                        crate::core::knowledge::ProjectKnowledge::load_or_create(
383                                            root,
384                                        );
385                                    for (cat, trigger, resolution, conf) in &promotions {
386                                        knowledge.remember(
387                                            &format!("gotcha-{cat}"),
388                                            trigger,
389                                            resolution,
390                                            &sid,
391                                            *conf,
392                                        );
393                                    }
394                                    let _ = knowledge.save();
395                                }
396                            }
397
398                            let _ = store.save(root);
399                        }
400                    }
401
402                    format!("{result}{savings_note}{tee_hint}")
403                }
404            }
405            "ctx_search" => {
406                let pattern = get_str(args, "pattern")
407                    .ok_or_else(|| ErrorData::invalid_params("pattern is required", None))?;
408                let path = crate::hooks::normalize_tool_path(
409                    &get_str(args, "path").unwrap_or_else(|| ".".to_string()),
410                );
411                let ext = get_str(args, "ext");
412                let max = get_int(args, "max_results").unwrap_or(20) as usize;
413                let no_gitignore = get_bool(args, "ignore_gitignore").unwrap_or(false);
414                let crp = self.crp_mode;
415                let respect = !no_gitignore;
416                let search_result = tokio::time::timeout(
417                    std::time::Duration::from_secs(30),
418                    tokio::task::spawn_blocking(move || {
419                        crate::tools::ctx_search::handle(
420                            &pattern,
421                            &path,
422                            ext.as_deref(),
423                            max,
424                            crp,
425                            respect,
426                        )
427                    }),
428                )
429                .await;
430                let (result, original) = match search_result {
431                    Ok(Ok(r)) => r,
432                    Ok(Err(e)) => {
433                        return Err(ErrorData::internal_error(
434                            format!("search task failed: {e}"),
435                            None,
436                        ))
437                    }
438                    Err(_) => {
439                        let msg = "ctx_search timed out after 30s. Try narrowing the search:\n\
440                                   • Use a more specific pattern\n\
441                                   • Specify ext= to limit file types\n\
442                                   • Specify a subdirectory in path=";
443                        self.record_call("ctx_search", 0, 0, None).await;
444                        return Ok(CallToolResult::success(vec![Content::text(msg)]));
445                    }
446                };
447                let sent = crate::core::tokens::count_tokens(&result);
448                let saved = original.saturating_sub(sent);
449                self.record_call("ctx_search", original, saved, None).await;
450                let savings_note = if saved > 0 {
451                    format!("\n[saved {saved} tokens vs native Grep]")
452                } else {
453                    String::new()
454                };
455                format!("{result}{savings_note}")
456            }
457            "ctx_compress" => {
458                let include_sigs = get_bool(args, "include_signatures").unwrap_or(true);
459                let cache = self.cache.read().await;
460                let result =
461                    crate::tools::ctx_compress::handle(&cache, include_sigs, self.crp_mode);
462                drop(cache);
463                self.record_call("ctx_compress", 0, 0, None).await;
464                result
465            }
466            "ctx_benchmark" => {
467                let path = get_str(args, "path")
468                    .map(|p| crate::hooks::normalize_tool_path(&p))
469                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
470                let action = get_str(args, "action").unwrap_or_default();
471                let result = if action == "project" {
472                    let fmt = get_str(args, "format").unwrap_or_default();
473                    let bench = crate::core::benchmark::run_project_benchmark(&path);
474                    match fmt.as_str() {
475                        "json" => crate::core::benchmark::format_json(&bench),
476                        "markdown" | "md" => crate::core::benchmark::format_markdown(&bench),
477                        _ => crate::core::benchmark::format_terminal(&bench),
478                    }
479                } else {
480                    crate::tools::ctx_benchmark::handle(&path, self.crp_mode)
481                };
482                self.record_call("ctx_benchmark", 0, 0, None).await;
483                result
484            }
485            "ctx_metrics" => {
486                let cache = self.cache.read().await;
487                let calls = self.tool_calls.read().await;
488                let result = crate::tools::ctx_metrics::handle(&cache, &calls, self.crp_mode);
489                drop(cache);
490                drop(calls);
491                self.record_call("ctx_metrics", 0, 0, None).await;
492                result
493            }
494            "ctx_analyze" => {
495                let path = get_str(args, "path")
496                    .map(|p| crate::hooks::normalize_tool_path(&p))
497                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
498                let result = crate::tools::ctx_analyze::handle(&path, self.crp_mode);
499                self.record_call("ctx_analyze", 0, 0, None).await;
500                result
501            }
502            "ctx_discover" => {
503                let limit = get_int(args, "limit").unwrap_or(15) as usize;
504                let history = crate::cli::load_shell_history_pub();
505                let result = crate::tools::ctx_discover::discover_from_history(&history, limit);
506                self.record_call("ctx_discover", 0, 0, None).await;
507                result
508            }
509            "ctx_smart_read" => {
510                let path = get_str(args, "path")
511                    .map(|p| crate::hooks::normalize_tool_path(&p))
512                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
513                let mut cache = self.cache.write().await;
514                let output = crate::tools::ctx_smart_read::handle(&mut cache, &path, self.crp_mode);
515                let original = cache.get(&path).map_or(0, |e| e.original_tokens);
516                let tokens = crate::core::tokens::count_tokens(&output);
517                drop(cache);
518                self.record_call(
519                    "ctx_smart_read",
520                    original,
521                    original.saturating_sub(tokens),
522                    Some("auto".to_string()),
523                )
524                .await;
525                output
526            }
527            "ctx_delta" => {
528                let path = get_str(args, "path")
529                    .map(|p| crate::hooks::normalize_tool_path(&p))
530                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
531                let mut cache = self.cache.write().await;
532                let output = crate::tools::ctx_delta::handle(&mut cache, &path);
533                let original = cache.get(&path).map_or(0, |e| e.original_tokens);
534                let tokens = crate::core::tokens::count_tokens(&output);
535                drop(cache);
536                {
537                    let mut session = self.session.write().await;
538                    session.mark_modified(&path);
539                }
540                self.record_call(
541                    "ctx_delta",
542                    original,
543                    original.saturating_sub(tokens),
544                    Some("delta".to_string()),
545                )
546                .await;
547                output
548            }
549            "ctx_edit" => {
550                let path = get_str(args, "path")
551                    .map(|p| crate::hooks::normalize_tool_path(&p))
552                    .ok_or_else(|| ErrorData::invalid_params("path is required", None))?;
553                let old_string = get_str(args, "old_string").unwrap_or_default();
554                let new_string = get_str(args, "new_string")
555                    .ok_or_else(|| ErrorData::invalid_params("new_string is required", None))?;
556                let replace_all = args
557                    .as_ref()
558                    .and_then(|a| a.get("replace_all"))
559                    .and_then(|v| v.as_bool())
560                    .unwrap_or(false);
561                let create = args
562                    .as_ref()
563                    .and_then(|a| a.get("create"))
564                    .and_then(|v| v.as_bool())
565                    .unwrap_or(false);
566
567                let mut cache = self.cache.write().await;
568                let output = crate::tools::ctx_edit::handle(
569                    &mut cache,
570                    crate::tools::ctx_edit::EditParams {
571                        path: path.clone(),
572                        old_string,
573                        new_string,
574                        replace_all,
575                        create,
576                    },
577                );
578                drop(cache);
579
580                {
581                    let mut session = self.session.write().await;
582                    session.mark_modified(&path);
583                }
584                self.record_call("ctx_edit", 0, 0, None).await;
585                output
586            }
587            "ctx_dedup" => {
588                let action = get_str(args, "action").unwrap_or_default();
589                if action == "apply" {
590                    let mut cache = self.cache.write().await;
591                    let result = crate::tools::ctx_dedup::handle_action(&mut cache, &action);
592                    drop(cache);
593                    self.record_call("ctx_dedup", 0, 0, None).await;
594                    result
595                } else {
596                    let cache = self.cache.read().await;
597                    let result = crate::tools::ctx_dedup::handle(&cache);
598                    drop(cache);
599                    self.record_call("ctx_dedup", 0, 0, None).await;
600                    result
601                }
602            }
603            "ctx_fill" => {
604                let paths = get_str_array(args, "paths")
605                    .ok_or_else(|| ErrorData::invalid_params("paths array is required", None))?
606                    .into_iter()
607                    .map(|p| crate::hooks::normalize_tool_path(&p))
608                    .collect::<Vec<_>>();
609                let budget = get_int(args, "budget")
610                    .ok_or_else(|| ErrorData::invalid_params("budget is required", None))?
611                    as usize;
612                let mut cache = self.cache.write().await;
613                let output =
614                    crate::tools::ctx_fill::handle(&mut cache, &paths, budget, self.crp_mode);
615                drop(cache);
616                self.record_call("ctx_fill", 0, 0, Some(format!("budget:{budget}")))
617                    .await;
618                output
619            }
620            "ctx_intent" => {
621                let query = get_str(args, "query")
622                    .ok_or_else(|| ErrorData::invalid_params("query is required", None))?;
623                let root = get_str(args, "project_root").unwrap_or_else(|| ".".to_string());
624                let mut cache = self.cache.write().await;
625                let output =
626                    crate::tools::ctx_intent::handle(&mut cache, &query, &root, self.crp_mode);
627                drop(cache);
628                {
629                    let mut session = self.session.write().await;
630                    session.set_task(&query, Some("intent"));
631                }
632                self.record_call("ctx_intent", 0, 0, Some("semantic".to_string()))
633                    .await;
634                output
635            }
636            "ctx_response" => {
637                let text = get_str(args, "text")
638                    .ok_or_else(|| ErrorData::invalid_params("text is required", None))?;
639                let output = crate::tools::ctx_response::handle(&text, self.crp_mode);
640                self.record_call("ctx_response", 0, 0, None).await;
641                output
642            }
643            "ctx_context" => {
644                let cache = self.cache.read().await;
645                let turn = self.call_count.load(std::sync::atomic::Ordering::Relaxed);
646                let result = crate::tools::ctx_context::handle_status(&cache, turn, self.crp_mode);
647                drop(cache);
648                self.record_call("ctx_context", 0, 0, None).await;
649                result
650            }
651            "ctx_graph" => {
652                let action = get_str(args, "action")
653                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
654                let path = get_str(args, "path").map(|p| crate::hooks::normalize_tool_path(&p));
655                let root = crate::hooks::normalize_tool_path(
656                    &get_str(args, "project_root").unwrap_or_else(|| ".".to_string()),
657                );
658                let mut cache = self.cache.write().await;
659                let result = crate::tools::ctx_graph::handle(
660                    &action,
661                    path.as_deref(),
662                    &root,
663                    &mut cache,
664                    self.crp_mode,
665                );
666                drop(cache);
667                self.record_call("ctx_graph", 0, 0, Some(action)).await;
668                result
669            }
670            "ctx_cache" => {
671                let action = get_str(args, "action")
672                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
673                let mut cache = self.cache.write().await;
674                let result = match action.as_str() {
675                    "status" => {
676                        let entries = cache.get_all_entries();
677                        if entries.is_empty() {
678                            "Cache empty — no files tracked.".to_string()
679                        } else {
680                            let mut lines = vec![format!("Cache: {} file(s)", entries.len())];
681                            for (path, entry) in &entries {
682                                let fref = cache
683                                    .file_ref_map()
684                                    .get(*path)
685                                    .map(|s| s.as_str())
686                                    .unwrap_or("F?");
687                                lines.push(format!(
688                                    "  {fref}={} [{}L, {}t, read {}x]",
689                                    crate::core::protocol::shorten_path(path),
690                                    entry.line_count,
691                                    entry.original_tokens,
692                                    entry.read_count
693                                ));
694                            }
695                            lines.join("\n")
696                        }
697                    }
698                    "clear" => {
699                        let count = cache.clear();
700                        format!("Cache cleared — {count} file(s) removed. Next ctx_read will return full content.")
701                    }
702                    "invalidate" => {
703                        let path = get_str(args, "path")
704                            .map(|p| crate::hooks::normalize_tool_path(&p))
705                            .ok_or_else(|| {
706                                ErrorData::invalid_params("path is required for invalidate", None)
707                            })?;
708                        if cache.invalidate(&path) {
709                            format!(
710                                "Invalidated cache for {}. Next ctx_read will return full content.",
711                                crate::core::protocol::shorten_path(&path)
712                            )
713                        } else {
714                            format!(
715                                "{} was not in cache.",
716                                crate::core::protocol::shorten_path(&path)
717                            )
718                        }
719                    }
720                    _ => "Unknown action. Use: status, clear, invalidate".to_string(),
721                };
722                drop(cache);
723                self.record_call("ctx_cache", 0, 0, Some(action)).await;
724                result
725            }
726            "ctx_session" => {
727                let action = get_str(args, "action")
728                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
729                let value = get_str(args, "value");
730                let sid = get_str(args, "session_id");
731                let mut session = self.session.write().await;
732                let result = crate::tools::ctx_session::handle(
733                    &mut session,
734                    &action,
735                    value.as_deref(),
736                    sid.as_deref(),
737                );
738                drop(session);
739                self.record_call("ctx_session", 0, 0, Some(action)).await;
740                result
741            }
742            "ctx_knowledge" => {
743                let action = get_str(args, "action")
744                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
745                let category = get_str(args, "category");
746                let key = get_str(args, "key");
747                let value = get_str(args, "value");
748                let query = get_str(args, "query");
749                let pattern_type = get_str(args, "pattern_type");
750                let examples = get_str_array(args, "examples");
751                let confidence: Option<f32> = args
752                    .as_ref()
753                    .and_then(|a| a.get("confidence"))
754                    .and_then(|v| v.as_f64())
755                    .map(|v| v as f32);
756
757                let session = self.session.read().await;
758                let session_id = session.id.clone();
759                let project_root = session.project_root.clone().unwrap_or_else(|| {
760                    std::env::current_dir()
761                        .map(|p| p.to_string_lossy().to_string())
762                        .unwrap_or_else(|_| "unknown".to_string())
763                });
764                drop(session);
765
766                if action == "gotcha" {
767                    let trigger = get_str(args, "trigger").unwrap_or_default();
768                    let resolution = get_str(args, "resolution").unwrap_or_default();
769                    let severity = get_str(args, "severity").unwrap_or_default();
770                    let cat = category.as_deref().unwrap_or("convention");
771
772                    if trigger.is_empty() || resolution.is_empty() {
773                        self.record_call("ctx_knowledge", 0, 0, Some(action)).await;
774                        return Ok(CallToolResult::success(vec![Content::text(
775                            "ERROR: trigger and resolution are required for gotcha action",
776                        )]));
777                    }
778
779                    let mut store = crate::core::gotcha_tracker::GotchaStore::load(&project_root);
780                    let msg = match store.report_gotcha(
781                        &trigger,
782                        &resolution,
783                        cat,
784                        &severity,
785                        &session_id,
786                    ) {
787                        Some(gotcha) => {
788                            let conf = (gotcha.confidence * 100.0) as u32;
789                            let label = gotcha.category.short_label();
790                            format!("Gotcha recorded: [{label}] {trigger} (confidence: {conf}%)")
791                        }
792                        None => format!(
793                            "Gotcha noted: {trigger} (evicted by higher-confidence entries)"
794                        ),
795                    };
796                    let _ = store.save(&project_root);
797                    self.record_call("ctx_knowledge", 0, 0, Some(action)).await;
798                    return Ok(CallToolResult::success(vec![Content::text(msg)]));
799                }
800
801                let result = crate::tools::ctx_knowledge::handle(
802                    &project_root,
803                    &action,
804                    category.as_deref(),
805                    key.as_deref(),
806                    value.as_deref(),
807                    query.as_deref(),
808                    &session_id,
809                    pattern_type.as_deref(),
810                    examples,
811                    confidence,
812                );
813                self.record_call("ctx_knowledge", 0, 0, Some(action)).await;
814                result
815            }
816            "ctx_agent" => {
817                let action = get_str(args, "action")
818                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
819                let agent_type = get_str(args, "agent_type");
820                let role = get_str(args, "role");
821                let message = get_str(args, "message");
822                let category = get_str(args, "category");
823                let to_agent = get_str(args, "to_agent");
824                let status = get_str(args, "status");
825
826                let session = self.session.read().await;
827                let project_root = session.project_root.clone().unwrap_or_else(|| {
828                    std::env::current_dir()
829                        .map(|p| p.to_string_lossy().to_string())
830                        .unwrap_or_else(|_| "unknown".to_string())
831                });
832                drop(session);
833
834                let current_agent_id = self.agent_id.read().await.clone();
835                let result = crate::tools::ctx_agent::handle(
836                    &action,
837                    agent_type.as_deref(),
838                    role.as_deref(),
839                    &project_root,
840                    current_agent_id.as_deref(),
841                    message.as_deref(),
842                    category.as_deref(),
843                    to_agent.as_deref(),
844                    status.as_deref(),
845                );
846
847                if action == "register" {
848                    if let Some(id) = result.split(':').nth(1) {
849                        let id = id.split_whitespace().next().unwrap_or("").to_string();
850                        if !id.is_empty() {
851                            *self.agent_id.write().await = Some(id);
852                        }
853                    }
854                }
855
856                self.record_call("ctx_agent", 0, 0, Some(action)).await;
857                result
858            }
859            "ctx_share" => {
860                let action = get_str(args, "action")
861                    .ok_or_else(|| ErrorData::invalid_params("action is required", None))?;
862                let to_agent = get_str(args, "to_agent");
863                let paths = get_str(args, "paths");
864                let message = get_str(args, "message");
865
866                let from_agent = self.agent_id.read().await.clone();
867                let cache = self.cache.read().await;
868                let result = crate::tools::ctx_share::handle(
869                    &action,
870                    from_agent.as_deref(),
871                    to_agent.as_deref(),
872                    paths.as_deref(),
873                    message.as_deref(),
874                    &cache,
875                );
876                drop(cache);
877
878                self.record_call("ctx_share", 0, 0, Some(action)).await;
879                result
880            }
881            "ctx_overview" => {
882                let task = get_str(args, "task");
883                let path = get_str(args, "path").map(|p| crate::hooks::normalize_tool_path(&p));
884                let cache = self.cache.read().await;
885                let result = crate::tools::ctx_overview::handle(
886                    &cache,
887                    task.as_deref(),
888                    path.as_deref(),
889                    self.crp_mode,
890                );
891                drop(cache);
892                self.record_call("ctx_overview", 0, 0, Some("overview".to_string()))
893                    .await;
894                result
895            }
896            "ctx_preload" => {
897                let task = get_str(args, "task").unwrap_or_default();
898                let path = get_str(args, "path").map(|p| crate::hooks::normalize_tool_path(&p));
899                let mut cache = self.cache.write().await;
900                let result = crate::tools::ctx_preload::handle(
901                    &mut cache,
902                    &task,
903                    path.as_deref(),
904                    self.crp_mode,
905                );
906                drop(cache);
907                self.record_call("ctx_preload", 0, 0, Some("preload".to_string()))
908                    .await;
909                result
910            }
911            "ctx_wrapped" => {
912                let period = get_str(args, "period").unwrap_or_else(|| "week".to_string());
913                let result = crate::tools::ctx_wrapped::handle(&period);
914                self.record_call("ctx_wrapped", 0, 0, Some(period)).await;
915                result
916            }
917            "ctx_semantic_search" => {
918                let query = get_str(args, "query")
919                    .ok_or_else(|| ErrorData::invalid_params("query is required", None))?;
920                let path = crate::hooks::normalize_tool_path(
921                    &get_str(args, "path").unwrap_or_else(|| ".".to_string()),
922                );
923                let top_k = get_int(args, "top_k").unwrap_or(10) as usize;
924                let action = get_str(args, "action").unwrap_or_default();
925                let result = if action == "reindex" {
926                    crate::tools::ctx_semantic_search::handle_reindex(&path)
927                } else {
928                    crate::tools::ctx_semantic_search::handle(&query, &path, top_k, self.crp_mode)
929                };
930                self.record_call("ctx_semantic_search", 0, 0, Some("semantic".to_string()))
931                    .await;
932                result
933            }
934            _ => {
935                return Err(ErrorData::invalid_params(
936                    format!("Unknown tool: {name}"),
937                    None,
938                ));
939            }
940        };
941
942        let mut result_text = result_text;
943
944        if let Some(ctx) = auto_context {
945            result_text = format!("{ctx}\n\n{result_text}");
946        }
947
948        if name == "ctx_read" {
949            let read_path =
950                crate::hooks::normalize_tool_path(&get_str(args, "path").unwrap_or_default());
951            let project_root = {
952                let session = self.session.read().await;
953                session.project_root.clone()
954            };
955            let mut cache = self.cache.write().await;
956            let enrich = crate::tools::autonomy::enrich_after_read(
957                &self.autonomy,
958                &mut cache,
959                &read_path,
960                project_root.as_deref(),
961            );
962            if let Some(hint) = enrich.related_hint {
963                result_text = format!("{result_text}\n{hint}");
964            }
965
966            crate::tools::autonomy::maybe_auto_dedup(&self.autonomy, &mut cache);
967        }
968
969        if name == "ctx_shell" {
970            let cmd = get_str(args, "command").unwrap_or_default();
971            let output_tokens = crate::core::tokens::count_tokens(&result_text);
972            let calls = self.tool_calls.read().await;
973            let last_original = calls.last().map(|c| c.original_tokens).unwrap_or(0);
974            drop(calls);
975            if let Some(hint) = crate::tools::autonomy::shell_efficiency_hint(
976                &self.autonomy,
977                &cmd,
978                last_original,
979                output_tokens,
980            ) {
981                result_text = format!("{result_text}\n{hint}");
982            }
983        }
984
985        let skip_checkpoint = matches!(
986            name,
987            "ctx_compress"
988                | "ctx_metrics"
989                | "ctx_benchmark"
990                | "ctx_analyze"
991                | "ctx_cache"
992                | "ctx_discover"
993                | "ctx_dedup"
994                | "ctx_session"
995                | "ctx_knowledge"
996                | "ctx_agent"
997                | "ctx_share"
998                | "ctx_wrapped"
999                | "ctx_overview"
1000                | "ctx_preload"
1001        );
1002
1003        if !skip_checkpoint && self.increment_and_check() {
1004            if let Some(checkpoint) = self.auto_checkpoint().await {
1005                let combined = format!(
1006                    "{result_text}\n\n--- AUTO CHECKPOINT (every {} calls) ---\n{checkpoint}",
1007                    self.checkpoint_interval
1008                );
1009                return Ok(CallToolResult::success(vec![Content::text(combined)]));
1010            }
1011        }
1012
1013        let tool_duration_ms = tool_start.elapsed().as_millis() as u64;
1014        if tool_duration_ms > 100 {
1015            LeanCtxServer::append_tool_call_log(
1016                name,
1017                tool_duration_ms,
1018                0,
1019                0,
1020                None,
1021                &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
1022            );
1023        }
1024
1025        let current_count = self.call_count.load(std::sync::atomic::Ordering::Relaxed);
1026        if current_count > 0 && current_count.is_multiple_of(100) {
1027            std::thread::spawn(crate::cloud_sync::cloud_background_tasks);
1028        }
1029
1030        Ok(CallToolResult::success(vec![Content::text(result_text)]))
1031    }
1032}
1033
1034pub fn build_instructions_for_test(crp_mode: CrpMode) -> String {
1035    crate::instructions::build_instructions(crp_mode)
1036}
1037
1038fn get_str_array(args: &Option<serde_json::Map<String, Value>>, key: &str) -> Option<Vec<String>> {
1039    let arr = args.as_ref()?.get(key)?.as_array()?;
1040    let mut out = Vec::with_capacity(arr.len());
1041    for v in arr {
1042        let s = v.as_str()?.to_string();
1043        out.push(s);
1044    }
1045    Some(out)
1046}
1047
1048fn get_str(args: &Option<serde_json::Map<String, Value>>, key: &str) -> Option<String> {
1049    args.as_ref()?.get(key)?.as_str().map(|s| s.to_string())
1050}
1051
1052fn get_int(args: &Option<serde_json::Map<String, Value>>, key: &str) -> Option<i64> {
1053    args.as_ref()?.get(key)?.as_i64()
1054}
1055
1056fn get_bool(args: &Option<serde_json::Map<String, Value>>, key: &str) -> Option<bool> {
1057    args.as_ref()?.get(key)?.as_bool()
1058}
1059
1060fn execute_command(command: &str) -> (String, i32) {
1061    let (shell, flag) = crate::shell::shell_and_flag();
1062    let output = std::process::Command::new(&shell)
1063        .arg(&flag)
1064        .arg(command)
1065        .env("LEAN_CTX_ACTIVE", "1")
1066        .output();
1067
1068    match output {
1069        Ok(out) => {
1070            let code = out.status.code().unwrap_or(1);
1071            let stdout = String::from_utf8_lossy(&out.stdout);
1072            let stderr = String::from_utf8_lossy(&out.stderr);
1073            let text = if stdout.is_empty() {
1074                stderr.to_string()
1075            } else if stderr.is_empty() {
1076                stdout.to_string()
1077            } else {
1078                format!("{stdout}\n{stderr}")
1079            };
1080            (text, code)
1081        }
1082        Err(e) => (format!("ERROR: {e}"), 1),
1083    }
1084}
1085
1086pub fn tool_descriptions_for_test() -> Vec<(&'static str, &'static str)> {
1087    crate::tool_defs::list_all_tool_defs()
1088        .into_iter()
1089        .map(|(name, desc, _)| (name, desc))
1090        .collect()
1091}
1092
1093pub fn tool_schemas_json_for_test() -> String {
1094    crate::tool_defs::list_all_tool_defs()
1095        .iter()
1096        .map(|(name, _, schema)| format!("{}: {}", name, schema))
1097        .collect::<Vec<_>>()
1098        .join("\n")
1099}
1100
1101#[cfg(test)]
1102mod tests {
1103    #[test]
1104    fn test_unified_tool_count() {
1105        let tools = crate::tool_defs::unified_tool_defs();
1106        assert_eq!(tools.len(), 5, "Expected 5 unified tools");
1107    }
1108
1109    #[test]
1110    fn test_granular_tool_count() {
1111        let tools = crate::tool_defs::granular_tool_defs();
1112        assert!(tools.len() >= 25, "Expected at least 25 granular tools");
1113    }
1114}