fix(search): prevent memory leaks and duplication in tantivy search index by explicitly deleting prior terms during updates and syncing store deletions

This commit is contained in:
Riz Ashraf committed 2026-09-21 19:16:55 +01:00
1 parent 7dc0c329ad
commit af2c4ab046
4 files changed
+42 -5

No files matched your search

+6 -1
View File
@@ -196,7 +196,7 @@ impl McpTool for DeleteEntitiesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> { 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 req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let to_delete: HashSet<_> = req.entity_names.into_iter().collect(); let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
state.modify_graph(|master| { state.modify_graph(|master| {
for name in &to_delete { for name in &to_delete {
master.entities.remove(name); master.entities.remove(name);
@@ -205,6 +205,11 @@ impl McpTool for DeleteEntitiesHandler {
.relations .relations
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to)); .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 {
let _ = idx.delete_document(&name).await;
}
Ok("Entities deleted".to_string()) Ok("Entities deleted".to_string())
} }
} }
+8 -2
View File
@@ -66,6 +66,7 @@ impl McpTool for DeleteTaskHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
let req: DeleteTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: DeleteTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut deleted_count = 0; let mut deleted_count = 0;
let mut actually_deleted = Vec::new();
state.tasks.modify(|tasks| { state.tasks.modify(|tasks| {
let initial_len = tasks.len(); let initial_len = tasks.len();
// Collect IDs of tasks to delete (this task + all its recursive children) // Collect IDs of tasks to delete (this task + all its recursive children)
@@ -93,12 +94,17 @@ impl McpTool for DeleteTaskHandler {
queue.extend(children.iter().cloned()); queue.extend(children.iter().cloned());
} }
} }
tasks.retain(|t| !to_delete.contains(&t.id)); actually_deleted = to_delete.into_iter().collect();
tasks.retain(|t| !actually_deleted.contains(&t.id));
deleted_count = initial_len - tasks.len(); deleted_count = initial_len - tasks.len();
}); });
if deleted_count > 0 { if deleted_count > 0 {
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
for id in actually_deleted {
let _ = idx.delete_document(&id).await;
}
Ok(vec![ Ok(vec![
format!("Deleted task and its children ({} total).", deleted_count).to_string(), format!("Deleted task and its children ({} total).", deleted_count).to_string(),
][0] ][0]
+2
View File
@@ -179,6 +179,8 @@ impl McpTool for DeleteSnippetHandler {
deleted = snippets.len() < orig; deleted = snippets.len() < orig;
}); });
if deleted { if deleted {
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
let _ = idx.delete_document(&req.name).await;
Ok("Snippet deleted.".to_string()) Ok("Snippet deleted.".to_string())
} else { } else {
Ok("Snippet not found.".to_string()) Ok("Snippet not found.".to_string())
+26 -2
View File
@@ -57,6 +57,8 @@ impl MemoryIndex {
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> { pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
let id_val = e.name.clone();
let doc = doc!( let doc = doc!(
self.id_field => e.name.as_str(), self.id_field => e.name.as_str(),
@@ -68,6 +70,7 @@ impl MemoryIndex {
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner()); let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
writer.add_document(doc)?; writer.add_document(doc)?;
Ok(()) Ok(())
}) })
@@ -75,22 +78,37 @@ impl MemoryIndex {
pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> { pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
let id_val = t.id.clone();
let doc = doc!( let doc = doc!(
self.id_field => t.id.as_str(), self.id_field => t.id.as_str(),
self.title_field => t.title.as_str(), self.title_field => t.title.as_str(),
self.body_field => t.description.as_str(), self.body_field => format!("{}\n{}", t.description, t.acceptance_criteria.iter().map(|c| c.description.as_str()).collect::<Vec<_>>().join("\n")),
self.type_field => "task", self.type_field => "task",
self.namespace_field => "global" self.namespace_field => "global"
); );
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner()); let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
writer.add_document(doc)?; writer.add_document(doc)?;
Ok(()) Ok(())
}) })
} }
pub fn delete_document(&self, id: &str) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
let id_val = id.to_string();
tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
Ok(())
})
}
pub async fn commit(&self) -> tantivy::Result<()> { pub async fn commit(&self) -> tantivy::Result<()> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
@@ -162,17 +180,20 @@ impl MemoryIndex {
pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> { pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
let id_val = s.name.clone();
let doc = doc!( let doc = doc!(
self.id_field => s.name.as_str(), self.id_field => s.name.as_str(),
self.title_field => s.name.as_str(), self.title_field => s.name.as_str(),
self.body_field => format!("{} {}", s.language, s.description), self.body_field => format!("{} {}\n{}", s.language, s.description, s.code),
self.type_field => "snippet", self.type_field => "snippet",
self.namespace_field => "global" self.namespace_field => "global"
); );
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner()); let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
writer.add_document(doc)?; writer.add_document(doc)?;
Ok(()) Ok(())
}) })
@@ -180,6 +201,8 @@ impl MemoryIndex {
pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> { pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
let id_val = a.id.clone();
let doc = doc!( let doc = doc!(
self.id_field => a.id.as_str(), self.id_field => a.id.as_str(),
@@ -191,6 +214,7 @@ impl MemoryIndex {
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner()); let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
writer.add_document(doc)?; writer.add_document(doc)?;
Ok(()) Ok(())
}) })