Skip to content

fix(engine): don't hold pooled DB connections across embedder/LLM calls (consolidation + mental model) - #3082

Merged
nicoloboschi merged 1 commit into
mainfrom
fix/db-conn-not-held-across-embed-2434
Jul 31, 2026
Merged

fix(engine): don't hold pooled DB connections across embedder/LLM calls (consolidation + mental model)#3082
nicoloboschi merged 1 commit into
mainfrom
fix/db-conn-not-held-across-embed-2434

Conversation

@nicoloboschi

Copy link
Copy Markdown
Collaborator

Problem

Several memory-engine paths held a pooled PostgreSQL connection checked out for the entire duration of a slow external call — an embedder or LLM round-trip. The pools are already bounded (asyncpg, default max_size=20) and per-process, so this is not a leak — it's saturation: enough concurrent operations park the pool on multi-second external calls and the poller / remaining work then block until acquire_timeout.

This is a reimplementation of the idea behind #2434 (which was opened before #2917 rewrote these paths around the pluggable memories store) onto current main.

What this PR fixes

update_mental_model — the embedding is now computed before acquiring a pooled connection. Its text depends only on the incoming name/content, never on DB state.

Consolidation_process_memory_batch and its executors / dedup helpers no longer receive a long-lived connection. Recall, the batch LLM call, every per-action embed, and dedup adjudication run with no connection held; each helper self-acquires a short-lived connection only around its own SQL.

Moving the slow calls off the connection widens the decision→write window, so the held-transaction serialization is replaced with explicit guards:

  • Each source-liveness check (FOR SHARE) is paired with its write in one short transaction, so a concurrent delete can't orphan an observation between check and write.
  • Dedup CREATE/UPDATE folds are RETURNING-gated and re-filter live sources inside the fold transaction (sources-before-observation lock order, matching the normal write paths). A twin deleted during the now connection-free window → the caller still CREATEs; a source deleted → it's dropped from the fold, never written back.
  • A cheap non-locking preflight restores the pre-refactor "skip before embedding when every source is already gone" short-circuit.

The separate-store (non-SQL) branches and the Oracle-safe search_vector clause (#3021) are preserved throughout.

Scope / deliberately deferred

update_memory_unit (and update_mental_model's sibling edit path) still embed inside the acquired connection, but their fix requires a two-phase read/embed-off-connection + short write transaction with re-lock/abort/retry that must be reconciled with the pluggable store's begin_txn/decide_txn cross-store transaction coordinator — and validated against a real database. That's a follow-up PR (linked once opened) rather than bundling an unvalidated concurrency rewrite of the interactive-edit path here.

Verification

  • ruff + ty clean.
  • 45 deterministic no-DB tests pass — they pin the dedup fold guards (RETURNING-gating, live-source filtering, twin-vanished → still-CREATE, created/skipped propagation) and the pre-embed short-circuits/rejections directly.
  • Live-DB consolidation / invalidation / document-transfer tests updated to the new short-acquire signatures, plus a new in-flight-delete race test (source deleted during the off-connection embed → the in-txn FOR SHARE guard skips the update). These run against live PostgreSQL in CI.

Refs #2434

Several memory-engine paths held a pooled PostgreSQL connection checked out
for the entire duration of a slow external call (embedder/LLM). The pools are
already bounded and per-process, so this is saturation, not a leak: enough
concurrent operations park the pool on multi-second calls and everything else
blocks on acquire.

This covers the two paths that can be fixed without widening the read→write
window unsafely:

- update_mental_model: compute the embedding BEFORE acquiring a connection.
  The embedding text depends only on the incoming name/content, never on DB
  state, so it needs no connection.

- consolidation: _process_memory_batch and its executors/dedup helpers no
  longer receive a long-lived connection. Recall, the batch LLM call, every
  per-action embed, and dedup adjudication run with NO connection held; each
  helper self-acquires a short-lived connection only around its own SQL.
  Moving the slow calls off the connection widens the decision→write window,
  so the held-transaction serialization is replaced with explicit guards:
  each source-liveness check (FOR SHARE) is paired with its write in one short
  transaction, and dedup CREATE/UPDATE folds are RETURNING-gated and re-filter
  live sources inside the fold transaction (sources-before-observation lock
  order, matching the normal write paths) so a twin or source deleted during
  the now connection-free window can't drop a CREATE or fold a dead source id.
  A cheap non-locking preflight restores the pre-refactor "skip before embed
  when every source is already gone" short-circuit. The separate-store
  (non-SQL) branches and the Oracle-safe search_vector clause are preserved.

Deterministic no-DB tests pin the fold guards (RETURNING gate, live-source
filtering, created/skipped propagation) and the pre-embed short-circuits;
live-DB curation/invalidation/document-transfer tests are updated to the new
short-acquire signatures.

The update_memory_unit hold-across-embed path is intentionally left for a
follow-up: its two-phase re-lock/abort/retry has to be reconciled with the
pluggable memories store's cross-store transaction coordinator and validated
against a real database.

Refs #2434
@nicoloboschi
nicoloboschi merged commit 2482520 into main Jul 31, 2026
101 of 103 checks passed
nicoloboschi added a commit that referenced this pull request Jul 31, 2026
…led connection (#3083)

Split update_memory_unit into a read/resolve/embed phase (no pooled connection held) and a short write transaction that re-reads the row, applies the precomputed embedding, and re-embeds in-txn only on a concurrent entity-set change; orphan entities from a failed edit are reclaimed by a forced graph-maintenance sweep. Preserves the pluggable store's begin_txn/decide_txn write-group.

Validated by CI (all three test-api shards incl. the live-PG curation suite pass). Follow-up to #3082. Refs #2434.
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