fix(index): recover unstable initial scans

reme/components/file_chunker/base_file_chunker.py

@@ -3,6 +3,7 @@
 import asyncio
 import contextvars
 import os
+import threading
 from abc import abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
@@ -42,26 +43,45 @@ class BaseFileChunker(BaseComponent):
         super().__init__(**kwargs)
         self.supported_extensions: list[str] = supported_extensions or []
         self._worker: ThreadPoolExecutor | None = None
-        self._worker_lock = asyncio.Lock()
+        self._worker_lock = threading.Lock()
+        self._worker_state_lock = threading.Lock()
         self._worker_generation = 0
         self._worker_closing = False
 
     async def _start(self) -> None:
-        self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker")
-        self._worker_generation += 1
-        self._worker_closing = False
+        await self._acquire_worker_lock()
+        try:
+            with self._worker_state_lock:
+                self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker")
+                self._worker_generation += 1
+                self._worker_closing = False
+        finally:
+            self._worker_lock.release()
 
     async def _close(self) -> None:
-        self._worker_closing = True
-        self._worker_generation += 1
+        with self._worker_state_lock:
+            self._worker_closing = True
+            self._worker_generation += 1
         await _finish_work(asyncio.create_task(self._close_worker()))
 
     async def _close_worker(self) -> None:
-        async with self._worker_lock:
-            if self._worker is not None:
-                # The admission lock is held through completion, so shutdown has no running work.
-                self._worker.shutdown(wait=True)
+        await self._acquire_worker_lock()
+        try:
+            with self._worker_state_lock:
+                self._worker_closing = True
+                self._worker_generation += 1
+                worker = self._worker
                 self._worker = None
+            if worker is not None:
+                worker.shutdown(wait=True)
+        finally:
+            self._worker_lock.release()
+
+    async def _acquire_worker_lock(self) -> None:
+        """Acquire cross-loop admission without blocking an event loop thread."""
+        # Ownership spans awaited work, so a synchronous lock context would block this loop.
+        while not self._worker_lock.acquire(blocking=False):  # pylint: disable=consider-using-with
+            await asyncio.sleep(0.01)
 
     async def _chunk_in_worker(
         self,
@@ -69,35 +89,50 @@ class BaseFileChunker(BaseComponent):
         chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]],
     ) -> tuple[FileNode, list[FileChunk]]:
         """Admit one parse at a time without accumulating an executor work queue."""
-        generation = self._worker_generation
+        with self._worker_state_lock:
+            generation = self._worker_generation
         file_path = Path(path)
         rel_path = self.to_workspace_relative(path)
-        async with self._worker_lock:
-            if self._worker_closing or generation != self._worker_generation:
-                raise RuntimeError("File chunker closed while waiting to parse")
-            worker = self._worker
-            temporary_worker = worker is None
-            if worker is None:
-                # Standalone chunk() remains usable without start(), with no persistent threads.
-                worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker")
-            try:
-                context = contextvars.copy_context()
-                future = asyncio.get_running_loop().run_in_executor(
-                    worker,
-                    context.run,
-                    self._chunk_unchanged_file,
-                    file_path,
-                    rel_path,
-                    chunk,
-                )
-                result, version = await _finish_work(future)
+        await self._acquire_worker_lock()
+        temporary_worker = False
+        try:
+            with self._worker_state_lock:
                 if self._worker_closing or generation != self._worker_generation:
-                    raise RuntimeError("File chunker closed while parsing")
-                check_file_version(file_path, version, file_path.stat())
-                return result
-            finally:
-                if temporary_worker:
-                    worker.shutdown(wait=True)
+                    raise RuntimeError("File chunker closed while waiting to parse")
+                worker = self._worker
+                temporary_worker = worker is None
+            if temporary_worker:
+                # Standalone chunk() remains usable without start(), with no persistent threads.
+                with ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") as worker:
+                    return await self._run_chunk(worker, generation, file_path, rel_path, chunk)
+            assert worker is not None
+            return await self._run_chunk(worker, generation, file_path, rel_path, chunk)
+        finally:
+            self._worker_lock.release()
+
+    async def _run_chunk(
+        self,
+        worker: ThreadPoolExecutor,
+        generation: int,
+        file_path: Path,
+        rel_path: str,
+        chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]],
+    ) -> tuple[FileNode, list[FileChunk]]:
+        context = contextvars.copy_context()
+        future = asyncio.get_running_loop().run_in_executor(
+            worker,
+            context.run,
+            self._chunk_unchanged_file,
+            file_path,
+            rel_path,
+            chunk,
+        )
+        result, version = await _finish_work(future)
+        with self._worker_state_lock:
+            if self._worker_closing or generation != self._worker_generation:
+                raise RuntimeError("File chunker closed while parsing")
+        check_file_version(file_path, version, file_path.stat())
+        return result
 
     @staticmethod
     def _chunk_unchanged_file(

reme/steps/index/init_changes.py

@@ -73,7 +73,11 @@ class InitChangesStep(BaseStep):
         self.context["changes"] = changes
         if changes:
             self.logger.info(f"[{self.name}] scan {self.monitor_type}:{self.monitor_name}: {counts}")
-            await self.dispatch_steps(self.dispatch_step_specs, changes=changes)
+            await self.dispatch_steps(
+                self.dispatch_step_specs,
+                changes=changes,
+                retry_source_changes_before_watch=True,
+            )
         else:
             self.logger.info(f"[{self.name}] {self.monitor_type}:{self.monitor_name} is up to date")
         self.context.response.metadata["counts"] = counts

reme/steps/index/update_changes.py

@@ -18,6 +18,9 @@ from ...enumeration import ComponentEnum
 from ...schema import FileChunk, FileNode
 from ...utils.file_utils import FileChangedDuringChunkingError, check_file_version
 
+_SOURCE_CHANGE_RETRY_MARKER = "_source_changed"
+_INITIAL_SCAN_RECOVERY_PASSES = 1
+
 
 class ChangeApplyStep(BaseStep):
     """Shared added/modified/deleted handling for index update targets."""
@@ -91,7 +94,40 @@ class ChangeApplyStep(BaseStep):
             path_exists=lambda p: self._to_abs_path(p).is_file(),
             path_key=self._change_path_key,
         )
-        results = await self._apply_existing(buckets)
+        results: list[dict] = []
+        pending_buckets = buckets
+        recover_source_changes = bool(self.context.get("retry_source_changes_before_watch"))
+        for recovery_pass in range(_INITIAL_SCAN_RECOVERY_PASSES + 1):
+            batch_results = await self._apply_existing(pending_buckets)
+            retryable: list[dict] = []
+            for result in batch_results:
+                if result.get(_SOURCE_CHANGE_RETRY_MARKER, False):
+                    retryable.append(result)
+            if not retryable:
+                for result in batch_results:
+                    result.pop(_SOURCE_CHANGE_RETRY_MARKER, None)
+                results.extend(batch_results)
+                break
+            if not recover_source_changes:
+                for result in batch_results:
+                    result.pop(_SOURCE_CHANGE_RETRY_MARKER, None)
+                results.extend(batch_results)
+                break
+            results.extend(result for result in batch_results if not result.pop(_SOURCE_CHANGE_RETRY_MARKER, False))
+            pending_buckets = bucket_changes(
+                [{"change": result["change"], "path": result["path"]} for result in retryable],
+                path_exists=lambda p: self._to_abs_path(p).is_file(),
+                path_key=self._change_path_key,
+            )
+            results.extend(await self._apply_deleted(pending_buckets[Change.deleted]))
+            if recovery_pass == _INITIAL_SCAN_RECOVERY_PASSES:
+                pending_paths = pending_buckets[Change.added] + pending_buckets[Change.modified]
+                if pending_paths:
+                    paths = ", ".join(pending_paths)
+                    raise FileChangedDuringChunkingError(
+                        f"Initial scan could not index stable source files after bounded retries: {paths}",
+                    )
+                break
         results.extend(await self._apply_deleted(buckets[Change.deleted]))
         if persist and results:
             await self.dump_target()
@@ -174,6 +210,7 @@ class ChangeApplyStep(BaseStep):
                         "path": path,
                         "success": False,
                         "error": "not a file",
+                        _SOURCE_CHANGE_RETRY_MARKER: True,
                     },
                 )
                 return None
@@ -197,6 +234,18 @@ class ChangeApplyStep(BaseStep):
                 )
                 return None
             return abs_path, size_bytes
