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

feat(embedding,vision): candle embeddings with offline fallback, on-demand clipboard vision capture, and concurrency audit

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-07 06:36:09 +01:00
1 parent 5bd8b1587a
commit e4a0fe72df
47 files changed
+6292 -3503
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/ast.rs
+135 -24
View File
@@ -6,6 +6,24 @@ use serde_json::Value;
use std::sync::Arc;
use tree_sitter::{Node, Parser};
fn validate_safe_path(path_str: &str) -> crate::error::Result<()> {
if path_str.trim().is_empty() || path_str.contains('\0') {
return Err(crate::error::AppError::Internal(
"Invalid file path: path is empty or contains null characters".to_string(),
));
}
let path = std::path::Path::new(path_str);
for component in path.components() {
if component == std::path::Component::ParentDir {
return Err(crate::error::AppError::Internal(format!(
"Path traversal forbidden: '{}' contains parent directory relative components",
path_str
)));
}
}
Ok(())
}
pub struct ReadFileSkeletonHandler;
#[async_trait]
@@ -26,13 +44,13 @@ impl McpTool for ReadFileSkeletonHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone();
validate_safe_path(&file_path)?;
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 mut parser = Parser::new();
let ext = std::path::Path::new(&file_path)
.extension()
.and_then(|s| s.to_str())
@@ -49,10 +67,10 @@ impl McpTool for ReadFileSkeletonHandler {
_ => return Ok(code),
};
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&language.into())
.map_err(|e| e.to_string())?;
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
@@ -150,6 +168,8 @@ impl McpTool for ReplaceAstNodeHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone();
validate_safe_path(&file_path)?;
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))
@@ -198,8 +218,14 @@ impl McpTool for ReplaceAstNodeHandler {
| "method_definition"
| "function"
),
"struct" => matches!(actual_kind, "struct_item" | "struct_declaration" | "struct_specifier"),
"class" => matches!(actual_kind, "class_declaration" | "class_definition" | "class_item"),
"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,
@@ -225,7 +251,11 @@ impl McpTool for ReplaceAstNodeHandler {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
let kind = child.kind();
if kind == "identifier" || kind == "name" || kind == "property_identifier" || kind == "field_identifier" {
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);
@@ -254,12 +284,38 @@ impl McpTool for ReplaceAstNodeHandler {
let start_byte = node.start_byte();
let end_byte = node.end_byte();
let mut new_file_content = String::new();
if !code.is_char_boundary(start_byte) || !code.is_char_boundary(end_byte) {
return Err(crate::error::AppError::Internal(format!(
"Byte offsets {}..{} do not fall on UTF-8 character boundaries in {}",
start_byte, end_byte, file_path
)));
}
let mut new_file_content =
String::with_capacity(code.len() + tool_args.new_content.len());
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())?;
let target_path = std::path::PathBuf::from(&file_path);
let parent_dir = target_path
.parent()
.unwrap_or_else(|| std::path::Path::new("."));
let temp_file_path =
parent_dir.join(format!(".tmp_ast_{}.tmp", uuid::Uuid::new_v4()));
std::fs::write(&temp_file_path, new_file_content).map_err(|e| e.to_string())?;
if let Err(e) = std::fs::rename(&temp_file_path, &target_path) {
// On Windows, std::fs::rename fails if the target file already exists.
// Fall back to copy-and-remove to ensure atomic-like overwrite behavior.
if let Err(copy_err) = std::fs::copy(&temp_file_path, &target_path) {
let _ = std::fs::remove_file(&temp_file_path);
return Err(crate::error::AppError::Internal(format!(
"Failed to atomically overwrite {}: rename failed ({}), copy failed ({})",
file_path, e, copy_err
)));
}
let _ = std::fs::remove_file(&temp_file_path);
}
Ok(format!(
"Successfully replaced node {} of type {} in {}",
tool_args.node_name, tool_args.node_type, file_path
@@ -278,20 +334,31 @@ impl McpTool for ReplaceAstNodeHandler {
}
}
fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
fn scan_workspace_for_symbol(
target_sym: &str,
limit: usize,
filter_fn_call: bool,
workspace_dir: Option<std::path::PathBuf>,
) -> Vec<serde_json::Value> {
let mut results = Vec::new();
let cwd = match std::env::current_dir() {
Ok(dir) => dir,
Err(_) => return results,
};
let scan_dir = workspace_dir.unwrap_or_else(|| {
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."))
});
let walker = ignore::WalkBuilder::new(&cwd)
if !scan_dir.exists() {
return results;
}
let walker = ignore::WalkBuilder::new(&scan_dir)
.hidden(true)
.git_ignore(true)
.build();
let mut scanned_files = 0;
let call_pattern = format!("{}(", target_sym);
let await_pattern = format!("{}.await", target_sym);
for result in walker {
let entry = match result {
Ok(e) => e,
@@ -301,7 +368,20 @@ fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: boo
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") {
if matches!(
ext,
"rs" | "ts"
| "tsx"
| "js"
| "jsx"
| "py"
| "go"
| "java"
| "c"
| "cpp"
| "h"
| "hpp"
) {
scanned_files += 1;
if scanned_files > 500 {
break;
@@ -309,8 +389,7 @@ fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: boo
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))
line.contains(&call_pattern) || line.contains(&await_pattern)
} else {
line.contains(target_sym)
};
@@ -371,10 +450,22 @@ impl McpTool for FindSymbolReferencesHandler {
Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
})?;
let custom_dir = if let Some(ref d) = req.workspace_dir {
Some(std::path::PathBuf::from(d))
} else {
state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
})
};
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)
scan_workspace_for_symbol(&target_sym, remaining, false, custom_dir)
})
.await
.unwrap_or_default();
@@ -405,13 +496,13 @@ impl McpTool for GetCallersHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10);
let target_fn = req.function_name.clone();
let call_pattern = format!("{}(", target_fn);
let await_pattern = format!("{}.await", target_fn);
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))
{
if snippet.code.contains(&call_pattern) || snippet.code.contains(&await_pattern) {
matching.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
@@ -425,10 +516,22 @@ impl McpTool for GetCallersHandler {
Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
})?;
let custom_dir = if let Some(ref d) = req.workspace_dir {
Some(std::path::PathBuf::from(d))
} else {
state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
})
};
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)
scan_workspace_for_symbol(&target_fn, remaining, true, custom_dir)
})
.await
.unwrap_or_default();
@@ -468,9 +571,17 @@ impl McpTool for AnalyzeImpactHandler {
}
});
let custom_dir = state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
});
let sym_clone = sym.clone();
let disk_refs = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&sym_clone, 20, false)
scan_workspace_for_symbol(&sym_clone, 20, false, custom_dir)
})
.await
.unwrap_or_default();