diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 2aa8f2da..32233ce4 100644 --- a/reme/config/default.yaml +++ b/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: "estimated bytes of changed file content per extraction call, taken from file size; it excludes prompt overhead and a single oversized file still gets its own call, so it is a batching budget rather than a hard request limit" + 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 diff --git a/reme/schema/dream.py b/reme/schema/dream.py index 105aa957..a7f211f6 100644 --- a/reme/schema/dream.py +++ b/reme/schema/dream.py @@ -63,7 +63,19 @@ class DreamState(BaseModel): ) failed_units: list[dict] = Field(default_factory=list) failed_paths: list[str] = Field(default_factory=list) + processed_paths: list[str] | None = Field( + default=None, + description=( + "Paths actually sent to the model. ``None`` means the run did not " + "record it (state written by an older version), and the finish step " + "then falls back to checkpointing every changed path." + ), + ) checkpoint_paths: list[str] = Field(default_factory=list) + deferred_paths: list[str] = Field( + default_factory=list, + description="Changed paths never sent to the model, left eligible for a later run.", + ) warnings: list[str] = Field(default_factory=list) errors: list[str] = Field(default_factory=list) summary: str = "" diff --git a/reme/steps/evolve/dream/extract.py b/reme/steps/evolve/dream/extract.py index d3d4c3b7..ba4fd4c1 100644 --- a/reme/steps/evolve/dream/extract.py +++ b/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,62 @@ 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.""" + """Scan changed daily files and globally extract merged memory units. - def __init__(self, scan_days: int = 2, max_units: int = 5, **kwargs): + ``max_chars_per_batch`` is a batching budget, not a hard request limit: the + total is estimated from file sizes in bytes, it excludes prompt overhead, + and a single file larger than the budget still gets a call of its own. It + lowers the chance that one call carries a whole scan window; it does not + guarantee that any billable request stays under a given size. + """ + + 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 @@ -99,6 +149,7 @@ class DreamExtractStep(BaseStep): files_deleted=len(deleted), changed_paths=changed, unchanged_paths=unchanged, + processed_paths=[], deleted_paths=deleted, existing=existing, indexed=indexed, @@ -113,57 +164,189 @@ 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 index, batch in enumerate(batches): + if len(state.units) >= max_units: + self.logger.info( + f"[{self.name}] stop batching max_units reached units={len(state.units)} " + f"deferred_batches={len(batches) - index}", + ) + break + await self._extract_batch( + state, + batch, + workspace=workspace, + day=day, + dates=dates, + hint=hint, + max_units=max_units, + 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_units, + 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. + + ``max_units`` is the global cap for the whole run, never a per-batch + allowance: ``clean_output`` is the only place that subtracts the units + already accepted, so the budget cannot be charged twice. + """ + if len(state.units) >= max_units: + self.logger.info( + f"[{self.name}] skip batch budget exhausted files={len(batch)}", + ) + return 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) + 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) + self._mark_processed(state, batch) + 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 + self._mark_processed(state, batch) - 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) + @staticmethod + def _mark_processed(state: DreamState, batch: list[str]) -> None: + """Record paths whose content the model actually accepted. + + Marking happens only after a successful reply, so a batch that fails or + is skipped because the unit budget is gone stays eligible for a + later run instead of being checkpointed unseen. + """ + if state.processed_paths is None: + state.processed_paths = [] + for path in batch: + if path not in state.processed_paths: + state.processed_paths.append(path) def _existing(self, workspace, files: list[str]) -> dict[str, float]: out: dict[str, float] = {} @@ -174,8 +357,23 @@ 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. + + ``max_units`` is the global cap for the run. The remaining allowance is + derived here exactly once, from the units already accepted, because the + callers pass the global cap rather than a pre-decremented slice. + """ 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 +392,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}) diff --git a/reme/steps/evolve/dream/finish.py b/reme/steps/evolve/dream/finish.py index 980f93be..b4d9d573 100644 --- a/reme/steps/evolve/dream/finish.py +++ b/reme/steps/evolve/dream/finish.py @@ -24,13 +24,32 @@ class DreamFinishStep(BaseStep): raise RuntimeError("dream_finish_step requires file_catalog") failed_paths = set(state.failed_paths) - checkpoint = [p for p in state.changed_paths if p not in failed_paths] + processed_paths = None if state.processed_paths is None else set(state.processed_paths) + if processed_paths is None: + # State written before processed paths were tracked: keep the old + # behaviour rather than leaving everything uncheckpointed. + checkpoint = [p for p in state.changed_paths if p not in failed_paths] + state.deferred_paths = [] + else: + checkpoint = [p for p in state.changed_paths if p in processed_paths and p not in failed_paths] + state.deferred_paths = [p for p in state.changed_paths if p not in processed_paths] day_index_paths = [f"{state.daily_dir}/{day}.md" for day in (state.dates or [state.date]) if day] - supplemental = [p for p in day_index_paths if p and p not in failed_paths] + deferred = set(state.deferred_paths) + supplemental = [] + for index_path in day_index_paths: + if index_path in failed_paths or index_path in deferred: + continue + day_prefix = f"{state.daily_dir}/{Path(index_path).stem}/" + if any(path.startswith(day_prefix) for path in deferred): + # The day's notes were never sent to the model, so its index + # must stay out of the catalog and be retried later. + continue + supplemental.append(index_path) upsert_paths = list(dict.fromkeys([*checkpoint, *supplemental])) self.logger.info( f"[{self.name}] start changed={len(state.changed_paths)} failed_paths={len(state.failed_paths)} " - f"checkpoint={len(checkpoint)} day_indexes={len(day_index_paths)} " + f"deferred={len(state.deferred_paths)} checkpoint={len(checkpoint)} " + f"day_indexes={len(day_index_paths)} " f"deleted={len(state.deleted_paths)} persist={self.persist}", ) upserts = self._nodes(workspace, upsert_paths, state.existing) @@ -93,6 +112,11 @@ def render_summary(state: DreamState) -> str: ), f"- Catalog: checkpointed {len(state.checkpoint_paths)} changed path(s)", ] + if state.deferred_paths: + lines.append( + f"- Deferred: {len(state.deferred_paths)} changed path(s) not sent this run " + f"({', '.join(state.deferred_paths)})", + ) if state.nodes_created: lines.append(f"- Created: {', '.join(state.nodes_created)}") if state.nodes_updated: diff --git a/reme/steps/evolve/dream/utils.py b/reme/steps/evolve/dream/utils.py index bce5c135..a465a4eb 100644 --- a/reme/steps/evolve/dream/utils.py +++ b/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/``).""" + 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): diff --git a/tests/unit/test_auto_dream.py b/tests/unit/test_auto_dream.py index 5a6a0be8..f4d1a0e1 100644 --- a/tests/unit/test_auto_dream.py +++ b/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 = ( @@ -44,10 +49,11 @@ def _touch(path: Path, text: str = "x") -> Path: class _Catalog(BaseFileCatalog): - def __init__(self): + def __init__(self, nodes=()): super().__init__() self.upserts = [] self.dumps = 0 + self.nodes = list(nodes) async def upsert(self, nodes): self.upserts.extend(nodes) @@ -56,7 +62,7 @@ class _Catalog(BaseFileCatalog): return None async def get_nodes(self, paths=None): - return [] + return list(self.nodes) async def dump(self): self.dumps += 1 @@ -1032,3 +1038,454 @@ 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, *, catalog=None, **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 or _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 + + +class _BatchAwareAgent(BaseAgentWrapper): + """Answer whatever batch it receives with units for that batch's own paths.""" + + def __init__(self, units_per_path: dict[str, int]): + super().__init__() + self.units_per_path = units_per_path + self.calls = 0 + self.requests = [] + + async def reply(self, _message, **_kwargs): + self.calls += 1 + self.requests.append((_message, _kwargs)) + units = [] + for path, count in self.units_per_path.items(): + if path not in _message: + continue + units.extend( + { + "name": f"{Path(path).stem}-{index}", + "bucket": "wiki", + "summary": f"unit {index} for {path}", + "paths": [path], + } + for index in range(count) + ) + return {"result": json.dumps({"units": units})} + + +@pytest.mark.asyncio +async def test_extract_uses_one_unit_budget_across_batches(tmp_path): + """Two batches of 3 + 2 units keep all 5 under one global budget. + + Regression for a double subtraction: the loop passed the remaining + allowance as the batch cap and ``clean_output`` subtracted the already + accepted units again, so later batches were silently truncated to nothing. + """ + first_path = "daily/2026-05-28/a.md" + second_path = "daily/2026-05-27/b.md" + _touch(tmp_path / first_path, "x" * 1000) + _touch(tmp_path / second_path, "x" * 1000) + + agent = _BatchAwareAgent({first_path: 3, second_path: 2}) + + 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) == 5, sorted(unit["name"] for unit in units) + + +async def _finish_dream(response, catalog): + """Run the finish step on the state the extract step produced.""" + finish = await DreamFinishStep()( + RuntimeContext( + dream=response.metadata["dream"], + file_catalog=catalog, + ), + ) + return finish.metadata["dream"] + + +@pytest.mark.asyncio +async def test_extract_keeps_a_bisected_half_inside_the_same_budget(tmp_path): + """Bisection must not hand out an extra allowance on top of the global cap. + + The parent call fails, so the batch is really split: the first child keeps + its units and the second is capped by what is left. + """ + first_path = "daily/2026-05-28/a.md" + second_path = "daily/2026-05-27/b.md" + _touch(tmp_path / first_path, "x" * 1000) + _touch(tmp_path / second_path, "x" * 1000) + + agent = _LateFailAgent({first_path: 3, second_path: 2}) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=4, + max_chars_per_batch=4000, + ) + + units = response.metadata["dream"]["units"] + assert agent.calls == 3, "the failing parent must really be split into two children" + assert len(units) == 4, sorted(unit["name"] for unit in units) + + +@pytest.mark.asyncio +async def test_extract_defers_batches_that_never_reach_the_model(tmp_path): + """A batch skipped after the budget fills must not be checkpointed. + + Regression: the loop broke once ``max_units`` was reached but the unseen + paths stayed in ``changed_paths`` and were not marked failed, so the finish + step checkpointed files that were never sent to the model and a later scan + skipped them. + """ + first_path = "daily/2026-05-28/a.md" + second_path = "daily/2026-05-27/b.md" + _touch(tmp_path / first_path, "x" * 1000) + _touch(tmp_path / second_path, "x" * 1000) + + agent = _BatchAwareAgent({first_path: 5, second_path: 5}) + + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=5, + max_chars_per_batch=1200, + ) + + assert agent.calls == 1 + sent = agent.requests[0][0] + unseen = second_path if first_path in sent else first_path + seen = first_path if unseen == second_path else second_path + + dream = await _finish_dream(response, _Catalog()) + + assert seen in set(dream["checkpoint_paths"]) + assert unseen not in set(dream["checkpoint_paths"]), "a file never sent to the model must stay eligible" + + +class _LateFailAgent(_BatchAwareAgent): + """Fail the first (whole-batch) call, then answer whatever batch it receives.""" + + async def reply(self, _message, **kwargs): + if self.calls == 0: + self.calls += 1 + self.requests.append((_message, kwargs)) + raise RuntimeError("prompt too large") + return await super().reply(_message, **kwargs) + + +@pytest.mark.asyncio +async def test_extract_does_not_mark_a_path_skipped_after_a_parent_failure(tmp_path): + """A path whose batch never reaches the model must stay rescan-eligible. + + Regression: the whole batch was recorded as processed *before* the agent + call, so a failed parent plus a child skipped for budget kept the skipped + half marked and the finish step checkpointed a file the model never saw. + """ + first_path = "daily/2026-05-28/a.md" + second_path = "daily/2026-05-27/b.md" + _touch(tmp_path / first_path, "x" * 1000) + _touch(tmp_path / second_path, "x" * 1000) + + agent = _LateFailAgent({first_path: 3, second_path: 3}) + response = await _run_dream_extract( + tmp_path, + agent, + scan_days=2, + max_units=3, + max_chars_per_batch=4000, + ) + + assert agent.calls == 2, "the parent fails once and one child fills the budget" + sent = agent.requests[1][0] + unseen = second_path if first_path in sent else first_path + seen = first_path if unseen == second_path else second_path + assert unseen not in response.metadata["dream"]["processed_paths"] + + dream = await _finish_dream(response, _Catalog()) + + assert seen in set(dream["checkpoint_paths"]) + assert unseen not in set(dream["checkpoint_paths"]), "a file never sent to the model must stay eligible" + assert unseen in dream["deferred_paths"] + + +@pytest.mark.asyncio +async def test_finish_defers_the_day_index_of_a_deferred_day(tmp_path): + """A deferred day must not get its index upserted into the catalog. + + Regression: ``supplemental`` day indexes bypassed the deferred filter, so a + day whose files never reached the model still landed in the catalog and + lost its retry eligibility. + """ + first_path = "daily/2026-05-28/a.md" + second_path = "daily/2026-05-27/b.md" + first_index, second_index = "daily/2026-05-28.md", "daily/2026-05-27.md" + _touch(tmp_path / first_path, "x" * 1000) + _touch(tmp_path / second_path, "x" * 1000) + _touch(tmp_path / first_index, "index") + _touch(tmp_path / second_index, "index") + + catalog = _Catalog( + [FileNode(path=rel, st_mtime=(tmp_path / rel).stat().st_mtime) for rel in (first_index, second_index)], + ) + agent = _BatchAwareAgent({first_path: 5, second_path: 5}) + response = await _run_dream_extract( + tmp_path, + agent, + catalog=catalog, + scan_days=2, + max_units=5, + max_chars_per_batch=1200, + ) + + assert agent.calls == 1, "the first batch fills the budget and the second is deferred" + sent = agent.requests[0][0] + unseen_day, seen_day = ("2026-05-27", "2026-05-28") if first_path in sent else ("2026-05-28", "2026-05-27") + + dream = await _finish_dream(response, catalog) + + assert any(path.startswith(f"daily/{unseen_day}/") for path in dream["deferred_paths"]) + upserted = {node.path for node in catalog.upserts} + assert f"daily/{seen_day}.md" in upserted + assert f"daily/{unseen_day}.md" not in upserted, "a deferred day index must not be upserted" + + +@pytest.mark.asyncio +async def test_extract_with_zero_budget_defers_every_changed_path(tmp_path): + """``max_units=0`` must not look like legacy state and checkpoint everything. + + Regression: with no extraction call at all ``processed_paths`` stayed + ``None``, which the finish step read as "state written by an older version" + and then checkpointed every changed path although the model never ran. + """ + path, index = "daily/2026-05-28/a.md", "daily/2026-05-28.md" + _touch(tmp_path / path, "x" * 1000) + _touch(tmp_path / index, "index") + + catalog = _Catalog([FileNode(path=index, st_mtime=(tmp_path / index).stat().st_mtime)]) + response = await _run_dream_extract(tmp_path, _BatchAwareAgent({path: 1}), catalog=catalog, max_units=0) + + assert response.metadata["dream"]["processed_paths"] == [] + dream = await _finish_dream(response, catalog) + + assert dream["checkpoint_paths"] == [] + assert path in dream["deferred_paths"] + assert not catalog.upserts, "nothing may be catalogued when the model never ran"