perf: fix memory leaks, optimize tracing logs, fix networking eof fragmentation bug in nvim-core and stub, and deduplicate json-rpc handlers

server/src/handlers/graph.rs

@@ -0,0 +1,579 @@
+use crate::router::McpTool;
+use crate::state::MemoryState;
+use crate::tools::*;
+use async_trait::async_trait;
+use serde_json::Value;
+use std::collections::HashSet;
+use std::sync::Arc;
+
+#[derive(serde::Serialize, Default)]
+struct BorrowedGraph<'a> {
+    entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
+    relations: Vec<&'a crate::models::Relation>,
+}
+
+pub struct QueryGraphPathHandler;
+
+#[async_trait]
+impl McpTool for QueryGraphPathHandler {
+    fn name(&self) -> &'static str {
+        "query_graph_path"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: crate::tools::QueryGraphPathTool =
+            serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.read_graph(|graph| {
+            let max_depth = req.max_depth.unwrap_or(5);
+            let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
+            let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
+            let mut parents: std::collections::HashMap<&str, (&str, std::borrow::Cow<'_, str>)> =
+                std::collections::HashMap::new();
+
+            queue.push_back(req.start_node.as_str());
+            visited.insert(req.start_node.as_str());
+
+            let mut found = false;
+            let mut current_depth = 0;
+            let mut nodes_at_current_depth = 1;
+            let mut nodes_at_next_depth = 0;
+
+            while let Some(current) = queue.pop_front() {
+                if current == req.end_node {
+                    found = true;
+                    break;
+                }
+                nodes_at_current_depth -= 1;
+                if current_depth < max_depth {
+                    for rel in &graph.relations {
+                        if rel.from == current && !visited.contains(rel.to.as_str()) {
+                            visited.insert(rel.to.as_str());
+                            parents.insert(
+                                rel.to.as_str(),
+                                (
+                                    current,
+                                    std::borrow::Cow::Borrowed(rel.relation_type.as_str()),
+                                ),
+                            );
+                            queue.push_back(rel.to.as_str());
+                            nodes_at_next_depth += 1;
+                        } else if rel.to == current && !visited.contains(rel.from.as_str()) {
+                            visited.insert(rel.from.as_str());
+                            parents.insert(
+                                rel.from.as_str(),
+                                (
+                                    current,
+                                    std::borrow::Cow::Owned(format!(
+                                        "inverse({})",
+                                        rel.relation_type
+                                    )),
+                                ),
+                            );
+                            queue.push_back(rel.from.as_str());
+                            nodes_at_next_depth += 1;
+                        }
+                    }
+                }
+                if nodes_at_current_depth == 0 {
+                    current_depth += 1;
+                    nodes_at_current_depth = nodes_at_next_depth;
+                    nodes_at_next_depth = 0;
+                }
+            }
+
+            if found {
+                let mut path = Vec::new();
+                let mut curr = req.end_node.as_str();
+                while curr != req.start_node {
+                    if let Some((parent, rel_type)) = parents.get(&curr) {
+                        path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
+                        curr = parent;
+                    } else {
+                        break;
+                    }
+                }
+                path.reverse();
+                Ok(format!("Path found:\n{}", path.join("\n")))
+            } else {
+                Ok(format!(
+                    "No path found between {} and {} within depth {}",
+                    req.start_node, req.end_node, max_depth
+                ))
+            }
+        })
+    }
+}
+
+pub struct CreateEntitiesHandler;
+
+#[async_trait]
+impl McpTool for CreateEntitiesHandler {
+    fn name(&self) -> &'static str {
+        "create_entities"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let mut inserted = Vec::new();
+        state.modify_graph(|g| {
+            for entity in req.entities {
+                if !entity.name.is_empty() {
+                    inserted.push(entity.clone());
+                    g.entities.insert(entity.name.clone(), entity);
+                }
+            }
+        });
+        let idx = state
+            .search_index
+            .read()
+            .unwrap_or_else(|e| e.into_inner())
+            .clone();
+        for entity in inserted {
+            drop(idx.index_entity(&entity));
+        }
+        Ok("Entities created".to_string())
+    }
+}
+
+pub struct CreateRelationsHandler;
+
+#[async_trait]
+impl McpTool for CreateRelationsHandler {
+    fn name(&self) -> &'static str {
+        "create_relations"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: CreateRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|g| {
+            for relation in req.relations {
+                if !relation.from.is_empty() && !relation.to.is_empty() {
+                    g.relations.push(relation);
+                }
+            }
+        });
+        Ok("Relations created".to_string())
+    }
+}
+
+pub struct AddObservationsHandler;
+
+#[async_trait]
+impl McpTool for AddObservationsHandler {
+    fn name(&self) -> &'static str {
+        "add_observations"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|g| {
+            for o in req.observations {
+                if let Some(e) = g.entities.get_mut(&o.entity_name) {
+                    e.observations.extend(o.contents);
+                }
+            }
+        });
+        Ok("Observations added".to_string())
+    }
+}
+
+pub struct DeleteEntitiesHandler;
+
+#[async_trait]
+impl McpTool for DeleteEntitiesHandler {
+    fn name(&self) -> &'static str {
+        "delete_entities"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<DeleteEntitiesTool>("delete_entities", "Execute delete_entities")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
+        state.modify_graph(|master| {
+            for name in &to_delete {
+                master.entities.remove(name);
+            }
+            master
+                .relations
+                .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
+        });
+
+        let idx = state
+            .search_index
+            .read()
+            .unwrap_or_else(|e| e.into_inner())
+            .clone();
+        for name in to_delete {
+            drop(idx.delete_document(&name));
+        }
+        Ok("Entities deleted".to_string())
+    }
+}
+
+pub struct DeleteObservationsHandler;
+
+#[async_trait]
+impl McpTool for DeleteObservationsHandler {
+    fn name(&self) -> &'static str {
+        "delete_observations"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<DeleteObservationsTool>(
+            "delete_observations",
+            "Execute delete_observations",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: DeleteObservationsTool =
+            serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|master| {
+            for d in req.deletions {
+                if let Some(e) = master.entities.get_mut(&d.entity_name) {
+                    let to_rem: HashSet<_> = d.observations.into_iter().collect();
+                    e.observations.retain(|o| !to_rem.contains(o));
+                }
+            }
+        });
+        Ok("Observations deleted".to_string())
+    }
+}
+
+pub struct DeleteRelationsHandler;
+
+#[async_trait]
+impl McpTool for DeleteRelationsHandler {
+    fn name(&self) -> &'static str {
+        "delete_relations"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<DeleteRelationsTool>("delete_relations", "Execute delete_relations")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|master| {
+            let to_rem: HashSet<_> = req.relations.into_iter().collect();
+            master.relations.retain(|r| !to_rem.contains(r));
+        });
+        Ok("Relations deleted".to_string())
+    }
+}
+
+pub struct ReadGraphHandler;
+
+#[async_trait]
+impl McpTool for ReadGraphHandler {
+    fn name(&self) -> &'static str {
+        "read_graph"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let data = state.read_graph(|full| {
+            if let Some(ns) = req.namespace {
+                let mut filtered = BorrowedGraph::default();
+                for (k, v) in &full.entities {
+                    if v.namespace == ns {
+                        filtered.entities.insert(k, v);
+                    }
+                }
+                for r in &full.relations {
+                    if r.namespace == ns {
+                        filtered.relations.push(r);
+                    }
+                }
+                serde_json::to_string(&filtered).unwrap_or_default()
+            } else {
+                serde_json::to_string(full).unwrap_or_default()
+            }
+        });
+        Ok(data)
+    }
+}
+
+pub struct SearchNodesHandler;
+
+#[async_trait]
+impl McpTool for SearchNodesHandler {
+    fn name(&self) -> &'static str {
+        "search_nodes"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<SearchNodesTool>("search_nodes", "Execute search_nodes")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let matches = if let Ok(idx) = state.search_index.read() {
+            idx.search(&req.query, req.namespace.as_deref())
+                .unwrap_or_default()
+        } else {
+            vec![]
+        };
+
+        let data = state.read_graph(|full| {
+            let mut result = BorrowedGraph::default();
+            for (id, doc_type, _, _, _) in &matches {
+                if doc_type == "entity"
+                    && let Some(e) = full.entities.get(id)
+                {
+                    result.entities.insert(id, e);
+                }
+            }
+            serde_json::to_string(&result).unwrap_or_default()
+        });
+        Ok(data)
+    }
+}
+
+pub struct OpenNodesHandler;
+
+#[async_trait]
+impl McpTool for OpenNodesHandler {
+    fn name(&self) -> &'static str {
+        "open_nodes"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<OpenNodesTool>("open_nodes", "Execute open_nodes")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let data = state.read_graph(|full| {
+            let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
+            let mut result = BorrowedGraph::default();
+            let mut connected: HashSet<&str> = HashSet::new();
+            for r in &full.relations {
+                if targets.contains(r.from.as_str()) {
+                    connected.insert(r.to.as_str());
+                    result.relations.push(r);
+                } else if targets.contains(r.to.as_str()) {
+                    connected.insert(r.from.as_str());
+                    result.relations.push(r);
+                }
+            }
+            for (name, e) in &full.entities {
+                if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
+                    result.entities.insert(name, e);
+                }
+            }
+            serde_json::to_string(&result).unwrap_or_default()
+        });
+        Ok(data)
+    }
+}
+
+pub struct VisualizeGraphHandler;
+
+#[async_trait]
+impl McpTool for VisualizeGraphHandler {
+    fn name(&self) -> &'static str {
+        "visualize_graph"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<VisualizeGraphTool>("visualize_graph", "Execute visualize_graph")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        let query = req.query.unwrap_or_default();
+        use std::fmt::Write;
+        let mut output = state.read_graph(|full| {
+            let mut included: HashSet<&str> = HashSet::new();
+            let mut to_draw = Vec::new();
+
+            for (name, e) in &full.entities {
+                if let Some(ns) = &req.namespace
+                    && e.namespace != *ns
+                {
+                    continue;
+                }
+                if query.is_empty()
+                    || contains_ignore_ascii_case(name, &query)
+                    || contains_ignore_ascii_case(&e.entity_type, &query)
+                {
+                    included.insert(name.as_str());
+                }
+            }
+
+            for r in &full.relations {
+                if let Some(ns) = &req.namespace
+                    && r.namespace != *ns
+                {
+                    continue;
+                }
+                if query.is_empty()
+                    || included.contains(r.from.as_str())
+                    || included.contains(r.to.as_str())
+                {
+                    included.insert(r.from.as_str());
+                    included.insert(r.to.as_str());
+                    to_draw.push(r);
+                }
+            }
+
+            let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
+            out.push_str("graph TD;\n");
+
+            let sanitize = |s: &str, id_mode: bool| -> String {
+                let mut o = String::with_capacity(s.len());
+                for c in s.chars() {
+                    if c != '"' && c != '(' && c != ')' {
+                        if id_mode && (c == ' ' || c == '-' || c == '.') {
+                            o.push('_');
+                        } else {
+                            o.push(c);
+                        }
+                    }
+                }
+                o
+            };
+
+            for name in &included {
+                let _ = writeln!(
+                    out,
+                    "    id_{}[\"{}\"];",
+                    sanitize(name, true),
+                    sanitize(name, false)
+                );
+            }
+            for r in to_draw {
+                let _ = writeln!(
+                    out,
+                    "    id_{}-->|\"{}\"|id_{};",
+                    sanitize(&r.from, true),
+                    r.relation_type.replace("\"", ""),
+                    sanitize(&r.to, true)
+                );
+            }
+            out
+        });
+        if output == "graph TD;\n" {
+            output = "No nodes found to visualize.".to_string();
+        }
+        Ok(output.to_string())
+    }
+}
+
+pub struct CondenseEntityHandler;
+
+#[async_trait]
+impl McpTool for CondenseEntityHandler {
+    fn name(&self) -> &'static str {
+        "condense_entity"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<CondenseEntityTool>("condense_entity", "Execute condense_entity")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|master| {
+            if let Some(e) = master.entities.get_mut(&req.entity_name) {
+                e.observations = req.summarized_observations;
+            }
+        });
+        Ok("Entity condensed".to_string())
+    }
+}
+
+pub struct MergeEntitiesHandler;
+
+#[async_trait]
+impl McpTool for MergeEntitiesHandler {
+    fn name(&self) -> &'static str {
+        "merge_entities"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<MergeEntitiesTool>("merge_entities", "Execute merge_entities")
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        state.modify_graph(|master| {
+            if let Some(src) = master.entities.remove(&req.source_entity) {
+                if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
+                    tgt.observations.extend(src.observations);
+                    MemoryState::deduplicate(&mut tgt.observations);
+                } else {
+                    let mut new_tgt = src.clone();
+                    new_tgt.name = req.target_entity.clone();
+                    master.entities.insert(req.target_entity.clone(), new_tgt);
+                }
+            }
+            for r in &mut master.relations {
+                if r.from == req.source_entity {
+                    r.from = req.target_entity.clone();
+                }
+                if r.to == req.source_entity {
+                    r.to = req.target_entity.clone();
+                }
+            }
+            MemoryState::deduplicate(&mut master.relations);
+        });
+        Ok("Entities merged".to_string())
+    }
+}
+
+pub struct FindOrphansHandler;
+
+#[async_trait]
+impl McpTool for FindOrphansHandler {
+    fn name(&self) -> &'static str {
+        "find_orphans"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<FindOrphansTool>("find_orphans", "Execute find_orphans")
+    }
+
+    async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> Result<String, String> {
+        let orphans = state.read_graph(|full| {
+            let mut connected = std::collections::HashSet::new();
+            for r in &full.relations {
+                connected.insert(r.from.as_str());
+                connected.insert(r.to.as_str());
+            }
+            full.entities
+                .keys()
+                .filter(|k| !connected.contains(k.as_str()))
+                .cloned()
+                .collect::<Vec<String>>()
+        });
+        let data = serde_json::to_string(&orphans).unwrap_or_default();
+        Ok(data.to_string())
+    }
+}
+
+use crate::handlers::utils::*;