Merge 0ccb364bd6 into 2767836b21

reme/config/default.yaml

@@ -77,6 +77,7 @@ jobs:
         file_catalog: dream
         scan_days: 2
         max_units: 5
+        max_chars_per_batch: 40000
       - backend: dream_integrate_step
       - backend: dream_finish_step
         file_catalog: dream
@@ -162,11 +163,16 @@ jobs:
           type: integer
           description: "maximum number of extracted memory units"
           default: 5
+        max_chars_per_batch:
+          type: integer
+          description: "maximum characters of changed material sent per extraction call"
+          default: 40000
     steps:
       - backend: dream_extract_step
         file_catalog: dream
         scan_days: 2
         max_units: 5
+        max_chars_per_batch: 40000
       - backend: dream_integrate_step
       - backend: dream_finish_step
         file_catalog: dream

reme/steps/evolve/dream/extract.py

@@ -1,6 +1,7 @@
 """Global dream extract step."""
 
 import json
+from pathlib import Path
 
 from ...base_step import BaseStep
 from ...file_io import refresh_day_index
@@ -9,6 +10,7 @@ from ....components import R
 from ....enumeration import DreamBucketEnum
 from ....schema import DreamExtractOutput, DreamState
 from .utils import (
+    batch_paths,
     clean_paths,
     daily_dir,
     llm_available,
@@ -24,14 +26,55 @@ from .utils import (
 _TOOLS = ("read",)
 
 
+class _UnusableReceipt(RuntimeError):
+    """The model returned an unusable receipt even after one bounded retry.
+
+    Bisecting the batch cannot help here: the input was delivered fine, the
+    model just failed to answer in the expected shape. Callers mark the batch
+    as failed instead of splitting it.
+    """
+
+
+def merge_units(units: list[dict]) -> list[dict]:
+    """Merge same-name units produced by different batches.
+
+    Batched extraction can surface the same abstraction twice; the integration
+    step writes one digest node per unit, so duplicates would double-write.
+    Paths are unioned and the richer summary wins.
+    """
+    merged: dict[str, dict] = {}
+    order: list[str] = []
+    for unit in units:
+        name = str(unit.get("name") or "").strip()
+        key = name.lower()
+        if not key:
+            continue
+        if key not in merged:
+            merged[key] = dict(unit)
+            order.append(key)
+            continue
+        current = merged[key]
+        current["paths"] = sorted({*current.get("paths", []), *unit.get("paths", [])})
+        if len(str(unit.get("summary") or "")) > len(str(current.get("summary") or "")):
+            current["summary"] = unit.get("summary")
+    return [merged[key] for key in order]
+
+
 @R.register("dream_extract_step")
 class DreamExtractStep(BaseStep):
     """Scan changed daily files and globally extract merged memory units."""
 
-    def __init__(self, scan_days: int = 2, max_units: int = 5, **kwargs):
+    def __init__(
+        self,
+        scan_days: int = 2,
+        max_units: int = 5,
+        max_chars_per_batch: int = 40000,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         self.scan_days = scan_days
         self.max_units = max_units
+        self.max_chars_per_batch = max_chars_per_batch
 
     async def execute(self):
         assert self.context is not None
@@ -113,57 +156,163 @@ 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)}")
+        output_schema = DreamExtractOutput.model_json_schema()
+        output_schema["required"] = ["units"]
+        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),
+        )
+        raw_max_chars = self.context.get("max_chars_per_batch", self.max_chars_per_batch)
+        max_chars = max(int(raw_max_chars or self.max_chars_per_batch), 0)
+        batches = batch_paths(workspace, changed, max_chars=max_chars)
+        self.logger.info(
+            f"[{self.name}] batching changed={len(changed)} batches={len(batches)} " f"max_chars_per_batch={max_chars}",
+        )
+        for batch in batches:
+            remaining = max_units - len(state.units)
+            if remaining <= 0:
+                self.logger.info(
+                    f"[{self.name}] stop batching max_units reached units={len(state.units)}",
+                )
+                break
+            await self._extract_batch(
+                state,
+                batch,
+                workspace=workspace,
+                day=day,
+                dates=dates,
+                hint=hint,
+                max_units=remaining,
+                system_prompt=system_prompt,
+            )
+        state.units = merge_units(state.units)
+
+        self.logger.info(
+            f"[{self.name}] parse done units={len(state.units)} batches={len(batches)} "
+            f"failed_paths={len(state.failed_paths)} errors={len(state.errors)}",
+        )
+        answer = f"Extracted {len(state.units)} unit(s) from {len(changed)} changed file(s) across {len(dates)} day(s)"
+        return self._finish(state, not state.errors, answer)
+
+    async def _extract_batch(
+        self,
+        state: DreamState,
+        batch: list[str],
+        *,
+        workspace: Path,
+        day: str,
+        dates: list[str],
+        hint: str,
+        max_units: int,
+        system_prompt: str,
+    ) -> None:
+        """Extract one size-capped batch, bisecting a batch that cannot be handled.
+
+        A transport failure (oversized prompt, backend crash) is never retried
+        with a *larger* message: the batch is split in half instead, so the
+        offending input is isolated. A single file that still fails is recorded
+        in ``failed_paths`` and the remaining batches keep going.
+        """
+        try:
+            await self._extract_once(
+                state,
+                batch,
+                workspace=workspace,
+                day=day,
+                dates=dates,
+                hint=hint,
+                max_units=max_units,
+                system_prompt=system_prompt,
+            )
+        except _UnusableReceipt as e:
+            error = f"dream extract agent failed after retry: {e}"
+            state.errors.append(error)
+            state.failed_paths.extend(batch)
+            self.logger.error(f"[{self.name}] {error}")
+        except Exception as e:  # noqa: BLE001
+            if len(batch) > 1:
+                middle = len(batch) // 2
+                self.logger.warning(
+                    f"[{self.name}] batch failed files={len(batch)}; splitting to isolate the "
+                    f"offending input: {type(e).__name__}: {e}",
+                )
+                for half in (batch[:middle], batch[middle:]):
+                    await self._extract_batch(
+                        state,
+                        half,
+                        workspace=workspace,
+                        day=day,
+                        dates=dates,
+                        hint=hint,
+                        max_units=max(max_units - len(state.units), 0),
+                        system_prompt=system_prompt,
+                    )
+                return
+            target = batch[0] if batch else "(empty batch)"
+            error = f"dream extract agent failed for {target}: {type(e).__name__}: {e}"
+            state.errors.append(error)
+            state.failed_paths.extend(batch)
+            self.logger.error(f"[{self.name}] {error}")
+
+    async def _extract_once(
+        self,
+        state: DreamState,
+        batch: list[str],
+        *,
+        workspace: Path,
+        day: str,
+        dates: list[str],
+        hint: str,
+        max_units: int,
+        system_prompt: str,
+    ) -> None:
+        """Run one LLM extraction for a single batch and merge its units."""
         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),
