fix(auto-dream): validate extraction before checkpointing inputs (#649)

* fix(auto-dream): validate extraction before checkpointing inputs

* fix(auto-dream): normalize bucket whitespace before validation

reme/steps/evolve/dream/extract.py

@@ -7,7 +7,7 @@ from ...file_io import refresh_day_index
 from .._evolve import agent_reply_result_text
 from ....components import R
 from ....enumeration import DreamBucketEnum
-from ....schema import DreamState
+from ....schema import DreamExtractOutput, DreamState
 from .utils import (
     clean_paths,
     daily_dir,
@@ -113,32 +113,43 @@ class DreamExtractStep(BaseStep):
             return self._finish(state, False, state.errors[-1])
 
         self.logger.info(f"[{self.name}] agent start changed={len(changed)} dates={len(dates)}")
-        raw_result, meta = "", {}
+        user_message = self.prompt_format(
+            "extract_user_message",
+            date=day,
+            dates_json=json.dumps(dates, ensure_ascii=False, indent=2),
+            hint=hint or "(none)",
+            max_units=max_units,
+            changed_paths_json=json.dumps(changed, ensure_ascii=False, indent=2),
+            material_blob=pack_paths(workspace, changed),
+        )
+        retry_feedback = ""
+        output_schema = DreamExtractOutput.model_json_schema()
+        output_schema["required"] = ["units"]
         for attempt in range(2):
             try:
                 result = await self.agent_wrapper.reply(
-                    self.prompt_format(
-                        "extract_user_message",
-                        date=day,
-                        dates_json=json.dumps(dates, ensure_ascii=False, indent=2),
-                        hint=hint or "(none)",
-                        max_units=max_units,
-                        changed_paths_json=json.dumps(changed, ensure_ascii=False, indent=2),
-                        material_blob=pack_paths(workspace, changed),
-                    ),
+                    user_message + retry_feedback,
                     system_prompt=self.prompt_format(
                         "extract_system_prompt",
                         workspace_dir=str(workspace),
                         buckets=", ".join(bucket.value for bucket in DreamBucketEnum),
                         max_units=max_units,
+                        output_schema_json=json.dumps(output_schema, ensure_ascii=False, indent=2),
                     ),
                     job_tools=list(_TOOLS),
                 )
                 self.logger.info(f"[{self.name}] agent done has_result={bool(result.get('result'))}")
                 raw_result = agent_reply_result_text(result)
-                meta = parse_structured_reply(raw_result)
+                state.extract_summary = raw_result
+                meta = parse_structured_reply(raw_result, strict=True)
+                self.clean_output(state, meta, max_units=max_units)
+                break
             except Exception as e:  # noqa: BLE001
                 if attempt == 0:
+                    retry_feedback = "\n\n" + self.prompt_format(
+                        "extract_retry_message",
+                        error=f"{type(e).__name__}: {e}",
+                    )
                     self.logger.warning(
                         f"[{self.name}] extract attempt 1 returned no usable receipt; retrying once: "
                         f"{type(e).__name__}: {e}",
@@ -150,20 +161,7 @@ class DreamExtractStep(BaseStep):
                 self.logger.error(f"[{self.name}] {error}")
                 return self._finish(state, False, error)
 
-            units = meta.get("units") if "units" in meta else meta.get("memory_units")
-            if isinstance(units, list):
-                break
-            if attempt == 0:
-                self.logger.warning(f"[{self.name}] extract attempt 1 returned an unusable receipt; retrying once")
-                continue
-            # Keep the warning-only result checkpointable after one retry so a bad source cannot loop forever.
-            warning = "dream extract skipped unusable agent receipt after retry; expected a units list"
-            state.warnings.append(warning)
-            self.logger.warning(f"[{self.name}] {warning}")
-
         self.logger.info(f"[{self.name}] parse done keys={','.join(sorted(meta.keys())) if meta else '(none)'}")
-        self.clean_output(state, meta, max_units=max_units)
-        state.extract_summary = raw_result
         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)
 
@@ -177,25 +175,34 @@ class DreamExtractStep(BaseStep):
         return out
 
     def clean_output(self, state: DreamState, meta: dict, max_units: int | None = None) -> None:
-        """Clean up output"""
+        """Validate the whole receipt before accepting any units or applying the cap."""
+        units = meta.get("units") if "units" in meta else meta.get("memory_units")
+        if not isinstance(units, list):
+            raise ValueError("expected a units list (use units: [] only when nothing is worth remembering)")
+        normalized = []
+        for raw in units:
+            if isinstance(raw, dict) and isinstance(raw.get("bucket"), str):
+                raw = {**raw, "bucket": raw["bucket"].strip()}
+            if (
+                isinstance(raw, dict)
+                and "bucket" in raw
+                and raw["bucket"] not in {bucket.value for bucket in DreamBucketEnum}
+            ):
+                self.logger.warning(
+                    f"[{self.name}] unit {raw.get('name')!r} emitted an unknown bucket; routing to wiki",
+                )
+                raw = {**raw, "bucket": DreamBucketEnum.WIKI.value}
+            normalized.append(raw)
+        output = DreamExtractOutput.model_validate({"units": normalized})
         allowed = set(state.changed_paths)
-        for raw in meta.get("units") or meta.get("memory_units") or []:
-            if max_units is not None and len(state.units) >= max_units:
-                break
-            if not isinstance(raw, dict):
-                continue
-            name = str(raw.get("name") or "").strip()
-            summary = str(raw.get("summary") or "").strip()
-            raw_bucket = str(raw.get("bucket") or "").strip()
-            paths = clean_paths(raw.get("paths"), allowed)
+        cleaned = []
+        for index, unit in enumerate(output.units):
+            name, summary = unit.name.strip(), unit.summary.strip()
+            paths = clean_paths(unit.paths, allowed)
             if not name or not summary or not paths:
-                continue
-            try:
-                bucket = DreamBucketEnum(raw_bucket).value
-            except ValueError:
-                self.logger.warning(f"[{self.name}] unit {name!r} emitted bucket {raw_bucket!r}; routing to wiki")
-                bucket = DreamBucketEnum.WIKI.value
-            state.units.append({"name": name, "bucket": bucket, "summary": summary, "paths": paths})
+                raise ValueError(f"units[{index}] requires non-empty name, summary, and paths from changed_paths")
+            cleaned.append({"name": name, "bucket": unit.bucket.value, "summary": summary, "paths": paths})
+        state.units.extend(cleaned if max_units is None else cleaned[: max(max_units - len(state.units), 0)])
 
     def _finish(self, state: DreamState, success: bool, answer: str):
         assert self.context is not None

reme/steps/evolve/dream/extract.yaml

@@ -61,12 +61,20 @@ extract_system_prompt: |
 
   ## Output Format
 
-  Return only one YAML or JSON object with this exact shape:
-  units:
-    - name: <short-name>
-      bucket: procedure|personal|wiki
-      summary: <grounded reusable abstraction>
-      paths: [<changed path>, ...]
+  Return only one valid JSON object, without Markdown fences or surrounding prose.
+  The units field is required and must be a list. Every unit must contain name,
+  bucket, summary, and paths. name and summary must be non-empty strings; paths
+  must be a non-empty list of strings from changed_paths.
+  Double-quote all keys and string values, and escape embedded double quotes,
+  backslashes, and newlines according to JSON rules. Code-like text containing
+  `key: value` must stay inside a quoted string; never emit YAML plain scalars.
+  Return {{"units": []}} only when there is no reusable memory to extract.
+
+  JSON schema:
+  {output_schema_json}
+
+  Valid example (replace paths with actual changed_paths):
+  {{"units": [{{"name": "ci-submodules", "bucket": "procedure", "summary": "CI checkout needs `submodules: recursive` to fetch nested dependencies.", "paths": ["daily/2026-10-20/ci.md"]}}]}}
 
 extract_system_prompt_zh: |
   你是 dream 全局抽取 agent。一起阅读所有发生变化的 daily 文件，输出一个由合并记忆 unit 组成的精简跨文件计划。
@@ -118,12 +126,39 @@ extract_system_prompt_zh: |
 
   ## 输出格式
 
-  只返回一个 YAML 或 JSON object，结构必须严格如下：
-  units:
-    - name: <短名称>
-      bucket: procedure|personal|wiki
-      summary: <基于证据的可复用抽象>
-      paths: [<changed path>, ...]
+  只返回一个合法的 JSON object，不要使用 Markdown 代码围栏或添加解释文字。
+  必须包含 units 字段，且它必须是 list。每个 unit 必须包含 name、bucket、summary 和 paths。
+  name 和 summary 必须是非空字符串；paths 必须是非空字符串列表，且只能使用 changed_paths 中的路径。
+  所有键和字符串值都必须使用双引号；内部双引号、反斜杠和换行必须按 JSON 规则转义。
+  含有 `key: value` 的代码文本必须放在带引号的字符串内，禁止输出 YAML 无引号标量。
+  只有在没有可复用记忆时才返回 {{"units": []}}。
+
+  JSON schema：
+  {output_schema_json}
+
+  合法示例（路径须替换为实际 changed_paths）：
+  {{"units": [{{"name": "ci-submodules", "bucket": "procedure", "summary": "CI checkout 需要 `submodules: recursive`，以获取嵌套依赖。", "paths": ["daily/2026-10-20/ci.md"]}}]}}
+
+extract_retry_message: |
+  The previous extraction failed parsing or validation:
+  {error}
+
+  Retry extraction from the changed material above and return a complete corrected
+  JSON object following the system prompt's schema. Include the units list and all
+  required unit fields. Double-quote every key and string value, including any
+  `key: value` text, and escape quotes, backslashes, and newlines. Do not use YAML,
+  Markdown fences, or explanatory prose. A format error does not mean there are no
+  reusable memories; use {{"units": []}} only when the material warrants it.
+
+extract_retry_message_zh: |
+  上次抽取发生解析或校验错误：
+  {error}
+
+  请根据上面的变化内容重新抽取，按系统提示的 schema 返回完整、修正后的 JSON object。
+  必须包含 units 列表和每个 unit 的必填字段。所有键和字符串值都要加双引号，
+  含 `key: value` 的文本也不例外；双引号、反斜杠和换行须转义。
+  不要使用 YAML、Markdown 代码围栏或解释文字。格式错误不代表没有可复用记忆；
+  只有材料确实不值得记忆时才返回 {{"units": []}}。
 
 extract_user_message: |
   date: {date}

reme/steps/evolve/dream/finish.py

@@ -74,8 +74,11 @@ class DreamFinishStep(BaseStep):
 def render_summary(state: DreamState) -> str:
     """Render a concise user-facing summary."""
     dates = ", ".join(state.dates or [state.date])
+    title = "AutoDream failed" if state.errors else "AutoDream completed"
+    if state.warnings and not state.errors:
+        title += " with warnings"
     lines = [
-        ("AutoDream completed with warnings" if state.warnings else "AutoDream completed"),
+        title,
         "",
         f"- Date: {state.date}",
         f"- Scan window: {dates}",

reme/steps/evolve/dream/integrate.py

@@ -61,7 +61,7 @@ class DreamIntegrateStep(BaseStep):
         state = state_from_context(self)
         if not state.units:
             self.logger.info(f"[{self.name}] skip no units")
-            return self._finish(state, True, "No dream units to integrate")
+            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)

reme/steps/evolve/dream/utils.py

@@ -134,19 +134,25 @@ def previous_dates(day: str, n_days: int) -> list[str]:
     return [(base - dt.timedelta(days=i)).isoformat() for i in range(1, max(n_days, 0) + 1)]
 
 
-def parse_structured_reply(text: str) -> dict:
-    """Parse a JSON/YAML object from an agent reply, including fenced blocks."""
+def parse_structured_reply(text: str, *, strict: bool = False) -> dict:
+    """Parse a JSON/YAML object; strict mode rejects lossy fallbacks with diagnostics."""
     candidates = [text.strip()]
     candidates.extend(m.group(1).strip() for m in re.finditer(r"```(?:json|ya?ml)?\s*(.*?)```", text, re.S | re.I))
+    error = "expected a JSON/YAML object"
     for raw in candidates:
         if not raw:
             continue
         try:
             data = yaml.safe_load(raw)
-        except yaml.YAMLError:
+        except yaml.YAMLError as exc:
+            error = f"{type(exc).__name__}: {exc}"
+            if strict:
+                continue
             data = _parse_scalar_mapping(raw)
-        if isinstance(data, dict) and data:
+        if isinstance(data, dict) and (data or strict):
             return data
+    if strict:
+        raise ValueError(f"Invalid structured reply: {error}")
     return {}
 
 

tests/unit/test_auto_dream.py

@@ -30,6 +30,12 @@ from reme.steps.evolve.dream.integrate import DreamIntegrateStep, _snapshot_dige
 from reme.steps.evolve.dream.utils import parse_structured_reply, recent_dates, scan_day_files
 from reme.steps.file_io.frontmatter_update import FrontmatterUpdateStep
 
+_MALFORMED_EXTRACT_REPLY = (
+    "units:\n- name: checkout\n  bucket: procedure\n"
+    "  summary: CI must add `with: submodules: recursive`\n"
+    "  paths: [daily/2026-05-28/session.md]\n"
+)
+
 
 def _touch(path: Path, text: str = "x") -> Path:
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -116,9 +122,11 @@ class _SequenceAgent(BaseAgentWrapper):
         super().__init__()
         self.outcomes = list(outcomes)
         self.calls = 0
+        self.requests = []
 
     async def reply(self, _message, **_kwargs):
         self.calls += 1
+        self.requests.append((_message, _kwargs))
         outcome = self.outcomes.pop(0)
         if callable(outcome):
             outcome = outcome()
@@ -245,6 +253,25 @@ def test_parse_structured_reply_handles_fenced_yaml_and_scalar_fallback():
     assert data["note"].startswith("Extended node")
 
 
+def test_extract_parser_preserves_yaml_diagnostics_and_accepts_quoted_colons():
+    """Strict extraction rejects the reported YAML while legacy parsing stays compatible."""
+    assert parse_structured_reply(_MALFORMED_EXTRACT_REPLY) == {}
+    with pytest.raises(ValueError, match="ScannerError") as exc:
+        parse_structured_reply(_MALFORMED_EXTRACT_REPLY, strict=True)
+    assert "line 4" in str(exc.value)
+    assert "submodules: recursive" in str(exc.value)
+    quoted = _MALFORMED_EXTRACT_REPLY.replace(
+        "summary: CI must add `with: submodules: recursive`",
+        'summary: "CI must add `with: submodules: recursive`"',
+    )
+    assert len(parse_structured_reply(quoted, strict=True)["units"]) == 1
+    assert len(parse_structured_reply(f"```yaml\n{quoted}```", strict=True)["units"]) == 1
+    assert parse_structured_reply("{}", strict=True) == {}
+    for reply in ("", "not an object", "[]"):
+        with pytest.raises(ValueError, match="expected a JSON/YAML object"):
+            parse_structured_reply(reply, strict=True)
+
+
 def test_integrate_prompts_require_wikilinks_in_contextual_prose():
     """Every digest bucket rejects standalone relation fields and bare links."""
     prompt_path = Path(__file__).parents[2] / "reme" / "steps" / "evolve" / "dream" / "integrate.yaml"
@@ -277,37 +304,171 @@ def test_extract_clean_output_respects_max_units():
     assert [unit["name"] for unit in state.units] == [f"unit-{i}" for i in range(5)]
 
 
-def test_extract_unusable_receipt_is_a_warning_not_a_failure(tmp_path):
-    """Malformed best-effort extraction output does not fail the nightly job."""
+@pytest.mark.parametrize("field", ["units", "memory_units"])
+@pytest.mark.parametrize(
+    ("raw_bucket", "expected_bucket"),
+    [
+        (" procedure ", "procedure"),
+        ("\tpersonal\n", "personal"),
+        (" wiki ", "wiki"),
+        (" unknown ", "wiki"),
+        (" \t\n ", "wiki"),
+    ],
+)
+def test_extract_validation_strips_bucket_whitespace_before_fallback(field, raw_bucket, expected_bucket):
+    """Padded known buckets retain their destination; unknown and blank values still use wiki."""
+    state = DreamState(changed_paths=["daily/a.md"])
+    unit = {"name": "unit", "bucket": raw_bucket, "summary": "memory", "paths": ["daily/a.md"]}
 
-    async def run():
-        _touch(tmp_path / "daily" / "2026-05-28" / "session.md")
-        step = DreamExtractStep(scan_days=1, app_context=ApplicationContext(workspace_dir=str(tmp_path)))
-        step.agent_wrapper = _ReplyAgent()
+    DreamExtractStep().clean_output(state, {field: [unit]})
 
-        with (
-            patch("reme.steps.evolve.dream.extract.refresh_day_index", return_value={}),
-            patch("reme.steps.evolve.dream.extract.llm_available", return_value=True),
-        ):
-            response = await step(
-                RuntimeContext(
-                    date="2026-05-28",
-                    agent_wrapper=step.agent_wrapper,
-                    file_catalog=_Catalog(),
-                    file_store=_FileStore(tmp_path),
-                ),
-            )
+    assert state.units == [{**unit, "bucket": expected_bucket}]
 
+
+def test_extract_validation_preserves_legacy_alias_and_unknown_bucket_fallback():
+    """Validation retains the documented bucket fallback and source-path cleanup."""
+    state = DreamState(changed_paths=["daily/a.md"])
+    DreamExtractStep().clean_output(
+        state,
+        {
+            "memory_units": [
+                {
+                    "name": " unit ",
+                    "bucket": "unknown",
+                    "summary": " memory ",
+                    "paths": ["daily/a.md", "daily/a.md", "outside.md"],
+                },
+            ],
+        },
+    )
+    assert state.units == [{"name": "unit", "bucket": "wiki", "summary": "memory", "paths": ["daily/a.md"]}]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("job_name", ["auto_dream", "dream_cron"])
+@pytest.mark.parametrize(
+    ("reply", "succeeds"),
+    [
+        (_MALFORMED_EXTRACT_REPLY, False),
+        ("{}", False),
+        ('{"units": null}', False),
+        ('{"units": {}}', False),
+        ('{"units": [42]}', False),
+        ('{"units": [{"name": "missing-fields"}]}', False),
+        (
+            '{"units": [{"name": "x", "summary": " ", "bucket": "wiki", "paths": ["daily/2026-05-28/session.md"]}]}',
+            False,
+        ),
+        ('{"units": [{"name": "x", "summary": "x", "bucket": "wiki", "paths": ["outside.md"]}]}', False),
+        (
+            '{"units": [{"name": "valid", "summary": "x", "bucket": "wiki", '
+            '"paths": ["daily/2026-05-28/session.md"]}, '
+            '{"name": "invalid", "summary": "x", "bucket": "wiki", "paths": []}]}',
+            False,
+        ),
+        ('{"units": []}', True),
+        ("units: []", True),
+        ("memory_units: []", True),
+    ],
+)
+async def test_dream_extraction_failure_preserves_changed_inputs(tmp_path, monkeypatch, job_name, reply, succeeds):
+    """Invalid receipts fail both configured pipelines without checkpointing notes or day indexes."""
+
+    class Catalog(_Catalog):
+        """Expose versions checkpointed by Finish to the next job invocation."""
+
+        async def get_nodes(self, paths=None):
+            return self.upserts
+
+    note = _touch(tmp_path / "daily/2026-05-28/session.md", "CI needs submodules: recursive")
+    _touch(tmp_path / "daily/2026-05-28.md", "day index")
+    app_context = ApplicationContext(workspace_dir=str(tmp_path))
+    catalog, agent = Catalog(), _ReplyAgent({"result": reply})
+    config = resolve_app_config(config="default", log_config=False)["jobs"][job_name]
+    job = BaseJob(name=job_name, steps=config["steps"], app_context=app_context)
+    monkeypatch.setattr("reme.steps.evolve.dream.extract.refresh_day_index", AsyncMock(return_value={}))
+    monkeypatch.setattr("reme.steps.evolve.dream.extract.llm_available", lambda _step: True)
+    await job.start()
+    try:
+        inputs = {
+            "date": "2026-05-28",
+            "scan_days": 1,
+            "agent_wrapper": agent,
+            "file_catalog": catalog,
+            "file_store": _FileStore(tmp_path),
+        }
+        response = await job(**inputs)
         dream = response.metadata["dream"]
-        assert response.success is True
+        assert response.success is succeeds
+        assert response.answer.startswith("AutoDream completed" if succeeds else "AutoDream failed")
         assert dream["units"] == []
-        assert dream["failed_paths"] == []
-        assert dream["warnings"] == [
-            "dream extract skipped unusable agent receipt after retry; expected a units list",
-        ]
-        assert step.agent_wrapper.calls == 2
+        assert dream["files_changed"] == 2
+        assert dream["failed_paths"] == ([] if succeeds else dream["changed_paths"])
+        assert bool(dream["errors"]) is not succeeds
+        assert set(dream["checkpoint_paths"]) == (set(dream["changed_paths"]) if succeeds else set())
+        assert catalog.dumps == int(succeeds)
+        assert agent.calls == (1 if succeeds else 2)
+        assert note.read_text(encoding="utf-8") == "CI needs submodules: recursive"
 
-    asyncio.run(run())
+        agent.result = {"result": '{"units": []}'}
+        response = await job(**inputs)
+        assert response.success is True
+        assert response.metadata["dream"]["files_changed"] == (0 if succeeds else 2)
+        assert agent.calls == (1 if succeeds else 3)
+        assert catalog.dumps == (2 if succeeds else 1)
+    finally:
+        await job.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("language", ["en", "zh"])
+@pytest.mark.parametrize(
+    ("first_reply", "error"),
+    [(_MALFORMED_EXTRACT_REPLY, "ScannerError"), ('{"units": [{"name": "missing-fields"}]}', "ValidationError")],
+)
+async def test_extract_retry_includes_parser_error_and_recovers_json(tmp_path, language, first_reply, error):
+    """Both localized prompts recover code-like strings using JSON and concrete error feedback."""
+    source = "daily/2026-05-28/session.md"
+    _touch(tmp_path / source)
+    unit = {
+        "name": "checkout",
+        "bucket": "procedure",
+        "summary": 'CI needs `submodules: recursive`, "quoted" text, and \\paths.\nKeep dependencies.',
+        "paths": [source],
+    }
+    agent = _SequenceAgent({"result": first_reply}, {"result": json.dumps({"units": [unit]})})
+    step = DreamExtractStep(
+        scan_days=1,
+        language=language,
+        app_context=ApplicationContext(workspace_dir=str(tmp_path)),
+        agent_wrapper=agent,
+    )
+    with (
+        patch("reme.steps.evolve.dream.extract.refresh_day_index", return_value={}),
+        patch("reme.steps.evolve.dream.extract.llm_available", return_value=True),
+    ):
+        response = await step(
+            RuntimeContext(date="2026-05-28", file_catalog=_Catalog(), file_store=_FileStore(tmp_path)),
+        )
+
+    assert response.success is True
+    assert response.metadata["dream"]["units"] == [unit]
+    assert not response.metadata["dream"]["errors"]
+    first_message, first_kwargs = agent.requests[0]
+    retry_message, retry_kwargs = agent.requests[1]
+    assert retry_message.startswith(first_message)
+    assert error in retry_message
+    if error == "ScannerError":
+        assert "line 4" in retry_message
+    else:
+        assert "Field required" in retry_message
+    assert "key: value" in retry_message
+    assert retry_kwargs["system_prompt"] == first_kwargs["system_prompt"]
+    system_prompt = first_kwargs["system_prompt"]
+    assert "JSON schema" in system_prompt
+    assert "YAML or JSON" not in system_prompt
+    example = json.loads(system_prompt.splitlines()[-1])
+    assert "submodules: recursive" in example["units"][0]["summary"]
 
 
 def test_extract_retries_one_unusable_receipt(tmp_path):
@@ -661,7 +822,7 @@ async def test_dream_jobs_tag_outputs_and_preserve_checkpoint_results(
         await job.close()
 
     assert response.success is not integrate_fails
-    assert response.answer.startswith("AutoDream completed")
+    assert response.answer.startswith("AutoDream failed" if integrate_fails else "AutoDream completed")
     assert response.metadata["modified"] is True
     dream, tagging = response.metadata["dream"], response.metadata["auto_tag"]
     assert (source in dream["checkpoint_paths"]) is not integrate_fails