Merge pull request 'fix: a backfill reads text when it embeds, not when it scanned (#4264)' (#176) from dev into main

src/scribe/services/embeddings.py

@@ -18,7 +18,7 @@ from collections.abc import Sequence
 
 from typing import TYPE_CHECKING
 
-from sqlalchemy import delete, or_, select
+from sqlalchemy import delete, func, or_, select
 
 from scribe.models import async_session
 from scribe.models.embedding import NoteEmbedding, RuleEmbedding
@@ -908,15 +908,135 @@ async def semantic_search_notes(
     return final
 
 
-async def backfill_note_embeddings() -> None:
+# --- backfill correctness: scan ids, read text at embed time (#4262) --------
-    """(Re-)embed every note that is missing vectors OR whose stored vectors
+#
-    were produced by an older chunker.
+# A backfill used to SELECT the text alongside the id and then iterate that
+# snapshot, sleeping between records. On a corpus of any size the run takes
+# minutes, and anything edited inside that window was re-embedded by its own
+# write path and then OVERWRITTEN with the pre-edit text the scan had captured
+# — stamped at the current version, so the next boot considered it current and
+# never repaired it. A record left findable by what it used to say, silently,
+# until someone happened to edit it again.
+#
+# The two functions below are the fix, and they are separate on purpose: one
+# stops it happening, the other repairs what already happened.
 
-    Runs as a background task at startup. Version-awareness is what makes a
+
-    document-shape change deployable: migration 0077 cleared the table once,
+async def _current_row(columns, id_column, row_id: int):
-    and every later CHUNKER_VERSION bump re-embeds the stale notes here — a
+    """A record's text AS IT IS NOW, read at the moment it is embedded.
-    version comparison instead of another wipe. Adds a small sleep between
+
-    notes so a large backfill doesn't peg CPU.
+    One indexed primary-key read, negligible beside the forward pass that
+    follows it — and it turns a backfill from "replay a snapshot" into "repair
+    to current truth". The residual race shrinks from the length of the whole
+    run to the width of one record's processing, which `_claim_parent_row` and
+    the single-transaction replacement then narrow further.
+
+    None when the row is gone: deleted between the scan and here, which is an
+    ordinary outcome of a long run and not an error.
+    """
+    try:
+        async with async_session() as session:
+            return (
+                await session.execute(select(*columns).where(id_column == row_id))
+            ).first()
+    except Exception:
+        logger.warning("Backfill could not re-read record %s", row_id, exc_info=True)
+        return None
+
+
+async def _vectors_older_than_their_record(
+    session, emb_id_column, emb_updated_column, rec_id_column, rec_updated_column
+) -> set[int]:
+    """Ids whose text changed AFTER their vectors were written.
+
+    The repair half, and a standing guard worth more than the bug that prompted
+    it. A version comparison alone asks "was this embedded by the current
+    chunker" and cannot see a record that was embedded by it from stale text —
+    nor one whose `embed_note` never ran at all, because the write happened
+    with no event loop, or the refresh was swallowed, or the process died
+    between the commit and the index.
+
+    Every one of those looks identical from the version column and identical
+    from the outside: a record that answers to what it used to say. This is the
+    only signal that separates them, and it costs one grouped read at startup.
+
+    A guard that refuses to write a bad value does not undo the bad values
+    already stored (#4202) — the rows are the thing that has to change, and
+    this is what changes them.
+    """
+    newest = (
+        select(
+            emb_id_column.label("rid"),
+            func.max(emb_updated_column).label("embedded_at"),
+        )
+        .group_by(emb_id_column)
+        .subquery()
+    )
+    rows = await session.execute(
+        select(rec_id_column)
+        .join(newest, newest.c.rid == rec_id_column)
+        .where(rec_updated_column > newest.c.embedded_at)
+    )
+    return {int(row[0]) for row in rows.fetchall()}
+
+
+async def _tasks_logged_since_embedding(session) -> set[int]:
+    """Task ids whose newest work log is newer than their newest vector.
+
+    The note corpus needs this on top of the plain timestamp comparison, and
+    the reason is a consequence of #4251 that is easy to miss: a task's
+    embedded document carries its work logs, but a log lives in its own table,
+    so writing one does NOT move `notes.updated_at`. A task re-embedded from
+    stale text during a backfill therefore looks current by every other
+    signal — the version is right and the note's own timestamp is older than
+    the vectors — while the document those vectors encode is out of date.
+
+    These are the exact records the backfill race was most likely to hit,
+    because a session logging work is what "forward work during a long
+    backfill" MEANS.
+    """
+    newest_vector = (
+        select(
+            NoteEmbedding.note_id.label("rid"),
+            func.max(NoteEmbedding.updated_at).label("embedded_at"),
+        )
+        .group_by(NoteEmbedding.note_id)
+        .subquery()
+    )
+    newest_log = (
+        select(
+            TaskLog.task_id.label("tid"),
+            func.max(TaskLog.updated_at).label("logged_at"),
+        )
+        .group_by(TaskLog.task_id)
+        .subquery()
+    )
+    rows = await session.execute(
+        select(newest_log.c.tid)
+        .join(newest_vector, newest_vector.c.rid == newest_log.c.tid)
+        .where(newest_log.c.logged_at > newest_vector.c.embedded_at)
+    )
+    return {int(row[0]) for row in rows.fetchall()}
+
+
+async def backfill_note_embeddings() -> None:
+    """(Re-)embed every note whose vectors are missing, outdated or WRONG.
+
+    Three conditions, and the third is the one a version number cannot see:
+
+    1. no vectors at all;
+    2. vectors from an older chunker — version-awareness is what makes a
+       document-shape change deployable. Migration 0077 cleared the table
+       once; every later CHUNKER_VERSION bump re-embeds here instead;
+    3. vectors OLDER THAN THE TEXT THEY CLAIM TO ENCODE. A record embedded by
+       the current chunker from text that has since changed is current by
+       every other signal and answers to what it used to say (#4262).
+
+    Runs as a background task at startup, after the serving flag (#4181), with
+    a small sleep between notes so a large backfill doesn't peg CPU. It is
+    resumable by construction: progress IS the version stamp on each record's
+    rows, so a restart re-queries and skips what finished — there is no
+    checkpoint to lose.
     """
     try:
         async with async_session() as session:
@@ -930,26 +1050,41 @@ async def backfill_note_embeddings() -> None:
                     )
                 ).fetchall()
             }
-            result = await session.execute(
+            # IDS ONLY. The text is read per note below, at the moment it is
-                select(Note.id, Note.user_id, Note.title, Note.body)
+            # embedded — a scan that carried the text would spend the whole run
-            )
+            # holding a snapshot and overwrite anything edited inside it.
-            notes_to_embed = [
+            all_ids = [
-                row for row in result.fetchall() if row[0] not in current
+                row[0]
+                for row in (await session.execute(select(Note.id))).fetchall()
             ]
+            stale = await _vectors_older_than_their_record(
+                session, NoteEmbedding.note_id, NoteEmbedding.updated_at,
+                Note.id, Note.updated_at,
+            )
+            # A work log is part of a task's document but lives in its own
+            # table, so writing one leaves `notes.updated_at` untouched and the
+            # comparison above blind to it.
+            stale |= await _tasks_logged_since_embedding(session)
     except Exception:
         logger.warning("Embedding backfill: failed to query notes", exc_info=True)
         return
 
+    notes_to_embed = [i for i in all_ids if i not in current or i in stale]
     if not notes_to_embed:
         logger.info("Embedding backfill: all notes current at chunker v%d", CHUNKER_VERSION)
         return
 
     logger.info(
-        "Embedding backfill: embedding %d notes at chunker v%d",
+        "Embedding backfill: embedding %d note(s) at chunker v%d (%d with vectors "
-        len(notes_to_embed), CHUNKER_VERSION,
+        "older than the record they describe)",
+        len(notes_to_embed), CHUNKER_VERSION, len(stale),
     )
     success = 0
-    for note_id, user_id, title, body in notes_to_embed:
+    for note_id in notes_to_embed:
+        row = await _current_row((Note.user_id, Note.title, Note.body), Note.id, note_id)
+        if row is None:
+            continue          # deleted between the scan and here
+        user_id, title, body = row
         if not chunk_document(title, body):
             continue
         await upsert_note_embedding(note_id, user_id, title, body)
@@ -1280,10 +1415,18 @@ async def backfill_rule_embeddings() -> None:
             current = select(RuleEmbedding.rule_id).where(
                 RuleEmbedding.chunker_version == CHUNKER_VERSION
             )
-            stale = (await session.execute(
+            # IDS ONLY — the text is re-read per rule below (#4262).
-                select(Rule.id, Rule.title, Rule.statement, Rule.when_to_apply)
+            by_version = {
+                row[0] for row in (await session.execute(
+                    select(Rule.id)
                     .where(Rule.deleted_at.is_(None), Rule.id.notin_(current))
-            )).all()
+                )).fetchall()
+            }
+            by_time = await _vectors_older_than_their_record(
+                session, RuleEmbedding.rule_id, RuleEmbedding.updated_at,
+                Rule.id, Rule.updated_at,
+            )
+            stale = sorted(by_version | by_time)
     except Exception:
         logger.warning("Rule embedding backfill: failed to query rules", exc_info=True)
         return
@@ -1291,9 +1434,17 @@ async def backfill_rule_embeddings() -> None:
     if not stale:
         logger.info("Rule embedding backfill: all rules current at chunker v%d", CHUNKER_VERSION)
         return
-    logger.info("Rule embedding backfill: embedding %d rule(s)", len(stale))
+    logger.info(
-    for rule_id, title, statement, when_to_apply in stale:
+        "Rule embedding backfill: embedding %d rule(s) (%d with stale vectors)",
-        await upsert_rule_embedding(rule_id, title, statement, when_to_apply)
+        len(stale), len(by_time),
+    )
+    for rule_id in stale:
+        row = await _current_row(
+            (Rule.title, Rule.statement, Rule.when_to_apply), Rule.id, rule_id
+        )
+        if row is None:
+            continue
+        await upsert_rule_embedding(rule_id, *row)
 
 
 # ── Milestones (milestone 415) ──────────────────────────────────────────
@@ -1641,10 +1792,18 @@ async def backfill_system_embeddings() -> None:
             current = select(SystemEmbedding.system_id).where(
                 SystemEmbedding.chunker_version == CHUNKER_VERSION
             )
-            stale = (await session.execute(
+            # IDS ONLY — the charter is re-read per System below (#4262).
-                select(System.id, System.name, System.description)
+            by_version = {
+                row[0] for row in (await session.execute(
+                    select(System.id)
                     .where(System.deleted_at.is_(None), System.id.notin_(current))
-            )).all()
+                )).fetchall()
+            }
+            by_time = await _vectors_older_than_their_record(
+                session, SystemEmbedding.system_id, SystemEmbedding.updated_at,
+                System.id, System.updated_at,
+            )
+            stale = sorted(by_version | by_time)
     except Exception:
         logger.warning("System embedding backfill: failed to query systems", exc_info=True)
         return
@@ -1652,9 +1811,17 @@ async def backfill_system_embeddings() -> None:
     if not stale:
         logger.info("System embedding backfill: all systems current at chunker v%d", CHUNKER_VERSION)
         return
-    logger.info("System embedding backfill: embedding %d system(s)", len(stale))
+    logger.info(
-    for system_id, name, description in stale:
+        "System embedding backfill: embedding %d system(s) (%d with stale vectors)",
-        await upsert_system_embedding(system_id, name, description)
+        len(stale), len(by_time),
+    )
+    for system_id in stale:
+        row = await _current_row(
+            (System.name, System.description), System.id, system_id
+        )
+        if row is None:
+            continue
+        await upsert_system_embedding(system_id, *row)
 
 
 async def backfill_milestone_embeddings() -> None:
@@ -1668,10 +1835,18 @@ async def backfill_milestone_embeddings() -> None:
             current = select(MilestoneEmbedding.milestone_id).where(
                 MilestoneEmbedding.chunker_version == CHUNKER_VERSION
             )
-            stale = (await session.execute(
+            # IDS ONLY — the plan is re-read per milestone below (#4262).
-                select(Milestone.id, Milestone.title, Milestone.description, Milestone.body)
+            by_version = {
+                row[0] for row in (await session.execute(
+                    select(Milestone.id)
                     .where(Milestone.deleted_at.is_(None), Milestone.id.notin_(current))
-            )).all()
+                )).fetchall()
+            }
+            by_time = await _vectors_older_than_their_record(
+                session, MilestoneEmbedding.milestone_id, MilestoneEmbedding.updated_at,
+                Milestone.id, Milestone.updated_at,
+            )
+            stale = sorted(by_version | by_time)
     except Exception:
         logger.warning("Milestone embedding backfill: failed to query milestones", exc_info=True)
         return
@@ -1679,6 +1854,15 @@ async def backfill_milestone_embeddings() -> None:
     if not stale:
         logger.info("Milestone embedding backfill: all milestones current at chunker v%d", CHUNKER_VERSION)
         return
-    logger.info("Milestone embedding backfill: embedding %d milestone(s)", len(stale))
+    logger.info(
-    for milestone_id, title, description, body in stale:
+        "Milestone embedding backfill: embedding %d milestone(s) (%d with stale vectors)",
-        await upsert_milestone_embedding(milestone_id, title, description, body)
+        len(stale), len(by_time),
+    )
+    for milestone_id in stale:
+        row = await _current_row(
+            (Milestone.title, Milestone.description, Milestone.body),
+            Milestone.id, milestone_id,
+        )
+        if row is None:
+            continue
+        await upsert_milestone_embedding(milestone_id, *row)

tests/test_backfill_reads_current_text.py

@@ -0,0 +1,186 @@
+"""A backfill repairs to current truth; it does not replay a snapshot (#4262).
+
+The bug these pin: every backfill used to SELECT the text alongside the id and
+then iterate that snapshot, sleeping between records. A run over a real corpus
+takes minutes, so anything edited inside that window was re-embedded by its own
+write path and then OVERWRITTEN with the pre-edit text the scan had captured —
+and stamped at the current chunker version, so the next boot considered it
+current and never repaired it. The record stayed findable by what it used to
+say, silently, until someone happened to edit it again.
+
+That window is exactly "forward work continues while a long backfill runs",
+which is the normal condition, not an unlucky one.
+"""
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from scribe.services import embeddings as emb
+
+
+def _ctx(session):
+    ctx = MagicMock()
+    ctx.__aenter__ = AsyncMock(return_value=session)
+    ctx.__aexit__ = AsyncMock(return_value=False)
+    return ctx
+
+
+def _scan_session(all_ids, current_ids, stale_ids=(), logged_ids=()):
+    """The opening scan: ids at the current version, all ids, and the two
+    staleness queries. Order matches backfill_note_embeddings."""
+    def _rows(values):
+        r = MagicMock()
+        r.fetchall.return_value = [(v,) for v in values]
+        return r
+
+    session = MagicMock()
+    session.execute = AsyncMock(side_effect=[
+        _rows(current_ids),     # note_ids already at CHUNKER_VERSION
+        _rows(all_ids),         # every note id
+        _rows(stale_ids),       # vectors older than the record
+        _rows(logged_ids),      # tasks logged since their vectors
+    ])
+    return session
+
+
+# --- the fix: text is read when the record is embedded, not when scanned -----
+
+
+@pytest.mark.asyncio
+async def test_the_backfill_embeds_the_text_as_it_is_now_not_as_it_was_scanned():
+    """THE regression. The scan sees the old body; by the time the loop reaches
+    the note, someone has edited it. The vectors must encode the edit."""
+    scan = _scan_session(all_ids=[1], current_ids=[])
+    edited = ("the body AFTER the edit",)
+
+    with (
+        patch.object(emb, "async_session", return_value=_ctx(scan)),
+        patch.object(emb, "_current_row",
+                     AsyncMock(return_value=(42, "T", *edited))),
+        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
+        patch.object(emb.asyncio, "sleep", AsyncMock()),
+    ):
+        await emb.backfill_note_embeddings()
+
+    upsert.assert_awaited_once()
+    assert upsert.await_args.args[3] == "the body AFTER the edit"
+
+
+@pytest.mark.asyncio
+async def test_a_record_deleted_between_the_scan_and_the_loop_is_skipped():
+    """An ordinary outcome of a long run, not an error — and embedding a row
+    that no longer exists is either a FK violation or a resurrected chunk."""
+    scan = _scan_session(all_ids=[1, 2], current_ids=[])
+
+    with (
+        patch.object(emb, "async_session", return_value=_ctx(scan)),
+        patch.object(emb, "_current_row",
+                     AsyncMock(side_effect=[None, (42, "T", "body")])),
+        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
+        patch.object(emb.asyncio, "sleep", AsyncMock()),
+    ):
+        await emb.backfill_note_embeddings()
+
+    assert upsert.await_count == 1
+    assert upsert.await_args.args[0] == 2
+
+
+# --- the repair: stale vectors are re-embedded even at the current version ---
+
+
+@pytest.mark.asyncio
+async def test_a_record_whose_text_outran_its_vectors_is_re_embedded():
+    """The repair half (#4202 — a guard that refuses to write a bad value does
+    not undo the bad values already stored). Note 5 IS at the current version,
+    so the version comparison alone would skip it forever; its text is newer
+    than its vectors, which is the only signal that separates "embedded from
+    stale text" from "embedded correctly"."""
+    scan = _scan_session(all_ids=[5], current_ids=[5], stale_ids=[5])
+
+    with (
+        patch.object(emb, "async_session", return_value=_ctx(scan)),
+        patch.object(emb, "_current_row", AsyncMock(return_value=(42, "T", "b"))),
+        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
+        patch.object(emb.asyncio, "sleep", AsyncMock()),
+    ):
+        await emb.backfill_note_embeddings()
+
+    upsert.assert_awaited_once()
+    assert upsert.await_args.args[0] == 5
+
+
+@pytest.mark.asyncio
+async def test_a_task_logged_since_its_vectors_is_re_embedded():
+    """A work log is part of a task's document but lives in its own table, so
+    writing one never moves `notes.updated_at` — the timestamp comparison is
+    blind to it. These are the records the race was MOST likely to hit, because
+    a session writing work logs is what "forward work" means."""
+    scan = _scan_session(all_ids=[9], current_ids=[9], stale_ids=[], logged_ids=[9])
+
+    with (
+        patch.object(emb, "async_session", return_value=_ctx(scan)),
+        patch.object(emb, "_current_row", AsyncMock(return_value=(42, "T", "b"))),
+        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
+        patch.object(emb.asyncio, "sleep", AsyncMock()),
+    ):
+        await emb.backfill_note_embeddings()
+
+    upsert.assert_awaited_once()
+    assert upsert.await_args.args[0] == 9
+
+
+@pytest.mark.asyncio
+async def test_a_corpus_that_is_genuinely_current_embeds_nothing():
+    """The guard must be able to stay quiet, or it is not a guard (#167) — and
+    a backfill that re-embeds everything every boot is a different bug."""
+    scan = _scan_session(all_ids=[1, 2, 3], current_ids=[1, 2, 3])
+
+    with (
+        patch.object(emb, "async_session", return_value=_ctx(scan)),
+        patch.object(emb, "_current_row", AsyncMock()) as fresh,
+        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
+        patch.object(emb.asyncio, "sleep", AsyncMock()),
+    ):
+        await emb.backfill_note_embeddings()
+
+    upsert.assert_not_awaited()
+    fresh.assert_not_awaited()
+
+
+# --- the scan must not carry text at all -------------------------------------
+
+
+def test_no_backfill_selects_a_records_text_in_its_opening_scan():
+    """A structural guard, because this is the shape of the bug rather than one
+    instance of it: all four backfills had it, and the fourth was written by
+    copying the third. A scan that carries text is a snapshot that goes stale
+    over the length of the run."""
+    import ast
+    import inspect
+
+    source = inspect.getsource(emb)
+    tree = ast.parse(source)
+    text_columns = {
+        "title", "body", "statement", "when_to_apply",
+        "description", "name",
+    }
+    offenders = []
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.AsyncFunctionDef)
+                and node.name.startswith("backfill_")):
+            continue
+        # Only the OPENING scan is at issue — it is the code inside the
+        # `async with async_session()` block, before the embed loop.
+        for inner in ast.walk(node):
+            if not (isinstance(inner, ast.Call)
+                    and getattr(inner.func, "id", None) == "select"):
+                continue
+            for arg in inner.args:
+                if (isinstance(arg, ast.Attribute)
+                        and arg.attr in text_columns):
+                    offenders.append(f"{node.name}: select(...{arg.attr})")
+
+    assert not offenders, (
+        "a backfill scan is carrying record text, which goes stale over the "
+        f"length of the run: {offenders}"
+    )

tests/test_chunking.py

@@ -287,18 +287,26 @@ async def test_backfill_reembeds_notes_with_a_stale_chunker_version():
 
     from scribe.services import embeddings as emb
 
-    current_rows = MagicMock()
+    def _ids(*values):
-    current_rows.fetchall.return_value = [(1,)]         # note 1 is current
+        r = MagicMock()
-    note_rows = MagicMock()
+        r.fetchall.return_value = [(v,) for v in values]
-    note_rows.fetchall.return_value = [
+        return r
-        (1, 42, "current", "body"),
+
-        (2, 42, "stale-version", "body"),
-    ]
     session, ctx = _session_ctx()
-    session.execute = AsyncMock(side_effect=[current_rows, note_rows])
+    # The opening scan, in order: ids already at the current version, every
+    # note id, vectors older than their note, tasks logged since their vectors.
+    # IDS ONLY — the text is re-read per note at embed time (#4262).
+    session.execute = AsyncMock(side_effect=[
+        _ids(1),            # note 1 is current
+        _ids(1, 2),         # the corpus
+        _ids(),             # nothing stale by timestamp
+        _ids(),             # no task logged since its vectors
+    ])
 
     with (
         patch.object(emb, "async_session", return_value=ctx),
+        patch.object(emb, "_current_row",
+                     AsyncMock(return_value=(42, "stale-version", "body"))),
         patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
         patch.object(emb.asyncio, "sleep", AsyncMock()),
     ):