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

fix(server): remove rwlock unwraps to prevent panic on lock poison and optimize omni search to serialize inner loops without allocations

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-21 16:49:21 +01:00
1 parent 1af21b3aec
commit 42d3638cf8
2 files changed
+46 -53
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers_v2/meta.rs
+43 -50
View File
@@ -354,80 +354,73 @@ impl McpTool for OmniSearchHandler {
vec![] vec![]
}; };
let mut kg = KnowledgeGraph::default(); let kg_json = state.read_graph(|full| {
let mut tasks = Vec::new(); let mut kg_entities = std::collections::HashMap::new();
let mut snippets = Vec::new();
let mut adrs = Vec::new();
state.read_graph(|full| {
for (id, doc_type, _, _, _) in &matches { for (id, doc_type, _, _, _) in &matches {
if doc_type == "entity" if doc_type == "entity"
&& let Some(e) = full.entities.get(id) && let Some(e) = full.entities.get(id)
{ {
kg.entities.insert(id.clone(), e.clone()); kg_entities.insert(id, e);
} }
} }
serde_json::to_value(&kg_entities).unwrap_or_default()
}); });
state.tasks.read_with(|all_tasks| {
for t in all_tasks { let tasks_json = state.tasks.read_with(|all_tasks| {
if matches let filtered: Vec<_> = all_tasks.iter().filter(|t| {
.iter() matches.iter().any(|(id, typ, _, _, _)| id == &t.id && typ == "task")
.any(|(id, typ, _, _, _)| id == &t.id && typ == "task") }).collect();
{ serde_json::to_value(&filtered).unwrap_or_default()
tasks.push(t.clone());
}
}
}); });
state.snippets.read_with(|all_snippets| {
for s in all_snippets { let snippets_json = state.snippets.read_with(|all_snippets| {
if matches let filtered: Vec<_> = all_snippets.iter().filter(|s| {
.iter() matches.iter().any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet")
.any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet") }).collect();
{ serde_json::to_value(&filtered).unwrap_or_default()
snippets.push(s.clone());
}
}
}); });
state.adrs.read_with(|all_adrs| {
for a in all_adrs { let adrs_json = state.adrs.read_with(|all_adrs| {
if matches let filtered: Vec<_> = all_adrs.iter().filter(|a| {
.iter() matches.iter().any(|(id, typ, _, _, _)| id == &a.id && typ == "adr")
.any(|(id, typ, _, _, _)| id == &a.id && typ == "adr") }).collect();
{ serde_json::to_value(&filtered).unwrap_or_default()
adrs.push(a.clone());
}
}
}); });
let q = req.query; let q = req.query;
let tech_debts = state.tech_debts.read_with(|debts| { let tech_debts_json = state.tech_debts.read_with(|debts| {
debts.iter().filter(|d| { let filtered: Vec<_> = debts.iter().filter(|d| {
req.namespace.as_ref().is_none_or(|ns| d.namespace == *ns) req.namespace.as_ref().is_none_or(|ns| d.namespace == *ns)
&& (contains_ignore_ascii_case(&d.description, &q) && (contains_ignore_ascii_case(&d.description, &q)
|| contains_ignore_ascii_case(&d.ideal_solution, &q)) || contains_ignore_ascii_case(&d.ideal_solution, &q))
}).cloned().collect::<Vec<_>>() }).collect();
serde_json::to_value(&filtered).unwrap_or_default()
}); });
let memos = state.handoff_memos.read_with(|memos| {
memos.iter().filter(|m| { let memos_json = state.handoff_memos.read_with(|memos| {
let filtered: Vec<_> = memos.iter().filter(|m| {
req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns) req.namespace.as_ref().is_none_or(|ns| m.namespace == *ns)
&& contains_ignore_ascii_case(&m.content, &q) && contains_ignore_ascii_case(&m.content, &q)
}).cloned().collect::<Vec<_>>() }).collect();
serde_json::to_value(&filtered).unwrap_or_default()
}); });
let error_fixes = state.error_fixes.read_with(|fixes| {
fixes.iter().filter(|f| { let error_fixes_json = state.error_fixes.read_with(|fixes| {
let filtered: Vec<_> = fixes.iter().filter(|f| {
contains_ignore_ascii_case(&f.signature, &q) contains_ignore_ascii_case(&f.signature, &q)
|| contains_ignore_ascii_case(&f.solution, &q) || contains_ignore_ascii_case(&f.solution, &q)
}).cloned().collect::<Vec<_>>() }).collect();
serde_json::to_value(&filtered).unwrap_or_default()
}); });
let report = serde_json::json!({ let report = serde_json::json!({
"knowledge_graph": kg.entities, "knowledge_graph": kg_json,
"tasks": tasks, "tasks": tasks_json,
"snippets": snippets, "snippets": snippets_json,
"adrs": adrs, "adrs": adrs_json,
"tech_debts": tech_debts, "tech_debts": tech_debts_json,
"handoff_memos": memos, "handoff_memos": memos_json,
"error_fixes": error_fixes "error_fixes": error_fixes_json
}); });
Ok(report.to_string()) Ok(report.to_string())
} }
server/src/main.rs
+3 -3
View File
@@ -446,7 +446,7 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
state state
.clients .clients
.write() .write()
.unwrap() .unwrap_or_else(|e| e.into_inner())
.insert(session_id.clone(), tx.clone()); .insert(session_id.clone(), tx.clone());
let (mut sender, mut receiver) = socket.split(); let (mut sender, mut receiver) = socket.split();
@@ -501,7 +501,7 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
let senders: Vec<_> = state_clone let senders: Vec<_> = state_clone
.clients .clients
.read() .read()
.unwrap() .unwrap_or_else(|e| e.into_inner())
.iter() .iter()
.filter_map(|(id, tx)| { .filter_map(|(id, tx)| {
if id != &session_id_clone { if id != &session_id_clone {
@@ -524,7 +524,7 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
let tx_opt = state_clone let tx_opt = state_clone
.clients .clients
.read() .read()
.unwrap() .unwrap_or_else(|e| e.into_inner())
.get(&session_id_clone) .get(&session_id_clone)
.cloned(); .cloned();
if let Some(client_tx) = tx_opt { if let Some(client_tx) = tx_opt {