diff --git a/docs/resource-budgets.md b/docs/resource-budgets.md index 3fe0675a..f01c02cd 100644 --- a/docs/resource-budgets.md +++ b/docs/resource-budgets.md @@ -647,3 +647,41 @@ resume. Synthetic SQLite tests cover precise REAL metadata rendered into separation text, signed zero, collations, storage types, rowid aliases, generated fields, repair failures and caller rollback. These tests do not measure native model behavior or add a process memory budget. + +### Foreground model admission before writer ownership + +Owned foreground add and incremental vector publication now acquire the model +identity fence before opening the SQLite write transaction. Add first owns its +engine write lock, so another mutation on that same engine cannot hold the +engine lock while waiting for a model fence held by the queued add. The model +fence validates the captured generation and remains held through commit or +rollback. There is one model acquisition, with no nested acquisition of the +nonreentrant model lock. + +This permits an owned foreground call to wait for a concurrent cluster source +snapshot without holding SQLite's writer or failing a valid vector publication. +Model loading and inference finish before owned writer admission. A model +generation change during loading, encoding or admission still rejects stale +publication; add and update propagate those typed errors instead of silently +falling back to a source-only write. Update retains its existing distinction: +capture failure precedes its source write, while a later vector-publication +failure can occur after the source update committed. + +Caller transactions use nonblocking model capture and publication admission. +Add and update classify the connection during model capture under their engine +write lock, so another same-engine mutation's transaction cannot be mistaken +for borrowed caller work. They release that lock before encoding. +Capture accepts only a cached model observed together with its identity under +the model lock. Busy ownership or an unloaded model produces a typed retry +error before any additional source write; the caller's transaction remains +intact. Cached-model inference can still run inside a caller-held transaction. +The publication fence protects the nested savepoint release, not a later outer +commit controlled by that caller. Streamed rebuilds retain nonblocking model +fences after writer admission and use the same cached-only borrowed capture. + +SQLite and event regressions cover actual clustering snapshot contention, +same-engine mutation ordering, stale generations, borrowed sentinels, atomic +add publication, and model release after admission/publication/commit failures. +This checkpoint changes admission ordering rather than bounding wait duration: +cluster snapshots may still hold the model lock while scanning the corpus. +Native integration and latency remain separate acceptance checks. diff --git a/tests/test-maintenance-foreground-admission-753.py b/tests/test-maintenance-foreground-admission-753.py new file mode 100644 index 00000000..d295387e --- /dev/null +++ b/tests/test-maintenance-foreground-admission-753.py @@ -0,0 +1,491 @@ +"""Foreground publication admission with real SQLite and synthetic models.""" + +import runpy +import sqlite3 +import threading +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +FOREGROUND = runpy.run_path(str(Path(__file__).with_name("test-rebuild-foreground-publication-756.py"))) +CLUSTERS = runpy.run_path(str(Path(__file__).with_name("test-cluster-publication-749.py"))) + + +class AdmissionFixture(unittest.TestCase): + setUp = FOREGROUND["TestForegroundPublication"].setUp + seed = FOREGROUND["TestForegroundPublication"].seed + ids = FOREGROUND["TestForegroundPublication"].ids + writer = FOREGROUND["TestForegroundPublication"].writer + engine = FOREGROUND["TestForegroundPublication"].engine + + def start(self, action, name="synthetic-worker"): + results, done = [], threading.Event() + + def run(): + try: + results.append(action()) + except BaseException as error: + results.append(error) + finally: + done.set() + + thread = threading.Thread(target=run, name=name, daemon=True) + thread.start() + self.addCleanup(thread.join, 3) + return thread, results, done + + def released(self): + self.assertTrue(self.vector._lock.acquire(blocking=False)) + self.vector._lock.release() + + def sentinel(self): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-caller','pending')") + + def preserved(self): + self.assertTrue(self.conn.in_transaction) + self.assertEqual(self.conn.execute("SELECT value FROM metadata WHERE key='synthetic-caller'").fetchone()[0], "pending") + + +class TestForegroundAdmission(AdmissionFixture): + def test_style_computation_runs_outside_capture_and_publication_locks(self) -> None: + engine = self.engine() + engine._has_style_vec = True + events = [] + style = [0.25, 0.75] + + def compute(content: str) -> list[float]: + self.assertFalse(engine._write_lock.locked()) + self.assertFalse(self.vector._lock.locked()) + self.assertFalse(self.conn.in_transaction) + events.append("compute") + return style + + def publish(conn: sqlite3.Connection, sender: str, content: str, + _pre_computed_vec: list[float]) -> None: + self.assertTrue(engine._write_lock.locked()) + self.assertTrue(self.vector._lock.locked()) + self.assertTrue(conn.in_transaction) + self.assertIs(_pre_computed_vec, style) + events.append("publish") + + module = types.ModuleType("synthetic_style") + module.compute_style_vector = compute + with patch.dict(self.modules, personality_style_vec=module), \ + patch.dict(engine.add.__globals__, _update_style_vec=publish): + engine.add("synthetic styled content", sender="synthetic sender") + self.assertEqual(events, ["compute", "publish"]) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + + def test_actual_cluster_snapshot_delays_add_before_writer_then_add_commits_both_vectors(self): + engine = self.engine() + cluster, runtime, _native = CLUSTERS["load_clustering"]() + encoding, snapshot, admission, release = (threading.Event() for _ in range(4)) + underlying = self.vector._lock + fixture = self + + class ObservedLock: + def acquire(self, blocking=True): + if threading.current_thread().name == "synthetic-add" and snapshot.is_set(): + fixture.assertTrue(blocking) + fixture.assertFalse(engine.conn.in_transaction) + fixture.assertTrue(engine._write_lock.locked()) + admission.set() + return underlying.acquire(blocking=blocking) + + def release(self): + underlying.release() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *_args): + self.release() + + self.vector._lock = runtime._lock = ObservedLock() + runtime.EMBEDDING_MODEL = self.vector.EMBEDDING_MODEL + cluster_conn, independent = self.writer(), self.writer() + original = cluster._get_all_embeddings + + def read(conn): + snapshot.set() + self.assertTrue(release.wait(3)) + return original(conn) + + def encode(texts): + if len(self.calls) == 1: + encoding.set() + self.assertTrue(snapshot.wait(3)) + return [[1., 0.] for _ in texts] + + def consolidate(): + self.assertTrue(encoding.wait(3)) + return cluster.cluster_messages(cluster_conn) + + self.on_encode = encode + with patch.object(cluster, "_get_all_embeddings", read): + cluster_thread, cluster_result, cluster_done = self.start(consolidate) + add_thread, result, done = self.start(lambda: engine.add("synthetic first add"), "synthetic-add") + try: + self.assertTrue(admission.wait(2)) + self.assertFalse(done.is_set()) + independent.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-independent','committed')") + independent.commit() + self.assertFalse(self.conn.in_transaction) + finally: + release.set() + add_thread.join(3) + cluster_thread.join(3) + self.assertTrue(done.is_set()) + self.assertTrue(cluster_done.is_set()) + self.assertIsInstance(cluster_result[0], (int, RuntimeError), cluster_result) + self.assertIsInstance(result[0], dict, result) + self.assertEqual(self.ids(), [1]) + self.assertEqual(self.ids("vec_messages_sep"), [1]) + self.assertEqual(self.conn.execute("SELECT content FROM messages").fetchone()[0], "synthetic first add") + self.released() + + def test_writer_owner_defers_nonblocking_model_fence_and_foreground_finishes(self): + owner = self.writer() + owner.execute("BEGIN IMMEDIATE") + admission, attempted = threading.Event(), threading.Event() + underlying = self.vector._lock + + class AdmittedLock: + def acquire(self, blocking=True): + acquired = underlying.acquire(blocking=blocking) + if acquired and threading.current_thread().name == "synthetic-publish": + admission.set() + return acquired + + def release(self): + underlying.release() + + self.vector._lock = AdmittedLock() + identity = (0, self.vector.EMBEDDING_MODEL, 2) + + def publish(): + with self.vector._foreground_vector_publication(self.conn, identity, 1): + self.conn.execute("INSERT INTO metadata(key,value) VALUES ('synthetic-published','yes')") + + thread, result, done = self.start(publish, "synthetic-publish") + try: + self.assertTrue(admission.wait(2)) + with self.assertRaisesRegex(RuntimeError, "model is busy"): + with self.vector._rebuild_model_fence(identity): + attempted.set() + self.assertFalse(attempted.is_set()) + finally: + owner.rollback() + thread.join(3) + self.assertTrue(done.is_set()) + self.assertEqual(result, [None]) + self.released() + + def test_add_waiting_for_same_engine_update_does_not_own_model(self): + self.seed((1,)) + engine = self.engine() + update_ready, add_waiting = threading.Event(), threading.Event() + underlying = engine._write_lock + + class EngineLock: + def __enter__(self): + if threading.current_thread().name == "synthetic-add": + add_waiting.set() + underlying.acquire() + return self + + def __exit__(self, *_args): + underlying.release() + + engine._write_lock = EngineLock() + original = engine.update.__globals__["update_message"] + + def update(conn, memory_id, **fields): + result = original(conn, memory_id, **fields) + update_ready.set() + self.assertTrue(add_waiting.wait(3)) + self.released() + return result + + with patch.dict(engine.update.__globals__, update_message=update): + update_thread, updated, update_done = self.start(lambda: engine.update(1, "synthetic updated"), "synthetic-update") + self.assertTrue(update_ready.wait(2)) + add_thread, added, add_done = self.start(lambda: engine.add("synthetic added"), "synthetic-add") + update_thread.join(3) + add_thread.join(3) + self.assertTrue(update_done.is_set()) + self.assertTrue(add_done.is_set()) + self.assertIsInstance(updated[0], dict, updated) + self.assertIsInstance(added[0], dict, added) + self.assertEqual(self.ids(), [1, 2]) + self.assertEqual(self.ids("vec_messages_sep"), [1, 2]) + + def test_same_engine_writer_finishes_before_capture_classifies_transaction(self) -> None: + self.seed((1,)) + for operation in ("add", "update"): + with self.subTest(operation=operation): + engine = self.engine() + writer_ready, waiting, release = (threading.Event() for _ in range(3)) + underlying = engine._write_lock + + class EngineLock: + def __enter__(self) -> "EngineLock": + if threading.current_thread().name == "synthetic-capture": + waiting.set() + underlying.acquire() + return self + + def __exit__(self, *_args: object) -> None: + underlying.release() + + engine._write_lock = EngineLock() + self.vector._model = None + loads = [] + + def load() -> object: + self.assertFalse(self.conn.in_transaction) + self.assertTrue(underlying.locked()) + loads.append(True) + self.vector._model = self.model + return self.model + + def writer() -> None: + with engine._write_lock: + self.conn.execute("BEGIN IMMEDIATE") + self.conn.execute("INSERT INTO metadata(key,value) VALUES (?, 'pending')", + ("synthetic-owner-" + operation,)) + writer_ready.set() + self.assertTrue(release.wait(3)) + self.conn.commit() + + with patch.object(self.vector, "get_model", load): + writer_thread, written, writer_done = self.start(writer) + self.assertTrue(writer_ready.wait(2)) + action = (lambda: engine.add("synthetic added")) if operation == "add" else ( + lambda: engine.update(1, "synthetic updated") + ) + mutation_thread, result, done = self.start(action, "synthetic-capture") + try: + self.assertTrue(waiting.wait(2)) + self.assertFalse(done.is_set()) + self.assertEqual(loads, []) + finally: + release.set() + mutation_thread.join(3) + writer_thread.join(3) + self.assertTrue(writer_done.is_set()) + self.assertEqual(written, [None]) + self.assertTrue(done.is_set()) + self.assertIsInstance(result[0], dict, result) + self.assertEqual(loads, [True]) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_generation_change_during_model_load_rejects_add_and_update_before_source_write(self): + self.seed((1,)) + engine = self.engine() + + def changed(): + self.vector._model_generation += 1 + return self.model + + with patch.object(self.vector, "get_model", changed): + for action in (lambda: engine.add("synthetic rejected"), lambda: engine.update(1, "synthetic rejected")): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "changed while loading"): + action() + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.assertEqual(self.ids(), []) + self.released() + + def test_generation_change_while_waiting_rejects_before_opening_transaction(self): + captured = (0, self.vector.EMBEDDING_MODEL, 2) + attempted = threading.Event() + underlying = self.vector._lock + + class WaitingLock: + def acquire(self, blocking=True): + attempted.set() + return underlying.acquire(blocking=blocking) + + def release(self): + underlying.release() + + self.vector._lock = WaitingLock() + underlying.acquire() + + def publish(): + with self.vector._foreground_vector_publication(self.conn, captured, 1): + self.fail("Changed generation entered publication") + + thread, results, done = self.start(publish) + try: + self.assertTrue(attempted.wait(2)) + self.assertFalse(done.is_set()) + self.assertFalse(self.conn.in_transaction) + self.vector._model_generation += 1 + finally: + underlying.release() + thread.join(3) + self.assertIsInstance(results[0], self.vector.VectorPublicationChanged) + self.assertIn("model changed", str(results[0])) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_add_failure_after_source_insert_rolls_back_source_and_vectors(self): + engine = self.engine() + with patch.object(self.vector, "_validate_foreground_vector_target", side_effect=self.vector.VectorPublicationChanged("synthetic target changed")): + with self.assertRaises(self.vector.VectorPublicationChanged): + engine.add("synthetic rollback") + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.assertEqual(self.ids("vec_messages_sep"), []) + self.assertFalse(self.conn.in_transaction) + self.released() + + def test_admission_begin_and_commit_failures_release_model(self): + real = self.conn + + class FailedConnection: + def __init__(self, phase): + self.phase = phase + + def __getattr__(self, name): + return getattr(real, name) + + def execute(self, sql, *args): + if self.phase == "begin" and sql == "BEGIN IMMEDIATE": + raise sqlite3.OperationalError("synthetic writer admission") + return real.execute(sql, *args) + + def commit(self): + if self.phase == "commit": + raise sqlite3.OperationalError("synthetic commit failure") + return real.commit() + + for phase in ("begin", "commit"): + with self.subTest(phase=phase): + with self.assertRaises(sqlite3.OperationalError): + self.engine(FailedConnection(phase)).add("synthetic rollback") + self.assertFalse(real.in_transaction) + self.assertEqual(real.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.released() + + def test_owned_model_load_and_encoding_finish_before_writer(self): + def load(): + self.assertFalse(self.conn.in_transaction) + return self.model + + def encode(texts): + self.assertFalse(self.conn.in_transaction) + self.released() + return [[1., 0.] for _ in texts] + + self.on_encode = encode + with patch.object(self.vector, "get_model", load): + self.engine().add("synthetic outside writer") + self.assertEqual(self.ids(), [1]) + + +class TestBorrowedAdmission(AdmissionFixture): + def actions(self): + engine = self.engine() + return ( + ("add", lambda: engine.add("synthetic rejected")), + ("update", lambda: engine.update(1, "synthetic rejected")), + ("embed_single", lambda: self.vector.embed_single(self.conn, 1, "synthetic-1")), + ("completion", lambda: self.vector.build_vectors(self.conn)), + ("separation", lambda: self.vector.build_separation_vectors(self.conn)), + ) + + def test_busy_capture_defers_every_mutation_without_waiting_or_loading(self): + self.seed((1,)) + self.sentinel() + calls = [] + + class BusyLock: + def acquire(self, blocking=True): + calls.append(blocking) + if blocking: + raise AssertionError("Borrowed caller attempted a blocking model acquisition") + return False + + with patch.object(self.vector, "_lock", BusyLock()), patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")): + for name, action in self.actions(): + with self.subTest(action=name): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model is busy"): + action() + self.preserved() + self.assertEqual(calls, [False] * 5) + self.assertEqual(self.calls, []) + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.assertEqual(self.ids(), []) + self.conn.rollback() + self.released() + + def test_unloaded_capture_defers_every_mutation_and_preserves_sentinel(self): + self.seed((1,)) + self.sentinel() + self.vector._model = None + with patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")) as loader: + for name, action in self.actions(): + with self.subTest(action=name): + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "not loaded"): + action() + self.preserved() + self.released() + loader.assert_not_called() + self.assertEqual(self.calls, []) + self.assertEqual(self.conn.execute("SELECT id,content FROM messages").fetchall(), [(1, "synthetic-1")]) + self.conn.rollback() + + def test_borrowed_loaded_snapshot_skips_loader_and_add_remains_rollbackable(self): + self.sentinel() + with patch.object(self.vector, "get_model", side_effect=AssertionError("Borrowed model load")) as loader: + item = self.engine().add("synthetic borrowed") + loader.assert_not_called() + self.preserved() + self.assertEqual(self.ids(), [item["id"]]) + self.assertEqual(self.ids("vec_messages_sep"), [item["id"]]) + self.conn.rollback() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + self.assertEqual(self.ids(), []) + self.released() + + def test_borrowed_publication_contention_preserves_caller_transaction(self): + self.sentinel() + held, release = threading.Event(), threading.Event() + + def hold(): + with self.vector._lock: + held.set() + self.assertTrue(release.wait(3)) + + holder = [] + + def encode(texts): + if len(self.calls) == 2: + holder.append(self.start(hold)[0]) + self.assertTrue(held.wait(2)) + return [[1., 0.] for _ in texts] + + self.on_encode = encode + try: + with self.assertRaisesRegex(self.vector.VectorPublicationChanged, "model is busy"): + self.engine().add("synthetic rejected") + self.preserved() + self.assertEqual(self.conn.execute("SELECT count(*) FROM messages").fetchone()[0], 0) + finally: + release.set() + for thread in holder: + thread.join(3) + self.conn.rollback() + self.released() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test-rebuild-foreground-publication-756.py b/tests/test-rebuild-foreground-publication-756.py index 6c71ac5e..5e8b0b7d 100644 --- a/tests/test-rebuild-foreground-publication-756.py +++ b/tests/test-rebuild-foreground-publication-756.py @@ -273,7 +273,7 @@ def __getattr__(self, 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()) + case.assertTrue(case.vector._lock.locked()) boundaries.append("writer") if sql == "RELEASE SAVEPOINT rebuild_source": case.assertTrue(case.vector._lock.locked()) @@ -296,13 +296,13 @@ def commit(self): 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): + def test_embed_caller_fence_precedes_writer_upgrade_and_contains_release(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): + def test_add_caller_fence_precedes_writer_upgrade_and_contains_release(self): self._assert_foreground_writer_and_terminal_fence("add", True) def test_denied_add_savepoint_release_rolls_back_only_add(self): diff --git a/tests/test-rebuild-source-streaming-756.py b/tests/test-rebuild-source-streaming-756.py index 369304e5..cf3ca3e8 100644 --- a/tests/test-rebuild-source-streaming-756.py +++ b/tests/test-rebuild-source-streaming-756.py @@ -78,6 +78,7 @@ def setUp(self): 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._model = self.model self.vector.get_model = lambda: self.model self.calls = [] self.on_encode = None diff --git a/tests/test_issue_511_write_path_opt.py b/tests/test_issue_511_write_path_opt.py index c2f1cbe4..0d1858a3 100644 --- a/tests/test_issue_511_write_path_opt.py +++ b/tests/test_issue_511_write_path_opt.py @@ -5,7 +5,9 @@ M16: Fallback embed_single inside write lock — removed in favor of skip-and-log """ +import ast import inspect +import textwrap import unittest @@ -30,13 +32,17 @@ class TestM15StyleVecOutsideLock(unittest.TestCase): def test_add_precomputes_style_vec(self): from truememory.engine import TrueMemoryEngine - source = inspect.getsource(TrueMemoryEngine.add) - lock_pos = source.find("self._write_lock") - precompute_pos = source.find("compute_style_vector") - self.assertGreater(precompute_pos, -1, - "add() should reference compute_style_vector for pre-computation") - self.assertLess(precompute_pos, lock_pos, - "compute_style_vector call should appear before _write_lock") + tree = ast.parse(textwrap.dedent(inspect.getsource(TrueMemoryEngine.add))) + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) and node.func.id == "compute_style_vector"] + self.assertEqual(len(calls), 1) + for block in (node for node in ast.walk(tree) if isinstance(node, ast.With)): + self.assertNotIn(calls[0], list(ast.walk(block)), + "Style computation must run outside lock and transaction contexts") + publications = [node for node in ast.walk(tree) if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) and node.func.id == "_update_style_vec"] + self.assertEqual(len(publications), 1) + self.assertLess(calls[0].lineno, publications[0].lineno) def test_style_vec_accepts_precomputed(self): from truememory.personality_style_vec import update_entity_style_vector_incremental diff --git a/truememory/engine.py b/truememory/engine.py index 0285b066..9e61f69f 100644 --- a/truememory/engine.py +++ b/truememory/engine.py @@ -724,14 +724,18 @@ def add( pre_embedding = None pre_sep_embedding = None if self._has_vectors: + from truememory.vector_search import VectorPublicationChanged try: from truememory.vector_search import ( _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model, pre_identity = _capture_rebuild_model() + with self._write_lock: + model, pre_identity = _capture_rebuild_model(self.conn) 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] + except VectorPublicationChanged: + raise except Exception: logger.debug("Failed to pre-compute embedding during add()", exc_info=True) @@ -743,13 +747,15 @@ def add( except Exception: logger.debug("Failed to pre-compute style vector during add()", exc_info=True) + from contextlib import nullcontext 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): + with self._write_lock, ( + _foreground_model_fence(pre_identity, blocking=not self.conn.in_transaction) + if pre_embedding is not None else nullcontext() + ), rebuild_transaction(self.conn, write=True): msg = { "content": content, "sender": sender, @@ -1176,11 +1182,13 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | pre_embedding = None pre_sep_embedding = None if content is not None and self._has_vectors: + from truememory.vector_search import VectorPublicationChanged try: from truememory.vector_search import ( _capture_rebuild_model, _encode_with_mps_fallback, _build_sep_text, ) - model, pre_identity = _capture_rebuild_model() + with self._write_lock: + model, pre_identity = _capture_rebuild_model(self.conn) pre_embedding = _encode_with_mps_fallback(model, [content])[0] row = self.conn.execute( "SELECT sender, recipient, timestamp FROM messages WHERE id = ?", @@ -1192,6 +1200,8 @@ def update(self, memory_id: int, content: str | None = None, **fields) -> dict | timestamp_val = fields.get("timestamp", row[2]) sep_text = _build_sep_text(sender_val, recipient_val, timestamp_val, content) pre_sep_embedding = _encode_with_mps_fallback(model, [sep_text])[0] + except VectorPublicationChanged: + raise except Exception: logger.debug("Failed to pre-compute embedding during update()", exc_info=True) diff --git a/truememory/vector_search.py b/truememory/vector_search.py index 507e4bb9..965ff0ab 100644 --- a/truememory/vector_search.py +++ b/truememory/vector_search.py @@ -739,13 +739,23 @@ def _encode_with_mps_fallback(model, texts, **kwargs): """ -def _capture_rebuild_model() -> tuple[object, tuple[int, str, int]]: +def _capture_rebuild_model(conn: sqlite3.Connection | None = None) -> tuple[object, tuple[int, str, int]]: + """Caller transactions can borrow a cached model, but cannot wait or load.""" + if conn is not None and conn.in_transaction: + if not _lock.acquire(blocking=False): + raise VectorPublicationChanged("Embedding model is busy; retry outside the caller transaction") + try: + if _model is None: + raise VectorPublicationChanged("Embedding model is not loaded; load it outside the caller transaction") + return _model, (_model_generation, EMBEDDING_MODEL, _embedding_dim) + finally: + _lock.release() 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") + raise VectorPublicationChanged("Embedding model changed while loading rebuild model") return model, (_model_generation, EMBEDDING_MODEL, _embedding_dim) @@ -767,8 +777,9 @@ class VectorPublicationChanged(RuntimeError): @contextmanager -def _foreground_model_fence(identity: tuple[int, str, int]) -> Iterator[None]: - if not _lock.acquire(blocking=False): +def _foreground_model_fence(identity: tuple[int, str, int], *, blocking: bool = False) -> Iterator[None]: + """Blocking admission is permitted only before an owned writer starts.""" + if not _lock.acquire(blocking=blocking): raise VectorPublicationChanged("Embedding model is busy; retry vector publication") try: if identity != (_model_generation, EMBEDDING_MODEL, _embedding_dim): @@ -815,7 +826,7 @@ def _foreground_vector_publication( ) -> Iterator[tuple[str, str]]: from truememory.rebuild_source import rebuild_transaction - with rebuild_transaction(conn, write=True, publication_fence=_foreground_model_fence(identity)): + with _foreground_model_fence(identity, blocking=not conn.in_transaction), rebuild_transaction(conn, write=True): yield _validate_foreground_vector_target(conn, identity, message_id) @@ -856,7 +867,7 @@ def _build_streamed_vectors( save_manifest(conn, key, empty) _clear_build_in_progress(conn, table) return 0 - model, identity = _capture_rebuild_model() + model, identity = _capture_rebuild_model(conn) 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 @@ -1191,7 +1202,7 @@ 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, identity = _capture_rebuild_model() + model, identity = _capture_rebuild_model(conn) embedding = _encode_with_mps_fallback(model, [content])[0] normed = _normalize_for_cosine(embedding) if normed is None: