fix(embeddings): recover from Ollama NaN-encoding 500 instead of failing whole batch (#2834)

src/openhuman/embeddings/ollama.rs

@@ -152,6 +152,104 @@ impl OllamaEmbedding {
             .map_err(|e| anyhow::anyhow!("invalid Ollama base_url `{}`: {e}", self.base_url))?;
         Ok(format!("{}/api/embed", self.base_url))
     }
+
+    /// Sends a single text to Ollama and returns either the embedding, or
+    /// `None` if Ollama produced the NaN-encoding 500 wire shape for that
+    /// individual input. Any other failure (transport error, non-2xx without
+    /// NaN signature, malformed JSON, dimension mismatch, count mismatch)
+    /// surfaces as an `Err` so genuine bugs are still loud.
+    ///
+    /// Used by [`Self::embed_per_text_fallback`] to recover a batch that
+    /// failed wholesale on NaN — see TAURI-RUST-AZ.
+    async fn embed_one_with_nan_recovery(&self, text: &str) -> anyhow::Result<Option<Vec<f32>>> {
+        let resp = self
+            .http_client()
+            .post(self.embed_url()?)
+            .json(&OllamaEmbedRequest {
+                model: self.model.clone(),
+                input: vec![text.to_string()],
+            })
+            .send()
+            .await
+            .map_err(|e| {
+                anyhow::anyhow!(
+                    "ollama embed request failed (is Ollama running at {}?): {e}",
+                    self.base_url
+                )
+            })?;
+
+        if !resp.status().is_success() {
+            let status = resp.status();
+            let body = resp.text().await.unwrap_or_default();
+            let detail = body.trim();
+            // Per-text NaN: substitute empty vector, log once.
+            if status.as_u16() == 500 && is_nan_encode_error(&body) {
+                tracing::warn!(
+                    target: "embeddings.ollama",
+                    "[embeddings] ollama produced NaN for a single text (model={}); \
+                     substituting empty embedding (downstream blob will be 0 bytes)",
+                    self.model
+                );
+                return Ok(None);
+            }
+            anyhow::bail!(
+                "ollama embed failed with status {status}{}",
+                if detail.is_empty() {
+                    String::new()
+                } else {
+                    format!(": {detail}")
+                }
+            );
+        }
+
+        let payload: OllamaEmbedResponse = resp
+            .json()
+            .await
+            .map_err(|e| anyhow::anyhow!("ollama embed response parse failed: {e}"))?;
+        if payload.embeddings.len() != 1 {
+            anyhow::bail!(
+                "ollama embed count mismatch: sent 1 text, got {} embeddings",
+                payload.embeddings.len()
+            );
+        }
+        let v = payload.embeddings.into_iter().next().unwrap();
+        if v.len() != self.dims {
+            anyhow::bail!(
+                "ollama embed dimension mismatch: expected {}, got {}",
+                self.dims,
+                v.len()
+            );
+        }
+        Ok(Some(v))
+    }
+
+    /// Recovery path invoked when a batch request fails with the Ollama
+    /// NaN-encoding 500 wire shape. Re-sends each live input one at a time;
+    /// per-text NaN failures are filled with `Vec::new()` so the overall
+    /// result vector still has length `total_len` and aligns with the
+    /// caller's original `texts` slice (same convention as blank inputs).
+    async fn embed_per_text_fallback(
+        &self,
+        total_len: usize,
+        live: &[(usize, String)],
+    ) -> anyhow::Result<Vec<Vec<f32>>> {
+        let mut result = vec![Vec::new(); total_len];
+        let mut nan_substituted = 0usize;
+        for (orig_idx, text) in live {
+            match self.embed_one_with_nan_recovery(text).await? {
+                Some(v) => result[*orig_idx] = v,
+                None => nan_substituted += 1,
+            }
+        }
+        tracing::warn!(
+            target: "embeddings.ollama",
+            "[embeddings] ollama per-text fallback complete: {} of {} inputs returned NaN and \
+             were substituted with empty embeddings",
+            nan_substituted,
+            live.len()
+        );
+        Ok(result)
+    }
 }
 
 impl Default for OllamaEmbedding {
@@ -179,6 +277,18 @@ struct OllamaEmbedResponse {
     embeddings: Vec<Vec<f32>>,
 }
 
+/// Detects the Ollama-side NaN-encoding 500 wire shape:
+///   `{"error":"failed to encode response: json: unsupported value: NaN"}`
+///
+/// Ollama produces NaN for some inputs (model bug / numerically degenerate
+/// token sequences) and then fails to encode the response as JSON. One bad
+/// input poisons the entire batch.
+///
+/// See TAURI-RUST-AZ on Sentry.
+fn is_nan_encode_error(body: &str) -> bool {
+    body.to_ascii_lowercase().contains("unsupported value: nan")
+}
+
 #[async_trait]
 impl EmbeddingProvider for OllamaEmbedding {
     fn name(&self) -> &str {
@@ -266,6 +376,40 @@ impl EmbeddingProvider for OllamaEmbedding {
                     format!(": {detail}")
                 }
             );
+
+            // TAURI-RUST-AZ: Ollama returns 500 `unsupported value: NaN` when the
+            // model produces NaN for some input in the batch. One bad input
+            // poisons the whole batch under the default code path. Recover by
+            // re-sending each live input individually; per-text NaN failures
+            // are replaced with an empty embedding (same convention used for
+            // blank inputs above), so the rest of the batch still succeeds.
+            if status.as_u16() == 500 && is_nan_encode_error(&body) {
+                tracing::warn!(
+                    target: "embeddings.ollama",
+                    "[embeddings] ollama returned NaN-encoding 500 for batch of {} (model={}); \
+                     falling back to per-text requests",
+                    live.len(),
+                    self.model
+                );
+                crate::core::observability::report_error_or_expected(
+                    &message,
+                    "embeddings",
+                    "ollama_embed",
+                    &[
+                        ("model", self.model.as_str()),
+                        ("status", status_str.as_str()),
+                        ("failure", "nan_batch_recovered"),
+                    ],
+                );
+                // Single-text NaN: re-issuing the same request would only
+                // reproduce the failure. Skip the wasted round-trip and
+                // substitute an empty embedding directly.
+                if live.len() == 1 {
+                    return Ok(vec![Vec::new(); texts.len()]);
+                }
+                return self.embed_per_text_fallback(texts.len(), &live).await;
+            }
+
             crate::core::observability::report_error_or_expected(
                 &message,
                 "embeddings",

src/openhuman/embeddings/ollama_tests.rs

@@ -441,6 +441,184 @@ fn ollama_dimension_mismatch_stays_unexpected() {
     );
 }
 
+// ── embed — NaN-encoding 500 recovery (TAURI-RUST-AZ) ───
+
+#[test]
+fn is_nan_encode_error_matches_ollama_wire_shape() {
+    // Canonical wire shape from Sentry TAURI-RUST-AZ.
+    assert!(is_nan_encode_error(
+        r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#
+    ));
+    // Case-insensitive on the "NaN" token.
+    assert!(is_nan_encode_error("unsupported value: nan"));
+    assert!(is_nan_encode_error("UNSUPPORTED VALUE: NAN"));
+    // Negative: unrelated 500 bodies must not match.
+    assert!(!is_nan_encode_error("model crashed"));
+    assert!(!is_nan_encode_error(
+        r#"{"error":"failed to encode response: json: unsupported value: +Inf"}"#
+    ));
+    assert!(!is_nan_encode_error(""));
+}
+
+#[tokio::test]
+async fn embed_nan_batch_recovers_via_per_text() {
+    // Ollama returns NaN-encoding 500 for any batch > 1, but succeeds on
+    // single-text requests. Recovery should re-issue per-text and assemble
+    // the full result.
+    let app = Router::new().route(
+        "/api/embed",
+        post(|Json(body): Json<serde_json::Value>| async move {
+            let n = body["input"].as_array().unwrap().len();
+            if n > 1 {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    String::from(
+                        r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#,
+                    ),
+                )
+            } else {
+                (
+                    StatusCode::OK,
+                    serde_json::json!({ "embeddings": [[1.0, 2.0]] }).to_string(),
+                )
+            }
+        }),
+    );
+    let url = start_mock(app).await;
+    let p = OllamaEmbedding::new(&url, "m", 2);
+
+    let result = p.embed(&["a", "b", "c"]).await.unwrap();
+    assert_eq!(result.len(), 3);
+    assert_eq!(result[0], vec![1.0, 2.0]);
+    assert_eq!(result[1], vec![1.0, 2.0]);
+    assert_eq!(result[2], vec![1.0, 2.0]);
+}
+
+#[tokio::test]
+async fn embed_nan_persists_per_text_substitutes_empty() {
+    // Every request — batch or single — returns NaN 500. Recovery
+    // substitutes empty vectors for each input rather than bailing.
+    let app = Router::new().route(
+        "/api/embed",
+        post(|| async {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                String::from(
+                    r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#,
+                ),
+            )
+        }),
+    );
+    let url = start_mock(app).await;
+    let p = OllamaEmbedding::new(&url, "m", 2);
+
+    let result = p.embed(&["a", "b"]).await.unwrap();
+    assert_eq!(result.len(), 2);
+    assert!(result[0].is_empty(), "NaN-failed entries must be empty vec");
+    assert!(result[1].is_empty(), "NaN-failed entries must be empty vec");
+}
+
+#[tokio::test]
+async fn embed_nan_single_text_short_circuits_to_empty() {
+    // Single-input NaN should NOT trigger a wasted retry — short-circuit
+    // straight to empty vector.
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::Arc;
+    let count = Arc::new(AtomicUsize::new(0));
+    let count_clone = count.clone();
+    let app = Router::new().route(
+        "/api/embed",
+        post(move || {
+            let c = count_clone.clone();
+            async move {
+                c.fetch_add(1, Ordering::SeqCst);
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    String::from(
+                        r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#,
+                    ),
+                )
+            }
+        }),
+    );
+    let url = start_mock(app).await;
+    let p = OllamaEmbedding::new(&url, "m", 2);
+
+    let result = p.embed(&["only-text"]).await.unwrap();
+    assert_eq!(result.len(), 1);
+    assert!(result[0].is_empty());
+    assert_eq!(
+        count.load(Ordering::SeqCst),
+        1,
+        "single-text NaN must not retry — exactly one request expected"
+    );
+}
+
+#[tokio::test]
+async fn embed_nan_mixed_per_text_outcomes() {
+    // Batch NaN-fails. Per-text retries: input "bad" returns NaN, others
+    // succeed. Result vector has empty slot for "bad" and embeddings for
+    // the rest, with positions preserved.
+    let app = Router::new().route(
+        "/api/embed",
+        post(|Json(body): Json<serde_json::Value>| async move {
+            let inputs = body["input"].as_array().unwrap();
+            if inputs.len() > 1 {
+                return (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    String::from(
+                        r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#,
+                    ),
+                );
+            }
+            let text = inputs[0].as_str().unwrap();
+            if text == "bad" {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    String::from(
+                        r#"{"error":"failed to encode response: json: unsupported value: NaN"}"#,
+                    ),
+                )
+            } else {
+                (
+                    StatusCode::OK,
+                    serde_json::json!({ "embeddings": [[9.0, 9.0]] }).to_string(),
+                )
+            }
+        }),
+    );
+    let url = start_mock(app).await;
+    let p = OllamaEmbedding::new(&url, "m", 2);
+
+    let result = p.embed(&["good1", "bad", "good2"]).await.unwrap();
+    assert_eq!(result.len(), 3);
+    assert_eq!(result[0], vec![9.0, 9.0]);
+    assert!(result[1].is_empty(), "NaN entry must be empty vec");
+    assert_eq!(result[2], vec![9.0, 9.0]);
+}
+
+#[tokio::test]
+async fn embed_non_nan_500_still_errors() {
+    // Real ollama 500s (anything that isn't the NaN wire shape) must still
+    // bail loudly — only NaN gets the recovery path.
+    let app = Router::new().route(
+        "/api/embed",
+        post(|| async {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                String::from("model crashed"),
+            )
+        }),
+    );
+    let url = start_mock(app).await;
+    let p = OllamaEmbedding::new(&url, "m", 2);
+
+    let err = p.embed(&["a", "b"]).await.unwrap_err();
+    let msg = err.to_string();
+    assert!(msg.contains("500"));
+    assert!(msg.contains("model crashed"));
+}
+
 // ── embed_one (trait default) ───────────────────────────
 
 #[tokio::test]