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/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)
}