From 011abc9485ac010260fd230d39049ba00cf7c1e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:23:54 +0000 Subject: [PATCH 1/2] Fix #97: move blocking embed out of SQLite write transactions The five write paths listed in #97 all opened an implicit BEGIN (or a SAVEPOINT) and then called SyncEmbedder inside it. In WAL mode that holds the single writer lock across a multi-second Ollama call, so any other connection times out on its 5s busy_timeout with 'database is locked'. Fix: compute the embedding BEFORE the write lock is taken, then pass the vector into _store_embedding once inside the transaction. For apply_decision, which fans out N embeds across _apply_new / _apply_augment inside one SAVEPOINT, the two helpers now collect (reflection_id, source_text) into an optional embed_tasks list; the caller commits first, then does one batched embed and a short second write. Missing vectors already self-heal on retrieval via _heal_missing_embeddings, preserving the existing best-effort contract. --- better_memory/services/reflection.py | 94 +++++++++++---- better_memory/services/semantic.py | 25 ++-- tests/services/test_reflection.py | 2 +- .../test_reflection_embedding_write.py | 107 ++++++++++++++++++ 4 files changed, 196 insertions(+), 32 deletions(-) diff --git a/better_memory/services/reflection.py b/better_memory/services/reflection.py index a38ad00..2fb07b0 100644 --- a/better_memory/services/reflection.py +++ b/better_memory/services/reflection.py @@ -773,7 +773,11 @@ def parse_response_dict(self, data: object) -> SynthesisResponse: # ---------------------------------------------------------------- _apply_new def _apply_new( - self, actions: list[NewAction], *, project: str + self, + actions: list[NewAction], + *, + project: str, + embed_tasks: list[tuple[str, str]] | None = None, ) -> int: """Insert new reflections + their source links + consume observations. @@ -781,6 +785,12 @@ def _apply_new( don't exist in the DB are dropped. Entries whose entire source list turns out to be invalid are skipped silently. + When ``embed_tasks`` is supplied, the blocking embed call is NOT + made here; instead ``(reflection_id, source_text)`` pairs are + appended for the caller to embed AFTER the write transaction + commits (see :meth:`apply_decision` and #97). When ``None``, the + embed still runs inline for direct callers. + Returns the count of reflections actually inserted (may be smaller than ``len(actions)`` when entries are dropped for lack of valid sources). @@ -835,12 +845,16 @@ def _apply_new( ) if self._sync_embedder is not None: - self._store_embedding( - reflection_id, - self._sync_embedder.embed_text(_embedding_source_text( - action.title, action.use_cases, action.hints, - )), + source_text = _embedding_source_text( + action.title, action.use_cases, action.hints, ) + if embed_tasks is not None: + embed_tasks.append((reflection_id, source_text)) + else: + self._store_embedding( + reflection_id, + self._sync_embedder.embed_text(source_text), + ) created += 1 return created @@ -887,7 +901,11 @@ def _filter_existing_observations( return [i for i in ids if i in existing] # ----------------------------------------------------------- _apply_augment - def _apply_augment(self, actions: list[AugmentAction]) -> int: + def _apply_augment( + self, + actions: list[AugmentAction], + embed_tasks: list[tuple[str, str]] | None = None, + ) -> int: """Apply augment actions: append hints, rewrite use_cases, bump confidence, link new sources, recompute evidence count. @@ -898,6 +916,12 @@ def _apply_augment(self, actions: list[AugmentAction]) -> int: - ``add_source_observation_ids`` filtered to existing obs; ``INSERT OR IGNORE`` dedupes against existing source rows. + When ``embed_tasks`` is supplied, the blocking embed call is NOT + made here; instead ``(reflection_id, source_text)`` pairs are + appended for the caller to embed AFTER the write transaction + commits (see :meth:`apply_decision` and #97). When ``None``, the + embed still runs inline for direct callers. + Returns the count of reflections actually augmented. """ augmented = 0 @@ -997,12 +1021,16 @@ def _apply_augment(self, actions: list[AugmentAction]) -> int: final_use_cases = (action.rewrite_use_cases if action.rewrite_use_cases is not None else row["use_cases"]) - self._store_embedding( - action.reflection_id, - self._sync_embedder.embed_text(_embedding_source_text( - row["title"], final_use_cases, merged_hints, - )), + source_text = _embedding_source_text( + row["title"], final_use_cases, merged_hints, ) + if embed_tasks is not None: + embed_tasks.append((action.reflection_id, source_text)) + else: + self._store_embedding( + action.reflection_id, + self._sync_embedder.embed_text(source_text), + ) augmented += 1 return augmented @@ -1277,6 +1305,12 @@ def apply_decision( f"(synthesized_at={row['synthesized_at']})" ) + # Collect embed inputs during the DB writes so the blocking Ollama + # call happens AFTER commit, outside the WAL writer lock — see #97. + # Missing vectors self-heal on first retrieval via + # _heal_missing_embeddings, so a crash between the two commits + # only costs one round of embedding on the next lookup. + embed_tasks: list[tuple[str, str]] = [] self._conn.execute("SAVEPOINT episode_synthesize") try: active_rows = self._conn.execute( @@ -1286,8 +1320,12 @@ def apply_decision( ).fetchall() active_ids = [r["id"] for r in active_rows] - created = self._apply_new(response.new, project=project) - augmented = self._apply_augment(response.augment) + created = self._apply_new( + response.new, project=project, embed_tasks=embed_tasks, + ) + augmented = self._apply_augment( + response.augment, embed_tasks=embed_tasks, + ) merged = self._apply_merge(response.merge) ignored = self._apply_ignore(response.ignore) auto_ignored = self._auto_ignore_unused(active_ids) @@ -1300,6 +1338,17 @@ def apply_decision( self._conn.execute("RELEASE SAVEPOINT episode_synthesize") self._conn.commit() + # Now that the main write transaction is committed, run the batched + # embed call. A slow or dead Ollama can no longer stall other + # connections; the vector write is a second, short transaction. + if embed_tasks and self._sync_embedder is not None: + texts = [t for _, t in embed_tasks] + vectors = self._sync_embedder.embed_batch(texts) + if vectors is not None: + for (rid, _), vec in zip(embed_tasks, vectors): + self._store_embedding(rid, vec) + self._conn.commit() + counts = { "created": created, "augmented": augmented, @@ -1771,19 +1820,20 @@ def update_text( f"Cannot edit reflection in status {status!r}" ) now = self._clock().isoformat() + # Compute the embedding BEFORE the UPDATE opens sqlite3's implicit + # write transaction — see #97. Ollama is blocking and would otherwise + # hold the WAL writer lock across a multi-second network call. + vector = None + if self._sync_embedder is not None: + vector = self._sync_embedder.embed_text(_embedding_source_text( + row["title"], use_cases, hint_list, + )) self._conn.execute( "UPDATE reflections SET use_cases = ?, hints = ?, updated_at = ? " "WHERE id = ?", (use_cases, json.dumps(hint_list), now, reflection_id), ) - if self._sync_embedder is not None: - _write_reflection_embedding( - self._conn, - reflection_id, - self._sync_embedder.embed_text(_embedding_source_text( - row["title"], use_cases, hint_list, - )), - ) + _write_reflection_embedding(self._conn, reflection_id, vector) self._conn.commit() def promote_to_general(self, *, reflection_id: str) -> None: diff --git a/better_memory/services/semantic.py b/better_memory/services/semantic.py index 68c0030..bea9441 100644 --- a/better_memory/services/semantic.py +++ b/better_memory/services/semantic.py @@ -88,6 +88,11 @@ def create( raise ValueError("content must not be empty") memory_id = uuid4().hex now = self._clock().isoformat() + # Compute the embedding BEFORE the INSERT opens sqlite3's implicit + # write transaction — see #97. A blocking Ollama call inside the + # WAL writer lock starves every other connection until commit. + vector = (self._sync_embedder.embed_text(content) + if self._sync_embedder is not None else None) self._conn.execute( """ INSERT INTO semantic_memories @@ -96,9 +101,7 @@ def create( """, (memory_id, content, project, scope, now, now), ) - if self._sync_embedder is not None: - self._store_embedding( - memory_id, self._sync_embedder.embed_text(content)) + self._store_embedding(memory_id, vector) self._conn.commit() return memory_id @@ -106,6 +109,10 @@ def update_text(self, *, id: str, content: str) -> None: if not content.strip(): raise ValueError("content must not be empty") now = self._clock().isoformat() + # Compute the embedding BEFORE the UPDATE opens sqlite3's implicit + # write transaction — see #97. + vector = (self._sync_embedder.embed_text(content) + if self._sync_embedder is not None else None) cur = self._conn.execute( "UPDATE semantic_memories SET content = ?, updated_at = ? " "WHERE id = ?", @@ -118,9 +125,7 @@ def update_text(self, *, id: str, content: str) -> None: # ObservationService.set_outcome (better_memory/services/observation.py:435). self._conn.rollback() raise ValueError(f"semantic memory not found: {id}") - if self._sync_embedder is not None: - self._store_embedding( - id, self._sync_embedder.embed_text(content)) + self._store_embedding(id, vector) self._conn.commit() def set_scope(self, *, id: str, scope: str) -> None: @@ -182,6 +187,10 @@ def create_from_observation( memory_id = uuid4().hex now = self._clock().isoformat() + # Compute the embedding BEFORE opening the SAVEPOINT — see #97. + # The blocking Ollama call must not run under the WAL writer lock. + vector = (self._sync_embedder.embed_text(row["content"]) + if self._sync_embedder is not None else None) self._conn.execute("SAVEPOINT promote_observation") try: self._conn.execute( @@ -198,9 +207,7 @@ def create_from_observation( "WHERE id = ?", (now, observation_id), ) - if self._sync_embedder is not None: - self._store_embedding( - memory_id, self._sync_embedder.embed_text(row["content"])) + self._store_embedding(memory_id, vector) except BaseException: self._conn.execute("ROLLBACK TO SAVEPOINT promote_observation") self._conn.execute("RELEASE SAVEPOINT promote_observation") diff --git a/tests/services/test_reflection.py b/tests/services/test_reflection.py index 2e4b0e8..93bdaca 100644 --- a/tests/services/test_reflection.py +++ b/tests/services/test_reflection.py @@ -2072,7 +2072,7 @@ def test_db_integrity_error_propagates_and_no_synthesized_at( ) conn.commit() - def boom(self, actions, *, project): + def boom(self, actions, *, project, embed_tasks=None): raise sqlite3.IntegrityError("simulated FK violation") monkeypatch.setattr( diff --git a/tests/services/test_reflection_embedding_write.py b/tests/services/test_reflection_embedding_write.py index 35cabea..a3496a6 100644 --- a/tests/services/test_reflection_embedding_write.py +++ b/tests/services/test_reflection_embedding_write.py @@ -20,6 +20,7 @@ NewAction, ReflectionService, ReflectionSynthesisService, + SynthesisResponse, _embedding_source_text, ) from tests.services._embedding_fakes import FakeEmbedder @@ -322,6 +323,112 @@ def test_no_embed_row_when_embedder_fails(self, conn, fixed_clock): assert _vec_count(conn) == 0 +class TestApplyDecisionDefersEmbeddingOutsideWriteLock: + """apply_decision must not hold the writer lock across Ollama calls. + + #97: the blocking embed is now performed AFTER the SAVEPOINT commits. + These tests exercise the end-to-end flow (apply_decision → embeds) + and prove the deferred path writes vectors correctly. + """ + + def _make_episode(self, conn, fixed_clock): + epsvc = EpisodeService(conn, clock=fixed_clock) + ep = epsvc.start_foreground(session_id="s1", project="p", goal="g") + epsvc.close_active( + session_id="s1", outcome="success", close_reason="goal_complete" + ) + return ep + + def test_apply_decision_writes_embedding_for_new(self, conn, fixed_clock): + ep = self._make_episode(conn, fixed_clock) + _insert_obs(conn, obs_id="obs-1", project="p", episode_id=ep) + conn.commit() + + fake = FakeEmbedder() + svc = ReflectionSynthesisService( + conn, clock=fixed_clock, sync_embedder=SyncEmbedder(lambda: fake), + ) + response = SynthesisResponse( + new=[NewAction( + title="Always test", phase="general", polarity="do", + use_cases="when writing code", hints=["write tests first"], + tech="python", confidence=0.6, + source_observation_ids=["obs-1"], + )], + augment=[], merge=[], ignore=[], + ) + svc.apply_decision(episode_id=ep, response=response, project="p") + + assert _vec_count(conn) == 1 + # FakeEmbedder.embed_batch appends the incoming list to calls; the + # per-row embed_text path would append a bare str instead. Assert + # the one recorded call is the batched form. + assert len(fake.calls) == 1 + assert isinstance(fake.calls[0], list) + assert len(fake.calls[0]) == 1 + assert "Always test" in fake.calls[0][0] + + def test_apply_decision_writes_embedding_for_augment( + self, conn, fixed_clock + ): + ep = self._make_episode(conn, fixed_clock) + _insert_reflection( + conn, refl_id="r1", project="p", + title="Existing", use_cases="old uc", hints='["old-hint"]', + ) + conn.commit() + + fake = FakeEmbedder() + svc = ReflectionSynthesisService( + conn, clock=fixed_clock, sync_embedder=SyncEmbedder(lambda: fake), + ) + response = SynthesisResponse( + new=[], augment=[AugmentAction( + reflection_id="r1", + add_hints=["new-hint"], + rewrite_use_cases=None, + confidence_delta=0.0, + add_source_observation_ids=[], + )], merge=[], ignore=[], + ) + svc.apply_decision(episode_id=ep, response=response, project="p") + + assert _vec_count(conn) == 1 + assert len(fake.calls) == 1 + assert isinstance(fake.calls[0], list) + assert "new-hint" in fake.calls[0][0] + + def test_apply_decision_survives_embedder_failure(self, conn, fixed_clock): + """Embed failure never rolls back the reflection writes. + + This is the pre-existing best-effort contract; the writer-lock fix + must not silently strengthen it into an atomicity guarantee. + """ + ep = self._make_episode(conn, fixed_clock) + _insert_obs(conn, obs_id="obs-1", project="p", episode_id=ep) + conn.commit() + + svc = ReflectionSynthesisService( + conn, clock=fixed_clock, + sync_embedder=SyncEmbedder(lambda: FakeEmbedder(fail=True)), + ) + response = SynthesisResponse( + new=[NewAction( + title="t", phase="general", polarity="do", + use_cases="uc", hints=[], tech=None, confidence=0.5, + source_observation_ids=["obs-1"], + )], + augment=[], merge=[], ignore=[], + ) + svc.apply_decision(episode_id=ep, response=response, project="p") + + refl = conn.execute( + "SELECT id FROM reflections WHERE title = 't'" + ).fetchone() + assert refl is not None + assert _vec_count(conn) == 0 + + class TestApplyMergeEmbedding: def test_merge_deletes_source_embedding_no_reembed_of_target( self, conn, fixed_clock From c4bf4c9b3672f75b567388fe8f6ae41f1b02e0f3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:29:20 +0000 Subject: [PATCH 2/2] Validate id before embed in SemanticMemoryService.update_text BugBot review on PR #125 caught a regression: after moving the embed call before the UPDATE (the #97 fix), a call with a stale/deleted id now pays the full ~15s Ollama worker timeout before discovering the row is gone. update_text is reachable from MCP (memory.update) and the UI drawer with an id chosen by the caller, so the fast-fail path matters. Restore fast-fail by doing a SELECT 1 before the embed. The rowcount check on UPDATE still handles the (rare) race where the row is deleted between the SELECT and the UPDATE. --- better_memory/services/semantic.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/better_memory/services/semantic.py b/better_memory/services/semantic.py index bea9441..35ea8e4 100644 --- a/better_memory/services/semantic.py +++ b/better_memory/services/semantic.py @@ -108,6 +108,13 @@ def create( def update_text(self, *, id: str, content: str) -> None: if not content.strip(): raise ValueError("content must not be empty") + # Validate existence BEFORE spending the (blocking) embed. Without + # this check a stale id from MCP / UI would pay the full ~15s + # Ollama worker timeout before failing — see #125 review. + if self._conn.execute( + "SELECT 1 FROM semantic_memories WHERE id = ?", (id,), + ).fetchone() is None: + raise ValueError(f"semantic memory not found: {id}") now = self._clock().isoformat() # Compute the embedding BEFORE the UPDATE opens sqlite3's implicit # write transaction — see #97. @@ -119,9 +126,9 @@ def update_text(self, *, id: str, content: str) -> None: (content, now, id), ) if cur.rowcount == 0: - # No row updated — roll back the implicit BEGIN that sqlite3 - # opened before the UPDATE so we don't strand the WAL write - # lock for callers sharing this connection. Mirrors + # Race: row was deleted between our SELECT and UPDATE. Rare, + # but roll back the implicit BEGIN so we don't strand the WAL + # write lock for callers sharing this connection. Mirrors # ObservationService.set_outcome (better_memory/services/observation.py:435). self._conn.rollback() raise ValueError(f"semantic memory not found: {id}")