scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-05 07:39:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

fix(index): retry and coalesce updated file changes

Browse source
This commit is contained in:
mikemikimike 2026-10-14 19:11:56 +08:00
parent 47c2b6a2d6
commit 1fb1606e57
No known key found for this signature in database
4 changed files with 159 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
reme/steps/index/_change_batch.py
View file

@ -15,22 +15,31 @@ def _normalize_change(raw) -> Change | None:
return None
def coalesce_changes(changes: list[dict], path_exists: Callable[[str], bool] | None = None) -> list[dict]:
def coalesce_changes(
changes: list[dict],
path_exists: Callable[[str], bool] | None = None,
*,
path_key: Callable[[str], str] | None = None,
) -> list[dict]:
"""Collapse duplicate path events to the final on-disk state."""
by_path: OrderedDict[str, set[Change]] = OrderedDict()
by_path: OrderedDict[str, tuple[str, set[Change]]] = OrderedDict()
for item in changes:
if not isinstance(item, dict) or "path" not in item:
continue
change = _normalize_change(item.get("change"))
if change not in (Change.added, Change.modified, Change.deleted):
continue
by_path.setdefault(item["path"], set()).add(change)
path = item["path"]
key = path_key(path) if path_key is not None else path
if key not in by_path:
by_path[key] = (path, set())
by_path[key][1].add(change)
def exists(path: str) -> bool:
return path_exists(path) if path_exists is not None else Path(path).is_file()
result: list[dict] = []
for path, seen in by_path.items():
for path, seen in by_path.values():
if not exists(path):
result.append({"change": Change.deleted.name, "path": path})
elif seen == {Change.added}:
@ -40,10 +49,15 @@ def coalesce_changes(changes: list[dict], path_exists: Callable[[str], bool] | N
return result
def bucket_changes(changes: list[dict], path_exists: Callable[[str], bool] | None = None) -> dict[Change, list[str]]:
def bucket_changes(
changes: list[dict],
path_exists: Callable[[str], bool] | None = None,
*,
path_key: Callable[[str], str] | None = None,
) -> dict[Change, list[str]]:
"""Group changes by watchfiles.Change."""
buckets: dict[Change, list[str]] = {Change.added: [], Change.modified: [], Change.deleted: []}
for item in coalesce_changes(changes, path_exists=path_exists):
for item in coalesce_changes(changes, path_exists=path_exists, path_key=path_key):
change = _normalize_change(item["change"])
if isinstance(change, Change) and change in buckets:
buckets[change].append(item["path"])

39
reme/steps/index/update_changes.py
View file

@ -86,7 +86,11 @@ class ChangeApplyStep(BaseStep):
assert self.context is not None
changes: list[dict] = self.context.get("changes") or []
persist = bool(self.context.get("persist", True)) if self.persist is None else self.persist
buckets = bucket_changes(changes, path_exists=lambda p: self._to_abs_path(p).is_file())
buckets = bucket_changes(
changes,
path_exists=lambda p: self._to_abs_path(p).is_file(),
path_key=self._change_path_key,
)
results = await self._apply_existing(buckets)
results.extend(await self._apply_deleted(buckets[Change.deleted]))
if persist and results:
@ -126,9 +130,11 @@ class ChangeApplyStep(BaseStep):
estimated_bytes = 0
memory_budget = self._batch_memory_budget()
item = await self._try_build_item(change, path, abs_path, results)
if item is None:
built = await self._try_build_item(change, path, abs_path, size_bytes, results)
if built is None:
continue
item, size_bytes = built
del built
item_bytes = self._try_estimate_memory(
path,
self.estimate_item_memory,
@ -201,12 +207,18 @@ class ChangeApplyStep(BaseStep):
change: Change,
path: str,
abs_path: Path,
size_bytes: int,
results: list[dict],
) -> Any | None:
) -> tuple[Any, int] | None:
for attempt in range(self.source_change_retries + 1):
if attempt:
inspected = await self._inspect_path(change, path, results)
if inspected is None:
return None
abs_path, size_bytes = inspected
try:
self.logger.debug(f"Processing {change.name} file: {path}")
return await self.build_item(abs_path)
return await self.build_item(abs_path), size_bytes
except FileChangedDuringChunkingError as e:
if attempt < self.source_change_retries:
self.logger.debug(f"Source changed while processing {path}; retrying")
@ -301,6 +313,10 @@ class ChangeApplyStep(BaseStep):
p = Path(path)
return p if p.is_absolute() else self.workspace_path / p
def _change_path_key(self, path: str) -> str:
"""Return the identity used to coalesce path aliases before batching."""
return path
@R.register("update_catalog_step")
class UpdateCatalogStep(ChangeApplyStep):
@ -362,6 +378,10 @@ class UpdateIndexStep(ChangeApplyStep):
)
self._source_versions: dict[str, os.stat_result] = {}
def _change_path_key(self, path: str) -> str:
"""Coalesce absolute and workspace-relative aliases of the same file."""
return os.path.normcase(os.fspath(self._to_abs_path(path).absolute()))
async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]:
try:
return await super()._apply_existing(buckets)
@ -434,9 +454,11 @@ class UpdateIndexStep(ChangeApplyStep):
inspected = await self._inspect_path(change, path, retry_results)
if inspected is None:
continue
abs_path, _ = inspected
item = await self._try_build_item(change, path, abs_path, retry_results)
if item is not None:
abs_path, size_bytes = inspected
built = await self._try_build_item(change, path, abs_path, size_bytes, retry_results)
if built is not None:
item, _ = built
del built
retried = await self._try_upsert(
change,
[item],
@ -444,6 +466,7 @@ class UpdateIndexStep(ChangeApplyStep):
retry_source_changes=False,
)
result_by_path.update({result["path"]: result for result in retried})
del item
result_by_path.update({result["path"]: result for result in retry_results})
return [result_by_path[path] for path in ok_paths]

