From eed9a9165c1d379aa8ef04034d45d5c03c80d784 Mon Sep 17 00:00:00 2001
From: Riz Ashraf <reazul.ashraf@inseinc.com>
Date: Sat, 10 Oct 2026 12:27:29 +0100
Subject: [PATCH] feat(graph): add ergonomic UpdateEntitiesTool for renaming
 and type changes

---
 server/src/handlers/graph.rs | 102 +++++++++++++++++++++++++++++++++++
 server/src/router.rs         |   1 +
 server/src/tools.rs          |  17 ++++++
 3 files changed, 120 insertions(+)

diff --git a/server/src/handlers/graph.rs b/server/src/handlers/graph.rs
index b5361ef..6e9ff0c 100644
--- a/server/src/handlers/graph.rs
+++ b/server/src/handlers/graph.rs
@@ -193,6 +193,108 @@ impl McpTool for CreateEntitiesHandler {
     }
 }
 
+pub struct UpdateEntitiesHandler;
+
+#[async_trait]
+impl McpTool for UpdateEntitiesHandler {
+    fn name(&self) -> &'static str {
+        "update_entities"
+    }
+
+    fn schema(&self) -> Value {
+        crate::mcp::tool_def::<UpdateEntitiesTool>(
+            "update_entities",
+            "Update existing entities in the knowledge graph (e.g., to rename an entity or change its type).",
+        )
+    }
+
+    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
+        let req: UpdateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
+        
+        let mut updated_entities = Vec::new();
+        let mut deleted_names = Vec::new();
+        let mut not_found = Vec::new();
+        let mut conflict_names = Vec::new();
+
+        state.modify_graph(|g| {
+            for update in req.updates {
+                if !g.entities.contains_key(&update.name) {
+                    not_found.push(update.name.clone());
+                    continue;
+                }
+
+                if let Some(new_name) = &update.new_name {
+                    if update.name != *new_name && g.entities.contains_key(new_name) {
+                        conflict_names.push(new_name.clone());
+                        continue;
+                    }
+                }
+
+                if let Some(mut entity) = g.entities.remove(&update.name) {
+                    let mut renamed = false;
+                    let old_name = entity.name.clone();
+
+                    if let Some(new_name) = update.new_name {
+                        if old_name != new_name {
+                            entity.name = new_name.clone();
+                            renamed = true;
+                            deleted_names.push(old_name.clone());
+                        }
+                    }
+
+                    if let Some(new_type) = update.new_entity_type {
+                        entity.entity_type = crate::models::normalize_entity_type(&new_type);
+                    }
+
+                    let current_name = entity.name.clone();
+                    
+                    if renamed {
+                        // update relations
+                        let mut seen = std::collections::HashSet::new();
+                        g.relations.retain_mut(|r| {
+                            if r.from == old_name {
+                                r.from = current_name.clone();
+                            }
+                            if r.to == old_name {
+                                r.to = current_name.clone();
+                            }
+                            // Keep unique relations
+                            seen.insert(r.clone())
+                        });
+                    }
+
+                    updated_entities.push(entity.clone());
+                    g.entities.insert(current_name, entity);
+                }
+            }
+        });
+
+        if !not_found.is_empty() {
+            return Err(crate::error::AppError::Internal(format!(
+                "Error: Entities not found: {}", not_found.join(", ")
+            )));
+        }
+
+        if !conflict_names.is_empty() {
+            return Err(crate::error::AppError::Internal(format!(
+                "Error: Cannot rename to existing entity names: {}", conflict_names.join(", ")
+            )));
+        }
+
+        let idx = state.get_search_index().await;
+        for old_name in deleted_names {
+            drop(idx.delete_document(&old_name));
+        }
+
+        if !updated_entities.is_empty() {
+            let _ = idx.index_entities_batch(&updated_entities).await;
+        }
+
+        let names: Vec<String> = updated_entities.iter().map(|e| e.name.clone()).collect();
+        Ok(format!("Successfully updated {} entity/entities: {}", names.len(), names.join(", ")))
+    }
+}
+
 pub struct CreateRelationsHandler;
 
 #[async_trait]
diff --git a/server/src/router.rs b/server/src/router.rs
index b4ac37b..1e9f503 100644
--- a/server/src/router.rs
+++ b/server/src/router.rs
@@ -479,6 +479,7 @@ impl MemoryHandler {
 
         register!(graph::QueryGraphPathHandler);
         register!(graph::CreateEntitiesHandler);
+        register!(graph::UpdateEntitiesHandler);
         register!(graph::CreateRelationsHandler);
         register!(graph::AddObservationsHandler);
         register!(graph::DeleteEntitiesHandler);
diff --git a/server/src/tools.rs b/server/src/tools.rs
index 4b3ed97..04c85a7 100644
--- a/server/src/tools.rs
+++ b/server/src/tools.rs
@@ -18,6 +18,23 @@ pub struct CreateEntitiesTool {
     pub entities: Vec<EntityInput>,
 }
 
+#[derive(Debug, Deserialize, Serialize, JsonSchema)]
+pub struct UpdateEntityInput {
+    /// The current name of the entity to update.
+    pub name: String,
+    /// Optional new name to rename the entity to.
+    pub new_name: Option<String>,
+    /// Optional new entity type (e.g. PascalCase).
+    pub new_entity_type: Option<String>,
+}
+
+/// Update existing entities in the knowledge graph (e.g., to rename an entity or change its type).
+#[derive(Debug, Deserialize, Serialize, JsonSchema)]
+pub struct UpdateEntitiesTool {
+    /// Array of entity updates to perform.
+    pub updates: Vec<UpdateEntityInput>,
+}
+
 #[derive(Debug, Deserialize, Serialize, JsonSchema)]
 pub struct RelationInput {
     pub from: String,