diff --git a/src/openhuman/memory/tree/jobs/handlers/mod.rs b/src/openhuman/memory/tree/jobs/handlers/mod.rs index 23b3ba59d..b69aaa8c3 100644 --- a/src/openhuman/memory/tree/jobs/handlers/mod.rs +++ b/src/openhuman/memory/tree/jobs/handlers/mod.rs @@ -567,6 +567,36 @@ const REEMBED_BACKFILL_REVISIT_MS: i64 = 750; /// /// Per-row read/embed failures are logged and skipped, never fail the /// chain — one unreadable row must not strand the rest of memory. +fn try_mark_chunk_reembed_skipped( + config: &Config, + chunk_id: &str, + model_signature: &str, + reason: &str, +) { + if let Err(e) = + chunk_store::mark_chunk_reembed_skipped(config, chunk_id, model_signature, reason) + { + log::warn!( + "[memory_tree::jobs] reembed_backfill: failed to persist chunk tombstone chunk_id={chunk_id} sig={model_signature}: {e}" + ); + } +} + +fn try_mark_summary_reembed_skipped( + config: &Config, + summary_id: &str, + model_signature: &str, + reason: &str, +) { + if let Err(e) = + summary_store::mark_summary_reembed_skipped(config, summary_id, model_signature, reason) + { + log::warn!( + "[memory_tree::jobs] reembed_backfill: failed to persist summary tombstone summary_id={summary_id} sig={model_signature}: {e}" + ); + } +} + async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result { let payload: ReembedBackfillPayload = serde_json::from_str(&job.payload_json).context("parse ReembedBackfill payload")?; @@ -659,8 +689,8 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result)> = Vec::new(); @@ -672,18 +702,13 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result { log::warn!( "[memory_tree::jobs] reembed_backfill: chunk {id} embed failed: {e}; skipping (sig={active_sig})" ); - let _ = chunk_store::mark_chunk_reembed_skipped( + try_mark_chunk_reembed_skipped( config, id, &active_sig, @@ -695,7 +720,7 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result Result { log::warn!( "[memory_tree::jobs] reembed_backfill: summary {id} embed failed: {e}; skipping (sig={active_sig})" ); - let _ = summary_store::mark_summary_reembed_skipped( + try_mark_summary_reembed_skipped( config, id, &active_sig, @@ -736,7 +756,7 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result bool { pub fn ensure_reembed_backfill(config: &crate::openhuman::config::Config) { let sig = crate::openhuman::memory::tree::store::tree_active_signature(config); let result = crate::openhuman::memory::tree::store::with_connection(config, |conn| { - let has_uncovered: bool = conn.query_row( - // The `NOT EXISTS … reembed_skipped` clauses match the worklist in - // `handle_reembed_backfill`: terminally-failed rows are sentinel- - // marked there and must NOT count as "uncovered" here, otherwise - // this probe keeps reporting "uncovered" → keeps re-enqueueing the - // backfill chain → infinite re-arming (#1574 §6 runaway-loop fix). - "SELECT EXISTS( - SELECT 1 FROM mem_tree_chunks c - WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e - WHERE e.chunk_id = c.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk - WHERE sk.chunk_id = c.id AND sk.model_signature = ?1)) - OR EXISTS( - SELECT 1 FROM mem_tree_summaries s - WHERE s.deleted = 0 - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e - WHERE e.summary_id = s.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk - WHERE sk.summary_id = s.id AND sk.model_signature = ?1))", - rusqlite::params![sig], - |r| r.get(0), - )?; - Ok(has_uncovered) + Ok(crate::openhuman::memory::tree::store::has_uncovered_reembed_work(conn, &sig)?) }); match result { Ok(true) => { diff --git a/src/openhuman/memory/tree/store.rs b/src/openhuman/memory/tree/store.rs index 3a18612ab..3a672a8a4 100644 --- a/src/openhuman/memory/tree/store.rs +++ b/src/openhuman/memory/tree/store.rs @@ -1348,29 +1348,7 @@ fn migrate_legacy_embeddings_to_sidecar(conn: &Connection, config: &Config) -> R // no-op job on every DB open — which would otherwise pollute the jobs // table for unrelated callers/tests. Enqueued atomically with the // migration; dedupe key = signature, so exactly one chain per space. - let has_uncovered: bool = tx.query_row( - // The `NOT EXISTS … reembed_skipped` clauses match the worklist in - // `handle_reembed_backfill`: terminally-failed rows (body missing, - // embed wrong dim / err) are sentinel-marked there and must NOT count - // as "uncovered" here, otherwise this migration probe keeps reporting - // "uncovered" → keeps enqueueing the backfill chain on every DB open → - // infinite re-arming (#1574 §6 runaway-loop fix). - "SELECT EXISTS( - SELECT 1 FROM mem_tree_chunks c - WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e - WHERE e.chunk_id = c.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk - WHERE sk.chunk_id = c.id AND sk.model_signature = ?1)) - OR EXISTS( - SELECT 1 FROM mem_tree_summaries s - WHERE s.deleted = 0 - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e - WHERE e.summary_id = s.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk - WHERE sk.summary_id = s.id AND sk.model_signature = ?1))", - rusqlite::params![sig], - |r| r.get(0), - )?; + let has_uncovered = has_uncovered_reembed_work(&*tx, &sig)?; if has_uncovered { let backfill_job = crate::openhuman::memory::tree::jobs::types::NewJob::reembed_backfill( &crate::openhuman::memory::tree::jobs::types::ReembedBackfillPayload { @@ -1632,6 +1610,34 @@ pub fn set_chunk_embedding_for_signature( }) } +/// `true` when at least one chunk or summary still needs an embedding at +/// `model_signature` and is not tombstoned as terminally unembeddable. +/// +/// Shared by `ensure_reembed_backfill`, the §7 migration enqueue probe, and +/// tests so the worklist and coverage probes cannot drift (#2358). +pub(crate) fn has_uncovered_reembed_work( + conn: &Connection, + model_signature: &str, +) -> rusqlite::Result { + conn.query_row( + "SELECT EXISTS( + SELECT 1 FROM mem_tree_chunks c + WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e + WHERE e.chunk_id = c.id AND e.model_signature = ?1) + AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk + WHERE sk.chunk_id = c.id AND sk.model_signature = ?1)) + OR EXISTS( + SELECT 1 FROM mem_tree_summaries s + WHERE s.deleted = 0 + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e + WHERE e.summary_id = s.id AND e.model_signature = ?1) + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk + WHERE sk.summary_id = s.id AND sk.model_signature = ?1))", + rusqlite::params![model_signature], + |r| r.get(0), + ) +} + /// Persistently record that `(chunk_id, signature)` cannot be re-embedded. /// /// Called by `handle_reembed_backfill` when the per-chunk body file is @@ -1647,6 +1653,8 @@ pub fn mark_chunk_reembed_skipped( model_signature: &str, reason: &str, ) -> Result<()> { + let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?; + let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; with_connection(config, |conn| { let now_ms = Utc::now().timestamp_millis(); conn.execute( @@ -1658,10 +1666,82 @@ pub fn mark_chunk_reembed_skipped( skipped_at_ms = excluded.skipped_at_ms", rusqlite::params![chunk_id, model_signature, reason, now_ms], )?; + log::debug!( + "[memory_tree::store] mark_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature} reason={reason}" + ); Ok(()) }) } +/// Remove a single chunk tombstone so re-embed backfill can retry the row. +/// +/// Idempotent: deleting a missing `(chunk_id, model_signature)` pair is a +/// no-op. Intended for operator recovery after environmental failures (moved +/// workspace, restored body files, fixed embedder config) — see #2358. +pub fn clear_chunk_reembed_skipped( + config: &Config, + chunk_id: &str, + model_signature: &str, +) -> Result<()> { + let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?; + let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; + with_connection(config, |conn| { + conn.execute( + "DELETE FROM mem_tree_chunk_reembed_skipped + WHERE chunk_id = ?1 AND model_signature = ?2", + rusqlite::params![chunk_id, model_signature], + )?; + log::debug!( + "[memory_tree::store] clear_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature}" + ); + Ok(()) + }) +} + +/// Clear all chunk and summary tombstones for a model signature. +/// +/// Returns the total number of rows removed across both tombstone tables. +/// Idempotent when no tombstones exist for the signature. +pub fn clear_reembed_skipped_for_signature( + config: &Config, + model_signature: &str, +) -> Result { + let model_signature = validate_reembed_skip_key("model_signature", model_signature)?; + with_connection(config, |conn| { + let chunk_deleted = conn.execute( + "DELETE FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1", + rusqlite::params![model_signature], + )?; + let summary_deleted = conn.execute( + "DELETE FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1", + rusqlite::params![model_signature], + )?; + log::debug!( + "[memory_tree::store] clear_reembed_skipped_for_signature sig={model_signature} chunk_rows={chunk_deleted} summary_rows={summary_deleted}" + ); + Ok(chunk_deleted + summary_deleted) + }) +} + +/// Bounds attacker-controlled ids/signatures passed to reembed-skipped admin +/// helpers without affecting legitimate rows (typical ids are well under 512 +/// chars). Rejects NUL bytes so SQLite bindings cannot be truncated. +const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048; + +pub(crate) fn validate_reembed_skip_key<'a>(label: &str, value: &'a str) -> Result<&'a str> { + let trimmed = value.trim(); + if trimmed.is_empty() { + anyhow::bail!("{label} must be non-empty"); + } + if trimmed.len() > REEMBED_SKIP_KEY_MAX_LEN { + anyhow::bail!("{label} exceeds maximum length ({REEMBED_SKIP_KEY_MAX_LEN})"); + } + if trimmed.as_bytes().contains(&0) { + anyhow::bail!("{label} must not contain NUL bytes"); + } + Ok(trimmed) +} + /// Transaction-scoped variant of [`set_chunk_embedding_for_signature`]. /// /// For callers that already hold a `Transaction` (e.g. the chunk-admission diff --git a/src/openhuman/memory/tree/store_tests.rs b/src/openhuman/memory/tree/store_tests.rs index 0f22753f2..60430b6f5 100644 --- a/src/openhuman/memory/tree/store_tests.rs +++ b/src/openhuman/memory/tree/store_tests.rs @@ -663,3 +663,108 @@ fn existing_wal_db_migrates_to_truncate() { "-wal must be gone after WAL→TRUNCATE migration" ); } + +#[test] +fn clear_chunk_reembed_skipped_is_idempotent() { + let (_tmp, cfg) = test_config(); + let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + let sig = tree_active_signature(&cfg); + mark_chunk_reembed_skipped(&cfg, &c.id, &sig, "test orphan").unwrap(); + clear_chunk_reembed_skipped(&cfg, &c.id, &sig).unwrap(); + clear_chunk_reembed_skipped(&cfg, &c.id, &sig).unwrap(); + let count: i64 = with_connection(&cfg, |conn| { + Ok(conn.query_row( + "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped + WHERE chunk_id = ?1 AND model_signature = ?2", + params![c.id, sig], + |r| r.get(0), + )?) + }) + .unwrap(); + assert_eq!(count, 0); +} + +#[test] +fn clear_reembed_skipped_for_signature_removes_all_tombstones_for_sig() { + let (_tmp, cfg) = test_config(); + let c1 = sample_chunk("slack:#a", 0, 1_700_000_000_000); + let c2 = sample_chunk("slack:#b", 1, 1_700_000_000_001); + upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap(); + let sig = tree_active_signature(&cfg); + let other_sig = "provider=other;model=x;dims=8"; + mark_chunk_reembed_skipped(&cfg, &c1.id, &sig, "r1").unwrap(); + mark_chunk_reembed_skipped(&cfg, &c2.id, &sig, "r2").unwrap(); + mark_chunk_reembed_skipped(&cfg, &c1.id, other_sig, "other").unwrap(); + let summary_id = "summary-bulk-clear-test"; + with_connection(&cfg, |conn| { + conn.execute( + "INSERT OR IGNORE INTO mem_tree_trees (id, kind, scope, created_at_ms) + VALUES ('tree-bulk-clear', 'source', 'bulk-clear', 0)", + [], + )?; + conn.execute( + "INSERT INTO mem_tree_summaries ( + id, tree_id, tree_kind, level, child_ids_json, content, token_count, + entities_json, topics_json, time_range_start_ms, time_range_end_ms, + score, sealed_at_ms, deleted + ) VALUES (?1, 'tree-bulk-clear', 'source', 0, '[]', 'x', 1, '[]', '[]', 0, 0, 0.0, 0, 0)", + params![summary_id], + )?; + Ok(()) + }) + .unwrap(); + crate::openhuman::memory::tree::tree_source::store::mark_summary_reembed_skipped( + &cfg, + summary_id, + &sig, + "summary tombstone", + ) + .unwrap(); + + let deleted = clear_reembed_skipped_for_signature(&cfg, &sig).unwrap(); + assert_eq!(deleted, 3); + + let remaining_chunks: i64 = with_connection(&cfg, |conn| { + Ok(conn.query_row( + "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1", + params![sig], + |r| r.get(0), + )?) + }) + .unwrap(); + assert_eq!(remaining_chunks, 0); + + let remaining_summaries: i64 = with_connection(&cfg, |conn| { + Ok(conn.query_row( + "SELECT COUNT(*) FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1", + params![sig], + |r| r.get(0), + )?) + }) + .unwrap(); + assert_eq!(remaining_summaries, 0); + + let other_kept: i64 = with_connection(&cfg, |conn| { + Ok(conn.query_row( + "SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped + WHERE chunk_id = ?1 AND model_signature = ?2", + params![c1.id, other_sig], + |r| r.get(0), + )?) + }) + .unwrap(); + assert_eq!(other_kept, 1); +} + +#[test] +fn validate_reembed_skip_key_rejects_empty_and_oversized() { + assert!(validate_reembed_skip_key("chunk_id", " ").is_err()); + let huge = "a".repeat(REEMBED_SKIP_KEY_MAX_LEN + 1); + assert!(validate_reembed_skip_key("chunk_id", &huge).is_err()); + assert!(validate_reembed_skip_key("chunk_id", "ok\0bad").is_err()); + assert_eq!( + validate_reembed_skip_key("chunk_id", " trimmed ").unwrap(), + "trimmed" + ); +} diff --git a/src/openhuman/memory/tree/tree_source/store.rs b/src/openhuman/memory/tree/tree_source/store.rs index f8e6f8bf2..a48b8a13e 100644 --- a/src/openhuman/memory/tree/tree_source/store.rs +++ b/src/openhuman/memory/tree/tree_source/store.rs @@ -348,6 +348,12 @@ pub fn mark_summary_reembed_skipped( model_signature: &str, reason: &str, ) -> Result<()> { + let summary_id = + crate::openhuman::memory::tree::store::validate_reembed_skip_key("summary_id", summary_id)?; + let model_signature = crate::openhuman::memory::tree::store::validate_reembed_skip_key( + "model_signature", + model_signature, + )?; with_connection(config, |conn| { let now_ms = Utc::now().timestamp_millis(); conn.execute( @@ -359,6 +365,36 @@ pub fn mark_summary_reembed_skipped( skipped_at_ms = excluded.skipped_at_ms", params![summary_id, model_signature, reason, now_ms], )?; + log::debug!( + "[memory_tree::store] mark_summary_reembed_skipped summary_id={summary_id} sig={model_signature} reason={reason}" + ); + Ok(()) + }) +} + +/// Remove a single summary tombstone so re-embed backfill can retry the row. +/// +/// Idempotent — see [`crate::openhuman::memory::tree::store::clear_chunk_reembed_skipped`]. +pub fn clear_summary_reembed_skipped( + config: &Config, + summary_id: &str, + model_signature: &str, +) -> Result<()> { + let summary_id = + crate::openhuman::memory::tree::store::validate_reembed_skip_key("summary_id", summary_id)?; + let model_signature = crate::openhuman::memory::tree::store::validate_reembed_skip_key( + "model_signature", + model_signature, + )?; + with_connection(config, |conn| { + conn.execute( + "DELETE FROM mem_tree_summary_reembed_skipped + WHERE summary_id = ?1 AND model_signature = ?2", + params![summary_id, model_signature], + )?; + log::debug!( + "[memory_tree::store] clear_summary_reembed_skipped summary_id={summary_id} sig={model_signature}" + ); Ok(()) }) }