diff --git a/src/everos/memory/reflection/orchestrator.py b/src/everos/memory/reflection/orchestrator.py index a1e07ddec..831cabdb1 100644 --- a/src/everos/memory/reflection/orchestrator.py +++ b/src/everos/memory/reflection/orchestrator.py @@ -10,7 +10,6 @@ from __future__ import annotations -import asyncio import datetime as _dt import json import uuid @@ -32,13 +31,29 @@ from everos.core.observability.tracing import memory_span from everos.core.persistence import MemoryRoot from everos.infra.ome.context import StrategyContext -from everos.infra.persistence.index import all_of, eq, is_null +from everos.infra.persistence.index import all_of, eq, is_null, one_of from everos.memory._partition_locks import get_partition_lock from everos.memory.events import EpisodeExtracted logger = get_logger(__name__) _MAX_CLUSTERS_PER_RUN = 10 + +# Upper bound on the episodes one merge consumes. A cluster with more members +# is merged incrementally: its existing merged episode plus the oldest sources +# go into this run, the remaining members stay in the cluster and are folded +# in by later runs in update mode. Keeps the merged narrative, and the +# cascade upsert it triggers, inside the LanceDB write-lock budget. +_MAX_SOURCES_PER_MERGE = 50 +# Orphaned merged episodes are retired on top of the source cap. A cluster +# that piled up failed runs keeps at most this many per run, and like a +# capped merge it converges over the following runs. +_MAX_ORPHANS_PER_RUN = 10 + +# Entry ids per LanceDB ``update`` when deprecating sources. One update per +# batch instead of one per row keeps the number of write-lock acquisitions +# small for large clusters. +_DEPRECATE_BATCH_SIZE = 100 _WAIT_TIMEOUT_SECONDS = 120.0 @@ -257,8 +272,6 @@ async def _process_cluster( Returns: A ReflectionReport on success, ``None`` on skip. """ - await self._detect_orphans(cluster_id, owner_id, app_id, project_id) - scope = dict(owner_id=owner_id, app_id=app_id, project_id=project_id) members, episodes = await self._load_cluster_episodes( cluster_id=cluster_id, **scope @@ -354,7 +367,52 @@ async def _load_cluster_episodes( app_id=app_id, project_id=project_id, ) - return members, episodes + return self._cap_cluster_sources( + cluster_id=cluster_id, members=members, episodes=episodes + ) + + @staticmethod + def _cap_cluster_sources( + *, + cluster_id: str, + members: list[tuple[str, str]], + episodes: list[Any], + ) -> tuple[list[tuple[str, str]], list[Any]]: + """Limit one merge to ``_MAX_SOURCES_PER_MERGE`` episodes. + + The cluster's existing merged episodes (``parent_type == "cluster"``) + are always kept so update mode still sees the previous narrative; + the remaining budget goes to the oldest source episodes. Members + whose episode is deferred are dropped from this run's snapshot so + they are neither merged nor deprecated now. + + Args: + cluster_id: Target cluster identifier (for logging). + members: Cluster members ``(member_id, member_type)``. + episodes: Episode rows sorted by timestamp ascending. + + Returns: + ``(members, episodes)`` restricted to this run. + """ + if len(episodes) <= _MAX_SOURCES_PER_MERGE: + return members, episodes + + merged = [e for e in episodes if e.parent_type == "cluster"] + sources = [e for e in episodes if e.parent_type != "cluster"] + budget = max(_MAX_SOURCES_PER_MERGE - len(merged), 1) + kept = merged + sources[:budget] + kept.sort(key=lambda e: e.timestamp) + deferred = sources[budget:] + deferred_keys = {e.entry_id for e in deferred} | {e.parent_id for e in deferred} + kept_members = [(mid, mt) for mid, mt in members if mid not in deferred_keys] + logger.info( + "reflection_cluster_capped", + cluster_id=cluster_id, + member_count=len(members), + merged_count=len(kept_members), + deferred_count=len(members) - len(kept_members), + ) + return kept_members, kept async def _write_and_reextract( self, @@ -508,14 +566,27 @@ async def _detect_orphans( owner_id: str, app_id: str, project_id: str, - ) -> None: - """Log warning if orphan merged episodes exist for this cluster. + *, + exclude: set[str], + ) -> list[Any]: + """Find live merged episodes of this cluster that no run committed. + + A run that fails after writing its merged episode (deprecation + failed and was reverted, or extraction timed out) leaves that + episode indexed and live but outside the cluster. It is never + loaded as a source, because sources come from cluster membership, + but it stays searchable next to the merge that replaces it. Args: cluster_id: Target cluster identifier. owner_id: Target owner identifier. app_id: Application scope. project_id: Project scope. + exclude: Entry IDs that are not orphans: the current cluster + members and the merged episode this run wrote. + + Returns: + The orphaned merged episode rows (possibly empty). """ where = all_of( eq("parent_type", "cluster"), @@ -525,13 +596,19 @@ async def _detect_orphans( eq("app_id", app_id), eq("project_id", project_id), ) - orphans = await self._episode_store.find_where(where, limit=10) + rows = await self._episode_store.find_where( + where, limit=_MAX_ORPHANS_PER_RUN + len(exclude) + ) + orphans = [row for row in rows if row.entry_id not in exclude][ + :_MAX_ORPHANS_PER_RUN + ] if orphans: logger.warning( "reflection_orphan_detected", cluster_id=cluster_id, orphan_entry_ids=[o.entry_id for o in orphans], ) + return orphans async def _fetch_episodes( self, @@ -733,16 +810,26 @@ async def _execute_deprecation( Returns: A ReflectionReport on success, ``None`` when no members to deprecate. """ - to_deprecate = await self._resolve_deprecation_targets( + to_deprecate, kept = await self._resolve_deprecation_targets( cluster_id=cluster_id, original_members=original_members, ) if not to_deprecate: return None + # Merged episodes left behind by earlier failed runs are retired + # toward this merge in the same writes (and reverted with them), so + # a retry ends with one live merged episode instead of two. + orphans = await self._detect_orphans( + cluster_id, + owner_id, + app_id, + project_id, + exclude=to_deprecate | kept | {merged_entry_id}, + ) dep_ep, dep_fact = await self._apply_deprecation_writes( - episodes=episodes, - to_deprecate=to_deprecate, + episodes=[*episodes, *orphans], + to_deprecate=to_deprecate | {o.entry_id for o in orphans}, owner_id=owner_id, app_id=app_id, project_id=project_id, @@ -754,6 +841,7 @@ async def _execute_deprecation( merged_entry_id=merged_entry_id, algo_result=algo_result, episodes=episodes, + member_count=len(kept) + 1, ) report = await self._create_reflection_report( cluster_id=cluster_id, @@ -771,6 +859,11 @@ async def _execute_deprecation( cluster_id=cluster_id, deprecated_episode_count=dep_ep, deprecated_fact_count=dep_fact, + # Members this merge left in the cluster (deferred by the source + # cap, or added after the snapshot): a later run merges them, so a + # large cluster converges over several runs rather than in one. + cluster_remaining_count=len(kept), + retired_orphan_count=len(orphans), ) return report @@ -784,11 +877,19 @@ async def _apply_deprecation_writes( project_id: str, merged_entry_id: str, ) -> tuple[int, int]: - """Patch md frontmatter and mark episodes/facts deprecated in LanceDB. + """Mark episodes/facts deprecated in LanceDB, then patch md frontmatter. + + The LanceDB writes go first and the markdown record last, so the + durable record only lists deprecations that were actually applied. + If any write fails, the ``deprecated_by`` values this run already + wrote are reverted before the error propagates, so a failed run + leaves no sources half-deprecated toward a merge that was never + committed to the cluster. Args: - episodes: Source episode rows (for md patching). - to_deprecate: Set of member IDs being deprecated. + episodes: Rows of the episodes being deprecated (for md patching). + to_deprecate: Entry IDs being deprecated: cluster members and + orphaned merged episodes. owner_id: Target owner identifier. app_id: Application scope. project_id: Project scope. @@ -797,33 +898,92 @@ async def _apply_deprecation_writes( Returns: ``(deprecated_episode_count, deprecated_fact_count)``. """ - await self._patch_md_frontmatter( - episodes=episodes, - to_deprecate=to_deprecate, - merged_entry_id=merged_entry_id, - ) - deprecated_ep_count = await self._deprecate_lance_episodes( - entry_ids=to_deprecate, - owner_id=owner_id, - app_id=app_id, - project_id=project_id, - merged_entry_id=merged_entry_id, - ) - deprecated_fact_count = await self._deprecate_lance_facts( - parent_ids=to_deprecate, - owner_id=owner_id, - app_id=app_id, - project_id=project_id, + try: + deprecated_ep_count = await self._deprecate_lance_episodes( + entry_ids=to_deprecate, + owner_id=owner_id, + app_id=app_id, + project_id=project_id, + merged_entry_id=merged_entry_id, + ) + deprecated_fact_count = await self._deprecate_lance_facts( + parent_ids=to_deprecate, + owner_id=owner_id, + app_id=app_id, + project_id=project_id, + merged_entry_id=merged_entry_id, + ) + await self._patch_md_frontmatter( + episodes=episodes, + to_deprecate=to_deprecate, + merged_entry_id=merged_entry_id, + ) + except BaseException: + await self._revert_lance_deprecation( + entry_ids=to_deprecate, + owner_id=owner_id, + app_id=app_id, + project_id=project_id, + merged_entry_id=merged_entry_id, + ) + raise + return deprecated_ep_count, deprecated_fact_count + + async def _revert_lance_deprecation( + self, + *, + entry_ids: set[str], + owner_id: str, + app_id: str, + project_id: str, + merged_entry_id: str, + ) -> None: + """Best-effort compensation: clear ``deprecated_by`` written by this run. + + Only rows that point at ``merged_entry_id`` are touched, so + deprecations from earlier, committed merges are left alone. + Failures are logged, never raised: the original error is the one + the caller needs to see. + """ + reverted = True + for store, field_name, label in ( + (self._episode_store, "entry_id", "episode"), + (self._atomic_fact_store, "parent_id", "atomic_fact"), + ): + for batch in _batched(sorted(entry_ids), _DEPRECATE_BATCH_SIZE): + # Entry ids restart per (app, project, day), so another + # project's rows can point at the same merged entry id. + where = all_of( + one_of(field_name, batch), + eq("owner_id", owner_id), + eq("app_id", app_id), + eq("project_id", project_id), + eq("deprecated_by", merged_entry_id), + ) + try: + await store.update({"deprecated_by": None}, where=where) + except Exception: + reverted = False + logger.error( + "reflection_deprecate_revert_failed", + table=label, + merged_entry_id=merged_entry_id, + batch_size=len(batch), + exc_info=True, + ) + logger.warning( + "reflection_deprecate_reverted", merged_entry_id=merged_entry_id, + source_count=len(entry_ids), + complete=reverted, ) - return deprecated_ep_count, deprecated_fact_count async def _resolve_deprecation_targets( self, *, cluster_id: str, original_members: list[tuple[str, str]], - ) -> set[str]: + ) -> tuple[set[str], set[str]]: """Re-read cluster members and intersect with the original snapshot. Args: @@ -831,12 +991,16 @@ async def _resolve_deprecation_targets( original_members: Snapshot ``(member_id, member_type)`` from selection. Returns: - Set of member IDs safe to deprecate (present in both snapshots). + ``(to_deprecate, kept)``: the member IDs safe to deprecate + (present in both snapshots) and the current members that stay + in the cluster (deferred by the per-merge cap, or added after + the snapshot). """ current_members = await self._cluster_repo.get_members_with_type(cluster_id) current_ids = {mid for mid, _ in current_members} original_ids = {mid for mid, _ in original_members} - return original_ids & current_ids + to_deprecate = original_ids & current_ids + return to_deprecate, current_ids - to_deprecate async def _deprecate_lance_episodes( self, @@ -849,24 +1013,25 @@ async def _deprecate_lance_episodes( ) -> int: """Mark deprecated episodes in LanceDB by entry_id. + One ``update`` per batch of ``_DEPRECATE_BATCH_SIZE`` ids, issued + sequentially: every call takes the table's write lock, so N + concurrent single-row updates queue behind each other and, on a + large cluster, blow the lock deadline. + Returns: - Number of LanceDB update calls issued. + Number of episodes targeted. """ - coros: list[Any] = [ - self._episode_store.update( + for batch in _batched(sorted(entry_ids), _DEPRECATE_BATCH_SIZE): + await self._episode_store.update( {"deprecated_by": merged_entry_id}, where=all_of( - eq("entry_id", eid), + one_of("entry_id", batch), eq("owner_id", owner_id), eq("app_id", app_id), eq("project_id", project_id), ), ) - for eid in entry_ids - ] - if coros: - await asyncio.gather(*coros) - return len(coros) + return len(entry_ids) async def _deprecate_lance_facts( self, @@ -891,26 +1056,23 @@ async def _deprecate_lance_facts( merged_entry_id: Entry ID of the replacement merged episode. Returns: - Total number of LanceDB update calls issued. + Number of parents whose facts were targeted. """ if not parent_ids: return 0 - coros = [ - self._atomic_fact_store.update( + for batch in _batched(sorted(parent_ids), _DEPRECATE_BATCH_SIZE): + await self._atomic_fact_store.update( {"deprecated_by": merged_entry_id}, where=all_of( - eq("parent_id", pid), + one_of("parent_id", batch), eq("owner_id", owner_id), eq("app_id", app_id), eq("project_id", project_id), is_null("deprecated_by"), ), ) - for pid in parent_ids - ] - await asyncio.gather(*coros) - return len(coros) + return len(parent_ids) async def _update_cluster_after_merge( self, @@ -920,6 +1082,7 @@ async def _update_cluster_after_merge( merged_entry_id: str, algo_result: AlgoEpisode, episodes: list[Any], + member_count: int = 1, ) -> None: """Remove old members, add merged, and recompute centroid. @@ -929,6 +1092,8 @@ async def _update_cluster_after_merge( merged_entry_id: Entry ID of the newly merged episode. algo_result: Algo reflector output (episode text for centroid). episodes: Source episode rows (for last timestamp). + member_count: Members left in the cluster after the merge + (the merged episode plus any deferred or newly added ones). """ await self._cluster_repo.remove_members(cluster_id, to_deprecate) await self._cluster_repo.add_member(cluster_id, merged_entry_id, "episode") @@ -939,7 +1104,7 @@ async def _update_cluster_after_merge( await self._cluster_repo.update_metadata( cluster_id, centroid_blob=centroid_blob, - count=1, + count=member_count, last_ts_ms=last_ts_ms, preview_json=json.dumps([algo_result.episode[:200]], ensure_ascii=False), ) @@ -1030,6 +1195,11 @@ async def _patch_md_frontmatter( ) +def _batched(items: list[str], size: int) -> list[list[str]]: + """Split ``items`` into consecutive lists of at most ``size`` elements.""" + return [items[i : i + size] for i in range(0, len(items), size)] + + def _to_algo_episodes(episodes: list[Any]) -> list[AlgoEpisode]: """Convert LanceDB episode rows to algo Episode objects. diff --git a/tests/integration/test_reflection_integration.py b/tests/integration/test_reflection_integration.py index 78c6c1088..7784f40ed 100644 --- a/tests/integration/test_reflection_integration.py +++ b/tests/integration/test_reflection_integration.py @@ -18,7 +18,9 @@ import datetime as _dt import hashlib import json +from collections.abc import Mapping, Sequence from pathlib import Path +from typing import Any import numpy as np import pytest @@ -27,8 +29,12 @@ from everalgo.user_memory.reflect import EpisodeReflector from sqlmodel import SQLModel +from everos.component.utils.datetime import from_iso_format from everos.config import LanceDBSettings, load_settings +from everos.core.errors import VectorStoreBusyError from everos.core.persistence import ( + EntryId, + MarkdownReader, MemoryRoot, open_lancedb_connection, ) @@ -36,6 +42,7 @@ from everos.infra.ome.testing import FakeStrategyContext from everos.infra.persistence.lancedb import LanceIndexContract from everos.infra.persistence.markdown.writers.episode_writer import EpisodeWriter +from everos.infra.persistence.predicate import Predicate from everos.infra.persistence.records.atomic_fact import AtomicFact from everos.infra.persistence.records.episode import Episode as LanceEpisode from everos.infra.persistence.sqlite import cluster_repo, reflection_report_repo @@ -850,3 +857,260 @@ async def test_reflected_episodes_visible_in_search_deprecated_excluded( finally: conn.close() await _teardown_sqlite() + + +# --------------------------------------------------------------------------- +# Retry after a failed deprecation step +# --------------------------------------------------------------------------- + + +class _IndexingEpisodeWriter(EpisodeWriter): + """EpisodeWriter that also plays the cascade and atomic fact extraction. + + ``FakeStrategyContext`` runs neither, so every entry appended here is + indexed into LanceDB right away together with one atomic fact derived + from it: the state the real pipeline reaches before Reflection gets to + its deprecation step. The first ``failing_patches`` frontmatter patches + raise. + """ + + def __init__( + self, + memory_root: MemoryRoot, + ep_repo: _EpisodeRepo, + af_repo: _AtomicFactRepo, + *, + failing_patches: int = 0, + ) -> None: + super().__init__(memory_root) + self._ep_repo = ep_repo + self._af_repo = af_repo + self._failing_patches = failing_patches + + async def append_entries( + self, + scope_id: str, + items: Sequence[tuple[Mapping[str, object], Mapping[str, str]]], + *, + date: _dt.date | None = None, + app_id: str = "default", + project_id: str = "default", + ) -> list[EntryId]: + eids = await super().append_entries( + scope_id, items, date=date, app_id=app_id, project_id=project_id + ) + for eid, (inline, sections) in zip(eids, items, strict=True): + entry_id = eid.format() + md_path = self._resolve_path(scope_id, eid.date, app_id, project_id) + timestamp = from_iso_format(str(inline["timestamp"])) + session_id = inline.get("session_id") + await self._ep_repo.add( + [ + _make_lance_episode( + entry_id=entry_id, + owner_id=scope_id, + episode=sections["Content"], + timestamp=timestamp, + parent_type=str(inline["parent_type"]), + parent_id=str(inline["parent_id"]), + session_id=None if session_id is None else str(session_id), + md_path=md_path.relative_to(self._root.root).as_posix(), + ) + ] + ) + await self._af_repo.add( + [ + _make_lance_fact( + entry_id=entry_id.replace("ep_", "af_", 1), + owner_id=scope_id, + fact=sections["Content"], + parent_id=entry_id, + parent_type="episode", + timestamp=timestamp, + ) + ] + ) + return eids + + async def patch_frontmatter(self, path: Path, updates: Mapping[str, Any]) -> None: + if self._failing_patches: + self._failing_patches -= 1 + raise OSError("No space left on device") + await super().patch_frontmatter(path, updates) + + +class _LockedFactStore(_AtomicFactRepo): + """Atomic fact index whose first deprecation write misses the lock deadline.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._armed = True + + async def update(self, updates: dict[str, Any], *, where: Predicate | str) -> None: + if self._armed and updates.get("deprecated_by") is not None: + self._armed = False + raise VectorStoreBusyError("write lock deadline") + await super().update(updates, where=where) + + +async def _deprecated_entries(memory_root: MemoryRoot, owner_id: str) -> dict[str, str]: + """Merge ``deprecated_entries`` across the owner's episode daily logs.""" + merged: dict[str, str] = {} + episodes_dir = memory_root.users_dir("default", "default") / owner_id / "episodes" + for path in sorted(episodes_dir.rglob("episode-*.md")): + parsed = await MarkdownReader.read(path) + merged.update(parsed.frontmatter.get("deprecated_entries") or {}) + return merged + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failing_step", ["lance", "markdown"]) +async def test_retry_after_failed_deprecation_leaves_one_merged_episode( + tmp_path: Path, + memory_root: MemoryRoot, + monkeypatch: pytest.MonkeyPatch, + failing_step: str, +) -> None: + """A run that fails after writing its merge is retried idempotently. + + Run 1 writes and indexes merged episode v1, then its deprecation step + fails after the LanceDB episode write (``lance``: the atomic fact update + misses the write-lock deadline; ``markdown``: the frontmatter patch + raises) and is compensated. v1 is left behind as an orphan: indexed, + live, and not a cluster member. + + Verifies that the retry: + - merges the three sources again in INIT mode, so v1 was neither + a source nor treated as the previous narrative + - deprecates the sources toward its own merged episode v2 only + - retires v1 and v1's atomic fact toward v2, in LanceDB and in + the md frontmatter, leaving exactly one live merged episode + """ + monkeypatch.setattr( + MemoryRoot, + "resolve", + classmethod(lambda cls: MemoryRoot(root=tmp_path)), + ) + monkeypatch.setenv("EVEROS_LLM__API_KEY", "fake-key") + monkeypatch.setenv("EVEROS_LLM__BASE_URL", "https://fake.example.com") + + load_settings.cache_clear() + + await _setup_sqlite(monkeypatch) + + try: + conn = await open_lancedb_connection(memory_root.lancedb_dir, LanceDBSettings()) + ep_table = await conn.create_table("episode", schema=LanceEpisode) + af_table = await conn.create_table("atomic_fact", schema=AtomicFact) + ep_repo = _EpisodeRepo(table=ep_table) + af_repo = _AtomicFactRepo(table=af_table) + + owner_id = "u_test" + cluster_id = "cl_test_retry" + episode_writer = _IndexingEpisodeWriter( + memory_root, + ep_repo, + af_repo, + failing_patches=1 if failing_step == "markdown" else 0, + ) + + # -- Three source episodes in md + LanceDB, one cluster holding them. + texts = ["Andrew has no pets", "Andrew adopted Toby", "Andrew adopted Buddy"] + source_eids = await episode_writer.append_entries( + owner_id, + [ + ( + { + "owner_id": owner_id, + "timestamp": f"2026-06-10T1{i}:00:00+00:00", + "parent_type": "memcell", + "parent_id": f"mc_00{i}", + "session_id": "s_test", + }, + {"Subject": "test", "Content": text}, + ) + for i, text in enumerate(texts, start=1) + ], + ) + source_ids = [eid.format() for eid in source_eids] + await cluster_repo.upsert_with_members( + AlgoCluster( + id=cluster_id, + centroid=np.zeros(1024, dtype=np.float32), + count=3, + last_ts=int( + _dt.datetime(2026, 6, 10, 13, 0, 0, tzinfo=_dt.UTC).timestamp() + * 1000 + ), + preview=texts[:2], + members=source_ids, + ), + owner_id=owner_id, + owner_type="user", + kind="user_memory", + member_type="episode", + ) + + fake_llm = FakeLLMClient( + responses=[ + json.dumps({"content": "Andrew pets narrative v1.", "title": "v1"}), + json.dumps({"content": "Andrew pets narrative v2.", "title": "v2"}), + ] + ) + orchestrator = ReflectionOrchestrator( + cluster_repo=cluster_repo, + episode_store=ep_repo, + atomic_fact_store=( + _LockedFactStore(table=af_table) if failing_step == "lance" else af_repo + ), + episode_writer=episode_writer, + report_repo=reflection_report_repo, + reflector=EpisodeReflector(llm=fake_llm), + embedder=_StubEmbedder(), + ) + live_episodes = f"owner_id = '{owner_id}' AND deprecated_by IS NULL" + + # -- Run 1: v1 is written and indexed, the deprecation step fails. + reports = await orchestrator.run(ctx=FakeStrategyContext(), owner_id=owner_id) + assert reports == [] + + merged_rows = await ep_repo.find_where( + f"parent_type = 'cluster' AND parent_id = '{cluster_id}'" + ) + assert len(merged_rows) == 1 + orphan_id = merged_rows[0].entry_id + live = await ep_repo.find_where(live_episodes) + assert sorted(r.entry_id for r in live) == sorted([*source_ids, orphan_id]) + members = await cluster_repo.get_members_with_type(cluster_id) + assert sorted(mid for mid, _ in members) == sorted(source_ids) + assert await _deprecated_entries(memory_root, owner_id) == {} + + # -- Run 2: the retry. + reports = await orchestrator.run(ctx=FakeStrategyContext(), owner_id=owner_id) + + assert len(reports) == 1 + report = reports[0] + assert report.mode == "init" + assert report.source_count == 3 + retry_prompt = " ".join(m.content for m in fake_llm.calls[1].messages) + assert "narrative v1" not in retry_prompt + merged_id = report.merged_entry_id + + live = await ep_repo.find_where(live_episodes) + assert [r.entry_id for r in live] == [merged_id] + superseded = await ep_repo.find_where( + f"owner_id = '{owner_id}' AND deprecated_by IS NOT NULL" + ) + expected = {eid: merged_id for eid in [*source_ids, orphan_id]} + assert {r.entry_id: r.deprecated_by for r in superseded} == expected + live_facts = await af_repo.find_where( + f"owner_id = '{owner_id}' AND deprecated_by IS NULL" + ) + assert [f.parent_id for f in live_facts] == [merged_id] + members = await cluster_repo.get_members_with_type(cluster_id) + assert members == [(merged_id, "episode")] + assert await _deprecated_entries(memory_root, owner_id) == expected + + finally: + conn.close() + await _teardown_sqlite() diff --git a/tests/unit/test_memory/test_reflection/test_orchestrator.py b/tests/unit/test_memory/test_reflection/test_orchestrator.py index 2bdb58be8..34d72ab02 100644 --- a/tests/unit/test_memory/test_reflection/test_orchestrator.py +++ b/tests/unit/test_memory/test_reflection/test_orchestrator.py @@ -16,8 +16,11 @@ import pytest +from everos.core.errors import VectorStoreBusyError from everos.infra.ome.testing import FakeStrategyContext +from everos.infra.persistence.predicate import All, Comparison, In, Predicate from everos.memory._partition_locks import _reset_for_tests +from everos.memory.reflection import orchestrator as orchestrator_module from everos.memory.reflection.orchestrator import ( _MAX_CLUSTERS_PER_RUN, ReflectionOrchestrator, @@ -199,7 +202,28 @@ async def test_deprecate_lance_facts_is_scoped_to_the_project() -> None: rendered = render_predicate(where) assert "(app_id = 'claude-code')" in rendered assert "(project_id = 'alpha')" in rendered - assert "(parent_id = 'ep_20260601_0001')" in rendered + assert "(parent_id IN ('ep_20260601_0001'))" in rendered + + +async def test_detect_orphans_is_capped_per_run(monkeypatch) -> None: + """Orphans are retired outside the source cap, so they have their own.""" + monkeypatch.setattr(orchestrator_module, "_MAX_ORPHANS_PER_RUN", 2) + episode_store = MagicMock() + rows = [ + _make_episode_row( + entry_id=f"ep_20260601_000{i}", parent_id="cl_abc", owner_id="u_alice" + ) + for i in range(1, 5) + ] + episode_store.find_where = AsyncMock(return_value=rows) + orch = _build_orchestrator(episode_store=episode_store) + + orphans = await orch._detect_orphans( + "cl_abc", "u_alice", "claude-code", "alpha", exclude={"ep_20260601_0001"} + ) + + assert [o.entry_id for o in orphans] == ["ep_20260601_0002", "ep_20260601_0003"] + assert episode_store.find_where.await_args.kwargs["limit"] == 3 async def test_run_init_mode_merges_and_deprecates() -> None: @@ -509,3 +533,217 @@ async def test_call_reflector_emits_consolidate_generation_span() -> None: spans["everos.reflect.consolidate"].attributes["langfuse.observation.type"] == "generation" ) + + +# ── Batched deprecation, compensation and per-merge cap ─────────────────── + + +def _in_values(predicate: Predicate, field_name: str) -> list[str]: + """Collect the values of every ``IN`` clause on ``field_name``.""" + if isinstance(predicate, In): + return list(predicate.values) if predicate.field == field_name else [] + if isinstance(predicate, All): + return [ + v for child in predicate.children for v in _in_values(child, field_name) + ] + return [] + + +def _eq_value(predicate: Predicate, field_name: str) -> object | None: + if isinstance(predicate, Comparison): + return predicate.value if predicate.field == field_name else None + if isinstance(predicate, All): + for child in predicate.children: + found = _eq_value(child, field_name) + if found is not None: + return found + return None + + +async def test_deprecate_lance_episodes_batches_updates(monkeypatch) -> None: + """One update per batch of ids instead of one update per row.""" + monkeypatch.setattr(orchestrator_module, "_DEPRECATE_BATCH_SIZE", 100) + episode_store = MagicMock() + episode_store.update = AsyncMock() + orch = _build_orchestrator(episode_store=episode_store) + entry_ids = {f"ep_{i:04d}" for i in range(250)} + + count = await orch._deprecate_lance_episodes( + entry_ids=entry_ids, + owner_id="u_alice", + app_id="default", + project_id="default", + merged_entry_id="ep_merged", + ) + + assert count == 250 + assert episode_store.update.await_count == 3 + seen: list[str] = [] + for call in episode_store.update.await_args_list: + updates, where = call.args[0], call.kwargs["where"] + assert updates == {"deprecated_by": "ep_merged"} + batch = _in_values(where, "entry_id") + assert 0 < len(batch) <= 100 + assert _eq_value(where, "owner_id") == "u_alice" + seen.extend(batch) + assert sorted(seen) == sorted(entry_ids) + + +async def test_deprecate_failure_reverts_applied_writes(monkeypatch) -> None: + """A failed LanceDB write reverts what this run already wrote and leaves + the markdown record untouched, then propagates the error.""" + monkeypatch.setattr(orchestrator_module, "_DEPRECATE_BATCH_SIZE", 100) + episode_store = MagicMock() + episode_store.update = AsyncMock( + side_effect=[None, VectorStoreBusyError("write lock deadline"), None, None] + ) + atomic_fact_store = MagicMock() + atomic_fact_store.update = AsyncMock() + episode_writer = MagicMock() + episode_writer.patch_frontmatter = AsyncMock() + orch = _build_orchestrator( + episode_store=episode_store, + atomic_fact_store=atomic_fact_store, + episode_writer=episode_writer, + ) + entry_ids = {f"ep_{i:04d}" for i in range(150)} + episodes = [_make_episode_row(entry_id=eid) for eid in sorted(entry_ids)] + + with pytest.raises(VectorStoreBusyError): + await orch._apply_deprecation_writes( + episodes=episodes, + to_deprecate=entry_ids, + owner_id="u_alice", + app_id="default", + project_id="default", + merged_entry_id="ep_merged", + ) + + # The markdown record was never patched. + episode_writer.patch_frontmatter.assert_not_awaited() + # Two forward batches (second failed) + two revert batches on episodes. + assert episode_store.update.await_count == 4 + reverts = episode_store.update.await_args_list[2:] + reverted: list[str] = [] + for call in reverts: + updates, where = call.args[0], call.kwargs["where"] + assert updates == {"deprecated_by": None} + assert _eq_value(where, "deprecated_by") == "ep_merged" + reverted.extend(_in_values(where, "entry_id")) + assert sorted(reverted) == sorted(entry_ids) + # Facts were never deprecated, only (harmlessly) reverted. + for call in atomic_fact_store.update.await_args_list: + assert call.args[0] == {"deprecated_by": None} + + +async def test_deprecation_revert_is_scoped_to_the_project(monkeypatch) -> None: + """Entry ids restart per (app, project, day), so the compensating revert + must not clear another project's rows that point at the same merged id.""" + episode_store = MagicMock() + episode_store.update = AsyncMock( + side_effect=[None, VectorStoreBusyError("write lock deadline"), None, None] + ) + atomic_fact_store = MagicMock() + atomic_fact_store.update = AsyncMock() + orch = _build_orchestrator( + episode_store=episode_store, + atomic_fact_store=atomic_fact_store, + episode_writer=MagicMock(), + ) + monkeypatch.setattr(orchestrator_module, "_DEPRECATE_BATCH_SIZE", 1) + entry_ids = {"ep_20260601_0001", "ep_20260601_0002"} + + with pytest.raises(VectorStoreBusyError): + await orch._apply_deprecation_writes( + episodes=[_make_episode_row(entry_id=eid) for eid in sorted(entry_ids)], + to_deprecate=entry_ids, + owner_id="u_alice", + app_id="claude-code", + project_id="alpha", + merged_entry_id="ep_20260614_0001", + ) + + reverts = episode_store.update.await_args_list[2:] + reverts += atomic_fact_store.update.await_args_list + assert reverts + for call in reverts: + assert call.args[0] == {"deprecated_by": None} + where = call.kwargs["where"] + assert _eq_value(where, "app_id") == "claude-code" + assert _eq_value(where, "project_id") == "alpha" + + +async def test_run_caps_sources_per_merge_and_keeps_the_rest(monkeypatch) -> None: + """A cluster above the cap merges only its oldest sources; the deferred + members stay in the cluster and the count reflects them.""" + monkeypatch.setattr(orchestrator_module, "_MAX_SOURCES_PER_MERGE", 2) + cluster_repo = MagicMock() + episode_store = MagicMock() + atomic_fact_store = MagicMock() + episode_writer = MagicMock() + report_repo = MagicMock() + reflector = MagicMock() + embedder = MagicMock() + + report_repo.list_reflected_cluster_ids = AsyncMock(return_value=set()) + cluster_repo.list_ids_and_member_counts = AsyncMock(return_value=[("cl_abc", 3)]) + episode_store.find_where = AsyncMock(return_value=[]) + members = [ + ("ep_20260601_0001", "episode"), + ("ep_20260601_0002", "episode"), + ("ep_20260601_0003", "episode"), + ] + cluster_repo.get_members_with_type = AsyncMock(return_value=members) + rows = [ + _make_episode_row( + entry_id=f"ep_20260601_000{i}", + parent_id=f"mc_00{i}", + timestamp=_dt.datetime(2026, 6, i, tzinfo=_dt.UTC), + ) + for i in (3, 1, 2) + ] + episode_store.find_by_owner_entries = AsyncMock(return_value=rows) + reflector.areflect = AsyncMock( + return_value=_FakeAlgoResult( + owner_id=None, episode="merged", subject="s", timestamp=1717200000000 + ) + ) + episode_writer.append_entries = AsyncMock( + return_value=[_make_entry_id("ep_20260614_0001")] + ) + episode_writer.patch_frontmatter = AsyncMock() + cluster_repo.remove_members = AsyncMock() + cluster_repo.add_member = AsyncMock() + cluster_repo.update_metadata = AsyncMock() + embedder.embed = AsyncMock(return_value=[0.1] * 1024) + episode_store.update = AsyncMock() + atomic_fact_store.update = AsyncMock() + report_repo.create = AsyncMock() + + orch = _build_orchestrator( + cluster_repo=cluster_repo, + episode_store=episode_store, + atomic_fact_store=atomic_fact_store, + episode_writer=episode_writer, + report_repo=report_repo, + reflector=reflector, + embedder=embedder, + ) + reports = await orch.run(ctx=FakeStrategyContext(), owner_id="u_alice") + + # Only the two oldest sources were reflected. + reflected = reflector.areflect.await_args.args[0] + assert [ep.timestamp for ep in reflected] == [ + _ts_to_ms(_dt.datetime(2026, 6, 1, tzinfo=_dt.UTC)), + _ts_to_ms(_dt.datetime(2026, 6, 2, tzinfo=_dt.UTC)), + ] + # ... and only those were deprecated and removed from the cluster. + cluster_repo.remove_members.assert_awaited_once_with( + "cl_abc", {"ep_20260601_0001", "ep_20260601_0002"} + ) + deprecated = _in_values(episode_store.update.await_args.kwargs["where"], "entry_id") + assert sorted(deprecated) == ["ep_20260601_0001", "ep_20260601_0002"] + # The deferred member plus the merged episode remain. + assert cluster_repo.update_metadata.await_args.kwargs["count"] == 2 + assert len(reports) == 1 + assert reports[0].source_count == 2