+            changed_paths_json=json.dumps(batch, ensure_ascii=False, indent=2),
+            material_blob=pack_paths(workspace, batch),
         )
-        retry_feedback = ""
-        output_schema = DreamExtractOutput.model_json_schema()
-        output_schema["required"] = ["units"]
-        for attempt in range(2):
-            try:
-                result = await self.agent_wrapper.reply(
-                    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)
-                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}",
-                    )
-                    continue
-                error = f"dream extract agent failed after retry: {type(e).__name__}: {e}"
-                state.errors.append(error)
-                state.failed_paths = list(changed)
-                self.logger.error(f"[{self.name}] {error}")
-                return self._finish(state, False, error)
-
-        self.logger.info(f"[{self.name}] parse done keys={','.join(sorted(meta.keys())) if meta else '(none)'}")
-        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)
+        allowed = set(batch)
+        result = await self.agent_wrapper.reply(
+            user_message,
+            system_prompt=system_prompt,
+            job_tools=list(_TOOLS),
+        )
+        self.logger.info(f"[{self.name}] agent done files={len(batch)} has_result={bool(result.get('result'))}")
+        raw_result = agent_reply_result_text(result)
+        state.extract_summary = "\n".join(filter(None, (state.extract_summary, raw_result)))
+        try:
+            meta = parse_structured_reply(raw_result, strict=True)
+            self.clean_output(state, meta, max_units=max_units, allowed=allowed)
+            return
+        except Exception as e:  # noqa: BLE001
+            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}",
+            )
+        result = await self.agent_wrapper.reply(
+            user_message + retry_feedback,
+            system_prompt=system_prompt,
+            job_tools=list(_TOOLS),
+        )
+        raw_result = agent_reply_result_text(result)
+        state.extract_summary = "\n".join(filter(None, (state.extract_summary, raw_result)))
+        try:
+            meta = parse_structured_reply(raw_result, strict=True)
+            self.clean_output(state, meta, max_units=max_units, allowed=allowed)
+        except Exception as e:  # noqa: BLE001
+            raise _UnusableReceipt(f"{type(e).__name__}: {e}") from e
 
     def _existing(self, workspace, files: list[str]) -> dict[str, float]:
         out: dict[str, float] = {}
