fix(memory): batch embedding retries (#3422)

src/openhuman/embeddings/retry_after.rs

@@ -4,8 +4,11 @@
 //! - A non-negative integer: delta-seconds from now (e.g. `Retry-After: 30`)
 //! - An HTTP-date string: absolute point in time (e.g. `Retry-After: Wed, 21 Oct 2015 07:28:00 GMT`)
 //!
-//! This module prefers the delta-seconds form and falls back to exponential
-//! backoff when the header is absent or unparseable. See RFC 9110 §10.2.4.
+//! This module prefers the delta-seconds form, accepts the common HTTP-date
+//! form, and falls back to exponential backoff when the header is absent or
+//! unparseable. See RFC 9110 §10.2.4.
+
+use chrono::{DateTime, Utc};
 
 /// Maximum number of 429/503 retries before giving up.
 ///
@@ -34,10 +37,25 @@ pub const MAX_BACKOFF_MS: u64 = 30_000;
 /// Returns `None` when the header is absent, empty, or not a valid
 /// non-negative integer.
 pub fn parse_retry_after_ms(header_value: Option<&str>) -> Option<u64> {
+    parse_retry_after_ms_at(header_value, Utc::now())
+}
+
+fn parse_retry_after_ms_at(header_value: Option<&str>, now: DateTime<Utc>) -> Option<u64> {
     let s = header_value?.trim();
-    // Only accept the delta-seconds form: a non-negative integer.
-    let secs: u64 = s.parse().ok()?;
-    Some(secs.saturating_mul(1_000).min(MAX_BACKOFF_MS))
+    if s.is_empty() {
+        return None;
+    }
+    if let Ok(secs) = s.parse::<u64>() {
+        return Some(secs.saturating_mul(1_000).min(MAX_BACKOFF_MS));
+    }
+    let retry_at = DateTime::parse_from_rfc2822(s).ok()?.with_timezone(&Utc);
+    let delay = retry_at.signed_duration_since(now);
+    if delay.num_milliseconds() <= 0 {
+        return Some(0);
+    }
+    u64::try_from(delay.num_milliseconds())
+        .ok()
+        .map(|ms| ms.min(MAX_BACKOFF_MS))
 }
 
 /// Compute the delay for attempt `n` (0-indexed) with optional `Retry-After`
@@ -81,10 +99,23 @@ mod tests {
 
     #[test]
     fn returns_none_for_http_date() {
-        // HTTP-date form — not parsed; fall through to exponential backoff.
+        let now = DateTime::parse_from_rfc2822("Wed, 21 Oct 2015 07:27:55 GMT")
+            .unwrap()
+            .with_timezone(&Utc);
         assert_eq!(
-            parse_retry_after_ms(Some("Wed, 21 Oct 2015 07:28:00 GMT")),
-            None
+            parse_retry_after_ms_at(Some("Wed, 21 Oct 2015 07:28:00 GMT"), now),
+            Some(5_000)
+        );
+    }
+
+    #[test]
+    fn http_date_in_the_past_returns_zero() {
+        let now = DateTime::parse_from_rfc2822("Wed, 21 Oct 2015 07:28:05 GMT")
+            .unwrap()
+            .with_timezone(&Utc);
+        assert_eq!(
+            parse_retry_after_ms_at(Some("Wed, 21 Oct 2015 07:28:00 GMT"), now),
+            Some(0)
         );
     }
 

src/openhuman/memory_queue/handlers/mod.rs

@@ -19,6 +19,7 @@ use crate::openhuman::memory_queue::types::{
     JobOutcome, NewJob, NodeRef, ReembedBackfillPayload, SealDocumentPayload, SealPayload,
 };
 use crate::openhuman::memory_store::chunks::store as chunk_store;
+use crate::openhuman::memory_store::chunks::types::Chunk;
 use crate::openhuman::memory_store::content::{
     self as content_store, read as content_read, tags as content_tags,
 };
@@ -32,6 +33,14 @@ use crate::openhuman::memory_tree::tree::{LeafRef, TreeFactory};
 /// 1 hour means low-volume sources get summaries within a working session.
 const L0_DEFAULT_FLUSH_AGE_SECS: i64 = 60 * 60;
 
+/// Maximum `extract_chunk` jobs to coalesce in one worker tick.
+///
+/// Scoring/extraction still runs per chunk, but kept chunks share one
+/// `Embedder::embed_batch` call. This is the high-volume memory-sync path:
+/// Slack/Gmail/etc. enqueue many extract jobs, and without coalescing each
+/// kept chunk produces its own embedding request.
+pub(crate) const EXTRACT_EMBED_BATCH: usize = 32;
+
 /// Derive the tree scope from a source_id. For GitHub per-item ids like
 /// `github:owner/repo:commit:sha` or `github:owner/repo:issue:42`,
 /// strips the item suffix and returns `github:owner/repo` so all items
@@ -168,6 +177,56 @@ async fn handle_seal_document(config: &Config, job: &Job) -> Result<JobOutcome>
 }
 
 async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
+    let mut results = handle_extract_batch(config, &[job.clone()]).await?;
+    results
+        .pop()
+        .expect("single extract batch returns one result")
+        .1
+}
+
+/// Handle a claimed run of `extract_chunk` jobs, batching the embedding
+/// sub-step while preserving per-job outcomes for worker settlement.
+pub async fn handle_extract_batch(
+    config: &Config,
+    jobs: &[Job],
+) -> Result<Vec<(Job, Result<JobOutcome>)>> {
+    let mut prepared = Vec::with_capacity(jobs.len());
+    let mut outcomes = Vec::new();
+
+    for job in jobs {
+        match prepare_extract(config, job).await {
+            Ok(Some(item)) => prepared.push(item),
+            Ok(None) => outcomes.push((job.clone(), Ok(JobOutcome::Done))),
+            Err(e) => outcomes.push((job.clone(), Err(e))),
+        }
+    }
+
+    attach_batched_embeddings(config, &mut prepared).await;
+
+    for mut item in prepared {
+        let job = item.job.clone();
+        let result = if let Some(err) = item.embedding_error.take() {
+            Err(err)
+        } else {
+            finalize_extract(config, item)
+        };
+        outcomes.push((job, result));
+    }
+
+    Ok(outcomes)
+}
+
+struct PreparedExtract {
+    job: Job,
+    chunk: Chunk,
+    body: String,
+    result: score::ScoreResult,
+    embedding: Option<Vec<f32>>,
+    embedding_error: Option<anyhow::Error>,
+    embedding_retry_needed: bool,
+}
+
+async fn prepare_extract(config: &Config, job: &Job) -> Result<Option<PreparedExtract>> {
     let payload: ExtractChunkPayload =
         serde_json::from_str(&job.payload_json).context("parse ExtractChunk payload")?;
     let Some(chunk) = chunk_store::get_chunk(config, &payload.chunk_id)? else {
@@ -175,7 +234,7 @@ async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
             "[memory::jobs] extract chunk missing chunk_id={}",
             payload.chunk_id
         );
-        return Ok(JobOutcome::Done);
+        return Ok(None);
     };
 
     // Read the full body from disk (the `content` column in SQLite holds a
