fix(index): retry changed source during indexing

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)."""

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)

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

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)
 

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