Skip to content

Commit 7553dd5

Browse files
phernandezclaude
andcommitted
fix(core): embed materialized notes from their own keyed pool
Awaiting the embed inside the materialization worker held that worker for the embed's duration, so a slow embedding provider delayed unrelated notes' file writes. Embeds now go to a separate bounded pool keyed by note, like cloud's separate index_embeddings job: one note's embeds still run one at a time in write order, and drain_pending_materializations waits for them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01APFUk2bjEwMptMQqpRhjea Signed-off-by: phernandez <paul@basicmachines.co>
1 parent 070a645 commit 7553dd5

2 files changed

Lines changed: 83 additions & 13 deletions

File tree

‎src/basic_memory/index/note_content_materialization.py‎

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -156,18 +156,26 @@ async def aclose(self) -> None:
156156

157157

158158
_materialization_pool = _MaterializationWorkerPool()
159+
# Embeds notes after their materialized files are indexed (#1732). A separate pool,
160+
# mirroring cloud's separate index_embeddings job: a slow embedding provider must not
161+
# hold a materialization worker and delay unrelated notes' file writes. Keyed by note
162+
# like the materialization pool, so one note's embeds still run one at a time in
163+
# write order.
164+
_vector_sync_pool = _MaterializationWorkerPool()
159165

160166

