Harden embedding: atomic model switch, consistent vectors, downloads

app/src-tauri/src/lib.rs

@@ -4424,20 +4424,20 @@ fn apply_embedding_model(app: &AppHandle, id: Option<String>) -> Result<(), Stri
                 return Err(format!("embedding model not downloaded: {model_id}"));
             }
 
-            load_embedder(app, model_id)?;
-            let _ = fs::write(&state.embed_model_path, model_id);
-
-            // Re-embed existing notes with the new model so semantic search covers them too (and so a
-            // model switch never mixes vectors of different dimensions). Best-effort: the model is
-            // already active for new saves even if this pass fails.
-            if let Err(e) = state
+            // Switch the embedder and re-embed existing notes in one locked, atomic step: either the
+            // new model becomes active with a freshly-rebuilt index, or — on any failure — the
+            // previous model and its index are left untouched and the error surfaces to the UI. This
+            // avoids both a half-applied switch (new embedder vs old-dimension vectors) and a
+            // silently-wiped index.
+            let embedder = resolve_embedder(&state, model_id)?;
+            state
                 .library
                 .lock()
                 .map_err(|_| "library lock poisoned".to_owned())?
-                .reindex_all()
-            {
-                eprintln!("reindex after embedding-model change failed: {e}");
-            }
+                .set_embedder_and_reindex(Some(embedder))
+                .map_err(|e| format!("indexing notes with {model_id} failed: {e}"))?;
+
+            let _ = fs::write(&state.embed_model_path, model_id);
         }
         None => {
             state

crates/wisp-embed/src/lib.rs

@@ -267,9 +267,24 @@ fn fetch_atomic(url: &str, dest: &Path) -> Result<()> {
         let resp = ureq::get(url)
             .call()
             .map_err(|e| LibraryError::Embed(format!("download {url}: {e}")))?;
+
+        // Hold the transfer to the declared length. Without this, a connection dropped mid-body on a
+        // length-less (close-delimited) response reads as a clean EOF, and the truncated file gets
+        // renamed into place as "complete" — then it fails to load forever, since the present (short)
+        // file makes the download skip it on every retry.
+        let expected: Option<u64> = resp.header("Content-Length").and_then(|h| h.parse().ok());
+
         let mut reader = resp.into_reader();
         let mut out = std::fs::File::create(&part).map_err(io_err)?;
-        std::io::copy(&mut reader, &mut out).map_err(io_err)?;
+        let written = std::io::copy(&mut reader, &mut out).map_err(io_err)?;
+
+        if let Some(total) = expected {
+            if written != total {
+                return Err(LibraryError::Embed(format!(
+                    "download {url}: truncated ({written} of {total} bytes)"
+                )));
+            }
+        }
         Ok(())
     })();
 

crates/wisp-embed/src/ort_embed.rs

@@ -130,6 +130,14 @@ impl OrtEmbedder {
         let seq = encodings[0].len();
         let batch = encodings.len();
 
+        // A zero-length encoding (a pathological input/tokenizer) would otherwise yield a 0-width
+        // tensor and an all-zero "embedding"; fail loudly instead.
+        if seq == 0 {
+            return Err(LibraryError::Embed(
+                "tokenizer produced an empty encoding".to_owned(),
+            ));
+        }
+
         let mut ids = Vec::with_capacity(batch * seq);
         let mut mask = Vec::with_capacity(batch * seq);
         let mut types = Vec::with_capacity(batch * seq);
@@ -279,11 +287,18 @@ impl Embedder for OrtEmbedder {
     }
 
     fn embed_passages(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
-        let prefixed: Vec<String> = texts
+        // Embed one text at a time (batch size 1). Batching pads the shorter rows, and a quantized
+        // model's per-batch scales make a padded row drift from the same text embedded alone — so a
+        // stored passage would not match a (single-text) query embedding, and two notes would index
+        // the same sentence differently depending on their neighbours. One at a time keeps every
+        // vector identical regardless of the batch it was embedded in.
+        texts
             .iter()
-            .map(|t| format!("{}{t}", self.passage_prefix))
-            .collect();
-        self.run(&prefixed)
+            .map(|t| {
+                let prefixed = format!("{}{t}", self.passage_prefix);
+                Ok(self.run(&[prefixed])?.pop().unwrap_or_default())
+            })
+            .collect()
     }
 
     fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
@@ -747,6 +762,52 @@ mod tests {
         );
     }
 
+    // A passage must embed identically whether alone or alongside others. `embed_passages` embeds one
+    // text at a time precisely so a quantized model's per-batch scales can't make a stored vector
+    // depend on its neighbours (which would also stop it matching a single-text query). Uses the
+    // cached Qwen3 model. `cargo test -- --ignored`.
+    #[test]
+    #[ignore]
+    fn qwen3_passage_embedding_is_independent_of_neighbours() {
+        let dir = std::env::temp_dir().join("wisp-ort-qwen3-0_6b");
+        let repo = "onnx-community/Qwen3-Embedding-0.6B-ONNX";
+        for f in [
+            "onnx/model_quantized.onnx",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "config.json",
+            "special_tokens_map.json",
+        ] {
+            download_file(repo, f, &dir.join(f));
+        }
+
+        let recipe = Recipe {
+            onnx_file: "onnx/model_quantized.onnx",
+            pooling: Pooling::LastToken,
+            passage_prefix: "",
+            query_prefix: "",
+            normalize: true,
+            dim: 1024,
+            max_length: 512,
+        };
+        let emb = OrtEmbedder::load(&dir, &recipe).unwrap();
+
+        let short = "budget";
+        let long = "the annual operating budget was approved by the board after a long and detailed \
+                    discussion covering every department and each contingency line item in the plan";
+
+        let alone = emb.embed_passages(&[short]).unwrap().pop().unwrap();
+        // Batching the short row with a much longer one forces heavy right-padding on the short row.
+        let batched = emb.embed_passages(&[short, long]).unwrap();
+
+        let cos = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>();
+        let sim = cos(&alone, &batched[0]);
+        assert!(
+            sim > 0.999,
+            "a passage embedded differently next to a neighbour (cos {sim}); embed_passages must embed one at a time"
+        );
+    }
+
     // BGE-M3 (XLM-RoBERTa encoder, CLS pooling, 1024-dim, 100+ languages) via the int8 ONNX (~600 MB).
     #[test]
     #[ignore]

crates/wisp-library/src/store.rs

@@ -113,6 +113,26 @@ impl Library {
         self.embedder = embedder;
     }
 
+    /// Atomically switches the active embedder and re-embeds the whole corpus with it. On any failure
+    /// the library is left exactly as it was — previous embedder, previous chunks — so a model switch
+    /// never wipes the index or leaves a new embedder pointed at stale-dimension vectors. Use this
+    /// (not [`set_embedder`] + [`reindex_all`]) whenever the model changes.
+    pub fn set_embedder_and_reindex(
+        &mut self,
+        embedder: Option<Box<dyn Embedder>>,
+    ) -> Result<usize> {
+        let previous = self.embedder.take();
+        self.embedder = embedder;
+
+        match self.reindex_all() {
+            Ok(n) => Ok(n),
+            Err(e) => {
+                self.embedder = previous;
+                Err(e)
+            }
+        }
+    }
+
     fn init(conn: Connection) -> Result<Self> {
         conn.execute_batch("PRAGMA foreign_keys = ON;")?;
         let mut lib = Self {
@@ -199,14 +219,13 @@ impl Library {
     /// Returns how many notes produced chunks; clears all chunks and returns 0 with no embedder set.
     /// A personal library is small, so it re-embeds the whole corpus in a single pass.
     pub fn reindex_all(&mut self) -> Result<usize> {
-        self.conn.execute("DELETE FROM chunk", [])?;
-
         if self.embedder.is_none() {
+            self.conn.execute("DELETE FROM chunk", [])?;
             return Ok(0);
         }
 
         // Embed every note before opening the write transaction — inference touches no DB and may be
-        // slow.
+        // slow. A note that fails to embed returns here with the existing index still intact.
         let mut indexed: Vec<(String, Vec<EmbeddedChunk>)> = Vec::new();
         for id in self.note_ids()? {
             let texts = self.segment_texts(&id)?;
@@ -219,7 +238,11 @@ impl Library {
             }
         }
 
+        // Clear the old chunks and insert the new ones in one transaction, so a failure (or a crash)
+        // never leaves the corpus de-indexed. The DELETE used to run + commit before embedding, so a
+        // mid-embed error wiped the whole index for good.
         let tx = self.conn.transaction()?;
+        tx.execute("DELETE FROM chunk", [])?;
         for (id, chunks) in &indexed {
             insert_chunks(&tx, id, chunks)?;
         }
@@ -707,6 +730,28 @@ mod tests {
         }
     }
 
+    /// A test embedder whose passage embedding always fails — used to drive a reindex error and prove
+    /// the corpus + previous embedder survive it.
+    struct FailingEmbedder;
+
+    impl Embedder for FailingEmbedder {
+        fn dim(&self) -> usize {
+            64
+        }
+
+        fn embed_passages(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>> {
+            Err(crate::LibraryError::Embed(
+                "simulated inference failure".to_owned(),
+            ))
+        }
+
+        fn embed_query(&self, _text: &str) -> Result<Vec<f32>> {
+            Err(crate::LibraryError::Embed(
+                "simulated inference failure".to_owned(),
+            ))
+        }
+    }
+
     /// An in-memory library with the deterministic test embedder configured.
     fn lib_with_embedder() -> Library {
         let mut lib = Library::open_in_memory().unwrap();
@@ -1179,6 +1224,45 @@ mod tests {
         );
     }
 
+    #[test]
+    fn a_failed_model_switch_keeps_the_old_index_and_embedder() {
+        let mut lib = lib_with_embedder();
+        lib.save_note(
+            "m1",
+            &meta("M"),
+            0,
+            &[seg(1, 0, 100, "budget review", AudioSourceKind::Microphone)],
+        )
+        .unwrap();
+        assert_eq!(lib.reindex_all().unwrap(), 1);
+        assert_eq!(
+            lib.search_semantic("budget", 10).unwrap()[0].meeting_id,
+            "m1"
+        );
+
+        // Switching to a model that fails to embed must NOT wipe the index or leave the failing model
+        // active: the call errors, and the previous (working) embedder + chunks are restored intact.
+        assert!(lib
+            .set_embedder_and_reindex(Some(Box::new(FailingEmbedder)))
+            .is_err());
+        assert_eq!(
+            lib.search_semantic("budget", 10).unwrap()[0].meeting_id,
+            "m1",
+            "a failed model switch must leave the old index + embedder intact"
+        );
+
+        // A switch to a working model still succeeds and rebuilds the index.
+        assert_eq!(
+            lib.set_embedder_and_reindex(Some(Box::new(HashingEmbedder { dim: 128 })))
+                .unwrap(),
+            1
+        );
+        assert_eq!(
+            lib.search_semantic("budget", 10).unwrap()[0].meeting_id,
+            "m1"
+        );
+    }
+
     #[test]
     fn reindex_without_an_embedder_clears_all_chunks() {
         let mut lib = lib_with_embedder();