@@ -174,8 +323,19 @@ class DreamExtractStep(BaseStep):
                 self.logger.error(f"[{self.name}] stat failed on {rel}: {e}")
         return out
 
-    def clean_output(self, state: DreamState, meta: dict, max_units: int | None = None) -> None:
-        """Validate the whole receipt before accepting any units or applying the cap."""
+    def clean_output(
+        self,
+        state: DreamState,
+        meta: dict,
+        max_units: int | None = None,
+        allowed: set[str] | None = None,
+    ) -> None:
+        """Validate the whole receipt before accepting any units or applying the cap.
+
+        ``allowed`` scopes the accepted paths to the current batch; by default
+        every changed path stays valid, which keeps single-batch behaviour
+        identical to the unbatched implementation.
+        """
         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)")
@@ -194,11 +354,11 @@ class DreamExtractStep(BaseStep):
                 raw = {**raw, "bucket": DreamBucketEnum.WIKI.value}
             normalized.append(raw)
         output = DreamExtractOutput.model_validate({"units": normalized})
-        allowed = set(state.changed_paths)
+        allowed_paths = set(state.changed_paths) if allowed is None else set(allowed)
         cleaned = []
         for index, unit in enumerate(output.units):
             name, summary = unit.name.strip(), unit.summary.strip()
-            paths = clean_paths(unit.paths, allowed)
+            paths = clean_paths(unit.paths, allowed_paths)
             if not name or not summary or not 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})

reme/steps/evolve/dream/utils.py

