mcp-memory/server/src/handlers/ast.rs

use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::ReadFileSkeletonTool;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use tree_sitter::{Parser, Node};

pub struct ReadFileSkeletonHandler;

#[async_trait]
impl McpTool for ReadFileSkeletonHandler {
    fn name(&self) -> &'static str {
        "read_file_skeleton"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ReadFileSkeletonTool>(
            "read_file_skeleton",
            "Read a source file and return only its AST structural skeleton, omitting implementation details to save tokens.",
        )
    }

    async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> Result<String, String> {
        let tool_args: ReadFileSkeletonTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        
        let file_path = tool_args.file_path.clone();
        let result = tokio::task::spawn_blocking(move || -> Result<String, String> {
            let code = std::fs::read_to_string(&file_path)
                .map_err(|e| format!("Failed to read file: {}", e))?;
                
            let mut parser = Parser::new();
            
            let ext = std::path::Path::new(&file_path)
                .extension()
                .and_then(|s| s.to_str())
                .unwrap_or("");
                
            let language = match ext {
                "rs" => tree_sitter_rust::LANGUAGE,
                "ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT,
                "py" => tree_sitter_python::LANGUAGE,
                "java" => tree_sitter_java::LANGUAGE,
                "c" | "h" => tree_sitter_c::LANGUAGE,
                "cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE,
                "go" => tree_sitter_go::LANGUAGE,
                _ => return Ok(code),
            };
            
            parser.set_language(&language.into()).map_err(|e| e.to_string())?;
            
            let tree = parser.parse(&code, None).ok_or("Failed to parse code")?;
            let mut result_skeleton = String::new();
            
            fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
                let kind = node.kind();
                
                let is_structural = match kind {
                    "use_declaration" | "import_statement" | "import_from_statement" => true,
                    "struct_item" | "enum_item" | "trait_item" | "impl_item" => true,
                    "function_item" | "function_declaration" | "function_definition" | "method_definition" => true,
                    "interface_declaration" | "type_alias_declaration" | "class_declaration" | "class_definition" => true,
                    _ => false,
                };

                if is_structural {
                    let indent = "  ".repeat(depth);
                    
                    let node_text = node.utf8_text(code.as_bytes()).unwrap_or("");
                    let mut signature = String::new();
                    for line in node_text.lines() {
                        let trimmed = line.trim();
                        if trimmed.ends_with('{') || trimmed.ends_with(':') {
                            signature.push_str(line);
                            signature.push_str(" ... }");
                            break;
                        } else {
                            signature.push_str(line);
                            signature.push('\n');
                        }
                    }
                    if signature.is_empty() {
                        signature = node_text.to_string();
                    }
                    
                    out.push_str(&indent);
                    out.push_str(&signature.trim());
                    out.push('\n');
                } else if node.is_named() {
                    let mut cursor = node.walk();
                    for child in node.named_children(&mut cursor) {
                        extract_skeleton(child, code, out, depth);
                    }
                }
            }
            
            extract_skeleton(tree.root_node(), &code, &mut result_skeleton, 0);
            
            if result_skeleton.is_empty() {
                Ok(code)
            } else {
                Ok(result_skeleton)
            }
        })
        .await
        .map_err(|e| format!("Task panic: {}", e))??;

        Ok(result)
    }
}
use crate::tools::ReplaceAstNodeTool;

pub struct ReplaceAstNodeHandler;

#[async_trait]
impl McpTool for ReplaceAstNodeHandler {
    fn name(&self) -> &'static str {
        "replace_ast_node"
    }

    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.",
        )
    }

    async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> Result<String, String> {
        let tool_args: ReplaceAstNodeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        
        let file_path = tool_args.file_path.clone();
        let result = tokio::task::spawn_blocking(move || -> Result<String, String> {
            let code = std::fs::read_to_string(&file_path)
                .map_err(|e| format!("Failed to read file: {}", e))?;
                
            let mut parser = Parser::new();
            
            let ext = std::path::Path::new(&file_path)
                .extension()
                .and_then(|s| s.to_str())
                .unwrap_or("");
                
            let language = match ext {
                "rs" => tree_sitter_rust::LANGUAGE,
                "ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT,
                "py" => tree_sitter_python::LANGUAGE,
                "java" => tree_sitter_java::LANGUAGE,
                "c" | "h" => tree_sitter_c::LANGUAGE,
                "cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE,
                "go" => tree_sitter_go::LANGUAGE,
                _ => return Err(format!("Unsupported language for AST replacement: {}", ext)),
            };
            
            parser.set_language(&language.into()).map_err(|e| e.to_string())?;
            let tree = parser.parse(&code, None).ok_or("Failed to parse code")?;
            
            // Search for the node
            fn find_node<'a>(node: Node<'a>, code: &str, target_type: &str, target_name: &str) -> Option<Node<'a>> {
                if 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" {
                            let name = child.utf8_text(code.as_bytes()).unwrap_or("");
                            if name == target_name {
                                return Some(node);
                            }
                        }
                    }
                }
                
                let mut cursor = node.walk();
                for child in node.children(&mut cursor) {
                    if let Some(found) = find_node(child, code, target_type, target_name) {
                        return Some(found);
                    }
                }
                None
            }
            
            let target_node = find_node(tree.root_node(), &code, &tool_args.node_type, &tool_args.node_name);
            
            if let Some(node) = target_node {
                let start_byte = node.start_byte();
                let end_byte = node.end_byte();
                
                let mut new_file_content = String::new();
                new_file_content.push_str(&code[..start_byte]);
                new_file_content.push_str(&tool_args.new_content);
                new_file_content.push_str(&code[end_byte..]);
                
                std::fs::write(&file_path, new_file_content).map_err(|e| e.to_string())?;
                Ok(format!("Successfully replaced node {} of type {} in {}", tool_args.node_name, tool_args.node_type, file_path))
            } else {
                Err(format!("Could not find node {} of type {}", tool_args.node_name, tool_args.node_type))
            }
        })
        .await
        .map_err(|e| format!("Task panic: {}", e))??;

        Ok(result)
    }
}