Fix embedding idempotency and test stability

Author: shisoft

src/apps/embedding/mod.rs

@@ -188,11 +188,13 @@ impl EmbeddingIndexer {
             vector,
         );
 
+        // Embedding cell IDs are deterministic per document/model pair, so retries and
+        // duplicate indexing requests must be idempotent.
         self.neb_client
-            .write_cell(emb_cell)
+            .upsert_cell(emb_cell)
             .await
-            .map_err(|e| IndexError::Other(format!("Failed to write embedding cell: {:?}", e)))?
-            .map_err(|e| IndexError::Other(format!("Failed to write embedding cell: {:?}", e)))?;
+            .map_err(|e| IndexError::Other(format!("Failed to persist embedding cell: {:?}", e)))?
+            .map_err(|e| IndexError::Other(format!("Failed to persist embedding cell: {:?}", e)))?;
 
         Ok(())
     }

src/apps/embedding/tests.rs

@@ -354,13 +354,23 @@ mod provider_tests {
 
 #[cfg(test)]
 mod integration_tests {
-    use dovahkiin::types::{Id, Map, OwnedMap, OwnedValue, Type};
+    use dovahkiin::types::{Id, Map, OwnedMap, OwnedPrimArray, OwnedValue, Type};
     use neb::{
         index::embedding::EmbeddingModel,
+        index::vector::HnswConfig,
         ram::schema::{Field, IndexType, Schema},
         ram::types::RandValue,
     };
 
+    use crate::apps::{
+        embedding::{
+            EmbeddingIndexer,
+            schema::{
+                create_embedding_cell_schema_for_client, embedding_cell_id, vector_field_id,
+            },
+        },
+        hnsw::measurements::MetricEncoding,
+    };
     use crate::tests::start_server;
 
     const TEST_SCHEMA_ID: u32 = 9001;
@@ -802,6 +812,78 @@ mod integration_tests {
         }
     }
 
+    #[tokio::test]
+    #[ignore] // Run with: cargo test test_embedding_cell_write_is_idempotent -- --ignored --exact --nocapture
+    async fn test_embedding_cell_write_is_idempotent() {
+        let server = start_server(6021, "embedding_upsert_regression").await.unwrap();
+
+        let _partition = server.init_hnsw_index_partition_service().await.unwrap();
+        let _hnsw = server.init_hnsw_index_service().await.unwrap();
+
+        let runtime = server.current_runtime();
+        let indexer = EmbeddingIndexer::new(&runtime).await.unwrap();
+
+        const DIRECT_SCHEMA_ID: u32 = 9002;
+        let field_id = bifrost_hasher::hash_str(TEST_FIELD);
+        let full_model = "test:test-model";
+        let dimensions = 4usize;
+
+        let emb_schema_id = create_embedding_cell_schema_for_client(
+            server.neb_client.as_ref(),
+            full_model,
+            DIRECT_SCHEMA_ID,
+            field_id,
+            dimensions,
+        )
+        .await
+        .unwrap();
+
+        indexer
+            .hnsw_coordinator
+            .new_index(
+                format!("EMB-{DIRECT_SCHEMA_ID}-{field_id}-{full_model}"),
+                emb_schema_id,
+                vector_field_id(),
+                HnswConfig::default(),
+            )
+            .await
+            .unwrap()
+            .unwrap();
+
+        let doc_id = Id::rand();
+        let vector = vec![0.5, 0.5, 0.5, 0.5];
+
+        indexer
+            .write_embedding_cell(&doc_id, DIRECT_SCHEMA_ID, field_id, full_model, vector.clone())
+            .await
+            .unwrap();
+
+        indexer
+            .write_embedding_cell(&doc_id, DIRECT_SCHEMA_ID, field_id, full_model, vector.clone())
+            .await
+            .unwrap();
+
+        let emb_cell_id = embedding_cell_id(&doc_id, DIRECT_SCHEMA_ID, field_id, full_model);
+        server.neb_client.read_cell(emb_cell_id).await.unwrap().unwrap();
+
+        let hits = indexer
+            .hnsw_coordinator
+            .query_top_k(
+                emb_schema_id,
+                vector_field_id(),
+                OwnedPrimArray::F32(vector),
+                10,
+                32,
+                MetricEncoding::Cosine,
+            )
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(hits.len(), 1, "duplicate writes should not create duplicate vector entries");
+        assert_eq!(hits[0].0, emb_cell_id);
+    }
+
     /// Test multi-model support on the same field
     #[tokio::test]
     #[ignore]

src/apps/recovery_integrity_tests.rs

@@ -242,6 +242,7 @@ pub async fn start_test_server(
         Service::RangedIndexer,
         Service::Query,
     ];
+    server_config.storage.disable_storage_locks = true;
 
     // Configure persistent storage
     server_config.storage.backup_storage = Some(config.backup_path());

src/apps/wikidata/tests.rs

@@ -81,6 +81,7 @@ fn create_test_config(port: u32, group_name: &str, enable_recovery: bool) -> Mor
         Service::RangedIndexer,
         Service::Query,
     ];
+    opts.storage.disable_storage_locks = true;
     opts.storage.backup_storage = Some(backup_path.clone());
     opts.storage.wal_storage = Some(wal_path);
     opts.storage.raft_storage = Some(raft_path);

src/config/mod.rs

@@ -101,6 +101,7 @@ impl From<HumanReadableMorphesOptions> for MorphesOptions {
                 index_enabled: hr_options.storage.index_enabled,
                 raft_storage: hr_options.storage.raft_storage,
                 undo_log_storage: hr_options.storage.undo_log_storage,
+                disable_storage_locks: false,
                 tiered_config: hr_options.storage.tiered_config.map(|tc| {
                     neb::ram::tiered::TieredConfig {
                         threshold: tc.threshold as f32,

src/config/tests.rs

@@ -48,6 +48,7 @@ storage:
         assert_eq!(config.storage.db_size, 1024 * 1024 * 1024); // Should be converted to bytes
         assert_eq!(config.server_addr, "127.0.0.1:5400");
         assert_eq!(config.group_name, "TestMorpheus");
+        assert!(!config.storage.disable_storage_locks);
     }
 
     #[test]
@@ -81,5 +82,6 @@ storage:
 
         // Verify that tiered memory configuration is None (disabled)
         assert!(config.storage.tiered_config.is_none());
+        assert!(!config.storage.disable_storage_locks);
     }
 }

src/tests/compute/mod.rs

@@ -175,7 +175,10 @@ impl GasProgram for TestProgram {
         accum: &<Self::C as GasCodec>::Accum,
         _ctx: &mut ApplyCtx,
     ) {
-        *state = *accum;
+        if *state != *accum {
+            *state = *accum;
+            _ctx.updated = true;
+        }
     }
 
     fn scatter(

src/tests/mod.rs

@@ -82,6 +82,7 @@ pub fn start_server_with_storage_and_recovery<'a>(
         Service::RangedIndexer,
         Service::Query,
     ];
+    config.storage.disable_storage_locks = true;
 
     // Configure persistent storage if paths provided
     // All storage paths are required for full persistence functionality
@@ -137,6 +138,7 @@ pub fn start_server_in_cluster<'a>(
         Service::Query,
     ];
     config.storage.index_enabled = true;
+    config.storage.disable_storage_locks = true;
     MorpheusServer::new(config)
 }