[Bug]: partition locks are keyed per strategy; profile clustering and reflection deprecation race on one cluster row

"""Replay of #542 item 8: partition locks are keyed per strategy, so
``trigger_profile_clustering`` (lock ``("trigger_profile_clustering",
app:project:owner)``) and reflection's deprecation (lock
``("reflection_deprecate", app:project:cluster_id)``) do not exclude each other
although both write the same ``cluster`` row.

Interleaving: clustering reads the owner's clusters (members ep1, ep2, count 2)
-> reflection removes ep1, ep2, adds the merged episode and writes count=1 +
a new centroid -> clustering folds ep3 into its stale snapshot and
``upsert_with_members`` overwrites count/centroid and re-inserts ep1, ep2
(members are insert-only, diffed against what is in the table now).

Real parts: the clustering strategy body end to end (everalgo geometry merge,
real SQLite cluster repo, real partition locks), reflection's
``_update_cluster_after_merge`` under its real lock key. Stand-ins: fixed
embedders; the md / LanceDB deprecation writes of reflection are skipped (they
do not touch the cluster row). The only injection is a pause after
clustering's read.

Run: uv run pytest .work_context/replay_542/test_replay_8_partition_locks_cross_strategy.py -q -s
"""

from __future__ import annotations

import asyncio
import datetime as dt
import importlib
import shutil
import tempfile
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from unittest import mock

import numpy as np
from everalgo.clustering import Cluster as AlgoCluster
from sqlmodel import SQLModel

import everos.component.embedding.accessor as acc
from everos.component.embedding import EmbeddingCapability, EmbeddingProvider
from everos.config import SqliteSettings
from everos.core.persistence import (
    MemoryRoot,
    create_session_factory,
    create_system_engine,
)
from everos.infra.ome.testing import FakeStrategyContext
from everos.infra.persistence.sqlite.repos.cluster import _ClusterRepo
from everos.memory._partition_locks import _reset_for_tests, get_partition_lock
from everos.memory.events import EpisodeExtracted
from everos.memory.reflection.orchestrator import ReflectionOrchestrator

# the package re-exports the Strategy object under the module's name; patch the module itself
tpc_mod = importlib.import_module("everos.memory.strategies.trigger_profile_clustering")

DIM = 1024
T0 = 1_700_000_000_000


class _Embedder(EmbeddingProvider):
    dim = DIM

    def __init__(self, value: float) -> None:
        self.value = value

    async def embed(self, text: str) -> list[float]:
        return [self.value] * DIM

    async def embed_batch(self, texts: Any) -> list[list[float]]:
        return [await self.embed(t) for t in texts]


class _GatingRepo:
    """The real cluster repo; ``list_for_owner`` pauses after returning its result."""

    def __init__(self, inner: Any) -> None:
        self._inner = inner
        self.read_done = asyncio.Event()
        self.go = asyncio.Event()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._inner, name)

    async def list_for_owner(self, *a: Any, **k: Any) -> Any:
        res = await self._inner.list_for_owner(*a, **k)
        self.read_done.set()
        await self.go.wait()
        return res


def _event(owner: str, entry_id: str, ts: int) -> EpisodeExtracted:
    return EpisodeExtracted(
        memcell_id="mc_1",
        episode_entry_id=entry_id,
        episode_text=f"text {entry_id}",
        episode_timestamp_ms=ts,
        owner_id=owner,
        session_id="s1",
    )


async def _seed(repo: _ClusterRepo, cid: str, ev: EpisodeExtracted) -> None:
    await repo.upsert_with_members(
        AlgoCluster(
            id=cid,
            centroid=np.full(DIM, 0.1, dtype=np.float32),
            count=2,
            last_ts=T0,
            preview=["p1", "p2"],
            members=["ep1", "ep2"],
        ),
        owner_id=ev.owner_id,
        owner_type="user",
        kind="user_memory",
        member_type="episode",
        app_id=ev.app_id,
        project_id=ev.project_id,
    )