@@ -197,66 +256,137 @@ async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
         None,
         None,
         None,
-        Some(format!(
-            "chunk {}",
-            &payload.chunk_id[..payload.chunk_id.len().min(16)]
-        )),
+        Some(format!("chunk {}", &chunk.id[..chunk.id.len().min(16)])),
     );
 
     let scoring_cfg = score::ScoringConfig::from_config(config);
     let result = score::score_chunk(&chunk_with_body, &scoring_cfg).await?;
-    let chunk_embedding: Option<Vec<f32>> = if result.kept {
-        // #002 (FR-002): when no usable embeddings provider is configured the
-        // write path returns None instead of an InertEmbedder — we SKIP
-        // embedding (the chunk is persisted embedding-less and re-embeddable
-        // later) rather than writing a fake all-zero vector that would
-        // silently poison semantic recall. `build_write_embedder` has already
-        // marked the process-global semantic-recall degraded flag with a typed
-        // cause for the status / doctor surface.
-        match build_write_embedder(config).context("build embedder in extract handler")? {
-            None => {
+    Ok(Some(PreparedExtract {
+        job: job.clone(),
+        chunk,
+        body,
+        result,
+        embedding: None,
+        embedding_error: None,
+        embedding_retry_needed: false,
+    }))
+}
+
+async fn attach_batched_embeddings(config: &Config, prepared: &mut [PreparedExtract]) {
+    let kept_count = prepared.iter().filter(|item| item.result.kept).count();
+    if kept_count == 0 {
+        return;
+    }
+
+    // #002 (FR-002): when no usable embeddings provider is configured the
+    // write path returns None instead of an InertEmbedder — we SKIP embedding
+    // rather than writing fake all-zero vectors.
+    let embedder = match build_write_embedder(config).context("build embedder in extract handler") {
+        Ok(Some(embedder)) => embedder,
+        Ok(None) => {
+            for item in prepared.iter().filter(|item| item.result.kept) {
                 log::warn!(
                     "[memory::jobs] extract chunk_id={} — embeddings unavailable, \
                      skipping embed (semantic recall degraded)",
-                    chunk.id
+                    item.chunk.id
                 );
-                None
-            }
-            Some(embedder) => {
-                // Reuse the body already read — avoid a second disk read.
-                let vector = match embedder.embed(&body).await {
-                    Ok(v) => v,
-                    Err(e) => {
-                        // #002: classify the embed failure so the worker can
-                        // fail fast on unrecoverable causes (budget/auth/dim)
-                        // and surface a typed reason, instead of burning the
-                        // retry budget. The typed failure is the outer
-                        // (downcast) error; the original chain is context.
-                        let failure =
-                            crate::openhuman::memory_tree::health::classify_embed_error(&e);
-                        return Err(anyhow::Error::new(failure).context(format!(
-                            "embed chunk_id={} in extract handler: {e:#}",
-                            chunk.id
-                        )));
-                    }
-                };
-                // Preserve the pre-cutover dimension guard (the job fails fast
-                // on a misconfigured embedder) even though #1574 no longer
-                // persists the packed blob to the legacy
-                // `mem_tree_chunks.embedding` column — the vector now goes to
-                // the per-model sidecar instead.
-                pack_checked(&vector).with_context(|| {
-                    format!("validate embedding dims for chunk_id={}", chunk.id)
-                })?;
-                // A real embed succeeded — recall is healthy again.
-                crate::openhuman::memory_tree::health::clear_semantic_recall_degraded();
-                Some(vector)
             }
+            return;
+        }
+        Err(e) => {
+            for item in prepared.iter_mut().filter(|item| item.result.kept) {
+                let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
+                item.embedding_retry_needed = !failure.is_unrecoverable();
+                log::warn!(
+                    "[memory::jobs] extract chunk_id={} — embedder build failed; \
+                     continuing without vector retry_needed={} err={e:#}",
+                    item.chunk.id,
+                    item.embedding_retry_needed
+                );
+            }
+            return;
         }
-    } else {
-        None
     };
 
+    let kept_indices: Vec<usize> = prepared
+        .iter()
+        .enumerate()
+        .filter_map(|(idx, item)| item.result.kept.then_some(idx))
+        .collect();
+    let texts: Vec<&str> = kept_indices
+        .iter()
+        .map(|idx| prepared[*idx].body.as_str())
+        .collect();
+
+    log::debug!(
+        "[memory::jobs] extract embedding batch: kept_chunks={} provider={}",
+        texts.len(),
+        embedder.name()
+    );
+    let vectors = embedder.embed_batch(&texts).await;
+    if vectors.len() != kept_indices.len() {
+        let err = anyhow::anyhow!(
+            "extract embed_batch returned {} results for {} texts",
+            vectors.len(),
+            kept_indices.len()
+        );
+        let failure = crate::openhuman::memory_tree::health::classify_embed_error(&err);
+        for idx in kept_indices {
+            prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
+            log::warn!(
+                "[memory::jobs] extract chunk_id={} — batch embedding contract failed; \
+                 continuing without vector retry_needed={} err={err:#}",
+                prepared[idx].chunk.id,
+                prepared[idx].embedding_retry_needed
+            );
+        }
+        return;
+    }
+
+    for (idx, vector_result) in kept_indices.into_iter().zip(vectors) {
+        match vector_result {
+            Ok(vector) => {
+                if let Err(e) = pack_checked(&vector).with_context(|| {
+                    format!(
+                        "validate embedding dims for chunk_id={}",
+                        prepared[idx].chunk.id
+                    )
+                }) {
+                    let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
+                    prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
+                    log::warn!(
+                        "[memory::jobs] extract chunk_id={} — embedding dim validation failed; \
+                         continuing without vector retry_needed={} err={e:#}",
+                        prepared[idx].chunk.id,
+                        prepared[idx].embedding_retry_needed
+                    );
+                } else {
+                    prepared[idx].embedding = Some(vector);
+                    crate::openhuman::memory_tree::health::clear_semantic_recall_degraded();
+                }
+            }
+            Err(e) => {
+                let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
+                prepared[idx].embedding_retry_needed = !failure.is_unrecoverable();
+                log::warn!(
+                    "[memory::jobs] extract chunk_id={} — embedding failed; \
+                     continuing without vector retry_needed={} err={e:#}",
+                    prepared[idx].chunk.id,
+                    prepared[idx].embedding_retry_needed
+                );
+            }
+        }
+    }
+}
+
+fn finalize_extract(config: &Config, item: PreparedExtract) -> Result<JobOutcome> {
+    let PreparedExtract {
+        chunk,
+        result,
+        embedding: chunk_embedding,
+        embedding_retry_needed,
+        ..
+    } = item;
     // Build follow-up job payloads before opening the tx — construction is
     // cheap and doesn't require a database connection. The two jobs are
     // enqueued inside the SAME transaction that commits the lifecycle update,
@@ -289,10 +419,7 @@ async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
         None,
         None,
         None,
-        Some(format!(
-            "chunk {}",
-            &payload.chunk_id[..payload.chunk_id.len().min(16)]
-        )),
+        Some(format!("chunk {}", &chunk.id[..chunk.id.len().min(16)])),
     );
 
     let active_sig = chunk_store::tree_active_signature(config);
@@ -393,6 +520,10 @@ async fn handle_extract(config: &Config, job: &Job) -> Result<JobOutcome> {
     if did_enqueue_source {
         super::worker::wake_workers();
     }
+    if result.kept && embedding_retry_needed {
+        crate::openhuman::memory_queue::ensure_reembed_backfill(config);
+        super::worker::wake_workers();
+    }
 
     Ok(JobOutcome::Done)
 }
@@ -776,8 +907,14 @@ async fn reembed_collect(
                 mark_skipped(config, id, active_sig, "embed wrong dim");
             }
             Err(e) => {
+                let failure = crate::openhuman::memory_tree::health::classify_embed_error(&e);
+                if !failure.is_unrecoverable() {
+                    return Err(anyhow::Error::new(failure).context(format!(
+                        "reembed_backfill: {label} {id} transient embed failed (sig={active_sig}): {e:#}"
+                    )));
+                }
                 log::warn!(
-                    "[memory::jobs] reembed_backfill: {label} {id} embed failed: {e}; skipping (sig={active_sig})"
+                    "[memory::jobs] reembed_backfill: {label} {id} embed failed with unrecoverable error: {e}; skipping (sig={active_sig})"
                 );
                 mark_skipped(config, id, active_sig, &format!("embed failed: {e}"));
             }

src/openhuman/memory_queue/store.rs

@@ -151,6 +151,59 @@ pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result<Option<Job>>
     })
 }
 
