From 549b663ad0c60f4d758136331bf94276fbad1c47 Mon Sep 17 00:00:00 2001 From: guoyuchen <2273345925@qq.com> Date: Sat, 24 Oct 2026 14:53:18 +0800 Subject: [PATCH] fix(auto-dream): checkpoint extraction and resume interrupted integration --- reme/steps/evolve/dream/checkpoint.py | 239 ++++++++++ reme/steps/evolve/dream/extract.py | 20 +- reme/steps/evolve/dream/finish.py | 13 + reme/steps/evolve/dream/integrate.py | 127 ++++-- reme/steps/evolve/dream/integrate.yaml | 12 + tests/unit/test_dream_checkpoint_versions.py | 8 +- tests/unit/test_dream_resume.py | 448 +++++++++++++++++++ 7 files changed, 827 insertions(+), 40 deletions(-) create mode 100644 reme/steps/evolve/dream/checkpoint.py create mode 100644 tests/unit/test_dream_resume.py diff --git a/reme/steps/evolve/dream/checkpoint.py b/reme/steps/evolve/dream/checkpoint.py new file mode 100644 index 00000000..ed2204cf --- /dev/null +++ b/reme/steps/evolve/dream/checkpoint.py @@ -0,0 +1,239 @@ +"""Workspace-local, atomic checkpoints for interrupted Dream invocations.""" + +import asyncio +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Literal +from uuid import uuid4 + +from pydantic import BaseModel, Field, model_validator + +from ....schema import DreamState, DreamUnit + + +def _receipt(state: DreamState, index: int) -> dict: + unit = state.units[index] + return next( + result + for result in state.integrate_results + if (result.get("unit"), result.get("bucket"), result.get("paths")) + == (unit["name"], unit["bucket"], unit["paths"]) + ) + + +def integration_lock(step) -> asyncio.Lock: + """Serialize digest integration and catalog finalization within one Application.""" + if step.app_context is None: + raise RuntimeError("dream integration requires app_context") + key = "dream_integration_lock" + lock = step.app_context.metadata.get(key) + if lock is None: + lock = asyncio.Lock() + step.app_context.metadata[key] = lock + if not isinstance(lock, asyncio.Lock): + raise TypeError(f"app_context.metadata[{key!r}] must be an asyncio.Lock") + return lock + + +class DreamCheckpoint(BaseModel): + """Versioned intermediate results; source watermarks remain in the dream catalog.""" + + version: Literal[1] = 1 + run_id: str = Field(default_factory=lambda: uuid4().hex) + sources: dict[str, tuple[int, int, str]] + state: DreamState + completed: list[int] = Field(default_factory=list) + targets: dict[str, str] = Field(default_factory=dict) + active: int | None = None + active_before: dict[str, tuple[int, int]] = Field(default_factory=dict) + + @model_validator(mode="after") + def validate_progress(self): + """Reject invalid progress rather than silently skipping unprocessed units.""" + for unit in self.state.units: + parsed = DreamUnit.model_validate(unit) + if not parsed.paths or any( + p not in self.sources or p not in self.state.changed_paths for p in parsed.paths + ): + raise ValueError("dream checkpoint unit references an unscanned source") + indices = [*self.completed, *([] if self.active is None else [self.active])] + if len(set(indices)) != len(indices) or any(i < 0 or i >= len(self.state.units) for i in indices): + raise ValueError("invalid dream checkpoint unit indices") + try: + for index in self.completed: + _receipt(self.state, index) + except StopIteration as error: + raise ValueError("completed dream unit has no integration receipt") from error + return self + + +def source_versions(workspace: Path, paths: dict[str, float]) -> dict[str, tuple[int, int, str]]: + """Fingerprint inputs before extraction, including edits that preserve mtimes.""" + versions = {} + for rel in paths: + path = (workspace / rel).resolve() + if not path.is_relative_to(workspace): + raise ValueError(f"dream source escapes workspace: {rel}") + stat = path.stat() + if stat.st_mtime != paths[rel]: + raise RuntimeError(f"dream source changed after scanning; retry the job: {rel}") + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + after = path.stat() + if (after.st_mtime_ns, after.st_size) != (stat.st_mtime_ns, stat.st_size): + raise RuntimeError(f"dream source changed while fingerprinting; retry the job: {rel}") + versions[rel] = (stat.st_mtime_ns, stat.st_size, digest) + return versions + + +class CheckpointStore: + """One checkpoint per scan window and extraction options, owned by this request.""" + + def __init__(self, step, state: DreamState, max_units: int): + workspace = Path(state.workspace).resolve() + self.workspace = workspace + options = [state.dates, state.daily_dir, state.hint, max_units, step.config_value("digest_dir")] + options.append(step.file_catalog.name) + key = hashlib.sha256(json.dumps(options, ensure_ascii=False).encode("utf-8")).hexdigest() + self.path = (workspace / step.config_value("metadata_dir") / "dream" / f"{key}.json").resolve() + if not self.path.is_relative_to(workspace): + raise ValueError("dream checkpoint escapes workspace") + self.logger = step.logger + self.record: DreamCheckpoint | None = None + self.resumed = False + + def read(self) -> DreamCheckpoint | None: + """Treat malformed/unsupported saved state as unusable, with a visible diagnostic.""" + try: + return DreamCheckpoint.model_validate_json(self.path.read_text(encoding="utf-8")) + except FileNotFoundError: + return None + except ValueError as error: + self.logger.warning(f"[AutoDream] ignoring invalid checkpoint {self.path.name}: {error}") + return None + + def restore(self, sources: dict[str, tuple[int, int, str]]) -> DreamState | None: + """Reuse extraction only when the entire scanned input window is unchanged.""" + record = self.read() + if record is None: + return None + if record.sources != sources: + self.logger.info("[AutoDream] checkpoint invalidated: source files changed; extracting from scratch") + return None + if Path(record.state.workspace).resolve() != self.workspace: + self.logger.warning("[AutoDream] ignoring checkpoint for another workspace") + return None + self.record = record + self.resumed = True + state = self._prepare() + self.logger.info( + f"[AutoDream] resuming: {len(record.completed)} of {len(state.units)} units completed from checkpoint, " + f"{len(sources)} sources unchanged; {len(state.units) - len(record.completed)} units pending", + ) + return state + + def publish(self, state: DreamState, sources: dict[str, tuple[int, int, str]]) -> DreamState: + """First extraction wins when identical invocations overlap within an Application.""" + restored = self.restore(sources) + if restored is not None: + return restored + self.record = DreamCheckpoint(sources=sources, state=state) + self._write() + return state + + def refresh(self) -> DreamState: + """Reload progress after acquiring the existing application-wide integration lock.""" + self.record = self._current() + return self._prepare() + + def complete(self, index: int, state: DreamState) -> None: + """Acknowledge a unit together with the content version of its actual digest target.""" + assert self.record is not None + target = _receipt(state, index)["target_path"] + version = self._target_version(target) + if version is None: + raise RuntimeError(f"Dream digest target disappeared before completion: {target}") + self.record.targets[target] = version + self.record.completed.append(index) + + def ensure_current(self, *, complete: bool) -> None: + """Reject stale finalizers and do not checkpoint missing or unacknowledged digest content.""" + current = self._current() + if complete: + if len(current.completed) != len(current.state.units): + raise RuntimeError("Dream integration progress changed before Finish; retry the job") + if self._changed_targets(current): + raise RuntimeError("Dream digest changed before Finish; retry the job") + + def _current(self) -> DreamCheckpoint: + current = self.read() + if current is None or self.record is None or current.run_id != self.record.run_id: + raise RuntimeError("Dream checkpoint was replaced or completed by another invocation; retry the job") + return current + + def _changed_targets(self, record: DreamCheckpoint) -> set[str]: + targets = {_receipt(record.state, index)["target_path"] for index in record.completed} + return { + path + for path in targets + if not record.targets.get(path) or record.targets[path] != self._target_version(path) + } + + def _target_version(self, relative: str) -> str | None: + path = (self.workspace / relative).resolve() + if not path.is_relative_to(self.workspace): + raise ValueError("dream digest target escapes workspace") + try: + with path.open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + except FileNotFoundError: + return None + + def _prepare(self) -> DreamState: + """Recheck landed content before skipping completed units, including after waiting for the lock.""" + assert self.record is not None + state = self.record.state + if Path(state.workspace).resolve() != self.workspace: + raise ValueError("dream checkpoint belongs to another workspace") + invalid = self._changed_targets(self.record) + if invalid: + self.logger.info(f"[AutoDream] resume recheck: completed digest targets changed: {sorted(invalid)}") + self.record.completed = [ + index for index in self.record.completed if _receipt(state, index)["target_path"] not in invalid + ] + state.integrate_results = [ + result for result in state.integrate_results if result["target_path"] not in invalid + ] + state.modified_paths = list(dict.fromkeys([*state.modified_paths, *sorted(invalid)])) + self.resumed = True + state.failed_units, state.failed_paths, state.errors, state.checkpoint_paths = [], [], [], [] + return state + + def save(self, state: DreamState) -> None: + """Fail closed on persistence errors so Finish cannot acknowledge lost progress.""" + self._current() + self.record.state = state + self._write() + + def remove(self) -> None: + """Only remove this generation, after the final catalog checkpoint is durable.""" + current = self.read() + if current is not None and self.record is not None and current.run_id == self.record.run_id: + self.path.unlink() + + def _write(self) -> None: + """Publish complete JSON with atomic replacement; never truncate the last checkpoint.""" + assert self.record is not None + self.path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary = tempfile.mkstemp(prefix=f".{self.path.stem}-", dir=self.path.parent) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + stream.write(self.record.model_dump_json(indent=2)) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, self.path) + finally: + Path(temporary).unlink(missing_ok=True) diff --git a/reme/steps/evolve/dream/extract.py b/reme/steps/evolve/dream/extract.py index d3d4c3b7..b359042f 100644 --- a/reme/steps/evolve/dream/extract.py +++ b/reme/steps/evolve/dream/extract.py @@ -8,6 +8,7 @@ from .._evolve import agent_reply_result_text from ....components import R from ....enumeration import DreamBucketEnum from ....schema import DreamExtractOutput, DreamState +from .checkpoint import CheckpointStore, source_versions from .utils import ( clean_paths, daily_dir, @@ -103,6 +104,20 @@ class DreamExtractStep(BaseStep): existing=existing, indexed=indexed, ) + checkpoint = CheckpointStore(self, state, max_units) + saved = checkpoint.read() + if changed or saved is not None: + versions = source_versions(workspace, existing) + restored = checkpoint.restore(versions) + if restored is not None: + self.context["_dream_checkpoint"] = checkpoint + return self._finish(restored, True, f"Resumed {len(restored.units)} extracted unit(s) from checkpoint") + if saved is not None: + # A previous failed dump may have advanced only in-memory mtimes. Validate content, too. + changed = list(dict.fromkeys([*changed, *(p for p in existing if versions[p] != saved.sources.get(p))])) + state.changed_paths = changed + state.unchanged_paths = [p for p in existing if p not in changed] + state.files_changed, state.files_unchanged = len(changed), len(state.unchanged_paths) if not changed: self.logger.info(f"[{self.name}] skip no changed input dates={','.join(dates)}") return self._finish(state, True, f"No changed dream input for {', '.join(dates)}") @@ -112,7 +127,7 @@ class DreamExtractStep(BaseStep): self.logger.warning(f"[{self.name}] skip no llm changed={len(changed)}") return self._finish(state, False, state.errors[-1]) - self.logger.info(f"[{self.name}] agent start changed={len(changed)} dates={len(dates)}") + self.logger.info(f"[{self.name}] extracting from scratch: changed={len(changed)} dates={len(dates)}") user_message = self.prompt_format( "extract_user_message", date=day, @@ -162,6 +177,9 @@ class DreamExtractStep(BaseStep): return self._finish(state, False, error) self.logger.info(f"[{self.name}] parse done keys={','.join(sorted(meta.keys())) if meta else '(none)'}") + state = checkpoint.publish(state, versions) + self.context["_dream_checkpoint"] = checkpoint + self.logger.info(f"[{self.name}] extraction checkpoint saved: {len(state.units)} units") answer = f"Extracted {len(state.units)} unit(s) from {len(changed)} changed file(s) across {len(dates)} day(s)" return self._finish(state, True, answer) diff --git a/reme/steps/evolve/dream/finish.py b/reme/steps/evolve/dream/finish.py index 980f93be..30aa6ca0 100644 --- a/reme/steps/evolve/dream/finish.py +++ b/reme/steps/evolve/dream/finish.py @@ -5,6 +5,7 @@ from pathlib import Path from ...base_step import BaseStep from ....components import R from ....schema import DreamState, FileNode +from .checkpoint import CheckpointStore, integration_lock from .utils import state_from_context, store_state, workspace_dir @@ -19,6 +20,15 @@ class DreamFinishStep(BaseStep): async def execute(self): assert self.context is not None state = state_from_context(self) + saved = self.context.get("_dream_checkpoint") + if isinstance(saved, CheckpointStore): + async with integration_lock(self): + saved.ensure_current(complete=not state.errors and not state.failed_units) + return await self._execute(state) + return await self._execute(state) + + async def _execute(self, state: DreamState): + """Persist catalog versions while holding the same lock as checkpointed integrations.""" workspace = Path(state.workspace).resolve() if state.workspace else workspace_dir(self) if self.file_catalog is None: raise RuntimeError("dream_finish_step requires file_catalog") @@ -46,6 +56,9 @@ class DreamFinishStep(BaseStep): self.logger.info(f"[{self.name}] catalog dump done") state.checkpoint_paths = [n.path for n in upserts if n.path in checkpoint] + saved = self.context.get("_dream_checkpoint") + if isinstance(saved, CheckpointStore) and self.persist and not state.errors and not state.failed_units: + saved.remove() state.summary = render_summary(state) store_state(self, state) self.context.response.success = not state.errors diff --git a/reme/steps/evolve/dream/integrate.py b/reme/steps/evolve/dream/integrate.py index f3946318..889a59f7 100644 --- a/reme/steps/evolve/dream/integrate.py +++ b/reme/steps/evolve/dream/integrate.py @@ -9,6 +9,7 @@ from .._evolve import agent_reply_result_text from ....components import R from ....enumeration import DreamBucketEnum from ....schema import IntegrateOutcome +from .checkpoint import CheckpointStore, integration_lock from .utils import llm_available, pack_paths, parse_structured_reply, state_from_context, store_state, workspace_dir _TOOLS = ("node_search", "read", "frontmatter_read", "write", "edit", "frontmatter_update") @@ -62,13 +63,8 @@ class DreamIntegrateStep(BaseStep): if not state.units: self.logger.info(f"[{self.name}] skip no units") return self._finish(state, not state.errors, "No dream units to integrate") - if not llm_available(self): - err = "no llm configured; dream integrate requires an LLM" - state.errors.append(err) - state.failed_units = state.units - state.failed_paths = sorted({p for u in state.units for p in u.get("paths", [])}) - self.logger.warning(f"[{self.name}] skip no llm units={len(state.units)}") - return self._finish(state, False, err) + if not isinstance(self.context.get("_dream_checkpoint"), CheckpointStore) and not llm_available(self): + return self._missing_llm(state, state.units) async with self._integration_lock(): return await self._execute_locked(state) @@ -76,6 +72,12 @@ class DreamIntegrateStep(BaseStep): async def _execute_locked(self, state): """Integrate units while serializing digest writes for this Application.""" + checkpoint = self.context.get("_dream_checkpoint") + if isinstance(checkpoint, CheckpointStore): + state = checkpoint.refresh() + pending = [unit for i, unit in enumerate(state.units) if i not in checkpoint.record.completed] + if pending and not llm_available(self): + return self._missing_llm(state, pending) workspace = Path(state.workspace).resolve() if state.workspace else workspace_dir(self) digest_dir = self.config_value("digest_dir") self.logger.info(f"[{self.name}] start units={len(state.units)} workspace={workspace} digest_dir={digest_dir}") @@ -85,12 +87,47 @@ class DreamIntegrateStep(BaseStep): # Compare the whole run so retries and later units updating a newly # created file still emit one "added" change for the final file. before = _snapshot_digest(workspace, digest_dir) + interrupted_index = None + if isinstance(checkpoint, CheckpointStore) and checkpoint.record.active is not None: + interrupted_index = checkpoint.record.active + # Preserve ambiguous writes even if an earlier unit now needs rechecking first. + _record_modified_paths(state, _changed_digest_paths(checkpoint.record.active_before, before)) for i, unit in enumerate(state.units, start=1): - await self._integrate_one(state, unit, i, workspace, digest_dir) + recovery_note = "" + if isinstance(checkpoint, CheckpointStore): + record = checkpoint.record + if i - 1 in record.completed: + self.logger.info(f"[{self.name}] resume skip completed unit {i}/{len(state.units)}") + continue + interrupted = interrupted_index == i - 1 + possible_writes = list(state.modified_paths) + if interrupted or checkpoint.resumed: + recovery_note = "\n\n" + self.prompt_format( + "integrate_resume_message", + possible_paths_json=json.dumps(sorted(set(possible_writes)), ensure_ascii=False), + ) + self.logger.info(f"[{self.name}] resume reconcile pending unit {i}/{len(state.units)}") + record.active = i - 1 + record.active_before = _snapshot_digest(workspace, digest_dir) + checkpoint.save(state) + failures = len(state.failed_units) + await self._integrate_one(state, unit, i, workspace, digest_dir, recovery_note=recovery_note) + if isinstance(checkpoint, CheckpointStore): + if len(state.failed_units) == failures: + checkpoint.complete(i - 1, state) + checkpoint.record.active = None + checkpoint.record.active_before = {} + checkpoint.save(state) + self.logger.info( + f"[{self.name}] progress checkpoint saved: " + f"{len(checkpoint.record.completed)}/{len(state.units)} units completed", + ) after = _snapshot_digest(workspace, digest_dir) + changed_paths = set(_changed_digest_paths(before, after)) | set(state.modified_paths) self.context["changes"] = [ - {"change": "modified" if path in before else "added", "path": path} - for path in _changed_digest_paths(before, after) + {"change": "added" if path in state.nodes_created or path not in before else "modified", "path": path} + for path in sorted(changed_paths) + if (workspace / path).is_file() ] state.failed_paths = sorted(set(state.failed_paths)) answer = ( @@ -99,20 +136,29 @@ class DreamIntegrateStep(BaseStep): ) return self._finish(state, not state.failed_units, answer) + def _missing_llm(self, state, pending): + """Keep pending sources eligible without invalidating already completed progress.""" + err = "no llm configured; dream integrate requires an LLM" + state.errors.append(err) + state.failed_units = pending + state.failed_paths = sorted({p for unit in pending for p in unit.get("paths", [])}) + self.logger.warning(f"[{self.name}] skip no llm units={len(pending)}") + return self._finish(state, False, err) + def _integration_lock(self) -> asyncio.Lock: """Return the Application-wide lock used by AutoDream integration.""" - if self.app_context is None: - raise RuntimeError("dream_integrate_step requires app_context") - key = "dream_integration_lock" - lock = self.app_context.metadata.get(key) - if lock is None: - lock = asyncio.Lock() - self.app_context.metadata[key] = lock - if not isinstance(lock, asyncio.Lock): - raise TypeError(f"app_context.metadata[{key!r}] must be an asyncio.Lock") - return lock + return integration_lock(self) - async def _integrate_one(self, state, unit: dict, index: int, workspace: Path, digest_dir: str) -> None: + async def _integrate_one( + self, + state, + unit: dict, + index: int, + workspace: Path, + digest_dir: str, + *, + recovery_note: str = "", + ) -> None: try: bucket = DreamBucketEnum(str(unit.get("bucket") or "")).value except ValueError: @@ -137,7 +183,8 @@ class DreamIntegrateStep(BaseStep): unit_summary=unit.get("summary", ""), unit_paths_json=json.dumps(paths, ensure_ascii=False, indent=2), material_blob=pack_paths(workspace, paths), - ), + ) + + recovery_note, system_prompt=self.prompt_format( f"integrate_system_prompt_{bucket}", workspace_dir=str(workspace), @@ -151,13 +198,17 @@ class DreamIntegrateStep(BaseStep): changed = _changed_digest_paths(before, after) _record_modified_paths(state, changed) created = len(changed) == 1 and changed[0] not in unit_before - if len(changed) == 1 and self._valid_target( - workspace, - digest_dir, - bucket, - changed[0], - action="CREATE" if created else "UPDATED", - existed_before=changed[0] in unit_before, + if ( + not recovery_note + and len(changed) == 1 + and self._valid_target( + workspace, + digest_dir, + bucket, + changed[0], + action="CREATE" if created else "UPDATED", + existed_before=changed[0] in unit_before, + ) ): self._record_recovered(state, unit, bucket, paths, changed[0], created, e) self.logger.warning( @@ -193,13 +244,17 @@ class DreamIntegrateStep(BaseStep): changed = _changed_digest_paths(before, after) _record_modified_paths(state, changed) created = len(changed) == 1 and changed[0] not in unit_before - if len(changed) == 1 and self._valid_target( - workspace, - digest_dir, - bucket, - changed[0], - action="CREATE" if created else "UPDATED", - existed_before=changed[0] in unit_before, + if ( + not recovery_note + and len(changed) == 1 + and self._valid_target( + workspace, + digest_dir, + bucket, + changed[0], + action="CREATE" if created else "UPDATED", + existed_before=changed[0] in unit_before, + ) ): self._record_recovered(state, unit, bucket, paths, changed[0], created, e) self.logger.warning( diff --git a/reme/steps/evolve/dream/integrate.yaml b/reme/steps/evolve/dream/integrate.yaml index 25d66250..6865575d 100644 --- a/reme/steps/evolve/dream/integrate.yaml +++ b/reme/steps/evolve/dream/integrate.yaml @@ -1,3 +1,15 @@ +integrate_resume_message: | + An earlier Dream integration attempt was interrupted or failed. Digest writes may already have landed without a + saved completion receipt. Inspect the existing digest and reconcile this unit before writing. Reuse existing + content and do not append duplicate evidence or create a duplicate node. Verify that the unit is fully integrated + before returning a successful receipt; do not assume that a changed file means the work is complete. + Files possibly written by earlier attempts: {possible_paths_json} + +integrate_resume_message_zh: | + 上一次 Dream 整合曾中断或失败,digest 文件可能已写入,但尚未保存完成回执。请在写入前检查已有 digest 并核对本 unit。 + 复用已有内容,不要重复追加证据或创建重复节点。确认本 unit 已完整整合后再返回成功回执,不能仅凭文件有变化就认定完成。 + 上一次尝试可能写入的文件:{possible_paths_json} + integrate_system_prompt_procedure: | You are the dream unit integration agent for the procedure bucket. The unit is a how-to-do-X abstraction: steps, methods, recipes, runbooks, executable diff --git a/tests/unit/test_dream_checkpoint_versions.py b/tests/unit/test_dream_checkpoint_versions.py index ec97e112..0806dca3 100644 --- a/tests/unit/test_dream_checkpoint_versions.py +++ b/tests/unit/test_dream_checkpoint_versions.py @@ -153,7 +153,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older index = f"daily/{DAY}.md" (tmp_path / index).write_text("day index") - def agent_for(app, write_attempt): + def agent_for(app, write_attempt, *, valid_receipt=False): calls = {"extract": 0, "integrate": 0} async def reply(_message, **kwargs): @@ -167,6 +167,8 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older if calls["integrate"] == write_attempt: written = await app.run_job("write", path=TARGET, content=OLD) assert written.success, written.answer + if valid_receipt: + return {"result": json.dumps({"action": "CREATE", "target_path": TARGET})} return {"result": "{}"} agent = AsyncMock(spec=BaseAgentWrapper) @@ -217,7 +219,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older catalog = app.context.components["file_catalog"]["dream"] if recover_on_attempt is None: assert {n.path: n.st_mtime for n in await catalog.get_nodes()} == previous - agent, calls = agent_for(app, 1) + agent, calls = agent_for(app, 1, valid_receipt=True) response = await app.run_job( "auto_dream", date=DAY, @@ -226,7 +228,7 @@ async def test_invalid_integration_preserves_retry_after_restart(tmp_path, older as_llm=Mock(spec=ChatModelBase), ) assert response.success, response.answer - assert calls == {"extract": int(recover_on_attempt is None), "integrate": int(recover_on_attempt is None)} + assert calls == {"extract": 0, "integrate": int(recover_on_attempt is None)} assert {SOURCE, index} <= {n.path for n in await catalog.get_nodes()} assert OLD in (tmp_path / TARGET).read_text() assert source.read_text() == original diff --git a/tests/unit/test_dream_resume.py b/tests/unit/test_dream_resume.py new file mode 100644 index 00000000..81e5516d --- /dev/null +++ b/tests/unit/test_dream_resume.py @@ -0,0 +1,448 @@ +"""Exercise Dream recovery with fresh Applications and a controlled local provider.""" + +# pylint: disable=protected-access + +import asyncio +import copy +import json +import os +from contextlib import asynccontextmanager, nullcontext +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from agentscope.model import ChatModelBase + +from reme.application import Application +from reme.components.agent_wrapper import BaseAgentWrapper +from reme.config import resolve_app_config +from reme.steps.evolve.dream.checkpoint import CheckpointStore, DreamCheckpoint, source_versions +from reme.steps.evolve.dream.finish import DreamFinishStep +from reme.steps.evolve.dream.integrate import DreamIntegrateStep + +DAY = "2026-10-23" +SOURCE = f"daily/{DAY}/source.md" +TARGETS = ["digest/wiki/first.md", "digest/wiki/second.md"] +UNITS = [ + {"name": name, "bucket": "wiki", "summary": f"{name} memory", "paths": [SOURCE]} for name in ("first", "second") +] + + +class DreamCase: + """Own isolated files, fresh application lifecycles, and provider call observations.""" + + def __init__(self, workspace): + self.workspace, self.metadata = workspace, "metadata" + self.source = workspace / SOURCE + self.source.parent.mkdir(parents=True) + self.source.write_text("Original source", encoding="utf-8") + self.units, self.calls, self.prompts = copy.deepcopy(UNITS), [], [] + + @property + def checkpoints(self): + """Return only committed JSON checkpoints, excluding atomic-write temporary files.""" + return list((self.workspace / self.metadata / "dream").glob("*.json")) + + @property + def saved(self): + """Read progress from disk, without borrowing any application state.""" + assert len(self.checkpoints) == 1 + return DreamCheckpoint.model_validate_json(self.checkpoints[0].read_text()) + + @asynccontextmanager + async def open(self): + """Start real local catalogs and jobs without external services or model providers.""" + defaults = resolve_app_config(config="default", log_config=False) + types = ("tokenizer", "keyword_index", "file_graph", "tag_index", "file_store", "file_catalog") + app = Application( + workspace_dir=str(self.workspace), + metadata_dir=self.metadata, + enable_logo=False, + log_to_file=False, + log_to_console=False, + service={"backend": "http", "web_enabled": False, "mcp_enabled": False}, + components={name: copy.deepcopy(defaults["components"][name]) for name in types}, + jobs={name: copy.deepcopy(defaults["jobs"][name]) for name in ("auto_dream", "write")}, + ) + self.calls.clear() + self.prompts.clear() + await app.start() + try: + yield app + finally: + await app.close() + + async def write(self, app, index, content): + """Exercise the normal workspace-contained write job.""" + response = await app.run_job("write", path=TARGETS[index], content=content) + assert response.success, response.answer + + @staticmethod + def receipt(index, action): + """Return a valid integration receipt for a controlled target.""" + return {"result": json.dumps({"action": action, "target_path": TARGETS[index]})} + + async def run(self, app, *, reply=None, interrupt="", fail=False, **options): + """Track extraction and integration independently; custom replies control only integration.""" + + async def provider(message, **kwargs): + if "list_tags" in kwargs["job_tools"]: + return {"result": "No tags added"} + if kwargs["job_tools"] == ["read"]: + self.calls.append("extract") + return {"result": json.dumps({"units": self.units})} + text = str(message) + index = 0 if "first memory" in text else 1 + self.calls.append(UNITS[index]["name"]) + self.prompts.append(text) + if reply: + return await reply(index, text) + if index == 1 and interrupt == "before_write": + raise asyncio.CancelledError + if index == 1 and fail: + return {"result": "{}"} + existed = (self.workspace / TARGETS[index]).exists() + if not existed: + await self.write(app, index, UNITS[index]["summary"]) + if index == 1 and interrupt == "after_write": + raise asyncio.CancelledError + return self.receipt(index, "CORROBORATE" if existed else "CREATE") + + agent = AsyncMock(spec=BaseAgentWrapper) + agent.reply.side_effect = provider + return await app.run_job( + "auto_dream", + date=DAY, + scan_days=1, + agent_wrapper=agent, + as_llm=Mock(spec=ChatModelBase), + **options, + ) + + async def interrupt(self, stage): + """Leave disk state at a named interruption boundary, then discard the Application.""" + boundary = { + "after_extract": (DreamIntegrateStep, "execute"), + "before_finish": (DreamFinishStep, "execute"), + "after_commit": (CheckpointStore, "remove"), + }.get(stage) + async with self.open() as app: + with ( + patch.object(*boundary, side_effect=asyncio.CancelledError) if boundary else nullcontext(), + pytest.raises(asyncio.CancelledError), + ): + await self.run(app, interrupt=stage) + if stage != "after_commit": + assert await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE]) == [] + + +@pytest.fixture(name="case") +def dream_case(tmp_path): + """Each case starts with user-owned source material in a temporary workspace.""" + return DreamCase(tmp_path) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["after_extract", "before_write", "after_write", "before_finish", "no_llm"]) +async def test_restart_reuses_extraction_and_completed_units(case, stage): + """Resume only pending units; completed work and source content stay untouched.""" + original = case.source.read_bytes() + await case.interrupt("before_finish" if stage == "no_llm" else stage) + saved, logs = case.saved, [] + assert saved.state.units == UNITS + assert saved.completed == ([] if stage == "after_extract" else [0] if "write" in stage else [0, 1]) + if "write" in stage: + assert saved.active == 1 + restore = CheckpointStore.restore + + def capture(store, sources): + store.logger = Mock(wraps=store.logger) + state = restore(store, sources) + logs.extend(call.args[0] for call in store.logger.info.call_args_list) + return state + + async with case.open() as app: + with ( + patch.object(CheckpointStore, "restore", capture), + ( + patch("reme.steps.evolve.dream.integrate.llm_available", return_value=False) + if stage == "no_llm" + else nullcontext() + ), + ): + response = await case.run(app) + assert response.success, response.answer + assert case.calls == ( + ["first", "second"] if stage == "after_extract" else ["second"] if "write" in stage else [] + ) + assert len(response.metadata["dream"]["integrate_results"]) == 2 + assert SOURCE in response.metadata["dream"]["checkpoint_paths"] + assert any(f"resuming: {len(saved.completed)} of 2 units completed" in line for line in logs) + assert not case.checkpoints + assert case.source.read_bytes() == original + for index, path in enumerate(TARGETS): + assert (case.workspace / path).read_text().strip() == UNITS[index]["summary"] + if stage == "after_write": + assert TARGETS[1] in case.prompts[0] and "do not append duplicate evidence" in case.prompts[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "change", + ["edit", "same_mtime", "add", "delete", "hint", "max_units", "corrupt", "version", "progress"], +) +async def test_incompatible_checkpoint_reextracts(case, change): + """Never silently reuse extraction for changed inputs or invalid saved progress.""" + await case.interrupt("after_extract") + options, path = {}, case.checkpoints[0] + if change in ("edit", "same_mtime"): + stat = case.source.stat() + case.source.write_text("Modified source") + if change == "same_mtime": + os.utime(case.source, ns=(stat.st_atime_ns, stat.st_mtime_ns)) + elif change == "add": + (case.source.parent / "another.md").write_text("New source") + elif change == "delete": + case.source.unlink() + elif change in ("hint", "max_units"): + options[change] = "New guidance" if change == "hint" else 1 + elif change == "corrupt": + path.write_text("{incomplete JSON") + else: + data = json.loads(path.read_text()) + data["version" if change == "version" else "completed"] = 999 if change == "version" else [0] + path.write_text(json.dumps(data)) + case.units = [] + async with case.open() as app: + response = await case.run(app, **options) + assert response.success, response.answer + assert case.calls == ["extract"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty", [False, True]) +async def test_restart_after_catalog_commit_only_cleans_up(case, empty): + """A durable final watermark needs no new model work, including empty extraction.""" + case.metadata = "custom-metadata" + if empty: + case.units = [] + await case.interrupt("after_commit") + assert case.checkpoints and not (case.workspace / "metadata").exists() + async with case.open() as app: + response = await case.run(app) + assert response.success, response.answer + assert not case.calls and not case.checkpoints + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sources", ["shared", "independent", "missing_target"]) +async def test_partial_failure_retries_only_pending_work(case, sources): + """Independent source watermarks cannot hide saved units needing target repair.""" + if sources != "shared": + other = f"daily/{DAY}/other.md" + (case.workspace / other).write_text("Another source") + case.units[1]["paths"] = [other] + async with case.open() as app: + response = await case.run(app, fail=True) + assert not response.success + dream = response.metadata["dream"] + assert (SOURCE in dream["checkpoint_paths"]) is (sources != "shared") + assert case.units[1]["paths"][0] not in dream["checkpoint_paths"] + if sources == "missing_target": + (case.workspace / TARGETS[0]).unlink() + async with case.open() as app: + response = await case.run(app) + assert response.success, response.answer + assert case.calls == (["first", "second"] if sources == "missing_target" else ["second"]) + assert len(response.metadata["dream"]["integrate_results"]) == 2 + assert not response.metadata["dream"]["failed_units"] and not response.metadata["dream"]["errors"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_receipt", ["invalid", "exception"]) +@pytest.mark.parametrize("stage", ["after_extract", "before_write"]) +async def test_resume_requires_receipt_despite_partial_writes(case, bad_receipt, stage): + """Changed files without valid receipts must not become durable completion markers.""" + await case.interrupt(stage) + async with case.open() as app: + + async def reply(index, text): + assert "do not append duplicate evidence" in text + await case.write(app, index, f"Partial memory, attempt {len(case.calls)}") + if bad_receipt == "exception": + raise RuntimeError("incomplete integration") + return {"result": "{}"} + + response = await case.run(app, reply=reply) + assert not response.success + assert len(case.calls) == (4 if stage == "after_extract" else 2) + assert SOURCE not in response.metadata["dream"]["checkpoint_paths"] + assert case.saved.completed == ([] if stage == "after_extract" else [0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["delete", "rewrite", "later_interrupted_write"]) +async def test_recheck_only_affected_digest_units(case, change): + """Rechecking an earlier target also preserves evidence of a later interrupted write.""" + interrupted = change == "later_interrupted_write" + await case.interrupt("after_write" if interrupted else "before_finish") + first, second = (case.workspace / path for path in TARGETS) + second_version = second.stat().st_mtime_ns + if change == "delete": + first.unlink() + else: + first.write_text("Overwritten memory") + async with case.open() as app: + + async def reply(index, text): + assert "do not append duplicate evidence" in text + if interrupted: + assert TARGETS[1] in text + if index == 0: + existed = first.exists() + await case.write(app, index, UNITS[index]["summary"]) + return case.receipt(index, "CORRECT" if existed else "CREATE") + assert interrupted + return case.receipt(index, "CORROBORATE") + + response = await case.run(app, reply=reply) + assert response.success, response.answer + assert case.calls == (["first", "second"] if interrupted else ["first"]) + assert first.read_text().strip() == UNITS[0]["summary"] + assert second.stat().st_mtime_ns == second_version + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["replacement", "catalog"]) +async def test_failed_persistence_keeps_recoverable_checkpoint(case, failure): + """Neither atomic replacement failure nor repeated catalog failure can discard progress.""" + async with case.open() as app: + catalog = app.context.components["file_catalog"]["dream"] + original_write, writes = CheckpointStore._write, 0 + + def replace(store): + nonlocal writes + writes += 1 + with ( + ( + patch("reme.steps.evolve.dream.checkpoint.os.replace", side_effect=OSError("disk full")) + if writes == 3 + else nullcontext() + ), + ): + return original_write(store) + + with ( + ( + patch.object(CheckpointStore, "_write", replace) + if failure == "replacement" + else patch.object(catalog, "dump", side_effect=OSError("disk full")) + ), + ): + response = await case.run(app) + assert not response.success and "disk full" in response.answer + assert case.checkpoints + if failure == "replacement": + assert case.saved.active == 0 and case.saved.completed == [] + assert await catalog.get_nodes([SOURCE]) == [] + assert len(list(case.checkpoints[0].parent.iterdir())) == 1 + else: + case.calls.clear() + assert not (await case.run(app)).success + assert not case.calls and case.checkpoints + case.calls.clear() + response = await case.run(app) + assert response.success, response.answer + assert "extract" not in case.calls and not case.checkpoints + + +@pytest.mark.asyncio +async def test_checkpoint_cannot_escape_workspace(case): + """A metadata symlink must not allow checkpoint writes outside the workspace.""" + outside = case.workspace / "outside" + outside.mkdir() + nested = DreamCase(case.workspace / "nested") + (nested.workspace / "metadata").symlink_to(outside, target_is_directory=True) + async with nested.open() as app: + response = await nested.run(app) + assert not response.success and "checkpoint escapes workspace" in response.answer + assert not nested.calls and not (outside / "dream").exists() + + +@pytest.mark.asyncio +async def test_edit_between_scan_and_extraction_keeps_source_eligible(case): + """Do not attach a newer extraction to the scan's older source watermark.""" + + def edit_then_hash(workspace, existing): + previous = case.source.stat().st_mtime + case.source.write_text("Modified after scanning") + os.utime(case.source, (previous + 2, previous + 2)) + return source_versions(workspace, existing) + + async with case.open() as app: + with patch("reme.steps.evolve.dream.extract.source_versions", edit_then_hash): + response = await case.run(app) + assert not response.success and "changed after scanning" in response.answer + assert not case.calls + assert await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE]) == [] + assert (await case.run(app)).success + assert case.calls == ["extract", "first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("overlap", ["integration", "finish"]) +async def test_overlapping_jobs_do_not_replay_or_regress_progress(case, overlap): + """Queued integrations refresh progress; stale finalizers cannot overwrite newer watermarks.""" + entered, release = asyncio.Event(), asyncio.Event() + boundary = DreamIntegrateStep if overlap == "integration" else DreamFinishStep + method = "_execute_locked" if overlap == "integration" else "execute" + original, invocations = getattr(boundary, method), 0 + + async def delay(step, *args): + nonlocal invocations + invocations += 1 + if invocations == 1: + entered.set() + await release.wait() + return await original(step, *args) + + async with case.open() as app: + tasks = [] + try: + with patch.object(boundary, method, delay): + tasks.append(asyncio.create_task(case.run(app))) + await asyncio.wait_for(entered.wait(), timeout=5) + if overlap == "finish": + previous = case.source.stat().st_mtime + case.source.write_text("New source version") + os.utime(case.source, (previous + 2, previous + 2)) + assert (await case.run(app)).success + else: + restored = asyncio.Event() + restore = CheckpointStore.restore + + def signal(store, sources): + result = restore(store, sources) + restored.set() + return result + + with patch.object(CheckpointStore, "restore", signal): + tasks.append(asyncio.create_task(case.run(app))) + await asyncio.wait_for(restored.wait(), timeout=5) + release.set() + responses = await asyncio.gather(*tasks) + if overlap == "finish": + assert not responses[0].success and "another invocation" in responses[0].answer + nodes = await app.context.components["file_catalog"]["dream"].get_nodes([SOURCE]) + assert nodes[0].st_mtime == case.source.stat().st_mtime + case.calls.clear() + assert (await case.run(app)).success and not case.calls + else: + assert any(r.success for r in responses) + assert all(r.success or "another invocation" in r.answer for r in responses) + assert case.calls == ["extract", "first", "second"] + finally: + release.set() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True)