From ebd35e44441663a74ebc8d30c29202eab887cf57 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:08:15 -0500 Subject: [PATCH 1/3] fix: Require MCP compatible with current Pydantic --- .github/workflows/ci.yml | 17 ++++++++++++++++- .github/workflows/publish.yml | 11 ++++++++++- pyproject.toml | 5 +++-- tests/mcp-wheel-smoke.py | 10 ++++++++++ 4 files changed, 39 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b3b8e2ee..762659bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -125,10 +125,22 @@ jobs: os: [ubuntu-latest, macos-latest, windows-latest] python-version: ["3.12"] mcp-version: [minimum, latest] + pydantic-version: [latest] include: - os: ubuntu-latest python-version: "3.10" mcp-version: minimum + pydantic-version: latest + # Exercise the MCP/Pydantic floor pair and the release that removed + # eval_type_backport. Other dependencies still use normal resolution. + - os: ubuntu-latest + python-version: "3.10" + mcp-version: minimum + pydantic-version: "2.11.0" + - os: ubuntu-latest + python-version: "3.12" + mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up Python ${{ matrix.python-version }} @@ -149,6 +161,7 @@ jobs: shell: python env: MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} run: | import os from pathlib import Path @@ -161,7 +174,9 @@ jobs: python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") requirements = [str(wheel.resolve())] if os.environ["MCP_VERSION"] == "minimum": - requirements.append("mcp[cli]==1.19.0") + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) subprocess.run([str(python), "-m", "pip", "check"], check=True) subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 9fda4e0e..90bdc63c 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -37,6 +37,12 @@ jobs: fail-fast: false matrix: mcp-version: [minimum, latest] + pydantic-version: [latest] + include: + - mcp-version: minimum + pydantic-version: "2.11.0" + - mcp-version: minimum + pydantic-version: "2.14.0" steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 @@ -54,6 +60,7 @@ jobs: shell: python env: MCP_VERSION: ${{ matrix.mcp-version }} + PYDANTIC_VERSION: ${{ matrix.pydantic-version }} run: | import os from pathlib import Path @@ -66,7 +73,9 @@ jobs: python = env_dir / ("Scripts/python.exe" if os.name == "nt" else "bin/python") requirements = [str(wheel.resolve())] if os.environ["MCP_VERSION"] == "minimum": - requirements.append("mcp[cli]==1.19.0") + requirements.append("mcp[cli]==1.21.1") + if os.environ["PYDANTIC_VERSION"] != "latest": + requirements.append("pydantic==" + os.environ["PYDANTIC_VERSION"]) subprocess.run([str(python), "-m", "pip", "install", *requirements], check=True) subprocess.run([str(python), "-m", "pip", "check"], check=True) subprocess.run([str(python), "-I", "tests/mcp-wheel-smoke.py"], check=True) diff --git a/pyproject.toml b/pyproject.toml index 50d5c171..d512714b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,9 @@ dependencies = [ "model2vec>=0.6.0,<1.0", "sqlite-vec>=0.1.0,<1.0", "numpy>=1.24.0,<3.0", - # Tool metadata requires 1.19; SDK 2 removes the FastMCP import path. - "mcp[cli]>=1.19.0,<2.0", + # 1.21.1 removes a private Pydantic import broken by Pydantic 2.14. + # SDK 2 removes the FastMCP import path. + "mcp[cli]>=1.21.1,<2.0", "httpx>=0.27.0,<1.0", "sentence-transformers>=3.0.0,<7.0", "torch>=2.0.0,<3.0", diff --git a/tests/mcp-wheel-smoke.py b/tests/mcp-wheel-smoke.py index 06388fe9..6a0f2522 100644 --- a/tests/mcp-wheel-smoke.py +++ b/tests/mcp-wheel-smoke.py @@ -182,6 +182,16 @@ async def _check_protocol(root: Path, env: dict[str, str], origin: Path) -> None def main() -> None: _require(bool(sys.flags.isolated), "Run this probe with python -I") distribution = importlib.metadata.distribution("truememory") + # Report public package versions before subprocess checks, whose output is + # intentionally redacted. Dependency import failures then remain traceable. + print(json.dumps({ + "stage": "dependency_versions", + "python_version": ".".join(str(part) for part in sys.version_info[:3]), + "truememory_version": distribution.version, + "mcp_version": importlib.metadata.version("mcp"), + "pydantic_version": importlib.metadata.version("pydantic"), + "pydantic_core_version": importlib.metadata.version("pydantic-core"), + }), flush=True) origin = Path(distribution.locate_file("truememory/__init__.py")).resolve() spec = importlib.util.find_spec("truememory") _require(origin.is_file() and spec is not None and spec.origin is not None, From e487d6ac201efd164ef9c276b637aacfd3c3c139 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:13:51 -0500 Subject: [PATCH 2/3] fix: Stream rebuild source pages with publication fences --- docs/resource-budgets.md | 121 ++ ...test-rebuild-foreground-publication-756.py | 374 ++++++ tests/test-rebuild-source-streaming-756.py | 1136 +++++++++++++++++ truememory/engine.py | 71 +- truememory/rebuild_source.py | 375 ++++++ truememory/storage.py | 8 +- truememory/vector_search.py | 468 +++---- 7 files changed, 2281 insertions(+), 272 deletions(-) create mode 100644 tests/test-rebuild-foreground-publication-756.py create mode 100644 tests/test-rebuild-source-streaming-756.py create mode 100644 truememory/rebuild_source.py diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index caa30888..c22ee937 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -501,3 +501,124 @@ responsible for its serialization. The engine's existing manual path is unchange in this foundation checkpoint. Per-layer freshness enforcement, scheduler routing, vector dependency generations and measured maintenance performance remain later integration work; initially pending rows are not success claims. + + +## Standalone rebuild source streaming (issue #756, first checkpoint) + +`build_vectors()` and `build_separation_vectors()` read database-backed inputs +in ascending-ID pages, using the existing native batch size. They capture a +fixed upper ID, row count and A10 source revision, including its epoch. Each +page read ends before inference. Publication checks the source revision, +model generation, target manifest and database schema version under writer +ownership. Writer acquisition precedes the nonblocking model fence; that +fence remains held through the terminal COMMIT or SAVEPOINT RELEASE. A source correction, deletion, lower-ID insertion, epoch change or +schema change rejects the unfinished generation. An unrelated schema change +is conservatively invalidating. New high-ID source rows remain outside the +captured range. + +A versioned metadata manifest contains model/dimension/target identity, +revision counters, the last **consumed input ID**, consumed count and inserted +vector count. It contains no source text. Skipped zero/nonfinite vectors still +advance consumed input, so restart cannot infer progress from `MAX(rowid)`. +A saved manifest without matching identity and revision is restarted once on +the next explicit call; the builder does not run an automatic restart loop. +Legacy boolean markers alone are insufficient resume evidence. + +The native batch size and `txn_batch` durability cadence remain configurable. +Each group is encoded and serialized before its owned writer transaction. +Serialized payload is bounded by +`max(1, txn_batch) * batch_size * dimension * 4` bytes per resident index, +plus Python objects, one input page, the current native output and allocator +memory. Defaults at 256 dimensions give: + +- CPU: `100 * 100 * 256 * 4 = 10,240,000` bytes (`9.765625 MiB`). +- MPS: `100 * 16 * 256 * 4 = 1,638,400` bytes (`1.5625 MiB`). +- Custom 384-dimensional CPU example: `100 * 100 * 384 * 4 = 15,360,000` + bytes (`14.6484375 MiB`). + +Completion and separation standalone calls stage their groups separately. +This bound is independent of corpus rows with fixed controls, but can increase +staging for small corpora compared with immediate insertion. It is not a +whole-process memory cap. A single record is not truncated and can itself be +large. Native model residency and thermal behavior are not measured by the +synthetic regression tests. + +One maintenance owner covers the complete operation. Same-thread nesting in +an existing A10 maintenance worker reuses its owner. A caller's open SQLite +transaction is preserved with savepoints; the rebuild cannot release locks +that the caller already owns. Such callers also retain control of durability. +Otherwise, no rebuild-owned read or writer transaction remains open during +inference. Short page snapshots avoid pinning a job-long WAL reader; they do +not promise globally bounded WAL when other readers exist. + +Explicit `messages` lists preserve caller order and text. A streamed digest +binds restart to that exact list without making another corpus-sized copy; +the consumed list position preserves out-of-order IDs. The caller's allocated +list is outside the database paging bound. Concurrent mutation of a supplied +list is not supported. Database-backed calls reuse the canonical source +tracker only with the exact built-in `messages` definition and verified tracking. +Older connections, including deprecated +`engine.open()` and the minimal five-column separation schema, instead install +a rebuild-only bridge when canonical tracking is unavailable. Completion also +supports a legacy table containing only `id INTEGER PRIMARY KEY` and `content`. + +The bridge adds one namespaced state table and three source triggers. It does +not add columns to `messages`, migrate other storage features, or mark the +canonical maintenance tracker ready. Installation and repair own the writer +before model loading or clearing vectors; caller transactions use savepoints. +The token explicitly identifies its tracker, epoch, covered source fields, +table schema and index definitions. Page/publication checks verify that source +schema and the exact trigger definitions; missing or altered tracking starts +a new epoch on the next call, invalidating untracked intervals. These checks +read schema metadata and scalar counters, not corpus text. Read-only legacy +connections and unsupported source schemas fail before model allocation. +Bridge updates compare each covered value's SQLite storage type and bytes, +so custom `NOCASE`/`RTRIM` collations cannot hide changed encoder input. +Other table definitions retain the bridge even after a full storage migration +reports canonical tracking ready. This conservative definition whitelist avoids +inferring equality semantics from arbitrary SQL; compatible custom definitions +remain usable with the additional bridge counters. + +Same-ID replacement invalidates a bridge generation even with recursive +triggers disabled. Custom schemas with unique indexes, including indexed primary keys that +are not rowid aliases, use a conservative bridge: every successful insert or update invalidates the +captured prefix. This includes expression and partial unique indexes, where +`OR REPLACE` can silently delete another row without a delete trigger. Such +schemas remain supported, but concurrent high-ID appends cannot preserve a +running generation. They retain this bridge even if the canonical tracker +reports ready, because its append certificate does not cover those hidden +replacement deletions. There is no repeated automatic restart loop. + +When the trusted built-in definition and canonical tracking become suitable, +the next rebuild changes tracker +identity and starts a fresh generation; it atomically retires only the bridge's +owned triggers. Its inactive state table remains. The bridge never certifies +the readiness of other maintenance layers. Normal canonical schemas need no +bridge table, triggers or additional writer commit during setup. + +Foreground engine add/update and `embed_single` bind inference to the active +model generation. Same-model appends above an in-progress database range are +published normally and cannot advance its consumed cursor. Initial clear and +source capture share one writer transaction, so a committed earlier append +joins the captured range. A correction within the range publishes its current +vector and invalidates the rebuild's source fence. A late model mismatch is +reported categorically: add rolls back its uncommitted source row; update's +existing storage helper has already committed its source, so its error says +to retry vector publication. Explicit-input or unattested in-progress targets +reject concurrent publication rather than silently reporting success. + +Foreground publication acquires the SQLite writer before its nonblocking +model fence, then validates the target manifest and retains both owners +through the owned commit or caller savepoint release. Inference finishes +before either is acquired. Add source insertion shares that transaction; +rejection rolls back only the new add, preserving unrelated pending caller +writes. Successful add and `embed_single` calls inside an existing caller +transaction also leave durability to that caller. Update retains its existing +source-helper commit behavior; the subsequent vector publication is fenced +separately and cannot claim the source update was rolled back. + +This checkpoint does **not** remove the tier-switch manager/worker's retained +message lists or solve tier configuration publication across SQLite, config +file and process state. Issue #756 remains open for that next checkpoint. +No global embedder-compatibility authority, models, precision, dimensions, +reranker selection, retrieval depth or deployment defaults change here. diff --git a/tests/test-rebuild-foreground-publication-756.py b/tests/test-rebuild-foreground-publication-756.py new file mode 100644 index 00000000..6c71ac5e --- /dev/null +++ b/tests/test-rebuild-foreground-publication-756.py @@ -0,0 +1,374 @@ +"""Actual foreground methods and rebuild publication on synthetic databases.""" + +import ast +import logging +import runpy +import sqlite3 +import threading +import types +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +BASE = types.SimpleNamespace(**runpy.run_path(str(ROOT / "tests/test-rebuild-source-streaming-756.py"))) + + +class TestForegroundPublication(unittest.TestCase): + setUp = BASE.TestStreamingBuilders.setUp + seed = BASE.TestStreamingBuilders.seed + ids = BASE.TestStreamingBuilders.ids + manifest = BASE.TestStreamingBuilders.manifest + writer = BASE.TestStreamingBuilders.writer + + def engine(self, conn=None): + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + methods = [node for node in original.body if isinstance(node, ast.FunctionDef) and node.name in ("add", "update")] + cls = ast.ClassDef(name="SyntheticEngine", bases=[], keywords=[], body=methods, decorator_list=[]) + namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], + logger=logging.getLogger(__name__), MAX_CONTENT_LENGTH=10000, + insert_message=self.modules["storage"].insert_message, + update_message=self.modules["storage"].update_message) + exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), "actual-engine-boundaries", "exec"), namespace) + engine = namespace["SyntheticEngine"]() + engine.conn = conn or self.conn + engine._write_lock = threading.Lock() + engine._ensure_connection = lambda: None + engine._maybe_auto_consolidate = lambda: None + engine._has_vectors = True + engine._has_personality = engine._has_style_vec = False + engine.get = lambda mid: self.modules["storage"].get_message(engine.conn, mid) + return engine + + def test_actual_add_above_active_range_survives_builder_completion(self): + self.seed(range(1, 5)) + other = self.writer() + engine = self.engine(other) + inserted = [] + def encode(texts): + if not inserted: + inserted.append(True) + item = engine.add("synthetic-late-append", sender="sender", recipient="recipient", timestamp="2026-01-02") + self.assertEqual(item["id"], 5) + inserted.append(item["id"]) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 4) + self.assertEqual(self.ids(), [1, 2, 3, 4, 5]) + self.assertEqual((self.manifest().cursor, self.manifest().consumed, self.manifest().source.high_id), (4, 4, 4)) + self.assertEqual(self.ids("vec_messages_sep"), [5]) + + def test_actual_update_invalidates_build_but_keeps_corrected_vector(self): + self.seed(range(1, 5)) + engine = self.engine(self.writer()) + updated = [] + def encode(texts): + if not updated: + updated.append(True) + self.assertEqual(engine.update(1, "synthetic-correction")["content"], "synthetic-correction") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.manifest().consumed, 0) + self.assertTrue(self.vector._build_in_progress(self.conn, "vec_messages")) + + def test_late_model_change_rolls_back_add_source_and_does_not_write_vectors(self): + engine = self.engine() + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model changed"): + engine.add("synthetic-stale-model") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertFalse(self.conn.in_transaction) + + def test_late_model_change_update_reports_already_committed_source(self): + self.seed((1,)) + engine = self.engine() + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "Source update committed"): + engine.update(1, "synthetic-committed-correction") + self.assertEqual(self.conn.execute("SELECT content FROM messages WHERE id=1").fetchone()[0], "synthetic-committed-correction") + self.assertEqual(self.ids(), []) + + def test_completed_foreign_model_manifest_rejects_old_process_write(self): + self.seed((1,)) + self.vector.build_vectors(self.conn) + manifest = self.manifest() + from dataclasses import replace + self.source.save_manifest(self.conn, self.source.manifest_key(("vec_messages",)), replace(manifest, model="different-synthetic-model")) + self.conn.commit() + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "different embedding model"): + self.engine().add("synthetic-old-process") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 1) + self.assertEqual(self.ids(), [1]) + + def test_unattested_marker_rejects_add_without_silent_success(self): + self.vector._mark_build_in_progress(self.conn, "vec_messages") + self.conn.commit() + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "without matching source"): + self.engine().add("synthetic-unattested") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_guard_releases_model_lock_when_sql_publication_raises(self): + self.seed((1,)) + identity = (0, "synthetic-model", 2) + with self.assertRaisesRegex(RuntimeError, "synthetic write failure"): + with self.vector._foreground_vector_publication(self.conn, identity, 1): + raise RuntimeError("synthetic write failure") + self.assertTrue(self.vector._lock.acquire(blocking=False)) + self.vector._lock.release() + + def test_embed_single_finishes_both_encodes_before_owned_database_write(self): + self.seed((1,)) + def encode(texts): + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + self.assertFalse(self.conn.in_transaction) + + def test_embed_single_stale_model_is_rejected_before_either_insert(self): + self.seed((1,)) + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.vector.VectorPublicationChanged): + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def _attempt_foreign_manifest_at_vector_insert(self, operation): + self.seed((1,)) + other = self.writer() + other.execute("PRAGMA busy_timeout = 0") + manifest = self.source.RebuildManifest.new( + "different-synthetic-model", 2, ("vec_messages",), total=0, + ) + attempts = [] + ownership = [] + real = self.conn + case = self + + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if not attempts and sql.startswith(("INSERT INTO vec_messages(", "DELETE FROM vec_messages WHERE")): + ownership.append((real.in_transaction, case.vector._lock.locked())) + try: + case.source.save_manifest(other, case.source.manifest_key(("vec_messages",)), manifest) + other.commit() + attempts.append("committed") + except sqlite3.OperationalError as error: + case.assertIn("locked", str(error).lower()) + attempts.append("blocked") + finally: + other.rollback() + return real.execute(sql, parameters) + + probe = Probe() + if operation == "embed": + self.vector.embed_single(probe, 1, "synthetic-1") + expected = [1] + elif operation == "update": + self.engine(probe).update(1, "synthetic-correction") + expected = [1] + else: + self.engine(probe).add("synthetic-append") + expected = [2] + self.assertEqual(attempts, ["blocked"]) + self.assertEqual(ownership, [(True, True)]) + self.assertIsNone(self.manifest()) + self.assertEqual(self.ids(), expected) + self.assertEqual(self.ids("vec_messages_sep"), expected) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()[0], + "synthetic-model") + self.assertFalse(self.conn.in_transaction) + + def test_embed_single_excludes_foreign_manifest_between_validation_and_insert(self): + self._attempt_foreign_manifest_at_vector_insert("embed") + + def test_update_excludes_foreign_manifest_after_source_helper_commits(self): + self._attempt_foreign_manifest_at_vector_insert("update") + + def test_add_excludes_foreign_manifest_through_vector_publication(self): + self._attempt_foreign_manifest_at_vector_insert("add") + + def test_rejected_add_preserves_unrelated_caller_write_and_transaction(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + def encode(texts): + if len(self.calls) == 2: + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.vector.VectorPublicationChanged): + self.engine().add("synthetic-rejected") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.ids(), []) + + def test_foreign_target_rejection_rolls_back_only_new_add_after_source_insert(self): + manifest = self.source.RebuildManifest.new("different-synthetic-model", 2, ("vec_messages",)) + self.source.save_manifest(self.conn, self.source.manifest_key(("vec_messages",)), manifest) + self.conn.commit() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "different embedding model"): + self.engine().add("synthetic-rejected") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.ids(), []) + + def test_successful_add_remains_inside_caller_transaction(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + item = self.engine().add("synthetic-pending-add") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), [item["id"]]) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + + def test_embed_single_keeps_caller_work_pending(self): + self.seed((1,)) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + self.vector.embed_single(self.conn, 1, "synthetic-1") + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), [1]) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def _assert_foreground_writer_and_terminal_fence(self, operation, caller_owned): + self.seed((1,)) + if caller_owned: + self.conn.execute("BEGIN") + real = self.conn + case = self + boundaries = [] + + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if sql == "BEGIN IMMEDIATE" or sql == "UPDATE messages SET id = id WHERE 0": + case.assertFalse(case.vector._lock.locked()) + boundaries.append("writer") + if sql == "RELEASE SAVEPOINT rebuild_source": + case.assertTrue(case.vector._lock.locked()) + boundaries.append("release") + return real.execute(sql, parameters) + + def commit(self): + case.assertTrue(case.vector._lock.locked()) + boundaries.append("commit") + real.commit() + + if operation == "embed": + self.vector.embed_single(Probe(), 1, "synthetic-1") + else: + self.engine(Probe()).add("synthetic-append") + self.assertEqual(boundaries, ["writer", "release" if caller_owned else "commit"]) + self.assertEqual(real.in_transaction, caller_owned) + self.assertFalse(self.vector._lock.locked()) + + def test_embed_owned_commit_stays_inside_model_fence(self): + self._assert_foreground_writer_and_terminal_fence("embed", False) + + def test_embed_caller_writer_upgrade_precedes_fence_and_release_stays_inside(self): + self._assert_foreground_writer_and_terminal_fence("embed", True) + + def test_add_owned_commit_stays_inside_model_fence(self): + self._assert_foreground_writer_and_terminal_fence("add", False) + + def test_add_caller_writer_upgrade_precedes_fence_and_release_stays_inside(self): + self._assert_foreground_writer_and_terminal_fence("add", True) + + def test_denied_add_savepoint_release_rolls_back_only_add(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + denied = [] + def authorize(action, first, second, _database, _source): + if action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" and second == "rebuild_source" and not denied: + self.assertTrue(self.vector._lock.locked()) + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + self.conn.set_authorizer(authorize) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.engine().add("synthetic-denied-release") + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def test_cancelled_add_rolls_back_only_add_and_releases_model_fence(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + real = self.conn + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, parameters=()): + if sql.startswith("INSERT INTO vec_messages_sep("): + raise KeyboardInterrupt("synthetic cancellation") + return real.execute(sql, parameters) + with self.assertRaises(KeyboardInterrupt): + self.engine(Probe()).add("synthetic-cancelled") + self.assertTrue(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + def test_failed_add_owned_commit_does_not_publish_on_later_commit(self): + real = self.conn + case = self + class Probe: + def __getattr__(self, name): + return getattr(real, name) + + def commit(self): + case.assertTrue(case.vector._lock.locked()) + raise sqlite3.OperationalError("synthetic commit failure") + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic commit failure"): + self.engine(Probe()).add("synthetic-uncommitted") + self.assertFalse(self.conn.in_transaction) + self.assertFalse(self.vector._lock.locked()) + self.conn.commit() + self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py new file mode 100644 index 00000000..48edfba4 --- /dev/null +++ b/tests/test-rebuild-source-streaming-756.py @@ -0,0 +1,1136 @@ +"""Synthetic SQLite and actual rebuild control flow; no model imports.""" + +import ast +import builtins +from collections.abc import Iterator +import datetime +import logging +import math +import os +import sqlite3 +import struct +import sys +import tempfile +import threading +import types +import unittest +from contextlib import closing, contextmanager +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + + +def load_modules(): + modules = {} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "torch": + raise ImportError("Native models are excluded from synthetic tests") + if name.startswith("truememory."): + short = name.removeprefix("truememory.") + if short in modules: + return modules[short] + raise AssertionError("Unexpected application import: " + name) + return builtins.__import__(name, globals, locals, fromlist, level) + + for name in ("storage", "_platform", "maintenance", "rebuild_source"): + module = types.ModuleType("synthetic_stream_" + name) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + vector = types.ModuleType("synthetic_stream_vector") + vector.__dict__.update( + __builtins__=dict(vars(builtins), __import__=safe_import), + contextmanager=contextmanager, closing=closing, Iterator=Iterator, replace=replace, datetime=datetime, + sqlite3=sqlite3, logger=logging.getLogger(__name__), _lock=threading.Lock(), + _model_generation=0, EMBEDDING_MODEL="synthetic-model", _embedding_dim=2, + _BUILD_VECTORS_TXN_BATCH=100, _BUILD_STATE_KEY_PREFIX="vec_build_state:", + _get_batch_size=lambda: 2, _flush_mps_cache=lambda: None, + _active_vec_table=lambda conn: "vec_messages", _active_sep_table=lambda conn: "vec_messages_sep", + serialize_f32=lambda vector: struct.pack("2f", *vector), + _normalize_for_cosine=lambda vector: vector if all(math.isfinite(v) for v in vector) and any(vector) else None, + ) + selected = {"_capture_rebuild_model", "_rebuild_model_fence", "_build_streamed_vectors", + "build_vectors", "build_separation_vectors", "_build_sep_text", "_ensure_metadata_table", + "_build_state_key", "_mark_build_in_progress", "_clear_build_in_progress", + "_build_in_progress", "_write_embedder_metadata_no_commit", "_write_embedder_metadata", + "_foreground_vector_publication", "_foreground_model_fence", "_validate_foreground_vector_target", + "VectorPublicationChanged", "embed_single"} + tree = ast.parse((ROOT / "truememory/vector_search.py").read_text(encoding="utf-8")) + body = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name in selected] + exec(compile(ast.Module(body=body, type_ignores=[]), "actual-vector-builders", "exec"), vector.__dict__) + modules["vector_search"] = vector + return modules + + +class TestStreamingBuilders(unittest.TestCase): + def setUp(self): + self.modules = load_modules() + self.source = self.modules["rebuild_source"] + self.vector = self.modules["vector_search"] + self.temp = tempfile.TemporaryDirectory(prefix="synthetic-streaming-") + self.addCleanup(self.temp.cleanup) + self.path = Path(self.temp.name) / "synthetic.sqlite" + self.conn = self.modules["storage"].create_db(self.path) + self.addCleanup(self.conn.close) + self.conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + self.model = object() + self.vector.get_model = lambda: self.model + self.calls = [] + self.on_encode = None + + def encode(model, texts, **kwargs): + self.assertIs(model, self.model) + self.calls.append(list(texts)) + if self.on_encode is not None: + return self.on_encode(texts) + return [[1., 0.] for _ in texts] + self.vector._encode_with_mps_fallback = encode + + def seed(self, ids=range(1, 10)): + for mid in ids: + self.conn.execute("INSERT INTO messages(id,content,sender,recipient,timestamp) VALUES (?,?,?,?,?)", + (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02")) + self.conn.commit() + + def ids(self, table="vec_messages"): + return [row[0] for row in self.conn.execute(f"SELECT rowid FROM {table} ORDER BY rowid")] + + def manifest(self, table="vec_messages"): + return self.source.load_manifest(self.conn, self.source.manifest_key((table,))) + + def writer(self): + conn = self.modules["storage"].create_db(self.path) + self.addCleanup(conn.close) + return conn + + def test_exact_ids_texts_and_separation_with_negative_zero_and_holes(self): + self.seed((-10, 0, 2, 8, 11)) + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=2), 5) + self.assertEqual(self.ids(), [-10, 0, 2, 8, 11]) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in (-10, 0, 2, 8, 11)]) + self.calls.clear() + self.assertEqual(self.vector.build_separation_vectors(self.conn, txn_batch=2), 5) + self.assertEqual(sum(self.calls, []), [f"sender to recipient on 2026-01-02: synthetic-{mid}" for mid in (-10, 0, 2, 8, 11)]) + + def test_pages_release_read_and_writer_transactions_before_each_encode(self): + self.seed(range(1, 22)) + writer = self.writer() + def encode(texts): + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('synthetic-writer','ok')") + writer.commit() + self.assertLessEqual(len(texts), 2) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.build_vectors(self.conn, txn_batch=3) + self.assertEqual(len(self.ids()), 21) + self.assertFalse(self.conn.in_transaction) + + def test_resume_uses_consumed_input_when_last_vector_skipped(self): + self.seed(range(1, 7)) + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.], [0., 0.]] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), [1]) + self.assertEqual((self.manifest().cursor, self.manifest().consumed), (2, 2)) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 4) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in range(3, 7)]) + self.assertEqual(self.ids(), [1, 3, 4, 5, 6]) + + def test_all_invalid_outputs_finish_with_nonzero_consumed_cursor(self): + self.seed() + self.on_encode = lambda texts: [[0., 0.] for _ in texts] + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=2), 0) + manifest = self.manifest() + self.assertEqual((manifest.consumed, manifest.cursor, manifest.outputs, manifest.complete), (9, 9, 0, True)) + self.assertFalse(self.vector._build_in_progress(self.conn, "vec_messages")) + + def test_explicit_list_keeps_order_and_text_without_database_rows(self): + messages = [{"id": 9, "content": "synthetic-nine"}, {"id": -1, "content": "synthetic-minus"}, + {"id": 3, "content": "synthetic-three"}] + self.assertEqual(self.vector.build_vectors(self.conn, messages, txn_batch=1), 3) + self.assertEqual(sum(self.calls, []), [message["content"] for message in messages]) + self.assertIsNone(self.manifest().source) + + def test_caller_transaction_is_preserved_on_success(self): + self.seed() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + reader = self.writer() + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(reader.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.conn.rollback() + self.assertEqual(self.ids(), []) + self.assertIsNone(self.manifest()) + + def test_correction_during_compute_rejects_publication_and_fresh_run_restarts(self): + self.seed() + writer = self.writer() + def encode(texts): + writer.execute("UPDATE messages SET content='synthetic-changed' WHERE id=1") + writer.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertTrue(self.vector._build_in_progress(self.conn, "vec_messages")) + self.on_encode = None + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(len(self.ids()), 9) + + def test_model_change_during_compute_rejects_rows_and_metadata(self): + self.seed() + def encode(texts): + self.vector._model_generation += 1 + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "model changed"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertIsNone(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()) + + def test_empty_database_finishes_without_loading_a_native_model(self): + def no_load(): + raise AssertionError("Empty source must not load a model") + self.vector.get_model = no_load + self.assertEqual(self.vector.build_vectors(self.conn), 0) + self.assertTrue(self.manifest().complete) + self.assertEqual(self.calls, []) + + def test_incomplete_model_output_is_rejected_before_writes(self): + self.seed() + self.on_encode = lambda texts: [] + with self.assertRaisesRegex(ValueError, "count"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_epoch_change_invalidates_saved_resume(self): + self.seed() + def encode(texts): + if len(self.calls) > 1: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(RuntimeError): + self.vector.build_vectors(self.conn, txn_batch=1) + old = self.manifest() + self.conn.execute("UPDATE maintenance_source_state SET epoch='synthetic-new-epoch'") + self.conn.commit() + self.on_encode = None + self.calls.clear() + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertNotEqual(self.manifest().generation, old.generation) + self.assertEqual(sum(self.calls, []), [f"synthetic-{mid}" for mid in range(1, 10)]) + + def test_source_mutations_after_fetch_are_classified_without_partial_publication(self): + cases = ( + ("hole", "INSERT INTO messages(id,content) VALUES (3,'synthetic-hole')", False), + ("delete", "DELETE FROM messages WHERE id=2", False), + ("rowid", "UPDATE messages SET rowid=3 WHERE id=2", False), + ("sender", "UPDATE messages SET sender='synthetic-change' WHERE id=2", False), + ("timestamp", "UPDATE messages SET timestamp='2025-01-01' WHERE id=2", False), + ("derived", "UPDATE messages SET episode_id=13 WHERE id=2", True), + ("append", "INSERT INTO messages(id,content) VALUES (100,'synthetic-tail')", True), + ) + for name, sql, accepted in cases: + with self.subTest(name=name): + # Each mutation gets an independent generation and database. + case = TestStreamingBuilders() + case.setUp() + try: + case.seed((2, 4, 6)) + writer = case.writer() + changed = [] + def encode(texts): + if not changed: + changed.append(True) + writer.execute(sql) + writer.commit() + return [[1., 0.] for _ in texts] + case.on_encode = encode + if accepted: + case.vector.build_vectors(case.conn, txn_batch=1) + self.assertEqual(case.ids(), [2, 4, 6]) + self.assertEqual(case.manifest().consumed, 3) + else: + with self.assertRaises(case.source.RebuildSourceChanged): + case.vector.build_vectors(case.conn, txn_batch=1) + self.assertEqual(case.ids(), []) + finally: + case.doCleanups() + + def test_rolled_back_source_correction_does_not_invalidate(self): + self.seed() + writer = self.writer() + def encode(texts): + writer.execute("UPDATE messages SET content='synthetic-rolled-back' WHERE id=1") + writer.rollback() + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 9) + + def test_short_pages_do_not_pin_a_wal_reader_during_native_work(self): + self.seed() + writer = self.writer() + checks = [] + def encode(texts): + self.assertFalse(self.conn.in_transaction) + writer.execute("INSERT OR REPLACE INTO metadata(key,value) VALUES ('checkpoint-probe','synthetic')") + writer.commit() + busy, log_frames, checkpointed = writer.execute("PRAGMA wal_checkpoint(PASSIVE)").fetchone() + checks.append((busy, log_frames, checkpointed)) + self.assertEqual(busy, 0) + self.assertEqual(log_frames, checkpointed) + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.vector.build_vectors(self.conn, txn_batch=2) + self.assertTrue(checks) + + def test_group_staging_and_commit_cadence_stay_bounded(self): + self.seed(range(1, 24)) + real = self.conn + sizes = [] + class Observed: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def execute(self, sql, parameters=()): + normalized = " ".join(sql.lower().split()) + if normalized.startswith("select id, content"): + self_check.assertIn(" limit ?", normalized) + self_check.assertLessEqual(parameters[-1], 2) + return real.execute(sql, parameters) + def executemany(self, sql, rows): + sizes.append((len(rows), sum(len(row[1]) for row in rows))) + return real.executemany(sql, rows) + def commit(self): + self.commits += 1 + real.commit() + self_check = self + observed = Observed() + self.vector.build_vectors(observed, txn_batch=3) + self.assertEqual([item[0] for item in sizes], [6, 6, 6, 5]) + self.assertLessEqual(max(item[1] for item in sizes), 3 * 2 * 2 * 4) + self.assertEqual(observed.commits, 1 + 4 + 1) + + def test_cancellation_before_group_commit_has_no_partial_group(self): + self.seed() + def encode(texts): + if len(self.calls) == 2: + raise KeyboardInterrupt("synthetic cancellation") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(KeyboardInterrupt): + self.vector.build_vectors(self.conn, txn_batch=3) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + self.assertFalse(self.conn.in_transaction) + + def test_failed_commit_cannot_leave_vectors_or_cursor_for_later_commit(self): + self.seed() + real = self.conn + class DeniedCommit: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + self.commits += 1 + if self.commits == 2: + raise sqlite3.OperationalError("synthetic commit failure") + real.commit() + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic commit failure"): + self.vector.build_vectors(DeniedCommit(), txn_batch=1) + real.commit() + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_savepoint_release_denial_rolls_back_group_but_keeps_caller_work(self): + self.seed() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + denied = [] + def authorize(action, first, second, *_args): + if action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" and second == "rebuild_source" and self.calls and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + def encode(texts): + # Reset the statement cache after initial savepoints were prepared. + self.conn.set_authorizer(authorize) + return [[1., 0.] for _ in texts] + self.on_encode = encode + try: + with self.assertRaises(sqlite3.DatabaseError): + self.vector.build_vectors(self.conn, txn_batch=1) + finally: + self.conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='caller'").fetchone()[0], "pending") + self.conn.commit() + self.assertEqual(self.ids(), []) + + def test_legacy_marker_with_high_output_id_is_not_a_resume_certificate(self): + self.seed((1, 2, 3)) + self.conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (999,?)", (struct.pack("2f", 1., 0.),)) + self.vector._mark_build_in_progress(self.conn, "vec_messages") + self.conn.commit() + self.assertEqual(self.vector.build_vectors(self.conn, txn_batch=1), 3) + self.assertEqual(self.ids(), [1, 2, 3]) + + def test_explicit_input_resume_tracks_position_not_sorted_id(self): + messages = [{"id": mid, "content": f"synthetic-{mid}"} for mid in (9, -4, 8, 2)] + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(RuntimeError): + self.vector.build_vectors(self.conn, messages, txn_batch=1) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(self.conn, messages, txn_batch=1), 2) + self.assertEqual(sum(self.calls, []), ["synthetic-8", "synthetic-2"]) + + def test_competing_owner_is_rejected_before_loading_or_clearing(self): + self.seed() + ready = threading.Event() + release = threading.Event() + errors = [] + def own(): + try: + with self.modules["maintenance"].maintenance_owner(self.path): + ready.set() + if not release.wait(3): + raise RuntimeError("synthetic owner timeout") + except BaseException as exc: + errors.append(exc) + owner = threading.Thread(target=own) + owner.start() + try: + self.assertTrue(ready.wait(3)) + with self.assertRaises(self.modules["maintenance"].MaintenanceBusyError): + self.vector.build_vectors(self.conn) + self.assertEqual(self.calls, []) + self.assertIsNone(self.manifest()) + finally: + release.set() + owner.join(3) + self.assertFalse(owner.is_alive()) + self.assertEqual(errors, []) + + def test_unmigrated_legacy_connection_keeps_bridge_after_full_storage_migration(self): + legacy_path = Path(self.temp.name) / "synthetic-legacy.sqlite" + legacy = sqlite3.connect(legacy_path) + self.addCleanup(legacy.close) + legacy.executescript("CREATE TABLE messages(id INTEGER PRIMARY KEY,content TEXT,sender TEXT,recipient TEXT,timestamp TEXT);" + "CREATE TABLE vec_messages(embedding BLOB);" + "INSERT INTO messages VALUES(1,'synthetic-legacy','sender','recipient','2026-01-01');") + legacy.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (1,?)", (struct.pack("2f", 1., 0.),)) + legacy.commit() + original_columns = legacy.execute("PRAGMA table_info(messages)").fetchall() + self.assertEqual(self.vector.build_vectors(legacy, txn_batch=1), 1) + self.assertEqual(legacy.execute("SELECT rowid FROM vec_messages").fetchall(), [(1,)]) + self.assertFalse(legacy.in_transaction) + self.assertEqual(legacy.execute("PRAGMA table_info(messages)").fetchall(), original_columns) + self.assertIsNone(legacy.execute("SELECT name FROM sqlite_master WHERE name='maintenance_source_state'").fetchone()) + before = self.source.capture_source(legacy) + self.assertEqual(before.tracker, "rebuild-bridge-v1") + migrated = self.modules["storage"].create_db(legacy_path) + self.addCleanup(migrated.close) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(migrated) + self.assertEqual(self.vector.build_vectors(migrated, txn_batch=1), 1) + self.assertEqual(self.source.capture_source(migrated).tracker, "rebuild-bridge-v1") + self.assertEqual(len(self.source._installed_triggers(migrated, self.source._BRIDGE_TRIGGERS)), 3) + self.modules["maintenance"].read_source_revision(migrated) + + def test_schema_replacement_during_compute_rejects_target_identity(self): + self.seed() + writer = self.writer() + changed = [] + def encode(texts): + if not changed: + changed.append(True) + writer.execute("CREATE TABLE synthetic_new_schema(value TEXT)") + writer.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(self.source.RebuildSourceChanged, "schema changed"): + self.vector.build_vectors(self.conn, txn_batch=1) + self.assertEqual(self.ids(), []) + self.assertEqual(self.manifest().consumed, 0) + + def test_model_fence_covers_initial_group_final_and_empty_commits(self): + self.seed((1, 2, 3)) + real = self.conn + fixture = self + class Observed: + commits = 0 + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + acquired = fixture.vector._lock.acquire(blocking=False) + if acquired: + fixture.vector._lock.release() + fixture.assertFalse(acquired, "Model identity was released before publication commit") + self.commits += 1 + real.commit() + observed = Observed() + self.vector.build_vectors(observed, txn_batch=1) + self.assertEqual(observed.commits, 4) + real.execute("DELETE FROM messages") + real.commit() + self.vector.build_vectors(observed, txn_batch=1) + self.assertEqual(observed.commits, 5) + + def test_competing_model_switch_waits_until_actual_final_commit(self): + self.seed((1, 2)) + underlying = threading.Lock() + attempted = threading.Event() + switched = threading.Event() + threads = [] + fixture = self + class ObservedLock: + def acquire(self, *args, **kwargs): + if threading.current_thread().name == "synthetic-model-switch": + attempted.set() + return underlying.acquire(*args, **kwargs) + def release(self): + underlying.release() + def __enter__(self): + self.acquire() + return self + def __exit__(self, *_args): + self.release() + self.vector._lock = ObservedLock() + real = self.conn + def switch(): + with fixture.vector._lock: + fixture.vector._model_generation += 1 + fixture.vector.EMBEDDING_MODEL = "synthetic-next-model" + switched.set() + class CommitBoundary: + def __getattr__(self, name): + return getattr(real, name) + def commit(self): + manifest = fixture.manifest() + if manifest is not None and manifest.complete: + thread = threading.Thread(target=switch, name="synthetic-model-switch") + threads.append(thread) + thread.start() + fixture.assertTrue(attempted.wait(3)) + fixture.assertFalse(switched.is_set(), "Model switched before final commit") + real.commit() + try: + self.vector.build_vectors(CommitBoundary(), txn_batch=1) + finally: + for thread in threads: + thread.join(3) + self.assertEqual(len(threads), 1) + self.assertFalse(threads[0].is_alive()) + self.assertTrue(switched.is_set()) + self.assertTrue(self.manifest().complete) + self.assertEqual(self.manifest().model, "synthetic-model") + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='embed_model'").fetchone()[0], "synthetic-model") + + def test_caller_writer_acquisition_precedes_model_fence_and_release_is_fenced(self): + self.seed((1, 2)) + real = self.conn + real.execute("INSERT INTO metadata(key,value) VALUES ('caller','pending')") + fixture = self + releases = [] + class Observed: + def __init__(self): + self.savepoints = [] + def __getattr__(self, name): + return getattr(real, name) + def execute(self, sql, parameters=()): + if sql == "SAVEPOINT rebuild_source": + self.savepoints.append(False) + if sql == "UPDATE messages SET id = id WHERE 0": + acquired = fixture.vector._lock.acquire(blocking=False) + fixture.assertTrue(acquired, "Writer acquisition must precede model fence") + fixture.vector._lock.release() + self.savepoints[-1] = True + if sql == "RELEASE SAVEPOINT rebuild_source": + publication = self.savepoints.pop() + if publication: + acquired = fixture.vector._lock.acquire(blocking=False) + if acquired: + fixture.vector._lock.release() + fixture.assertFalse(acquired, "Publication RELEASE must retain the model fence") + releases.append(True) + return real.execute(sql, parameters) + self.vector.build_vectors(Observed(), txn_batch=1) + self.assertGreaterEqual(len(releases), 3) + self.assertTrue(real.in_transaction) + real.rollback() + self.assertEqual(self.ids(), []) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_TEST_NATIVE_VEC") == "1", + "Actual sqlite-vec verification is an explicit isolated-host gate") +class TestStreamingSQLiteVec(TestStreamingBuilders): + def setUp(self): + super().setUp() + import sqlite_vec + self.sqlite_vec = sqlite_vec + self.conn.enable_load_extension(True) + sqlite_vec.load(self.conn) + self.conn.enable_load_extension(False) + for table in ("vec_messages", "vec_messages_sep"): + self.conn.execute(f"DROP TABLE {table}") + self.conn.execute(f"CREATE VIRTUAL TABLE {table} USING vec0(embedding float[2] distance_metric=cosine)") + self.conn.commit() + + def writer(self): + conn = super().writer() + conn.enable_load_extension(True) + self.sqlite_vec.load(conn) + conn.enable_load_extension(False) + return conn + + +class TestLegacyRebuildBridge(unittest.TestCase): + setUp = TestStreamingBuilders.setUp + + def legacy(self, *, minimal=False, unique_content=False, ids=(-3, 0, 2, 5)): + path = Path(self.temp.name) / f"synthetic-legacy-{len(getattr(self, 'legacy_paths', []))}.sqlite" + self.legacy_paths = [*getattr(self, "legacy_paths", []), path] + conn = sqlite3.connect(path, cached_statements=0) + self.addCleanup(conn.close) + fields = "id INTEGER PRIMARY KEY,content TEXT" + if unique_content: + fields += " UNIQUE" + if not minimal: + fields += ",sender TEXT,recipient TEXT,timestamp TEXT" + conn.execute(f"CREATE TABLE messages({fields})") + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + for mid in ids: + if minimal: + conn.execute("INSERT INTO messages VALUES(?,?)", (mid, f"synthetic-{mid}")) + else: + conn.execute("INSERT INTO messages VALUES(?,?,?,?,?)", (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02")) + conn.commit() + return conn, path + + def reopen(self, path): + conn = sqlite3.connect(path) + self.addCleanup(conn.close) + return conn + + def test_original_minimal_separation_shape_reaches_encoding_without_full_migration(self): + conn, _ = self.legacy(ids=range(1, 6)) + columns = conn.execute("PRAGMA table_info(messages)").fetchall() + self.assertEqual(self.vector.build_separation_vectors(conn, txn_batch=1), 5) + self.assertEqual(sum(self.calls, []), [f"sender to recipient on 2026-01-02: synthetic-{mid}" for mid in range(1, 6)]) + self.assertEqual(conn.execute("PRAGMA table_info(messages)").fetchall(), columns) + tables = {row[0] for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")} + self.assertEqual(tables, {"messages", "vec_messages", "vec_messages_sep", "metadata", self.source._BRIDGE_TABLE}) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + + def test_id_content_only_completion_preserves_negative_zero_and_hole_ids(self): + conn, _ = self.legacy(minimal=True) + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages ORDER BY rowid").fetchall(), [(-3,), (0,), (2,), (5,)]) + first = self.source.capture_source(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(first, self.source.capture_source(conn)) + self.assertEqual(len(conn.execute("PRAGMA table_info(messages)").fetchall()), 2) + + def test_bridge_install_and_successful_build_preserve_caller_transaction(self): + conn, _ = self.legacy() + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertTrue(conn.in_transaction) + self.assertEqual(conn.execute("SELECT value FROM synthetic_caller").fetchone(), ("pending",)) + conn.rollback() + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM synthetic_caller").fetchone()[0], 0) + + def test_bridge_resume_across_connection_uses_consumed_cursor(self): + conn, path = self.legacy() + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.assertEqual((before.cursor, before.consumed), (0, 2)) + self.on_encode = None + self.calls.clear() + reopened = self.reopen(path) + self.assertEqual(self.vector.build_vectors(reopened, txn_batch=1), 2) + after = self.source.load_manifest(reopened, key) + self.assertEqual(before.generation, after.generation) + self.assertEqual(sum(self.calls, []), ["synthetic-2", "synthetic-5"]) + + def test_bridge_rejects_concurrent_corrections_deletes_low_inserts_and_replace(self): + mutations = ("UPDATE messages SET content='synthetic-corrected' WHERE id=2", + "UPDATE messages SET sender='synthetic-new-sender' WHERE id=2", + "UPDATE messages SET rowid=3 WHERE id=2", "DELETE FROM messages WHERE id=2", + "INSERT INTO messages(id,content) VALUES(1,'synthetic-hole')", + "INSERT OR REPLACE INTO messages(id,content) VALUES(2,'synthetic-replacement')") + for statement in mutations: + with self.subTest(statement=statement): + conn, path = self.legacy() + other = self.reopen(path) + other.execute("PRAGMA recursive_triggers=OFF") + changed = [] + def encode(texts): + self.assertFalse(conn.in_transaction) + if not changed: + changed.append(True) + other.execute(statement) + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + + def test_bridge_allows_high_id_append_outside_captured_range(self): + conn, path = self.legacy() + other = self.reopen(path) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + other.execute("INSERT INTO messages(id,content) VALUES(9,'synthetic-tail')") + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages ORDER BY rowid").fetchall(), [(-3,), (0,), (2,), (5,)]) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual((manifest.source.high_id, manifest.consumed), (5, 4)) + manifest.source.check(conn) + + def test_bridge_noop_and_derived_only_writes_do_not_dirty_source(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages ADD COLUMN episode_id INTEGER") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content=content,episode_id=99") + conn.commit() + self.assertEqual(before.revision, self.source.capture_source(conn).revision) + + def test_missing_replaced_trigger_new_column_or_missing_state_rotates_epoch(self): + for fault in ("missing_trigger", "replaced_trigger", "new_column", "missing_state"): + with self.subTest(fault=fault): + conn, _ = self.legacy() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + trigger = self.source._BRIDGE_TRIGGERS[1] + if fault in ("missing_trigger", "replaced_trigger"): + conn.execute(f"DROP TRIGGER {trigger}") + if fault == "replaced_trigger": + conn.execute(f"CREATE TRIGGER {trigger} AFTER UPDATE ON messages BEGIN SELECT 1; END") + conn.execute("UPDATE messages SET content='synthetic-untracked' WHERE id=2") + elif fault == "new_column": + conn.execute("ALTER TABLE messages ADD COLUMN metadata TEXT") + else: + conn.execute(f"DELETE FROM {self.source._BRIDGE_TABLE}") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + self.source.ensure_rebuild_tracking(conn) + after = self.source.capture_source(conn) + self.assertNotEqual(before.revision.epoch, after.revision.epoch) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_rolled_back_bridge_source_change_does_not_invalidate(self): + conn, _ = self.legacy() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content='synthetic-rolled-back' WHERE id=2") + conn.rollback() + self.assertEqual(before, self.source.capture_source(conn)) + + def test_denied_setup_or_release_preserves_target_and_caller_data_before_model(self): + for fault in ("ddl", "release"): + with self.subTest(fault=fault): + conn, _ = self.legacy() + conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (99,?)", (struct.pack("2f", 1., 0.),)) + conn.execute("CREATE TABLE synthetic_caller(value TEXT)") + conn.commit() + conn.execute("INSERT INTO synthetic_caller VALUES('pending')") + denied = [] + def authorize(action, first, second, _database, _source): + is_target = ((fault == "ddl" and action == sqlite3.SQLITE_CREATE_TRIGGER) + or (fault == "release" and action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" + and second == "rebuild_source" and conn.execute( + "SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,), + ).fetchone() is not None)) + if is_target and not denied: + denied.append(True) + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + conn.set_authorizer(authorize) + before_calls = len(self.calls) + try: + with self.assertRaises(sqlite3.DatabaseError): + self.vector.build_vectors(conn, txn_batch=1) + finally: + conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertTrue(conn.in_transaction) + self.assertEqual(len(self.calls), before_calls) + conn.commit() + self.assertEqual(conn.execute("SELECT value FROM synthetic_caller").fetchone(), ("pending",)) + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages").fetchall(), [(99,)]) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_bridge_commit_failure_rolls_back_install_before_model_or_clear(self): + conn, _ = self.legacy() + conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (99,?)", (struct.pack("2f", 1., 0.),)) + conn.commit() + class Probe: + def __getattr__(self, name): + return getattr(conn, name) + def commit(self): + raise sqlite3.OperationalError("synthetic setup commit failure") + with self.assertRaisesRegex(sqlite3.OperationalError, "synthetic setup commit failure"): + self.vector.build_vectors(Probe(), txn_batch=1) + self.assertEqual(self.calls, []) + conn.commit() + self.assertEqual(conn.execute("SELECT rowid FROM vec_messages").fetchall(), [(99,)]) + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_readonly_legacy_install_fails_before_model_or_clear(self): + conn, path = self.legacy() + readonly = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + self.addCleanup(readonly.close) + with self.assertRaises(sqlite3.OperationalError): + self.vector.build_vectors(readonly, txn_batch=1) + self.assertEqual(self.calls, []) + self.assertFalse(readonly.in_transaction) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 4) + + def test_unsupported_source_schema_fails_before_model(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages RENAME TO synthetic_old_messages") + conn.execute("CREATE TABLE messages(id TEXT PRIMARY KEY,content TEXT)") + with self.assertRaisesRegex(self.modules["maintenance"].MaintenanceUnavailableError, "INTEGER PRIMARY KEY"): + self.vector.build_vectors(conn, txn_batch=1) + self.assertEqual(self.calls, []) + + def test_empty_legacy_generation_never_loads_a_model(self): + conn, _ = self.legacy(ids=()) + def no_model(): + self.fail("Empty legacy build loaded a model") + self.vector.get_model = no_model + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 0) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertTrue(manifest.complete) + self.assertEqual((manifest.source.tracker, manifest.total), ("rebuild-bridge-v1", 0)) + + def test_bridge_schema_change_during_encode_rejects_publication_and_restarts(self): + for fault in ("drop_trigger", "add_source_column"): + with self.subTest(fault=fault): + conn, path = self.legacy() + other = self.reopen(path) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + if fault == "drop_trigger": + other.execute(f"DROP TRIGGER {self.source._BRIDGE_TRIGGERS[1]}") + other.execute("UPDATE messages SET content='synthetic-untracked-correction' WHERE id=2") + else: + other.execute("ALTER TABLE messages ADD COLUMN metadata TEXT") + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + after = self.source.load_manifest(conn, key) + self.assertNotEqual(before.source.revision.epoch, after.source.revision.epoch) + self.assertNotEqual(before.generation, after.generation) + self.assertEqual(len(sum(self.calls, [])), 4) + + def test_canonical_handoff_does_not_resume_an_incomplete_bridge_prefix(self): + conn = self.conn + conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", + [(mid, f"synthetic-{mid}") for mid in (-3, 0, 2, 5)]) + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + conn.commit() + def encode(texts): + if len(self.calls) == 2: + raise RuntimeError("synthetic interruption") + return [[1., 0.] for _ in texts] + self.on_encode = encode + with self.assertRaisesRegex(RuntimeError, "synthetic interruption"): + self.vector.build_vectors(conn, txn_batch=1) + key = self.source.manifest_key(("vec_messages",)) + before = self.source.load_manifest(conn, key) + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + after = self.source.load_manifest(conn, key) + self.assertNotEqual(before.generation, after.generation) + self.assertEqual((before.source.tracker, after.source.tracker), ("rebuild-bridge-v1", "canonical-v1")) + self.assertEqual(sum(self.calls, []), ["synthetic--3", "synthetic-0", "synthetic-2", "synthetic-5"]) + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + self.assertEqual(conn.execute(f"SELECT tracking_ready FROM {self.source._BRIDGE_TABLE}").fetchone()[0], 0) + + def test_actual_deprecated_open_rebuilds_existing_legacy_connection(self): + import warnings + from unittest.mock import patch + conn, path = self.legacy() + tree = ast.parse((ROOT / "truememory/engine.py").read_text(encoding="utf-8")) + original = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "TrueMemoryEngine") + method = next(node for node in original.body if isinstance(node, ast.FunctionDef) and node.name == "open") + cls = ast.ClassDef(name="LegacyEngine", bases=[], keywords=[], body=[method], decorator_list=[]) + self.vector._check_embedder_compatibility = lambda _conn: None + self.vector.vectors_are_built = lambda _conn, _table: False + namespace = dict(__builtins__=self.vector.__dict__["__builtins__"], + warnings=warnings, sqlite3=sqlite3, DEFAULT_BUSY_TIMEOUT_MS=1000, + logger=logging.getLogger(__name__), _HAS_VECTOR=True, _HAS_STYLE_VEC=False, + _HAS_HYBRID=False, _HAS_CLUSTERING=False, resolve_tier=lambda: "base", + _check_rebuild_allowed=lambda _conn: None, init_vec_table=lambda _conn: None, + build_vectors=self.vector.build_vectors, + TrueMemoryMigrationError=type("SyntheticMigrationError", (RuntimeError,), {})) + exec(compile(ast.fix_missing_locations(ast.Module(body=[cls], type_ignores=[])), "actual-legacy-open", "exec"), namespace) + engine = namespace["LegacyEngine"]() + engine.db_path, engine.stats = path, {} + engine._purge_legacy_entity_profile_summaries = lambda: None + fake_vec = types.ModuleType("sqlite_vec") + fake_vec.load = lambda _conn: None + with patch.dict(sys.modules, {"sqlite_vec": fake_vec}), warnings.catch_warnings(record=True): + self.assertIs(engine.open(), engine) + self.addCleanup(engine.conn.close) + self.assertTrue(engine._has_vectors) + self.assertTrue(engine.ready) + self.assertEqual(engine.stats["message_count"], 4) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 4) + self.assertEqual(self.source.capture_source(engine.conn).tracker, "rebuild-bridge-v1") + self.assertIsNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name='maintenance_source_state'").fetchone()) + + def test_custom_unique_replacements_invalidate_conservative_bridge_even_with_canonical_ready(self): + for statement in ( + "INSERT OR REPLACE INTO messages(id,content,synthetic_unique) VALUES(9,'synthetic-high-replacement','slot-2')", + "UPDATE OR REPLACE messages SET synthetic_unique='slot-2' WHERE id=1", + ): + with self.subTest(statement=statement): + path = Path(self.temp.name) / ("synthetic-canonical-insert.sqlite" if statement.startswith("INSERT") + else "synthetic-canonical-update.sqlite") + conn = self.modules["storage"].create_db(path) + self.addCleanup(conn.close) + conn.execute("ALTER TABLE messages ADD COLUMN synthetic_unique TEXT") + conn.execute("CREATE UNIQUE INDEX synthetic_unique_key ON messages(synthetic_unique)") + conn.executemany("INSERT INTO messages(id,content,synthetic_unique) VALUES(?,?,?)", + [(1, "synthetic-first", "slot-1"), (2, "synthetic-second", "slot-2")]) + conn.commit() + other = self.reopen(path) + other.execute("PRAGMA recursive_triggers=OFF") + canonical_before = self.modules["maintenance"].read_source_revision(conn) + changed = [] + def encode(texts): + if not changed: + changed.append(True) + other.execute(statement) + other.commit() + return [[1., 0.] for _ in texts] + self.on_encode = encode + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + # Negative control: canonical counters alone do not see the + # hidden conflict deletion under these custom constraints. + canonical_after = self.modules["maintenance"].read_source_revision(conn) + self.assertTrue(canonical_after.is_append_only_since(canonical_before)) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual(manifest.source.tracker, "rebuild-bridge-conservative-v1") + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + + def test_expression_partial_unique_index_is_conservative_and_index_changes_rotate_epoch(self): + conn, _ = self.legacy() + conn.execute("ALTER TABLE messages ADD COLUMN synthetic_unique TEXT") + self.source.ensure_rebuild_tracking(conn) + normal = self.source.capture_source(conn) + conn.execute("CREATE UNIQUE INDEX synthetic_unique_expression ON messages(lower(synthetic_unique)) WHERE synthetic_unique IS NOT NULL") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + normal.check(conn) + self.source.ensure_rebuild_tracking(conn) + conservative = self.source.capture_source(conn) + self.assertEqual(conservative.tracker, "rebuild-bridge-conservative-v1") + self.assertNotEqual(normal.revision.epoch, conservative.revision.epoch) + conn.execute("INSERT INTO messages(id,content) VALUES(99,'synthetic-high-append')") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + conservative.check(conn) + + def test_canonical_handoff_retires_only_bridge_owned_triggers(self): + conn = self.conn + conn.execute("UPDATE maintenance_source_state SET tracking_ready=0") + conn.commit() + self.source.ensure_rebuild_tracking(conn) + conn.execute("CREATE TRIGGER synthetic_unrelated AFTER UPDATE ON messages BEGIN SELECT 1; END") + conn.commit() + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "canonical-v1") + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + self.assertIsNotNone(conn.execute("SELECT 1 FROM sqlite_master WHERE name='synthetic_unrelated'").fetchone()) + + def test_unowned_reserved_trigger_is_not_replaced_or_silently_adopted(self): + conn, _ = self.legacy() + trigger = self.source._BRIDGE_TRIGGERS[0] + conn.execute(f"CREATE TRIGGER {trigger} AFTER INSERT ON messages BEGIN SELECT 1; END") + original = conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (trigger,)).fetchone()[0] + with self.assertRaisesRegex(self.modules["maintenance"].MaintenanceUnavailableError, "already in use"): + self.vector.build_vectors(conn) + self.assertEqual(self.calls, []) + self.assertEqual(conn.execute("SELECT sql FROM sqlite_master WHERE name=?", (trigger,)).fetchone()[0], original) + + def test_schema_attestation_accepts_caller_row_factory(self): + conn, _ = self.legacy() + conn.row_factory = sqlite3.Row + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 4) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + + def test_unique_column_constraint_detects_hidden_replacement_deletion(self): + conn, _ = self.legacy(unique_content=True) + conn.execute("PRAGMA recursive_triggers=OFF") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("INSERT OR REPLACE INTO messages(id,content) VALUES(99,'synthetic-2')") + conn.commit() + self.assertIsNone(conn.execute("SELECT 1 FROM messages WHERE id=2").fetchone()) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0], 4) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_indexed_integer_primary_key_replacement_on_rowid_invalidates_prefix(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY DESC, content TEXT)") + conn.execute("INSERT INTO messages(rowid,id,content) VALUES(1,1,'synthetic-first')") + conn.commit() + conn.execute("PRAGMA recursive_triggers=OFF") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("INSERT OR REPLACE INTO messages(rowid,id,content) VALUES(1,9,'synthetic-replacement')") + conn.commit() + self.assertEqual(conn.execute("SELECT id FROM messages").fetchall(), [(9,)]) + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_removing_custom_unique_index_enables_safe_canonical_handoff(self): + conn = self.conn + conn.execute("CREATE UNIQUE INDEX synthetic_content_key ON messages(content)") + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + self.assertEqual(before.tracker, "rebuild-bridge-conservative-v1") + conn.execute("DROP INDEX synthetic_content_key") + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "canonical-v1") + self.assertEqual(self.source._installed_triggers(conn, self.source._BRIDGE_TRIGGERS), {}) + + def test_collation_equivalent_changes_reject_actual_builder_publication(self): + for canonical_ready in (False, True): + for collation, changed_text in (("NOCASE", "synthetic"), ("RTRIM", "Synthetic ")): + with self.subTest(canonical_ready=canonical_ready, collation=collation): + path = Path(self.temp.name) / f"synthetic-collation-{canonical_ready}-{collation}.sqlite" + conn = sqlite3.connect(path) + self.addCleanup(conn.close) + if canonical_ready: + schema = self.modules["storage"]._SCHEMA_SQL.replace( + "content TEXT NOT NULL,", f"content TEXT COLLATE {collation} NOT NULL,", 1, + ) + conn.executescript(schema) + self.modules["storage"]._initialize_maintenance_tracking(conn) + else: + conn.execute(f"CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT COLLATE {collation})") + conn.executescript("CREATE TABLE vec_messages(embedding BLOB); CREATE TABLE vec_messages_sep(embedding BLOB);") + conn.execute("INSERT INTO messages(id,content) VALUES(1,'Synthetic')") + conn.commit() + other = self.reopen(path) + if canonical_ready: + canonical_before = self.modules["maintenance"].read_source_revision(conn) + def encode(texts): + self.assertEqual(texts, ["Synthetic"]) + self.assertFalse(conn.in_transaction) + other.execute("UPDATE messages SET content=? WHERE id=1", (changed_text,)) + other.commit() + return [[1., 0.]] + self.on_encode = encode + with self.assertRaises(self.source.RebuildSourceChanged): + self.vector.build_vectors(conn, txn_batch=1) + if canonical_ready: + # These custom collations remain outside canonical + # tracking's guarantees even when it reports ready. + self.assertEqual(self.modules["maintenance"].read_source_revision(conn), canonical_before) + self.assertEqual(conn.execute("SELECT COUNT(*) FROM vec_messages").fetchone()[0], 0) + manifest = self.source.load_manifest(conn, self.source.manifest_key(("vec_messages",))) + self.assertEqual(manifest.source.tracker, "rebuild-bridge-v1") + self.assertEqual((manifest.complete, manifest.consumed), (False, 0)) + self.on_encode = None + self.calls.clear() + self.assertEqual(self.vector.build_vectors(conn, txn_batch=1), 1) + self.assertEqual(self.calls, [[changed_text]]) + + def test_bridge_compares_storage_type_when_blob_cast_bytes_match(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + conn.execute("CREATE TABLE messages(id INTEGER PRIMARY KEY, content)") + conn.execute("INSERT INTO messages VALUES(1,?)", ("synthetic",)) + conn.commit() + self.source.ensure_rebuild_tracking(conn) + before = self.source.capture_source(conn) + conn.execute("UPDATE messages SET content=? WHERE id=1", (b"synthetic",)) + conn.commit() + self.assertEqual(conn.execute("SELECT CAST(content AS BLOB) FROM messages").fetchone()[0], b"synthetic") + with self.assertRaises(self.source.RebuildSourceChanged): + before.check(conn) + + def test_standard_messages_definition_reuses_canonical_without_bridge(self): + self.source.ensure_rebuild_tracking(self.conn) + self.assertEqual(self.source.capture_source(self.conn).tracker, "canonical-v1") + self.assertIsNone(self.conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,)).fetchone()) + + def test_compatible_custom_definition_keeps_bridge_despite_ready_canonical(self): + conn = sqlite3.connect(":memory:") + self.addCleanup(conn.close) + schema = self.modules["storage"]._SCHEMA_SQL.replace("sender TEXT DEFAULT '',", "sender TEXT DEFAULT '' COLLATE BINARY,", 1) + conn.executescript(schema) + self.modules["storage"]._initialize_maintenance_tracking(conn) + self.source.ensure_rebuild_tracking(conn) + self.assertEqual(self.source.capture_source(conn).tracker, "rebuild-bridge-v1") + self.modules["maintenance"].read_source_revision(conn) + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/engine.py b/truememory/engine.py index 8f212b54..0285b066 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -726,9 +726,9 @@ def add( if self._has_vectors: try: from truememory.vector_search import ( - get_model, _encode_with_mps_fallback, _build_sep_text, + _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model = get_model() + model, pre_identity = _capture_rebuild_model() pre_embedding = _encode_with_mps_fallback(model, [content])[0] sep_text = _build_sep_text(sender, recipient, timestamp, content) pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] @@ -743,7 +743,13 @@ def add( except Exception: logger.debug("Failed to pre-compute style vector during add()", exc_info=True) - with self._write_lock: + from truememory.rebuild_source import rebuild_transaction + publication_fence = None + if pre_embedding is not None: + from truememory.vector_search import _foreground_model_fence + publication_fence = _foreground_model_fence(pre_identity) + + with self._write_lock, rebuild_transaction(self.conn, write=True, publication_fence=publication_fence): msg = { "content": content, "sender": sender, @@ -757,23 +763,25 @@ def add( new_id = insert_message(self.conn, msg) if pre_embedding is not None: + from truememory.vector_search import VectorPublicationChanged, _validate_foreground_vector_target try: from truememory.vector_search import ( - serialize_f32, _active_vec_table, _active_sep_table, - _write_embedder_metadata, + serialize_f32, + _write_embedder_metadata_no_commit, ) - vec_tbl = _active_vec_table(self.conn) + vec_tbl, sep_tbl = _validate_foreground_vector_target(self.conn, pre_identity, new_id) self.conn.execute( f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_embedding)), ) if pre_sep_embedding is not None: - sep_tbl = _active_sep_table(self.conn) self.conn.execute( f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (new_id, serialize_f32(pre_sep_embedding)), ) - _write_embedder_metadata(self.conn) + _write_embedder_metadata_no_commit(self.conn) + except VectorPublicationChanged: + raise except Exception: logger.warning("Failed to store embedding for message %s during add()", new_id, exc_info=True) @@ -795,8 +803,6 @@ def add( except Exception: logger.debug("Failed to update style vector for %s during add()", sender, exc_info=True) - self.conn.commit() - self._maybe_auto_consolidate() return { @@ -1172,9 +1178,9 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | if content is not None and self._has_vectors: try: from truememory.vector_search import ( - get_model, _encode_with_mps_fallback, _build_sep_text, + _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model = get_model() + model, pre_identity = _capture_rebuild_model() pre_embedding = _encode_with_mps_fallback(model, [content])[0] row = self.conn.execute( "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", @@ -1198,36 +1204,33 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | return None if pre_embedding is not None: + from truememory.vector_search import VectorPublicationChanged, _foreground_vector_publication try: - from truememory.vector_search import ( - serialize_f32, _active_vec_table, _active_sep_table, - _write_embedder_metadata, - ) - vec_tbl = _active_vec_table(self.conn) - try: - self.conn.execute(f"DELETE FROM {vec_tbl} WHERE rowid = ?", (memory_id,)) - except Exception: - logger.debug("Failed to delete old vector embedding for message %d", memory_id, exc_info=True) - self.conn.execute( - f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", - (memory_id, serialize_f32(pre_embedding)), - ) - if pre_sep_embedding is not None: - sep_tbl = _active_sep_table(self.conn) + from truememory.vector_search import serialize_f32, _write_embedder_metadata_no_commit + with _foreground_vector_publication(self.conn, pre_identity, memory_id) as (vec_tbl, sep_tbl): try: - self.conn.execute(f"DELETE FROM {sep_tbl} WHERE rowid = ?", (memory_id,)) + self.conn.execute(f"DELETE FROM {vec_tbl} WHERE rowid = ?", (memory_id,)) except Exception: - logger.debug("Failed to delete old sep vector for message %d", memory_id, exc_info=True) + logger.debug("Failed to delete old vector embedding for message %d", memory_id, exc_info=True) self.conn.execute( - f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", - (memory_id, serialize_f32(pre_sep_embedding)), + f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", + (memory_id, serialize_f32(pre_embedding)), ) - _write_embedder_metadata(self.conn) + if pre_sep_embedding is not None: + try: + self.conn.execute(f"DELETE FROM {sep_tbl} WHERE rowid = ?", (memory_id,)) + except Exception: + logger.debug("Failed to delete old sep vector for message %d", memory_id, exc_info=True) + self.conn.execute( + f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", + (memory_id, serialize_f32(pre_sep_embedding)), + ) + _write_embedder_metadata_no_commit(self.conn) + except VectorPublicationChanged as exc: + raise VectorPublicationChanged("Source update committed; retry update to publish vectors") from exc except Exception: logger.warning("Vector embedding failed for message %d", memory_id, exc_info=True) - self.conn.commit() - return self.get(memory_id) def get(self, memory_id: int) -> dict | None: diff --git a/truememory/rebuild_source.py b/truememory/rebuild_source.py new file mode 100644 index 00000000..e44bf2cf --- /dev/null +++ b/truememory/rebuild_source.py @@ -0,0 +1,375 @@ +"""Bounded rebuild inputs and durable, source-attested progress. + +Manifests contain identifiers and counters only. Page reads finish before the +caller computes embeddings; publication must revalidate under writer ownership. +""" + +import datetime +import hashlib +import json +import sqlite3 +import uuid +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import asdict, dataclass, replace + +from truememory.maintenance import MaintenanceUnavailableError, SourceRevision, read_source_revision +from truememory.storage import _MAINTENANCE_SOURCE_FIELDS, _MESSAGES_TABLE_SQL, _maintenance_trigger_definitions + +_BRIDGE_TABLE = "truememory_rebuild_source_v1" +_BRIDGE_OWNER = "truememory.rebuild-source.v1" +_BRIDGE_TRIGGERS = tuple(f"truememory_rebuild_messages_v1_{suffix}" for suffix in ("ai", "au", "ad")) +_BRIDGE_SQL = f"""CREATE TABLE {_BRIDGE_TABLE} ( + singleton INTEGER PRIMARY KEY CHECK (singleton = 1), + owner TEXT NOT NULL CHECK (owner = '{_BRIDGE_OWNER}'), + epoch TEXT NOT NULL, + revision INTEGER NOT NULL CHECK (typeof(revision) = 'integer' AND revision >= 0), + insert_count INTEGER NOT NULL CHECK (typeof(insert_count) = 'integer' AND insert_count >= 0), + correction_count INTEGER NOT NULL CHECK (typeof(correction_count) = 'integer' AND correction_count >= 0), + max_seen_message_id INTEGER NOT NULL, + nonappend_revision INTEGER NOT NULL, + source_signature TEXT NOT NULL, + covered_columns TEXT NOT NULL, + tracking_ready INTEGER NOT NULL CHECK (tracking_ready IN (0, 1)) +)""" + + +class RebuildSourceChanged(RuntimeError): + """The captured source or generation can no longer be published.""" + + +def _normalized_sql(sql: str) -> str: + return " ".join(sql.strip().rstrip(";").split()) + + +def _source_schema(conn: sqlite3.Connection) -> tuple[tuple[str, ...], str, bool]: + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'messages'").fetchone() + columns = [tuple(row) for row in conn.execute("PRAGMA table_xinfo(messages)")] + ordinary = {row[1]: row for row in columns if row[6] == 0} + present = {row[1] for row in columns} + identifier = ordinary.get("id") + if (table is None or not table[0] or table[0].upper().startswith("CREATE VIRTUAL TABLE") + or "content" not in present or identifier is None + or identifier[2].upper() != "INTEGER" or identifier[5] != 1 + or sum(row[5] > 0 for row in columns) != 1): + raise MaintenanceUnavailableError("Rebuild tracking requires messages with an INTEGER PRIMARY KEY id and content") + indexes = [] + conservative = False + for index in conn.execute("PRAGMA index_list(messages)"): + index = tuple(index) + name = index[1] + definition = conn.execute("SELECT sql FROM sqlite_master WHERE type='index' AND name=?", (name,)).fetchone() + quoted = '"' + name.replace('"', '""') + '"' + indexes.append([index, tuple(definition) if definition else None, + [tuple(row) for row in conn.execute(f"PRAGMA index_xinfo({quoted})")]]) + # An indexed INTEGER PRIMARY KEY (for example inline DESC) is not + # a rowid alias. REPLACE on its separate rowid can delete a prefix row. + conservative |= bool(index[2]) + covered = tuple(name for name in _MAINTENANCE_SOURCE_FIELDS if name in present) + signature = hashlib.sha256(json.dumps([table[0], columns, indexes], separators=(",", ":")).encode()).hexdigest() + return covered, signature, conservative + + +def _installed_triggers(conn: sqlite3.Connection, names: tuple[str, ...]) -> dict[str, str]: + placeholders = ",".join("?" for _ in names) + return {name: _normalized_sql(sql) for name, sql in conn.execute( + f"SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND name IN ({placeholders})", names, + )} + + +def _canonical_revision(conn: sqlite3.Connection, covered: tuple[str, ...], conservative: bool) -> SourceRevision | None: + if conservative or covered != _MAINTENANCE_SOURCE_FIELDS: + return None + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'messages'").fetchone() + # Canonical triggers use declared column comparisons. Only our trusted + # definition attests their byte-sensitive semantics; custom DDL stays on + # the bridge without trying to parse arbitrary collations or affinities. + trusted = _MESSAGES_TABLE_SQL.replace("CREATE TABLE IF NOT EXISTS ", "CREATE TABLE ", 1) + if table is None or _normalized_sql(table[0]) != _normalized_sql(trusted): + return None + required = {"epoch", "revision", "insert_count", "correction_count", "max_seen_message_id", + "nonappend_revision", "tracking_ready", "singleton"} + if not required.issubset({row[1] for row in conn.execute("PRAGMA table_info(maintenance_source_state)")}): + return None + definitions = _maintenance_trigger_definitions() + if _installed_triggers(conn, tuple(definitions)) != {name: _normalized_sql(sql) for name, sql in definitions.items()}: + return None + try: + return read_source_revision(conn) + except MaintenanceUnavailableError: + return None + + +def _bridge_definitions(covered: tuple[str, ...], conservative: bool) -> dict[str, str]: + changed = " OR ".join( + f"typeof(old.{name}) IS NOT typeof(new.{name}) " + f"OR CAST(old.{name} AS BLOB) IS NOT CAST(new.{name} AS BLOB)" + for name in covered + ) + update_guard = "" if conservative else f" WHEN {changed}" + insert_fence = ("revision + 1" if conservative else + "CASE WHEN new.id <= max_seen_message_id THEN revision + 1 ELSE nonappend_revision END") + insert, update, delete = _BRIDGE_TRIGGERS + return { + insert: f"""CREATE TRIGGER {insert} AFTER INSERT ON messages BEGIN + UPDATE {_BRIDGE_TABLE} SET + nonappend_revision = {insert_fence}, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, insert_count = insert_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + update: f"""CREATE TRIGGER {update} AFTER UPDATE ON messages{update_guard} BEGIN + UPDATE {_BRIDGE_TABLE} SET nonappend_revision = revision + 1, + max_seen_message_id = MAX(max_seen_message_id, new.id), + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + delete: f"""CREATE TRIGGER {delete} AFTER DELETE ON messages BEGIN + UPDATE {_BRIDGE_TABLE} SET nonappend_revision = revision + 1, + revision = revision + 1, correction_count = correction_count + 1 + WHERE singleton = 1 AND tracking_ready = 1; + END""", + } + + +def _bridge_row(conn: sqlite3.Connection) -> tuple | None: + table = conn.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() + if table is None: + return None + if _normalized_sql(table[0]) != _normalized_sql(_BRIDGE_SQL): + raise MaintenanceUnavailableError("Reserved rebuild tracking table has an unsupported definition") + return conn.execute( + f"SELECT epoch, revision, insert_count, correction_count, max_seen_message_id, nonappend_revision, " + f"source_signature, covered_columns, tracking_ready FROM {_BRIDGE_TABLE} WHERE singleton = 1 AND owner = ?", + (_BRIDGE_OWNER,), + ).fetchone() + + +def _valid_bridge(conn: sqlite3.Connection, covered: tuple[str, ...], signature: str, + conservative: bool) -> SourceRevision | None: + row = _bridge_row(conn) + definitions = _bridge_definitions(covered, conservative) + if (row is None or not row[8] or row[6] != signature or row[7] != json.dumps(covered) + or _installed_triggers(conn, _BRIDGE_TRIGGERS) != {name: _normalized_sql(sql) for name, sql in definitions.items()}): + return None + return SourceRevision(*row[:6]) + + +def read_rebuild_revision(conn: sqlite3.Connection) -> tuple[str, SourceRevision, str]: + """Read only; a scoped bridge never certifies unrelated maintenance layers.""" + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + if canonical is not None: + return "canonical-v1", canonical, signature + bridge = _valid_bridge(conn, covered, signature, conservative) + if bridge is None: + raise MaintenanceUnavailableError("Rebuild source tracking is unavailable or its schema identity changed") + return "rebuild-bridge-conservative-v1" if conservative else "rebuild-bridge-v1", bridge, signature + + +def ensure_rebuild_tracking(conn: sqlite3.Connection) -> None: + """Install only this rebuild's tracker, under writer ownership before models.""" + with rebuild_transaction(conn): + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + bridge_row = _bridge_row(conn) + installed = _installed_triggers(conn, _BRIDGE_TRIGGERS) + if canonical is not None and not installed and (bridge_row is None or not bridge_row[8]): + return + if canonical is None and _valid_bridge(conn, covered, signature, conservative) is not None: + return + with rebuild_transaction(conn, write=True): + covered, signature, conservative = _source_schema(conn) + canonical = _canonical_revision(conn, covered, conservative) + bridge_row = _bridge_row(conn) + installed = _installed_triggers(conn, _BRIDGE_TRIGGERS) + owned_table = conn.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() + if not owned_table and installed: + raise MaintenanceUnavailableError("Reserved rebuild trigger names are already in use without owned tracking state") + if canonical is not None: + if owned_table: + for name in installed: + conn.execute(f"DROP TRIGGER {name}") + conn.execute(f"UPDATE {_BRIDGE_TABLE} SET tracking_ready = 0 WHERE singleton = 1") + return + if _valid_bridge(conn, covered, signature, conservative) is not None: + return + if conn.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (_BRIDGE_TABLE,)).fetchone() is None: + conn.execute(_BRIDGE_SQL) + for name in installed: + conn.execute(f"DROP TRIGGER {name}") + for definition in _bridge_definitions(covered, conservative).values(): + conn.execute(definition) + # A repair cannot certify the interval before installation. A fresh + # epoch also prevents a later fallback from reviving a retired bridge. + conn.execute( + f"INSERT OR REPLACE INTO {_BRIDGE_TABLE} VALUES (1, ?, ?, 0, 0, 0, " + "(SELECT MAX(0, COALESCE(MAX(id), 0)) FROM messages), 0, ?, ?, 1)", + (_BRIDGE_OWNER, uuid.uuid4().hex, signature, json.dumps(covered)), + ) + + +@contextmanager +def rebuild_transaction(conn: sqlite3.Connection, *, write: bool = False, + publication_fence: AbstractContextManager | None = None) -> Iterator[None]: + owned = not conn.in_transaction + conn.execute(("BEGIN IMMEDIATE" if write else "BEGIN") if owned else "SAVEPOINT rebuild_source") + try: + if write and not owned: + conn.execute("UPDATE messages SET id = id WHERE 0") + with publication_fence if publication_fence is not None else nullcontext(): + yield + if owned: + if write: + conn.commit() + else: + conn.execute("COMMIT") + else: + conn.execute("RELEASE SAVEPOINT rebuild_source") + except BaseException: + if owned: + conn.rollback() + else: + conn.execute("ROLLBACK TO SAVEPOINT rebuild_source") + conn.execute("RELEASE SAVEPOINT rebuild_source") + raise + + +@dataclass(frozen=True) +class RebuildSource: + revision: SourceRevision + after_id: int | None + high_id: int | None + total: int + tracker: str = "canonical-v1" + schema_signature: str | None = None + + def check(self, conn: sqlite3.Connection) -> None: + try: + tracker, revision, signature = read_rebuild_revision(conn) + except MaintenanceUnavailableError as error: + raise RebuildSourceChanged("Rebuild source tracking changed; start a fresh generation") from error + if (tracker != self.tracker or signature != self.schema_signature + or not revision.is_append_only_since(self.revision)): + raise RebuildSourceChanged("Rebuild source changed; start a fresh generation") + + def page(self, conn: sqlite3.Connection, after_id: int | None, size: int, *, + include_metadata: bool = True) -> list[dict]: + if size < 1: + raise ValueError("Rebuild page size must be positive") + with rebuild_transaction(conn): + self.check(conn) + if self.high_id is None: + return [] + lower = " AND id > ?" if after_id is not None else "" + parameters = (self.high_id, after_id, size) if lower else (self.high_id, size) + columns = (("id", "content", "sender", "recipient", "timestamp") + if include_metadata else ("id", "content")) + cursor = conn.execute( + f"SELECT {', '.join(columns)} FROM messages WHERE id <= ?{lower} ORDER BY id LIMIT ?", parameters, + ) + try: + return [dict(zip(columns, row)) for row in cursor] + finally: + cursor.close() + + +def capture_source(conn: sqlite3.Connection, *, after_id: int | None = None) -> RebuildSource: + with rebuild_transaction(conn): + tracker, revision, signature = read_rebuild_revision(conn) + predicate = " WHERE id > ?" if after_id is not None else "" + count, high = conn.execute( + f"SELECT COUNT(*), MAX(id) FROM messages{predicate}", + (after_id,) if predicate else (), + ).fetchone() + return RebuildSource(revision, after_id, high, count, tracker, signature) + + +def input_fingerprint(messages: list[dict], *, separation: bool) -> str: + """Bind explicit caller input without copying its corpus or changing order.""" + digest = hashlib.sha256() + fields = ("id", "content", "sender", "recipient", "timestamp") if separation else ("id", "content") + for message in messages: + payload = json.dumps([message.get(field) for field in fields], ensure_ascii=False, + separators=(",", ":")).encode("utf-8") + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + return digest.hexdigest() + + +@dataclass(frozen=True) +class RebuildManifest: + generation: str + model: str + dimension: int + targets: tuple[str, ...] + source: RebuildSource | None + input_digest: str | None + total: int + cursor: int | None = None + consumed: int = 0 + outputs: int = 0 + complete: bool = False + schema_version: int | None = None + + @classmethod + def new(cls, model: str, dimension: int, targets: tuple[str, ...], *, + source: RebuildSource | None = None, input_digest: str | None = None, + total: int = 0) -> "RebuildManifest": + return cls(uuid.uuid4().hex, model, dimension, targets, source, input_digest, + source.total if source is not None else total, + source.after_id if source is not None else None) + + def advance(self, last_id: int, consumed: int, outputs: int) -> "RebuildManifest": + return replace(self, cursor=last_id, consumed=self.consumed + consumed, + outputs=self.outputs + outputs) + + def check(self, conn: sqlite3.Connection, key: str) -> None: + if self.schema_version != conn.execute("PRAGMA schema_version").fetchone()[0]: + raise RebuildSourceChanged("Rebuild database schema changed; start a fresh generation") + if self.source is not None: + self.source.check(conn) + saved = load_manifest(conn, key) + if saved != self: + raise RebuildSourceChanged("Rebuild generation changed; retry with its current owner") + + +def manifest_key(targets: tuple[str, ...]) -> str: + return "vec_source_v1:" + ":".join(targets) + + +def load_manifest(conn: sqlite3.Connection, key: str) -> RebuildManifest | None: + try: + row = conn.execute("SELECT value FROM metadata WHERE key = ?", (key,)).fetchone() + except sqlite3.OperationalError as error: + if str(error) == "no such table: metadata": + return None + raise + if row is None: + return None + try: + payload = json.loads(row[0]) + if not isinstance(payload, dict) or payload.pop("version") != 1: + return None + source = payload["source"] + if source is not None: + source["revision"] = SourceRevision(*source["revision"]) + payload["source"] = RebuildSource(**source) + payload["targets"] = tuple(payload["targets"]) + manifest = RebuildManifest(**payload) + if not 0 <= manifest.consumed <= manifest.total or manifest.outputs < 0: + return None + return manifest + except (TypeError, ValueError, KeyError): + return None + + +def save_manifest(conn: sqlite3.Connection, key: str, manifest: RebuildManifest) -> None: + conn.execute("CREATE TABLE IF NOT EXISTS metadata (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT)") + payload = asdict(manifest) + payload["version"] = 1 + conn.execute( + "INSERT OR REPLACE INTO metadata(key,value,updated_at) VALUES (?,?,?)", + (key, json.dumps(payload, separators=(",", ":")), + datetime.datetime.now(datetime.timezone.utc).isoformat()), + ) diff --git a/truememory/storage.py b/truememory/storage.py index d57f2798..061de4ad 100644 --- a/truememory/storage.py +++ b/truememory/storage.py @@ -42,9 +42,7 @@ END; """ -_SCHEMA_SQL = """ --- Core messages table -CREATE TABLE IF NOT EXISTS messages ( +_MESSAGES_TABLE_SQL = """CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, content TEXT NOT NULL, sender TEXT DEFAULT '', @@ -57,7 +55,9 @@ embedding_separation BLOB DEFAULT NULL, directive INTEGER DEFAULT 0, metadata TEXT DEFAULT '{}' -); +);""" + +_SCHEMA_SQL = "\n-- Core messages table\n" + _MESSAGES_TABLE_SQL + """ -- FTS5 virtual table for full-text search CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5( diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 537623d1..507e4bb9 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -45,6 +45,9 @@ import struct import sqlite3 import threading +from collections.abc import Iterator +from contextlib import closing, contextmanager +from dataclasses import replace from typing import TYPE_CHECKING import numpy as np @@ -163,6 +166,7 @@ def resolve_tier() -> str: _model = None _embedding_dim: int = _MODEL_DIMS.get(EMBEDDING_MODEL, 256) _lock = threading.Lock() +_model_generation = 0 def _active_tier_group() -> str: @@ -205,8 +209,9 @@ def set_embedding_model(name: str) -> None: Accepts tier names ("base", "pro") or internal model names. """ - global EMBEDDING_MODEL, _model, _embedding_dim + global EMBEDDING_MODEL, _model, _embedding_dim, _model_generation with _lock: + _model_generation += 1 _model = None # Force reload EMBEDDING_MODEL = _resolve_model_name(name) _embedding_dim = get_embedding_dim(EMBEDDING_MODEL) @@ -223,8 +228,9 @@ def get_embedding_dim(name: str | None = None) -> int: def unload_model() -> None: """Release the embedding model from memory.""" - global _model + global _model, _model_generation with _lock: + _model_generation += 1 _model = None @@ -728,12 +734,208 @@ def _encode_with_mps_fallback(model, texts, **kwargs): """Number of embedding batches to accumulate before committing. Each embedding batch is ``_get_batch_size()`` messages (100 on CPU, 16 on -MPS). Every ``_BUILD_VECTORS_TXN_BATCH`` of those batches we commit the -transaction and release the DB write lock so other writers (MCP servers, -ingest pipeline) can proceed. +MPS). Serialize a finite group before acquiring an owned DB writer transaction; +publish the group and consumed-input checkpoint in one commit. """ +def _capture_rebuild_model() -> tuple[object, tuple[int, str, int]]: + with _lock: + before = (_model_generation, EMBEDDING_MODEL) + model = get_model() + with _lock: + if before != (_model_generation, EMBEDDING_MODEL): + raise RuntimeError("Embedding model changed while loading rebuild model") + return model, (_model_generation, EMBEDDING_MODEL, _embedding_dim) + + +@contextmanager +def _rebuild_model_fence(identity: tuple[int, str, int]) -> Iterator[None]: + # A writer must never wait for a model load that needs this state lock. + if not _lock.acquire(blocking=False): + raise RuntimeError("Embedding model is busy; retry rebuild publication") + try: + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise RuntimeError("Embedding model changed during rebuild") + yield + finally: + _lock.release() + + +class VectorPublicationChanged(RuntimeError): + """Precomputed vectors no longer match the active publication target.""" + + +@contextmanager +def _foreground_model_fence(identity: tuple[int, str, int]) -> Iterator[None]: + if not _lock.acquire(blocking=False): + raise VectorPublicationChanged("Embedding model is busy; retry vector publication") + try: + if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): + raise VectorPublicationChanged("Embedding model changed; retry vector publication") + yield + finally: + _lock.release() + + +def _validate_foreground_vector_target( + conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, +) -> tuple[str, str]: + """Validate while the caller owns both the database writer and model fence.""" + from truememory.rebuild_source import RebuildSourceChanged, load_manifest, manifest_key + + tables = (_active_vec_table(conn), _active_sep_table(conn)) + for table in tables: + manifest = load_manifest(conn, manifest_key((table,))) + if manifest is not None and (manifest.model != identity[1] or manifest.dimension != identity[2]): + raise VectorPublicationChanged("Vector target uses a different embedding model; reopen with the active tier") + if not _build_in_progress(conn, table): + continue + if manifest is None: + raise VectorPublicationChanged("Vector target is rebuilding without matching source identity") + if manifest.source is None: + # Explicit caller input has no database prefix certificate. + raise VectorPublicationChanged("Vector target is rebuilding explicit input; retry after completion") + high = manifest.source.high_id + if high is None or message_id > high: + continue + try: + manifest.source.check(conn) + except RebuildSourceChanged: + # A corrected source invalidates this generation. Publish the + # correction normally; the rebuild cannot later certify it. + continue + raise VectorPublicationChanged("Vector row belongs to an active rebuild; retry after completion") + return tables + + +@contextmanager +def _foreground_vector_publication( + conn: sqlite3.Connection, identity: tuple[int, str, int], message_id: int, +) -> Iterator[tuple[str, str]]: + from truememory.rebuild_source import rebuild_transaction + + with rebuild_transaction(conn, write=True, publication_fence=_foreground_model_fence(identity)): + yield _validate_foreground_vector_target(conn, identity, message_id) + + +def _build_streamed_vectors( + conn: sqlite3.Connection, messages: list[dict] | None, *, table_name: str | None, + txn_batch: int | None, separation: bool, +) -> int: + from truememory.maintenance import connection_database_path, maintenance_owner + from truememory.rebuild_source import ( + RebuildManifest, RebuildSourceChanged, capture_source, input_fingerprint, + ensure_rebuild_tracking, load_manifest, manifest_key, rebuild_transaction, save_manifest, + ) + + if messages is not None and not messages: + return 0 + table = table_name or (_active_sep_table(conn) if separation else _active_vec_table(conn)) + quoted_table = '"' + table.replace('"', '""') + '"' + key = manifest_key((table,)) + digest = input_fingerprint(messages, separation=separation) if messages is not None else None + with maintenance_owner(connection_database_path(conn)): + source_empty = False + if messages is None: + ensure_rebuild_tracking(conn) + with closing(conn.execute("SELECT 1 FROM messages LIMIT 1")) as cursor: + source_empty = cursor.fetchone() is None + if source_empty: + identity = (_model_generation, EMBEDDING_MODEL, _embedding_dim) + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + empty_source = capture_source(conn) + if empty_source.total == 0: + # An empty range needs no native model allocation. Recheck + # under the writer so an append cannot slip through clear. + conn.execute(f"DELETE FROM {quoted_table}") + _ensure_metadata_table(conn) + empty = RebuildManifest.new(identity[1], identity[2], (table,), source=empty_source) + empty = replace(empty, complete=True, + schema_version=conn.execute("PRAGMA schema_version").fetchone()[0]) + save_manifest(conn, key, empty) + _clear_build_in_progress(conn, table) + return 0 + model, identity = _capture_rebuild_model() + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + manifest = load_manifest(conn, key) + resumable = (manifest is not None and not manifest.complete + and manifest.model == identity[1] and manifest.dimension == identity[2] + and manifest.targets == (table,) and manifest.input_digest == digest + and manifest.schema_version == conn.execute("PRAGMA schema_version").fetchone()[0] + and (manifest.source is not None) == (messages is None)) + if resumable and manifest.source is not None: + try: + manifest.source.check(conn) + except RebuildSourceChanged: + resumable = False + if not resumable: + source = capture_source(conn) if messages is None else None + manifest = RebuildManifest.new(identity[1], identity[2], (table,), source=source, + input_digest=digest, total=len(messages) if messages is not None else 0) + conn.execute(f"DELETE FROM {quoted_table}") + _mark_build_in_progress(conn, table) + manifest = replace(manifest, schema_version=conn.execute("PRAGMA schema_version").fetchone()[0]) + save_manifest(conn, key, manifest) + + batch_size = _get_batch_size() + commit_every = max(1, txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH) + total = 0 + try: + import torch + no_grad = torch.no_grad() + except ImportError: + from contextlib import nullcontext + no_grad = nullcontext() + + with no_grad: + while manifest.consumed < manifest.total: + next_manifest = manifest + rows_to_insert = [] + for _ in range(commit_every): + if next_manifest.consumed >= next_manifest.total: + break + if manifest.source is not None: + batch = manifest.source.page(conn, next_manifest.cursor, batch_size, include_metadata=separation) + else: + batch = messages[next_manifest.consumed:next_manifest.consumed + batch_size] + if not batch: + raise RebuildSourceChanged("Rebuild source ended before its captured row count") + ids = [message["id"] for message in batch] + texts = ([_build_sep_text(message.get("sender", "?"), message.get("recipient", "?"), + message.get("timestamp", "?"), message["content"]) for message in batch] + if separation else [message["content"] for message in batch]) + embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) + if len(embeddings) != len(ids): + raise ValueError("Embedding count does not match the rebuild batch") + valid = 0 + for mid, embedding in zip(ids, embeddings): + normed = _normalize_for_cosine(embedding) + if normed is None: + logger.warning("Skipping undefined cosine embedding during rebuild") + continue + rows_to_insert.append((mid, serialize_f32(normed))) + valid += 1 + next_manifest = next_manifest.advance(ids[-1], len(batch), valid) + del embeddings, batch, texts, ids + _flush_mps_cache() + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + manifest.check(conn, key) + if rows_to_insert: + conn.executemany(f"INSERT INTO {quoted_table}(rowid, embedding) VALUES (?, ?)", rows_to_insert) + save_manifest(conn, key, next_manifest) + total += next_manifest.outputs - manifest.outputs + manifest = next_manifest + del rows_to_insert + + with rebuild_transaction(conn, write=True, publication_fence=_rebuild_model_fence(identity)): + manifest.check(conn, key) + save_manifest(conn, key, replace(manifest, complete=True)) + _clear_build_in_progress(conn, table) + _write_embedder_metadata_no_commit(conn) + return total + + def build_vectors( conn: sqlite3.Connection, messages: list[dict] | None = None, @@ -748,14 +950,13 @@ def build_vectors( ``messages`` table. Otherwise it uses the supplied list (each dict must have an ``"id"`` and ``"content"`` key). - Batch size adapts to the device: 100 on CPU (Edge tier), 16 on MPS - (Base/Pro tier) to keep GPU memory under ~4GB on 8GB machines. - MPS cache is flushed between batches to prevent memory accumulation. + Database inputs are read in bounded ID pages. A source revision and model + fence protect every publication. Resume uses consumed input, including + skipped invalid vectors, rather than the maximum stored vector ID. - The write transaction is committed every *txn_batch* embedding batches - (default ``_BUILD_VECTORS_TXN_BATCH`` = 100) so that the DB write lock - is not held for the entire re-embedding run. This lets other writers - (MCP servers, ingest pipeline) make progress between commits. + Serialize *txn_batch* embedding batches (default 100) before acquiring an + owned writer transaction. Caller-owned transactions remain caller-owned; + their existing locks cannot be released by this function. Args: conn: Open database connection with sqlite-vec already loaded @@ -769,102 +970,9 @@ def build_vectors( Returns: Number of vectors inserted. """ - if messages is None: - rows = conn.execute( - "SELECT id, content FROM messages ORDER BY id" - ).fetchall() - messages = [{"id": row[0], "content": row[1]} for row in rows] - - if not messages: - return 0 - - tbl = table_name or _active_vec_table(conn) - - # Resume vs. fresh build (issue #647, M-21). - # - # If a prior run was interrupted it left an ``in_progress`` marker. The - # rows already written are valid (each is L2-normalized at write time), so - # we resume from ``max(rowid)`` instead of re-wiping — no data confusion, - # no redundant re-embedding. On a fresh build we wipe and set the marker in - # the SAME transaction, so a crash before completion never leaves an - # untracked partial/empty table that engine.open() would trust. - resume_after = 0 - if _build_in_progress(conn, tbl): - try: - resume_after = conn.execute( - f"SELECT MAX(rowid) FROM {tbl}" - ).fetchone()[0] or 0 - except sqlite3.OperationalError: - resume_after = 0 - if resume_after > 0: - messages = [m for m in messages if m["id"] > resume_after] - if not messages: - # Everything was already embedded before the interruption; just - # finalize (clear marker + write metadata) so the table is trusted. - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata(conn) - return 0 - else: - conn.execute(f"DELETE FROM {tbl}") - _mark_build_in_progress(conn, tbl) - conn.commit() # DELETE + marker land atomically - - model = get_model() - batch_size = _get_batch_size() - commit_every = txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH - total = 0 - batches_since_commit = 0 + return _build_streamed_vectors(conn, messages, table_name=table_name, txn_batch=txn_batch, + separation=False) - try: - import torch - no_grad = torch.no_grad() - except ImportError: - from contextlib import nullcontext - no_grad = nullcontext() - - with no_grad: - for start in range(0, len(messages), batch_size): - batch = messages[start : start + batch_size] - texts = [m["content"] for m in batch] - ids = [m["id"] for m in batch] - - embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) - - rows_to_insert = [] - for mid, emb in zip(ids, embeddings): - normed = _normalize_for_cosine(emb) - if normed is None: - # Zero/NaN vector: undefined cosine — skip so it can't - # poison the cosine table with NULL distances (C2-8). - logger.warning( - "Skipping message %s: zero-norm/non-finite embedding " - "(undefined cosine)", mid, - ) - continue - rows_to_insert.append((mid, serialize_f32(normed))) - if rows_to_insert: - conn.executemany( - f"INSERT INTO {tbl}(rowid, embedding) VALUES (?, ?)", - rows_to_insert, - ) - - total += len(rows_to_insert) - batches_since_commit += 1 - del embeddings - _flush_mps_cache() - - if batches_since_commit >= commit_every: - conn.commit() - batches_since_commit = 0 - - # Final commit: clear the in-progress marker and stamp embedder metadata in - # the SAME transaction so the table is only ever "trusted" once it is fully - # built under the current embedder (issue #647: no NEW vectors under OLD - # metadata, no empty/partial table accepted as built). - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata_no_commit(conn) - conn.commit() - return total # --------------------------------------------------------------------------- @@ -1058,102 +1166,8 @@ def build_separation_vectors( Returns: Number of separation vectors inserted. """ - if messages is None: - rows = conn.execute( - "SELECT id, content, sender, recipient, timestamp FROM messages ORDER BY id" - ).fetchall() - messages = [ - {"id": r[0], "content": r[1], "sender": r[2], "recipient": r[3], "timestamp": r[4]} - for r in rows - ] - - if not messages: - return 0 - - tbl = table_name or _active_sep_table(conn) - - # Resume vs. fresh build — same durability contract as build_vectors - # (issue #647, M-21 + M-45). Marker + DELETE land atomically; batched - # commits below release the write lock during the run (#619 treatment). - resume_after = 0 - if _build_in_progress(conn, tbl): - try: - resume_after = conn.execute( - f"SELECT MAX(rowid) FROM {tbl}" - ).fetchone()[0] or 0 - except sqlite3.OperationalError: - resume_after = 0 - if resume_after > 0: - messages = [m for m in messages if m["id"] > resume_after] - if not messages: - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata(conn) - return 0 - else: - conn.execute(f"DELETE FROM {tbl}") - _mark_build_in_progress(conn, tbl) - conn.commit() # DELETE + marker land atomically - - model = get_model() - batch_size = _get_batch_size() - commit_every = txn_batch if txn_batch is not None else _BUILD_VECTORS_TXN_BATCH - total = 0 - batches_since_commit = 0 - - try: - import torch - no_grad = torch.no_grad() - except ImportError: - from contextlib import nullcontext - no_grad = nullcontext() - - with no_grad: - for start in range(0, len(messages), batch_size): - batch = messages[start : start + batch_size] - - texts = [] - ids = [] - for m in batch: - sep_text = _build_sep_text( - m.get("sender", "?"), m.get("recipient", "?"), - m.get("timestamp", "?"), m["content"], - ) - texts.append(sep_text) - ids.append(m["id"]) - - embeddings = _encode_with_mps_fallback(model, texts, show_progress_bar=False) - - rows_to_insert = [] - for mid, emb in zip(ids, embeddings): - normed = _normalize_for_cosine(emb) - if normed is None: - logger.warning( - "Skipping separation vector for message %s: " - "zero-norm/non-finite embedding (undefined cosine)", mid, - ) - continue - rows_to_insert.append((mid, serialize_f32(normed))) - if rows_to_insert: - conn.executemany( - f"INSERT INTO {tbl}(rowid, embedding) VALUES (?, ?)", - rows_to_insert, - ) - - total += len(rows_to_insert) - batches_since_commit += 1 - del embeddings - _flush_mps_cache() - - if batches_since_commit >= commit_every: - conn.commit() - batches_since_commit = 0 - - # Final commit: clear the in-progress marker + stamp embedder metadata - # atomically (issue #647). - _clear_build_in_progress(conn, tbl) - _write_embedder_metadata_no_commit(conn) - conn.commit() - return total + return _build_streamed_vectors(conn, messages, table_name=table_name, txn_batch=txn_batch, + separation=True) def _build_sep_text(sender: str, recipient: str, timestamp: str, content: str) -> str: @@ -1177,48 +1191,34 @@ def embed_single(conn: sqlite3.Connection, message_id: int, content: str) -> Non message_id: The ``messages.id`` of the row being embedded. content: The text to embed. """ - model = get_model() - embedding = _encode_with_mps_fallback(model, [content])[0] # shape (dim,) + model, identity = _capture_rebuild_model() + embedding = _encode_with_mps_fallback(model, [content])[0] normed = _normalize_for_cosine(embedding) if normed is None: - # Zero/NaN vector has undefined cosine — skip rather than poison the - # cosine table with NULL distances (C2-8). - logger.warning( - "Skipping message %d: zero-norm/non-finite embedding " - "(undefined cosine)", message_id, - ) + logger.warning("Skipping undefined cosine embedding during incremental publication") return - vec_tbl = _active_vec_table(conn) - conn.execute( - f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", - (message_id, serialize_f32(normed)), - ) - + completion = serialize_f32(normed) + separation = None try: row = conn.execute( - "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", - (message_id,), + "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", (message_id,), ).fetchone() if row: sep_text = _build_sep_text(row[0], row[1], row[2], content) sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] sep_normed = _normalize_for_cosine(sep_embedding) - if sep_normed is None: - logger.warning( - "Skipping separation vector for message %d: " - "zero-norm/non-finite embedding (undefined cosine)", - message_id, - ) - return - sep_tbl = _active_sep_table(conn) - conn.execute( - f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", - (message_id, serialize_f32(sep_normed)), - ) + if sep_normed is not None: + separation = serialize_f32(sep_normed) except Exception: - logger.warning("Failed to create separation vector for message %d", message_id, exc_info=True) - _write_embedder_metadata(conn) - # Caller is responsible for committing + logger.warning("Failed to prepare separation vector during incremental publication", exc_info=True) + with _foreground_vector_publication(conn, identity, message_id) as (vec_tbl, sep_tbl): + conn.execute(f"INSERT INTO {vec_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, completion)) + if separation is not None: + try: + conn.execute(f"INSERT INTO {sep_tbl}(rowid, embedding) VALUES (?, ?)", (message_id, separation)) + except Exception: + logger.warning("Failed to create separation vector during incremental publication", exc_info=True) + _write_embedder_metadata_no_commit(conn) def search_vector_separation( From 6fc5def689e860dfa072f4f83edbd78497dcaabe Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:39:10 -0500 Subject: [PATCH 3/3] test: Make rebuild release fault injection portable --- tests/test-rebuild-source-streaming-756.py | 41 ++++++++++++++++------ 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py index 48edfba4..369304e5 100644 --- a/tests/test-rebuild-source-streaming-756.py +++ b/tests/test-rebuild-source-streaming-756.py @@ -607,10 +607,10 @@ def writer(self): class TestLegacyRebuildBridge(unittest.TestCase): setUp = TestStreamingBuilders.setUp - def legacy(self, *, minimal=False, unique_content=False, ids=(-3, 0, 2, 5)): + def legacy(self, *, minimal=False, unique_content=False, ids=(-3, 0, 2, 5), cached_statements=0): path = Path(self.temp.name) / f"synthetic-legacy-{len(getattr(self, 'legacy_paths', []))}.sqlite" self.legacy_paths = [*getattr(self, "legacy_paths", []), path] - conn = sqlite3.connect(path, cached_statements=0) + conn = sqlite3.connect(path, cached_statements=cached_statements) self.addCleanup(conn.close) fields = "id INTEGER PRIMARY KEY,content TEXT" if unique_content: @@ -767,9 +767,10 @@ def test_rolled_back_bridge_source_change_does_not_invalidate(self): self.assertEqual(before, self.source.capture_source(conn)) def test_denied_setup_or_release_preserves_target_and_caller_data_before_model(self): - for fault in ("ddl", "release"): - with self.subTest(fault=fault): - conn, _ = self.legacy() + # Python 3.10 keeps at least five cached statements even when given zero. + for fault, cached_statements in ((fault, size) for fault in ("ddl", "release") for size in (0, 5, 100)): + with self.subTest(fault=fault, cached_statements=cached_statements): + conn, _ = self.legacy(cached_statements=cached_statements) conn.execute("INSERT INTO vec_messages(rowid,embedding) VALUES (99,?)", (struct.pack("2f", 1., 0.),)) conn.execute("CREATE TABLE synthetic_caller(value TEXT)") conn.commit() @@ -778,20 +779,40 @@ def test_denied_setup_or_release_preserves_target_and_caller_data_before_model(s def authorize(action, first, second, _database, _source): is_target = ((fault == "ddl" and action == sqlite3.SQLITE_CREATE_TRIGGER) or (fault == "release" and action == sqlite3.SQLITE_SAVEPOINT and first == "RELEASE" - and second == "rebuild_source" and conn.execute( - "SELECT 1 FROM sqlite_master WHERE name=?", (self.source._BRIDGE_TABLE,), - ).fetchone() is not None)) + and second == "rebuild_source")) if is_target and not denied: denied.append(True) return sqlite3.SQLITE_DENY return sqlite3.SQLITE_OK - conn.set_authorizer(authorize) + bridge_table = self.source._BRIDGE_TABLE + armed = [] + + class Probe: + def __getattr__(self, name: str) -> object: + return getattr(conn, name) + + def execute(self, sql: str, *args: object) -> sqlite3.Cursor: + if fault == "release" and sql == "RELEASE SAVEPOINT rebuild_source" and not armed: + installed = conn.execute( + "SELECT 1 FROM sqlite_master WHERE name=?", (bridge_table,), + ).fetchone() is not None + if installed: + # Authorizers run at prepare time. Re-arm outside + # the callback so cached RELEASE is reauthorized. + conn.set_authorizer(authorize) + armed.append(True) + return conn.execute(sql, *args) + + if fault == "ddl": + conn.set_authorizer(authorize) before_calls = len(self.calls) try: with self.assertRaises(sqlite3.DatabaseError): - self.vector.build_vectors(conn, txn_batch=1) + self.vector.build_vectors(Probe(), txn_batch=1) finally: conn.set_authorizer(lambda *_args: sqlite3.SQLITE_OK) + self.assertEqual(denied, [True]) + self.assertEqual(armed, [True] if fault == "release" else []) self.assertTrue(conn.in_transaction) self.assertEqual(len(self.calls), before_calls) conn.commit()