This commit is contained in:
mikemikimike 2026-10-28 14:15:38 +08:00 • committed by GitHub
commit 921ca5d470
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1864 additions and 44 deletions

View file

@ -1,11 +1,37 @@
"""Abstract base for file chunkers."""
import asyncio
import contextvars
import os
import threading
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
from ...utils.file_utils import check_file_version
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):
@ -16,6 +42,109 @@ 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 = threading.Lock()
self._worker_state_lock = threading.Lock()
self._worker_generation = 0
self._worker_closing = False
async def _start(self) -> None:
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:
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:
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,
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."""
with self._worker_state_lock:
generation = self._worker_generation
file_path = Path(path)
rel_path = self.to_workspace_relative(path)
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 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(
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()
check_file_version(path, before, after)
return result, after
@abstractmethod
async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:

View file

@ -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), []

View file

@ -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():

View file

@ -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():

View file

@ -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,11 +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)
# A write during the read must stay detectable by the next scan.
stat = file_path.stat()
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] = []
@ -144,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:
@ -159,7 +162,7 @@ class MarkdownFileChunker(DefaultFileChunker):
node = FileNode(
path=rel_path,
st_mtime=stat.st_mtime,
st_mtime=file_path.stat().st_mtime,
chunk_ids=[chunk.id for chunk in chunks],
links=links,
front_matter=front_matter,

View file

@ -7,7 +7,7 @@ from pathlib import Path
from watchfiles import Change
def _normalize_change(raw) -> Change | None:
def _normalize_change(raw: object) -> Change | None:
if isinstance(raw, Change):
return raw
if isinstance(raw, str):
@ -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"])

View file

@ -73,7 +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_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

View file

@ -2,6 +2,7 @@
import asyncio
import math
import os
from abc import abstractmethod
from pathlib import Path
from typing import Any, Callable
@ -15,6 +16,10 @@ 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
_SOURCE_CHANGE_RETRY_MARKER = "_source_changed"
_INITIAL_SCAN_RECOVERY_PASSES = 1
class ChangeApplyStep(BaseStep):
@ -22,6 +27,7 @@ class ChangeApplyStep(BaseStep):
target_name = "target"
reads_file_content = False
source_change_retries: int = 0
def __init__(
self,
@ -83,8 +89,45 @@ 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())
results = await self._apply_existing(buckets)
buckets = bucket_changes(
changes,
path_exists=lambda p: self._to_abs_path(p).is_file(),
path_key=self._change_path_key,
)
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()
@ -123,9 +166,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,
@ -144,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
@ -165,6 +210,7 @@ class ChangeApplyStep(BaseStep):
"path": path,
"success": False,
"error": "not a file",
_SOURCE_CHANGE_RETRY_MARKER: True,
},
)
return None
@ -188,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)})
@ -198,15 +256,52 @@ class ChangeApplyStep(BaseStep):
change: Change,
path: str,
abs_path: Path,
size_bytes: int,
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
) -> 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), size_bytes
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),
_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}")
results.append({"change": change.name, "path": path, "success": False, "error": str(e)})
return None
def _try_estimate_memory(
self,
@ -290,6 +385,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):
@ -322,9 +421,143 @@ 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] = {}
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)
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()
check_file_version(path, before, after)
self._source_versions[item[0].path] = after
return item
async def _try_upsert(
self,
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] = []
item: tuple[FileNode, list[FileChunk]] | None = None
for item, path in zip(items, ok_paths):
node = item[0]
try:
check_file_version(
self._to_abs_path(node.path),
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),
_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] = {
"change": change.name,
"path": path,
"success": False,
"error": str(exc),
}
else:
current_items.append(item)
current_paths.append(path)
finally:
self._source_versions.pop(node.path)
if current_items:
result_by_path.update(
{
result["path"]: result
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:
inspected = await self._inspect_path(change, path, retry_results)
if inspected is None:
continue
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],
[path],
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]
async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None:
await self.file_store.upsert(items)

15
reme/utils/file_utils.py Normal file
View file

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

View file

@ -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": {}
}
]
}
}

View file

@ -15,8 +15,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
@ -29,8 +31,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
@ -51,6 +55,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")
@ -944,6 +949,83 @@ async def test_update_index_handles_file_removed_before_stat():
]
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())
@pytest.mark.asyncio
async 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."""
@ -998,6 +1080,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)
@ -1009,6 +1092,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"),
@ -1026,6 +1110,133 @@ 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")
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] = []
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 count_chunk(file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]:
result = original_chunk(file_path, rel_path)
chunk_calls.append(file_path)
attempt = len(chunk_calls) - 1
if attempt in {0, 1, 3, 4}:
parse_started[attempt].set()
if not resume_parse[attempt].wait(timeout=5):
raise TimeoutError(f"parse attempt {attempt + 1} was not resumed")
return result
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]:
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()
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:
await store.start()
await chunker.start()
await job.start()
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[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 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()
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."""
@ -1313,12 +1524,12 @@ async 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)

View file

@ -0,0 +1,793 @@
"""Worker admission, output compatibility, and real indexing/search regressions."""
import asyncio
import contextvars
import gc
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
import pytest
from httpx import ASGITransport, AsyncClient
from watchfiles import Change
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
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:
"""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 is (outcome in ("complete", "change"))
nodes = await store.get_nodes()
chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids]
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()
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", "oversize", "cancel"])
async def test_changed_batched_file_is_rebuilt_before_publication_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)
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(runtime_context))
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()
elif change == "oversize":
first_path.write_text("content exceeds the configured limit")
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
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()}
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"
)
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
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)
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(source_path)])
source_path.write_text("parsed source")
flush_waiting = asyncio.Event()
original_guard = store._maintenance_guard
@asynccontextmanager
async def track_store_guard():
flush_waiting.set()
async with original_guard():
yield
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
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 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 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)
await store._maintenance_lock.acquire()
try:
task = asyncio.create_task(
step(
RuntimeContext(
changes=[{"change": "modified", "path": str(source_path)}],
retry_source_changes_before_watch=True,
),
),
)
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
@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(
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 True
assert not step._source_versions
nodes = await store.get_nodes()
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