refactor: apply 5-pass audit optimizations across mcp-memory codebase

server/src/search.rs

@@ -3,6 +3,15 @@ use std::sync::{Arc, Mutex};
 use tantivy::schema::*;
 use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct SearchResult {
+    pub id: String,
+    pub doc_type: String,
+    pub title: String,
+    pub body: String,
+    pub score: f32,
+}
+
 pub type SearchResultTuple = (String, String, String, String, f32);
 
 #[derive(Clone)]
@@ -89,23 +98,35 @@ impl MemoryIndex {
     }
 
     pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
+        self.index_entities_batch(std::slice::from_ref(e))
+    }
+
+    pub fn index_entities_batch(&self, entities: &[Entity]) -> tokio::task::JoinHandle<tantivy::Result<()>> {
         let writer = Arc::clone(&self.writer);
         let id_field = self.id_field;
-        let id_val = e.name.clone();
         let needs_commit = Arc::clone(&self.needs_commit);
 
-        let doc = doc!(
-            self.id_field => e.name.as_str(),
-            self.title_field => e.name.as_str(),
-            self.body_field => e.observations.join(" "),
-            self.type_field => "entity",
-            self.namespace_field => e.namespace.as_str()
-        );
+        let docs: Vec<(String, tantivy::TantivyDocument)> = entities
+            .iter()
+            .map(|e| {
+                let id_val = e.name.clone();
+                let doc = doc!(
+                    self.id_field => e.name.as_str(),
+                    self.title_field => e.name.as_str(),
+                    self.body_field => e.observations.join(" "),
+                    self.type_field => "entity",
+                    self.namespace_field => e.namespace.as_str()
+                );
+                (id_val, doc)
+            })
+            .collect();
 
         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));
-            writer.add_document(doc)?;
+            let writer_guard = writer.lock().unwrap_or_else(|e| e.into_inner());
+            for (id_val, doc) in docs {
+                writer_guard.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
+                writer_guard.add_document(doc)?;
+            }
             needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
             Ok(())
         })
@@ -172,6 +193,15 @@ impl MemoryIndex {
     }
 
 
