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

refactor: address 5-pass audit findings for antipatterns, bottlenecks, memory efficiency, and LLM handlers

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-05 21:44:10 +01:00
1 parent 626403900f
commit 924b6d09fa
30 files changed
+1120 -503
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/ast.rs
+4 -1
View File
@@ -59,6 +59,9 @@ impl McpTool for ReadFileSkeletonHandler {
let mut result_skeleton = String::new();
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
if depth > 128 {
return;
}
let kind = node.kind();
let is_structural = matches!(
@@ -106,7 +109,7 @@ impl McpTool for ReadFileSkeletonHandler {
} 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(child, code, out, depth + 1);
}
}
}
server/src/handlers/graph.rs
+92 -51
View File
@@ -29,6 +29,18 @@ impl McpTool for QueryGraphPathHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.read_graph(|graph| {
let max_depth = req.max_depth.unwrap_or(5);
// Pre-index relations into an adjacency map for O(1) neighbor lookups
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> = std::collections::HashMap::new();
for rel in &graph.relations {
adj.entry(rel.from.as_str())
.or_default()
.push((rel.to.as_str(), rel.relation_type.as_str(), false));
adj.entry(rel.to.as_str())
.or_default()
.push((rel.from.as_str(), rel.relation_type.as_str(), true));
}
let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
@@ -49,23 +61,14 @@ impl McpTool for QueryGraphPathHandler {
}
nodes_at_current_depth -= 1;
if current_depth < max_depth {
for rel in &graph.relations {
if rel.from == current && !visited.contains(rel.to.as_str()) {
visited.insert(rel.to.as_str());
parents.insert(
rel.to.as_str(),
(current, rel.relation_type.as_str(), false),
);
queue.push_back(rel.to.as_str());
nodes_at_next_depth += 1;
} else if rel.to == current && !visited.contains(rel.from.as_str()) {
visited.insert(rel.from.as_str());
parents.insert(
rel.from.as_str(),
(current, rel.relation_type.as_str(), true),
);
queue.push_back(rel.from.as_str());
nodes_at_next_depth += 1;
if let Some(neighbors) = adj.get(current) {
for &(neighbor, rel_type, is_inverse) in neighbors {
if !visited.contains(neighbor) {
visited.insert(neighbor);
parents.insert(neighbor, (current, rel_type, is_inverse));
queue.push_back(neighbor);
nodes_at_next_depth += 1;
}
}
}
}
@@ -165,25 +168,21 @@ impl McpTool for CreateRelationsHandler {
}
};
let mut missing_nodes = std::collections::HashSet::new();
state.modify_graph(|g| {
for mut relation in req.relations {
state.read_graph(|g| {
for relation in &req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() {
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
let from_exists = g.entities.contains_key(&relation.from);
let to_exists = g.entities.contains_key(&relation.to);
if from_exists && to_exists {
g.relations.push(relation);
} else {
if !from_exists {
missing_nodes.insert(relation.from);
}
if !to_exists {
missing_nodes.insert(relation.to);
}
if !from_exists {
missing_nodes.insert(relation.from.clone());
}
if !to_exists {
missing_nodes.insert(relation.to.clone());
}
}
}
});
if !missing_nodes.is_empty() {
let missing: Vec<_> = missing_nodes.into_iter().collect();
return Err(crate::error::AppError::Internal(format!(
@@ -191,6 +190,15 @@ impl McpTool for CreateRelationsHandler {
missing.join(", ")
)));
}
state.modify_graph(|g| {
for mut relation in req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() {
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
g.relations.push(relation);
}
}
});
Ok("Relations created".to_string())
}
}
@@ -210,21 +218,28 @@ impl McpTool for AddObservationsHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut missing_entities = Vec::new();
state.modify_graph(|g| {
for o in req.observations {
if let Some(e) = g.entities.get_mut(&o.entity_name) {
e.observations.extend(o.contents);
} else {
missing_entities.push(o.entity_name);
state.read_graph(|g| {
for o in &req.observations {
if !g.entities.contains_key(&o.entity_name) {
missing_entities.push(o.entity_name.clone());
}
}
});
if !missing_entities.is_empty() {
return Err(crate::error::AppError::Internal(format!(
"Error: Observations dropped for missing entities: {}",
missing_entities.join(", ")
)));
}
state.modify_graph(|g| {
for o in req.observations {
if let Some(e) = g.entities.get_mut(&o.entity_name) {
e.observations.extend(o.contents);
}
}
});
Ok("Observations added".to_string())
}
}
@@ -245,15 +260,12 @@ impl McpTool for DeleteEntitiesHandler {
let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
let mut missing = Vec::new();
state.modify_graph(|master| {
state.read_graph(|g| {
for name in &to_delete {
if master.entities.remove(name).is_none() {
if !g.entities.contains_key(name) {
missing.push(name.clone());
}
}
master
.relations
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
});
if !missing.is_empty() {
@@ -263,6 +275,15 @@ impl McpTool for DeleteEntitiesHandler {
)));
}
state.modify_graph(|master| {
for name in &to_delete {
master.entities.remove(name);
}
master
.relations
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
});
let idx = state.get_search_index();
for name in to_delete {
drop(idx.delete_document(&name));
@@ -290,22 +311,29 @@ impl McpTool for DeleteObservationsHandler {
let req: DeleteObservationsTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut missing = Vec::new();
state.modify_graph(|master| {
for d in req.deletions {
if let Some(e) = master.entities.get_mut(&d.entity_name) {
let to_rem: HashSet<_> = d.observations.into_iter().collect();
e.observations.retain(|o| !to_rem.contains(o));
} else {
missing.push(d.entity_name);
state.read_graph(|g| {
for d in &req.deletions {
if !g.entities.contains_key(&d.entity_name) {
missing.push(d.entity_name.clone());
}
}
});
if !missing.is_empty() {
return Err(crate::error::AppError::Internal(format!(
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
missing.join(", ")
)));
}
state.modify_graph(|master| {
for d in req.deletions {
if let Some(e) = master.entities.get_mut(&d.entity_name) {
let to_rem: HashSet<_> = d.observations.into_iter().collect();
e.observations.retain(|o| !to_rem.contains(o));
}
}
});
Ok("Observations deleted".to_string())
}
}
@@ -401,6 +429,9 @@ impl McpTool for SearchNodesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10);
let include_body = req.include_body.unwrap_or(false);
let matches = if let Ok(idx) = state.search_index.read() {
idx.search(&req.query, req.namespace.as_deref())
.unwrap_or_default()
@@ -409,15 +440,25 @@ impl McpTool for SearchNodesHandler {
};
let data = state.read_graph(|full| -> crate::error::Result<String> {
let mut result = BorrowedGraph::default();
for (id, doc_type, _, _, _) in &matches {
let mut matched_entities = Vec::new();
for (id, doc_type, _, _, _) in matches.iter().take(limit) {
if doc_type == "entity"
&& let Some(e) = full.entities.get(id)
{
result.entities.insert(id, e);
if include_body {
matched_entities.push(serde_json::to_value(e)?);
} else {
matched_entities.push(serde_json::json!({
"name": e.name,
"entity_type": e.entity_type,
"namespace": e.namespace,
"git_branch": e.git_branch,
"observations_count": e.observations.len()
}));
}
}
}
Ok::<String, crate::error::AppError>(serde_json::to_string(&result)?)
Ok::<String, crate::error::AppError>(serde_json::to_string(&matched_entities)?)
})?;
Ok(data)
}
server/src/handlers/meta.rs
+8 -3
View File
@@ -23,14 +23,14 @@ impl McpTool for LogDecisionHandler {
let idx = state.get_search_index();
let mut final_id = String::new();
let mut adrs_to_index = Vec::new();
state.code.adrs.modify(|adrs| {
if let Some(superseded_id) = &req.supersedes {
for old_adr in adrs.iter_mut() {
if old_adr.id == *superseded_id {
old_adr.status = "superseded".to_string();
// Re-index the modified old ADR
drop(idx.index_adr(old_adr));
adrs_to_index.push(old_adr.clone());
break;
}
}
@@ -48,10 +48,15 @@ impl McpTool for LogDecisionHandler {
timestamp: crate::handlers::utils::now_secs(),
};
drop(idx.index_adr(&a));
adrs_to_index.push(a.clone());
adrs.push(a);
});
// Index in Tantivy outside the store write lock
for adr in &adrs_to_index {
drop(idx.index_adr(adr));
}
state.record_activity("decision", &format!("Logged {}: {}", final_id, req.title), Some(&req.decision));
Ok(format!("Logged decision {}: {}", final_id, req.title))
}
server/src/handlers/vision.rs
+119 -41
View File
@@ -84,6 +84,96 @@ impl McpTool for WriteClipboardHandler {
}
}
pub fn get_native_clipboard_text() -> Option<String> {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() {
return Some(text);
}
}
}
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.arg("--no-newline");
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
}
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
if !text.trim().is_empty() {
return Some(text);
}
}
}
}
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
if !text.trim().is_empty() {
return Some(text);
}
}
}
}
None
}
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
image_data.width as u32,
image_data.height as u32,
image_data.bytes.into_owned(),
) {
return Some(image::DynamicImage::ImageRgba8(img));
}
}
}
for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] {
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.args(["--type", mime]);
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
}
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
}
}
}
}
for mime in &["image/png", "image/jpeg", "image/bmp"] {
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
}
}
}
}
None
}
pub struct ReadClipboardHandler;
#[async_trait]
@@ -104,51 +194,40 @@ impl McpTool for ReadClipboardHandler {
tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
let mut out = serde_json::Map::new();
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() {
out.insert("text".into(), json!(text));
}
if let Some(text) = get_native_clipboard_text() {
out.insert("text".into(), json!(text));
}
if let Some(dynamic_img) = get_native_clipboard_image() {
let mut img = dynamic_img;
let max_dim = 1024;
if img.width() > max_dim || img.height() > max_dim {
img = img.resize(max_dim, max_dim, FilterType::Lanczos3);
}
let rgb_img = img.into_rgb8();
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
image_data.width as u32,
image_data.height as u32,
image_data.bytes.into_owned(),
) {
let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
let max_dim = 1024;
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim {
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3);
}
let rgb_img = dynamic_img.into_rgb8();
let cache_dir = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/clipboard");
let _ = std::fs::create_dir_all(&cache_dir);
let cache_dir = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/clipboard");
let _ = std::fs::create_dir_all(&cache_dir);
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
if rgb_img
.save_with_format(&file_path, image::ImageFormat::Jpeg)
.is_ok()
{
let path_str = file_path.to_string_lossy().to_string();
out.insert("image_path".into(), json!(path_str));
if rgb_img
.save_with_format(&file_path, image::ImageFormat::Jpeg)
.is_ok()
{
let path_str = file_path.to_string_lossy().to_string();
out.insert("image_path".into(), json!(path_str));
// Read image bytes for base64 encoding if needed by vision
if let Ok(bytes) = std::fs::read(&file_path) {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
out.insert("image_base64".into(), json!(b64));
}
}
if let Ok(bytes) = std::fs::read(&file_path) {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
out.insert("image_base64".into(), json!(b64));
}
}
}
@@ -161,7 +240,6 @@ impl McpTool for ReadClipboardHandler {
let mut final_obj = result;
if let Some(b64) = final_obj.get("image_base64").and_then(|v| v.as_str()) {
let b64_str = b64.to_string();
// Remove huge base64 string from final user output
if let Some(obj) = final_obj.as_object_mut() {
obj.remove("image_base64");
}
server/src/handlers/workspaces.rs
+7 -7
View File
@@ -630,8 +630,8 @@ impl McpTool for SnippetsHandler {
let req: SnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
SnippetAction::Store => {
let name = req.query.or(req.id).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'query' or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'query' field and retry.".to_string())
let name = req.name.or(req.query).or(req.id).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'query', or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'name' or 'query' field and retry.".to_string())
})?;
let lang = req.language.unwrap_or_else(|| "text".to_string());
let code = req.code.unwrap_or_default();
@@ -648,7 +648,7 @@ impl McpTool for SnippetsHandler {
).await
}
SnippetAction::Search => {
let q = req.query.unwrap_or_default();
let q = req.query.or(req.name).unwrap_or_default();
if req.hybrid.unwrap_or(false) {
crate::handlers::meta::SearchSnippetsHybridHandler.execute(serde_json::json!({"query": q, "tags": req.tags}), state).await
} else {
@@ -656,14 +656,14 @@ impl McpTool for SnippetsHandler {
}
}
SnippetAction::Delete => {
let id = req.id.or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string())
let id = req.name.or(req.id).or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string())
})?;
DeleteSnippetHandler.execute(serde_json::json!({"name": id}), state).await
}
SnippetAction::Tag => {
let id = req.id.or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'id' or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string())
let id = req.name.or(req.id).or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string())
})?;
let tags = req.tags.unwrap_or_default();
TagSnippetHandler.execute(serde_json::json!({"name": id, "tags": tags}), state).await