async def _reflection_update(
    repo: _ClusterRepo, cid: str, ev: EpisodeExtracted
) -> None:
    stub = SimpleNamespace(_cluster_repo=repo, _embedder=_Embedder(0.2))
    algo_result = SimpleNamespace(episode="merged episode text")
    # the merged episode keeps the cluster inside the 7-day merge window of T0
    episodes = [
        SimpleNamespace(timestamp=dt.datetime.fromtimestamp(T0 / 1000, tz=dt.UTC))
    ]
    partition = f"{ev.app_id}:{ev.project_id}:{cid}"
    async with get_partition_lock(
        "reflection_deprecate", partition
    ):  # orchestrator.py:680
        await ReflectionOrchestrator._update_cluster_after_merge(
            stub,
            cluster_id=cid,
            to_deprecate={"ep1", "ep2"},
            merged_entry_id="ep_merged",
            algo_result=algo_result,
            episodes=episodes,
        )


def _describe(c: AlgoCluster) -> str:
    return f"count={c.count} members={sorted(c.members)} centroid[0]={float(c.centroid[0]):.3f}"


async def _replay() -> None:
    tmp = Path(tempfile.mkdtemp(prefix="replay542_8_"))
    root = MemoryRoot(tmp / "root")
    root.ensure()
    engine = create_system_engine(root.system_db, SqliteSettings())
    factory = create_session_factory(engine)
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    repo = _ClusterRepo(session_factory=factory)
    acc._capability = EmbeddingCapability(provider=_Embedder(0.1))
    _reset_for_tests()
    try:
        # control first (owner u1): the two operations one after the other
        ev_ctl = _event("u1", "ep3", T0 + 1000)
        await _seed(repo, "cl_serial", ev_ctl)
        await _reflection_update(repo, "cl_serial", ev_ctl)
        with mock.patch.object(tpc_mod, "cluster_repo", repo):
            await tpc_mod.trigger_profile_clustering(ev_ctl, FakeStrategyContext())
        serial = await repo.get_with_members("cl_serial")
        print(f"  serialized: {_describe(serial)}")
        assert serial.count == len(serial.members) and not {"ep1", "ep2"} & set(
            serial.members
        )

        # the race (owner u2): clustering reads, reflection runs to completion, clustering writes
        ev3 = _event("u2", "ep3", T0 + 1000)
        cid = "cl_race"
        await _seed(repo, cid, ev3)
        gate = _GatingRepo(repo)
        with mock.patch.object(tpc_mod, "cluster_repo", gate):
            task = asyncio.create_task(
                tpc_mod.trigger_profile_clustering(ev3, FakeStrategyContext())
            )
            await asyncio.wait_for(gate.read_done.wait(), timeout=10)
            print(
                "  clustering has read {ep1, ep2}, count 2; now reflection deprecates them under its own lock"
            )
            await asyncio.wait_for(
                _reflection_update(repo, cid, ev3), timeout=10
            )  # would hang if the locks were shared
            mid = await repo.get_with_members(cid)
            print(f"  after reflection: {_describe(mid)}")
            assert (
                mid.count == 1
                and sorted(mid.members) == ["ep_merged"]
                and abs(float(mid.centroid[0]) - 0.2) < 1e-6
            )
            gate.go.set()
            await asyncio.wait_for(task, timeout=10)
        final = await repo.get_with_members(cid)
        print(f"  after clustering's stale write: {_describe(final)}")
        resurrected = {"ep1", "ep2"} & set(final.members)
        assert resurrected == {"ep1", "ep2"}, (
            "deprecated members are back in the cluster"
        )
        assert final.count != len(final.members), (
            f"count {final.count} vs {len(final.members)} members"
        )
        assert abs(float(final.centroid[0]) - 0.2) > 1e-3, (
            "reflection's centroid was overwritten"
        )
    finally:
        await engine.dispose()
        shutil.rmtree(tmp, ignore_errors=True)


def test_cross_strategy_writes_to_one_cluster_row_are_not_serialized() -> None:
    asyncio.run(_replay())