+/// Claim additional ready `extract_chunk` jobs for the same worker tick.
+///
+/// The caller has already claimed one job through [`claim_next`]. This helper
+/// leases up to `limit` more extract jobs so the handler can batch the
+/// embedding sub-step without widening non-extract scheduling semantics.
+pub fn claim_ready_extract_batch(
+    config: &Config,
+    lock_duration_ms: i64,
+    limit: usize,
+) -> Result<Vec<Job>> {
+    if limit == 0 {
+        return Ok(Vec::new());
+    }
+    with_connection(config, |conn| {
+        let now_ms = Utc::now().timestamp_millis();
+        let lock_until = now_ms.saturating_add(lock_duration_ms);
+        let mut stmt = conn
+            .prepare(
+                "UPDATE mem_tree_jobs
+                    SET status = 'running',
+                        attempts = attempts + 1,
+                        started_at_ms = ?1,
+                        locked_until_ms = ?2,
+                        last_error = NULL
+                  WHERE id IN (
+                      SELECT id FROM mem_tree_jobs
+                       WHERE status = 'ready'
+                         AND kind = 'extract_chunk'
+                         AND available_at_ms <= ?1
+                       ORDER BY available_at_ms ASC
+                       LIMIT ?3
+                  )
+              RETURNING id, kind, payload_json, dedupe_key, status, attempts,
+                        max_attempts, available_at_ms, locked_until_ms, last_error,
+                        created_at_ms, started_at_ms, completed_at_ms,
+                        failure_reason, failure_class",
+            )
+            .context("prepare claim_ready_extract_batch")?;
+        let jobs = stmt
+            .query_map(params![now_ms, lock_until, limit as i64], row_to_job)?
+            .collect::<rusqlite::Result<Vec<_>>>()
+            .context("query claim_ready_extract_batch")?;
+        if !jobs.is_empty() {
+            log::debug!(
+                "[memory::jobs] claimed extract batch count={} lock_until_ms={}",
+                jobs.len(),
+                lock_until
+            );
+        }
+        Ok(jobs)
+    })
+}
+
 /// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`.
 ///
 /// The UPDATE is gated on `attempts` and `started_at_ms` matching the values

