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

fix(mcp-memory,nvim-core): eliminate lock contention in get_nvim_connection and resolve CPU bound loop in graph.rs

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-11 18:21:58 +01:00
1 parent e91b53fdcc
commit 0afba74b4c
2 files changed
+104 -102
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nvim-core/src/lib.rs
+4 -3
View File
@@ -391,6 +391,8 @@ fn handle_nvim_notification(params: &[rmpv::Value]) {
} }
async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> { async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
let target_socket = read_active_socket_file().await.unwrap_or_default();
// Fast path: If NVIM_CONN is already active and healthy, check if active_nvim.txt matches // Fast path: If NVIM_CONN is already active and healthy, check if active_nvim.txt matches
// without spawning any external nvim.exe child process! // without spawning any external nvim.exe child process!
{ {
@@ -400,9 +402,8 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
{ {
let active_lock = ACTIVE_SOCKET.lock().await; let active_lock = ACTIVE_SOCKET.lock().await;
if let Some(ref current_socket) = *active_lock if let Some(ref current_socket) = *active_lock
&& let Ok(target) = read_active_socket_file().await && !target_socket.is_empty()
&& !target.is_empty() && &target_socket == current_socket
&& &target == current_socket
{ {
return Ok(sender.clone()); return Ok(sender.clone());
} }
server/src/handlers/graph.rs
+100 -99
View File
@@ -632,119 +632,120 @@ impl McpTool for ReadGraphHandler {
let limit = req.limit; let limit = req.limit;
let offset = req.offset.unwrap_or(0); let offset = req.offset.unwrap_or(0);
let data = state.read_graph(|full| -> crate::error::Result<String> { tokio::task::spawn_blocking(move || {
let filtered_entities: Vec<(&String, &crate::models::Entity)> = full state.read_graph(|full| -> crate::error::Result<String> {
.entities let filtered_entities: Vec<(&String, &crate::models::Entity)> = full
.iter() .entities
.filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns)) .iter()
.filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et)) .filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns))
.collect(); .filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et))
let total_count = filtered_entities.len();
let (entities_map, relations_vec, meta_json) = if limit.is_some() || req.offset.is_some() {
let effective_limit = limit.unwrap_or(50);
let paged_entities: std::collections::HashMap<&String, &crate::models::Entity> = filtered_entities
.into_iter()
.skip(offset)
.take(effective_limit)
.collect(); .collect();
let has_more = offset + paged_entities.len() < total_count; let total_count = filtered_entities.len();
let relevant_relations: Vec<&crate::models::Relation> = full let (entities_map, relations_vec, meta_json) = if limit.is_some() || req.offset.is_some() {
.relations let effective_limit = limit.unwrap_or(50);
.iter() let paged_entities: std::collections::HashMap<&String, &crate::models::Entity> = filtered_entities
.filter(|r| { .into_iter()
let ns_match = req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns); .skip(offset)
ns_match && (paged_entities.contains_key(&r.from) || paged_entities.contains_key(&r.to)) .take(effective_limit)
.collect();
let has_more = offset + paged_entities.len() < total_count;
let relevant_relations: Vec<&crate::models::Relation> = full
.relations
.iter()
.filter(|r| {
let ns_match = req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns);
ns_match && (paged_entities.contains_key(&r.from) || paged_entities.contains_key(&r.to))
})
.collect();
let meta = serde_json::json!({
"total_count": total_count,
"offset": offset,
"limit": effective_limit,
"has_more": has_more,
});
(paged_entities, relevant_relations, Some(meta))
} else {
let ent_map: std::collections::HashMap<&String, &crate::models::Entity> =
filtered_entities.into_iter().collect();
let rel_vec: Vec<&crate::models::Relation> = full
.relations
.iter()
.filter(|r| req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns))
.collect();
(ent_map, rel_vec, None)
};
let serialized = if let Some(meta) = &meta_json {
serde_json::json!({
"entities": entities_map,
"relations": relations_vec,
"total_count": meta["total_count"],
"offset": meta["offset"],
"limit": meta["limit"],
"has_more": meta["has_more"],
}) })
.collect(); } else {
serde_json::json!({
"entities": entities_map,
"relations": relations_vec,
})
};
let meta = serde_json::json!({ let mut result_json = serde_yaml::to_string(&serialized)?;
"total_count": total_count,
"offset": offset,
"limit": effective_limit,
"has_more": has_more,
});
(paged_entities, relevant_relations, Some(meta))
} else {
let ent_map: std::collections::HashMap<&String, &crate::models::Entity> =
filtered_entities.into_iter().collect();
let rel_vec: Vec<&crate::models::Relation> = full
.relations
.iter()
.filter(|r| req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns))
.collect();
(ent_map, rel_vec, None)
};
let serialized = if let Some(meta) = &meta_json { if let Some(max_tok) = max_tokens {
serde_json::json!({ let max_chars = max_tok * 4;
"entities": entities_map, if result_json.len() > max_chars {
"relations": relations_vec, let total_ent = entities_map.len();
"total_count": meta["total_count"], let total_rel = relations_vec.len();
"offset": meta["offset"], let frac = (max_chars as f64) / (result_json.len() as f64);
"limit": meta["limit"], let mut keep_count = (total_ent as f64 * frac * 0.85).max(1.0) as usize;
"has_more": meta["has_more"],
})
} else {
serde_json::json!({
"entities": entities_map,
"relations": relations_vec,
})
};
let mut result_json = serde_yaml::to_string(&serialized)?; let ent_vec: Vec<_> = entities_map.into_iter().collect();
if let Some(max_tok) = max_tokens { while keep_count > 0 {
let max_chars = max_tok * 4; let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect();
if result_json.len() > max_chars { let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect();
let total_ent = entities_map.len();
let total_rel = relations_vec.len();
let frac = (max_chars as f64) / (result_json.len() as f64);
let mut keep_count = (total_ent as f64 * frac * 0.85).max(1.0) as usize;
let ent_vec: Vec<_> = entities_map.into_iter().collect(); let mut truncated_val = serde_json::json!({
"entities": sliced_ent,
"relations": sliced_rel,
"_meta": {
"truncated": true,
"total_entities": total_ent,
"returned_entities": sliced_ent.len(),
"total_relations": total_rel,
"returned_relations": sliced_rel.len(),
"note": "Output structurally truncated to fit within max_tokens. Use pagination (limit, offset), search_nodes, or get_subgraph for targeted discovery."
}
});
while keep_count > 0 { if let Some(meta) = &meta_json
let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect(); && let Some(obj) = truncated_val.as_object_mut() {
let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect(); obj.insert("total_count".to_string(), meta["total_count"].clone());
obj.insert("offset".to_string(), meta["offset"].clone());
obj.insert("limit".to_string(), meta["limit"].clone());
obj.insert("has_more".to_string(), serde_json::Value::Bool(true));
}
let mut truncated_val = serde_json::json!({ if let Ok(trunc_str) = serde_yaml::to_string(&truncated_val)
"entities": sliced_ent, && (trunc_str.len() <= max_chars || keep_count == 1) {
"relations": sliced_rel, result_json = trunc_str;
"_meta": { break;
"truncated": true, }
"total_entities": total_ent, keep_count = keep_count.saturating_sub(keep_count / 3 + 1);
"returned_entities": sliced_ent.len(), }
"total_relations": total_rel,
"returned_relations": sliced_rel.len(),
"note": "Output structurally truncated to fit within max_tokens. Use pagination (limit, offset), search_nodes, or get_subgraph for targeted discovery."
}
});
if let Some(meta) = &meta_json
&& let Some(obj) = truncated_val.as_object_mut() {
obj.insert("total_count".to_string(), meta["total_count"].clone());
obj.insert("offset".to_string(), meta["offset"].clone());
obj.insert("limit".to_string(), meta["limit"].clone());
obj.insert("has_more".to_string(), serde_json::Value::Bool(true));
}
if let Ok(trunc_str) = serde_yaml::to_string(&truncated_val)
&& (trunc_str.len() <= max_chars || keep_count == 1) {
result_json = trunc_str;
break;
}
keep_count = keep_count.saturating_sub(keep_count / 3 + 1);
} }
} }
}
Ok(result_json) Ok(result_json)
})?; })
Ok(data) }).await.map_err(|e| crate::error::AppError::Internal(format!("Graph read task panicked: {}", e)))?
} }
} }
pub struct SearchNodesHandler; pub struct SearchNodesHandler;