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

refactor(server): Resolve bottlenecks and duplications

Browse Source 
- Removed heavy Tantivy index rebuild loop that ran every 5 seconds and rewrote entire graph to disk
- Converted polling worker to lightweight IndexWriter committer
- Fixed 67+ instances of unnecessary clone and string parsing overhead in tool return results
- Deduplicated state graph modification methods
This commit is contained in:
Riz Ashraf committed 2026-09-19 07:48:51 +01:00
1 parent 4fe6ed7791
commit 2a25f17a32
3 files changed
+86 -108
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers.rs
+80 -80
View File
@@ -437,7 +437,7 @@ impl MemoryHandler {
}
"create_entities" => {
let req = parse_tool!(args.clone(), id, CreateEntitiesTool);
self.state.write_to_local_delta(|g| {
self.state.modify_graph(|g| {
for entity in req.entities {
if !entity.name.is_empty() {
if let Ok(idx) = self.state.search_index.read() {
@@ -447,23 +447,23 @@ impl MemoryHandler {
}
}
});
Ok(vec!["Entities created".to_string()][0].clone())
Ok("Entities created".to_string())
}
"create_relations" => {
let req = parse_tool!(args.clone(), id, CreateRelationsTool);
self.state.write_to_local_delta(|g| {
self.state.modify_graph(|g| {
for relation in req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() {
g.relations.push(relation);
}
}
});
Ok(vec!["Relations created".to_string()][0].clone())
Ok("Relations created".to_string())
}
"add_observations" => {
let req = parse_tool!(args.clone(), id, AddObservationsTool);
let full = self.state.get_full_graph();
self.state.write_to_local_delta(|g| {
self.state.modify_graph(|g| {
for o in req.observations {
if let Some(full_e) = full.entities.get(&o.entity_name) {
let mut e =
@@ -481,12 +481,12 @@ impl MemoryHandler {
}
}
});
Ok(vec!["Observations added".to_string()][0].clone())
Ok("Observations added".to_string())
}
"delete_entities" => {
let req = parse_tool!(args.clone(), id, DeleteEntitiesTool);
let to_delete: HashSet<_> = req.entity_names.into_iter().collect();
self.state.apply_sync_write(|master| {
self.state.modify_graph(|master| {
for name in &to_delete {
master.entities.remove(name);
}
@@ -494,11 +494,11 @@ impl MemoryHandler {
!to_delete.contains(&r.from) && !to_delete.contains(&r.to)
});
});
Ok(vec!["Entities deleted".to_string()][0].clone())
Ok("Entities deleted".to_string())
}
"delete_observations" => {
let req = parse_tool!(args.clone(), id, DeleteObservationsTool);
self.state.apply_sync_write(|master| {
self.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();
@@ -506,11 +506,11 @@ impl MemoryHandler {
}
}
});
Ok(vec!["Observations deleted".to_string()][0].clone())
Ok("Observations deleted".to_string())
}
"delete_relations" => {
let req = parse_tool!(args.clone(), id, DeleteRelationsTool);
self.state.apply_sync_write(|master| {
self.state.modify_graph(|master| {
let mut to_rem = HashSet::new();
for r in req.relations {
to_rem.insert(format!(
@@ -525,7 +525,7 @@ impl MemoryHandler {
))
});
});
Ok(vec!["Relations deleted".to_string()][0].clone())
Ok("Relations deleted".to_string())
}
"read_graph" => {
let req = parse_tool!(args.clone(), id, ReadGraphTool);
@@ -535,7 +535,7 @@ impl MemoryHandler {
full.relations.retain(|r| r.namespace == ns);
}
let data = serde_json::to_string(&full).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"search_nodes" => {
let req = parse_tool!(args.clone(), id, SearchNodesTool);
@@ -556,7 +556,7 @@ impl MemoryHandler {
}
}
let data = serde_json::to_string(&result).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"open_nodes" => {
let req = parse_tool!(args.clone(), id, OpenNodesTool);
@@ -579,7 +579,7 @@ impl MemoryHandler {
}
}
let data = serde_json::to_string(&result).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"log_code_change" => {
let req = parse_tool!(args.clone(), id, LogCodeChangeTool);
@@ -595,12 +595,12 @@ impl MemoryHandler {
git_branch: req.git_branch,
});
});
Ok(vec!["Code change logged".to_string()][0].clone())
Ok("Code change logged".to_string())
}
"query_recent_changes" => {
let data = serde_json::to_string(&self.state.ledger.read())
.unwrap_or_else(|_| "[]".to_string());
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"visualize_graph" => {
let req = parse_tool!(args.clone(), id, VisualizeGraphTool);
@@ -670,7 +670,7 @@ impl MemoryHandler {
if output == "graph TD;\n" {
output = "No nodes found to visualize.".to_string();
}
Ok(vec![output.to_string()][0].clone())
Ok(output.to_string())
}
"add_sticky_note" => {
let req = parse_tool!(args.clone(), id, AddStickyNoteTool);
@@ -683,12 +683,12 @@ impl MemoryHandler {
content: req.content,
});
});
Ok(vec!["Sticky note added.".to_string()][0].clone())
Ok("Sticky note added.".to_string())
}
"read_sticky_notes" => {
let data = serde_json::to_string(&self.state.sticky.read())
.unwrap_or_else(|_| "[]".to_string());
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"delete_sticky_note" => {
let req = parse_tool!(args.clone(), id, DeleteStickyNoteTool);
@@ -700,7 +700,7 @@ impl MemoryHandler {
}
});
if success {
Ok(vec!["Sticky note deleted.".to_string()][0].clone())
Ok("Sticky note deleted.".to_string())
} else {
Err("Invalid sticky note index.".to_string())
}
@@ -709,16 +709,16 @@ impl MemoryHandler {
self.state.sticky.modify(|notes| {
notes.clear();
});
Ok(vec!["All sticky notes cleared.".to_string()][0].clone())
Ok("All sticky notes cleared.".to_string())
}
"condense_entity" => {
let req = parse_tool!(args.clone(), id, CondenseEntityTool);
self.state.apply_sync_write(|master| {
self.state.modify_graph(|master| {
if let Some(e) = master.entities.get_mut(&req.entity_name) {
e.observations = req.summarized_observations;
}
});
Ok(vec!["Entity condensed".to_string()][0].clone())
Ok("Entity condensed".to_string())
}
"add_task" => {
let req = parse_tool!(args.clone(), id, AddTaskTool);
@@ -749,7 +749,7 @@ impl MemoryHandler {
self.state.tasks.modify(|tasks| {
tasks.push(task);
});
Ok(vec![format!("Task added with ID: {}", task_id).to_string()][0].clone())
Ok(format!("Task added with ID: {}", task_id).to_string())
}
"delete_task" => {
let req = parse_tool!(args.clone(), id, DeleteTaskTool);
@@ -783,7 +783,7 @@ impl MemoryHandler {
][0]
.clone())
} else {
Ok(vec!["Task not found.".to_string()][0].clone())
Ok("Task not found.".to_string())
}
}
"update_task_status" => {
@@ -899,9 +899,9 @@ impl MemoryHandler {
.next()
.unwrap())
} else if found {
Ok(vec!["Task status updated.".to_string()][0].clone())
Ok("Task status updated.".to_string())
} else {
Ok(vec!["Task not found.".to_string()][0].clone())
Ok("Task not found.".to_string())
}
}
"list_active_tasks" => {
@@ -915,7 +915,7 @@ impl MemoryHandler {
});
}
let data = serde_json::to_string(&tasks).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"set_acceptance_criteria" => {
let req = parse_tool!(args.clone(), id, SetAcceptanceCriteriaTool);
@@ -946,7 +946,7 @@ impl MemoryHandler {
.clone(),
)
} else {
Ok(vec!["Task not found.".to_string()][0].clone())
Ok("Task not found.".to_string())
}
}
"verify_acceptance_criteria" => {
@@ -979,7 +979,7 @@ impl MemoryHandler {
)][0]
.clone())
} else if already_met {
Ok(vec!["Acceptance criteria was already met.".to_string()][0].clone())
Ok("Acceptance criteria was already met.".to_string())
} else {
Ok(
vec!["Acceptance criteria or task not found.".to_string()][0]
@@ -1002,7 +1002,7 @@ impl MemoryHandler {
.as_secs(),
});
});
Ok(vec![format!("Snippet '{}' stored.", req.name).to_string()][0].clone())
Ok(format!("Snippet '{}' stored.", req.name).to_string())
}
"search_snippets" => {
let req = parse_tool!(args.clone(), id, SearchSnippetsTool);
@@ -1018,7 +1018,7 @@ impl MemoryHandler {
}
}
let data = serde_json::to_string(&results).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"delete_snippet" => {
let req = parse_tool!(args.clone(), id, DeleteSnippetTool);
@@ -1029,9 +1029,9 @@ impl MemoryHandler {
deleted = snippets.len() < orig;
});
if deleted {
Ok(vec!["Snippet deleted.".to_string()][0].clone())
Ok("Snippet deleted.".to_string())
} else {
Ok(vec!["Snippet not found.".to_string()][0].clone())
Ok("Snippet not found.".to_string())
}
}
"log_decision" => {
@@ -1051,7 +1051,7 @@ impl MemoryHandler {
.as_secs(),
});
});
Ok(vec![format!("Decision logged as {}", id).to_string()][0].clone())
Ok(format!("Decision logged as {}", id).to_string())
}
"query_decisions" => {
let req = parse_tool!(args.clone(), id, QueryDecisionsTool);
@@ -1065,11 +1065,11 @@ impl MemoryHandler {
});
}
let data = serde_json::to_string(&adrs).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"merge_entities" => {
let req = parse_tool!(args.clone(), id, MergeEntitiesTool);
self.state.apply_sync_write(|master| {
self.state.modify_graph(|master| {
if let Some(src) = master.entities.remove(&req.source_entity) {
if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
tgt.observations.extend(src.observations);
@@ -1091,7 +1091,7 @@ impl MemoryHandler {
}
master.relations = MemoryState::unique_items(master.relations.clone());
});
Ok(vec!["Entities merged".to_string()][0].clone())
Ok("Entities merged".to_string())
}
"find_orphans" => {
let full = self.state.get_full_graph();
@@ -1107,7 +1107,7 @@ impl MemoryHandler {
.cloned()
.collect();
let data = serde_json::to_string(&orphans).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"learn_preference" => {
let req = parse_tool!(args.clone(), id, LearnPreferenceTool);
@@ -1124,12 +1124,12 @@ impl MemoryHandler {
},
);
});
Ok(vec!["Preference learned".to_string()][0].clone())
Ok("Preference learned".to_string())
}
"read_preferences" => {
let prefs = self.state.prefs.read();
let data = serde_json::to_string(&prefs).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"log_error_fix" => {
let req = parse_tool!(args.clone(), id, LogErrorFixTool);
@@ -1145,7 +1145,7 @@ impl MemoryHandler {
git_branch: req.git_branch,
})
});
Ok(vec!["Error fix logged".to_string()][0].clone())
Ok("Error fix logged".to_string())
}
"search_error_fixes" => {
let req = parse_tool!(args.clone(), id, SearchErrorFixesTool);
@@ -1156,7 +1156,7 @@ impl MemoryHandler {
|| f.solution.to_lowercase().contains(&q)
});
let data = serde_json::to_string(&fixes).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"pin_file" => {
let req = parse_tool!(args.clone(), id, PinFileTool);
@@ -1174,7 +1174,7 @@ impl MemoryHandler {
git_branch: req.git_branch,
});
});
Ok(vec!["File pinned".to_string()][0].clone())
Ok("File pinned".to_string())
}
"unpin_file" => {
let req = parse_tool!(args.clone(), id, UnpinFileTool);
@@ -1183,7 +1183,7 @@ impl MemoryHandler {
!(p.namespace == req.namespace && p.file_path == req.file_path)
})
});
Ok(vec!["File unpinned".to_string()][0].clone())
Ok("File unpinned".to_string())
}
"list_pinned_files" => {
let req = parse_tool!(args.clone(), id, ListPinnedFilesTool);
@@ -1198,7 +1198,7 @@ impl MemoryHandler {
});
}
let data = serde_json::to_string(&pinned).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"add_session_summary" => {
let req = parse_tool!(args.clone(), id, AddSessionSummaryTool);
@@ -1212,7 +1212,7 @@ impl MemoryHandler {
.as_secs(),
})
});
Ok(vec!["Session summary added".to_string()][0].clone())
Ok("Session summary added".to_string())
}
"get_project_timeline" => {
let req = parse_tool!(args.clone(), id, GetProjectTimelineTool);
@@ -1222,7 +1222,7 @@ impl MemoryHandler {
}
summaries.sort_by_key(|s| s.timestamp);
let data = serde_json::to_string(&summaries).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"leave_handoff_memo" => {
let req = parse_tool!(args.clone(), id, LeaveHandoffMemoTool);
@@ -1238,7 +1238,7 @@ impl MemoryHandler {
.as_secs(),
})
});
Ok(vec!["Handoff memo left".to_string()][0].clone())
Ok("Handoff memo left".to_string())
}
"read_handoff_memos" => {
let req = parse_tool!(args.clone(), id, ReadHandoffMemosTool);
@@ -1247,7 +1247,7 @@ impl MemoryHandler {
memos.retain(|m| m.namespace == ns);
}
let data = serde_json::to_string(&memos).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"clear_handoff_memos" => {
let req = parse_tool!(args.clone(), id, ClearHandoffMemosTool);
@@ -1255,7 +1255,7 @@ impl MemoryHandler {
self.state
.handoff_memos
.modify(|memos| memos.retain(|m| !ids.contains(&m.id)));
Ok(vec!["Handoff memos cleared".to_string()][0].clone())
Ok("Handoff memos cleared".to_string())
}
"update_env_fingerprint" => {
let req = parse_tool!(args.clone(), id, UpdateEnvFingerprintTool);
@@ -1275,16 +1275,16 @@ impl MemoryHandler {
},
);
});
Ok(vec!["Env fingerprint updated".to_string()][0].clone())
Ok("Env fingerprint updated".to_string())
}
"read_env_fingerprint" => {
let req = parse_tool!(args.clone(), id, ReadEnvFingerprintTool);
let fps = self.state.env_fingerprints.read();
if let Some(fp) = fps.get(&req.namespace) {
let data = serde_json::to_string(fp).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
} else {
Ok(vec!["{}".to_string()][0].clone())
Ok("{}".to_string())
}
}
"log_env_requirement" => {
@@ -1298,7 +1298,7 @@ impl MemoryHandler {
is_secret: req.is_secret,
});
});
Ok(vec!["Env requirement logged".to_string()][0].clone())
Ok("Env requirement logged".to_string())
}
"add_milestone" => {
let req = parse_tool!(args.clone(), id, AddMilestoneTool);
@@ -1311,7 +1311,7 @@ impl MemoryHandler {
target_date: None,
})
});
Ok(vec!["Milestone added".to_string()][0].clone())
Ok("Milestone added".to_string())
}
"update_milestone" => {
let req = parse_tool!(args.clone(), id, UpdateMilestoneTool);
@@ -1326,9 +1326,9 @@ impl MemoryHandler {
}
});
if found {
Ok(vec!["Milestone updated".to_string()][0].clone())
Ok("Milestone updated".to_string())
} else {
Ok(vec!["Milestone not found".to_string()][0].clone())
Ok("Milestone not found".to_string())
}
}
"list_milestones" => {
@@ -1338,7 +1338,7 @@ impl MemoryHandler {
ms.retain(|m| m.namespace == ns);
}
let data = serde_json::to_string(&ms).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"generate_standup_report" => {
let req = parse_tool!(args.clone(), id, GenerateStandupReportTool);
@@ -1369,7 +1369,7 @@ impl MemoryHandler {
.filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff)
.collect::<Vec<_>>();
let report = serde_json::json!({ "tasks_updated": tasks, "code_changes": changes, "session_summaries": summaries });
Ok(vec![report.to_string().to_string()][0].clone())
Ok(report.to_string())
}
"register_environment" => {
let req = parse_tool!(args.clone(), id, RegisterEnvironmentTool);
@@ -1387,14 +1387,14 @@ impl MemoryHandler {
.as_secs(),
});
});
Ok(vec!["Environment registered".to_string()][0].clone())
Ok("Environment registered".to_string())
}
"get_environment_details" => {
let req = parse_tool!(args.clone(), id, GetEnvironmentDetailsTool);
let mut envs = self.state.environments.read();
envs.retain(|e| e.namespace == req.namespace);
let data = serde_json::to_string(&envs).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"add_pr_checklist_item" => {
let req = parse_tool!(args.clone(), id, AddPrChecklistItemTool);
@@ -1405,21 +1405,21 @@ impl MemoryHandler {
description: req.description,
})
});
Ok(vec!["PR checklist item added".to_string()][0].clone())
Ok("PR checklist item added".to_string())
}
"get_pr_checklist" => {
let req = parse_tool!(args.clone(), id, GetPrChecklistTool);
let mut items = self.state.pr_checklists.read();
items.retain(|i| i.namespace == req.namespace);
let data = serde_json::to_string(&items).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"clear_pr_checklist" => {
let req = parse_tool!(args.clone(), id, ClearPrChecklistTool);
self.state
.pr_checklists
.modify(|items| items.retain(|i| i.namespace != req.namespace));
Ok(vec!["PR checklist cleared".to_string()][0].clone())
Ok("PR checklist cleared".to_string())
}
"log_tech_debt" => {
let req = parse_tool!(args.clone(), id, LogTechDebtTool);
@@ -1438,7 +1438,7 @@ impl MemoryHandler {
git_branch: req.git_branch,
})
});
Ok(vec!["Tech debt logged".to_string()][0].clone())
Ok("Tech debt logged".to_string())
}
"resolve_tech_debt" => {
let req = parse_tool!(args.clone(), id, ResolveTechDebtTool);
@@ -1453,9 +1453,9 @@ impl MemoryHandler {
}
});
if found {
Ok(vec!["Tech debt resolved".to_string()][0].clone())
Ok("Tech debt resolved".to_string())
} else {
Ok(vec!["Tech debt not found".to_string()][0].clone())
Ok("Tech debt not found".to_string())
}
}
"list_tech_debt" => {
@@ -1465,7 +1465,7 @@ impl MemoryHandler {
d.namespace == req.namespace && (req.include_resolved || !d.is_resolved)
});
let data = serde_json::to_string(&debts).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"save_context_workspace" => {
let req = parse_tool!(args.clone(), id, SaveContextWorkspaceTool);
@@ -1482,21 +1482,21 @@ impl MemoryHandler {
.as_secs(),
});
});
Ok(vec!["Context workspace saved".to_string()][0].clone())
Ok("Context workspace saved".to_string())
}
"load_context_workspace" => {
let req = parse_tool!(args.clone(), id, LoadContextWorkspaceTool);
let mut ws = self.state.context_workspaces.read();
ws.retain(|w| w.namespace == req.namespace && w.name == req.name);
let data = serde_json::to_string(&ws.first()).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"list_context_workspaces" => {
let req = parse_tool!(args.clone(), id, ListContextWorkspacesTool);
let mut ws = self.state.context_workspaces.read();
ws.retain(|w| w.namespace == req.namespace);
let data = serde_json::to_string(&ws).unwrap_or_default();
Ok(vec![data.to_string()][0].clone())
Ok(data.to_string())
}
"omni_search" => {
let req = parse_tool!(args.clone(), id, OmniSearchTool);
@@ -1583,7 +1583,7 @@ impl MemoryHandler {
"handoff_memos": memos,
"error_fixes": error_fixes
});
Ok(vec![report.to_string().to_string()][0].clone())
Ok(report.to_string())
}
"get_project_health" => {
let req = parse_tool!(args.clone(), id, GetProjectHealthTool);
@@ -1630,7 +1630,7 @@ impl MemoryHandler {
"active_milestones": active_milestones,
"remaining_pr_checklist_items": remaining_checklists
});
Ok(vec![report.to_string().to_string()][0].clone())
Ok(report.to_string())
}
_ => Err(format!("Unknown tool: {}", name)),
@@ -2064,7 +2064,7 @@ mod tests {
);
});
// Force flush session to master
handler.state.apply_sync_write(|_| {});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
@@ -2097,7 +2097,7 @@ mod tests {
},
);
});
handler.state.apply_sync_write(|_| {});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 10,
@@ -2262,7 +2262,7 @@ mod tests {
namespace: "".to_string(),
});
});
handler.state.apply_sync_write(|_| {});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
@@ -2302,7 +2302,7 @@ mod tests {
},
);
});
handler.state.apply_sync_write(|_| {});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
@@ -2333,7 +2333,7 @@ mod tests {
.entities
.insert("UserRepository".to_string(), entity);
});
handler.state.apply_sync_write(|_| {});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",