161167
async def drain_pending_materializations() -> None:
162-
"""Block until queued local materializations finish writing + indexing.
168+
"""Block until queued local materializations finish writing, indexing and embedding.
163169
164170
One-shot clients (``bm tool write-note``, importers) return right after the
165171
accept enqueues the markdown write/index; without this drain the event loop can
166172
close before the worker writes the source-of-truth file, silently losing the
167173
write even though the API already reported it accepted. Long-lived servers keep
168-
the loop alive and don't need it.
174+
the loop alive and don't need it. Materializations queue the embeds, so the
175+
materialization pool drains first.
169176
"""
170177
await _materialization_pool.join()
178+
await _vector_sync_pool.join()
171179

172180

173181
# --- Startup Recovery ---
@@ -677,9 +685,9 @@ class LocalNoteContentMaterializationProvider:
677685
# target (#1351). Optional so providers wired without it (cloud, focused
678686
# tests) keep working; cloud relies on its orphan sweeper instead.
679687
relation_cleanup_refresher: ProjectIndexMovedEntitySearchRefresher | None = None
680-
# Embeds the note once its materialized file is indexed (#1732). Optional so
681-
# providers wired without semantic work (cloud, focused tests) keep working;
682-
# cloud embeds from its materialize_note_file job instead.
688+
# Embeds the note once its materialized file is indexed (#1732), through the
689+
# vector pool. Optional so providers wired without semantic work (cloud, focused
690+
# tests) keep working; cloud embeds from its materialize_note_file job instead.
683691
entity_vector_sync: EntityVectorSync | None = None
684692

685693
async def materialize_write_change(
@@ -827,27 +835,48 @@ async def _materialize_write_now(
827835
or "note-content-materialization",
828836
),
829837
)
830-
await self._sync_indexed_note_vectors(accepted.materialization.entity_id)
838+
await self._embed_indexed_note(accepted.materialization)
831839
return indexed
832840
return accepted
833841

834-
async def _sync_indexed_note_vectors(self, entity_id: int) -> None:
842+
async def _embed_indexed_note(self, materialization: RuntimePendingNoteMaterialization) -> None:
835843
"""Embed a note from the search rows its materialized file just produced.
836844
837845
Trigger: the accepted note's file was written and indexed.
838846
Why: the accept path only publishes the note body as a temporary search row;
839847
observations and relations exist only after this index. A request-time
840848
embed raced this job, embedded the body alone, and nothing re-embedded
841-
the rest (#1732). This worker runs one note's jobs in order, so each
842-
accepted write embeds exactly once, after its own index. Cloud does the
843-
same from its materialize_note_file job.
844-
Outcome: vectors match the indexed note. The outer invalidation scope
845-
retires reads cached while vectors were stale.
849+
the rest (#1732). Cloud likewise queues index_embeddings from its
850+
materialize_note_file job.
851+
Outcome: one embed per accepted write, queued after its own index. Test mode
852+
runs it inline so tests can assert vectors synchronously; production
853+
hands it to the vector pool so the materialization worker is free for
854+
the next file write.
846855
"""
847856
if self.entity_vector_sync is None:
848857
return
858+
if self.test_mode:
859+
await self._sync_note_vectors(materialization.entity_id)
860+
return
861+
_vector_sync_pool.submit(
862+
self._sync_note_vectors(materialization.entity_id),
863+
workers=self.materialization_workers,
864+
key=(materialization.project_id, materialization.entity_id),
865+
)
866+
867+
async def _sync_note_vectors(self, entity_id: int) -> None:
868+
if self.entity_vector_sync is None: # pragma: no cover - guarded by caller
869+
return
870+
# Vector publication can change VECTOR/HYBRID results after the note's own
871+
# invalidation, so retire reads cached while vectors were stale.
872+
invalidation_scope = (
873+
invalidate_cache(self.read_cache, self.project_external_id)
874+
if self.read_cache is not None
875+
else nullcontext()
876+
)
849877
try:
850-
await self.entity_vector_sync.sync_entity_vectors(entity_id)
878+
async with invalidation_scope:
879+
await self.entity_vector_sync.sync_entity_vectors(entity_id)
851880
except Exception:
852881
# Vectors are derived state: the accepted note and its file are already
853882
# published, and the next write or a reindex re-embeds the note. Keep

‎tests/cloud/test_note_content_materialization.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,47 @@ async def test_local_materialization_embeds_once_after_index(
595595
assert result.payload == indexed_payload
596596

597597

598+
@pytest.mark.asyncio
599+
async def test_production_materialization_hands_embed_to_vector_pool(
600+
monkeypatch: pytest.MonkeyPatch,
601+
) -> None:
602+
"""A slow embed must not hold the materialization worker (#1732 review).
603+
604+
The file write finishes and returns while the embed waits in the vector pool;
605+
the shutdown drain still waits for that embed.
606+
"""
607+
accepted = accepted_materialization_change()
608+
written_materialization(monkeypatch, accepted.payload)
609+
vector_pool = note_content_materialization._MaterializationWorkerPool()
610+
monkeypatch.setattr(note_content_materialization, "_vector_sync_pool", vector_pool)
611+
indexer = RecordingFileIndexer()
612+
release_embed = asyncio.Event()
613+
embedded: list[int] = []
614+
615+
class SlowVectorSync:
616+
async def sync_entity_vectors(self, entity_id: int) -> VectorSyncBatchResult:
617+
await release_embed.wait()
618+
embedded.append(entity_id)
619+
return VectorSyncBatchResult(entities_total=1, entities_synced=1, entities_failed=0)
620+
621+
read_cache = RecordingReadCache()
622+
provider = replace(
623+
local_materialization_provider(indexer, test_mode=False, read_cache=read_cache),
624+
entity_vector_sync=SlowVectorSync(),
625+
)
626+
627+
await provider._materialize_write_now(accepted)
628+
assert indexer.calls == [("notes/test.md", "note-content-materialization")]
629+
assert embedded == []
630+
631+
release_embed.set()
632+
await note_content_materialization.drain_pending_materializations()
633+
assert embedded == [42]
634+
# Two invalidations close the materialization; the embed adds its own.
635+
assert read_cache.invalidated_project_ids == [PROJECT_EXTERNAL_ID] * 3
636+
await vector_pool.aclose()
637+
638+
598639
@pytest.mark.asyncio
599640
async def test_local_materialization_keeps_indexed_result_when_embed_fails(
600641
monkeypatch: pytest.MonkeyPatch,

0 commit comments

Comments
 (0)