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