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

refactor: apply rust best practices and fix memory optimizations

Browse Source 
- Optimized memory allocation in router.rs by offloading JSON serialization to spawn_blocking and using references.
- Prevented full graph duplication on startup in state.rs index rebuild.
- Eliminated massive String allocations in QueryGraphPathHandler BFS loops.
- Avoided temporary Strings in VisualizeGraphHandler via inline writing.
- Fixed O(N) full-graph deduplication in MergeEntitiesHandler to scale efficiently.
This commit is contained in:
Riz Ashraf committed 2026-09-22 22:01:38 +01:00
1 parent 251757f8fc
commit 37003be620
17 files changed
+399 -198
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/graph.rs
+113 -56
View File
@@ -31,7 +31,7 @@ impl McpTool for QueryGraphPathHandler {
let max_depth = req.max_depth.unwrap_or(5);
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, std::borrow::Cow<'_, str>)> =
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
std::collections::HashMap::new();
queue.push_back(req.start_node.as_str());
@@ -54,10 +54,7 @@ impl McpTool for QueryGraphPathHandler {
visited.insert(rel.to.as_str());
parents.insert(
rel.to.as_str(),
(
current,
std::borrow::Cow::Borrowed(rel.relation_type.as_str()),
),
(current, rel.relation_type.as_str(), false),
);
queue.push_back(rel.to.as_str());
nodes_at_next_depth += 1;
@@ -65,13 +62,7 @@ impl McpTool for QueryGraphPathHandler {
visited.insert(rel.from.as_str());
parents.insert(
rel.from.as_str(),
(
current,
std::borrow::Cow::Owned(format!(
"inverse({})",
rel.relation_type
)),
),
(current, rel.relation_type.as_str(), true),
);
queue.push_back(rel.from.as_str());
nodes_at_next_depth += 1;
@@ -89,8 +80,12 @@ impl McpTool for QueryGraphPathHandler {
let mut path = Vec::new();
let mut curr = req.end_node.as_str();
while curr != req.start_node {
if let Some((parent, rel_type)) = parents.get(&curr) {
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
if *is_inverse {
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
} else {
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
}
curr = parent;
} else {
break;
@@ -406,7 +401,6 @@ impl McpTool for VisualizeGraphHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let query = req.query.unwrap_or_default();
use std::fmt::Write;
let mut output = state.read_graph(|full| {
let mut included: HashSet<&str> = HashSet::new();
let mut to_draw = Vec::new();
@@ -444,36 +438,33 @@ impl McpTool for VisualizeGraphHandler {
let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
out.push_str("graph TD;\n");
let sanitize = |s: &str, id_mode: bool| -> String {
let mut o = String::with_capacity(s.len());
let sanitize_to = |out_str: &mut String, s: &str, id_mode: bool| {
for c in s.chars() {
if c != '"' && c != '(' && c != ')' {
if id_mode && (c == ' ' || c == '-' || c == '.') {
o.push('_');
out_str.push('_');
} else {
o.push(c);
out_str.push(c);
}
}
}
o
};
for name in &included {
let _ = writeln!(
out,
" id_{}[\"{}\"];",
sanitize(name, true),
sanitize(name, false)
);
out.push_str(" id_");
sanitize_to(&mut out, name, true);
out.push_str("[\"");
sanitize_to(&mut out, name, false);
out.push_str("\"];\n");
}
for r in to_draw {
let _ = writeln!(
out,
" id_{}-->|\"{}\"|id_{};",
sanitize(&r.from, true),
r.relation_type.replace("\"", ""),
sanitize(&r.to, true)
);
out.push_str(" id_");
sanitize_to(&mut out, &r.from, true);
out.push_str("-->|\"");
out.push_str(&r.relation_type.replace("\"", ""));
out.push_str("\"|id_");
sanitize_to(&mut out, &r.to, true);
out.push_str(";\n");
}
out
});
@@ -532,15 +523,21 @@ impl McpTool for MergeEntitiesHandler {
master.entities.insert(req.target_entity.clone(), new_tgt);
}
}
for r in &mut master.relations {
let mut seen = std::collections::HashSet::new();
master.relations.retain_mut(|r| {
if r.from == req.source_entity {
r.from = req.target_entity.clone();
}
if r.to == req.source_entity {
r.to = req.target_entity.clone();
}
}
MemoryState::deduplicate(&mut master.relations);
if r.from == req.target_entity || r.to == req.target_entity {
seen.insert(r.clone())
} else {
true
}
});
});
Ok("Entities merged".to_string())
}
@@ -581,24 +578,24 @@ use crate::handlers::utils::*;
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
use serde_json::json;
use tempfile::tempdir;
#[tokio::test]
async fn test_create_and_read_entities() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let create_handler = CreateEntitiesHandler;
let args = json!({
"entities": [
{"name": "Alice", "entityType": "Person", "observations": ["Likes Bob"]}
]
});
let res = create_handler.execute(args, state.clone()).await.unwrap();
assert_eq!(res, "Entities created");
// Ensure graph contains the entity
state.graph.read_with(|g| {
assert!(g.entities.contains_key("Alice"));
@@ -610,11 +607,29 @@ mod tests {
async fn test_create_relations() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
// Needs entities first
state.graph.modify(|g| {
g.entities.insert("A".to_string(), crate::models::Entity { name: "A".to_string(), entity_type: "Node".to_string(), observations: vec![], namespace: "global".to_string(), git_branch: None });
g.entities.insert("B".to_string(), crate::models::Entity { name: "B".to_string(), entity_type: "Node".to_string(), observations: vec![], namespace: "global".to_string(), git_branch: None });
g.entities.insert(
"A".to_string(),
crate::models::Entity {
name: "A".to_string(),
entity_type: "Node".to_string(),
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
},
);
g.entities.insert(
"B".to_string(),
crate::models::Entity {
name: "B".to_string(),
entity_type: "Node".to_string(),
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
},
);
});
let handler = CreateRelationsHandler;
@@ -631,10 +646,19 @@ mod tests {
async fn test_observations_and_reads() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
// Need entity first
state.graph.modify(|g| {
g.entities.insert("A".to_string(), crate::models::Entity { name: "A".to_string(), entity_type: "Node".to_string(), observations: vec![], namespace: "global".to_string(), git_branch: None });
g.entities.insert(
"A".to_string(),
crate::models::Entity {
name: "A".to_string(),
entity_type: "Node".to_string(),
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
},
);
});
let add_obs = AddObservationsHandler;
@@ -647,15 +671,24 @@ mod tests {
assert_eq!(res1, "Observations added");
let read_graph = ReadGraphHandler;
let res2 = read_graph.execute(json!({"namespace": "global"}), state.clone()).await.unwrap();
let res2 = read_graph
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(res2.contains("Obs 1"));
assert!(res2.contains("Obs 2"));
let del_entity = DeleteEntitiesHandler;
let res4 = del_entity.execute(json!({"entityNames": ["A"]}), state.clone()).await.unwrap();
let res4 = del_entity
.execute(json!({"entityNames": ["A"]}), state.clone())
.await
.unwrap();
assert_eq!(res4, "Entities deleted");
let res5 = read_graph.execute(json!({"namespace": "global"}), state.clone()).await.unwrap();
let res5 = read_graph
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(!res5.contains("A"));
}
@@ -663,7 +696,7 @@ mod tests {
async fn test_advanced_graph_operations() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let create_handler = CreateEntitiesHandler;
let args_ent = json!({
"entities": [
@@ -671,7 +704,10 @@ mod tests {
{"name": "Y", "entityType": "File", "observations": ["Obs Y"], "namespace": "global"}
]
});
create_handler.execute(args_ent, state.clone()).await.unwrap();
create_handler
.execute(args_ent, state.clone())
.await
.unwrap();
let rel_handler = CreateRelationsHandler;
let args_rel = json!({
@@ -682,24 +718,45 @@ mod tests {
rel_handler.execute(args_rel, state.clone()).await.unwrap();
let read_handler = ReadGraphHandler;
let res_read = read_handler.execute(json!({"namespace": "global"}), state.clone()).await.unwrap();
let res_read = read_handler
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(res_read.contains("X"));
assert!(res_read.contains("depends_on"));
let open_handler = OpenNodesHandler;
let res_open = open_handler.execute(json!({"names": ["X"]}), state.clone()).await.unwrap();
let res_open = open_handler
.execute(json!({"names": ["X"]}), state.clone())
.await
.unwrap();
assert!(res_open.contains("Y"));
let viz_handler = VisualizeGraphHandler;
let res_viz = viz_handler.execute(json!({"query": "X"}), state.clone()).await.unwrap();
assert!(res_viz.len() > 0);
let res_viz = viz_handler
.execute(json!({"query": "X"}), state.clone())
.await
.unwrap();
assert!(!res_viz.is_empty());
let condense = CondenseEntityHandler;
let res_cond = condense.execute(json!({"entityName": "X", "summarized_observations": ["X condensed"]}), state.clone()).await.unwrap();
let res_cond = condense
.execute(
json!({"entityName": "X", "summarized_observations": ["X condensed"]}),
state.clone(),
)
.await
.unwrap();
assert_eq!(res_cond, "Entity condensed");
let merge = MergeEntitiesHandler;
let res_merge = merge.execute(json!({"sourceEntity": "X", "targetEntity": "Y"}), state.clone()).await.unwrap();
let res_merge = merge
.execute(
json!({"sourceEntity": "X", "targetEntity": "Y"}),
state.clone(),
)
.await
.unwrap();
assert_eq!(res_merge, "Entities merged");
let orphans = FindOrphansHandler;