From 1fb1606e5746df83f1cae6b5de812a34fe25bbc0 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Wed, 14 Oct 2026 19:11:56 +0800 Subject: [PATCH] fix(index): retry and coalesce updated file changes --- reme/steps/index/_change_batch.py | 26 +++++-- reme/steps/index/update_changes.py | 39 +++++++--- tests/unit/test_background_steps.py | 32 ++++----- tests/unit/test_file_chunker_workers.py | 94 +++++++++++++++++++++++++ 4 files changed, 159 insertions(+), 32 deletions(-) diff --git a/reme/steps/index/_change_batch.py b/reme/steps/index/_change_batch.py index 5000ff27..0d67057f 100644 --- a/reme/steps/index/_change_batch.py +++ b/reme/steps/index/_change_batch.py @@ -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"]) diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index c74e2848..26492e5a 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -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] diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 5365fc02..f5987908 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -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() diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index 1ef3b441..47c04e28 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -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(