diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index bb477c95..5f4c9081 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -11,6 +11,7 @@ from typing import Callable, TypeVar from ..base_component import BaseComponent from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode +from ...utils.file_utils import check_file_version WorkResult = TypeVar("WorkResult") @@ -92,7 +93,7 @@ class BaseFileChunker(BaseComponent): result, version = await _finish_work(future) if self._worker_closing or generation != self._worker_generation: raise RuntimeError("File chunker closed while parsing") - self._check_file_version(file_path, version, file_path.stat()) + check_file_version(file_path, version, file_path.stat()) return result finally: if temporary_worker: @@ -107,15 +108,9 @@ class BaseFileChunker(BaseComponent): before = path.stat() result = chunk(path, rel_path) after = path.stat() - BaseFileChunker._check_file_version(path, before, after) + check_file_version(path, before, after) return result, after - @staticmethod - def _check_file_version(path: Path, before: os.stat_result, after: os.stat_result) -> None: - fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") - if any(getattr(before, field) != getattr(after, field) for field in fields): - raise RuntimeError(f"File changed while chunking: {path}") - @abstractmethod async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: """Chunk a file into (node, chunks).""" diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 3dbcce4c..c74e2848 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -16,6 +16,7 @@ from ...components import R from ...components.file_chunker import BaseFileChunker from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode +from ...utils.file_utils import FileChangedDuringChunkingError, check_file_version class ChangeApplyStep(BaseStep): @@ -23,6 +24,7 @@ class ChangeApplyStep(BaseStep): target_name = "target" reads_file_content = False + source_change_retries: int = 0 def __init__( self, @@ -47,7 +49,6 @@ class ChangeApplyStep(BaseStep): self.chunk_memory_overhead_bytes = int(chunk_memory_overhead_bytes) self.estimated_chunk_bytes = int(estimated_chunk_bytes) self.float16_bytes = int(float16_bytes) - self._source_versions: dict[int, os.stat_result] = {} if self.batch_max_files <= 0: raise ValueError("batch_max_files must be greater than zero") if not math.isfinite(self.batch_available_memory_ratio) or not 0 < self.batch_available_memory_ratio <= 1: @@ -201,14 +202,22 @@ class ChangeApplyStep(BaseStep): path: str, abs_path: Path, results: list[dict], - ): - try: - self.logger.debug(f"Processing {change.name} file: {path}") - return await self.build_item(abs_path) - except Exception as e: - self.logger.exception(f"Failed to process {path}") - results.append({"change": change.name, "path": path, "success": False, "error": str(e)}) - return None + ) -> Any | None: + for attempt in range(self.source_change_retries + 1): + try: + self.logger.debug(f"Processing {change.name} file: {path}") + return await self.build_item(abs_path) + except FileChangedDuringChunkingError as e: + if attempt < self.source_change_retries: + self.logger.debug(f"Source changed while processing {path}; retrying") + continue + self.logger.exception(f"Failed to process {path}") + results.append({"change": change.name, "path": path, "success": False, "error": str(e)}) + return None + except Exception as e: + self.logger.exception(f"Failed to process {path}") + results.append({"change": change.name, "path": path, "success": False, "error": str(e)}) + return None def _try_estimate_memory( self, @@ -324,6 +333,34 @@ class UpdateIndexStep(ChangeApplyStep): target_name = "file_store" reads_file_content = True + source_change_retries: int = 1 + + def __init__( + self, + persist: bool | None = None, + batch_max_files: int = 100, + batch_available_memory_ratio: float = 0.10, + batch_memory_target_bytes: int = 512 * 1024 * 1024, + batch_memory_expansion_factor: float = 8.0, + file_memory_overhead_bytes: int = 32 * 1024, + chunk_memory_overhead_bytes: int = 4 * 1024, + estimated_chunk_bytes: int = 10_000, + float16_bytes: int = 2, + **kwargs: Any, + ) -> None: + super().__init__( + persist=persist, + batch_max_files=batch_max_files, + batch_available_memory_ratio=batch_available_memory_ratio, + batch_memory_target_bytes=batch_memory_target_bytes, + batch_memory_expansion_factor=batch_memory_expansion_factor, + file_memory_overhead_bytes=file_memory_overhead_bytes, + chunk_memory_overhead_bytes=chunk_memory_overhead_bytes, + estimated_chunk_bytes=estimated_chunk_bytes, + float16_bytes=float16_bytes, + **kwargs, + ) + self._source_versions: dict[str, os.stat_result] = {} async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]: try: @@ -336,8 +373,8 @@ class UpdateIndexStep(ChangeApplyStep): before = path.stat() item = await self.chunk_file(path) after = path.stat() - BaseFileChunker._check_file_version(path, before, after) # pylint: disable=protected-access - self._source_versions[id(item[0])] = after + check_file_version(path, before, after) + self._source_versions[item[0].path] = after return item async def _try_upsert( @@ -345,31 +382,70 @@ class UpdateIndexStep(ChangeApplyStep): change: Change, items: list[tuple[FileNode, list[FileChunk]]], ok_paths: list[str], + *, + retry_source_changes: bool = True, ) -> list[dict]: # Recheck after waiting for the same guard used by every store mutation. + retry_paths: list[str] = [] + result_by_path: dict[str, dict] = {} async with self.file_store._maintenance_guard(): # pylint: disable=protected-access current_items: list[tuple[FileNode, list[FileChunk]]] = [] current_paths: list[str] = [] - results: list[dict] = [] for item, path in zip(items, ok_paths): node = item[0] try: - BaseFileChunker._check_file_version( # pylint: disable=protected-access + check_file_version( self._to_abs_path(node.path), - self._source_versions[id(node)], + self._source_versions[node.path], self._to_abs_path(node.path).stat(), ) + except FileChangedDuringChunkingError as exc: + if retry_source_changes: + retry_paths.append(path) + else: + result_by_path[path] = { + "change": change.name, + "path": path, + "success": False, + "error": str(exc), + } except (OSError, RuntimeError) as exc: - results.append({"change": change.name, "path": path, "success": False, "error": str(exc)}) + result_by_path[path] = { + "change": change.name, + "path": path, + "success": False, + "error": str(exc), + } else: current_items.append(item) current_paths.append(path) finally: - self._source_versions.pop(id(node)) + self._source_versions.pop(node.path) if current_items: - results.extend(await super()._try_upsert(change, current_items, current_paths)) - by_path = {result["path"]: result for result in results} - return [by_path[path] for path in ok_paths] + result_by_path.update( + { + result["path"]: result + for result in await super()._try_upsert(change, current_items, current_paths) + }, + ) + + retry_results: list[dict] = [] + for path in retry_paths: + 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: + retried = await self._try_upsert( + change, + [item], + [path], + retry_source_changes=False, + ) + result_by_path.update({result["path"]: result for result in retried}) + result_by_path.update({result["path"]: result for result in retry_results}) + return [result_by_path[path] for path in ok_paths] async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None: await self.file_store.upsert(items) diff --git a/reme/utils/file_utils.py b/reme/utils/file_utils.py new file mode 100644 index 00000000..440f1635 --- /dev/null +++ b/reme/utils/file_utils.py @@ -0,0 +1,15 @@ +"""Shared file-version checks used by file readers and index updates.""" + +import os +from pathlib import Path + + +class FileChangedDuringChunkingError(RuntimeError): + """Raised when a source file changes while it is being parsed.""" + + +def check_file_version(path: Path, before: os.stat_result, after: os.stat_result) -> None: + """Raise when the source file's identity or content metadata changed.""" + fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + if any(getattr(before, field) != getattr(after, field) for field in fields): + raise FileChangedDuringChunkingError(f"File changed while chunking: {path}") diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 87867a5b..9fd5b58e 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -14,8 +14,10 @@ import asyncio import datetime import os import tempfile +import threading import warnings import weakref +from collections.abc import AsyncGenerator, Callable from pathlib import Path from unittest.mock import MagicMock, patch @@ -28,8 +30,10 @@ from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.file_chunker import DefaultFileChunker from reme.components.file_catalog import LocalFileCatalog from reme.components.file_store import LocalFileStore +from reme.components.job.background_job import BackgroundJob from reme.components.runtime_context import RuntimeContext from reme.enumeration import ComponentEnum +from reme.schema import FileChunk, FileNode from reme.steps.evolve.auto_memory import AutoMemoryStep from reme.steps.evolve.base_auto_resource import _compute_note_stem from reme.steps.evolve.auto_resource import AutoResourceStep @@ -1044,6 +1048,116 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches(): asyncio.run(run()) +def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatch: pytest.MonkeyPatch) -> None: + """The real background job indexes the latest source before its watcher baseline.""" + + async def run() -> None: + 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() + watcher_started = asyncio.Event() + chunk_calls: list[Path] = [] + + store = LocalFileStore(name="default", embedding_store="") + app_ctx = _make_app_context(cwd) + chunker = DefaultFileChunker(app_context=app_ctx) + app_ctx.components = { + ComponentEnum.FILE_STORE: {"default": store}, + ComponentEnum.FILE_CHUNKER: {"default": chunker}, + } + job = BackgroundJob( + name="index_update_loop", + app_context=app_ctx, + watch_dirs=["daily_dir"], + watch_suffixes=["md"], + steps=[ + { + "backend": "init_changes_step", + "monitor_type": "file_store", + "monitor_name": "default", + "dispatch_steps": ["update_index_step"], + }, + { + "backend": "watch_changes_step", + "step": 20, + "debounce": 50, + "poll_delay_ms": 10, + "dispatch_steps": [{"backend": "update_index_step", "persist": False}], + }, + ], + ) + + 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) + 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 + + from reme.steps.index import watch_changes as watch_changes_module + + original_awatch = watch_changes_module.awatch + + async def track_watcher_start( + *paths: str | Path, + watch_filter: Callable[[Change, str], bool], + recursive: bool, + force_polling: bool, + debounce: int, + step: int, + poll_delay_ms: int, + stop_event: asyncio.Event, + ) -> AsyncGenerator[set[tuple[Change, str]], None]: + watcher_started.set() + async for changes in original_awatch( + *paths, + watch_filter=watch_filter, + recursive=recursive, + force_polling=force_polling, + debounce=debounce, + step=step, + poll_delay_ms=poll_delay_ms, + stop_event=stop_event, + ): + yield changes + + monkeypatch.setattr(chunker, "_chunk_sync", change_during_first_parse) + 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() + + await asyncio.wait_for(watcher_started.wait(), timeout=5) + nodes = await store.get_nodes() + assert len(nodes) == 1 + chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids] + 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() + + asyncio.run(run()) + + def test_digest_watch_loop_init_dispatch_updates_named_catalog_and_logs(): """digest_watch_loop style config updates the digest catalog without touching resource/default catalogs.""" @@ -1349,12 +1463,12 @@ def test_auto_resource_handles_file_removed_before_stat(): original_stat = Path.stat def existing_is_file(path, *args, **kwargs): - if path == source: + if Path(path).name == source.name: return True return original_is_file(path, *args, **kwargs) def disappearing_stat(path, *args, **kwargs): - if path == source: + if Path(path).name == source.name: raise FileNotFoundError("file disappeared") return original_stat(path, *args, **kwargs) diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index 7047b9c5..1ef3b441 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -411,11 +411,17 @@ async def test_existing_content_searchable_during_indexing( await updating else: response = await updating - assert response.success == (outcome == "complete") + assert response.success is (outcome in ("complete", "change")) nodes = await store.get_nodes() chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids] - assert chunks[0].text == ("new searchable content" if outcome == "complete" else "old searchable content") - if outcome != "complete": + expected_text = { + "complete": "new searchable content", + "change": "latest source content", + "cancel": "old searchable content", + "error": "old searchable content", + }[outcome] + assert chunks[0].text == expected_text + if outcome in ("cancel", "error"): recovered = await step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}])) assert recovered.success is True nodes = await store.get_nodes() @@ -424,8 +430,8 @@ async def test_existing_content_searchable_during_indexing( @pytest.mark.asyncio -@pytest.mark.parametrize("change", ["modify", "delete", "replace", "cancel"]) -async def test_changed_batched_file_is_not_published_and_valid_files_update( +@pytest.mark.parametrize("change", ["modify", "delete", "replace", "oversize", "cancel"]) +async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_files_update( tmp_path: Path, change: str, monkeypatch: pytest.MonkeyPatch, @@ -454,8 +460,13 @@ async def test_changed_batched_file_is_not_published_and_valid_files_update( changes = [{"change": "modified", "path": str(path)} for path in (first_path, second_path)] step = UpdateIndexStep(file_store=store, persist=False, app_context=context) + runtime_context = ( + RuntimeContext(changes=changes, max_file_bytes=20) + if change == "oversize" + else RuntimeContext(changes=changes) + ) with patch.object(chunker, "chunk", side_effect=blocking_second): - task = asyncio.create_task(step(RuntimeContext(changes=changes))) + task = asyncio.create_task(step(runtime_context)) await _wait_started(entered) try: if change == "modify": @@ -464,6 +475,8 @@ async def test_changed_batched_file_is_not_published_and_valid_files_update( first_path.unlink() elif change == "cancel": task.cancel() + elif change == "oversize": + first_path.write_text("content exceeds the configured limit") else: replacement = tmp_path / "replacement" replacement.write_text("replacement first") @@ -475,11 +488,26 @@ async def test_changed_batched_file_is_not_published_and_valid_files_update( await task else: response = await task - assert response.success is False - assert [result["success"] for result in response.answer] == [False, True] + expected_success = change in ("modify", "replace", "oversize") + assert response.success is expected_success + assert [result["success"] for result in response.answer] == ( + [True, True] if expected_success else [False, True] + ) + if change == "oversize": + assert response.answer[0]["skipped"] is True + assert response.answer[0]["reason"] == "file_too_large" assert not step._source_versions nodes = {node.path: node for node in await store.get_nodes()} - assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == "original first" + if change == "oversize": + assert "first.txt" not in nodes + else: + expected_first = { + "modify": "latest first", + "delete": "original first", + "replace": "replacement first", + "cancel": "original first", + }[change] + assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == expected_first assert store.file_chunks[nodes["second.txt"].chunk_ids[0]].text == ( "original second" if change == "cancel" else "updated second" ) @@ -536,12 +564,14 @@ async def test_version_is_checked_after_waiting_for_store_lock( path.write_text("latest source") if outcome == "change": response = await task - assert response.success is False + assert response.success is True assert not step._source_versions nodes = await store.get_nodes() - assert store.file_chunks[nodes[0].chunk_ids[0]].text == "original source" - recovered = await step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}])) - assert recovered.success is True - nodes = await store.get_nodes() - assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text() + expected_text = "latest source" if outcome == "change" else "original source" + assert store.file_chunks[nodes[0].chunk_ids[0]].text == expected_text + if outcome == "cancel": + recovered = await step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}])) + assert recovered.success is True + nodes = await store.get_nodes() + assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text() assert not step._source_versions