refactor: apply 5-pass audit optimizations across mcp-memory codebase

server/src/handlers/ast.rs

@@ -56,7 +56,7 @@ impl McpTool for ReadFileSkeletonHandler {
             let tree = parser.parse(&code, None).ok_or_else(|| {
                 crate::error::AppError::Internal("Failed to parse code".to_string())
             })?;
-            let mut result_skeleton = String::new();
+            let mut result_skeleton = String::with_capacity(code.len() / 2);
 
             fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
                 if depth > 128 {
@@ -141,7 +141,7 @@ impl McpTool for ReplaceAstNodeHandler {
     fn schema(&self) -> Value {
         crate::mcp::tool_def::<ReplaceAstNodeTool>(
             "replace_ast_node",
-            "Replace a specific AST node (e.g., function, struct) entirely using tree-sitter for robust structural editing.",
+            "Replace a specific AST node (e.g., function, struct, enum, class, trait) entirely using tree-sitter for robust structural editing. Supported node_type values include: 'function_item' (or 'function'/'fn'/'method'), 'struct_item' (or 'struct'), 'class_declaration' (or 'class'), 'enum_item' (or 'enum'), 'trait_item' (or 'trait'/'interface'), 'type_alias_declaration' (or 'type').",
         )
     }
 
@@ -185,6 +185,34 @@ impl McpTool for ReplaceAstNodeHandler {
                 crate::error::AppError::Internal("Failed to parse code".to_string())
             })?;
 
+            fn matches_node_type(actual_kind: &str, requested_type: &str) -> bool {
+                if actual_kind == requested_type {
+                    return true;
+                }
+                match requested_type.to_lowercase().as_str() {
+                    "function" | "func" | "fn" | "method" | "def" => matches!(
+                        actual_kind,
+                        "function_item"
+                            | "function_declaration"
+                            | "function_definition"
+                            | "method_definition"
+                            | "function"
+                    ),
+                    "struct" => matches!(actual_kind, "struct_item" | "struct_declaration" | "struct_specifier"),
+                    "class" => matches!(actual_kind, "class_declaration" | "class_definition" | "class_item"),
+                    "enum" => matches!(actual_kind, "enum_item" | "enum_declaration"),
+                    "trait" | "interface" => matches!(
+                        actual_kind,
+                        "trait_item" | "interface_declaration" | "interface_item"
+                    ),
+                    "type" | "type_alias" => matches!(
+                        actual_kind,
+                        "type_alias_declaration" | "type_item" | "type_definition"
+                    ),
+                    _ => false,
+                }
+            }
+
             // Search for the node
             fn find_node<'a>(
                 node: Node<'a>,
@@ -192,12 +220,12 @@ impl McpTool for ReplaceAstNodeHandler {
                 target_type: &str,
                 target_name: &str,
             ) -> Option<Node<'a>> {
-                if node.kind() == target_type {
+                if matches_node_type(node.kind(), target_type) {
                     // Try to find the name/identifier
                     let mut cursor = node.walk();
                     for child in node.children(&mut cursor) {
                         let kind = child.kind();
-                        if kind == "identifier" || kind == "name" {
+                        if kind == "identifier" || kind == "name" || kind == "property_identifier" || kind == "field_identifier" {
                             let name = child.utf8_text(code.as_bytes()).unwrap_or("");
                             if name == target_name {
                                 return Some(node);
@@ -250,6 +278,61 @@ impl McpTool for ReplaceAstNodeHandler {
     }
 }
 
+fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
+    let mut results = Vec::new();
+    let cwd = match std::env::current_dir() {
+        Ok(dir) => dir,
+        Err(_) => return results,
+    };
+
+    let walker = ignore::WalkBuilder::new(&cwd)
+        .hidden(true)
+        .git_ignore(true)
+        .build();
+
+    let mut scanned_files = 0;
+
+    for result in walker {
+        let entry = match result {
+            Ok(e) => e,
+            Err(_) => continue,
+        };
+
+        if entry.file_type().is_some_and(|ft| ft.is_file()) {
+            let path = entry.path();
+            let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
+            if matches!(ext, "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "h" | "hpp") {
+                scanned_files += 1;
+                if scanned_files > 500 {
+                    break;
+                }
+                if let Ok(content) = std::fs::read_to_string(path) {
+                    for (line_num, line) in content.lines().enumerate() {
+                        let is_match = if filter_fn_call {
+                            line.contains(&format!("{}(", target_sym))
+                                || line.contains(&format!("{}.await", target_sym))
+                        } else {
+                            line.contains(target_sym)
+                        };
+
+                        if is_match {
+                            results.push(serde_json::json!({
+                                "file_path": path.to_string_lossy(),
+                                "line": line_num + 1,
+                                "content": line.trim(),
+                            }));
+                            if results.len() >= limit {
+                                return results;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    results
+}
+
 pub struct FindSymbolReferencesHandler;
 
 #[async_trait]
@@ -269,22 +352,36 @@ impl McpTool for FindSymbolReferencesHandler {
         let req: crate::tools::FindSymbolReferencesTool =
             serde_json::from_value(args).map_err(|e| e.to_string())?;
         let limit = req.limit.unwrap_or(10);
-        let target_sym = req.symbol;
+        let target_sym = req.symbol.clone();
 
-        let matches = state.code.snippets.read_with(|snippets| {
+        let mut matches = state.code.snippets.read_with(|snippets| {
             let mut refs = Vec::new();
             for snippet in snippets {
                 if snippet.code.contains(&target_sym) || snippet.name.contains(&target_sym) {
-                    refs.push(snippet.clone());
+                    refs.push(serde_json::json!({
+                        "source": "snippet",
+                        "name": snippet.name,
+                        "code": snippet.code,
+                    }));
                     if refs.len() >= limit {
                         break;
                     }
                 }
             }
-            Ok::<String, crate::error::AppError>(serde_json::to_string(&refs)?)
+            Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
         })?;
 
-        Ok(matches)
+        if matches.len() < limit {
+            let remaining = limit - matches.len();
+            let disk_matches = tokio::task::spawn_blocking(move || {
+                scan_workspace_for_symbol(&target_sym, remaining, false)
+            })
+            .await
+            .unwrap_or_default();
+            matches.extend(disk_matches);
+        }
+
+        Ok(serde_json::to_string_pretty(&matches)?)
     }
 }
 
@@ -307,24 +404,38 @@ impl McpTool for GetCallersHandler {
         let req: crate::tools::GetCallersTool =
             serde_json::from_value(args).map_err(|e| e.to_string())?;
         let limit = req.limit.unwrap_or(10);
-        let target_fn = req.function_name;
+        let target_fn = req.function_name.clone();
 
-        let callers = state.code.snippets.read_with(|snippets| {
+        let mut callers = state.code.snippets.read_with(|snippets| {
             let mut matching = Vec::new();
             for snippet in snippets {
                 if snippet.code.contains(&format!("{}(", target_fn))
                     || snippet.code.contains(&format!("{}.await", target_fn))
                 {
-                    matching.push(snippet.clone());
+                    matching.push(serde_json::json!({
+                        "source": "snippet",
+                        "name": snippet.name,
+                        "code": snippet.code,
+                    }));
                     if matching.len() >= limit {
                         break;
                     }
                 }
             }
-            Ok::<String, crate::error::AppError>(serde_json::to_string(&matching)?)
+            Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
         })?;
 
-        Ok(callers)
+        if callers.len() < limit {
+            let remaining = limit - callers.len();
+            let disk_callers = tokio::task::spawn_blocking(move || {
+                scan_workspace_for_symbol(&target_fn, remaining, true)
+            })
+            .await
+            .unwrap_or_default();
+            callers.extend(disk_callers);
+        }
+
+        Ok(serde_json::to_string_pretty(&callers)?)
     }
 }
 
@@ -357,6 +468,20 @@ impl McpTool for AnalyzeImpactHandler {
             }
         });
 
+        let sym_clone = sym.clone();
+        let disk_refs = tokio::task::spawn_blocking(move || {
+            scan_workspace_for_symbol(&sym_clone, 20, false)
+        })
+        .await
+        .unwrap_or_default();
+
+        for r in &disk_refs {
+            if let Some(path) = r.get("file_path").and_then(|p| p.as_str()) {
+                let line = r.get("line").and_then(|l| l.as_u64()).unwrap_or(0);
+                callers.push(format!("{}:{}", path, line));
+            }
+        }
+
         let mut kg_connected = Vec::new();
         state.read_graph(|g| {
             for rel in &g.relations {
@@ -384,7 +509,7 @@ impl McpTool for AnalyzeImpactHandler {
             "target_symbol": sym,
             "risk_level": risk_level,
             "ast_callers_count": caller_count,
-            "ast_callers_sample": callers.into_iter().take(5).collect::<Vec<_>>(),
+            "ast_callers_sample": callers.into_iter().take(10).collect::<Vec<_>>(),
             "graph_relations_count": graph_count,
             "graph_relations": kg_connected,
             "recommendation": match risk_level {