From 4f4c25b32aff8a1be25bdd27d1ec76d679ee28eb Mon Sep 17 00:00:00 2001 From: mikemikimike <186855910+mikemikimike@users.noreply.github.com> Date: Sat, 10 Oct 2026 09:02:51 +0800 Subject: [PATCH 1/6] fix(index): move file parsing off the service event loop Signed-off-by: mikemikimike <186855910+mikemikimike@users.noreply.github.com> --- .../file_chunker/base_file_chunker.py | 99 ++++ .../file_chunker/default_file_chunker.py | 16 +- .../file_chunker/json_file_chunker.py | 5 +- .../file_chunker/jsonl_file_chunker.py | 5 +- .../file_chunker/markdown_file_chunker.py | 13 +- reme/steps/index/update_changes.py | 47 +- tests/unit/fixtures/file_chunker_outputs.json | 412 +++++++++++++ tests/unit/test_file_chunker_workers.py | 547 ++++++++++++++++++ 8 files changed, 1126 insertions(+), 18 deletions(-) create mode 100644 tests/unit/fixtures/file_chunker_outputs.json create mode 100644 tests/unit/test_file_chunker_workers.py diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 49312989..bb477c95 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -1,12 +1,36 @@ """Abstract base for file chunkers.""" +import asyncio +import contextvars +import os from abc import abstractmethod +from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from typing import Callable, TypeVar from ..base_component import BaseComponent from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode +WorkResult = TypeVar("WorkResult") + + +async def _finish_work(future: asyncio.Future[WorkResult]) -> WorkResult: + """Defer cancellation until submitted work finishes, then discard its result.""" + cancellation: asyncio.CancelledError | None = None + while not future.done(): + try: + await asyncio.shield(future) + except asyncio.CancelledError as exc: + cancellation = exc + except Exception: + break + if cancellation is not None: + if not future.cancelled(): + future.exception() + raise cancellation + return future.result() + class BaseFileChunker(BaseComponent): """Abstract base for file chunkers. Subclasses implement `chunk`.""" @@ -16,6 +40,81 @@ class BaseFileChunker(BaseComponent): def __init__(self, supported_extensions: list[str] | None = None, **kwargs): super().__init__(**kwargs) self.supported_extensions: list[str] = supported_extensions or [] + self._worker: ThreadPoolExecutor | None = None + self._worker_lock = asyncio.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 + + async def _close(self) -> None: + 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) + self._worker = None + + async def _chunk_in_worker( + self, + path: str | Path, + 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 + 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) + 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()) + return result + finally: + if temporary_worker: + worker.shutdown(wait=True) + + @staticmethod + def _chunk_unchanged_file( + path: Path, + rel_path: str, + chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]], + ) -> tuple[tuple[FileNode, list[FileChunk]], os.stat_result]: + before = path.stat() + result = chunk(path, rel_path) + after = path.stat() + BaseFileChunker._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]]: diff --git a/reme/components/file_chunker/default_file_chunker.py b/reme/components/file_chunker/default_file_chunker.py index 615e1171..89e816bb 100644 --- a/reme/components/file_chunker/default_file_chunker.py +++ b/reme/components/file_chunker/default_file_chunker.py @@ -4,8 +4,6 @@ from bisect import bisect_right from pathlib import Path from typing import Literal -import aiofiles - from .base_file_chunker import BaseFileChunker from ..component_registry import R from ...schema import FileChunk, FileNode @@ -37,10 +35,9 @@ class DefaultFileChunker(BaseFileChunker): self.chunk_byte_size = max(100, chunk_byte_size) self.overlap_byte_size = max(4, overlap_byte_size) - async def _read_text_for_indexing(self, file_path: Path) -> str: + def _read_text_for_indexing(self, file_path: Path) -> str: """Decode text for a derived index without modifying the source file.""" - async with aiofiles.open(file_path, "rb") as f: - data = await f.read() + data = file_path.read_bytes() try: text = data.decode(self.encoding) except UnicodeDecodeError as exc: @@ -61,11 +58,12 @@ class DefaultFileChunker(BaseFileChunker): return text.replace("\r\n", "\n").replace("\r", "\n") async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: - file_path = Path(path) - stat = file_path.stat() - rel_path = self.to_workspace_relative(path) + return await self._chunk_in_worker(path, self._chunk_sync) - text = await self._read_text_for_indexing(file_path) + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + stat = file_path.stat() + + text = self._read_text_for_indexing(file_path) if not text: return FileNode(path=rel_path, st_mtime=stat.st_mtime), [] diff --git a/reme/components/file_chunker/json_file_chunker.py b/reme/components/file_chunker/json_file_chunker.py index 78f8e617..e2c97686 100644 --- a/reme/components/file_chunker/json_file_chunker.py +++ b/reme/components/file_chunker/json_file_chunker.py @@ -113,9 +113,10 @@ class JsonFileChunker(BaseFileChunker): async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: """Read and chunk a JSON file at *path*.""" - file_path = Path(path) + return await self._chunk_in_worker(path, self._chunk_sync) + + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: stat = file_path.stat() - rel_path = self.to_workspace_relative(path) raw_text = file_path.read_text(encoding=self.encoding) if not raw_text.strip(): diff --git a/reme/components/file_chunker/jsonl_file_chunker.py b/reme/components/file_chunker/jsonl_file_chunker.py index 9417282c..0326b299 100644 --- a/reme/components/file_chunker/jsonl_file_chunker.py +++ b/reme/components/file_chunker/jsonl_file_chunker.py @@ -124,9 +124,10 @@ class JsonlFileChunker(BaseFileChunker): async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: """Read and chunk a JSONL file at *path*.""" - file_path = Path(path) + return await self._chunk_in_worker(path, self._chunk_sync) + + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: stat = file_path.stat() - rel_path = self.to_workspace_relative(path) text = file_path.read_text(encoding=self.encoding) if not text.strip(): diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index 06a07793..a43e8879 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -16,6 +16,7 @@ the single source of truth for ``[[...]]`` syntax. """ import re +import threading from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -39,6 +40,9 @@ _ATX_HEADING_RE = re.compile(r"^ {0,3}#{1,6}(?:[ \t]+|$)") _SETEXT_HEADING_RE = re.compile(r"^ {0,3}(?:=+|-+)[ \t]*$") _FENCE_RE = re.compile(r"^ {0,3}(`{3,}|~{3,})") +# MarkdownRenderer changes mistletoe's process-wide token registry. +_RENDERER_LOCK = threading.Lock() + @dataclass class MdNode: @@ -124,9 +128,10 @@ class MarkdownFileChunker(DefaultFileChunker): self.include_frontmatter_keys_in_metadata = list(include_frontmatter_keys_in_metadata or []) async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: - file_path = Path(path) - rel_path = self.to_workspace_relative(path) - text = await self._read_text_for_indexing(file_path) + return await self._chunk_in_worker(path, self._chunk_sync) + + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + text = self._read_text_for_indexing(file_path) front_matter, content, line_offset = self._parse_front_matter(text) chunks: list[FileChunk] = [] @@ -142,7 +147,7 @@ class MarkdownFileChunker(DefaultFileChunker): from mistletoe.markdown_renderer import MarkdownRenderer from mistletoe.block_token import Document - with MarkdownRenderer() as renderer: + with _RENDERER_LOCK, MarkdownRenderer() as renderer: tree = self._build_tree(Document(content), renderer, line_offset=line_offset) chunks = self._chunk_node(tree, (), rel_path, renderer) if self.include_frontmatter_in_metadata: diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 6ce461f0..3dbcce4c 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -2,6 +2,7 @@ import asyncio import math +import os from abc import abstractmethod from pathlib import Path from typing import Any, Callable @@ -46,6 +47,7 @@ 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: @@ -323,8 +325,51 @@ class UpdateIndexStep(ChangeApplyStep): target_name = "file_store" reads_file_content = True + async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]: + try: + return await super()._apply_existing(buckets) + finally: + # Steps are invocation-scoped; cancelled batches must release their snapshots too. + self._source_versions.clear() + async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]: - return await self.chunk_file(path) + 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 + return item + + async def _try_upsert( + self, + change: Change, + items: list[tuple[FileNode, list[FileChunk]]], + ok_paths: list[str], + ) -> list[dict]: + # Recheck after waiting for the same guard used by every store mutation. + 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 + self._to_abs_path(node.path), + self._source_versions[id(node)], + self._to_abs_path(node.path).stat(), + ) + except (OSError, RuntimeError) as exc: + results.append({"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)) + 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] async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None: await self.file_store.upsert(items) diff --git a/tests/unit/fixtures/file_chunker_outputs.json b/tests/unit/fixtures/file_chunker_outputs.json new file mode 100644 index 00000000..19ad5a49 --- /dev/null +++ b/tests/unit/fixtures/file_chunker_outputs.json @@ -0,0 +1,412 @@ +{ + "plain.txt": { + "node": { + "path": "plain.txt", + "links": [], + "chunk_ids": [ + "27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8", + "text": "first\nabc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 1, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23", + "text": "c abc ", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + } + ] + }, + "note.md": { + "node": { + "path": "note.md", + "links": [ + { + "source_path": "note.md", + "target_path": "other.md", + "target_anchor": null + } + ], + "chunk_ids": [ + "c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3", + "03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c", + "f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7" + ], + "front_matter": { + "name": "", + "description": "", + "title": "Example", + "tags": [ + "one", + "two" + ] + } + }, + "chunks": [ + { + "id": "c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3", + "text": "# First\n\nbody [[other.md]]", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 5, + "end_line": 7, + "scores": {} + }, + { + "id": "03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c", + "text": "# First\n\n## Second", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 9, + "end_line": 9, + "scores": {} + }, + { + "id": "f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7", + "text": "text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 11, + "end_line": 11, + "scores": {} + } + ] + }, + "fallback.md": { + "node": { + "path": "fallback.md", + "links": [], + "chunk_ids": [ + "22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab", + "text": "# First\ntext text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 1, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19", + "text": "xt text ", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + } + ] + }, + "data.json": { + "node": { + "path": "data.json", + "links": [], + "chunk_ids": [ + "cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1", + "562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76", + "147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799", + "79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62", + "352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba", + "04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 0}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 4, + "end_line": 5, + "scores": {} + }, + { + "id": "562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 1}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 8, + "end_line": 9, + "scores": {} + }, + { + "id": "147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 2}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 12, + "end_line": 13, + "scores": {} + }, + { + "id": "79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 3}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 16, + "end_line": 17, + "scores": {} + }, + { + "id": "352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 4}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 20, + "end_line": 21, + "scores": {} + }, + { + "id": "04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 5}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 24, + "end_line": 25, + "scores": {} + } + ] + }, + "bad.json": { + "node": { + "path": "bad.json", + "links": [], + "chunk_ids": [ + "e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3", + "text": "{invalid\ntext\n", + "embedding": null, + "metadata": {}, + "path": "bad.json", + "start_line": 1, + "end_line": 3, + "scores": {} + } + ] + }, + "events.jsonl": { + "node": { + "path": "events.jsonl", + "links": [], + "chunk_ids": [ + "fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f", + "7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6", + "48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b", + "116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634", + "c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93", + "489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 1, + "end_line": 3, + "scores": {} + }, + { + "id": "7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 3, + "end_line": 5, + "scores": {} + }, + { + "id": "48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 5, + "end_line": 7, + "scores": {} + }, + { + "id": "116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 7, + "end_line": 9, + "scores": {} + }, + { + "id": "c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 9, + "end_line": 11, + "scores": {} + }, + { + "id": "489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 11, + "end_line": 12, + "scores": {} + } + ] + } +} diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py new file mode 100644 index 00000000..7047b9c5 --- /dev/null +++ b/tests/unit/test_file_chunker_workers.py @@ -0,0 +1,547 @@ +"""Worker admission, output compatibility, and real indexing/search regressions.""" + +import asyncio +import contextvars +import json +import os +import threading +from pathlib import Path +from typing import Callable +from unittest.mock import patch + +import pytest +from httpx import ASGITransport, AsyncClient + +from reme.application import Application +from reme.components.application_context import ApplicationContext +from reme.components.file_chunker import DefaultFileChunker, JsonFileChunker, JsonlFileChunker, MarkdownFileChunker +from reme.components.file_store import LocalFileStore +from reme.components.job import BaseJob, StreamJob +from reme.components.runtime_context import RuntimeContext +from reme.components.service import HttpService +from reme.enumeration import ChunkEnum, ComponentEnum +from reme.schema import FileChunk, FileNode, Response, StreamChunk +from reme.steps.index import UpdateIndexStep + +# Worker tests intentionally inspect admission and teardown at the internal boundary. +# pylint: disable=protected-access + + +class _StoreSearchJob(BaseJob): + """Use the actual keyword store behind the HTTP route and MCP tool.""" + + def __init__(self, store: LocalFileStore) -> None: + super().__init__(name="search", parameters={"type": "object", "properties": {"query": {"type": "string"}}}) + self.store = store + + async def __call__(self, query: str, **kwargs) -> Response: + del kwargs + chunks = await self.store.keyword_search(query, limit=10, search_filter={}) + return Response(answer=[chunk.text for chunk in chunks]) + + +class _StoreSearchStreamJob(StreamJob): + """Stream real keyword matches through the existing SSE endpoint.""" + + def __init__(self, store: LocalFileStore) -> None: + super().__init__(name="search_stream") + self.store = store + + async def __call__(self, stream_queue: asyncio.Queue[StreamChunk], query: str, **kwargs) -> None: + del kwargs + chunks = await self.store.keyword_search(query, limit=10, search_filter={}) + for chunk in chunks: + await stream_queue.put(StreamChunk(chunk=chunk.text)) + await asyncio.sleep(0) + await stream_queue.put(StreamChunk(chunk_type=ChunkEnum.DONE, done=True)) + + +@pytest.mark.asyncio +async def test_outputs_equal_pre_worker_fixtures(tmp_path: Path) -> None: + """Compare every serialized field, including hashes, links, metadata and ranges.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + cases = [ + (DefaultFileChunker(chunk_byte_size=100), "plain.txt", "first\r\n" + "abc " * 100), + ( + MarkdownFileChunker(chunk_byte_size=100, include_frontmatter_in_metadata=True), + "note.md", + "---\ntitle: Example\ntags: [one, two]\n---\n# First\n\nbody [[other.md]]\n\n## Second\n\n" + "text " * 70, + ), + (MarkdownFileChunker(chunk_byte_size=100, max_ast_sections=0), "fallback.md", "# First\n" + "text " * 100), + ( + JsonFileChunker(chunk_chars=256), + "data.json", + json.dumps({"items": [{"name": "abc" * 35, "n": i} for i in range(6)]}, indent=2), + ), + (JsonFileChunker(chunk_chars=256), "bad.json", "{invalid\ntext\n"), + (JsonlFileChunker(max_chars=64, max_overlap_chars=30), "events.jsonl", '{"value": "alpha"}\n' * 12), + ] + expected = json.loads((Path(__file__).parent / "fixtures" / "file_chunker_outputs.json").read_text()) + for chunker, name, text in cases: + chunker.app_context = context + path = tmp_path / name + path.write_bytes(text.encode()) + node, chunks = await chunker.chunk(path) + output = {"node": node.model_dump(mode="json"), "chunks": [c.model_dump(mode="json") for c in chunks]} + assert output["node"].pop("st_mtime") == path.stat().st_mtime + assert output == expected[name] + + +async def _wait_started(event: threading.Event) -> None: + for _ in range(500): + if event.is_set(): + return + await asyncio.sleep(0.01) + raise AssertionError("worker did not start") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunker_class", [DefaultFileChunker, MarkdownFileChunker, JsonFileChunker, JsonlFileChunker]) +async def test_parsing_leaves_event_loop_responsive( + tmp_path: Path, + chunker_class: type[DefaultFileChunker | MarkdownFileChunker | JsonFileChunker | JsonlFileChunker], +) -> None: + """A parser can block while unrelated async requests keep progressing.""" + chunker = chunker_class() + path = tmp_path / "source" + path.write_text("{}") + entered, release = threading.Event(), threading.Event() + main_thread = threading.get_ident() + original = chunker._chunk_sync + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + assert threading.get_ident() != main_thread + entered.set() + assert release.wait(5) + return original(file_path, rel_path) + + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + task = asyncio.create_task(chunker.chunk(path)) + try: + await _wait_started(entered) + for _ in range(10): + await asyncio.sleep(0) + assert not task.done() + finally: + release.set() + await task + assert chunker._worker is None + + +@pytest.mark.asyncio +async def test_cancellation_drains_worker_and_preserves_admission(tmp_path: Path) -> None: + """Repeated cancellation cannot free admission while parsing is still running.""" + path = tmp_path / "source.txt" + path.write_text("source") + chunker = DefaultFileChunker() + entered, release = threading.Event(), threading.Event() + calls = [] + original = chunker._chunk_sync + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + calls.append(rel_path) + entered.set() + assert release.wait(5) + return original(file_path, rel_path) + + async with chunker: + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + first = asyncio.create_task(chunker.chunk(path)) + await _wait_started(entered) + first.cancel() + second = asyncio.create_task(chunker.chunk(path)) + await asyncio.sleep(0.01) + first.cancel() + await asyncio.sleep(0.01) + try: + assert not first.done() + assert len(calls) == 1 + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await first + await second + assert chunker._worker is None + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_shutdown_rejects_queued_work_and_drains_despite_cancellation(tmp_path: Path) -> None: + """Closing waits for running work, rejects queued parsing, and permits restart.""" + path = tmp_path / "source.txt" + path.write_text("source") + chunker = DefaultFileChunker() + entered, release = threading.Event(), threading.Event() + 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) + + await chunker.start() + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + running = asyncio.create_task(chunker.chunk(path)) + await _wait_started(entered) + queued = asyncio.create_task(chunker.chunk(path)) + await asyncio.sleep(0) + closing = asyncio.create_task(chunker.close()) + await asyncio.sleep(0.01) + closing.cancel() + try: + assert not closing.done() + finally: + release.set() + with pytest.raises(RuntimeError, match="closed"): + await running + with pytest.raises(RuntimeError, match="closed"): + await queued + with pytest.raises(asyncio.CancelledError): + await closing + assert chunker._worker is None + assert not chunker.is_started + async with chunker: + assert (await chunker.chunk(path))[1][0].text == "source" + + +@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: + """Changes after reading cannot be mistaken for successfully parsed current content.""" + path = tmp_path / "source.txt" + path.write_text("old source") + chunker = DefaultFileChunker() + original = chunker._chunk_sync + entered, release = threading.Event(), threading.Event() + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + result = original(file_path, rel_path) + entered.set() + assert release.wait(5) + return result + + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + task = asyncio.create_task(chunker.chunk(path)) + await _wait_started(entered) + try: + if change == "modify": + path.write_text("new source with different size") + elif change == "replace": + replacement = tmp_path / "replacement" + replacement.write_text("replacement source") + replacement.replace(path) + else: + path.unlink() + finally: + release.set() + with pytest.raises((RuntimeError, FileNotFoundError)): + await task + + +@pytest.mark.asyncio +async def test_change_before_event_loop_receives_result_is_rejected(tmp_path: Path) -> None: + """Validate the worker snapshot again when its result reaches the calling loop.""" + path = tmp_path / "source.txt" + path.write_text("old source") + chunker = DefaultFileChunker() + original = chunker._chunk_unchanged_file + loop = asyncio.get_running_loop() + changed = threading.Event() + + def change_on_loop() -> None: + path.write_text("changed after the worker snapshot") + changed.set() + + def change_after_parse( + file_path: Path, + rel_path: str, + chunk: Callable[[Path, str], tuple[FileNode, list[FileChunk]]], + ) -> tuple[tuple[FileNode, list[FileChunk]], os.stat_result]: + result = original(file_path, rel_path, chunk) + loop.call_soon_threadsafe(change_on_loop) + assert changed.wait(5) + return result + + with patch.object(chunker, "_chunk_unchanged_file", side_effect=change_after_parse): + with pytest.raises(RuntimeError, match="File changed"): + await chunker.chunk(path) + + +@pytest.mark.asyncio +async def test_queued_cancellation_does_not_submit_and_context_is_preserved(tmp_path: Path) -> None: + """A cancelled waiter uses no worker slot; caller context is copied into the worker.""" + path = tmp_path / "source.txt" + path.write_text("source") + chunker = DefaultFileChunker() + entered, release = threading.Event(), threading.Event() + marker = contextvars.ContextVar("chunk_test_marker", default="unset") + original = chunker._chunk_sync + calls = [] + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + calls.append(marker.get()) + entered.set() + assert release.wait(5) + return original(file_path, rel_path) + + token = marker.set("caller") + try: + async with chunker: + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + first = asyncio.create_task(chunker.chunk(path)) + await _wait_started(entered) + queued = asyncio.create_task(chunker.chunk(path)) + await asyncio.sleep(0) + queued.cancel() + try: + with pytest.raises(asyncio.CancelledError): + await queued + finally: + release.set() + await first + finally: + marker.reset(token) + assert calls == ["caller"] + + +@pytest.mark.asyncio +async def test_markdown_renderer_is_serialized_across_components(tmp_path: Path) -> None: + """Independent chunker instances cannot overlap renderer registry changes.""" + from mistletoe.markdown_renderer import MarkdownRenderer + + path = tmp_path / "source.md" + path.write_text("# Title\n\nparagraph **bold** [[other.md]]\n") + entered, release = threading.Event(), threading.Event() + calls = [] + original_enter = MarkdownRenderer.__enter__ + + def blocking_enter(renderer: MarkdownRenderer) -> MarkdownRenderer: + calls.append(threading.get_ident()) + entered.set() + assert release.wait(5) + return original_enter(renderer) + + with patch.object(MarkdownRenderer, "__enter__", blocking_enter): + first = asyncio.create_task(MarkdownFileChunker().chunk(path)) + await _wait_started(entered) + second = asyncio.create_task(MarkdownFileChunker().chunk(path)) + try: + await asyncio.sleep(0.05) + assert len(calls) == 1 + finally: + release.set() + first_output, second_output = await asyncio.gather(first, second) + assert first_output == second_output + assert len(calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ["complete", "cancel", "change", "error"]) +async def test_existing_content_searchable_during_indexing( + tmp_path: Path, + outcome: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exercise UpdateIndexStep and LocalFileStore while parsing is suspended.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "source.txt" + path.write_text("old searchable content") + context = ApplicationContext(workspace_dir=str(tmp_path)) + chunker = DefaultFileChunker(app_context=context) + context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}} + store = LocalFileStore(embedding_store="") + async with store, chunker: + await store.upsert([await chunker.chunk(path)]) + path.write_text("new searchable content") + entered, release = threading.Event(), threading.Event() + original = chunker._chunk_sync + + def blocking_parse(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + result = original(file_path, rel_path) + entered.set() + assert release.wait(5) + if outcome == "error": + raise ValueError("parser failed") + return result + + step = UpdateIndexStep(file_store=store, persist=False, app_context=context) + app = Application( + workspace_dir=str(tmp_path / "service"), + service={"backend": "http", "web_enabled": False}, + components={}, + jobs={}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + ) + service = app.context.service + assert isinstance(service, HttpService) + service.build_service(app) + service.add_job(_StoreSearchJob(store)) + service.add_job(_StoreSearchStreamJob(store)) + with patch.object(chunker, "_chunk_sync", side_effect=blocking_parse): + updating = asyncio.create_task(step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}]))) + await _wait_started(entered) + try: + nodes = await store.get_nodes() + assert len(nodes) == 1 + matches = await store.keyword_search("searchable", limit=10, search_filter={}) + assert matches[0].text == "old searchable content" + async with AsyncClient(transport=ASGITransport(app=service.service), base_url="http://test") as client: + response = await client.post("/search", json={"query": "searchable"}) + assert response.status_code == 200 + assert response.json()["answer"] == ["old searchable content"] + streaming = await client.post("/search_stream", json={"query": "searchable"}) + assert streaming.status_code == 200 + assert "old searchable content" in streaming.text + assert streaming.text.endswith("data:[DONE]\n\n") + assert service.mcp_server is not None + tool = await service.mcp_server.get_tool("search") + assert tool is not None + result = await tool.run({"query": "searchable"}) + assert "old searchable content" in str(result.content) + if outcome == "cancel": + updating.cancel() + elif outcome == "change": + path.write_text("latest source content") + finally: + release.set() + if outcome == "cancel": + with pytest.raises(asyncio.CancelledError): + await updating + else: + response = await updating + assert response.success == (outcome == "complete") + 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": + 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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["modify", "delete", "replace", "cancel"]) +async def test_changed_batched_file_is_not_published_and_valid_files_update( + tmp_path: Path, + change: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An earlier parsed file can change while the batch waits for the next worker.""" + 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("updated first") + second_path.write_text("updated second") + entered, release = threading.Event(), threading.Event() + original = chunker.chunk + + async def blocking_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(file_path) + + changes = [{"change": "modified", "path": str(path)} for path in (first_path, second_path)] + step = UpdateIndexStep(file_store=store, persist=False, app_context=context) + with patch.object(chunker, "chunk", side_effect=blocking_second): + task = asyncio.create_task(step(RuntimeContext(changes=changes))) + await _wait_started(entered) + try: + if change == "modify": + first_path.write_text("latest first") + elif change == "delete": + first_path.unlink() + elif change == "cancel": + task.cancel() + else: + replacement = tmp_path / "replacement" + replacement.write_text("replacement first") + replacement.replace(first_path) + finally: + release.set() + if change == "cancel": + with pytest.raises(asyncio.CancelledError): + await task + else: + response = await task + assert response.success is False + assert [result["success"] for result in response.answer] == [False, True] + 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" + assert store.file_chunks[nodes["second.txt"].chunk_ids[0]].text == ( + "original second" if change == "cancel" else "updated second" + ) + recovered = await step( + RuntimeContext( + changes=[{"change": "deleted" if change == "delete" else "modified", "path": str(first_path)}], + ), + ) + assert recovered.success is True + nodes = {node.path: node for node in await store.get_nodes()} + if change == "delete": + assert "first.txt" not in nodes + else: + assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == first_path.read_text() + 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( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + outcome: str, +) -> None: + """A delayed mutation lock cannot admit a result that became stale while waiting.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "source.txt" + path.write_text("original 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: + await store.upsert([await chunker.chunk(path)]) + path.write_text("updated source") + built = asyncio.Event() + step = UpdateIndexStep(file_store=store, persist=False, app_context=context) + original = step.build_item + + async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk]]: + item = await original(file_path) + built.set() + return item + + with patch.object(step, "build_item", side_effect=record_build): + async with store._maintenance_guard(): + task = asyncio.create_task(step(RuntimeContext(changes=[{"change": "modified", "path": str(path)}]))) + await asyncio.wait_for(built.wait(), 5) + await asyncio.sleep(0) + if outcome == "cancel": + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + path.write_text("latest source") + if outcome == "change": + response = await task + assert response.success is False + 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() + assert not step._source_versions From 47c2b6a2d615d3fa2d919cdec6771b64f6aeaf2b Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 13 Oct 2026 15:38:03 +0800 Subject: [PATCH 2/6] fix(index): retry changed source during indexing --- .../file_chunker/base_file_chunker.py | 11 +- reme/steps/index/update_changes.py | 114 ++++++++++++++--- reme/utils/file_utils.py | 15 +++ tests/unit/test_background_steps.py | 118 +++++++++++++++++- tests/unit/test_file_chunker_workers.py | 60 ++++++--- 5 files changed, 274 insertions(+), 44 deletions(-) create mode 100644 reme/utils/file_utils.py 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 eff113ff..5365fc02 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, JsonFileChunker 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 @@ -1097,6 +1101,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.""" @@ -1402,12 +1516,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 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 3/6] 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( From c54ce03653f9a443b35671fa381288a75c702240 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Wed, 14 Oct 2026 21:02:55 +0800 Subject: [PATCH 4/6] test(index): cover initial scan source race --- tests/unit/test_background_steps.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index f5987908..ef0bb30a 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -14,6 +14,7 @@ import asyncio import datetime import os import tempfile +import threading import warnings import weakref from collections.abc import AsyncGenerator, Callable @@ -1107,7 +1108,8 @@ 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") - item_built = asyncio.Event() + parse_paused = threading.Event() + resume_parse = threading.Event() watcher_started = asyncio.Event() chunk_calls: list[Path] = [] @@ -1141,18 +1143,16 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc ) original_chunk = chunker._chunk_sync - original_build_item = UpdateIndexStep.build_item 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") 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 @@ -1181,17 +1181,16 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc yield changes 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() + assert await asyncio.to_thread(parse_paused.wait, 5) latest = "latest contents after the initial scan" - async with store._maintenance_guard(): - await job.start() - await asyncio.wait_for(item_built.wait(), timeout=5) - path.write_text(latest, encoding="utf-8") + path.write_text(latest, encoding="utf-8") + resume_parse.set() await asyncio.wait_for(watcher_started.wait(), timeout=5) nodes = await store.get_nodes() @@ -1200,6 +1199,7 @@ 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: + resume_parse.set() await job.close() await chunker.close() await store.close() From ae01c5eb9c80b8bf2e7e7c5203b52358482a3e7c Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Wed, 21 Oct 2026 05:09:47 +0800 Subject: [PATCH 5/6] fix(index): recover unstable initial scans --- .../file_chunker/base_file_chunker.py | 107 ++++++++++------ reme/steps/index/init_changes.py | 6 +- reme/steps/index/update_changes.py | 93 +++++++++++++- tests/unit/test_background_steps.py | 109 ++++++++++++++-- tests/unit/test_file_chunker_workers.py | 120 ++++++++++++++++++ 5 files changed, 386 insertions(+), 49 deletions(-) diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 5f4c9081..a4ad0d57 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/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( diff --git a/reme/steps/index/init_changes.py b/reme/steps/index/init_changes.py index 40711637..fe52e25e 100644 --- a/reme/steps/index/init_changes.py +++ b/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 diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 26492e5a..cdb904eb 100644 --- a/reme/steps/index/update_changes.py +++ b/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: diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index ef0bb30a..126dc8f7 100644 --- a/tests/unit/test_background_steps.py +++ b/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() diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index 47c04e28..be64e117 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/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( From 970b7a3b461b97e15070a6ea77d9c8a9f547a2b3 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Thu, 22 Oct 2026 02:12:54 +0800 Subject: [PATCH 6/6] Fix initial scan retry scope and stale batches --- reme/steps/index/init_changes.py | 16 +++++--- reme/steps/index/update_changes.py | 2 +- tests/unit/test_background_steps.py | 50 ++++++++++++++--------- tests/unit/test_file_chunker_workers.py | 54 +++++++++++++------------ 4 files changed, 71 insertions(+), 51 deletions(-) diff --git a/reme/steps/index/init_changes.py b/reme/steps/index/init_changes.py index fe52e25e..365305d3 100644 --- a/reme/steps/index/init_changes.py +++ b/reme/steps/index/init_changes.py @@ -73,11 +73,17 @@ 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, - retry_source_changes_before_watch=True, - ) + retry_key = "retry_source_changes_before_watch" + had_retry_value = retry_key in self.context + previous_retry_value = self.context.get(retry_key) + self.context[retry_key] = True + try: + await self.dispatch_steps(self.dispatch_step_specs, changes=changes) + finally: + if had_retry_value: + self.context[retry_key] = previous_retry_value + else: + del self.context[retry_key] else: self.logger.info(f"[{self.name}] {self.monitor_type}:{self.monitor_name} is up to date") self.context.response.metadata["counts"] = counts diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index cdb904eb..0be08e6c 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -189,13 +189,13 @@ class ChangeApplyStep(BaseStep): memory_budget = self._batch_memory_budget() items.append(item) + del item ok_paths.append(path) estimated_bytes += item_bytes if self._batch_is_full(len(items), estimated_bytes, 0, memory_budget): await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results) estimated_bytes = 0 memory_budget = self._batch_memory_budget() - del item if items: await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results) return results diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 126dc8f7..b70acd40 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -1150,6 +1150,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches(): first = await init_step(ctx) assert first.metadata["counts"] == {"added": 2, "modified": 0, "deleted": 0} + assert "retry_source_changes_before_watch" not in ctx nodes = {n.path: n for n in await fs.get_nodes()} assert set(nodes) == {"daily/a.md", "digest/report.md"} assert all(nodes[p].chunk_ids for p in nodes) @@ -1161,6 +1162,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches(): second = await init_step(ctx) assert second.metadata["counts"] == {"added": 1, "modified": 1, "deleted": 1} + assert "retry_source_changes_before_watch" not in ctx assert {(c["change"], Path(c["path"]).name) for c in ctx["changes"]} == { ("modified", "a.md"), ("deleted", "report.md"), @@ -1186,9 +1188,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_started = [threading.Event(), threading.Event()] - resume_parse = [threading.Event(), threading.Event()] + parse_started = [threading.Event() for _ in range(5)] + resume_parse = [threading.Event() for _ in range(5)] watcher_started = asyncio.Event() + watcher_update_requested = asyncio.Event() + watcher_update_completed = asyncio.Event() chunk_calls: list[Path] = [] watcher_contents: list[str] = [] @@ -1227,16 +1231,12 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc result = original_chunk(file_path, rel_path) chunk_calls.append(file_path) attempt = len(chunk_calls) - 1 - if attempt < len(parse_started): + if attempt in {0, 1, 3, 4}: parse_started[attempt].set() if not resume_parse[attempt].wait(timeout=5): - raise TimeoutError(f"initial parse {attempt + 1} was not resumed") + raise TimeoutError(f"parse attempt {attempt + 1} was not resumed") 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], @@ -1247,24 +1247,20 @@ 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]: + del paths, watch_filter, recursive, force_polling, debounce, step, poll_delay_ms 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, - 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 + await watcher_update_requested.wait() + yield {(Change.modified, str(path))} + watcher_update_completed.set() + await stop_event.wait() monkeypatch.setattr(chunker, "_chunk_sync", count_chunk) + from reme.steps.index import watch_changes as watch_changes_module + monkeypatch.setattr(watch_changes_module, "awatch", track_watcher_start) try: @@ -1286,6 +1282,22 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc assert "".join(chunk.text for chunk in chunks) == latest assert watcher_contents == [latest] assert len(chunk_calls) == 3 + + path.write_text("watcher update before first parse", encoding="utf-8") + watcher_update_requested.set() + assert await asyncio.to_thread(parse_started[3].wait, 5) + path.write_text("watcher update before retry parse", encoding="utf-8") + resume_parse[3].set() + assert await asyncio.to_thread(parse_started[4].wait, 5) + path.write_text("latest watcher update source", encoding="utf-8") + resume_parse[4].set() + + await asyncio.wait_for(watcher_update_completed.wait(), timeout=5) + assert len(chunk_calls) == 5 + assert job._task is not None and not job._task.done() + nodes = await store.get_nodes() + chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids] + assert "".join(chunk.text for chunk in chunks) == latest finally: for event in resume_parse: event.set() diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py index be64e117..1d0a86e7 100644 --- a/tests/unit/test_file_chunker_workers.py +++ b/tests/unit/test_file_chunker_workers.py @@ -7,6 +7,7 @@ import json import os import threading import weakref +from contextlib import asynccontextmanager from pathlib import Path from typing import Callable from unittest.mock import patch @@ -582,27 +583,25 @@ async def test_stale_batch_chunks_are_released_before_source_retry( ) -> 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") + source_path = tmp_path / "source.txt" + source_path.write_text("original 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: - 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 + await store.upsert([await chunker.chunk(source_path)]) + source_path.write_text("parsed source") + flush_waiting = asyncio.Event() + original_guard = store._maintenance_guard - 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) + @asynccontextmanager + async def track_store_guard(): + flush_waiting.set() + async with original_guard(): + yield - step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=2, app_context=context) + monkeypatch.setattr(store, "_maintenance_guard", track_store_guard) + step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=1, app_context=context) original_build = step._try_build_item stale_chunk_released: threading.Event | None = None @@ -614,35 +613,38 @@ async def test_stale_batch_chunks_are_released_before_source_retry( 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: + if 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: + if 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): + await store._maintenance_lock.acquire() + try: task = asyncio.create_task( step( RuntimeContext( - changes=[ - {"change": "modified", "path": str(first_path)}, - {"change": "modified", "path": str(second_path)}, - ], + changes=[{"change": "modified", "path": str(source_path)}], + retry_source_changes_before_watch=True, ), ), ) - await _wait_started(entered) - first_path.write_text("latest first") - release.set() - response = await task + await asyncio.wait_for(flush_waiting.wait(), timeout=5) + source_path.write_text("latest source with a different size") + finally: + store._maintenance_lock.release() + 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 + nodes = await store.get_nodes() + assert len(nodes) == 1 + assert store.file_chunks[nodes[0].chunk_ids[0]].text == "latest source with a different size" @pytest.mark.asyncio