From 0afba74b4c9dd506e8d11b72a1361e8f73d108ae Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Sun, 11 Oct 2026 18:21:58 +0100 Subject: [PATCH] fix(mcp-memory,nvim-core): eliminate lock contention in get_nvim_connection and resolve CPU bound loop in graph.rs --- nvim-core/src/lib.rs | 7 +- server/src/handlers/graph.rs | 199 ++++++++++++++++++----------------- 2 files changed, 104 insertions(+), 102 deletions(-) diff --git a/nvim-core/src/lib.rs b/nvim-core/src/lib.rs index ccefd92..ed7c993 100644 --- a/nvim-core/src/lib.rs +++ b/nvim-core/src/lib.rs @@ -391,6 +391,8 @@ fn handle_nvim_notification(params: &[rmpv::Value]) { } async fn get_nvim_connection() -> Result, 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 // without spawning any external nvim.exe child process! { @@ -400,9 +402,8 @@ async fn get_nvim_connection() -> Result, String> { { let active_lock = ACTIVE_SOCKET.lock().await; if let Some(ref current_socket) = *active_lock - && let Ok(target) = read_active_socket_file().await - && !target.is_empty() - && &target == current_socket + && !target_socket.is_empty() + && &target_socket == current_socket { return Ok(sender.clone()); } diff --git a/server/src/handlers/graph.rs b/server/src/handlers/graph.rs index 7c06bb5..2c800a6 100644 --- a/server/src/handlers/graph.rs +++ b/server/src/handlers/graph.rs @@ -632,119 +632,120 @@ impl McpTool for ReadGraphHandler { let limit = req.limit; let offset = req.offset.unwrap_or(0); - let data = state.read_graph(|full| -> crate::error::Result { - let filtered_entities: Vec<(&String, &crate::models::Entity)> = full - .entities - .iter() - .filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns)) - .filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et)) - .collect(); - - 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) + tokio::task::spawn_blocking(move || { + state.read_graph(|full| -> crate::error::Result { + let filtered_entities: Vec<(&String, &crate::models::Entity)> = full + .entities + .iter() + .filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns)) + .filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et)) .collect(); - let has_more = offset + paged_entities.len() < total_count; + let total_count = filtered_entities.len(); - 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)) + 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(); + + 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!({ - "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 mut result_json = serde_yaml::to_string(&serialized)?; - 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"], - }) - } else { - serde_json::json!({ - "entities": entities_map, - "relations": relations_vec, - }) - }; + if let Some(max_tok) = max_tokens { + let max_chars = max_tok * 4; + if result_json.len() > max_chars { + 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 mut result_json = serde_yaml::to_string(&serialized)?; + let ent_vec: Vec<_> = entities_map.into_iter().collect(); - if let Some(max_tok) = max_tokens { - let max_chars = max_tok * 4; - if result_json.len() > max_chars { - 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; + while keep_count > 0 { + let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect(); + let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect(); - 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 { - let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect(); - let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect(); + 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)); + } - 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." - } - }); - - 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); + 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(data) + Ok(result_json) + }) + }).await.map_err(|e| crate::error::AppError::Internal(format!("Graph read task panicked: {}", e)))? } } pub struct SearchNodesHandler;