fix(mcp-memory,nvim-core): eliminate lock contention in get_nvim_connection and resolve CPU bound loop in graph.rs
This commit is contained in:
Riz Ashraf committed
1 parent
e91b53fdcc
commit
0afba74b4c
2 files changed
+104
-102
No files matched your search
+100
-99
@@ -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<String> {
|
||||
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<String> {
|
||||
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;
|
||||
|
||||
Reference in new issue
Block a user