src/openhuman/memory_queue/types.rs

@@ -405,7 +405,7 @@ impl NewJob {
             payload_json: serde_json::to_string(p)?,
             dedupe_key: Some(p.dedupe_key()),
             available_at_ms: None,
-            max_attempts: None,
+            max_attempts: Some(3),
         })
     }
 
@@ -636,6 +636,7 @@ mod tests {
         let job = NewJob::reembed_backfill(&payload).unwrap();
         assert_eq!(job.kind, JobKind::ReembedBackfill);
         assert_eq!(job.dedupe_key.as_deref(), Some("reembed_backfill:embed-v2"));
+        assert_eq!(job.max_attempts, Some(3));
         let roundtrip: ReembedBackfillPayload = serde_json::from_str(&job.payload_json).unwrap();
         assert_eq!(roundtrip.signature, "embed-v2");
     }

src/openhuman/memory_queue/worker.rs

@@ -18,10 +18,10 @@ use crate::openhuman::config::Config;
 use crate::openhuman::memory_queue::handlers;
 use crate::openhuman::memory_queue::redact::scrub_for_log;
 use crate::openhuman::memory_queue::store::{
-    claim_next, mark_deferred, mark_done, mark_failed_typed, recover_stale_locks,
-    release_running_locks, DEFAULT_LOCK_DURATION_MS,
+    claim_next, claim_ready_extract_batch, mark_deferred, mark_done, mark_failed_typed,
+    recover_stale_locks, release_running_locks, DEFAULT_LOCK_DURATION_MS,
 };