+        except FileNotFoundError as e:
+            self.logger.debug(f"Source disappeared while processing {path}; reconciling its current state")
+            results.append(
+                {
+                    "change": change.name,
+                    "path": path,
+                    "success": False,
+                    "error": str(e),
+                    _SOURCE_CHANGE_RETRY_MARKER: True,
+                },
+            )
+            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)})
@@ -224,7 +273,30 @@ class ChangeApplyStep(BaseStep):
                     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)})
+                results.append(
+                    {
+                        "change": change.name,
+                        "path": path,
+                        "success": False,
+                        "error": str(e),
+                        _SOURCE_CHANGE_RETRY_MARKER: True,
+                    },
+                )
+                return None
+            except FileNotFoundError as e:
+                if attempt < self.source_change_retries:
+                    self.logger.debug(f"Source disappeared while processing {path}; retrying")
+                    continue
+                self.logger.debug(f"Source disappeared while processing {path}; reconciling its current state")
+                results.append(
+                    {
+                        "change": change.name,
+                        "path": path,
+                        "success": False,
+                        "error": str(e),
+                        _SOURCE_CHANGE_RETRY_MARKER: True,
+                    },
+                )
                 return None
             except Exception as e:
                 self.logger.exception(f"Failed to process {path}")
@@ -411,6 +483,7 @@ class UpdateIndexStep(ChangeApplyStep):
         async with self.file_store._maintenance_guard():  # pylint: disable=protected-access
             current_items: list[tuple[FileNode, list[FileChunk]]] = []
             current_paths: list[str] = []
+            item: tuple[FileNode, list[FileChunk]] | None = None
             for item, path in zip(items, ok_paths):
                 node = item[0]
                 try:
@@ -428,6 +501,18 @@ class UpdateIndexStep(ChangeApplyStep):
                             "path": path,
                             "success": False,
                             "error": str(exc),
+                            _SOURCE_CHANGE_RETRY_MARKER: True,
+                        }
+                except FileNotFoundError 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),
+                            _SOURCE_CHANGE_RETRY_MARKER: True,
                         }
                 except (OSError, RuntimeError) as exc:
                     result_by_path[path] = {
@@ -448,6 +533,10 @@ class UpdateIndexStep(ChangeApplyStep):
                         for result in await super()._try_upsert(change, current_items, current_paths)
                     },
                 )
+            if item is not None:
+                del item
+            items.clear()
+            current_items.clear()
 
         retry_results: list[dict] = []
         for path in retry_paths:

tests/unit/test_background_steps.py

@@ -54,6 +54,7 @@ from reme.steps.index import (
 )
 from reme.steps.index._change_batch import bucket_changes
 from reme.steps.index._watch_rules import WatchRule, build_watch_rules, collect_existing, match_file
+from reme.utils.file_utils import FileChangedDuringChunkingError
 
 warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba")
 warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources")
@@ -1014,6 +1015,83 @@ def test_update_index_handles_file_removed_before_stat():
     asyncio.run(run())
 
 
