riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(server): refactor handlers, router, state management, and memory tools

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-02 07:27:37 +01:00
1 parent 87ddb01063
commit a083719cf1
36 files changed
+1899 -597
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/ast.rs
+92 -42
View File
@@ -4,7 +4,7 @@ use crate::tools::ReadFileSkeletonTool;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use tree_sitter::{Parser, Node};
use tree_sitter::{Node, Parser};
pub struct ReadFileSkeletonHandler;
@@ -22,20 +22,22 @@ impl McpTool for ReadFileSkeletonHandler {
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ReadFileSkeletonTool = serde_json::from_value(args).map_err(|e| e.to_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 || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path)
.map_err(|e| crate::error::AppError::Internal(format!("Failed to read file: {}", e)))?;
let code = std::fs::read_to_string(&file_path).map_err(|e| {
crate::error::AppError::Internal(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,
@@ -46,20 +48,41 @@ impl McpTool for ReadFileSkeletonHandler {
"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_else(|| crate::error::AppError::Internal("Failed to parse code".to_string()))?;
parser
.set_language(&language.into())
.map_err(|e| e.to_string())?;
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();
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
let kind = node.kind();
let is_structural = matches!(kind, "use_declaration" | "import_statement" | "import_from_statement" | "struct_item" | "enum_item" | "trait_item" | "impl_item" | "function_item" | "function_declaration" | "function_definition" | "method_definition" | "interface_declaration" | "type_alias_declaration" | "class_declaration" | "class_definition");
let is_structural = matches!(
kind,
"use_declaration"
| "import_statement"
| "import_from_statement"
| "struct_item"
| "enum_item"
| "trait_item"
| "impl_item"
| "function_item"
| "function_declaration"
| "function_definition"
| "method_definition"
| "interface_declaration"
| "type_alias_declaration"
| "class_declaration"
| "class_definition"
);
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() {
@@ -76,7 +99,7 @@ impl McpTool for ReadFileSkeletonHandler {
if signature.is_empty() {
signature = node_text.to_string();
}
out.push_str(&indent);
out.push_str(signature.trim());
out.push('\n');
@@ -87,9 +110,9 @@ impl McpTool for ReadFileSkeletonHandler {
}
}
}
extract_skeleton(tree.root_node(), &code, &mut result_skeleton, 0);
if result_skeleton.is_empty() {
Ok(code)
} else {
@@ -120,20 +143,22 @@ impl McpTool for ReplaceAstNodeHandler {
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ReplaceAstNodeTool = serde_json::from_value(args).map_err(|e| e.to_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 || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path)
.map_err(|e| crate::error::AppError::Internal(format!("Failed to read file: {}", e)))?;
let code = std::fs::read_to_string(&file_path).map_err(|e| {
crate::error::AppError::Internal(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,
@@ -142,14 +167,28 @@ impl McpTool for ReplaceAstNodeHandler {
"c" | "h" => tree_sitter_c::LANGUAGE,
"cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE,
"go" => tree_sitter_go::LANGUAGE,
_ => return Err(crate::error::AppError::Internal(format!("Unsupported language for AST replacement: {}", ext))),
_ => {
return Err(crate::error::AppError::Internal(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_else(|| crate::error::AppError::Internal("Failed to parse code".to_string()))?;
parser
.set_language(&language.into())
.map_err(|e| e.to_string())?;
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
// Search for the node
fn find_node<'a>(node: Node<'a>, code: &str, target_type: &str, target_name: &str) -> Option<Node<'a>> {
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();
@@ -163,7 +202,7 @@ impl McpTool for ReplaceAstNodeHandler {
}
}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if let Some(found) = find_node(child, code, target_type, target_name) {
@@ -172,22 +211,33 @@ impl McpTool for ReplaceAstNodeHandler {
}
None
}
let target_node = find_node(tree.root_node(), &code, &tool_args.node_type, &tool_args.node_name);
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))
Ok(format!(
"Successfully replaced node {} of type {} in {}",
tool_args.node_name, tool_args.node_type, file_path
))
} else {
Err(crate::error::AppError::Internal(format!("Could not find node {} of type {}", tool_args.node_name, tool_args.node_type)))
Err(crate::error::AppError::Internal(format!(
"Could not find node {} of type {}",
tool_args.node_name, tool_args.node_type
)))
}
})
.await
@@ -208,7 +258,7 @@ mod tests {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let file_path = dir.path().join("test_skeleton.rs");
let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}";
std::fs::write(&file_path, code).unwrap();
@@ -227,7 +277,7 @@ mod tests {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let file_path = dir.path().join("test_replace.rs");
let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}";
std::fs::write(&file_path, code).unwrap();