-use crate::openhuman::memory_queue::types::JobOutcome;
+use crate::openhuman::memory_queue::types::{Job, JobKind, JobOutcome};
 use crate::openhuman::memory_tree::health::PipelineFailure;
 
 /// Number of concurrent job-worker tasks. Each worker claims one job
@@ -214,7 +214,28 @@ pub async fn run_once(config: &Config) -> Result<bool> {
         None
     };
 
-    let result = handlers::handle_job(config, &job).await;
+    let mut jobs = vec![job];
+    if jobs[0].kind == JobKind::ExtractChunk {
+        let extra_limit = handlers::EXTRACT_EMBED_BATCH.saturating_sub(1);
+        let mut extra = claim_ready_extract_batch(config, DEFAULT_LOCK_DURATION_MS, extra_limit)?;
+        if !extra.is_empty() {
+            log::debug!(
+                "[memory::jobs] running extract batch count={}",
+                extra.len() + 1
+            );
+            jobs.append(&mut extra);
+        }
+    }
+
+    let results = if jobs.len() > 1 && jobs[0].kind == JobKind::ExtractChunk {
+        handlers::handle_extract_batch(config, &jobs).await?
+    } else {
+        let job = jobs
+            .pop()
+            .expect("worker has exactly one claimed job in non-batch path");
+        let result = handlers::handle_job(config, &job).await;
+        vec![(job, result)]
+    };
     drop(llm_permit);
 
     // A failed settle (`mark_done` / `mark_failed` / `mark_deferred` below)
