scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-14 09:01:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

fix(dream): batch oversized extraction input (#664)

Browse source 
Split changed dream inputs into configurable file-size batches and bisect failed batches. Scope source references to each batch and merge same-name units across batches.

Keep one global unit budget and checkpoint only successfully processed, non-failed paths. Preserve retry eligibility for deferred inputs, including day indexes, and initialize new runs separately from legacy state.

Validation: 192 focused tests passed; isolated QwenPaw tests with a real model covered batching, restart deduplication, deferred continuation, and injected-failure recovery.

Closes #659
This commit is contained in:
北沐城歌 2026-11-02 11:54:34 +08:00 • committed by GitHub
parent 966c691d12
commit 7f2fa864c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 818 additions and 55 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
reme/config/default.yaml
View file

@ -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

12
reme/schema/dream.py
View file

@ -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 = ""

296
reme/steps/evolve/dream/extract.py
View file

@ -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})

30
reme/steps/evolve/dream/finish.py
View file

@ -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:

66
reme/steps/evolve/dream/utils.py
View file

@ -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):

463
tests/unit/test_auto_dream.py
View file

@ -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"