32
tests/unit/test_background_steps.py
View file

@ -14,7 +14,6 @@ import asyncio
import datetime
import os
import tempfile
import threading
import warnings
import weakref
from collections.abc import AsyncGenerator, Callable
@ -1108,7 +1107,7 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
cwd = Path.cwd()
path = write_file(cwd / "daily" / "source.md", "initial contents")
entered, release = threading.Event(), threading.Event()
item_built = asyncio.Event()
watcher_started = asyncio.Event()
chunk_calls: list[Path] = []
@ -1142,16 +1141,18 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
)
original_chunk = chunker._chunk_sync
original_build_item = UpdateIndexStep.build_item
def change_during_first_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
def count_chunk(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
result = original_chunk(file_path, rel_path)
chunk_calls.append(file_path)
if len(chunk_calls) == 1:
entered.set()
if not release.wait(5):
raise TimeoutError("test did not release the initial parse")
return result
async def signal_item_built(step: UpdateIndexStep, file_path: Path) -> tuple[FileNode, list[FileChunk]]:
item = await original_build_item(step, file_path)
item_built.set()
return item
from reme.steps.index import watch_changes as watch_changes_module
original_awatch = watch_changes_module.awatch
@ -1179,22 +1180,18 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
):
yield changes
monkeypatch.setattr(chunker, "_chunk_sync", change_during_first_parse)
monkeypatch.setattr(chunker, "_chunk_sync", count_chunk)
monkeypatch.setattr(UpdateIndexStep, "build_item", signal_item_built)
monkeypatch.setattr(watch_changes_module, "awatch", track_watcher_start)
try:
await store.start()
await chunker.start()
await job.start()
for _ in range(500):
if entered.is_set():
break
await asyncio.sleep(0.01)
assert entered.is_set(), "initial file parse did not start"
latest = "latest contents after the initial scan"
path.write_text(latest, encoding="utf-8")
release.set()
async with store._maintenance_guard():
await job.start()
await asyncio.wait_for(item_built.wait(), timeout=5)
path.write_text(latest, encoding="utf-8")
await asyncio.wait_for(watcher_started.wait(), timeout=5)
nodes = await store.get_nodes()
@ -1203,7 +1200,6 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
assert "".join(chunk.text for chunk in chunks) == latest
assert len(chunk_calls) == 2
finally:
release.set()
await job.close()
await chunker.close()
await store.close()

94
tests/unit/test_file_chunker_workers.py
View file

@ -525,6 +525,100 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file
assert not step._source_versions
@pytest.mark.asyncio
@pytest.mark.parametrize("new_size", [12, 100])
async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate(
tmp_path: Path,
new_size: int,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A parse retry rechecks the size cap and estimates memory from the new size."""
monkeypatch.chdir(tmp_path)
path = tmp_path / "source.txt"
path.write_text("small")
context = ApplicationContext(workspace_dir=str(tmp_path))
chunker = DefaultFileChunker(app_context=context)
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
async with LocalFileStore(embedding_store="") as store, chunker:
await store.upsert([await chunker.chunk(path)])
entered, release = threading.Event(), threading.Event()
calls: list[Path] = []
original_chunk = chunker._chunk_sync
def change_during_first_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
result = original_chunk(file_path, rel_path)
calls.append(file_path)
if len(calls) == 1:
entered.set()
assert release.wait(5)
return result
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
estimated_sizes: list[int] = []
original_estimate = step.estimate_item_memory
def record_estimate(item: tuple[FileNode, list[FileChunk]], file_path: Path, size_bytes: int) -> int:
estimated_sizes.append(size_bytes)
return original_estimate(item, file_path, size_bytes)
monkeypatch.setattr(step, "estimate_item_memory", record_estimate)
with patch.object(chunker, "_chunk_sync", side_effect=change_during_first_parse):
task = asyncio.create_task(
step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}], max_file_bytes=20)),
)
await _wait_started(entered)
path.write_text("x" * new_size)
release.set()
response = await task
assert response.success is True
assert len(response.answer) == 1
if new_size > 20:
assert response.answer[0]["skipped"] is True
assert response.answer[0]["reason"] == "file_too_large"
assert calls == [path]
assert not estimated_sizes
assert await store.get_nodes() == []
else:
assert response.answer[0]["success"] is True
assert len(calls) == 2
assert estimated_sizes == [new_size]
nodes = await store.get_nodes()
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
assert not step._source_versions
@pytest.mark.asyncio
async def test_index_updates_deduplicate_absolute_and_relative_paths(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Absolute and workspace-relative aliases of one file enter a batch once."""
monkeypatch.chdir(tmp_path)
path = tmp_path / "source.txt"
path.write_text("source")
context = ApplicationContext(workspace_dir=str(tmp_path))
chunker = DefaultFileChunker(app_context=context)
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
async with LocalFileStore(embedding_store="") as store, chunker:
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
response = await step(
RuntimeContext(
changes=[
{"change": "modified", "path": str(path)},
{"change": "modified", "path": "source.txt"},
],
),
)
assert response.success is True
assert response.answer == [{"change": "modified", "path": str(path), "success": True}]
nodes = await store.get_nodes()
assert len(nodes) == 1
assert nodes[0].path == "source.txt"
assert not step._source_versions
@pytest.mark.asyncio
@pytest.mark.parametrize("outcome", ["change", "cancel"])
async def test_version_is_checked_after_waiting_for_store_lock(