fix: make text index reinserts atomic (#234)

Author: presempathy-awb

src/git/versioned_store/namespaced.rs

@@ -2680,25 +2680,38 @@ where
     M: MetadataBackend,
 {
     /// Insert or update `(id, text)`. The text is split through the
-    /// configured chunker; each chunk is embedded and stored under its own
-    /// chunk-id. Previous chunks for this `id` are removed first so that a
-    /// re-insert with fewer chunks doesn't leak stale ones.
+    /// configured chunker; replacement chunks are embedded first, then previous
+    /// chunks for this `id` are removed so failed embeds leave the old document
+    /// searchable. A successful empty chunk set removes the document.
     pub fn insert(&mut self, id: &[u8], text: &str) -> Result<(), TextIndexError> {
-        // Drop any previous chunks for this doc_id.
-        self.delete_chunks_for_doc(id);
-
         let chunks = self.chunker.split(text);
         if chunks.is_empty() {
+            self.delete_chunks_for_doc(id);
             return Ok(());
         }
+        if self.embedder.dim() == 0 {
+            return Err(TextIndexError::Proximity(ProximityError::ZeroDim));
+        }
+        let mut replacements = Vec::with_capacity(chunks.len());
+        for (chunk_idx, chunk_text) in chunks.iter().enumerate() {
+            let vec = self.embedder.embed(chunk_text)?;
+            if vec.len() as u16 != self.embedder.dim() {
+                return Err(TextIndexError::DimensionMismatch {
+                    stored: self.embedder.dim(),
+                    got: vec.len() as u16,
+                });
+            }
+            let chunk_id = make_chunk_id(id, chunk_idx as u32);
+            replacements.push((chunk_id, vec));
+        }
+
+        self.delete_chunks_for_doc(id);
         let idx = self
             .store
             .proximity_indexes
             .get_mut(&self.inner_idx_key)
             .expect("inner proximity index must be loaded");
-        for (chunk_idx, chunk_text) in chunks.iter().enumerate() {
-            let vec = self.embedder.embed(chunk_text)?;
-            let chunk_id = make_chunk_id(id, chunk_idx as u32);
+        for (chunk_id, vec) in replacements {
             idx.insert(chunk_id, vec)?;
         }
         self.store

src/proximity/text_index.rs

@@ -446,22 +446,23 @@ impl<const N: usize, E: Embedder, S: NodeStorage<N>> TextIndex<N, E, S> {
     /// `Chunker` into one or more chunks; each chunk is embedded and stored
     /// under its own chunk-id.
     ///
-    /// **Re-insert semantics:** existing chunks for this `id` are removed
-    /// first, so that switching to a chunker that produces fewer chunks
-    /// doesn't leak stale ones.
+    /// **Re-insert semantics:** replacement chunks are embedded first, then
+    /// existing chunks for this `id` are removed, so a failed embed leaves the
+    /// previous document searchable. A successful empty chunk set removes the
+    /// document from the index.
     pub fn insert(&mut self, id: &[u8], text: &str) -> Result<(), TextIndexError> {
-        // Drop any previous chunks for this doc_id so the new chunk count
-        // overwrites cleanly (handles "re-chunk after switching chunkers"
-        // as well as plain upserts).
-        self.delete_chunks_for_doc(id);
-
         let chunks = self.chunker.split(text);
         if chunks.is_empty() {
             // The chunker explicitly opted out (e.g. an empty document under
             // LineChunker). Treat as "don't index" rather than as an error —
             // matches the cascade transformer's `None` semantics.
+            self.delete_chunks_for_doc(id);
             return Ok(());
         }
+        if self.embedder.dim() == 0 {
+            return Err(TextIndexError::Proximity(ProximityError::ZeroDim));
+        }
+        let mut replacements = Vec::with_capacity(chunks.len());
         for (idx, chunk_text) in chunks.iter().enumerate() {
             let vec = self.embedder.embed(chunk_text)?;
             if vec.len() != usize::from(self.embedder.dim()) {
@@ -471,6 +472,11 @@ impl<const N: usize, E: Embedder, S: NodeStorage<N>> TextIndex<N, E, S> {
                 }));
             }
             let chunk_id = make_chunk_id(id, idx as u32);
+            replacements.push((chunk_id, vec));
+        }
+
+        self.delete_chunks_for_doc(id);
+        for (chunk_id, vec) in replacements {
             self.inner.insert(chunk_id, vec)?;
         }
         Ok(())
@@ -604,6 +610,42 @@ mod tests {
         TextIndexConfig::new(HashEmbedder::new(dim, 0))
     }
 
+    #[derive(Debug, Clone)]
+    struct FailsOnNeedleEmbedder {
+        inner: HashEmbedder,
+        needle: &'static str,
+    }
+
+    impl FailsOnNeedleEmbedder {
+        fn new(dim: u16, needle: &'static str) -> Self {
+            Self {
+                inner: HashEmbedder::new(dim, 0),
+                needle,
+            }
+        }
+    }
+
+    impl Embedder for FailsOnNeedleEmbedder {
+        fn id(&self) -> &str {
+            self.inner.id()
+        }
+
+        fn version(&self) -> &str {
+            self.inner.version()
+        }
+
+        fn dim(&self) -> u16 {
+            self.inner.dim()
+        }
+
+        fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
+            if text.contains(self.needle) {
+                return Err(EmbedError::Failure("forced insert failure".to_string()));
+            }
+            self.inner.embed(text)
+        }
+    }
+
     #[test]
     fn insert_and_search_finds_exact_match() {
         let storage = InMemoryNodeStorage::<32>::new();
@@ -623,6 +665,23 @@ mod tests {
         );
     }
 
+    #[test]
+    fn failed_reinsert_preserves_existing_document() {
+        let storage = InMemoryNodeStorage::<32>::new();
+        let mut idx = TextIndex::new(
+            storage,
+            TextIndexConfig::new(FailsOnNeedleEmbedder::new(8, "fail-insert")),
+        );
+        idx.insert(b"doc:1", "stable text").unwrap();
+
+        let err = idx.insert(b"doc:1", "please fail-insert").unwrap_err();
+        assert!(err.to_string().contains("forced insert failure"));
+
+        let hits = idx.search("stable text", 1).unwrap();
+        assert_eq!(hits[0].id, b"doc:1".to_vec());
+        assert!(hits[0].score < 1e-4);
+    }
+
     #[test]
     fn delete_removes_document_from_search() {
         let storage = InMemoryNodeStorage::<32>::new();

tests/text_index_namespaced.rs

@@ -25,14 +25,54 @@ mod common;
 
 use common::setup_repo_and_dataset;
 use prollytree::git::versioned_store::FileNamespacedKvStore;
-use prollytree::proximity::{HashEmbedder, TextIndexConfig, TextIndexError};
+use prollytree::proximity::{EmbedError, Embedder, HashEmbedder, TextIndexConfig, TextIndexError};
 
 const N: usize = 32;
 
 fn cfg(dim: u16, seed: u64) -> TextIndexConfig<HashEmbedder> {
     TextIndexConfig::new(HashEmbedder::new(dim, seed))
 }
 
+#[derive(Debug, Clone)]
+struct FailsOnNeedleEmbedder {
+    inner: HashEmbedder,
+    needle: &'static str,
+}
+
+impl FailsOnNeedleEmbedder {
+    fn new(dim: u16, seed: u64, needle: &'static str) -> Self {
+        Self {
+            inner: HashEmbedder::new(dim, seed),
+            needle,
+        }
+    }
+}
+
+impl Embedder for FailsOnNeedleEmbedder {
+    fn id(&self) -> &str {
+        self.inner.id()
+    }
+
+    fn version(&self) -> &str {
+        self.inner.version()
+    }
+
+    fn dim(&self) -> u16 {
+        self.inner.dim()
+    }
+
+    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
+        if text.contains(self.needle) {
+            return Err(EmbedError::Failure("forced insert failure".to_string()));
+        }
+        self.inner.embed(text)
+    }
+}
+
+fn failing_cfg(dim: u16, seed: u64) -> TextIndexConfig<FailsOnNeedleEmbedder> {
+    TextIndexConfig::new(FailsOnNeedleEmbedder::new(dim, seed, "fail-insert"))
+}
+
 #[test]
 fn text_index_insert_search_basic() {
     let (_temp, dataset) = setup_repo_and_dataset();
@@ -50,6 +90,22 @@ fn text_index_insert_search_basic() {
     assert!(hits[0].score < 1e-4);
 }
 
+#[test]
+fn text_index_failed_reinsert_preserves_existing_document() {
+    let (_temp, dataset) = setup_repo_and_dataset();
+    let mut store = FileNamespacedKvStore::<N>::init(&dataset).unwrap();
+    let mut personal = store.namespace("personal");
+    let mut docs = personal.text_index("docs", failing_cfg(8, 0)).unwrap();
+
+    docs.insert(b"doc:1", "stable text").unwrap();
+    let err = docs.insert(b"doc:1", "please fail-insert").unwrap_err();
+    assert!(err.to_string().contains("forced insert failure"));
+
+    let hits = docs.search("stable text", 1).unwrap();
+    assert_eq!(hits[0].id, b"doc:1".to_vec());
+    assert!(hits[0].score < 1e-4);
+}
+
 #[test]
 fn text_index_survives_commit_and_reopen() {
     let (_temp, dataset) = setup_repo_and_dataset();
@@ -79,8 +135,6 @@ fn text_index_survives_commit_and_reopen() {
 #[test]
 fn reopen_with_different_embedder_id_returns_mismatch() {
     // Different embedder family at re-open time → EmbedderMismatch.
-    use prollytree::proximity::{EmbedError, Embedder};
-
     struct DifferentFamily(HashEmbedder);
     impl Embedder for DifferentFamily {
         fn id(&self) -> &str {