[Bug]: upsert_entity_with_link's exact-match/semantic-dedup paths still have the #383 lost-update race

[Vignesh-P-C]

Affected Component

Backend API / Server

Bug Description

upsert_entity_with_link (server/db/repositories.py) has 3 dedup paths
for merging a memory_id into a subject_entities row. Issue #383
(fixed by #384) made the fresh-insert path (Step 3) atomic via
INSERT ... ON CONFLICT DO UPDATE, closing a race where concurrent
writers could create duplicate rows.

However, Steps 1 (exact-match) and 2 (semantic/cosine-similarity
match) — the paths taken when a row for the entity already exists —
were never migrated to the same atomic pattern. They still do a
Python-side read-modify-write:

exact = (await session.execute(exact_stmt)).scalar_one_or_none()
if exact is not None:
    if memory_id not in exact.linked_memory_ids:
        exact.linked_memory_ids = [*exact.linked_memory_ids, memory_id]
    return exact

Two concurrent transactions that both read the same row before either
commits will each independently compute a new array and flush it —
whichever commits last silently overwrites the other's append. This
is a lost-update race, the same class of bug #383 was originally
about, just on the other two branches.

This is why test_concurrent_upserts_converge_to_one_row — the
regression test #384 itself added to prove #383 was fixed — is
intermittently flaky in CI: when concurrent writers race the
already-exists path (Step 1) rather than the fresh-insert path
(Step 3), the fix doesn't apply and memory_ids get silently dropped.

Expected Behavior

All N concurrent writers linking memory_ids to the same normalized
entity should converge to exactly one row, with linked_memory_ids
containing the union of every writer's memory_id — regardless of
whether the race lands on the fresh-insert path or the
already-exists/merge path.

Steps to Reproduce

1. Pick a subject_id and an entity (e.g. "acme corp") with no existing
   subject_entities row.
2. Fire 5+ concurrent calls to upsert_entity_with_link for that same
   (subject_id, entity_normalized), each with a distinct memory_id, in
   separate sessions/transactions (see
   tests/integration/test_subject_entities_uniqueness.py::
   test_concurrent_upserts_converge_to_one_row for the exact pattern).
3. Query subject_entities for that subject_id afterward.
4. Some memory_ids are missing from linked_memory_ids — the surviving
   set depends on commit ordering, not on which writers actually ran.

Code Sample

import asyncio, uuid
from server.db import repositories as repo

async def upsert(session_factory, subject_id, memory_id):
    async with session_factory() as session:
        await repo.upsert_entity_with_link(
            session, subject_id=subject_id, tenant_id=None,
            entity_text="Acme Corp", entity_normalized="acme corp",
            entity_kind="ORG", embedding=None, memory_id=memory_id,
        )
        await session.commit()

memory_ids = [uuid.uuid4() for _ in range(5)]
await asyncio.gather(*[upsert(session_factory, "s1", m) for m in memory_ids])
# Expect: one row, linked_memory_ids == set(memory_ids). Observed: some
# memory_ids missing under this exact-match/merge race.

Error Output / Logs

AssertionError: assert {UUID('1ddf8dd0-...'), UUID('c7016ea1-...')} == {UUID('1abd9552-...'), UUID('1b5e1716-...'), UUID('1ddf8dd0-...'), UUID('c7016ea1-...'), UUID('f1cd9afe-...')}

Extra items in the right set:
UUID('f1cd9afe-82ea-4145-9cc7-6891e444248c')
UUID('1b5e1716-a74c-4add-a0fc-a113b2b9dcba')
UUID('1abd9552-9338-4e45-a63e-31d43ad202db')

tests/integration/test_subject_entities_uniqueness.py:64: AssertionError

Environment

• OS: ubuntu-latest (GitHub Actions runner)
• Python version: 3.11.16
• Statewave version: 1.5.0
• Deployment: CI (docker compose postgres service, pgvector/pgvector:pg16)

Severity

High - Major feature broken, no workaround

Additional Context

Related to #383 / #384, which fixed the fresh-insert path only. I have
a fix ready that applies the same atomic UPDATE ... RETURNING pattern
Step 3 already uses to Steps 1 and 2, closing the race on all three
paths — happy to open a PR.

Checklist

• I have searched existing issues to ensure this is not a duplicate
• I have provided all required information above

[Vignesh-P-C]

I'd like to take this on — have a fix ready. Applied the same atomic UPDATE ... RETURNING pattern Step 3 already uses (CASE WHEN memory_id = ANY(linked_memory_ids) THEN linked_memory_ids ELSE array_append(...) END) to Steps 1 and 2, replacing the SELECT-then-mutate-in-Python read-modify-write on each. Postgres's row-level locking then serializes concurrent writers targeting the same row instead of losing updates, the same way it already does for the ON CONFLICT path. Will verify against test_concurrent_upserts_converge_to_one_row (run repeatedly, since concurrency bugs don't reproduce every run) before opening the PR.