Skip to content

fix(memory-engine): compute update_memory_unit embeddings off the pooled connection - #3083

Merged
nicoloboschi merged 1 commit into
mainfrom
fix/update-memory-unit-off-connection-embed
Jul 31, 2026
Merged

fix(memory-engine): compute update_memory_unit embeddings off the pooled connection#3083
nicoloboschi merged 1 commit into
mainfrom
fix/update-memory-unit-off-connection-embed

Conversation

@nicoloboschi

Copy link
Copy Markdown
Collaborator

Draft — needs live-PostgreSQL validation before review/merge. Lint + ty are clean and the change imports, but this is a concurrency rewrite of the interactive-edit path and I have not run the live-DB curation suite locally (no database available in my environment). Do not merge until test_memory_curation.py (+ test_observation_invalidation.py, test_document_transfer.py) pass green in CI against real PostgreSQL.

Problem

update_memory_unit re-embedded inside the write transaction (edit path and revert path), so an interactive curation edit held a pooled PostgreSQL connection checked out across the embedder call. Same saturation class as #3082 — the connection sits idle mid-embed while other work blocks on acquire.

This is the follow-up promised in #3082, which fixed the two paths that didn't need to widen the read→write window (update_mental_model + consolidation). update_memory_unit is harder because #2917 rewrote it around the pluggable memories store (apply_edit/restore_memory + the begin_txn/decide_txn cross-store write-group), so the embed-off-connection split has to be reconciled with that coordinator.

Approach — two phases

Phase 1 (no write transaction held): read the live/archived record; for an entity-changing edit, resolve the corrected set (resolve_entities_only, idempotent find-or-create, autocommitted on a short connection) and read its canonical names; for revert, read the archive's entity_ids snapshot names. Then compute the embedding with no connection held, into a typed _MemoryEditPlan / _MemoryRevertPlan.

Phase 2 (short write transaction):

  • Re-read the row; skip cleanly (return None) if it was concurrently moved/deleted between the phases.
  • Reassert + link the resolved entities (reassert-before-link keeps the retain FK race: prune_orphan_entities deletes entities between retain Phase-1 resolve and Phase-2 unit_entities insert — silent memory loss on re-ingest (v0.8.4) #2662 prune race safe), apply the field edit / archive move through the store, and write the precomputed embedding.
  • Bounded in-txn re-embed fallback: when entities are not part of the edit, the current linked names are re-read under the write lock; only on a mismatch (a concurrent entity-only edit landed between the phases) is the vector re-embedded in-txn, so it stays consistent with the committed unit_entities. The revert path does the same if the surviving entity set differs from the Phase-1 snapshot.
  • The store's begin_txn/decide_txn write-group is opened inside the Phase-2 transaction and decided once Postgres commits.

Orphan reclaim: entities resolved (and autocommitted) in Phase 1 whose edit didn't durably apply — row concurrently invalidated, or Phase 2 raised — are reclaimed by forcing a graph-maintenance sweep in the finally block.

Why no hard abort/retry

Unlike the original #2434 raw-SQL design (SELECT … FOR UPDATE + abort/retry), this relies on last-writer-wins for concurrent text/date edits — which is already the current single-transaction semantics — and only needs the bounded in-txn re-embed for the one case that would otherwise corrupt (embedding naming a stale entity set). This avoids threading a re-lock/abort/retry loop through the store's transaction coordinator. Open to switching to strict abort/retry if reviewers prefer a "never embed in-transaction" guarantee.

Validation status

  • ruff + ty clean; module imports; the 23 existing test_memory_curation.py cases (edit / revert / invalidate / entity-swap round-trips) collect and will exercise this end-to-end in CI on live PostgreSQL — they are the regression gate for this refactor.
  • Not run locally. Reviewer/CI must confirm green before this leaves draft.

Refs #2434

…led connection

update_memory_unit re-embedded inside the write transaction, so an interactive
curation edit (or revert) held a pooled connection across the embedder call.
This splits it into two phases against the pluggable memories store:

- Phase 1 (no write transaction held): read the live/archived record, resolve
  the corrected entity set (idempotent find-or-create, autocommitted on a short
  connection), and read the canonical entity names — then compute the embedding
  with NO connection held. Carried into typed _MemoryEditPlan / _MemoryRevertPlan.

- Phase 2 (short write transaction): re-read the row (skip cleanly if it was
  concurrently moved/deleted), reassert + link the resolved entities, apply the
  field edit / archive move through the store, and write the precomputed
  embedding. When entities are NOT part of the edit, the current linked names are
  re-read under the write lock and, only on a mismatch (a concurrent entity-only
  edit landed between the phases), the vector is re-embedded in-txn so it stays
  consistent with the committed unit_entities. The revert path re-embeds in-txn
  the same way if the surviving entity set differs from the Phase-1 snapshot.

- Orphan reclaim: entities resolved (autocommitted) in Phase 1 whose edit did not
  durably apply — row concurrently invalidated, or Phase 2 raised — are reclaimed
  by forcing a graph-maintenance sweep in the finally block.

The store's begin_txn/decide_txn cross-store write-group is preserved: it is
opened inside the Phase-2 transaction and decided once Postgres commits.

Follow-up to #3082 (which handled the consolidation + update_mental_model paths).
Refs #2434

NOT YET VALIDATED against a live database — see PR description.
@nicoloboschi
nicoloboschi marked this pull request as ready for review July 31, 2026 16:08
@nicoloboschi nicoloboschi changed the title fix(memory-engine): compute update_memory_unit embeddings off the pooled connection [DRAFT] fix(memory-engine): compute update_memory_unit embeddings off the pooled connection Jul 31, 2026
@nicoloboschi
nicoloboschi merged commit d571199 into main Jul 31, 2026
102 of 103 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant