diff --git a/src/openhuman/embeddings/retry_after.rs b/src/openhuman/embeddings/retry_after.rs index b799f77a6..db0bf8096 100644 --- a/src/openhuman/embeddings/retry_after.rs +++ b/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 { + parse_retry_after_ms_at(header_value, Utc::now()) +} + +fn parse_retry_after_ms_at(header_value: Option<&str>, now: DateTime) -> Option { 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::() { + 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) ); } diff --git a/src/openhuman/memory_queue/handlers/mod.rs b/src/openhuman/memory_queue/handlers/mod.rs index e39f6dbc4..a091d0505 100644 --- a/src/openhuman/memory_queue/handlers/mod.rs +++ b/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 } async fn handle_extract(config: &Config, job: &Job) -> Result { + 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)>> { + 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>, + embedding_error: Option, + embedding_retry_needed: bool, +} + +async fn prepare_extract(config: &Config, job: &Job) -> Result> { 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 { "[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 { 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> = 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 = 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 { + 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 { 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 { 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}")); } diff --git a/src/openhuman/memory_queue/store.rs b/src/openhuman/memory_queue/store.rs index eabafa09b..cba8eeb1d 100644 --- a/src/openhuman/memory_queue/store.rs +++ b/src/openhuman/memory_queue/store.rs @@ -151,6 +151,59 @@ pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result> }) } +/// 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> { + 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::>>() + .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 diff --git a/src/openhuman/memory_queue/types.rs b/src/openhuman/memory_queue/types.rs index dabcd29ca..ba82c17dc 100644 --- a/src/openhuman/memory_queue/types.rs +++ b/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"); } diff --git a/src/openhuman/memory_queue/worker.rs b/src/openhuman/memory_queue/worker.rs index 95315e516..642e8eac7 100644 --- a/src/openhuman/memory_queue/worker.rs +++ b/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 { 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 { // 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) -> Result<()> { match result { Ok(JobOutcome::Done) => { log::debug!( @@ -274,8 +303,7 @@ pub async fn run_once(config: &Config) -> Result { mark_failed_typed(config, &job, &message, typed)?; } } - - Ok(true) + Ok(()) } /// Classify whether an error is a transient I/O failure that should be diff --git a/src/openhuman/memory_tree/health/mod.rs b/src/openhuman/memory_tree/health/mod.rs index 8742897ad..4fa1d0dff 100644 --- a/src/openhuman/memory_tree/health/mod.rs +++ b/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] diff --git a/tests/embeddings_ollama_raw_coverage_e2e.rs b/tests/embeddings_ollama_raw_coverage_e2e.rs index 7c62ad0ab..442535735 100644 --- a/tests/embeddings_ollama_raw_coverage_e2e.rs +++ b/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);