diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index a768579e..4bb389bd 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -1506,3 +1506,47 @@ retain their existing savepoint and caller-transaction semantics. This bounds repeated nonprogress, not one native call's duration, allocator bytes, the complete source list retained by the worker, or whole-process memory. Native Mac/MPS calibration and the original memory incident remain unresolved. + +### Tier source adapter: bounded checkpoints for an inactive pair + +`tier_switch/source.py` supplies model-free plan, initialize, read, publish-prefix, +and finish operations. It is not yet called by the tier manager or worker, so it +does not remove their retained source list or change serving activation. Callers +must own an inactive completion/separation pair, validate its exact model and +schema beforehand, install source tracking and metadata before planning, and +hold the canonical maintenance lease plus exclusive cooperating pair-writer +ownership. Database identity on background reopen and snapshot consistency are +caller obligations. The adapter does not establish these preconditions. + +Full initialization accepts only an already empty pair and never clears tables. +A registry row, matching row count, generic manifest, or schema alone cannot +certify existing vector contents. Resume and delta require both a strict source +manifest and this adapter's receipt. The receipt binds the model, dimension, +target names, schema generation, consumed source range, and an ordered SHA256 +chain over signed 64-bit row IDs and both serialized vector blobs. Restart +recomputes that chain once in pages of 64 rows. Dimensions are limited to 1..4096; +each stored blob must contain exactly `4 * dimension` bytes. Two vector pages +therefore contain at most `2 * 64 * 4096 * 4 = 2,097,152` bytes of serialized +payload, excluding Python objects, SQLite buffers and transient hash input. +Receipts detect changed bytes under cooperating ownership; they do not authenticate +data against a writer that can forge both metadata and vectors. + +Restart verification takes O(N) time in one read snapshot and can retain WAL +pages until that snapshot ends. Source pages default to 64 rows and honor the +explicit caller row limit. A row limit is not a text-byte bound: one message can +still be large. Reads finish before encoding. Every encoded prefix publishes both +target rows, the consumed cursor and the receipt in one owned transaction. The +returned immutable pending suffix supports shrinking retries without skipping +inputs. Failed database writes require replanning because SQLite total_changes +also counts rolled-back writes. Borrowed transactions are refused untouched. + +The plan pins its connection, data_version and total_changes. Any intervening +write, including a legitimate concurrent tail append or unrelated status write, +requires explicit replanning and bounded restart verification. TEMP shadows and +side-effect triggers on pair/metadata tables are refused. Source corrections, +deletions and low-ID insertions invalidate captured provenance; a separately +prepared empty pair is then required. Negative and zero IDs remain valid cursors. +Finish distinguishes complete captured-range coverage from whole-source coverage +at its own transaction snapshot. Neither result activates a target or promises +future freshness. Final concurrent-tail caller adoption, native memory budgets, +and the original Mac memory incident remain unresolved. diff --git a/tests/test-tier-rebuild-source-756.py b/tests/test-tier-rebuild-source-756.py new file mode 100644 index 00000000..883ae16b --- /dev/null +++ b/tests/test-tier-rebuild-source-756.py @@ -0,0 +1,524 @@ +"""Synthetic in-memory source/pair lifecycle; no package or model initialization.""" + +import builtins +import json +import os +import sqlite3 +import struct +import sys +import types +import unittest +from dataclasses import replace +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +TARGETS = ("vec_messages_synthetic", "vec_messages_sep_synthetic") + + +def load_modules(): + modules = {} + + def safe_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in {"numpy", "torch", "sentence_transformers", "sqlite_vec"}: + raise AssertionError("Native/model imports are excluded from the adapter") + 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", "tier_switch.source"): + module = types.ModuleType("synthetic_tier_source_" + name.replace(".", "_")) + sys.modules[module.__name__] = module + module.__dict__["__builtins__"] = dict(vars(builtins), __import__=safe_import) + path = ROOT / "truememory" / (name.replace(".", "/") + ".py") + exec(compile(path.read_text(encoding="utf-8"), str(path), "exec"), module.__dict__) + modules[name] = module + return modules + + +class FaultConnection(sqlite3.Connection): + fail_sql = None + fail_commit = False + + def execute(self, sql, parameters=()): + if self.fail_sql is not None and self.fail_sql in sql: + self.fail_sql = None + raise sqlite3.OperationalError("synthetic pair write failure") + return super().execute(sql, parameters) + + def commit(self): + if self.fail_commit: + self.fail_commit = False + raise sqlite3.OperationalError("synthetic commit failure") + return super().commit() + + +class TestTierSource(unittest.TestCase): + def setUp(self): + self.modules = load_modules() + self.api = self.modules["tier_switch.source"] + self.source = self.modules["rebuild_source"] + self.conn = sqlite3.connect(":memory:", factory=FaultConnection) + self.addCleanup(self.conn.close) + self.create_schema(self.conn) + + def create_schema(self, conn): + conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT); + CREATE TABLE vec_messages_synthetic(embedding BLOB); + CREATE TABLE vec_messages_sep_synthetic(embedding BLOB); + """) + self.source.ensure_rebuild_tracking(conn) + + def seed(self, ids=(-3, 0, 2, 8)): + self.conn.executemany("INSERT INTO messages VALUES (?,?,?,?,?)", [ + (mid, f"synthetic-{mid}", "sender", "recipient", "2026-01-02") for mid in ids + ]) + self.conn.commit() + + def plan(self, **kwargs): + return self.api.plan_tier_source(self.conn, model_id="synthetic-model", dimension=2, targets=TARGETS, **kwargs) + + def start(self): + return self.api.initialize_tier_source(self.conn, self.plan()) + + def vectors(self, count): + return [struct.pack("2f", 1.0, 0.0)] * count, [struct.pack("2f", 0.0, 1.0)] * count + + def publish(self, plan, page, count=None): + left, right = self.vectors(len(page.rows) if count is None else count) + return self.api.publish_tier_prefix(self.conn, plan, page, completion=left, separation=right) + + def build(self, plan=None): + plan = self.start() if plan is None else plan + while plan.manifest.consumed < plan.manifest.total: + page = self.api.read_tier_page(self.conn, plan, size=2) + plan, pending = self.publish(plan, page) + self.assertEqual(pending.rows, ()) + return self.api.finish_tier_source(self.conn, plan) + + def ids(self, table=TARGETS[0]): + return [row[0] for row in self.conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')] + + def test_plan_is_read_only_and_initialization_never_clears(self): + self.seed() + before = self.conn.total_changes + plan = self.plan() + self.assertEqual(plan.action, "full") + self.assertEqual(self.conn.total_changes, before) + self.assertIsNone(self.source.load_manifest(self.conn, self.source.manifest_key(TARGETS))) + plan = self.api.initialize_tier_source(self.conn, plan) + self.assertEqual(plan.action, "resume") + self.assertEqual(self.ids(), []) + self.assertFalse(self.conn.in_transaction) + + def test_negative_zero_sparse_ids_and_exact_metadata_survive(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=2) + self.assertEqual(page.messages(), [ + {"id": mid, "content": f"synthetic-{mid}", "sender": "sender", "recipient": "recipient", "timestamp": "2026-01-02"} + for mid in (-3, 0) + ]) + self.assertFalse(self.conn.in_transaction) + result = self.build(plan) + self.assertTrue(result.captured_range_complete) + self.assertTrue(result.current_source_complete) + self.assertEqual(self.ids(), [-3, 0, 2, 8]) + self.assertEqual(self.ids(TARGETS[1]), [-3, 0, 2, 8]) + self.assertEqual((result.plan.manifest.cursor, result.plan.manifest.consumed, result.plan.manifest.outputs), (8, 4, 4)) + self.assertEqual(self.plan().action, "complete") + + def test_prefix_commit_retains_exact_pending_suffix_for_reduced_batches(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=4) + updated, pending = self.publish(plan, page, 1) + self.assertEqual(plan.manifest.consumed, 0) + self.assertEqual(updated.manifest.cursor, -3) + self.assertEqual([row[0] for row in pending.rows], [0, 2, 8]) + self.assertEqual(page.rows[1:], pending.rows) + updated, pending = self.publish(updated, pending, 2) + self.assertEqual(updated.manifest.cursor, 2) + self.assertEqual([row[0] for row in pending.rows], [8]) + updated, pending = self.publish(updated, pending, 1) + self.assertEqual(pending.rows, ()) + self.assertEqual(self.ids(), [-3, 0, 2, 8]) + self.assertTrue(self.api.finish_tier_source(self.conn, updated).current_source_complete) + + def test_resume_recomputes_digest_once_and_not_on_each_page(self): + self.seed(range(1, 140)) + plan = self.start() + plan, _ = self.publish(plan, self.api.read_tier_page(self.conn, plan, size=70)) + calls = [] + original = self.api._audit_pair + + def audit(conn, candidate): + calls.append(candidate.manifest.consumed) + return original(conn, candidate) + + self.api._audit_pair = audit + resumed = self.plan() + self.assertEqual(calls, [70]) + self.assertEqual(resumed.manifest.generation, plan.manifest.generation) + result = self.build(resumed) + self.assertEqual(calls, [70]) + self.assertEqual(result.plan.manifest.consumed, 139) + + def test_restart_on_new_in_memory_connection_revalidates_receipt(self): + uri = "file:tier_source_reopen_synthetic?mode=memory&cache=shared" + self.conn = sqlite3.connect(uri, uri=True, factory=FaultConnection) + self.addCleanup(self.conn.close) + self.create_schema(self.conn) + self.seed() + plan = self.start() + plan, _ = self.publish(plan, self.api.read_tier_page(self.conn, plan, size=2)) + other = sqlite3.connect(uri, uri=True, factory=FaultConnection) + self.addCleanup(other.close) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.api.read_tier_page(other, plan) + resumed = self.api.plan_tier_source(other, model_id="synthetic-model", dimension=2, targets=TARGETS) + self.assertEqual(resumed.manifest, plan.manifest) + self.assertEqual([row[0] for row in self.api.read_tier_page(other, resumed).rows], [2, 8]) + + def test_empty_source_finishes_without_model_or_vector_work(self): + result = self.build() + self.assertTrue(result.current_source_complete) + self.assertIsNone(result.plan.manifest.cursor) + self.assertEqual(self.ids(), []) + self.assertEqual(self.plan().action, "complete") + + def test_high_tail_requires_replan_and_fixed_range_is_not_current_coverage(self): + self.seed((1, 2)) + plan = self.start() + self.seed((3, 4)) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "explicitly replan"): + self.api.read_tier_page(self.conn, plan) + resumed = self.plan() + self.assertEqual(resumed.manifest.source.high_id, 2) + result = self.build(resumed) + self.assertTrue(result.captured_range_complete) + self.assertFalse(result.current_source_complete) + delta = self.plan() + self.assertEqual((delta.action, delta.prefix_count, delta.prefix_cursor, delta.manifest.total), ("delta", 2, 2, 2)) + delta = self.api.initialize_tier_source(self.conn, delta) + self.assertEqual([row[0] for row in self.api.read_tier_page(self.conn, delta).rows], [3, 4]) + result = self.build(delta) + self.assertTrue(result.current_source_complete) + self.assertEqual(self.ids(), [1, 2, 3, 4]) + self.assertEqual(self.plan().action, "complete") + + def test_empty_complete_pair_can_begin_first_delta(self): + self.build() + self.seed((1, 2)) + delta = self.plan() + self.assertEqual((delta.action, delta.prefix_count, delta.prefix_cursor), ("delta", 0, None)) + self.assertTrue(self.build(self.api.initialize_tier_source(self.conn, delta)).current_source_complete) + self.assertEqual(self.plan().action, "complete") + + def test_correction_delete_and_low_id_insert_refuse_resume(self): + for sql in ("UPDATE messages SET content='corrected' WHERE id=2", "DELETE FROM messages WHERE id=2", + "INSERT INTO messages(id,content) VALUES (1,'late-low-id')"): + with self.subTest(sql=sql): + conn = sqlite3.connect(":memory:") + try: + self.create_schema(conn) + conn.execute("INSERT INTO messages(id,content) VALUES (2,'synthetic')") + conn.commit() + plan = self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + self.api.initialize_tier_source(conn, plan) + conn.execute(sql) + conn.commit() + with self.assertRaises(self.source.RebuildSourceChanged): + self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + finally: + conn.close() + + def test_existing_unreceipted_rows_are_never_certified_or_cleared(self): + self.seed((1,)) + blob = self.vectors(1)[0][0] + for table in TARGETS: + self.conn.execute(f'INSERT INTO "{table}"(rowid,embedding) VALUES (1,?)', (blob,)) + self.conn.commit() + for force in (False, True): + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "empty inactive pair"): + self.plan(force=force) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids(TARGETS[1]), [1]) + + def test_same_count_tampered_vectors_and_wrong_width_refuse_restart(self): + self.seed((1,)) + self.build() + for blob in (struct.pack("2f", 0.5, 0.5), b"bad"): + with self.subTest(blob=blob): + self.conn.execute(f'UPDATE "{TARGETS[0]}" SET embedding=? WHERE rowid=1', (blob,)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + self.assertEqual(self.ids(), [1]) + + def test_pair_row_deletion_or_addition_refuses_restart(self): + self.seed((1, 2)) + result = self.build() + self.conn.execute(f'DELETE FROM "{TARGETS[1]}" WHERE rowid=1') + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + self.assertTrue(result.plan.manifest.complete) + + def test_pair_schema_drop_recreate_is_not_matching_content_proof(self): + self.seed((1,)) + self.build() + self.conn.executescript(f'DROP TABLE "{TARGETS[0]}"; CREATE TABLE "{TARGETS[0]}"(embedding BLOB);') + self.conn.execute(f'INSERT INTO "{TARGETS[0]}"(rowid,embedding) VALUES (1,?)', (self.vectors(1)[0][0],)) + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "schema"): + self.plan() + + def test_temp_case_variant_shadows_and_temp_triggers_refuse(self): + for statement in ( + 'CREATE TEMP TABLE MeSsAgEs(id INTEGER)', 'CREATE TEMP TABLE MeTaDaTa(key TEXT)', + 'CREATE TEMP TABLE VEC_MESSAGES_SYNTHETIC(embedding BLOB)', + 'CREATE TEMP TABLE maintenance_source_state(singleton INTEGER)', + 'CREATE TEMP TABLE truememory_rebuild_source_v1(singleton INTEGER)', + 'CREATE TEMP TRIGGER synthetic_shadow AFTER INSERT ON main.messages BEGIN SELECT 1; END', + ): + with self.subTest(statement=statement): + conn = sqlite3.connect(":memory:") + try: + self.create_schema(conn) + conn.execute(statement) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "TEMP"): + self.api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + finally: + conn.close() + + def test_shadow_added_after_planning_refuses_publication(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.execute('CREATE TEMP TABLE VEC_MESSAGES_SEP_SYNTHETIC(embedding BLOB)') + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "TEMP"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + + def test_unrelated_write_invalidates_fence_without_changing_checkpoint(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('unrelated','value')") + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.publish(plan, page) + resumed = self.plan() + self.assertEqual(resumed.manifest, plan.manifest) + self.assertEqual(self.publish(resumed, page)[0].manifest.consumed, 1) + + def test_external_commit_data_version_requires_explicit_revalidation(self): + uri = "file:tier_source_synthetic?mode=memory&cache=shared" + first = sqlite3.connect(uri, uri=True) + second = sqlite3.connect(uri, uri=True) + try: + self.create_schema(first) + plan = self.api.plan_tier_source(first, model_id="synthetic-model", dimension=2, targets=TARGETS) + plan = self.api.initialize_tier_source(first, plan) + second.execute("INSERT INTO metadata(key,value) VALUES ('external','synthetic')") + second.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Connection changed"): + self.api.read_tier_page(first, plan) + self.assertEqual(first.total_changes, plan.changes) + finally: + second.close() + first.close() + + def test_second_pair_write_failure_rolls_back_both_rows_and_checkpoint(self): + self.seed((1, 2)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.fail_sql = f'INSERT INTO main."{TARGETS[1]}"' + with self.assertRaisesRegex(sqlite3.OperationalError, "pair write"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids(TARGETS[1]), []) + self.assertFalse(self.conn.in_transaction) + resumed = self.plan() + self.assertEqual(resumed.manifest, plan.manifest) + updated, pending = self.publish(resumed, page, 1) + self.assertEqual(updated.manifest.cursor, 1) + self.assertEqual([row[0] for row in pending.rows], [2]) + + def test_commit_failure_rolls_back_pair_receipt_and_manifest(self): + self.seed((1,)) + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + self.conn.fail_commit = True + with self.assertRaisesRegex(sqlite3.OperationalError, "commit failure"): + self.publish(plan, page) + self.assertEqual(self.ids(), []) + self.assertEqual(self.plan().manifest, plan.manifest) + self.assertFalse(self.conn.in_transaction) + + def test_pending_inputs_prevent_finish_and_bad_prefixes_do_not_advance(self): + self.seed() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan, size=2) + with self.assertRaises(self.api.TierSourceUntrusted): + self.api.finish_tier_source(self.conn, plan) + for left, right in (([], []), (self.vectors(1)[0], []), (self.vectors(3)[0], self.vectors(3)[1])): + with self.assertRaises(self.api.TierSourceUntrusted): + self.api.publish_tier_prefix(self.conn, plan, page, completion=left, separation=right) + with self.assertRaises(ValueError): + self.api.publish_tier_prefix(self.conn, plan, page, completion=[b"bad"], separation=[b"bad"]) + changed = replace(page, rows=((page.rows[0][0], "wrong-text", *page.rows[0][2:]), *page.rows[1:])) + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Pending page"): + self.publish(plan, changed, 1) + self.assertEqual(self.plan().manifest, plan.manifest) + self.assertEqual(self.ids(), []) + + def test_borrowed_transaction_is_never_committed_or_rolled_back(self): + plan = self.start() + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('borrowed','pending')") + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "owned, clean"): + self.api.read_tier_page(self.conn, plan) + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='borrowed'").fetchone()[0], "pending") + self.conn.rollback() + + def test_page_fields_require_exact_sqlite_type_and_signed_real_zero(self): + self.conn = sqlite3.connect(":memory:", factory=FaultConnection) + self.addCleanup(self.conn.close) + self.conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content, sender, recipient, timestamp); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL); + CREATE TABLE vec_messages_synthetic(embedding BLOB); + CREATE TABLE vec_messages_sep_synthetic(embedding BLOB); + """) + self.source.ensure_rebuild_tracking(self.conn) + self.conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", [(1, 1), (2, -0.0)]) + self.conn.commit() + plan = self.start() + page = self.api.read_tier_page(self.conn, plan) + integer_as_real = replace(page, rows=((1, 1.0, None, None, None), page.rows[1])) + signed_zero = replace(page, rows=(page.rows[0], (2, 0.0, None, None, None))) + for changed in (integer_as_real, signed_zero): + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Pending page"): + self.publish(plan, changed) + self.assertEqual(self.ids(), []) + self.assertEqual(self.publish(plan, page)[0].manifest.consumed, 2) + + def test_malformed_manifest_types_ranges_and_complete_flags_refuse(self): + self.seed() + plan = self.start() + key = self.source.manifest_key(TARGETS) + original = json.loads(plan.saved_manifest) + mutations = ( + {"version": True}, {"dimension": True}, {"consumed": -1}, {"consumed": True}, + {"outputs": 1}, {"consumed": 5, "outputs": 5}, {"cursor": 0}, {"complete": True}, + {"complete": 1}, {"total": 5}, {"source": None}, {"schema_version": None}, + {"targets": list(reversed(TARGETS))}, {"generation": "not-a-generation"}, + ) + for update in mutations: + with self.subTest(update=update): + self.conn.execute("UPDATE metadata SET value=? WHERE key=?", (json.dumps(dict(original, **update)), key)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + + def test_missing_receipt_is_not_registry_or_manifest_compatibility(self): + self.start() + self.conn.execute("DELETE FROM metadata WHERE key=?", (self.api._receipt_key(TARGETS),)) + self.conn.commit() + with self.assertRaisesRegex(self.api.TierSourceUntrusted, "Both manifest and pair receipt"): + self.plan() + + def test_oversized_and_duplicate_metadata_are_rejected(self): + plan = self.start() + key = self.source.manifest_key(TARGETS) + for raw in ("x" * 65537, plan.saved_manifest[:-1] + ',"version":1}'): + self.conn.execute("UPDATE metadata SET value=? WHERE key=?", (raw, key)) + self.conn.commit() + with self.assertRaises(self.api.TierSourceUntrusted): + self.plan() + + def test_wrong_model_dimension_and_target_identifiers_refuse(self): + self.start() + for model, dimension, targets in ( + ("other", 2, TARGETS), ("synthetic-model", 3, TARGETS), + ("synthetic-model", 4097, TARGETS), + ("synthetic-model", True, TARGETS), ("synthetic-model", 2, ("vec_messages_bad;--", TARGETS[1])), + ): + with self.assertRaises((ValueError, self.api.TierSourceUntrusted)): + self.api.plan_tier_source(self.conn, model_id=model, dimension=dimension, targets=targets) + + def test_restart_digest_reader_uses_bounded_pages(self): + self.seed(range(1, 170)) + self.build() + batches = [] + original = self.api._bounded_rows + + def bounded(cursor): + class Reader: + def fetchmany(inner, size): + batches.append(size) + return cursor.fetchmany(size) + + def close(inner): + cursor.close() + yield from original(Reader()) + + self.api._bounded_rows = bounded + self.assertEqual(self.plan().action, "complete") + self.assertGreater(len(batches), 3) + self.assertEqual(set(batches), {64}) + + +@unittest.skipUnless(os.environ.get("TRUEMEMORY_PREPARED_NATIVE_SQLITE_TEST") == "1", "GPUBox-only native SQLite gate") +class TestNativeTierSource(unittest.TestCase): + def test_real_vec0_pair_restart_and_delta(self): + import sqlite_vec + + modules = load_modules() + api, source = modules["tier_switch.source"], modules["rebuild_source"] + conn = sqlite3.connect(":memory:") + try: + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + conn.executescript(""" + CREATE TABLE messages(id INTEGER PRIMARY KEY, content TEXT, sender TEXT, recipient TEXT, timestamp TEXT); + CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TEXT); + CREATE VIRTUAL TABLE vec_messages_synthetic USING vec0(embedding float[2] distance_metric=cosine); + CREATE VIRTUAL TABLE vec_messages_sep_synthetic USING vec0(embedding float[2] distance_metric=cosine); + """) + source.ensure_rebuild_tracking(conn) + conn.executemany("INSERT INTO messages(id,content) VALUES (?,?)", [(-3, "synthetic-negative"), (0, "synthetic-zero")]) + conn.commit() + + def plan(): + return api.plan_tier_source(conn, model_id="synthetic-model", dimension=2, targets=TARGETS) + + current = api.initialize_tier_source(conn, plan()) + page = api.read_tier_page(conn, current) + left, right = [struct.pack("2f", 1., 0.)] * 2, [struct.pack("2f", 0., 1.)] * 2 + current, pending = api.publish_tier_prefix(conn, current, page, completion=left[:1], separation=right[:1]) + self.assertEqual(plan().manifest, current.manifest) + current, pending = api.publish_tier_prefix(conn, current, pending, completion=left[1:], separation=right[1:]) + self.assertTrue(api.finish_tier_source(conn, current).current_source_complete) + conn.execute("INSERT INTO messages(id,content) VALUES (2,'synthetic-tail')") + conn.commit() + current = api.initialize_tier_source(conn, plan()) + page = api.read_tier_page(conn, current) + current, _ = api.publish_tier_prefix(conn, current, page, completion=left[:1], separation=right[:1]) + self.assertTrue(api.finish_tier_source(conn, current).current_source_complete) + self.assertEqual(plan().action, "complete") + for table in TARGETS: + self.assertEqual([row[0] for row in conn.execute(f'SELECT rowid FROM "{table}" ORDER BY rowid')], [-3, 0, 2]) + finally: + conn.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/truememory/tier_switch/source.py b/truememory/tier_switch/source.py new file mode 100644 index 00000000..3e2d4c48 --- /dev/null +++ b/truememory/tier_switch/source.py @@ -0,0 +1,464 @@ +"""Model-free, bounded source checkpoints for caller-owned inactive tier pairs. + +The caller owns the maintenance lease, exclusive cooperating pair writer, target +schema/model validation, and database identity across reopen. Tracking, metadata, +and both target tables must exist before planning. This adapter neither clears +tables nor activates them. It rejects borrowed transactions and concurrent writes +until the caller replans. Receipts detect changed stored bytes, not a hostile SQL +writer capable of forging both data and receipts. +""" + +import hashlib +import json +import re +import sqlite3 +import struct +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from dataclasses import asdict, dataclass, replace + +from truememory.maintenance import SourceRevision +from truememory.rebuild_source import RebuildManifest, RebuildSource, RebuildSourceChanged, capture_source, manifest_key + +_PAGE_ROWS = 64 +_MAX_METADATA_BYTES = 65536 +_SEED = hashlib.sha256(b"truememory-tier-pair-chain-v1").hexdigest() +_FIELDS = ("id", "content", "sender", "recipient", "timestamp") +_ID_MIN, _ID_MAX = -(2**63), 2**63 - 1 +_HEX = re.compile(r"[0-9a-f]{64}\Z") +SQLiteValue = str | bytes | int | float | None + + +class TierSourceUntrusted(RebuildSourceChanged): + """A fresh empty pair or explicit revalidation is required.""" + + +@dataclass(frozen=True) +class TierSourcePlan: + action: str + manifest: RebuildManifest + schema: str + digest: str + prefix_count: int + prefix_cursor: int | None + connection: sqlite3.Connection + data_version: int + changes: int + saved_manifest: str | None + saved_receipt: str | None + + +@dataclass(frozen=True) +class TierSourcePage: + generation: str + cursor: int | None + consumed: int + rows: tuple[tuple[SQLiteValue, ...], ...] + + def messages(self) -> list[dict[str, SQLiteValue]]: + """A bounded copy suitable for the existing completion/separation format.""" + return [dict(zip(_FIELDS, row)) for row in self.rows] + + +@dataclass(frozen=True) +class TierSourceCompletion: + plan: TierSourcePlan + captured_range_complete: bool + current_source_complete: bool + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _hash(value: object) -> str: + return hashlib.sha256(_json(value).encode("utf-8")).hexdigest() + + +def _integer(value: object, *, minimum: int = 0, maximum: int = _ID_MAX) -> bool: + return type(value) is int and minimum <= value <= maximum + + +def _identifier(value: object) -> bool: + return _integer(value, minimum=_ID_MIN) + + +def _quote(name: str) -> str: + return '"' + name.replace('"', '""') + '"' + + +def _identity(model_id: str, dimension: int, targets: tuple[str, str]) -> None: + if (not isinstance(model_id, str) or not model_id or len(model_id) > 4096 + or not _integer(dimension, minimum=1, maximum=4096)): + raise ValueError("A concrete model and integer dimension in 1..4096 are required") + if (type(targets) is not tuple or len(targets) != 2 + or not all(isinstance(name, str) for name in targets)): + raise ValueError("An exact ordered completion/separation pair is required") + match = re.fullmatch(r"vec_messages_([a-z][a-z0-9_]{0,63})", targets[0]) + if match is None or targets[1] != "vec_messages_sep_" + match[1]: + raise ValueError("Target names must be a canonical paired tier group") + + +def _guard_schema(conn: sqlite3.Connection, targets: tuple[str, ...]) -> str: + guarded = (*targets, "messages", "metadata", "maintenance_source_state", "truememory_rebuild_source_v1") + marks = ",".join("?" for _ in guarded) + if conn.execute( + f"SELECT 1 FROM temp.sqlite_master WHERE name COLLATE NOCASE IN ({marks}) OR " + f"(type='trigger' AND tbl_name COLLATE NOCASE IN ({marks})) LIMIT 1", (*guarded, *guarded), + ).fetchone(): + raise TierSourceUntrusted("TEMP shadows or triggers invalidate tier source ownership") + if conn.execute( + "SELECT 1 FROM main.sqlite_master WHERE type='trigger' AND tbl_name COLLATE NOCASE IN (?,?,?) LIMIT 1", + (*targets, "metadata"), + ).fetchone(): + raise TierSourceUntrusted("Tier pair and receipt tables must have no side-effect triggers") + definitions = [] + for name in (*targets, "metadata"): + row = conn.execute("SELECT type, rootpage, sql FROM main.sqlite_master WHERE name=?", (name,)).fetchone() + if row is None or row[0] != "table" or not row[2]: + raise TierSourceUntrusted("Initialized MAIN target and metadata tables are required") + columns = [tuple(item) for item in conn.execute(f"PRAGMA main.table_xinfo({_quote(name)})")] + if name in targets and "embedding" not in {column[1] for column in columns}: + raise TierSourceUntrusted("Target table has no embedding column") + if name == "metadata" and ([column[1] for column in columns if column[5]] != ["key"] + or "value" not in {column[1] for column in columns}): + raise TierSourceUntrusted("Metadata requires a unique key and value column") + definitions.append((name, tuple(row), columns)) + return _hash([conn.execute("PRAGMA main.schema_version").fetchone()[0], definitions]) + + +def _metadata(conn: sqlite3.Connection, key: str) -> str | None: + shape = conn.execute( + "SELECT typeof(value), length(CAST(value AS BLOB)) FROM main.metadata WHERE key=?", (key,), + ).fetchone() + if shape is None: + return None + if shape[0] != "text" or shape[1] > _MAX_METADATA_BYTES: + raise TierSourceUntrusted("Tier source metadata has an unsupported type or size") + return conn.execute("SELECT value FROM main.metadata WHERE key=?", (key,)).fetchone()[0] + + +def _receipt_key(targets: tuple[str, ...]) -> str: + return "tier_pair_receipt_v1:" + ":".join(targets) + + +def _manifest_json(manifest: RebuildManifest) -> str: + return _json(dict(asdict(manifest), version=1)) + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result = dict(pairs) + if len(result) != len(pairs): + raise ValueError("Duplicate metadata field") + return result + + +def _parse_manifest(raw: str, model: str, dimension: int, targets: tuple[str, str]) -> RebuildManifest: + try: + value = json.loads(raw, object_pairs_hook=_unique_object) + expected = set(RebuildManifest.__dataclass_fields__) | {"version"} + if (type(value) is not dict or set(value) != expected + or type(value["version"]) is not int or value.pop("version") != 1): + raise ValueError("manifest shape") + source = value["source"] + if type(source) is not dict or set(source) != set(RebuildSource.__dataclass_fields__): + raise ValueError("source shape") + revision = source["revision"] + if (type(revision) is not list or len(revision) != 6 or not isinstance(revision[0], str) + or not revision[0] or not all(_integer(revision[index]) for index in (1, 2, 3, 5)) + or not _identifier(revision[4]) or revision[5] > revision[1] + or revision[2] + revision[3] != revision[1]): + raise ValueError("source revision") + source["revision"] = SourceRevision(*revision) + if (source["tracker"] not in ("canonical-v1", "rebuild-bridge-v1", "rebuild-bridge-conservative-v1") + or not isinstance(source["schema_signature"], str) or not _HEX.fullmatch(source["schema_signature"]) + or not _integer(source["total"])): + raise ValueError("source descriptor") + for field in ("after_id", "high_id"): + if source[field] is not None and not _identifier(source[field]): + raise ValueError("source identifier") + if ((source["total"] == 0) != (source["high_id"] is None) + or (source["high_id"] is not None and source["after_id"] is not None + and source["high_id"] <= source["after_id"])): + raise ValueError("source range") + if (value["model"] != model or type(value["dimension"]) is not int or value["dimension"] != dimension + or value["targets"] != list(targets) or value["input_digest"] is not None + or not isinstance(value["generation"], str) or not re.fullmatch(r"[0-9a-f]{32}", value["generation"]) + or not all(_integer(value[field]) for field in ("total", "consumed", "outputs", "schema_version")) + or value["total"] != source["total"] or value["outputs"] != value["consumed"] + or value["consumed"] > value["total"] or type(value["complete"]) is not bool + or (value["complete"] and value["consumed"] != value["total"])): + raise ValueError("manifest invariant") + cursor = value["cursor"] + if value["consumed"] == 0: + if cursor != source["after_id"] or type(cursor) is not type(source["after_id"]): + raise ValueError("initial cursor") + elif (not _identifier(cursor) or cursor > source["high_id"] + or (source["after_id"] is not None and cursor <= source["after_id"]) + or ((value["consumed"] == value["total"]) != (cursor == source["high_id"]))): + raise ValueError("consumed cursor") + value["source"] = RebuildSource(**source) + value["targets"] = tuple(value["targets"]) + return RebuildManifest(**value) + except (ValueError, TypeError, KeyError, OverflowError, RecursionError) as error: + raise TierSourceUntrusted("Malformed or incompatible tier source manifest") from error + + +def _receipt(plan: TierSourcePlan) -> str: + return _json({"version": 1, "manifest": _hash(json.loads(_manifest_json(plan.manifest))), + "schema": plan.schema, "digest": plan.digest, "prefix_count": plan.prefix_count, + "prefix_cursor": plan.prefix_cursor}) + + +def _parse_receipt(raw: str, manifest: RebuildManifest, schema: str) -> tuple[str, int, int | None]: + try: + value = json.loads(raw, object_pairs_hook=_unique_object) + if (type(value) is not dict or set(value) != {"version", "manifest", "schema", "digest", "prefix_count", "prefix_cursor"} + or type(value["version"]) is not int or value["version"] != 1 + or value["manifest"] != _hash(json.loads(_manifest_json(manifest))) or value["schema"] != schema + or not isinstance(value["digest"], str) or not _HEX.fullmatch(value["digest"]) + or not _integer(value["prefix_count"])): + raise ValueError("receipt invariant") + cursor = value["prefix_cursor"] + if ((cursor is not None and not _identifier(cursor)) or cursor != manifest.source.after_id + or (value["prefix_count"] == 0) != (cursor is None)): + raise ValueError("receipt prefix") + return value["digest"], value["prefix_count"], cursor + except (ValueError, TypeError, KeyError, RecursionError) as error: + raise TierSourceUntrusted("Malformed or incompatible tier pair receipt") from error + + +def _chain(digest: str, row_id: int, completion: bytes, separation: bytes) -> str: + return hashlib.sha256(bytes.fromhex(digest) + struct.pack(">qQQ", row_id, len(completion), len(separation)) + + completion + separation).hexdigest() + + +def _bounded_rows(cursor: sqlite3.Cursor) -> Iterator[tuple]: + try: + while rows := cursor.fetchmany(_PAGE_ROWS): + yield from rows + finally: + cursor.close() + + +def _audit_pair(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + """One snapshot, bounded resident rows; called only by plan/resume.""" + captured = plan.manifest.source + lower = " AND id > ?" if captured.after_id is not None else "" + args = (captured.high_id, captured.after_id) if lower else (captured.high_id,) + total, high = conn.execute( + f"SELECT COUNT(*), MAX(id) FROM main.messages WHERE id <= ?{lower}", args, + ).fetchone() + if (total, high) != (captured.total, captured.high_id): + raise TierSourceUntrusted("Captured source bounds do not match the manifest") + width = plan.manifest.dimension * 4 + readers = [_bounded_rows(conn.execute( + f"SELECT rowid, CASE WHEN typeof(embedding)='blob' AND length(embedding)=? THEN embedding END " + f"FROM main.{_quote(table)} ORDER BY rowid", (width,), + )) for table in plan.manifest.targets] + high = plan.manifest.cursor + source = _bounded_rows(conn.execute( + "SELECT id FROM main.messages WHERE id <= ? ORDER BY id", (high,), + )) + readers.append(source) + digest, count = _SEED, 0 + try: + while True: + left, right, original = (next(reader, None) for reader in readers) + if left is None and right is None and original is None: + break + if (left is None or right is None or original is None or left[0] != right[0] or left[0] != original[0] + or not _identifier(left[0]) or type(left[1]) is not bytes or type(right[1]) is not bytes + or len(left[1]) != width or len(right[1]) != width): + raise TierSourceUntrusted("Tier pair bytes or source coverage do not match the receipt") + digest = _chain(digest, left[0], left[1], right[1]) + count += 1 + finally: + for reader in readers: + reader.close() + if digest != plan.digest or count != plan.prefix_count + plan.manifest.outputs: + raise TierSourceUntrusted("Tier pair digest or consumed count changed") + + +def _fence(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + if (conn is not plan.connection or conn.total_changes != plan.changes + or conn.execute("PRAGMA main.data_version").fetchone()[0] != plan.data_version): + raise TierSourceUntrusted("Connection changed; explicitly replan and revalidate the tier pair") + + +@contextmanager +def _transaction(conn: sqlite3.Connection, *, write: bool = False) -> Iterator[None]: + if conn.in_transaction: + raise TierSourceUntrusted("Tier source work requires an owned, clean connection") + conn.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + yield + conn.commit() + except BaseException: + conn.rollback() + raise + + +def _check(conn: sqlite3.Connection, plan: TierSourcePlan) -> None: + _fence(conn, plan) + if _guard_schema(conn, plan.manifest.targets) != plan.schema: + raise TierSourceUntrusted("Tier target schema changed") + if (_metadata(conn, manifest_key(plan.manifest.targets)) != plan.saved_manifest + or _metadata(conn, _receipt_key(plan.manifest.targets)) != plan.saved_receipt): + raise TierSourceUntrusted("Tier generation or receipt changed") + plan.manifest.source.check(conn) + + +def plan_tier_source(conn: sqlite3.Connection, *, model_id: str, dimension: int, + targets: tuple[str, str], force: bool = False) -> TierSourcePlan: + """Plan full, resume, delta or complete without changing persistent state. + + A full plan can only initialize an already empty inactive pair. Existing + unreceipted rows are refused, even with force. Snapshot consistency and + exclusive cooperating writer ownership are required throughout this call. + """ + _identity(model_id, dimension, targets) + version, changes = conn.execute("PRAGMA main.data_version").fetchone()[0], conn.total_changes + with _transaction(conn): + schema = _guard_schema(conn, targets) + saved = _metadata(conn, manifest_key(targets)) + receipt = _metadata(conn, _receipt_key(targets)) + if force or (saved is None and receipt is None): + if any(conn.execute(f"SELECT 1 FROM main.{_quote(table)} LIMIT 1").fetchone() for table in targets): + raise TierSourceUntrusted("Full initialization requires a separately owned empty inactive pair") + source = capture_source(conn) + manifest = replace(RebuildManifest.new(model_id, dimension, targets, source=source), + schema_version=conn.execute("PRAGMA main.schema_version").fetchone()[0]) + plan = TierSourcePlan("full", manifest, schema, _SEED, 0, None, conn, version, changes, saved, receipt) + else: + if saved is None or receipt is None: + raise TierSourceUntrusted("Both manifest and pair receipt are required for compatibility") + manifest = _parse_manifest(saved, model_id, dimension, targets) + if manifest.schema_version != conn.execute("PRAGMA main.schema_version").fetchone()[0]: + raise TierSourceUntrusted("Manifest schema generation changed") + manifest.source.check(conn) + digest, prefix, cursor = _parse_receipt(receipt, manifest, schema) + plan = TierSourcePlan("resume", manifest, schema, digest, prefix, cursor, conn, version, changes, saved, receipt) + _audit_pair(conn, plan) + if manifest.complete: + source = capture_source(conn, after_id=manifest.cursor) + if source.total: + updated = replace(RebuildManifest.new(model_id, dimension, targets, source=source), + schema_version=manifest.schema_version) + plan = replace(plan, action="delta", manifest=updated, + prefix_count=prefix + manifest.outputs, prefix_cursor=manifest.cursor) + else: + plan = replace(plan, action="complete") + _fence(conn, plan) + _fence(conn, plan) + return plan + + +def _persist(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourcePlan: + saved, receipt = _manifest_json(plan.manifest), _receipt(plan) + for key, value in ((manifest_key(plan.manifest.targets), saved), (_receipt_key(plan.manifest.targets), receipt)): + conn.execute("INSERT OR REPLACE INTO main.metadata(key,value) VALUES (?,?)", (key, value)) + return replace(plan, saved_manifest=saved, saved_receipt=receipt, changes=conn.total_changes) + + +def initialize_tier_source(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourcePlan: + """Persist a planned generation; never create or clear target tables.""" + with _transaction(conn, write=True): + _check(conn, plan) + if plan.action not in ("full", "delta"): + return plan + if plan.action == "full" and any( + conn.execute(f"SELECT 1 FROM main.{_quote(table)} LIMIT 1").fetchone() for table in plan.manifest.targets + ): + raise TierSourceUntrusted("Full target pair is no longer empty") + updated = _persist(conn, replace(plan, action="resume")) + return updated + + +def _read_rows(conn: sqlite3.Connection, plan: TierSourcePlan, size: int) -> tuple[tuple[SQLiteValue, ...], ...]: + source = plan.manifest.source + if source.high_id is None: + return () + lower = " AND id > ?" if plan.manifest.cursor is not None else "" + args = (source.high_id, plan.manifest.cursor, size) if lower else (source.high_id, size) + return tuple(tuple(row) for row in conn.execute( + f"SELECT {', '.join(_FIELDS)} FROM main.messages WHERE id <= ?{lower} ORDER BY id LIMIT ?", args, + )) + + +def _same_rows(left: tuple[tuple[SQLiteValue, ...], ...], right: tuple[tuple[SQLiteValue, ...], ...]) -> bool: + if len(left) != len(right): + return False + for current, pending in zip(left, right): + if len(current) != len(pending): + return False + for actual, captured in zip(current, pending): + if type(actual) is not type(captured): + return False + if isinstance(actual, float): + if struct.pack(">d", actual) != struct.pack(">d", captured): + return False + elif actual != captured: + return False + return True + + +def read_tier_page(conn: sqlite3.Connection, plan: TierSourcePlan, *, size: int = _PAGE_ROWS) -> TierSourcePage: + """Release the read snapshot before returning any texts to the encoder.""" + if not _integer(size, minimum=1): + raise ValueError("Page size must be a positive integer") + if plan.action not in ("resume", "complete"): + raise TierSourceUntrusted("Initialize the generation before reading") + with _transaction(conn): + _check(conn, plan) + rows = _read_rows(conn, plan, size) + if len(rows) != min(size, plan.manifest.total - plan.manifest.consumed): + raise TierSourceUntrusted("Source range no longer matches its captured count") + _fence(conn, plan) + return TierSourcePage(plan.manifest.generation, plan.manifest.cursor, plan.manifest.consumed, rows) + + +def publish_tier_prefix(conn: sqlite3.Connection, plan: TierSourcePlan, page: TierSourcePage, *, + completion: Sequence[bytes], separation: Sequence[bytes]) -> tuple[TierSourcePlan, TierSourcePage]: + """Atomically publish only an encoded page prefix and retain its pending suffix. + + The caller supplies already serialized vectors from the captured model. This + function checks byte widths, but does not infer model identity from vectors. + On failure the original immutable plan/page remain the retry checkpoint; + after database write failures replan because total_changes includes rollbacks. + """ + count = len(completion) + if (plan.action != "resume" or plan.manifest.complete or page.generation != plan.manifest.generation + or page.cursor != plan.manifest.cursor or page.consumed != plan.manifest.consumed + or count < 1 or count != len(separation) or count > len(page.rows)): + raise TierSourceUntrusted("Encoded prefix does not belong to the current pending page") + width = plan.manifest.dimension * 4 + if any(type(blob) is not bytes or len(blob) != width for vectors in (completion, separation) for blob in vectors): + raise ValueError("Both serialized vector widths must match the captured dimension") + digest = plan.digest + with _transaction(conn, write=True): + _check(conn, plan) + rows = _read_rows(conn, plan, len(page.rows)) + if not _same_rows(rows, page.rows): + raise TierSourceUntrusted("Pending page does not match the source under writer ownership") + for row, left, right in zip(rows[:count], completion, separation): + for table, vector in zip(plan.manifest.targets, (left, right)): + conn.execute(f"INSERT INTO main.{_quote(table)}(rowid,embedding) VALUES (?,?)", (row[0], vector)) + digest = _chain(digest, row[0], left, right) + manifest = plan.manifest.advance(rows[count - 1][0], count, count) + updated = _persist(conn, replace(plan, manifest=manifest, digest=digest)) + pending = TierSourcePage(manifest.generation, manifest.cursor, manifest.consumed, page.rows[count:]) + return updated, pending + + +def finish_tier_source(conn: sqlite3.Connection, plan: TierSourcePlan) -> TierSourceCompletion: + """Certify captured coverage only; no registry, model or serving activation.""" + if plan.action not in ("resume", "complete") or plan.manifest.consumed != plan.manifest.total: + raise TierSourceUntrusted("The captured source range still has pending inputs") + with _transaction(conn, write=True): + _check(conn, plan) + current = capture_source(conn) + covers_current = (current.total == plan.prefix_count + plan.manifest.outputs + and current.high_id == plan.manifest.cursor) + updated = _persist(conn, replace(plan, action="complete", manifest=replace(plan.manifest, complete=True))) + return TierSourceCompletion(updated, True, covers_current)