+    pub fn commit_if_needed(&self) -> tantivy::Result<()> {
+        if self.needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
+            let mut writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+            writer.commit()?;
+            let _ = self.reader.reload();
+        }
+        Ok(())
+    }
+
     pub fn search(
         &self,
         query: &str,
@@ -182,13 +212,28 @@ impl MemoryIndex {
             &self.index,
             vec![self.title_field, self.body_field],
         );
-        let q = query_parser.parse_query(query)?;
+        let q = query_parser.parse_query(query).or_else(|_| {
+            let sanitized: String = query
+                .chars()
+                .map(|c| match c {
+                    '+' | '-' | '&' | '|' | '!' | '(' | ')' | '{' | '}' | '[' | ']' | '^'
+                    | '"' | '~' | '*' | '?' | ':' | '\\' | '/' => ' ',
+                    _ => c,
+                })
+                .collect();
+            let safe_query = sanitized.trim();
+            if safe_query.is_empty() {
+                query_parser.parse_query("*")
+            } else {
+                query_parser.parse_query(safe_query)
+            }
+        })?;
 
         let top_docs = searcher.search(
             &q,
             &tantivy::collector::TopDocs::with_limit(50).order_by_score(),
         )?;
-        let mut results = Vec::new();
+        let mut results = Vec::with_capacity(top_docs.len());
         for (score, doc_address) in top_docs {
             let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
             let id = retrieved_doc
@@ -273,72 +318,62 @@ impl MemoryIndex {
     }
 
     pub fn add_entity_sync(&self, e: &Entity) {
-        if let Ok(writer) = self.writer.lock() {
-            let _ = writer.add_document(doc!(
-                self.id_field => e.name.as_str(),
-                self.title_field => e.name.as_str(),
-                self.body_field => e.observations.join(" "),
-                self.type_field => "entity",
-                self.namespace_field => e.namespace.as_str()
-            ));
-            self.needs_commit
-                .store(true, std::sync::atomic::Ordering::SeqCst);
-        }
+        let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+        let _ = writer.add_document(doc!(
+            self.id_field => e.name.as_str(),
+            self.title_field => e.name.as_str(),
+            self.body_field => e.observations.join(" "),
+            self.type_field => "entity",
+            self.namespace_field => e.namespace.as_str()
+        ));
+        self.needs_commit
+            .store(true, std::sync::atomic::Ordering::SeqCst);
     }
 
     pub fn delete_all(&self) {
-        if let Ok(writer) = self.writer.lock() {
-            let _ = writer.delete_all_documents();
-            self.needs_commit
-                .store(true, std::sync::atomic::Ordering::SeqCst);
-        }
+        let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+        let _ = writer.delete_all_documents();
+        self.needs_commit
+            .store(true, std::sync::atomic::Ordering::SeqCst);
     }
 
     pub fn add_task_sync(&self, t: &Task) {
-        // println!("add_task_sync called for task: {}", t.id);
-        if let Ok(writer) = self.writer.lock() {
-            let _res = writer.add_document(doc!(
-                self.id_field => t.id.as_str(),
-                self.title_field => t.title.as_str(),
-                self.body_field => t.description.as_str(),
-                self.type_field => "task",
-                self.namespace_field => "global"
-            ));
-            // println!("Writer add_document returned id/result");
-            self.needs_commit
-                .store(true, std::sync::atomic::Ordering::SeqCst);
-            // println!("Needs_commit set to true in add_task_sync");
-        } else {
-            // println!("Failed to acquire writer lock in add_task_sync");
-        }
+        let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+        let _res = writer.add_document(doc!(
+            self.id_field => t.id.as_str(),
+            self.title_field => t.title.as_str(),
+            self.body_field => t.description.as_str(),
+            self.type_field => "task",
+            self.namespace_field => "global"
+        ));
+        self.needs_commit
+            .store(true, std::sync::atomic::Ordering::SeqCst);
     }
 
     pub fn add_snippet_sync(&self, s: &Snippet) {
-        if let Ok(writer) = self.writer.lock() {
-            let _ = writer.add_document(doc!(
-                self.id_field => s.name.as_str(),
-                self.title_field => s.name.as_str(),
-                self.body_field => format!("{} {}", s.language, s.description),
-                self.type_field => "snippet",
-                self.namespace_field => "global"
-            ));
-            self.needs_commit
-                .store(true, std::sync::atomic::Ordering::SeqCst);
-        }
+        let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+        let _ = writer.add_document(doc!(
+            self.id_field => s.name.as_str(),
+            self.title_field => s.name.as_str(),
+            self.body_field => format!("{} {}", s.language, s.description),
+            self.type_field => "snippet",
+            self.namespace_field => "global"
+        ));
+        self.needs_commit
+            .store(true, std::sync::atomic::Ordering::SeqCst);
     }
 
     pub fn add_adr_sync(&self, a: &Adr) {
-        if let Ok(writer) = self.writer.lock() {
-            let _ = writer.add_document(doc!(
-                self.id_field => a.id.as_str(),
-                self.title_field => a.title.as_str(),
-                self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
-                self.type_field => "adr",
-                self.namespace_field => "global"
-            ));
-            self.needs_commit
-                .store(true, std::sync::atomic::Ordering::SeqCst);
-        }
+        let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
+        let _ = writer.add_document(doc!(
+            self.id_field => a.id.as_str(),
+            self.title_field => a.title.as_str(),
+            self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
+            self.type_field => "adr",
+            self.namespace_field => "global"
+        ));
+        self.needs_commit
+            .store(true, std::sync::atomic::Ordering::SeqCst);
     }
 }
 
@@ -423,12 +458,12 @@ mod tests {
         let temp_dir = TempDir::new().unwrap();
         let index = MemoryIndex::new(temp_dir.path()).unwrap();
 
-        // Malformed lucene query (unclosed parenthesis)
+        // Malformed lucene query (unclosed parenthesis) - sanitized gracefully
         let result = index.search("title: (unclosed", None);
-        assert!(result.is_err());
+        assert!(result.is_ok());
 
-        // Another malformed query (unclosed quote)
+        // Another malformed query (unclosed quote) - sanitized gracefully
         let result2 = index.search("title: \"unclosed", None);
-        assert!(result2.is_err());
+        assert!(result2.is_ok());
     }
 }