+def test_update_index_initial_scan_removes_file_disappeared_during_parse():
+    """The initial scan removes stale indexed content when a changed source vanishes."""
+
+    async def run():
+        with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
+            cwd = Path.cwd()
+            source = write_file(cwd / "daily" / "a.md", "old contents")
+            store = LocalFileStore(name="default", embedding_store="")
+            chunker = DefaultFileChunker(supported_extensions=["md"])
+            await store.start()
+            await chunker.start()
+            try:
+                app_ctx = _make_app_context(cwd)
+                app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
+                await store.upsert([await chunker.chunk(source)])
+                source.write_text("changed before initial scan", encoding="utf-8")
+                step = UpdateIndexStep(file_store=store, persist=False, app_context=app_ctx)
+
+                async def disappear_during_parse(path: Path):
+                    path.unlink()
+                    raise FileNotFoundError("file disappeared during parse")
+
+                with patch.object(step, "build_item", side_effect=disappear_during_parse):
+                    response = await step(
+                        RuntimeContext(
+                            changes=[{"change": "modified", "path": str(source)}],
+                            retry_source_changes_before_watch=True,
+                        ),
+                    )
+
+                assert response.success is True
+                assert response.answer == [{"change": "deleted", "path": str(source), "success": True}]
+                assert await store.get_nodes() == []
+            finally:
+                await chunker.close()
+                await store.close()
+
+    asyncio.run(run())
+
+
+def test_update_index_initial_scan_fails_when_source_never_stabilizes():
+    """A watcher cannot start from an initial scan that exhausted source retries."""
+
+    async def run():
+        with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
+            cwd = Path.cwd()
+            source = write_file(cwd / "daily" / "a.md", "changing source")
+            store = LocalFileStore(name="default", embedding_store="")
+            chunker = DefaultFileChunker(supported_extensions=["md"])
+            await store.start()
+            await chunker.start()
+            try:
+                app_ctx = _make_app_context(cwd)
+                app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
+                step = UpdateIndexStep(file_store=store, persist=False, app_context=app_ctx)
+
+                async def source_keeps_changing(path: Path):
+                    raise FileChangedDuringChunkingError("source keeps changing")
+
+                with patch.object(step, "build_item", side_effect=source_keeps_changing) as build_item:
+                    with pytest.raises(FileChangedDuringChunkingError, match="after bounded retries"):
+                        await step(
+                            RuntimeContext(
+                                changes=[{"change": "added", "path": str(source)}],
+                                retry_source_changes_before_watch=True,
+                            ),
+                        )
+
+                assert build_item.call_count == 4
+                assert await store.get_nodes() == []
+            finally:
+                await chunker.close()
+                await store.close()
+
+    asyncio.run(run())
+
+
 def test_update_index_reports_memory_estimation_failure_without_aborting():
     """A missing chunker is isolated to the affected file and returned as a normal failure."""
 
@@ -1108,10 +1186,11 @@ 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")
-            parse_paused = threading.Event()
-            resume_parse = threading.Event()
+            parse_started = [threading.Event(), threading.Event()]
+            resume_parse = [threading.Event(), threading.Event()]
             watcher_started = asyncio.Event()
             chunk_calls: list[Path] = []
+            watcher_contents: list[str] = []
 
             store = LocalFileStore(name="default", embedding_store="")
             app_ctx = _make_app_context(cwd)
@@ -1147,10 +1226,11 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
             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:
-                    parse_paused.set()
-                    if not resume_parse.wait(timeout=5):
-                        raise TimeoutError("initial parse was not resumed")
+                attempt = len(chunk_calls) - 1
+                if attempt < len(parse_started):
+                    parse_started[attempt].set()
+                    if not resume_parse[attempt].wait(timeout=5):
+                        raise TimeoutError(f"initial parse {attempt + 1} was not resumed")
                 return result
 
             from reme.steps.index import watch_changes as watch_changes_module
@@ -1167,6 +1247,10 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
                 poll_delay_ms: int,
                 stop_event: asyncio.Event,
             ) -> AsyncGenerator[set[tuple[Change, str]], None]:
+                nodes = await store.get_nodes()
+                watcher_contents.append(
+                    "".join(store.file_chunks[chunk_id].text for node in nodes for chunk_id in node.chunk_ids),
+                )
                 watcher_started.set()
                 async for changes in original_awatch(
                     *paths,
@@ -1187,19 +1271,24 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
                 await store.start()
                 await chunker.start()
                 await job.start()
-                assert await asyncio.to_thread(parse_paused.wait, 5)
+                assert await asyncio.to_thread(parse_started[0].wait, 5)
+                path.write_text("contents after the first parse", encoding="utf-8")
+                resume_parse[0].set()
+                assert await asyncio.to_thread(parse_started[1].wait, 5)
                 latest = "latest contents after the initial scan"
                 path.write_text(latest, encoding="utf-8")
-                resume_parse.set()
+                resume_parse[1].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
+                assert watcher_contents == [latest]
+                assert len(chunk_calls) == 3
             finally:
-                resume_parse.set()
+                for event in resume_parse:
+                    event.set()
                 await job.close()
                 await chunker.close()
                 await store.close()

tests/unit/test_file_chunker_workers.py

@@ -2,15 +2,18 @@
 
 import asyncio
 import contextvars
+import gc
 import json
 import os
 import threading
+import weakref
 from pathlib import Path
 from typing import Callable
 from unittest.mock import patch
 
 import pytest
 from httpx import ASGITransport, AsyncClient
+from watchfiles import Change
 
 from reme.application import Application
 from reme.components.application_context import ApplicationContext
@@ -204,6 +207,53 @@ async def test_shutdown_rejects_queued_work_and_drains_despite_cancellation(tmp_
         assert (await chunker.chunk(path))[1][0].text == "source"
 
 
+@pytest.mark.asyncio
+async def test_shutdown_from_another_event_loop_waits_for_worker_and_closes_executor(tmp_path: Path) -> None:
+    """Worker admission and shutdown can cross event loops without stranding the executor."""
+    path = tmp_path / "source.txt"
+    path.write_text("source")
+    chunker = DefaultFileChunker()
+    entered, release = threading.Event(), threading.Event()
+    worker_errors: list[BaseException] = []
+    original = chunker._chunk_sync
+
+    def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
+        entered.set()
+        assert release.wait(5)
+        return original(file_path, rel_path)
+
+    def run_on_worker_loop() -> None:
+        try:
+            asyncio.run(chunker.chunk(path))
+        except BaseException as exc:
+            worker_errors.append(exc)
+
+    await chunker.start()
+    with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse):
+        worker_thread = threading.Thread(target=run_on_worker_loop)
+        worker_thread.start()
+        await _wait_started(entered)
+        closing = asyncio.create_task(chunker.close())
+        try:
+            for _ in range(100):
+                if chunker._worker_closing:
+                    break
+                await asyncio.sleep(0.01)
+            assert chunker._worker_closing
+            assert not closing.done()
+        finally:
+            release.set()
+        await closing
+        await asyncio.to_thread(worker_thread.join, 5)
+
+    assert not worker_thread.is_alive()
+    assert len(worker_errors) == 1
+    assert isinstance(worker_errors[0], RuntimeError)
+    assert "different event loop" not in str(worker_errors[0])
+    assert chunker._worker is None
+    assert not chunker.is_started
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("change", ["modify", "replace", "delete"])
 async def test_source_changes_do_not_return_stale_chunks(tmp_path: Path, change: str) -> None:
@@ -525,6 +575,76 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file
         assert not step._source_versions
 
 
+@pytest.mark.asyncio
+async def test_stale_batch_chunks_are_released_before_source_retry(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The retry parse does not overlap the retained stale chunk batch in memory."""
+    monkeypatch.chdir(tmp_path)
+    first_path, second_path = tmp_path / "first.txt", tmp_path / "second.txt"
+    first_path.write_text("original first")
+    second_path.write_text("original second")
+    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(first_path), await chunker.chunk(second_path)])
+        first_path.write_text("parsed first")
+        second_path.write_text("parsed second")
+        entered, release = threading.Event(), threading.Event()
+        original_chunk = chunker.chunk
+
+        async def pause_second(file_path: Path) -> tuple[FileNode, list[FileChunk]]:
+            if file_path == second_path:
+                entered.set()
+                while not release.is_set():
+                    await asyncio.sleep(0.01)
+            return await original_chunk(file_path)
+
+        step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=2, app_context=context)
+        original_build = step._try_build_item
+        stale_chunk_released: threading.Event | None = None
+
+        async def capture_and_check_stale_chunk(
+            change: Change,
+            path: str,
+            abs_path: Path,
+            size_bytes: int,
+            results: list[dict],
+        ) -> tuple[tuple[FileNode, list[FileChunk]], int] | None:
+            nonlocal stale_chunk_released
+            if Path(path) == first_path and stale_chunk_released is not None:
+                gc.collect()
+                assert stale_chunk_released.is_set()
+            built = await original_build(change, path, abs_path, size_bytes, results)
+            if Path(path) == first_path and stale_chunk_released is None and built is not None:
+                stale_chunk_released = threading.Event()
+                weakref.finalize(built[0][1][0], stale_chunk_released.set)
+            return built
+
+        monkeypatch.setattr(step, "_try_build_item", capture_and_check_stale_chunk)
+        with patch.object(chunker, "chunk", side_effect=pause_second):
+            task = asyncio.create_task(
+                step(
+                    RuntimeContext(
+                        changes=[
+                            {"change": "modified", "path": str(first_path)},
+                            {"change": "modified", "path": str(second_path)},
+                        ],
+                    ),
+                ),
+            )
+            await _wait_started(entered)
+            first_path.write_text("latest first")
+            release.set()
+            response = await task
+
+        assert response.success is True
+        assert stale_chunk_released is not None and stale_chunk_released.is_set()
+        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(