diff --git a/src/scribe/services/embeddings.py b/src/scribe/services/embeddings.py index eba3c5f..1361c18 100644 --- a/src/scribe/services/embeddings.py +++ b/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: - """(Re-)embed every note that is missing vectors OR whose stored vectors - were produced by an older chunker. +# --- backfill correctness: scan ids, read text at embed time (#4262) -------- +# +# 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, - and every later CHUNKER_VERSION bump re-embeds the stale notes here — a - version comparison instead of another wipe. Adds a small sleep between - notes so a large backfill doesn't peg CPU. + +async def _current_row(columns, id_column, row_id: int): + """A record's text AS IT IS NOW, read at the moment it is embedded. + + 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( - select(Note.id, Note.user_id, Note.title, Note.body) - ) - notes_to_embed = [ - row for row in result.fetchall() if row[0] not in current + # IDS ONLY. The text is read per note below, at the moment it is + # embedded — a scan that carried the text would spend the whole run + # holding a snapshot and overwrite anything edited inside it. + all_ids = [ + 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", - len(notes_to_embed), CHUNKER_VERSION, + "Embedding backfill: embedding %d note(s) at chunker v%d (%d with vectors " + "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( - select(Rule.id, Rule.title, Rule.statement, Rule.when_to_apply) - .where(Rule.deleted_at.is_(None), Rule.id.notin_(current)) - )).all() + # IDS ONLY — the text is re-read per rule below (#4262). + by_version = { + row[0] for row in (await session.execute( + select(Rule.id) + .where(Rule.deleted_at.is_(None), Rule.id.notin_(current)) + )).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)) - for rule_id, title, statement, when_to_apply in stale: - await upsert_rule_embedding(rule_id, title, statement, when_to_apply) + logger.info( + "Rule embedding backfill: embedding %d rule(s) (%d with stale vectors)", + 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( - select(System.id, System.name, System.description) - .where(System.deleted_at.is_(None), System.id.notin_(current)) - )).all() + # IDS ONLY — the charter is re-read per System below (#4262). + by_version = { + row[0] for row in (await session.execute( + select(System.id) + .where(System.deleted_at.is_(None), System.id.notin_(current)) + )).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)) - for system_id, name, description in stale: - await upsert_system_embedding(system_id, name, description) + logger.info( + "System embedding backfill: embedding %d system(s) (%d with stale vectors)", + 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( - select(Milestone.id, Milestone.title, Milestone.description, Milestone.body) - .where(Milestone.deleted_at.is_(None), Milestone.id.notin_(current)) - )).all() + # IDS ONLY — the plan is re-read per milestone below (#4262). + by_version = { + row[0] for row in (await session.execute( + select(Milestone.id) + .where(Milestone.deleted_at.is_(None), Milestone.id.notin_(current)) + )).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)) - for milestone_id, title, description, body in stale: - await upsert_milestone_embedding(milestone_id, title, description, body) + logger.info( + "Milestone embedding backfill: embedding %d milestone(s) (%d with stale vectors)", + 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) diff --git a/tests/test_backfill_reads_current_text.py b/tests/test_backfill_reads_current_text.py new file mode 100644 index 0000000..999ca5d --- /dev/null +++ b/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}" + ) diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 618ec96..e775119 100644 --- a/tests/test_chunking.py +++ b/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() - current_rows.fetchall.return_value = [(1,)] # note 1 is current - note_rows = MagicMock() - note_rows.fetchall.return_value = [ - (1, 42, "current", "body"), - (2, 42, "stale-version", "body"), - ] + def _ids(*values): + r = MagicMock() + r.fetchall.return_value = [(v,) for v in values] + return r + 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()), ):