@@ -223,6 +244,14 @@ pub async fn run_once(config: &Config) -> Result<bool> {
     // report) via [`is_sqlite_busy`]. On a stale settle the row's
     // `locked_until_ms` eventually elapses and `recover_stale_locks`
     // requeues it, so dropping the error here is at-most a re-run.
+    for (job, result) in results {
+        settle_job(config, &job, result)?;
+    }
+
+    Ok(true)
+}
+
+fn settle_job(config: &Config, job: &Job, result: Result<JobOutcome>) -> Result<()> {
     match result {
         Ok(JobOutcome::Done) => {
             log::debug!(
@@ -274,8 +303,7 @@ pub async fn run_once(config: &Config) -> Result<bool> {
             mark_failed_typed(config, &job, &message, typed)?;
         }
     }
-
-    Ok(true)
+    Ok(())
 }
 
 /// Classify whether an error is a transient I/O failure that should be

src/openhuman/memory_tree/health/mod.rs

@@ -234,9 +234,10 @@ pub fn classify_embed_error_str(msg: &str) -> PipelineFailure {
             401 | 403 => {
                 PipelineFailure::new(FailureCode::AuthInvalid).with_detail(truncate_detail(msg))
             }
-            402 | 429 => {
+            402 => {
                 PipelineFailure::new(FailureCode::BudgetExhausted).with_detail(truncate_detail(msg))
             }
+            429 => PipelineFailure::new(FailureCode::Transient).with_detail(truncate_detail(msg)),
             // 4xx other than the above is a hard client error retrying won't
             // fix (malformed request, model not found); fail fast but tag it
             // generically as auth_invalid's sibling — use Transient only for
@@ -573,15 +574,17 @@ mod tests {
     }
 
     #[test]
-    fn classify_budget_from_402_and_429() {
-        for status in ["402 Payment Required", "429 Too Many Requests"] {
-            let f = classify_embed_error_str(&format!("Embedding API error ({status}): nope"));
-            assert_eq!(
-                f.code,
-                FailureCode::BudgetExhausted,
-                "status {status} should map to budget_exhausted"
-            );
-        }
+    fn classify_budget_from_402() {
+        let f = classify_embed_error_str("Embedding API error (402 Payment Required): nope");
+        assert_eq!(f.code, FailureCode::BudgetExhausted);
+        assert!(f.is_unrecoverable());
+    }
+
+    #[test]
+    fn classify_429_rate_limit_as_transient() {
+        let f = classify_embed_error_str("Embedding API error (429 Too Many Requests): nope");
+        assert_eq!(f.code, FailureCode::Transient);
+        assert!(!f.is_unrecoverable());
     }
 
     #[test]

tests/embeddings_ollama_raw_coverage_e2e.rs

@@ -564,7 +564,7 @@ async fn embedding_catalog_factory_retry_noop_and_cloud_empty_paths_are_reachabl
     assert_eq!(parse_retry_after_ms(Some(" 5 ")), Some(5_000));
     assert_eq!(
         parse_retry_after_ms(Some("Wed, 21 Oct 2015 07:28:00 GMT")),
-        None
+        Some(0)
     );
     assert_eq!(parse_retry_after_ms(Some("99999")), Some(MAX_BACKOFF_MS));
     assert_eq!(backoff_ms_for_attempt(2, Some("1")), 1_000);