Feat/2358 clear reembed skipped (#2443)

Co-authored-by: Vladimir Yastreboff <vlad@Vladimirs-MacBook-Air.local>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
MrMrVlad
2026-05-23 01:05:28 -07:00
committed by GitHub
co-authored by Vladimir Yastreboff Cursor
parent 2f913a248c
commit f62bc968a6
5 changed files with 355 additions and 82 deletions
+110 -36
View File
@@ -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<JobOutcome> {
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<JobOutcom
// unembeddable row is attempted at most ONCE per signature instead of
// re-selected on every batch forever (the original bug: 16 orphans
// generating ~128k warns across ~8k defers, observed in the wild).
// The mark itself is best-effort — if its own SQLite write fails the
// row will be retried on a later batch, which is the desired fallback.
// Tombstone writes are best-effort: failures are logged so the row can
// be retried on a later batch instead of spinning forever.
let embedder =
build_embedder_from_config(config).context("build embedder in reembed_backfill")?;
let mut chunk_vecs: Vec<(String, Vec<f32>)> = Vec::new();
@@ -672,18 +702,13 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result<JobOutcom
log::warn!(
"[memory_tree::jobs] reembed_backfill: chunk {id} embed wrong dim, skipping (sig={active_sig})"
);
let _ = chunk_store::mark_chunk_reembed_skipped(
config,
id,
&active_sig,
"embed wrong dim",
);
try_mark_chunk_reembed_skipped(config, id, &active_sig, "embed wrong dim");
}
Err(e) => {
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<JobOutcom
log::warn!(
"[memory_tree::jobs] reembed_backfill: chunk {id} body read failed: {e}; skipping (sig={active_sig})"
);
let _ = chunk_store::mark_chunk_reembed_skipped(
try_mark_chunk_reembed_skipped(
config,
id,
&active_sig,
@@ -713,18 +738,13 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result<JobOutcom
log::warn!(
"[memory_tree::jobs] reembed_backfill: summary {id} embed wrong dim, skipping (sig={active_sig})"
);
let _ = summary_store::mark_summary_reembed_skipped(
config,
id,
&active_sig,
"embed wrong dim",
);
try_mark_summary_reembed_skipped(config, id, &active_sig, "embed wrong dim");
}
Err(e) => {
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<JobOutcom
log::warn!(
"[memory_tree::jobs] reembed_backfill: summary {id} body read failed: {e}; skipping (sig={active_sig})"
);
let _ = summary_store::mark_summary_reembed_skipped(
try_mark_summary_reembed_skipped(
config,
id,
&active_sig,
@@ -1350,24 +1370,8 @@ mod tests {
// (3) Migration probe in `ensure_reembed_backfill` must agree the
// space is covered, otherwise the chain re-arms on every config
// save and we're back to the original infinite-loop bug.
let probe_uncovered: bool = with_connection(&cfg, |conn| {
Ok(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))",
params![sig],
|r| r.get(0),
)?)
let probe_uncovered = with_connection(&cfg, |conn| {
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
})
.unwrap();
assert!(
@@ -1376,6 +1380,76 @@ mod tests {
);
}
/// #2358: clearing a tombstone re-opens the row for the backfill worklist.
#[tokio::test]
async fn clear_chunk_reembed_skipped_reopens_worklist() {
use crate::openhuman::memory::tree::store::{
clear_chunk_reembed_skipped, get_chunk_content_path, mark_chunk_reembed_skipped,
tree_active_signature, upsert_chunks, upsert_staged_chunks_tx,
};
use crate::openhuman::memory::tree::types::{
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
};
let (_tmp, cfg) = test_config();
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let chunk = Chunk {
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "clear-tombstone-seed"),
content: "memory content for clear tombstone test".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: Some(SourceRef::new("slack://x")),
},
token_count: 12,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
let staged_rel = get_chunk_content_path(&cfg, &chunk.id)
.unwrap()
.expect("staged body path");
std::fs::remove_file(content_root.join(&staged_rel)).unwrap();
let sig = tree_active_signature(&cfg);
mark_chunk_reembed_skipped(&cfg, &chunk.id, &sig, "orphan").unwrap();
let covered_before_clear = with_connection(&cfg, |conn| {
Ok(!chunk_store::has_uncovered_reembed_work(conn, &sig)?)
})
.unwrap();
assert!(
covered_before_clear,
"tombstone must hide orphan from uncovered probe"
);
clear_chunk_reembed_skipped(&cfg, &chunk.id, &sig).unwrap();
let uncovered_after_clear = with_connection(&cfg, |conn| {
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
})
.unwrap();
assert!(
uncovered_after_clear,
"clearing tombstone must re-include chunk in worklist probe"
);
}
/// #1574 §4: `ensure_reembed_backfill` (the switch-path trigger) enqueues
/// exactly one chain when there is uncovered work, is idempotent on
/// re-call (per-signature dedupe), and enqueues nothing for an
+1 -23
View File
@@ -73,29 +73,7 @@ pub fn backfill_in_progress() -> 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) => {
+103 -23
View File
@@ -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<bool> {
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<usize> {
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
+105
View File
@@ -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"
);
}
@@ -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(())
})
}