From cd3cbbadd84b533229aebba8003590d863224336 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:08:02 +0800 Subject: [PATCH] fix(memory): preserve file consistency and decision client lifecycle (#672) * fix(memory): keep writes and startup indexes consistent * fix(release): verify current Auto Fin plugin backends * fix(memory): preserve write ownership and deduplicate startup changes * fix(compaction): offload decision client initialization safely * fix(file-io): preserve cancellation origin and file permissions --- .github/workflows/release-auto-fin.yml | 3 +- docs/en/context_compaction.md | 8 + docs/zh/context_compaction.md | 6 + .../as_decision_model/agentscope.py | 22 ++- .../file_chunker/markdown_file_chunker.py | 4 +- reme/config/default.yaml | 5 + reme/steps/compaction/__init__.py | 5 +- reme/steps/evolve/auto_tag.py | 2 + reme/steps/file_io/_file_io.py | 34 +++- reme/steps/index/watch_changes.py | 23 +++ reme/utils/async_utils.py | 22 ++- tests/unit/test_async_utils.py | 100 ++++++++++ tests/unit/test_atomic_file_write.py | 184 ++++++++++++++++++ tests/unit/test_auto_tag.py | 34 ++++ tests/unit/test_compact_context.py | 3 +- tests/unit/test_decision_model.py | 55 ++++++ tests/unit/test_markdown_file_chunker.py | 22 +++ tests/unit/test_threaded_watch.py | 75 +++++++ 18 files changed, 586 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_async_utils.py create mode 100644 tests/unit/test_atomic_file_write.py diff --git a/.github/workflows/release-auto-fin.yml b/.github/workflows/release-auto-fin.yml index 6ebbf48d..ccb9f52e 100644 --- a/.github/workflows/release-auto-fin.yml +++ b/.github/workflows/release-auto-fin.yml @@ -133,7 +133,8 @@ jobs: assert set(manifest.backends) == { "auto_fin_data_step", "auto_fin_topic_step", - "auto_fin_merge_step", + "auto_fin_research_step", + "auto_fin_digest_step", } assert set(manifest.application_defaults["jobs"]) == { "auto_fin", diff --git a/docs/en/context_compaction.md b/docs/en/context_compaction.md index c6677668..48e705be 100644 --- a/docs/en/context_compaction.md +++ b/docs/en/context_compaction.md @@ -176,3 +176,11 @@ keep the instance alive until ReMe has finished its calls, then close it itself. replacing an internally constructed model or any model with active requests. Shutdown waits for active calls and owned-client cleanup even when cancelled, then propagates cancellation. Component shutdown detaches the instance; inject again before restarting a borrowed-model component. Standalone credential-based behavior is unchanged. + +Owned clients initialize in a worker thread so certificate loading does not block the event loop. Cancelling a request +during construction waits for initialization to finish and leaves the client owned by the component until shutdown; +the cancelled request is not sent. + +Hosts can import `CompactContextStep`, `ToolPair`, `collect_pairs`, and `preview_result` from `reme.steps.compaction`. +The pairing and preview helpers share the Step's rules for validating proposals. Archiving, resolving references, +and committing proposals remain host responsibilities. diff --git a/docs/zh/context_compaction.md b/docs/zh/context_compaction.md index 33accb2e..78328411 100644 --- a/docs/zh/context_compaction.md +++ b/docs/zh/context_compaction.md @@ -162,3 +162,9 @@ YAML 中的 `model` 仍表示供应商模型名称；运行时 `model` 属性接 宿主应等待 ReMe 调用结束后自行关闭实例。替换 ReMe 自建模型或存在进行中请求的实例之前，须先关闭组件。 关闭过程即使被取消，也会等待活动请求和自建客户端清理完成，再传播取消信号。组件关闭时会解除实例引用； 重新启动使用外部模型的组件前需要再次注入。独立运行时按凭据创建模型的行为保持不变。 + +自建客户端在线程中初始化，证书加载不会阻塞事件循环。构造期间取消请求时，组件会等待构造完成并保留客户端所有权， +由后续组件关闭负责释放；此时不会发送该请求。 + +宿主可从 `reme.steps.compaction` 导入 `CompactContextStep`、`ToolPair`、`collect_pairs` 和 `preview_result`。 +配对识别与预览构造函数和压缩 Step 共享相同规则，供宿主校验提案使用；归档、引用解析和提交仍由宿主负责。 diff --git a/reme/components/as_decision_model/agentscope.py b/reme/components/as_decision_model/agentscope.py index 640e3a93..a5699acf 100644 --- a/reme/components/as_decision_model/agentscope.py +++ b/reme/components/as_decision_model/agentscope.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter from .base_as_decision_model import BaseAsDecisionModel from ..component_registry import R +from ...utils.async_utils import complete_in_thread if TYPE_CHECKING: from agentscope.classifier import ClassifierModelBase, ClassifierResponse @@ -117,18 +118,19 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): raise ValueError("Decision model inference requires DECISION_MODEL_API_KEY or credential.api_key.") if not credential.get("base_url"): credential["base_url"] = None + # Reserve ownership before yielding so host injection cannot race + # initialization. The worker publishes the model to this component + # before cancellation can release the lifecycle lock. + self._owns_model = True try: - credential_cls = import_module("agentscope.credential").TypeSafeCredential - self._model = model_cls( - credential=credential_cls(**credential), - **self.options.model_dump(exclude_none=True), - ) - self._owns_model = True + await complete_in_thread(self._create_model, model_cls, credential) except ImportError as exc: raise ImportError( "Decision model inference requires a compatible typesafe-sdk. " "Install it with `pip install 'reme-ai[decision]'`.", ) from exc + finally: + self._owns_model = self._model is not None model = self._model self._active_calls += 1 self._idle.clear() @@ -141,6 +143,14 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): if self._active_calls == 0: self._idle.set() + def _create_model(self, model_cls, credential) -> None: + """Construct and retain the owned client, including synchronous TLS setup.""" + credential_cls = import_module("agentscope.credential").TypeSafeCredential + self._model = model_cls( + credential=credential_cls(**credential), + **self.options.model_dump(exclude_none=True), + ) + async def _close(self) -> None: """Finish draining and cleanup before propagating shutdown cancellation.""" cleanup = asyncio.create_task(self._release_model()) diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index 06a07793..e8fe3277 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -126,6 +126,8 @@ class MarkdownFileChunker(DefaultFileChunker): 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) front_matter, content, line_offset = self._parse_front_matter(text) @@ -157,7 +159,7 @@ class MarkdownFileChunker(DefaultFileChunker): node = FileNode( path=rel_path, - st_mtime=file_path.stat().st_mtime, + st_mtime=stat.st_mtime, chunk_ids=[chunk.id for chunk in chunks], links=links, front_matter=front_matter, diff --git a/reme/config/default.yaml b/reme/config/default.yaml index c57c0b3c..347f5a1d 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -16,6 +16,7 @@ jobs: dispatch_steps: [update_index_step] - backend: watch_changes_step watch_mode: awatch # Set to watch to isolate the blocking wait from AnyIO's shared thread limiter. + monitor_type: file_store dispatch_steps: - backend: update_index_step persist: False @@ -37,6 +38,8 @@ jobs: - auto_text_resource_step - backend: watch_changes_step watch_mode: awatch + monitor_type: file_catalog + monitor_name: resource dispatch_steps: - backend: update_catalog_step file_catalog: resource @@ -59,6 +62,8 @@ jobs: - backend: log_changes_step - backend: watch_changes_step watch_mode: awatch + monitor_type: file_catalog + monitor_name: digest dispatch_steps: - backend: update_catalog_step file_catalog: digest diff --git a/reme/steps/compaction/__init__.py b/reme/steps/compaction/__init__.py index 57aeeab5..eaeb79b4 100644 --- a/reme/steps/compaction/__init__.py +++ b/reme/steps/compaction/__init__.py @@ -1,5 +1,6 @@ -"""Tool-history compaction Steps.""" +"""Tool-history compaction and shared host validation helpers.""" from .compact_context import CompactContextStep +from .messages import ToolPair, collect_pairs, preview_result -__all__ = ["CompactContextStep"] +__all__ = ["CompactContextStep", "ToolPair", "collect_pairs", "preview_result"] diff --git a/reme/steps/evolve/auto_tag.py b/reme/steps/evolve/auto_tag.py index 05ea0604..823d59cb 100644 --- a/reme/steps/evolve/auto_tag.py +++ b/reme/steps/evolve/auto_tag.py @@ -123,6 +123,8 @@ class AutoTagStep(BaseStep): path = self.workspace_path / target.path metadata = dict(frontmatter.loads(path.read_text(encoding="utf-8")).metadata or {}) + if tag_key not in metadata: + raise RuntimeError(f"Auto-tag agent did not write the required {tag_key!r} field") normalized = normalize_memory_tags( metadata.get(tag_key), max_tags_per_file=tag_index.max_tags_per_file, diff --git a/reme/steps/file_io/_file_io.py b/reme/steps/file_io/_file_io.py index 82c5264d..22e945b0 100644 --- a/reme/steps/file_io/_file_io.py +++ b/reme/steps/file_io/_file_io.py @@ -1,7 +1,10 @@ """Encoding-aware file IO, output truncation, and per-path write locks.""" import asyncio +import os from pathlib import Path +from secrets import token_hex +from stat import S_IMODE from typing import Iterable import aiofiles @@ -10,6 +13,7 @@ import aiofiles.os from ...constants import DEFAULT_MAX_BYTES, MAX_FILE_READ_BYTES, TRUNCATION_NOTICE_MARKER from ...schema import FileChunk from ...utils import get_logger +from ...utils.async_utils import complete_in_thread logger = get_logger(log_to_file=False) @@ -152,7 +156,7 @@ async def detect_file_encoding(file_path, sniff_bytes: int = 8192) -> str: async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = "utf-8") -> None: - """Write ``content`` to ``file_path`` in binary mode; creates parent dirs.""" + """Publish complete bytes atomically while retaining normal file permissions.""" file_path.parent.mkdir(parents=True, exist_ok=True) if isinstance(content, str): try: @@ -165,8 +169,32 @@ async def write_file_safe(file_path: Path, content: str | bytes, encoding: str = payload = content.encode("utf-8") else: payload = content - async with aiofiles.open(str(file_path), "wb") as f: - await f.write(payload) + + def publish(): + flags = os.O_WRONLY | getattr(os, "O_BINARY", 0) + # Check OS write access without truncating the existing bytes. + try: + existing_fd = os.open(file_path, flags) + except FileNotFoundError: + mode = None + else: + try: + mode = S_IMODE(os.fstat(existing_fd).st_mode) + finally: + os.close(existing_fd) + + temporary = file_path.parent / f".reme-{token_hex(16)}.tmp" + fd = os.open(temporary, flags | os.O_CREAT | os.O_EXCL, 0o600 if mode is not None else 0o666) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + if mode is not None: + os.chmod(temporary, mode) + os.replace(temporary, file_path) + finally: + Path(temporary).unlink(missing_ok=True) + + await complete_in_thread(publish) # --------------------------------------------------------------------------- diff --git a/reme/steps/index/watch_changes.py b/reme/steps/index/watch_changes.py index a01d205f..a39eb467 100644 --- a/reme/steps/index/watch_changes.py +++ b/reme/steps/index/watch_changes.py @@ -60,6 +60,8 @@ class WatchChangesStep(BaseStep): step: int = DEFAULT_WATCH_STEP_MS, poll_delay_ms: int = DEFAULT_LOW_POWER_POLL_MS, watch_mode: Literal["awatch", "watch"] = "awatch", + monitor_type: Literal["file_store", "file_catalog"] | None = None, + monitor_name: str = "default", **kwargs, ): if watch_mode not in ("awatch", "watch"): @@ -71,6 +73,8 @@ class WatchChangesStep(BaseStep): self.debounce: int = debounce self.step: int = step self.poll_delay_ms: int = poll_delay_ms + self.monitor_type = monitor_type + self.monitor_name = monitor_name self._rules: list[WatchRule] = [] def _get_watch_rules(self) -> list[WatchRule]: @@ -112,7 +116,9 @@ class WatchChangesStep(BaseStep): step=self.step, poll_delay_ms=self.poll_delay_ms, stop_event=stop_event, + yield_on_timeout=self.monitor_type is not None, ) + needs_initial_scan = self.monitor_type is not None async with aclosing(watcher): async for raw_changes in watcher: if stop_event.is_set(): @@ -122,6 +128,23 @@ class WatchChangesStep(BaseStep): for c, p in raw_changes if c in (Change.added, Change.modified, Change.deleted) ] + if needs_initial_scan: + # Subscribe first, then recover writes made after the + # preceding scan but before watcher initialization. + await self.dispatch_steps( + [ + { + "backend": "init_changes_step", + "monitor_type": self.monitor_type, + "monitor_name": self.monitor_name, + "recursive": self.recursive, + }, + ], + ) + # The scan only produces changes. Merge overlapping events + # before running handlers, which may have model side effects. + changes.extend(self.context.get("changes") or []) + needs_initial_scan = False changes = coalesce_changes(changes) if changes: self.logger.info(f"Detected {len(changes)} change(s)") diff --git a/reme/utils/async_utils.py b/reme/utils/async_utils.py index 290a4456..8f54fbcb 100644 --- a/reme/utils/async_utils.py +++ b/reme/utils/async_utils.py @@ -7,11 +7,21 @@ from typing import Any async def complete_in_thread(func: Callable[..., Any], /, *args) -> Any: - """Finish a side-effecting thread call before propagating cancellation.""" + """Finish a side-effecting thread call, including through repeated cancellation.""" task = asyncio.create_task(asyncio.to_thread(func, *args)) - try: - return await asyncio.shield(task) - except asyncio.CancelledError: + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + except Exception: + if cancellation is None: + raise + break + if cancellation is not None: with suppress(Exception): - await task - raise + task.result() + raise cancellation + return task.result() diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py new file mode 100644 index 00000000..70b90adc --- /dev/null +++ b/tests/unit/test_async_utils.py @@ -0,0 +1,100 @@ +"""Side-effecting thread calls retain ownership through cancellation.""" + +import asyncio +import threading +import time + +import anyio +import pytest + +from reme.utils.async_utils import complete_in_thread + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancellations", [0, 1, 3]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_complete_in_thread_finishes_before_propagating_result(cancellations, fail): + """Repeated cancellation waits for success or failure without detaching work.""" + started, release, finished = threading.Event(), threading.Event(), threading.Event() + + def operation(): + started.set() + try: + assert release.wait(5) + if fail: + raise RuntimeError("thread failed") + return "complete" + finally: + finished.set() + + task = asyncio.create_task(complete_in_thread(operation)) + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(cancellations): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + assert not finished.is_set() + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + + assert finished.is_set() + if cancellations: + with pytest.raises(asyncio.CancelledError): + task.result() + elif fail: + with pytest.raises(RuntimeError, match="thread failed"): + task.result() + else: + assert task.result() == "complete" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope_name", ["fail_after", "move_on_after"]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_complete_in_thread_preserves_anyio_cancel_scope(scope_name, fail): + """AnyIO must recognize its original cancellation after the thread finishes.""" + finished = threading.Event() + + def operation(): + time.sleep(0.05) + finished.set() + if fail: + raise RuntimeError("thread failed after cancellation") + + async def run_scope(): + with getattr(anyio, scope_name)(0.01) as scope: + await complete_in_thread(operation) + assert scope.cancelled_caught + + if scope_name == "fail_after": + with pytest.raises(TimeoutError): + await run_scope() + else: + await run_scope() + assert finished.is_set() + + +@pytest.mark.asyncio +async def test_complete_in_thread_retains_first_cancellation_message(): + """Repeated cancellation must retain the original caller's reason.""" + started, release = threading.Event(), threading.Event() + + def operation(): + started.set() + assert release.wait(5) + + task = asyncio.create_task(complete_in_thread(operation)) + try: + assert await asyncio.to_thread(started.wait, 5) + for message in ("original cancellation", "cancel again", "cancel once more"): + task.cancel(message) + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError, match="original cancellation"): + await task + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/unit/test_atomic_file_write.py b/tests/unit/test_atomic_file_write.py new file mode 100644 index 00000000..b2a68dbd --- /dev/null +++ b/tests/unit/test_atomic_file_write.py @@ -0,0 +1,184 @@ +"""Memory writers publish complete files to concurrent index readers.""" + +import asyncio +import os +from pathlib import Path +import stat +import subprocess +import sys +import threading + +import pytest + +from reme.steps.file_io._file_io import write_file_safe +from reme.components.application_context import ApplicationContext +from reme.steps.file_io.save import SaveStep + + +@pytest.mark.asyncio +async def test_overwrite_preserves_complete_old_content_until_commit(tmp_path, monkeypatch): + """Readers retain the prior complete revision until atomic publication.""" + target = tmp_path / "memory.md" + target.write_text("old complete memory", encoding="utf-8") + target.chmod(0o640) + replace = os.replace + + def observe_commit(source, destination): + assert target.read_text(encoding="utf-8") == "old complete memory" + assert source != destination + replace(source, destination) + + monkeypatch.setattr(os, "replace", observe_commit) + await write_file_safe(target, "new complete memory") + + assert target.read_text(encoding="utf-8") == "new complete memory" + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + assert sorted(path.name for path in tmp_path.iterdir()) == ["memory.md"] + + +@pytest.mark.asyncio +async def test_failed_commit_preserves_original_and_cleans_temporary_file(tmp_path, monkeypatch): + """Failed publication leaves source bytes intact and no temporary assets.""" + target = tmp_path / "memory.md" + target.write_bytes(b"original memory") + + def fail_commit(_source, _destination): + raise OSError("controlled commit failure") + + monkeypatch.setattr(os, "replace", fail_commit) + with pytest.raises(OSError, match="controlled commit failure"): + await write_file_safe(target, "replacement") + + assert target.read_bytes() == b"original memory" + assert sorted(path.name for path in tmp_path.iterdir()) == ["memory.md"] + + +@pytest.mark.asyncio +async def test_new_encoded_file_and_binary_payload(tmp_path): + """Atomic writes preserve the existing encoding and binary contracts.""" + target = tmp_path / "nested" / "memory.csv" + await write_file_safe(target, "中文", encoding="gbk") + assert target.read_bytes() == "中文".encode("gbk") + await write_file_safe(target, b"\x00\xff") + assert target.read_bytes() == b"\x00\xff" + + +@pytest.mark.asyncio +async def test_atomic_write_accepts_long_valid_filename(tmp_path): + """Temporary names must not grow with the destination's filename.""" + target = tmp_path / ("a" * 241 + ".md") + target.write_text("original", encoding="utf-8") + + await write_file_safe(target, "replacement") + + assert target.read_text(encoding="utf-8") == "replacement" + assert list(tmp_path.iterdir()) == [target] + + +@pytest.mark.asyncio +async def test_cancelled_save_keeps_path_lock_until_publication(tmp_path, monkeypatch): + """A repeatedly cancelled writer cannot overwrite a later successful save.""" + target = tmp_path / "memory.md" + target.write_text("original", encoding="utf-8") + context = ApplicationContext(workspace_dir=str(tmp_path)) + started, release = threading.Event(), threading.Event() + replace = os.replace + + def delayed_replace(source, destination): + if not started.is_set(): + started.set() + assert release.wait(5) + replace(source, destination) + + monkeypatch.setattr(os, "replace", delayed_replace) + first = asyncio.create_task(SaveStep(app_context=context)(path="memory.md", content="cancelled write")) + second = None + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(3): + first.cancel() + await asyncio.sleep(0) + second = asyncio.create_task(SaveStep(app_context=context)(path="memory.md", content="successful write")) + await asyncio.sleep(0) + assert not first.done() + assert not second.done() + finally: + release.set() + results = await asyncio.gather(first, *([second] if second is not None else []), return_exceptions=True) + + assert isinstance(results[0], asyncio.CancelledError) + assert results[1].success + assert target.read_text(encoding="utf-8") == "successful write" + assert list(tmp_path.iterdir()) == [target] + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="POSIX file permissions") +async def test_save_rejects_read_only_target_without_replacing_it(tmp_path): + """A writable parent must not let saving bypass the target's access rights.""" + if os.geteuid() == 0: + pytest.skip("Root bypasses ordinary file write permissions") + target = tmp_path / "readonly.md" + target.write_bytes(b"original") + target.chmod(0o444) + try: + response = await SaveStep(app_context=ApplicationContext(workspace_dir=str(tmp_path)))( + path=target.name, + content="replacement", + ) + assert response.success is False + assert target.read_bytes() == b"original" + assert stat.S_IMODE(target.stat().st_mode) == 0o444 + assert list(tmp_path.iterdir()) == [target] + finally: + target.chmod(0o644) + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="POSIX file permissions") +async def test_atomic_write_does_not_require_target_read_access(tmp_path): + """Permission checks retain the existing write-only file contract.""" + target = tmp_path / "writeonly.md" + target.write_bytes(b"original") + target.chmod(0o200) + try: + await write_file_safe(target, b"replacement") + assert stat.S_IMODE(target.stat().st_mode) == 0o200 + finally: + target.chmod(0o600) + assert target.read_bytes() == b"replacement" + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX umask semantics") +@pytest.mark.parametrize("mask, expected", [(0o002, 0o664), (0o077, 0o600)]) +def test_new_save_respects_process_umask(tmp_path, mask, expected): + """Check real creation modes without changing the test process's umask.""" + script = """ +import asyncio +import os +from pathlib import Path +import stat +import sys +from reme.components.application_context import ApplicationContext +from reme.steps.file_io.save import SaveStep + +workspace = Path(sys.argv[1]) +os.umask(int(sys.argv[2])) +response = asyncio.run(SaveStep(app_context=ApplicationContext(workspace_dir=str(workspace)))( + path="new.md", content="memory" +)) +assert response.success, response.answer +target = workspace / "new.md" +assert target.read_text() == "memory" +assert list(workspace.iterdir()) == [target] +print(stat.S_IMODE(target.stat().st_mode)) +""" + result = subprocess.run( + [sys.executable, "-c", script, str(tmp_path), str(mask)], + cwd=Path(__file__).resolve().parents[2], + capture_output=True, + text=True, + check=True, + timeout=15, + ) + assert int(result.stdout.splitlines()[-1]) == expected diff --git a/tests/unit/test_auto_tag.py b/tests/unit/test_auto_tag.py index 3336a381..69d6a6f1 100644 --- a/tests/unit/test_auto_tag.py +++ b/tests/unit/test_auto_tag.py @@ -3,6 +3,7 @@ # pylint: disable=missing-function-docstring from pathlib import Path +from unittest.mock import AsyncMock import frontmatter import pytest @@ -55,6 +56,39 @@ def _file_store(*, name: str = "default", **tag_index_kwargs) -> LocalFileStore: return store +@pytest.mark.asyncio +async def test_auto_tag_reports_missing_agent_write_without_inventing_empty_tags(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + note = tmp_path / "memory/note.md" + _write_note(note) + original = note.read_bytes() + wrapper = _TaggingWrapper(tmp_path) + monkeypatch.setattr(wrapper, "reply", AsyncMock(return_value={"result": "", "last_message": {}})) + step = AutoTagStep(agent_wrapper=wrapper, file_store=_file_store()) + + response = await step(RuntimeContext(changes=[{"change": "added", "path": "memory/note.md"}])) + + assert response.metadata["auto_tag"]["failed"] == 1 + assert response.metadata["auto_tag"]["succeeded"] == 0 + assert "did not write" in response.metadata["auto_tag"]["results"][0]["error"] + assert note.read_bytes() == original + + +@pytest.mark.asyncio +async def test_auto_tag_accepts_explicitly_written_empty_tags(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + note = tmp_path / "memory/note.md" + _write_note(note) + wrapper = _TaggingWrapper(tmp_path, tags=[]) + step = AutoTagStep(agent_wrapper=wrapper, file_store=_file_store()) + + response = await step(RuntimeContext(changes=[{"change": "added", "path": "memory/note.md"}])) + + assert response.metadata["auto_tag"]["succeeded"] == 1 + assert response.metadata["auto_tag"]["failed"] == 0 + assert frontmatter.loads(note.read_text(encoding="utf-8")).metadata["memory_tags"] == [] + + @pytest.mark.asyncio async def test_auto_tag_handles_noop_invalid_changes_and_no_tag_index(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) diff --git a/tests/unit/test_compact_context.py b/tests/unit/test_compact_context.py index 5d10fae0..3f48f2a4 100644 --- a/tests/unit/test_compact_context.py +++ b/tests/unit/test_compact_context.py @@ -14,8 +14,7 @@ from reme.components.as_decision_model import BaseAsDecisionModel from reme.components.job import BaseJob from reme.enumeration import ComponentEnum from reme.schema.compaction import CompactionOptions, CompactionRequest -from reme.steps.compaction import CompactContextStep -from reme.steps.compaction.messages import collect_pairs +from reme.steps.compaction import CompactContextStep, collect_pairs from reme.steps.compaction.state import build_batches, estimate_decision_tokens, serialize diff --git a/tests/unit/test_decision_model.py b/tests/unit/test_decision_model.py index 1a4c3beb..b7416127 100644 --- a/tests/unit/test_decision_model.py +++ b/tests/unit/test_decision_model.py @@ -4,6 +4,7 @@ import asyncio import sys +import threading from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, Mock @@ -101,6 +102,13 @@ def test_valid_questions_reuse_model_and_restart_closes_clients(provider, provid } async def run(): + loop_thread = threading.get_ident() + + def construct(**_kwargs): + assert threading.get_ident() != loop_thread + return model + + factory.side_effect = construct component = AgentScopeAsDecisionModel( credential={"api_key": "test-key", "base_url": ""}, model="decision-test", @@ -138,6 +146,53 @@ def test_valid_questions_reuse_model_and_restart_closes_clients(provider, provid asyncio.run(run()) +@pytest.mark.parametrize("construction_fails", [False, True]) +def test_cancelled_construction_keeps_ownership_until_close(provider, construction_fails): + """Cancellation cannot abandon a client or let shutdown overtake its worker.""" + factory, model = provider + + async def run(): + loop = asyncio.get_running_loop() + entered, release = asyncio.Event(), threading.Event() + + def construct(**_kwargs): + loop.call_soon_threadsafe(entered.set) + assert release.wait(5), "Construction was not released" + if construction_fails: + raise RuntimeError("construction failed") + return model + + factory.side_effect = construct + component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"}) + call = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}})) + closing = None + try: + await asyncio.wait_for(entered.wait(), 2) + with pytest.raises(RuntimeError, match="Close the decision component"): + component.model = AsyncMock() + for _ in range(2): + call.cancel() + await asyncio.sleep(0) + assert not call.done() + closing = asyncio.create_task(component.close()) + await asyncio.sleep(0) + assert not closing.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(call, 2) + await asyncio.wait_for(closing, 2) + assert component.model is None + assert not component._owns_model + model.assert_not_awaited() + assert model.client.aclose.await_count == (0 if construction_fails else 1) + finally: + release.set() + await asyncio.gather(call, *([closing] if closing is not None else []), return_exceptions=True) + await component.close() + + asyncio.run(run()) + + @pytest.mark.parametrize( "question", [ diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index 3b18601b..73926bf5 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -28,6 +28,28 @@ def _write_md(tmpdir: str, name: str, body: str) -> str: return name +@pytest.mark.asyncio +async def test_chunk_keeps_pre_read_mtime_when_source_changes(tmp_path, monkeypatch): + """A concurrent write must remain visible to the next filesystem scan.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "changing.md" + path.write_text("", encoding="utf-8") + os.utime(path, (1000, 1000)) + chunker = MarkdownFileChunker() + + async def read_then_finish_write(_path): + path.write_text("Completed memory note.", encoding="utf-8") + os.utime(path, (2000, 2000)) + return "" + + monkeypatch.setattr(chunker, "_read_text_for_indexing", read_then_finish_write) + node, chunks = await chunker.chunk("changing.md") + + assert chunks == [] + assert node.st_mtime == 1000 + assert path.stat().st_mtime != node.st_mtime + + @pytest.mark.asyncio async def test_parse_empty_file(): """An empty .md → FileNode, no chunks, no links.""" diff --git a/tests/unit/test_threaded_watch.py b/tests/unit/test_threaded_watch.py index d93e7c02..5a56bf48 100644 --- a/tests/unit/test_threaded_watch.py +++ b/tests/unit/test_threaded_watch.py @@ -10,7 +10,10 @@ import pytest from watchfiles import Change from reme.components.application_context import ApplicationContext +from reme.components.file_catalog import LocalFileCatalog from reme.components.runtime_context import RuntimeContext +from reme.schema import FileNode +from reme.steps.base_step import BaseStep from reme.steps.index import _threaded_watch from reme.steps.index._threaded_watch import watch_in_thread from reme.steps.index.watch_changes import WatchChangesStep @@ -49,6 +52,78 @@ def test_watch_mode_defaults_and_validation(): WatchChangesStep(watch_mode=value) +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["watch", "awatch"]) +async def test_watcher_reconciles_startup_gap_on_first_idle_batch(tmp_path, monkeypatch, mode): + """Both watchers reconcile once even when no subsequent change arrives.""" + + async def idle_batches(*_paths, **kwargs): + assert kwargs["yield_on_timeout"] is True + yield set() + yield set() + + monkeypatch.setattr("reme.steps.index.watch_changes.awatch", idle_batches) + monkeypatch.setattr("reme.steps.index.watch_changes.watch_in_thread", idle_batches) + step = _step(tmp_path, asyncio.Event(), mode=mode) + step.monitor_type = "file_store" + step.dispatch_step_specs = ["update_index_step"] + + await step.execute() + + step.dispatch_steps.assert_awaited_once_with( + [ + { + "backend": "init_changes_step", + "monitor_type": "file_store", + "monitor_name": "default", + "recursive": True, + }, + ], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["watch", "awatch"]) +async def test_watcher_merges_startup_scan_and_first_change_batch(tmp_path, monkeypatch, mode): + """Overlapping events and missed additions/deletions dispatch once together.""" + step = _step(tmp_path, asyncio.Event(), mode=mode) + catalog = LocalFileCatalog(app_context=step.app_context) + step.app_context.components["file_catalog"] = {"default": catalog} + await catalog.upsert([FileNode(path="daily/deleted.md", st_mtime=1)]) + step.monitor_type = "file_catalog" + step.dispatch_step_specs = ["update_catalog_step"] + changed = tmp_path / "daily/changed.md" + missed = tmp_path / "daily/missed.md" + + async def first_batch(*_paths, **_kwargs): + changed.write_text("new", encoding="utf-8") + missed.write_text("missed", encoding="utf-8") + yield {(Change.added, str(changed))} + yield set() + + batches = [] + + async def dispatch(specs, **kwargs): + if specs[0] == "update_catalog_step": + batches.append(kwargs["changes"]) + return [] + return await BaseStep.dispatch_steps(step, specs, **kwargs) + + monkeypatch.setattr("reme.steps.index.watch_changes.awatch", first_batch) + monkeypatch.setattr("reme.steps.index.watch_changes.watch_in_thread", first_batch) + step.dispatch_steps = dispatch + await step.execute() + + assert len(batches) == 1 + assert sorted((item["path"], item["change"]) for item in batches[0]) == sorted( + [ + (str(changed), "added"), + (str(missed), "added"), + (str(tmp_path / "daily/deleted.md"), "deleted"), + ], + ) + + @pytest.mark.asyncio async def test_default_mode_uses_awatch_and_closes_on_dispatch_failure(tmp_path, monkeypatch): """An exception in downstream processing closes the default generator."""