@@ -113,6 +113,72 @@ def pack_paths(
     return "\n".join(blocks)
 
 
+def _file_size(workspace: Path, rel: str) -> int:
+    """Best-effort byte size of one workspace-relative path."""
+    try:
+        return (workspace / rel).stat().st_size
+    except OSError:
+        return 0
+
+
+def _day_key(rel: str) -> str:
+    """Group key that keeps one scanned day together (``daily/<date>``)."""
+    parts = rel.split("/")
+    return "/".join(parts[:2]) if len(parts) >= 2 else rel
+
+
+def batch_paths(
+    workspace: Path,
+    paths: list[str],
+    *,
+    max_chars: int,
+) -> list[list[str]]:
+    """Split paths into size-capped batches without dropping any path.
+
+    Files belonging to the same scanned day are kept in one batch whenever the
+    day fits the budget, so cross-file extraction still sees a whole day. A day
+    larger than the budget is split by file size, and a single oversized file
+    keeps a batch of its own (``pack_paths`` still applies its per-file cap).
+
+    Every path lands in exactly one batch; callers must never mark an unseen
+    path as extracted, which is why this returns a partition and not a slice.
+    """
+    if not paths:
+        return []
+    if max_chars <= 0:
+        return [list(paths)]
+
+    groups: list[list[str]] = []
+    for rel in paths:
+        if groups and _day_key(groups[-1][0]) == _day_key(rel):
+            groups[-1].append(rel)
+        else:
+            groups.append([rel])
+
+    batches: list[list[str]] = []
+    current: list[str] = []
+    total = 0
+    for group in groups:
+        group_size = sum(_file_size(workspace, rel) for rel in group)
+        if current and total + group_size > max_chars:
+            batches.append(current)
+            current, total = [], 0
+        if group_size <= max_chars:
+            current.extend(group)
+            total += group_size
+            continue
+        for rel in group:
+            size = _file_size(workspace, rel)
+            if current and total + size > max_chars:
+                batches.append(current)
+                current, total = [], 0
+            current.append(rel)
+            total += size
+    if current:
+        batches.append(current)
+    return batches
+
+
 def clean_paths(raw_paths, allowed: set[str]) -> list[str]:
     """Clean paths."""
     if not isinstance(raw_paths, list):

tests/unit/test_auto_dream.py

@@ -27,7 +27,12 @@ from reme.steps.evolve.auto_tag import AutoTagStep
 from reme.steps.evolve.dream.extract import DreamExtractStep
 from reme.steps.evolve.dream.finish import DreamFinishStep
 from reme.steps.evolve.dream.integrate import DreamIntegrateStep, _snapshot_digest
-from reme.steps.evolve.dream.utils import parse_structured_reply, recent_dates, scan_day_files
+from reme.steps.evolve.dream.utils import (  # noqa: E501
+    batch_paths,
+    parse_structured_reply,
+    recent_dates,
+    scan_day_files,
+)
 from reme.steps.file_io.frontmatter_update import FrontmatterUpdateStep
 
 _MALFORMED_EXTRACT_REPLY = (
@@ -1032,3 +1037,216 @@ async def test_dream_initializes_started_provider_before_availability_check(tmp_
     finally:
         await agent.close()
         await llm.close()
+
+
+# ---------------------------------------------------------------------------
+# Dream extract batching (#659): oversized material must never be sent in one
+# call, and nothing may be silently dropped.
+# ---------------------------------------------------------------------------
+
+
+def test_batch_paths_partitions_without_dropping(tmp_path):
+    """Every path lands in exactly one batch and same-day files stay together."""
+    day1 = [f"daily/2026-05-28/f{i}.md" for i in range(3)]
+    day2 = [f"daily/2026-05-27/g{i}.md" for i in range(2)]
+    for rel in [*day1, *day2]:
+        _touch(tmp_path / rel, "x" * 1000)
+
+    batches = batch_paths(tmp_path, [*day1, *day2], max_chars=4000)
+
+    assert [sorted(b) for b in batches] == [sorted(day1), sorted(day2)]
+    flat = [rel for batch in batches for rel in batch]
+    assert sorted(flat) == sorted([*day1, *day2])
+    assert len(flat) == len(set(flat))
+
+
+def test_batch_paths_keeps_oversize_file_alone(tmp_path):
+    """A single file larger than the budget gets its own batch, never dropped."""
+    _touch(tmp_path / "daily/2026-05-28/huge.md", "x" * 9000)
+    _touch(tmp_path / "daily/2026-05-28/small.md", "x" * 100)
+
+    batches = batch_paths(
+        tmp_path,
+        ["daily/2026-05-28/huge.md", "daily/2026-05-28/small.md"],
+        max_chars=1000,
+    )
+
+    assert ["daily/2026-05-28/huge.md"] in [list(b) for b in batches]
+    flat = [rel for batch in batches for rel in batch]
+    assert sorted(flat) == ["daily/2026-05-28/huge.md", "daily/2026-05-28/small.md"]
+
+
+async def _run_dream_extract(tmp_path, agent, **kwargs):
+    step = DreamExtractStep(
+        app_context=ApplicationContext(workspace_dir=str(tmp_path)),
+        agent_wrapper=agent,
+        **kwargs,
+    )
+    with (
+        patch("reme.steps.evolve.dream.extract.refresh_day_index", return_value={}),
+        patch("reme.steps.evolve.dream.extract.llm_available", return_value=True),
+    ):
+        return await step(
+            RuntimeContext(
+                date="2026-05-28",
+                file_catalog=_Catalog(),
+                file_store=_FileStore(tmp_path),
+            ),
+        )
+
+
+@pytest.mark.asyncio
+async def test_extract_splits_oversized_material_across_calls(tmp_path):
+    """8 changed files must not be packed into a single extraction request."""
+    day1 = [f"daily/2026-05-28/f{i}.md" for i in range(4)]
+    day2 = [f"daily/2026-05-27/g{i}.md" for i in range(4)]
+    for rel in [*day1, *day2]:
+        _touch(tmp_path / rel, "x" * 1000)
+
+    agent = _SequenceAgent({"result": '{"units": []}'}, {"result": '{"units": []}'})
+    response = await _run_dream_extract(
+        tmp_path,
+        agent,
+        scan_days=2,
+        max_units=5,
+        max_chars_per_batch=4000,
+    )
+
+    assert response.success is True
+    assert agent.calls == 2
+    seen = [rel for message, _kwargs in agent.requests for rel in [*day1, *day2] if rel in message]
+    assert sorted(seen) == sorted([*day1, *day2])
+    assert not response.metadata["dream"]["failed_paths"]
+
+
+@pytest.mark.asyncio
+async def test_extract_bisects_a_failing_batch(tmp_path):
+    """A transport failure splits the batch instead of retrying a larger prompt."""
+    day1 = [f"daily/2026-05-28/f{i}.md" for i in range(2)]
+    for rel in day1:
+        _touch(tmp_path / rel, "x" * 1000)
+
+    agent = _SequenceAgent(
+        RuntimeError("prompt too large"),
+        {"result": '{"units": []}'},
+        {"result": '{"units": []}'},
+    )
+    response = await _run_dream_extract(
+        tmp_path,
+        agent,
+        scan_days=2,
+        max_units=5,
+        max_chars_per_batch=9000,
+    )
+
+    assert agent.calls == 3
+    assert response.success is True
+    assert not response.metadata["dream"]["failed_paths"]
+
+
+@pytest.mark.asyncio
+async def test_extract_isolates_one_offending_file(tmp_path):
+    """A file that always fails only blocks itself, not the whole run."""
+    day1 = [f"daily/2026-05-28/f{i}.md" for i in range(2)]
+    for rel in day1:
+        _touch(tmp_path / rel, "x" * 1000)
+
+    good = {
+        "result": json.dumps(
+            {
+                "units": [
+                    {
+                        "name": "kept",
+                        "bucket": "wiki",
+                        "summary": "kept unit",
+                        "paths": ["daily/2026-05-28/f0.md"],
+                    },
+                ],
+            },
+        ),
+    }
+    agent = _SequenceAgent(
+        RuntimeError("prompt too large"),
+        good,
+        RuntimeError("engine crashed"),
+    )
+    response = await _run_dream_extract(
+        tmp_path,
+        agent,
+        scan_days=2,
+        max_units=5,
+        max_chars_per_batch=9000,
+    )
+
+    dream = response.metadata["dream"]
+    assert [u["name"] for u in dream["units"]] == ["kept"]
+    assert dream["failed_paths"] == ["daily/2026-05-28/f1.md"]
+    assert response.success is False
+
+
+@pytest.mark.asyncio
+async def test_extract_merges_same_name_units_across_batches(tmp_path):
+    """The same abstraction surfaced by two batches becomes one unit."""
+    _touch(tmp_path / "daily/2026-05-28/a.md", "x" * 1000)
+    _touch(tmp_path / "daily/2026-05-27/b.md", "x" * 1000)
+
+    first = {
+        "result": json.dumps(
+            {
+                "units": [
+                    {
+                        "name": "ci-submodules",
+                        "bucket": "procedure",
+                        "summary": "short",
+                        "paths": ["daily/2026-05-27/b.md"],
+                    },
+                ],
+            },
+        ),
+    }
+    second = {
+        "result": json.dumps(
+            {
+                "units": [
+                    {
+                        "name": "CI-Submodules",
+                        "bucket": "procedure",
+                        "summary": "a much longer summary",
+                        "paths": ["daily/2026-05-28/a.md"],
+                    },
+                ],
+            },
+        ),
+    }
+    agent = _SequenceAgent(first, second)
+    response = await _run_dream_extract(
+        tmp_path,
+        agent,
+        scan_days=2,
+        max_units=5,
+        max_chars_per_batch=1200,
+    )
+
+    units = response.metadata["dream"]["units"]
+    assert agent.calls == 2
+    assert len(units) == 1
+    assert units[0]["paths"] == ["daily/2026-05-27/b.md", "daily/2026-05-28/a.md"]
+    assert units[0]["summary"] == "a much longer summary"
+
+
+@pytest.mark.asyncio
+async def test_extract_single_small_input_stays_one_call(tmp_path):
+    """Batching must not change behaviour for inputs that already fit."""
+    _touch(tmp_path / "daily/2026-05-28/only.md", "x" * 100)
+    agent = _SequenceAgent({"result": '{"units": []}'})
+
+    response = await _run_dream_extract(
+        tmp_path,
+        agent,
+        scan_days=1,
+        max_units=5,
+        max_chars_per_batch=40000,
+    )
+
+    assert response.success is True
+    assert agent.calls == 1