From bc7e16b813c3bfb4c39babceb8b95f13f8ffaaa6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Mon, 26 Oct 2026 20:34:34 +0800 Subject: [PATCH] Cache operation receipts and avoid rereading completed journals --- .../src/reme_procedure_memory/store.py | 91 +++++++++++----- .../procedure-memory/tests/test_recovery.py | 101 +++++++++++++++++- 2 files changed, 162 insertions(+), 30 deletions(-) diff --git a/plugins/procedure-memory/src/reme_procedure_memory/store.py b/plugins/procedure-memory/src/reme_procedure_memory/store.py index 310715ee..98160040 100644 --- a/plugins/procedure-memory/src/reme_procedure_memory/store.py +++ b/plugins/procedure-memory/src/reme_procedure_memory/store.py @@ -66,6 +66,9 @@ class ProcedureMemoryStore(BaseComponent): self.index_error: str | None = None self.rebuild_pending = False self.pending_indexes: dict[str, dict] = {} + self._operations: dict[str, dict] = {} + self._operations_by_source: dict[str, dict[int, set[str]]] = {} + self._pending_operations: dict[str, set[str]] = {} @property def views(self) -> dict: @@ -86,6 +89,7 @@ class ProcedureMemoryStore(BaseComponent): # A crash can lose cached index dumps, even for an acknowledged write. # Rebuild from user files at startup, never from cached chunks alone. async with self.data_lock: + self._load_operations() await self._recover_unlocked() result = await self._rebuild_unlocked() for error in result["errors"]: @@ -143,18 +147,50 @@ class ProcedureMemoryStore(BaseComponent): self.scopes[token] = {"source": source, "segment": segment} return token + def _remember_operation(self, journal: Path, record: dict) -> None: + """Cache receipts and lookup keys, never the full experience payload.""" + name = journal.name + self._operations[name] = { + key: value for key, value in record.items() if key not in {"document", "previous", "signature"} + } + segments = self._operations_by_source.setdefault(record["source"], {}) + segments.setdefault(record["segment"], set()).add(name) + pending = self._pending_operations.setdefault(record["path"], set()) + if record["status"] in {"indexed", "superseded"}: + pending.discard(name) + if not pending: + self._pending_operations.pop(record["path"]) + else: + pending.add(name) + + def _write_operation(self, journal: Path, record: dict) -> None: + atomic_json(journal, record) + self._remember_operation(journal, record) + + def _load_operations(self) -> None: + """Rebuild the in-memory directory once at startup from durable journals.""" + self._operations.clear() + self._operations_by_source.clear() + self._pending_operations.clear() + root = self.state_root / "operations" + for path in sorted(root.glob("*.json")): + self._remember_operation(path, load_json(path)) + # Completed replacements can still name older intents. Resolve all + # chains before replay without rereading completed journal payloads. + for name, record in list(self._operations.items()): + self._mark_superseded(root / name, record) + def operation_records(self, source: str, segment: int | None = None) -> list[dict]: - """Recover acknowledged and in-flight tool context for this segment.""" + """Return acknowledged and in-flight tool context without reading journals.""" + segments = self._operations_by_source.get(source, {}) + names = set().union(*segments.values()) if segment is None else segments.get(segment, set()) records = [] - for path in sorted((self.state_root / "operations").glob("*.json")): - record = load_json(path) - if record["source"] == source and (segment is None or record["segment"] == segment): - records.append( - {key: record[key] for key in ("operation_key", "path", "version", "status", "action")}, - ) - for key in ("index_error", "superseded_by"): - if key in record: - records[-1][key] = record[key] + for name in sorted(names): + record = self._operations[name] + records.append({key: record[key] for key in ("operation_key", "path", "version", "status", "action")}) + for key in ("index_error", "superseded_by"): + if key in record: + records[-1][key] = record[key] return records async def apply_mutation(self, action: str, arguments: dict, scope_token: str) -> dict: @@ -186,15 +222,10 @@ class ProcedureMemoryStore(BaseComponent): # Persist the replacement's dependency list before touching either # file. Recovery can then skip the old intent even if we crash # before marking it superseded or publishing the replacement. - record["supersedes"] = [ - journal.name - for journal in (self.state_root / "operations").glob("*.json") - if (prior := load_json(journal))["path"] == record["path"] - and prior["status"] not in {"indexed", "superseded"} - ] + record["supersedes"] = sorted(self._pending_operations.get(record["path"], set())) path = self.experience_path(record["path"]) async with await get_path_lock(path): - atomic_json(operation_path, record) + self._write_operation(operation_path, record) self._mark_superseded(operation_path, record) await self._commit_record(operation_path, record) return self._operation_result(record) @@ -256,19 +287,19 @@ class ProcedureMemoryStore(BaseComponent): raise ValueError(f"Journal conflicts with current experience: {record['path']}") atomic_json(path, record["document"]) record["status"] = "committed" - atomic_json(journal, record) + self._write_operation(journal, record) try: await self._sync_unlocked([record["path"]]) except Exception as exc: record["index_error"] = str(exc) - atomic_json(journal, record) + self._write_operation(journal, record) await self._quarantine_unlocked(record["path"], str(exc)) raise ProcedureIndexError( {**self._operation_result(record), "status": "index_pending", "error": str(exc)}, ) from exc record["status"] = "indexed" record.pop("index_error", None) - atomic_json(journal, record) + self._write_operation(journal, record) self.pending_indexes.pop(record["path"], None) def _mark_superseded(self, replacement: Path, record: dict) -> None: @@ -277,13 +308,14 @@ class ProcedureMemoryStore(BaseComponent): prior_path = replacement.parent / name if prior_path.parent != replacement.parent or prior_path.suffix != ".json": raise ValueError("Invalid superseded operation path") - prior = load_json(prior_path) + prior = self._operations[name] if prior["path"] != record["path"]: raise ValueError("A replacement cannot supersede a different experience") if prior["status"] != "superseded": + prior = load_json(prior_path) prior["status"] = "superseded" prior["superseded_by"] = record["operation_key"] - atomic_json(prior_path, prior) + self._write_operation(prior_path, prior) @staticmethod def _operation_result(record: dict) -> dict: @@ -293,18 +325,19 @@ class ProcedureMemoryStore(BaseComponent): return result async def _recover_unlocked(self) -> None: - journals = sorted((self.state_root / "operations").glob("*.json")) + root = self.state_root / "operations" + names = sorted(set().union(*self._pending_operations.values())) + records = {name: load_json(root / name) for name in names} # Resolve replacement chains before replay, independent of hashed file # ordering and of how far the interrupted replacement had progressed. - for path in journals: - self._mark_superseded(path, load_json(path)) + for name, record in records.items(): + self._mark_superseded(root / name, record) attempted = set() - for path in journals: - record = load_json(path) - if record["status"] not in {"indexed", "superseded"}: + for name, record in records.items(): + if self._operations[name]["status"] not in {"indexed", "superseded"}: attempted.add(record["path"]) try: - await self._commit_record(path, record) + await self._commit_record(root / name, record) except ProcedureIndexError: pass # Keep this source file repairable without blocking other files. if self.rebuild_pending: diff --git a/plugins/procedure-memory/tests/test_recovery.py b/plugins/procedure-memory/tests/test_recovery.py index 509a736a..eeb5a70c 100644 --- a/plugins/procedure-memory/tests/test_recovery.py +++ b/plugins/procedure-memory/tests/test_recovery.py @@ -497,14 +497,17 @@ async def test_pending_payload_survives_restart_and_can_still_be_repaired( @pytest.mark.asyncio @pytest.mark.parametrize("stage", ["supersede", "publish", "indexed"]) +@pytest.mark.parametrize("restart", [False, True]) async def test_repair_intent_recovers_without_replaying_old_payload( runtime, experience, reject_payload, monkeypatch, stage, + restart, + app_factory, ): - _, store, scope = runtime + app, store, scope = runtime reject_payload(store.views["when_to_use"].embedding_store) with pytest.raises(RuntimeError, match="indexes pending"): await store.apply_mutation("create", {**experience, "content": "unembeddable-payload"}, scope) @@ -530,6 +533,10 @@ async def test_repair_intent_recovers_without_replaying_old_payload( with pytest.raises(RuntimeError, match="repair interrupted"): await store.apply_mutation("update", correction, scope) monkeypatch.setattr(store_module, "atomic_json", original) + if restart: + await app.close() + _, store = await app_factory(workspace=store.workspace_path) + scope = store.open_scope("session/trajectories/input.log", 0) await store.reconcile() assert (await store.apply_mutation("update", correction, scope))["status"] == "indexed" assert (await store.read_experience(saved["path"]))["experience"]["content"] == "valid replacement" @@ -706,3 +713,95 @@ async def test_failed_update_hides_previous_projection_and_allows_another_correc result = await store.search("repair", "when_to_use_content") assert result["index_complete"] assert result["results"][0]["experience"]["content"] == "valid procedure" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("restart", [False, True]) +async def test_completed_history_is_not_read_on_search_or_write( + runtime, + experience, + monkeypatch, + app_factory, + restart, +): + app, store, scope = runtime + source = "session/trajectories/input.log" + current = await store.apply_mutation("create", experience, scope) + for number in range(32): + current = await store.apply_mutation( + "update", + { + "path": current["path"], + "expected_version": current["version"], + "operation_key": f"update-{number}", + "changes": {"content": f"repair index revision {number}"}, + }, + store.open_scope(source, number % 3), + ) + root = store.state_root / "operations" + expected = {segment: store.operation_records(source, segment) for segment in (None, 0, 1, 2)} + reads = [] + original = store_module.load_json + + def track(path): + if path.parent == root: + reads.append(path.name) + return original(path) + + monkeypatch.setattr(store_module, "load_json", track) + if restart: + await app.close() + _, store = await app_factory(workspace=store.workspace_path) + scope = store.open_scope(source, 0) + # Indexed journals are decoded only once during startup. + assert len(reads) == len(set(reads)) == 33 + reads.clear() + for segment, records in expected.items(): + assert store.operation_records(source, segment) == records + assert store.operation_records("unknown") == [] + assert all("document" not in receipt for receipt in store._operations.values()) + for _ in range(3): + assert (await store.search("repair", "when_to_use"))["index_complete"] + await store.reconcile(source=source, segment=0) + await store.apply_mutation( + "update", + { + "path": current["path"], + "expected_version": current["version"], + "operation_key": "new-update", + "changes": {"content": "repair latest index"}, + }, + scope, + ) + assert not reads + await store.apply_mutation("create", experience, scope) + assert len(reads) == 1 # An idempotent retry reads just its own journal. + + +@pytest.mark.asyncio +async def test_recovery_reads_only_pending_history(runtime, experience, reject_payload, monkeypatch): + _, store, scope = runtime + await store.apply_mutation("create", experience, scope) + reject_payload(store.views["when_to_use"].embedding_store) + with pytest.raises(RuntimeError, match="indexes pending"): + await store.apply_mutation( + "create", + {**experience, "filename": "pending.json", "operation_key": "bad", "content": "unembeddable-payload"}, + scope, + ) + root = store.state_root / "operations" + pending = next(path.name for path in root.glob("*.json") if load_json(path)["status"] == "committed") + original = store_module.load_json + reads = [] + + def track(path): + if path.parent == root: + reads.append(path.name) + return original(path) + + monkeypatch.setattr(store_module, "load_json", track) + for _ in range(3): + result = await store.search("repair", "when_to_use_content") + assert not result["index_complete"] + assert len(result["results"]) == 1 + assert reads == [pending] * 3