fix: a backfill reads text when it embeds, not when it scanned (#4264) #176

Merged
bvandeusen merged 1 commits from dev into main 2026-09-21 14:06:35 -04:00
3 changed files with 424 additions and 46 deletions
+222 -38
View File
@@ -18,7 +18,7 @@ from collections.abc import Sequence
from typing import TYPE_CHECKING 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 import async_session
from scribe.models.embedding import NoteEmbedding, RuleEmbedding from scribe.models.embedding import NoteEmbedding, RuleEmbedding
@@ -908,15 +908,135 @@ async def semantic_search_notes(
return final 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: try:
async with async_session() as session: async with async_session() as session:
@@ -930,26 +1050,41 @@ async def backfill_note_embeddings() -> None:
) )
).fetchall() ).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: except Exception:
logger.warning("Embedding backfill: failed to query notes", exc_info=True) logger.warning("Embedding backfill: failed to query notes", exc_info=True)
return return
notes_to_embed = [i for i in all_ids if i not in current or i in stale]
if not notes_to_embed: if not notes_to_embed:
logger.info("Embedding backfill: all notes current at chunker v%d", CHUNKER_VERSION) logger.info("Embedding backfill: all notes current at chunker v%d", CHUNKER_VERSION)
return return
logger.info( 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 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): if not chunk_document(title, body):
continue continue
await upsert_note_embedding(note_id, user_id, title, body) 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( current = select(RuleEmbedding.rule_id).where(
RuleEmbedding.chunker_version == CHUNKER_VERSION 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 = {
.where(Rule.deleted_at.is_(None), Rule.id.notin_(current)) row[0] for row in (await session.execute(
)).all() 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: except Exception:
logger.warning("Rule embedding backfill: failed to query rules", exc_info=True) logger.warning("Rule embedding backfill: failed to query rules", exc_info=True)
return return
@@ -1291,9 +1434,17 @@ async def backfill_rule_embeddings() -> None:
if not stale: if not stale:
logger.info("Rule embedding backfill: all rules current at chunker v%d", CHUNKER_VERSION) logger.info("Rule embedding backfill: all rules current at chunker v%d", CHUNKER_VERSION)
return 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) ────────────────────────────────────────── # ── Milestones (milestone 415) ──────────────────────────────────────────
@@ -1641,10 +1792,18 @@ async def backfill_system_embeddings() -> None:
current = select(SystemEmbedding.system_id).where( current = select(SystemEmbedding.system_id).where(
SystemEmbedding.chunker_version == CHUNKER_VERSION 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 = {
.where(System.deleted_at.is_(None), System.id.notin_(current)) row[0] for row in (await session.execute(
)).all() 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: except Exception:
logger.warning("System embedding backfill: failed to query systems", exc_info=True) logger.warning("System embedding backfill: failed to query systems", exc_info=True)
return return
@@ -1652,9 +1811,17 @@ async def backfill_system_embeddings() -> None:
if not stale: if not stale:
logger.info("System embedding backfill: all systems current at chunker v%d", CHUNKER_VERSION) logger.info("System embedding backfill: all systems current at chunker v%d", CHUNKER_VERSION)
return 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: async def backfill_milestone_embeddings() -> None:
@@ -1668,10 +1835,18 @@ async def backfill_milestone_embeddings() -> None:
current = select(MilestoneEmbedding.milestone_id).where( current = select(MilestoneEmbedding.milestone_id).where(
MilestoneEmbedding.chunker_version == CHUNKER_VERSION 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 = {
.where(Milestone.deleted_at.is_(None), Milestone.id.notin_(current)) row[0] for row in (await session.execute(
)).all() 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: except Exception:
logger.warning("Milestone embedding backfill: failed to query milestones", exc_info=True) logger.warning("Milestone embedding backfill: failed to query milestones", exc_info=True)
return return
@@ -1679,6 +1854,15 @@ async def backfill_milestone_embeddings() -> None:
if not stale: if not stale:
logger.info("Milestone embedding backfill: all milestones current at chunker v%d", CHUNKER_VERSION) logger.info("Milestone embedding backfill: all milestones current at chunker v%d", CHUNKER_VERSION)
return 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)
+186
View File
@@ -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}"
)
+16 -8
View File
@@ -287,18 +287,26 @@ async def test_backfill_reembeds_notes_with_a_stale_chunker_version():
from scribe.services import embeddings as emb 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, 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 ( with (
patch.object(emb, "async_session", return_value=ctx), 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, "upsert_note_embedding", AsyncMock()) as upsert,
patch.object(emb.asyncio, "sleep", AsyncMock()), patch.object(emb.asyncio, "sleep", AsyncMock()),
): ):