fix(retain): close to_unit_id deferred-FK race on memory_links inserts (#1882)

[nicoloboschi]

Fixes #1882.

Problem

The memory_links → memory_units FKs are DEFERRABLE INITIALLY DEFERRED (migration 9f8e7d6c5b4a), so an INSERT into memory_links takes no lock on the referenced parent row until COMMIT. Temporal and ANN link inserts reference a pre-existing neighbor unit as to_unit_id (graph maintenance also references a pre-existing from_unit_id). A concurrent transaction that commits a DELETE of that unit in the window between the link INSERT and our COMMIT — consolidation pruning observation units (consolidator.py), document re-tracking (fact_storage.py) — makes the deferred check fail at COMMIT:

insert or update on table "memory_links" violates foreign key constraint
"fk_memory_links_to_unit_id_memory_units"
DETAIL:  Key (to_unit_id)=(<uuid>) is not present in table "memory_units".

The async op is left status='failed' with no auto-retry.

#1795 / #1805 only removed one deleter (sibling async children sharing a document_id) for the from_unit_id side. The to_unit_id side — and any other deleter — stayed uncovered, which is the mirror case the issue reports.

Fix

Before inserting links, lock the referenced parent units FOR KEY SHARE (held until the surrounding transaction commits, which blocks a concurrent DELETE) and drop any link whose endpoint already vanished.

• New DataAccessOps.lock_referenced_units — a no-op default for backends with immediate FKs (Oracle has no deferred-FK window), overridden on PostgreSQL to take the row-share lock and report which rows survived.
• _bulk_insert_links calls it before the insert and filters vanished endpoints.
• The final ANN pass ran its insert in autocommit, so its insert is now wrapped in a transaction to hold the lock through the COMMIT-time check. The other callers (Phase-2 temporal/causal, graph maintenance) already run inside a write transaction.

Test

Adds a deterministic regression test (test_memory_links_to_unit_id_concurrent_delete.py) that hand-drives the connection interleaving with no sleeps:

1. conn A inserts a temporal link via the real _bulk_insert_links path — uncommitted.
2. conn B deletes the neighbor unit (bounded by statement_timeout).
3. conn A commits.

Pre-fix: step 2 succeeds instantly and step 3 raises the FK violation. Post-fix: step 2 blocks on A's FOR KEY SHARE lock and is cut off by the timeout; the neighbor survives and the link commits cleanly.

Verification

• New regression test passes; fails (reproduces fk_memory_links_to_unit_id_memory_units violation under sustained concurrent retain (mirror of #1795) #1882) without the fix.
• test_memory_links_deferred_fk.py, test_link_utils.py, test_retain.py, test_document_tracking.py, test_async_batch_retain.py, test_causal_relations.py, test_graph_maintenance.py, test_db_abstraction.py all pass (233 tests).
• ./scripts/hooks/lint.sh and uv run ty check pass.

[nicoloboschi]

Perf evaluation of the fix

Revised the implementation to a single-statement locking CTE (instead of a separate SELECT … FOR KEY SHARE + Python filter + a transaction wrap around the autocommit ANN insert). This keeps the link insert to one round-trip, no new transaction — so retain's round-trip/transaction shape is unchanged. The change is now contained entirely in PostgreSQLOps.bulk_insert_links.

End-to-end retain throughput (mock LLM + pg0, --scale medium, 1000 items, A/B toggling only this file)

	run 1	run 2
baseline	24.72 it/s	26.08 it/s
this fix	24.09 it/s	24.28 it/s

~0–4%, within run-to-run noise — link inserts are a small slice of retain (embeddings + fact insert dominate).

Micro-benchmark of the changed statement (isolates the link insert; 1000 links/call, real schema, 60 iters)

Typical batch (~70 distinct referenced units):

variant	p50	mean
baseline skip_exists_check=True (temporal/causal)	5.39ms	6.21ms
baseline WHERE EXISTS (semantic/graph)	6.18ms	6.29ms
this fix (locking CTE)	5.70ms	5.86ms

Worst case (~1050 distinct referenced units locked):

variant	p50	mean
baseline skip_exists_check=True	5.45ms	5.85ms
this fix (locking CTE)	7.53ms	8.35ms

Takeaways

• vs the cheapest baseline path (skip_exists_check=True, used by temporal/causal links): +~0.3ms per insert for typical batches, up to +~2ms for large 1000-link batches with ~1000 distinct neighbors (the FOR KEY SHARE lock acquisition + semi-join filter). Cost scales with the number of distinct referenced units.
• vs the WHERE EXISTS path (semantic ANN / graph maintenance): roughly neutral to faster — one CTE lock-scan replaces two per-row EXISTS subqueries.
• No extra round-trips, no extra transactions; the burst-write scenario from the issue uses single-item retains (small link batches), so the practical overhead there is the ~0.3ms end.

Net: a small, bounded cost on the link-insert step (invisible end-to-end within noise) in exchange for eliminating the FK-violation-at-COMMIT failure under concurrent retain.