From 99ecdd24d14d0cf3f59ddcd5ffff738138fe6a2e Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 3 Nov 2026 13:39:14 +0800 Subject: [PATCH 1/5] fix(index): move file parsing off the service event loop (#599) * fix(index): move file parsing off the service event loop Signed-off-by: mikemikimike <186855910+mikemikimike@users.noreply.github.com> * fix(index): retry changed source during indexing * fix(index): retry and coalesce updated file changes * test(index): cover initial scan source race * fix(index): recover unstable initial scans * Fix initial scan retry scope and stale batches * test(index): update review regressions * fix(index): resolve deletion aliases before applying * refactor(index): simplify review follow-ups * refactor(index): carry source versions with batch items --------- Signed-off-by: mikemikimike <186855910+mikemikimike@users.noreply.github.com> Co-authored-by: mikemikimike <186855910+mikemikimike@users.noreply.github.com> Co-authored-by: jinli.yl --- .../file_chunker/base_file_chunker.py | 108 +++ .../file_chunker/default_file_chunker.py | 16 +- .../file_chunker/json_file_chunker.py | 5 +- .../file_chunker/jsonl_file_chunker.py | 5 +- .../file_chunker/markdown_file_chunker.py | 16 +- reme/steps/index/_change_batch.py | 28 +- reme/steps/index/init_changes.py | 12 +- reme/steps/index/update_changes.py | 208 ++++- reme/utils/async_utils.py | 8 +- reme/utils/file_utils.py | 15 + tests/unit/fixtures/file_chunker_outputs.json | 412 +++++++++ tests/unit/test_background_steps.py | 217 ++++- tests/unit/test_file_chunker_workers.py | 824 ++++++++++++++++++ tests/unit/test_markdown_file_chunker.py | 14 +- 14 files changed, 1831 insertions(+), 57 deletions(-) create mode 100644 reme/utils/file_utils.py create mode 100644 tests/unit/fixtures/file_chunker_outputs.json create mode 100644 tests/unit/test_file_chunker_workers.py diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 49312989..6d9fed61 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -1,11 +1,19 @@ """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 from ..base_component import BaseComponent from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode +from ...utils.async_utils import finish_cleanup +from ...utils.file_utils import check_file_version class BaseFileChunker(BaseComponent): @@ -16,6 +24,106 @@ 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_cleanup(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() + 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 + if worker is None: + # 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) + 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_cleanup(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]]: diff --git a/reme/components/file_chunker/default_file_chunker.py b/reme/components/file_chunker/default_file_chunker.py index 615e1171..89e816bb 100644 --- a/reme/components/file_chunker/default_file_chunker.py +++ b/reme/components/file_chunker/default_file_chunker.py @@ -4,8 +4,6 @@ from bisect import bisect_right from pathlib import Path from typing import Literal -import aiofiles - from .base_file_chunker import BaseFileChunker from ..component_registry import R from ...schema import FileChunk, FileNode @@ -37,10 +35,9 @@ class DefaultFileChunker(BaseFileChunker): self.chunk_byte_size = max(100, chunk_byte_size) self.overlap_byte_size = max(4, overlap_byte_size) - async def _read_text_for_indexing(self, file_path: Path) -> str: + def _read_text_for_indexing(self, file_path: Path) -> str: """Decode text for a derived index without modifying the source file.""" - async with aiofiles.open(file_path, "rb") as f: - data = await f.read() + data = file_path.read_bytes() try: text = data.decode(self.encoding) except UnicodeDecodeError as exc: @@ -61,11 +58,12 @@ class DefaultFileChunker(BaseFileChunker): return text.replace("\r\n", "\n").replace("\r", "\n") async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: - file_path = Path(path) - stat = file_path.stat() - rel_path = self.to_workspace_relative(path) + return await self._chunk_in_worker(path, self._chunk_sync) - text = await self._read_text_for_indexing(file_path) + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: + stat = file_path.stat() + + text = self._read_text_for_indexing(file_path) if not text: return FileNode(path=rel_path, st_mtime=stat.st_mtime), [] diff --git a/reme/components/file_chunker/json_file_chunker.py b/reme/components/file_chunker/json_file_chunker.py index 78f8e617..e2c97686 100644 --- a/reme/components/file_chunker/json_file_chunker.py +++ b/reme/components/file_chunker/json_file_chunker.py @@ -113,9 +113,10 @@ class JsonFileChunker(BaseFileChunker): async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: """Read and chunk a JSON file at *path*.""" - file_path = Path(path) + return await self._chunk_in_worker(path, self._chunk_sync) + + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: stat = file_path.stat() - rel_path = self.to_workspace_relative(path) raw_text = file_path.read_text(encoding=self.encoding) if not raw_text.strip(): diff --git a/reme/components/file_chunker/jsonl_file_chunker.py b/reme/components/file_chunker/jsonl_file_chunker.py index 9417282c..0326b299 100644 --- a/reme/components/file_chunker/jsonl_file_chunker.py +++ b/reme/components/file_chunker/jsonl_file_chunker.py @@ -124,9 +124,10 @@ class JsonlFileChunker(BaseFileChunker): async def chunk(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]: """Read and chunk a JSONL file at *path*.""" - file_path = Path(path) + return await self._chunk_in_worker(path, self._chunk_sync) + + def _chunk_sync(self, file_path: Path, rel_path: str) -> tuple[FileNode, list[FileChunk]]: stat = file_path.stat() - rel_path = self.to_workspace_relative(path) text = file_path.read_text(encoding=self.encoding) if not text.strip(): diff --git a/reme/components/file_chunker/markdown_file_chunker.py b/reme/components/file_chunker/markdown_file_chunker.py index e8fe3277..9d5742c9 100644 --- a/reme/components/file_chunker/markdown_file_chunker.py +++ b/reme/components/file_chunker/markdown_file_chunker.py @@ -16,6 +16,7 @@ the single source of truth for ``[[...]]`` syntax. """ import re +import threading from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -39,6 +40,9 @@ _ATX_HEADING_RE = re.compile(r"^ {0,3}#{1,6}(?:[ \t]+|$)") _SETEXT_HEADING_RE = re.compile(r"^ {0,3}(?:=+|-+)[ \t]*$") _FENCE_RE = re.compile(r"^ {0,3}(`{3,}|~{3,})") +# MarkdownRenderer changes mistletoe's process-wide token registry. +_RENDERER_LOCK = threading.Lock() + @dataclass class MdNode: @@ -123,12 +127,8 @@ class MarkdownFileChunker(DefaultFileChunker): self.include_frontmatter_in_metadata = include_frontmatter_in_metadata 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) + 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 +144,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 +159,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, diff --git a/reme/steps/index/_change_batch.py b/reme/steps/index/_change_batch.py index 5000ff27..e1457603 100644 --- a/reme/steps/index/_change_batch.py +++ b/reme/steps/index/_change_batch.py @@ -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"]) diff --git a/reme/steps/index/init_changes.py b/reme/steps/index/init_changes.py index 40711637..365305d3 100644 --- a/reme/steps/index/init_changes.py +++ b/reme/steps/index/init_changes.py @@ -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 diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 6ce461f0..11a74b01 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -2,6 +2,7 @@ import asyncio import math +import os from abc import abstractmethod from pathlib import Path from typing import Any, Callable @@ -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,40 @@ 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 or 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 +161,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 +184,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 +205,7 @@ class ChangeApplyStep(BaseStep): "path": path, "success": False, "error": "not a file", + _SOURCE_CHANGE_RETRY_MARKER: True, }, ) return None @@ -188,6 +229,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 +251,41 @@ 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, FileNotFoundError) as e: + if attempt < self.source_change_retries: + reason = "changed" if isinstance(e, FileChangedDuringChunkingError) else "disappeared" + self.logger.debug(f"Source {reason} while processing {path}; retrying") + continue + if isinstance(e, FileChangedDuringChunkingError): + self.logger.exception(f"Failed to process {path}") + else: + 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, @@ -279,7 +358,7 @@ class ChangeApplyStep(BaseStep): await asyncio.sleep(0) batch = deleted[start : start + self.batch_max_files] try: - await self.delete_paths([self.to_workspace_relative(p) for p in batch]) + await self.delete_paths([self.to_workspace_relative(self._to_abs_path(p)) for p in batch]) results.extend({"change": "deleted", "path": p, "success": True} for p in batch) except Exception as e: self.logger.exception(f"Failed to delete {len(batch)} file(s) from {self.target_name}") @@ -290,6 +369,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,12 +405,99 @@ class UpdateIndexStep(ChangeApplyStep): target_name = "file_store" reads_file_content = True + source_change_retries: int = 1 - async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]: - return await self.chunk_file(path) + 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 upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None: - await self.file_store.upsert(items) + async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk], os.stat_result]: + """Keep each source version with its parsed item until publication.""" + before = path.stat() + node, chunks = await self.chunk_file(path) + after = path.stat() + check_file_version(path, before, after) + return node, chunks, after + + async def _try_upsert( + self, + change: Change, + items: list[tuple[FileNode, list[FileChunk], os.stat_result]], + 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], os.stat_result]] = [] + current_paths: list[str] = [] + item: tuple[FileNode, list[FileChunk], os.stat_result] | None = None + for item, path in zip(items, ok_paths): + node = item[0] + try: + check_file_version( + self._to_abs_path(node.path), + item[2], + self._to_abs_path(node.path).stat(), + ) + except (FileChangedDuringChunkingError, 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) + 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], os.stat_result]]) -> None: + await self.file_store.upsert([(node, chunks) for node, chunks, _ in items]) async def delete_paths(self, paths: list[str]) -> None: await self.file_store.delete(paths) @@ -343,7 +513,7 @@ class UpdateIndexStep(ChangeApplyStep): def estimate_item_memory( self, - item: tuple[FileNode, list[FileChunk]], + item: tuple[FileNode, list[FileChunk], os.stat_result], path: Path, size_bytes: int, ) -> int: diff --git a/reme/utils/async_utils.py b/reme/utils/async_utils.py index e80cb7ea..d2dd9a85 100644 --- a/reme/utils/async_utils.py +++ b/reme/utils/async_utils.py @@ -3,10 +3,12 @@ import asyncio from collections.abc import Awaitable, Callable from contextlib import suppress -from typing import Any +from typing import Any, TypeVar + +T = TypeVar("T") -async def finish_cleanup(operation: Awaitable): +async def finish_cleanup(operation: Awaitable[T]) -> T: """Finish cleanup before propagating any cancellation received while waiting.""" task = asyncio.ensure_future(operation) cancellation: asyncio.CancelledError | None = None @@ -21,7 +23,7 @@ async def finish_cleanup(operation: Awaitable): raise break if cancellation is not None: - with suppress(Exception): + with suppress(Exception, asyncio.CancelledError): task.result() raise cancellation return task.result() diff --git a/reme/utils/file_utils.py b/reme/utils/file_utils.py new file mode 100644 index 00000000..440f1635 --- /dev/null +++ b/reme/utils/file_utils.py @@ -0,0 +1,15 @@ +"""Shared file-version checks used by file readers and index updates.""" + +import os +from pathlib import Path + + +class FileChangedDuringChunkingError(RuntimeError): + """Raised when a source file changes while it is being parsed.""" + + +def check_file_version(path: Path, before: os.stat_result, after: os.stat_result) -> None: + """Raise when the source file's identity or content metadata changed.""" + fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + if any(getattr(before, field) != getattr(after, field) for field in fields): + raise FileChangedDuringChunkingError(f"File changed while chunking: {path}") diff --git a/tests/unit/fixtures/file_chunker_outputs.json b/tests/unit/fixtures/file_chunker_outputs.json new file mode 100644 index 00000000..19ad5a49 --- /dev/null +++ b/tests/unit/fixtures/file_chunker_outputs.json @@ -0,0 +1,412 @@ +{ + "plain.txt": { + "node": { + "path": "plain.txt", + "links": [], + "chunk_ids": [ + "27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "27beade9f75899604a5c33dc2a2cad48c68782309d1961d34aa416fa4a9f44e8", + "text": "first\nabc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 1, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "339017807ea80235b88b9a55e5f38a0eef93d0e5515ef7b5ae5d5586551f38c7", + "text": "c abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc abc ab", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "5fec11d28b94df329491529d38cc187104d52ea3500622503103a89feae73a23", + "text": "c abc ", + "embedding": null, + "metadata": {}, + "path": "plain.txt", + "start_line": 2, + "end_line": 2, + "scores": {} + } + ] + }, + "note.md": { + "node": { + "path": "note.md", + "links": [ + { + "source_path": "note.md", + "target_path": "other.md", + "target_anchor": null + } + ], + "chunk_ids": [ + "c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3", + "03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c", + "f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7" + ], + "front_matter": { + "name": "", + "description": "", + "title": "Example", + "tags": [ + "one", + "two" + ] + } + }, + "chunks": [ + { + "id": "c135f79d777949fa45848ad6302adcf6368a348910b2bcc4d99369b1acee75c3", + "text": "# First\n\nbody [[other.md]]", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 5, + "end_line": 7, + "scores": {} + }, + { + "id": "03ee7e6359843ec17c78ecc037d232e473d1b6b750ea544a0b1333a479a35f9c", + "text": "# First\n\n## Second", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 9, + "end_line": 9, + "scores": {} + }, + { + "id": "f74ee6f8929ad297e613b0341f0307ff1c9ff9656ec0b28b53de9ae58373c2d7", + "text": "text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text text", + "embedding": null, + "metadata": { + "title": "Example", + "tags": [ + "one", + "two" + ] + }, + "path": "note.md", + "start_line": 11, + "end_line": 11, + "scores": {} + } + ] + }, + "fallback.md": { + "node": { + "path": "fallback.md", + "links": [], + "chunk_ids": [ + "22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "22b74c7a4dcd76d21c6aa9b3717065fb8328a8d77215fb7cb37003a9ecafbbab", + "text": "# First\ntext text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 1, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "adfc277818b38936075c7aaf1fbc53d397ac6f9c2818dfa507d4598d67c2f203", + "text": "xt text text text text text text text text text text text text text text text text text text text te", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + }, + { + "id": "2560053ffa2a0c6decb9265aa1e33c0b264792535d3c9cf2a7b3a992084ccd19", + "text": "xt text ", + "embedding": null, + "metadata": {}, + "path": "fallback.md", + "start_line": 2, + "end_line": 2, + "scores": {} + } + ] + }, + "data.json": { + "node": { + "path": "data.json", + "links": [], + "chunk_ids": [ + "cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1", + "562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76", + "147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799", + "79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62", + "352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba", + "04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "cacaea11e2cb2ee3a14b64163ef61cb4033f937a9e1926b8e5c59f938ab883a1", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 0}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 4, + "end_line": 5, + "scores": {} + }, + { + "id": "562530d77b54d6f7ea397db77ca543023f6f082bc9a2d17d0693a9057f6f9c76", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 1}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 8, + "end_line": 9, + "scores": {} + }, + { + "id": "147aa7fe035030c8a7632aefa473adbbe447bcb4ded262234c53c65c120eb799", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 2}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 12, + "end_line": 13, + "scores": {} + }, + { + "id": "79cbd086821831544adee04ddc6e61d4ea8a052c21941710c754f8a8248eed62", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 3}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 16, + "end_line": 17, + "scores": {} + }, + { + "id": "352cc7ee59226ac87644680d8c33f239bc5a8a30b0d088383ebfbf4559ba5bba", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 4}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 20, + "end_line": 21, + "scores": {} + }, + { + "id": "04c5f9658b948d4f9f8eb9cb98f2096e1dd6507b4429211925a7eb5f4aed0f7e", + "text": "{\"items\": [{\"name\": \"abcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabcabc\", \"n\": 5}]}", + "embedding": null, + "metadata": {}, + "path": "data.json", + "start_line": 24, + "end_line": 25, + "scores": {} + } + ] + }, + "bad.json": { + "node": { + "path": "bad.json", + "links": [], + "chunk_ids": [ + "e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "e974bca33c3651cf36a114da6f575a5af54134648f0368b34b9e1d2bc52acbb3", + "text": "{invalid\ntext\n", + "embedding": null, + "metadata": {}, + "path": "bad.json", + "start_line": 1, + "end_line": 3, + "scores": {} + } + ] + }, + "events.jsonl": { + "node": { + "path": "events.jsonl", + "links": [], + "chunk_ids": [ + "fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f", + "7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6", + "48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b", + "116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634", + "c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93", + "489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3" + ], + "front_matter": { + "name": "", + "description": "" + } + }, + "chunks": [ + { + "id": "fe4880a04aedd78ae5743cb495b0d9d8500ab9d3ba79515f4230e1b2f5f7889f", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 1, + "end_line": 3, + "scores": {} + }, + { + "id": "7725431a86471c427c1a123803ccb0218fec1fd05877a6a70902380d50999fc6", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 3, + "end_line": 5, + "scores": {} + }, + { + "id": "48c0edd450850c2382a954df3a907b6ee87f210de7a21a6341797438a30dfd3b", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 5, + "end_line": 7, + "scores": {} + }, + { + "id": "116a24007dabf192a80d4d011a8cbbcf29d6831cc4a0c5941233fadac054e634", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 7, + "end_line": 9, + "scores": {} + }, + { + "id": "c91f51a356df943b80ba09cf654f1a042aa08bbe8912f393f13ac99211b19b93", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 9, + "end_line": 11, + "scores": {} + }, + { + "id": "489e1d724b739ddc864755e3d8e9a8d0c669dd1962102d525056ea265f8915a3", + "text": "{\"value\": \"alpha\"}\n{\"value\": \"alpha\"}\n", + "embedding": null, + "metadata": {}, + "path": "events.jsonl", + "start_line": 11, + "end_line": 12, + "scores": {} + } + ] + } +} diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index 7c5f31d5..1be377ba 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -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,135 @@ 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, + yield_on_timeout: bool, + ) -> AsyncGenerator[set[tuple[Change, str]], None]: + assert yield_on_timeout is False + 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 +1526,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) diff --git a/tests/unit/test_file_chunker_workers.py b/tests/unit/test_file_chunker_workers.py new file mode 100644 index 00000000..f34900c5 --- /dev/null +++ b/tests/unit/test_file_chunker_workers.py @@ -0,0 +1,824 @@ +"""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("original cancellation") + second = asyncio.create_task(chunker.chunk(path)) + await asyncio.sleep(0.01) + first.cancel("second cancellation") + await asyncio.sleep(0.01) + try: + assert not first.done() + assert len(calls) == 1 + finally: + release.set() + with pytest.raises(asyncio.CancelledError, match="original cancellation"): + await first + await second + assert chunker._worker is None + assert len(calls) == 2 + + +@pytest.mark.asyncio +async def test_relative_delete_alias_removes_indexed_content(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A relative-first alias batch resolves against the workspace, not the working directory.""" + workspace = tmp_path / "workspace" + workspace.mkdir() + monkeypatch.chdir(tmp_path) + source = workspace / "source.txt" + source.write_text("unique searchable content") + context = ApplicationContext(workspace_dir=str(workspace)) + 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)]) + assert await store.keyword_search("unique", limit=10, search_filter={}) + source.unlink() + step = UpdateIndexStep(file_store=store, persist=False, app_context=context) + response = await step( + RuntimeContext( + changes=[ + {"change": "deleted", "path": "source.txt"}, + {"change": "deleted", "path": str(source)}, + ], + ), + ) + assert response.success is True + assert await store.get_nodes() == [] + assert not await store.keyword_search("unique", limit=10, search_filter={}) + + +@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() + + +@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" + 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() + + +@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], os.stat_result], 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() + 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], os.stat_result], + 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() + + +@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" + + +@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() + pending_chunks: list[weakref.ReferenceType[FileChunk]] = [] + 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], os.stat_result]: + item = await original(file_path) + pending_chunks.append(weakref.ref(item[1][0])) + 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 + if outcome == "cancel": + await asyncio.sleep(0) + gc.collect() + assert pending_chunks and all(ref() is None for ref in pending_chunks) + 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() diff --git a/tests/unit/test_markdown_file_chunker.py b/tests/unit/test_markdown_file_chunker.py index 73926bf5..ad866452 100644 --- a/tests/unit/test_markdown_file_chunker.py +++ b/tests/unit/test_markdown_file_chunker.py @@ -17,6 +17,7 @@ from unittest.mock import patch import pytest from reme.components.file_chunker import DefaultFileChunker, MarkdownFileChunker +from reme.utils.file_utils import FileChangedDuringChunkingError def _write_md(tmpdir: str, name: str, body: str) -> str: @@ -37,17 +38,22 @@ async def test_chunk_keeps_pre_read_mtime_when_source_changes(tmp_path, monkeypa os.utime(path, (1000, 1000)) chunker = MarkdownFileChunker() - async def read_then_finish_write(_path): + original_read = chunker._read_text_for_indexing + + 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) + with pytest.raises(FileChangedDuringChunkingError): + await chunker.chunk("changing.md") + + monkeypatch.setattr(chunker, "_read_text_for_indexing", original_read) node, chunks = await chunker.chunk("changing.md") - assert chunks == [] - assert node.st_mtime == 1000 - assert path.stat().st_mtime != node.st_mtime + assert node.st_mtime == 2000 + assert "Completed memory note." in "\n".join(chunk.text for chunk in chunks) @pytest.mark.asyncio From e5bbb85e5468162e77bd8fd7972bd7cb66e97159 Mon Sep 17 00:00:00 2001 From: xyf2020 <75460675+xyf2020@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:31:59 +0800 Subject: [PATCH 2/5] fix(python): sandbox Python execution and disable benchmark service exposure (#708) * fix: sandbox Python execution and disable benchmark service exposure * fix: preserve optional sandbox policy for AutoBash * fix: sandbox AutoBash with a read-only workspace by default * fix: preserve sandboxed Python timeout diagnostics * Simplify Python timeout handling and sandbox workspace defaults * fix: require patched core for home and allow sandboxed Apple Git * fix: preserve Python execution workspace directory * fix: cap Python execution timeouts and allow Xcode Git runtime * Fix macOS sandbox Git resolution and add Xcode CI coverage * Make Seatbelt runtime grant test independent of host Xcode symlinks --- .github/workflows/ci-macos-sandbox.yml | 93 +++++ benchmark/beam/README.md | 7 + benchmark/longmemeval/README.md | 4 + plugins/home/README.md | 35 +- plugins/home/pyproject.toml | 2 +- plugins/home/src/reme_home/auto_bash.py | 17 +- plugins/home/src/reme_home/plugin.yaml | 12 +- plugins/home/tests/test_auto_bash.py | 123 +++++-- plugins/home/tests/test_auto_bash_seatbelt.py | 95 +++++ reme/config/benchmark.yaml | 9 +- reme/steps/common/file_sandbox.py | 24 +- reme/steps/common/python_execute.py | 36 +- .../file_sandbox/file_sandbox_backend.py | 17 +- .../file_sandbox/path_seatbelt_backend.py | 103 +++++- reme/utils/local_backend.py | 5 +- tests/integration/test_home_wheel.py | 97 ++++++ tests/integration/test_seatbelt_backend.py | 60 ++++ tests/unit/test_common_steps.py | 10 + tests/unit/test_file_sandbox_backend.py | 324 ++++++++++++++++-- tests/unit/test_local_backend.py | 30 +- tests/unit/test_package_versions.py | 7 +- tests/unit/test_path_bubble_backend.py | 68 +++- tests/unit/test_process_output_limits.py | 6 +- tests/unit/test_python_execute_timeout.py | 96 ++++++ 24 files changed, 1158 insertions(+), 122 deletions(-) create mode 100644 .github/workflows/ci-macos-sandbox.yml create mode 100644 plugins/home/tests/test_auto_bash_seatbelt.py create mode 100644 tests/integration/test_home_wheel.py create mode 100644 tests/unit/test_python_execute_timeout.py diff --git a/.github/workflows/ci-macos-sandbox.yml b/.github/workflows/ci-macos-sandbox.yml new file mode 100644 index 00000000..c587d31a --- /dev/null +++ b/.github/workflows/ci-macos-sandbox.yml @@ -0,0 +1,93 @@ +name: CI / macOS sandbox + +on: + push: + branches: [main] + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + seatbelt: + name: Seatbelt - ${{ matrix.toolchain }} + runs-on: macos-15 + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + toolchain: [clt, xcode] + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + + - name: Configure selected developer tools + shell: bash + env: + TOOLCHAIN: ${{ matrix.toolchain }} + run: | + set -euo pipefail + if [[ "$TOOLCHAIN" == xcode ]]; then + # Hosted images commonly use a symlink named Xcode.app. Test the + # real standard installation supported by the sandbox's contract. + selected=$(/usr/bin/env -u DEVELOPER_DIR /usr/bin/xcode-select -p) + app=$(python3 -c 'import os, sys; print(os.path.realpath(sys.argv[1]))' "${selected%/Contents/Developer}") + [[ "$app" == /Applications/*.app ]] + [[ -x "$app/Contents/Developer/usr/bin/git" ]] + if [[ -L /Applications/Xcode.app ]]; then + sudo unlink /Applications/Xcode.app + fi + if [[ "$app" != /Applications/Xcode.app ]]; then + # Fail rather than overwrite any existing real installation. + [[ ! -e /Applications/Xcode.app ]] + sudo mv "$app" /Applications/Xcode.app + fi + developer=/Applications/Xcode.app/Contents/Developer + [[ -f /Applications/Xcode.app/Contents/Frameworks/libxcodebuildLoader.dylib ]] + [[ -f /Applications/Xcode.app/Contents/version.plist ]] + else + developer=/Library/Developer/CommandLineTools + fi + sudo /usr/bin/xcode-select --switch "$developer" + echo "DEVELOPER_DIR=$developer" >> "$GITHUB_ENV" + /usr/bin/env -i /usr/bin/xcode-select -p + /usr/bin/env -i /usr/bin/git --version + "$developer/usr/bin/git" --version + # Require real Seatbelt and a GUI launchd domain for process cleanup + # tests; missing infrastructure must fail this job, not skip tests. + /usr/bin/sandbox-exec -p '(version 1)(allow default)' /usr/bin/true + /bin/launchctl print "gui/$(id -u)" > /dev/null + + - name: Set up Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.13' + cache: 'pip' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip setuptools wheel + python -m pip install -e reme_studio -e '.[dev,as]' + python -m pip install 'safecmd>=0.1.16,<0.2' shfmt-py + python -m pip install --no-deps -e plugins/home + + - name: Run sandbox and executor regressions + env: + FASTMCP_MCP_CAMELCASE_COMPAT: 'false' + run: | + python -m pytest \ + tests/unit/test_file_sandbox_backend.py \ + tests/unit/test_sandbox_startup_cancellation.py \ + tests/unit/test_local_backend.py \ + tests/unit/test_path_bubble_backend.py \ + plugins/home/tests/test_auto_bash.py \ + tests/integration/test_seatbelt_backend.py \ + plugins/home/tests/test_auto_bash_seatbelt.py \ + -v --tb=long --junitxml=macos-sandbox-results.xml diff --git a/benchmark/beam/README.md b/benchmark/beam/README.md index 287616b0..e8425956 100644 --- a/benchmark/beam/README.md +++ b/benchmark/beam/README.md @@ -31,6 +31,13 @@ Custom application config paths still work through `reme.config` and can use `ex This directory continues to own the runner, evaluation settings, dataset and outputs. Model credentials use the environment variables declared by the shared benchmark configuration. +The shared preset keeps `python_execute` private (`enable_serve: false`); in-process benchmark jobs +remain callable. Python always uses a sandbox with `workspace_dir` read-only by default, no network, +and a minimal environment. Linux requires Bubblewrap and, for the default memory cap, a writable +delegated cgroup v2; macOS requires Seatbelt. Unavailable isolation fails the call. Custom path grants +must be configured on the Python Step, not passed as job arguments. Read-only workspace files, +including any `.env` files, remain readable: keep credentials outside the workspace. + ## 1. Get the Dataset BEAM is a public repository, cloned into `benchmark/beam/dataset/`: diff --git a/benchmark/longmemeval/README.md b/benchmark/longmemeval/README.md index 09185db7..7ddda3cb 100644 --- a/benchmark/longmemeval/README.md +++ b/benchmark/longmemeval/README.md @@ -28,6 +28,10 @@ Custom application config paths still work through `reme.config` and can use `ex This directory continues to own the runner, evaluation settings, dataset and outputs. Model credentials use the environment variables declared by the shared benchmark configuration. +The shared preset keeps `python_execute` private (`enable_serve: false`); in-process benchmark jobs +remain callable. Python sandbox requirements and credential placement are described in the +[BEAM benchmark guide](../beam/README.md). + ## 1. Get the Dataset ReMe uses only the **cleaned-S** split, hosted on HuggingFace: diff --git a/plugins/home/README.md b/plugins/home/README.md index c3470af3..b76f07f1 100644 --- a/plugins/home/README.md +++ b/plugins/home/README.md @@ -9,16 +9,33 @@ reme plugins install plugins/home ``` Enable it in an application config with `plugins: [home]`. The plugin then -contributes the `auto_bash` job, which can be run with -`reme start config= job=auto_bash command=pwd`. +contributes the `auto_bash` job. Start `reme start config=` +in one terminal, then call `reme auto_bash command=pwd` from another. -The job defaults to read-only command validation and a 60-second timeout. Set -`is_read_only=false` to disable command validation. Supplying `read_only_paths` -or `writable_paths` enables a file sandbox (Bubblewrap on Linux, Seatbelt on -macOS), with sandbox paths identical to their host paths. On Linux, the -`memory_limit_mib` cgroup limit defaults to 4096 MiB; 0 disables it. Supplying -the limit requires a path policy. macOS has no memory limit and warns when a -nonzero value is supplied. +The job always uses a file sandbox (Bubblewrap on Linux, Seatbelt on macOS), +with `workspace_dir` read-only by default, network access blocked, and a minimal +child environment. Missing or unavailable isolation fails without host execution. +Files inside the workspace remain readable, including any `.env` files placed there. +Keep credentials outside the workspace. + +The job also defaults to read-only command validation and a 60-second timeout. +`is_read_only=false` disables command validation while keeping the sandbox active. +Configure `read_only_paths` and `writable_paths` only on the Step; request-level +path grants are rejected. Existing path arguments must move to Step configuration. +For example, a trusted Step can grant workspace writes: + +```yaml +steps: + - backend: auto_bash_step + read_only_paths: [] + writable_paths: ["/absolute/path/to/workspace"] +``` + +On Linux, Bubblewrap and a writable delegated cgroup v2 are required. The +`memory_limit_mib` cgroup limit defaults to 4096 MiB; 0 disables it. macOS has +no memory limit and warns when a nonzero value is supplied. Direct `AutoBash` +instances also default to a sandbox with their working directory read-only; +trusted callers may inject a different backend. Standard output and standard error each have a 1 MiB (1,048,576-byte) limit on all backends. Exceeding either limit terminates the command and returns a diff --git a/plugins/home/pyproject.toml b/plugins/home/pyproject.toml index 77f8c174..0da629b4 100644 --- a/plugins/home/pyproject.toml +++ b/plugins/home/pyproject.toml @@ -5,7 +5,7 @@ description = "Local command execution jobs for ReMe." readme = "README.md" requires-python = ">=3.11" dependencies = [ - "reme-ai[as]>=0.4.2.1", + "reme-ai[as]>=0.4.2.2", "safecmd>=0.1.16,<0.2", ] diff --git a/plugins/home/src/reme_home/auto_bash.py b/plugins/home/src/reme_home/auto_bash.py index 0f4a2d25..3f03d452 100644 --- a/plugins/home/src/reme_home/auto_bash.py +++ b/plugins/home/src/reme_home/auto_bash.py @@ -4,6 +4,7 @@ import asyncio import os import shlex from functools import cache +from pathlib import Path from agentscope.message import TextBlock, ToolResultState from agentscope.permission import PermissionBehavior, PermissionContext, PermissionDecision @@ -12,7 +13,7 @@ from agentscope.tool import Bash, ToolChunk from reme.steps.base_step import BaseStep from reme.steps.common.file_sandbox import file_sandbox_for_step -from reme.utils.local_backend import LocalProcessBackend +from reme.utils.file_sandbox import create_file_sandbox_backend _READ_COMMANDS = { ":": {}, @@ -153,7 +154,7 @@ class AutoBash(Bash): Shell expansions, function definitions and for loops are rejected because their runtime arguments or environment can differ from the parsed tokens. Git external filters require unrestricted mode; other Git helpers are - disabled. Use path injection for OS file isolation. + disabled. The default backend isolates the working directory read-only. """ name = "AutoBash" @@ -172,7 +173,13 @@ class AutoBash(Bash): def __init__(self, *, is_read_only: bool = True, **kwargs): if kwargs.get("backend") is None: - kwargs["backend"] = LocalProcessBackend() + workspace = Path(kwargs.get("cwd") or Path.cwd()).expanduser().resolve() + kwargs["cwd"] = workspace + kwargs["backend"] = create_file_sandbox_backend( + cwd=workspace, + read_only_paths=[], + writable_paths=[], + ) super().__init__(**kwargs) self.is_read_only = is_read_only @@ -247,7 +254,7 @@ class AutoBash(Bash): class AutoBashStep(BaseStep): - """Expose AgentScope Bash with configurable read-only mode as a ReMe job.""" + """Run sandboxed Bash with a read-only workspace and optional command validation.""" async def execute(self): """Run the command with the selected permission policy and backend.""" @@ -263,7 +270,7 @@ class AutoBashStep(BaseStep): if not isinstance(is_read_only, bool): raise ValueError("is_read_only must be a boolean") backend = file_sandbox_for_step(self) - except (TypeError, ValueError, RuntimeError) as exc: + except (TypeError, ValueError, OSError, RuntimeError) as exc: self.context.response.success = False self.context.response.answer = str(exc) return self.context.response diff --git a/plugins/home/src/reme_home/plugin.yaml b/plugins/home/src/reme_home/plugin.yaml index 77ad1584..5ef67f39 100644 --- a/plugins/home/src/reme_home/plugin.yaml +++ b/plugins/home/src/reme_home/plugin.yaml @@ -5,7 +5,7 @@ application_defaults: jobs: auto_bash: backend: base - description: "Run a shell command in the workspace; read-only mode is enabled by default" + description: "Run a sandboxed shell command with a read-only workspace and command validation by default" parameters: type: object properties: @@ -22,18 +22,10 @@ application_defaults: default: 60000 minimum: 1 maximum: 600000 - read_only_paths: - type: array - items: { type: string } - description: "Read-only host paths; providing either path list enables the platform file sandbox" - writable_paths: - type: array - items: { type: string } - description: "Writable host paths; sandbox paths match host paths" memory_limit_mib: type: integer minimum: 0 - description: "Linux cgroup memory limit in MiB (default 4096; 0 disables it); requires a path policy; macOS warns and ignores nonzero values" + description: "Linux cgroup memory limit in MiB (default 4096; 0 disables it); macOS warns and ignores nonzero values" required: - command steps: diff --git a/plugins/home/tests/test_auto_bash.py b/plugins/home/tests/test_auto_bash.py index c8f7ddfc..a10f518d 100644 --- a/plugins/home/tests/test_auto_bash.py +++ b/plugins/home/tests/test_auto_bash.py @@ -19,6 +19,7 @@ from reme_home.auto_bash import AutoBash, AutoBashStep from reme.utils.local_backend import LocalProcessBackend from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend +from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend from reme.components.application_context import ApplicationContext from reme.config.config_parser import _load_config from reme.enumeration import ComponentEnum @@ -31,6 +32,16 @@ def linux_backend_selection(monkeypatch): monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", "linux") +@pytest.fixture(autouse=True) +def sandbox_process_boundary(monkeypatch): + """Use real local pipes for command tests; policy tests inspect the selected sandbox.""" + + async def local_exec(_self, command, *, cwd=None, timeout=None): + return await LocalProcessBackend().exec_shell(command, cwd=cwd, timeout=timeout) + + monkeypatch.setattr(PathBubbleBackend, "exec_shell", local_exec) + + @pytest.fixture(autouse=True) def isolated_git_config(tmp_path, monkeypatch): """Keep Git tests independent of the developer's global and system config.""" @@ -38,6 +49,29 @@ def isolated_git_config(tmp_path, monkeypatch): monkeypatch.setenv("GIT_CONFIG_NOSYSTEM", "1") +@pytest.mark.asyncio +@pytest.mark.parametrize("platform", ["linux", "darwin"]) +async def test_direct_auto_bash_defaults_to_read_only_sandbox(tmp_path, monkeypatch, platform): + """Direct tool construction has the same default grants as the job.""" + monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", platform) + tool = AutoBash(cwd=tmp_path, is_read_only=False) + backend = tool._backend # pylint: disable=protected-access + if platform == "linux": + assert isinstance(backend, PathBubbleBackend) + argv = backend._bwrap_argv(["/bin/sh", "-c", "echo ok"], str(tmp_path)) # pylint: disable=protected-access + assert argv[argv.index(str(tmp_path)) - 1] == "--ro-bind" + assert "--bind" not in argv + else: + assert isinstance(backend, PathSeatbeltBackend) + profile = backend._profile(str(tmp_path / "temporary")) # pylint: disable=protected-access + assert f'(allow file-read* (subpath "{tmp_path}"))' in profile + assert f'(allow file-write* (subpath "{tmp_path}"))' not in profile + + monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", "unsupported") + with pytest.raises(RuntimeError, match="No file sandbox backend"): + AutoBash(cwd=tmp_path, is_read_only=False) + + @pytest.mark.asyncio async def test_auto_bash_runs_in_workspace(tmp_path): """Approved commands run relative to the configured workspace.""" @@ -212,8 +246,12 @@ async def test_step_returns_permission_error(tmp_path): @pytest.mark.asyncio async def test_step_can_run_unrestricted_commands(tmp_path): - """Job input can disable the guard for commands that write files.""" - step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + """Trusted Step grants allow writes when command validation is disabled.""" + step = AutoBashStep( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + read_only_paths=[], + writable_paths=[str(tmp_path)], + ) response = await step(command="echo hello > written.txt", is_read_only=False) @@ -222,34 +260,66 @@ async def test_step_can_run_unrestricted_commands(tmp_path): @pytest.mark.asyncio -async def test_unsupported_platform_runs_without_paths_but_rejects_path_policy(tmp_path, monkeypatch): - """AutoBash never drops an explicitly requested sandbox on unsupported systems.""" +async def test_unsupported_platform_never_runs_on_host(tmp_path, monkeypatch): + """AutoBash requires isolation even without an explicit path policy.""" monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", "unsupported") step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - plain = await step(command="echo host", is_read_only=False) - assert plain.success is True - assert plain.answer.strip() == "host" - - restricted = await step(command="echo blocked", is_read_only=False, read_only_paths=[]) - assert restricted.success is False - assert restricted.answer == "No file sandbox backend for unsupported" + result = await step(command="touch host-marker", is_read_only=False) + assert result.success is False + assert result.answer == "No file sandbox backend for unsupported" + assert not (tmp_path / "host-marker").exists() @pytest.mark.asyncio -@pytest.mark.parametrize("limit", [0, 1, -1]) -async def test_step_rejects_memory_limit_without_path_policy(tmp_path, limit): - """An explicitly supplied limit cannot be ignored by host execution.""" +@pytest.mark.parametrize("limit", [None, 0, 1, -1]) +@pytest.mark.parametrize("is_read_only", [False, True]) +async def test_step_defaults_to_read_only_workspace(tmp_path, monkeypatch, limit, is_read_only): + """Sandbox grants and limits apply independently of command validation.""" + seen = [] + + async def inspect_exec(backend, command, *, cwd=None, timeout=None): + assert timeout == 60 + seen.append(backend._bwrap_argv(command, cwd)) # pylint: disable=protected-access + assert backend._memory_limit_mib == (4096 if limit is None else limit) # pylint: disable=protected-access + return ExecResult(exit_code=0, stdout=b"ok", stderr=b"") + + monkeypatch.setattr(PathBubbleBackend, "exec_shell", inspect_exec) step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(command="echo should-not-run", is_read_only=False, memory_limit_mib=limit) + response = await step(command="echo ok", is_read_only=is_read_only, memory_limit_mib=limit) + if limit == -1: + assert response.success is False + assert response.answer == "memory_limit_mib must be a non-negative integer" + assert not seen + else: + assert response.success is True + argv = seen[0] + mount = argv.index(str(tmp_path)) + assert argv[mount - 1] == "--ro-bind" + assert "--bind" not in argv + assert "--clearenv" in argv + assert "--unshare-all" in argv + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", ["read_only_paths", "writable_paths"]) +@pytest.mark.parametrize("value", [None, [], ["/"]]) +async def test_step_rejects_request_path_grants(tmp_path, name, value): + """Request payloads cannot widen or replace trusted sandbox grants.""" + step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + response = await step(command="echo blocked", **{name: value}) assert response.success is False - assert response.answer == "memory_limit_mib requires read_only_paths or writable_paths to enable the sandbox" + assert response.answer == f"{name} must be configured on the Step, not supplied by the caller" @pytest.mark.asyncio async def test_step_configuration_can_disable_read_only(tmp_path): """Step configuration applies when the job does not supply the flag.""" - step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path)), is_read_only=False) + step = AutoBashStep( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + is_read_only=False, + writable_paths=[str(tmp_path)], + ) response = await step(command="touch written.txt") @@ -341,14 +411,20 @@ async def test_auto_bash_uses_path_backend_when_paths_supplied(tmp_path, monkeyp return ExecResult(exit_code=0, stdout=b"ok", stderr=b"") monkeypatch.setattr(PathBubbleBackend, "exec_shell", fake_exec) - step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(command="printf ok", is_read_only=False, read_only_paths=[], writable_paths=[str(tmp_path)]) + step = AutoBashStep( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + read_only_paths=[], + writable_paths=[str(tmp_path)], + ) + response = await step(command="printf ok", is_read_only=False) assert response.success is True assert response.answer == "ok" assert seen["command"] == ["/bin/sh", "-c", "printf ok"] assert seen["cwd"] == str(tmp_path) assert seen["timeout"] == 60 + argv = seen["backend"]._bwrap_argv(seen["command"], seen["cwd"]) # pylint: disable=protected-access + assert argv[argv.index(str(tmp_path)) - 1] == "--bind" @pytest.mark.asyncio @@ -358,12 +434,13 @@ async def test_sandbox_start_failure_does_not_run_on_host(tmp_path, monkeypatch) async def unavailable(*_args, **_kwargs): raise RuntimeError("bwrap unavailable") - monkeypatch.setattr(PathBubbleBackend, "_start_process", unavailable) + monkeypatch.setattr(PathBubbleBackend, "exec_shell", unavailable) step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - result = await step(command="printf host", is_read_only=False, read_only_paths=[]) + result = await step(command="touch host-marker", is_read_only=False) assert result.success is False assert "bwrap unavailable" in result.answer + assert not (tmp_path / "host-marker").exists() def test_home_manifest_declares_step_and_job(): @@ -377,8 +454,8 @@ def test_home_manifest_declares_step_and_job(): assert AutoBash.input_schema["properties"]["timeout"]["default"] == 60000 assert job["parameters"]["properties"]["timeout"]["default"] == 60000 assert job["parameters"]["properties"]["is_read_only"]["default"] is True - assert job["parameters"]["properties"]["read_only_paths"]["type"] == "array" - assert job["parameters"]["properties"]["writable_paths"]["type"] == "array" + assert "read_only_paths" not in job["parameters"]["properties"] + assert "writable_paths" not in job["parameters"]["properties"] assert job["parameters"]["properties"]["memory_limit_mib"]["minimum"] == 0 diff --git a/plugins/home/tests/test_auto_bash_seatbelt.py b/plugins/home/tests/test_auto_bash_seatbelt.py new file mode 100644 index 00000000..50594ef2 --- /dev/null +++ b/plugins/home/tests/test_auto_bash_seatbelt.py @@ -0,0 +1,95 @@ +"""Real macOS execution, deliberately separate from the mocked command tests.""" + +import os +import shlex +import subprocess +import sys + +import pytest +from reme_home.auto_bash import AutoBash, AutoBashStep + +from reme.components.application_context import ApplicationContext +from reme.steps.common.python_execute import PythonExecuteStep +from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend + +pytestmark = pytest.mark.skipif(sys.platform != "darwin", reason="Seatbelt requires macOS") + + +@pytest.mark.asyncio +async def test_auto_bash_git_preserves_sandbox(tmp_path, monkeypatch): + """Apple Git works under the default sandbox without granting external access.""" + root = tmp_path.resolve() + workspace = root / "workspace" + workspace.mkdir() + # Use the Apple dispatcher rather than a developer's Homebrew Git. + subprocess.run(["/usr/bin/git", "init", "-q", str(workspace)], check=True) + (workspace / "note.txt").write_text("hello", encoding="utf-8") + outside = root / "outside.txt" + outside.write_text("private", encoding="utf-8") + monkeypatch.setenv("REME_SANDBOX_TEST_SECRET", "must-not-inherit") + context = ApplicationContext(workspace_dir=str(workspace)) + step = AutoBashStep(app_context=context) + + # The default PATH may select Homebrew Git. Explicitly exercise Apple's + # dispatcher argv too, which the backend resolves to trusted developer tools. + backend = AutoBash(cwd=workspace)._backend # pylint: disable=protected-access + assert isinstance(backend, PathSeatbeltBackend) + for command, expected in (("status", "?? note.txt"), ("rev-parse", str(workspace))): + flag = "--short" if command == "status" else "--show-toplevel" + result = await backend.exec_shell(["/usr/bin/git", command, flag], timeout=10) + assert result.exit_code == 0, result.stderr.decode(errors="replace") + assert result.stderr == b"" + assert expected in result.stdout.decode() + + for read_only in (True, False): + status = await step(command="git status --short", is_read_only=read_only) + assert status.success, status.answer + assert "?? note.txt" in status.answer + toplevel = await step(command="git rev-parse --show-toplevel", is_read_only=read_only) + assert toplevel.success, toplevel.answer + assert str(workspace) in toplevel.answer + + denied = await step(command=f"cat {shlex.quote(str(outside))}") + assert not denied.success + denied = await step(command=f"echo changed > {shlex.quote(str(outside))}", is_read_only=False) + assert not denied.success + assert outside.read_text(encoding="utf-8") == "private" + denied = await step(command="echo changed > note.txt", is_read_only=False) + assert not denied.success + assert (workspace / "note.txt").read_text(encoding="utf-8") == "hello" + denied = await step(command="true", is_read_only=False, writable_paths=[str(root)]) + assert not denied.success + assert "must be configured on the Step" in denied.answer + + # Both public executors share the backend affected by the runtime grants. + code = ( + "import os; from pathlib import Path; " + "assert 'REME_SANDBOX_TEST_SECRET' not in os.environ; " + f"Path({str(outside)!r}).read_text()" + ) + python = PythonExecuteStep(app_context=context) + denied = await python(code=code) + assert not denied.success + assert "PermissionError" in denied.metadata["stderr"] + clean = await step( + command=shlex.join([sys.executable, "-c", "import os; print(os.getenv('REME_SANDBOX_TEST_SECRET', 'clean'))"]), + is_read_only=False, + ) + assert clean.success, clean.answer + assert clean.answer.strip() == "clean" + + # Make the trace destination writable so absence of output proves that + # environment cleanup works independently of filesystem denials. + trace = workspace / "trace" + trace.mkdir() + marker = trace / "git-trace" + monkeypatch.setenv("GIT_TRACE", str(marker)) + step = AutoBashStep(app_context=context, writable_paths=[str(trace)]) + for read_only in (True, False): + result = await step(command="git rev-parse --show-toplevel", is_read_only=read_only) + assert result.success, result.answer + assert not marker.exists() + assert os.environ["GIT_TRACE"] == str(marker) + assert os.environ["REME_SANDBOX_TEST_SECRET"] == "must-not-inherit" + subprocess.run(["/usr/bin/git", "rev-parse", "--show-toplevel"], cwd=workspace, check=True, capture_output=True) + assert marker.stat().st_size > 0 diff --git a/reme/config/benchmark.yaml b/reme/config/benchmark.yaml index 743475ed..79a67ab4 100644 --- a/reme/config/benchmark.yaml +++ b/reme/config/benchmark.yaml @@ -132,7 +132,10 @@ jobs: python_execute: backend: base - description: "Execute Python code and return printed stdout." + enable_serve: false + # Local-only execution; the workspace is read-only unless Step config grants writes. + # Readable workspace files (including .env) are visible to executed code. + description: "Execute sandboxed Python with a read-only workspace and return printed stdout." parameters: type: object properties: @@ -141,11 +144,11 @@ jobs: description: "Python code to execute. Print the final result to stdout." python_timeout: type: number - description: "Execution timeout in seconds; defaults to 60." + description: "Positive execution timeout in seconds; defaults to 60, capped at 86400 or the shorter Step max_python_timeout." memory_limit_mib: type: integer minimum: 0 - description: "Linux memory limit in MiB (default 4096; 0 disables it); requires a path policy; macOS ignores nonzero values" + description: "Linux sandbox memory limit in MiB (default 4096; 0 disables it); macOS ignores nonzero values" required: - code steps: diff --git a/reme/steps/common/file_sandbox.py b/reme/steps/common/file_sandbox.py index fd516899..450cf44a 100644 --- a/reme/steps/common/file_sandbox.py +++ b/reme/steps/common/file_sandbox.py @@ -1,23 +1,23 @@ -"""Resolve an invocation's optional file sandbox from injected job parameters.""" +"""Resolve executor sandbox policies.""" from ...utils.file_sandbox import create_file_sandbox_backend from ...utils.file_sandbox.file_sandbox_backend import FileSandboxBackend from ..base_step import BaseStep -def file_sandbox_for_step(step: BaseStep) -> FileSandboxBackend | None: - """Keep path and memory policy selection identical across executor steps.""" +def file_sandbox_for_step(step: BaseStep) -> FileSandboxBackend: + """Default to a read-only workspace; callers cannot supply path grants.""" assert step.context is not None - read_only_paths = step.context.get("read_only_paths", step.kwargs.get("read_only_paths")) - writable_paths = step.context.get("writable_paths", step.kwargs.get("writable_paths")) + for name in ("read_only_paths", "writable_paths"): + if name in step.context: + raise ValueError(f"{name} must be configured on the Step, not supplied by the caller") + workspace = step.workspace_path.resolve() + read_only_paths = step.kwargs.get("read_only_paths", []) + writable_paths = step.kwargs.get("writable_paths", []) memory_limit_mib = step.context.get("memory_limit_mib", step.kwargs.get("memory_limit_mib")) - if read_only_paths is None and writable_paths is None: - if memory_limit_mib is not None: - raise ValueError("memory_limit_mib requires read_only_paths or writable_paths to enable the sandbox") - return None return create_file_sandbox_backend( - cwd=step.workspace_path, - read_only_paths=[] if read_only_paths is None else read_only_paths, - writable_paths=[] if writable_paths is None else writable_paths, + cwd=workspace, + read_only_paths=read_only_paths, + writable_paths=writable_paths, memory_limit_mib=memory_limit_mib, ) diff --git a/reme/steps/common/python_execute.py b/reme/steps/common/python_execute.py index 14ac974c..7d613524 100644 --- a/reme/steps/common/python_execute.py +++ b/reme/steps/common/python_execute.py @@ -1,27 +1,38 @@ """Execute Python code and return printed stdout.""" +import math import sys from typing import Any -from ...utils.local_backend import LocalProcessBackend, ProcessTimeout +from ...utils.local_backend import ProcessTimeout from ..base_step import BaseStep from ...components import R from .file_sandbox import file_sandbox_for_step DEFAULT_TIMEOUT = 60.0 +MAX_TIMEOUT = 86400.0 @R.register("python_execute_step") class PythonExecuteStep(BaseStep): - """Run Python code and return stdout, limiting each output stream to 1 MiB.""" + """Run sandboxed Python with a read-only workspace and bounded output. + + Timeouts default to 60 seconds and are capped at one day. A trusted Step + configuration can set max_python_timeout to a shorter limit; request data + cannot override that limit. Positive infinity is capped; NaN is rejected. + """ async def execute(self): assert self.context is not None code = self.context.get("code", "") + maximum, maximum_error = self._parse_timeout(self.kwargs.get("max_python_timeout", MAX_TIMEOUT)) timeout, timeout_error = self._parse_timeout( self.context.get("python_timeout", self.context.get("timeout", DEFAULT_TIMEOUT)), + maximum=maximum, ) + if maximum_error: + timeout_error = maximum_error.replace("python_timeout", "max_python_timeout") if not isinstance(code, str) or not code.strip(): self.context.response.success = False self.context.response.answer = "code is required" @@ -34,11 +45,11 @@ class PythonExecuteStep(BaseStep): timed_out = False try: backend = file_sandbox_for_step(self) - result = await (backend or LocalProcessBackend()).exec_shell( + result = await backend.exec_shell( [sys.executable, "-c", code], - cwd=str(self.workspace_path), + cwd=str(self.workspace_path.resolve()), timeout=timeout, - **({"raise_errors": True} if backend is None else {}), + raise_errors=True, ) except ProcessTimeout as exc: result, timed_out = exc.result, True @@ -46,7 +57,6 @@ class PythonExecuteStep(BaseStep): self.context.response.success = False self.context.response.answer = str(exc) return self.context.response - timed_out = timed_out or (result.exit_code == -1 and result.stderr == b"timed out") stdout = result.stdout.decode(errors="replace") stderr = result.stderr.decode(errors="replace") if timed_out: @@ -65,11 +75,19 @@ class PythonExecuteStep(BaseStep): return self.context.response @staticmethod - def _parse_timeout(raw: Any) -> tuple[float, str]: + def _parse_timeout(raw: Any, *, maximum: float = MAX_TIMEOUT) -> tuple[float, str]: try: + if isinstance(raw, bool): + raise ValueError("boolean timeout") timeout = float(raw) + except OverflowError: + # JSON integers can exceed float's range. They still have a known + # sign, so large positive integers follow the normal cap policy. + if isinstance(raw, int) and raw > 0: + return maximum, "" + return DEFAULT_TIMEOUT, "python_timeout must be a positive number" except (TypeError, ValueError): return DEFAULT_TIMEOUT, "python_timeout must be a positive number" - if timeout <= 0: + if math.isnan(timeout) or timeout <= 0: return DEFAULT_TIMEOUT, "python_timeout must be a positive number" - return timeout, "" + return min(timeout, maximum), "" diff --git a/reme/utils/file_sandbox/file_sandbox_backend.py b/reme/utils/file_sandbox/file_sandbox_backend.py index 7926eafd..4dc50ce3 100644 --- a/reme/utils/file_sandbox/file_sandbox_backend.py +++ b/reme/utils/file_sandbox/file_sandbox_backend.py @@ -111,18 +111,27 @@ class FileSandboxBackend(LocalProcessBackend): *, cwd: str | None = None, timeout: float | None = None, + raise_errors: bool = False, ) -> ExecResult: - """Execute an argv list with a timeout inside the platform sandbox.""" - return await self._run_job(command, cwd or self._cwd, data=None, timeout=timeout) + """Execute inside the sandbox; raise_errors preserves output in ProcessTimeout.""" + return await self._run_job(command, cwd or self._cwd, data=None, timeout=timeout, raise_errors=raise_errors) - async def _run_job(self, command: list[str], cwd: str, *, data: bytes | None, timeout: float | None) -> ExecResult: + async def _run_job( + self, + command: list[str], + cwd: str, + *, + data: bytes | None, + timeout: float | None, + raise_errors: bool = False, + ) -> ExecResult: """Run with optional stdin; platform backends may own a different job lifecycle.""" process = await self._start_process( command, cwd, stdin=asyncio.subprocess.PIPE if data is not None else asyncio.subprocess.DEVNULL, ) - return await self._communicate(process, data=data, timeout=timeout) + return await self._communicate(process, data=data, timeout=timeout, raise_errors=raise_errors) async def read_file(self, path: str) -> bytes: """Read through a sandboxed command so file rules still apply.""" diff --git a/reme/utils/file_sandbox/path_seatbelt_backend.py b/reme/utils/file_sandbox/path_seatbelt_backend.py index f2b9ee5a..5fd420ca 100644 --- a/reme/utils/file_sandbox/path_seatbelt_backend.py +++ b/reme/utils/file_sandbox/path_seatbelt_backend.py @@ -21,6 +21,7 @@ from ..async_utils import finish_cleanup from ..local_backend import ( LocalProcessBackend, OutputLimitExceeded, + ProcessTimeout, capture_output, output_limit_result, ) @@ -37,6 +38,8 @@ class LaunchctlError(RuntimeError): class PathSeatbeltBackend(FileSandboxBackend): """Run each command under a deny-by-default Seatbelt profile.""" + _DEVELOPER_ROOTS = ("/Library/Developer/CommandLineTools", "/Applications/Xcode.app/Contents/Developer") + def __init__( self, *, @@ -73,6 +76,28 @@ class PathSeatbeltBackend(FileSandboxBackend): """Allow only configured files plus the system runtime needed to launch tools.""" system_roots = ("/System", "/usr", "/bin", "/sbin", "/private/etc", "/dev") read_paths = [path for path in system_roots if os.path.exists(path)] + # Apple Git dispatches into the installed developer tools. Only trust + # standard runtime locations, never symlinks to arbitrary user data. + read_paths += [path for path in self._DEVELOPER_ROOTS if os.path.isdir(path) and os.path.realpath(path) == path] + # The Apple dispatcher also loads runtime files next to Developer. + # Keep these grants separate from Contents and reject redirected paths. + xcode_runtime = ( + "/Applications/Xcode.app/Contents/Info.plist", + "/Applications/Xcode.app/Contents/version.plist", + "/Applications/Xcode.app/Contents/Frameworks", + "/Applications/Xcode.app/Contents/SharedFrameworks", + ) + read_paths += [path for path in xcode_runtime if os.path.exists(path) and os.path.realpath(path) == path] + # Selectors may be absent: xcode-select must be allowed to probe them + # before falling back to the standard installation. Keep these literal. + runtime_files = ( + "/etc", + "/var", + "/var/select/developer_dir", + "/private/var/select/developer_dir", + "/var/db/xcode_select_link", + "/private/var/db/xcode_select_link", + ) read_paths += self._read_only_paths + self._writable_paths + [tempdir] lines = [ "(version 1)", @@ -84,13 +109,14 @@ class PathSeatbeltBackend(FileSandboxBackend): '(allow file-read* (literal "/") (literal "/private/var/select/sh"))', ] ancestors = set() - for path in read_paths: + for path in read_paths + list(runtime_files): parent = os.path.dirname(path) while parent != "/": ancestors.add(parent) parent = os.path.dirname(parent) lines.extend(f"(allow file-read-metadata (literal {json.dumps(path)}))" for path in sorted(ancestors)) lines.extend(self._runtime_link_rules()) + lines.extend(f"(allow file-read* (literal {json.dumps(path)}))" for path in runtime_files) lines.extend(f"(allow file-read* {self._path_rule(path)})" for path in read_paths) # Shell redirection to the null device is safe and is accepted by # AutoBash's read-only policy; do not grant writes to the rest of /dev. @@ -123,14 +149,60 @@ class PathSeatbeltBackend(FileSandboxBackend): lines.append(f"(deny file-write-create {rule})") return "\n".join(lines) + def _environment(self, tempdir: str, developer_bin: str | None = None) -> dict[str, str]: + """Expose runtime settings only, never service credentials or Python hooks.""" + paths = [str(Path(sys.prefix) / "bin"), "/usr/local/bin"] + if developer_bin is not None: + paths.append(developer_bin) + paths.extend(("/usr/bin", "/bin")) + return { + "HOME": self._cwd, + "TMPDIR": tempdir, + "PATH": ":".join(paths), + } + + async def _prepare_command(self, command: list[str], tempdir: str) -> tuple[list[str], dict[str, str]]: + """Resolve Apple's Git shim before sandboxing, without running user commands on the host.""" + # xcode-select only reports the system selection. In particular, do not + # inherit DEVELOPER_DIR, PATH or loader hooks from the service process. + backend = LocalProcessBackend() + try: + process = await backend._spawn_process( # pylint: disable=protected-access + ["/usr/bin/xcode-select", "--print-path"], + cwd="/", + env={}, + stdin=asyncio.subprocess.DEVNULL, + ) + result = await backend._communicate(process, timeout=5) # pylint: disable=protected-access + except OSError: + return command, self._environment(tempdir) + developer = result.stdout.decode(errors="replace").strip() if result.ok() else "" + git = os.path.join(developer, "usr", "bin", "git") + # Match the runtime grants: selection never authorizes an arbitrary + # directory or a standard installation redirected to user-owned data. + if ( + developer not in self._DEVELOPER_ROOTS + or os.path.realpath(git) != git + or not os.path.isfile(git) + or not os.access(git, os.X_OK) + ): + return command, self._environment(tempdir) + environment = self._environment(tempdir, os.path.dirname(git)) + # Shells find the real Git through PATH; explicit dispatcher argv must + # also bypass xcrun's loader and per-user cache outside sandbox TMPDIR. + if command[0] == "/usr/bin/git": + command = [git, *command[1:]] + return command, environment + async def _start_process(self, command: list[str], cwd: str, *, stdin=asyncio.subprocess.DEVNULL): """Start Seatbelt directly when the GUI launchd domain is unavailable.""" tempdir = tempfile.mkdtemp(prefix="reme-seatbelt-", dir=os.path.realpath(tempfile.gettempdir())) try: + command, environment = await self._prepare_command(command, tempdir) process = await self._spawn_process( ["/usr/bin/sandbox-exec", "-p", self._profile(tempdir), *command], cwd=cwd, - env={**os.environ, "HOME": self._cwd, "TMPDIR": tempdir}, + env=environment, stdin=stdin, ) except BaseException: @@ -202,7 +274,15 @@ class PathSeatbeltBackend(FileSandboxBackend): capture_output(buffer, chunk, name, exceeded) await asyncio.sleep(0) - async def _run_job(self, command: list[str], cwd: str, *, data: bytes | None, timeout: float | None) -> ExecResult: + async def _run_job( + self, + command: list[str], + cwd: str, + *, + data: bytes | None, + timeout: float | None, + raise_errors: bool = False, + ) -> ExecResult: if sys.platform != "darwin": raise RuntimeError("PathSeatbeltBackend requires macOS") if shutil.which("sandbox-exec") is None: @@ -217,7 +297,7 @@ class PathSeatbeltBackend(FileSandboxBackend): RuntimeWarning, stacklevel=2, ) - return await super()._run_job(command, cwd, data=data, timeout=timeout) + return await super()._run_job(command, cwd, data=data, timeout=timeout, raise_errors=raise_errors) control = Path(tempfile.mkdtemp(prefix="reme-seatbelt-", dir=os.path.realpath(tempfile.gettempdir()))) sandbox_tmp = control / "sandbox-tmp" sandbox_tmp.mkdir() @@ -249,6 +329,7 @@ class PathSeatbeltBackend(FileSandboxBackend): pass try: + command, environment = await self._prepare_command(command, str(sandbox_tmp)) for path, buffer, name in ((stdout_path, stdout, "stdout"), (stderr_path, stderr, "stderr")): os.mkfifo(path, 0o600) # Holding both ends open prevents a transient EOF before launchd @@ -269,13 +350,18 @@ class PathSeatbeltBackend(FileSandboxBackend): "-c", script, str(status_path), + # launchd can also supply its own environment. Clear it at + # the payload boundary, rather than merely adding safe keys. + "/usr/bin/env", + "-i", + *(f"{key}={value}" for key, value in environment.items()), "/usr/bin/sandbox-exec", "-p", self._profile(str(sandbox_tmp)), *command, ], "WorkingDirectory": cwd, - "EnvironmentVariables": {**os.environ, "HOME": self._cwd, "TMPDIR": str(sandbox_tmp)}, + "EnvironmentVariables": environment, "StandardInPath": str(stdin_path) if data is not None else "/dev/null", "StandardOutPath": str(stdout_path), "StandardErrorPath": str(stderr_path), @@ -297,6 +383,13 @@ class PathSeatbeltBackend(FileSandboxBackend): try: exit_code = await self._execute(run(), timeout=timeout, cleanup=cleanup) except TimeoutError: + if raise_errors: + # launchd does not expose a subprocess returncode. Use the + # wrapper's status when available, otherwise retain -1. + exit_code = int(status_path.read_text(encoding="ascii")) if status_path.exists() else -1 + raise ProcessTimeout( + ExecResult(exit_code=exit_code, stdout=bytes(stdout), stderr=bytes(stderr)), + ) from None return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") except OutputLimitExceeded as exc: return output_limit_result(bytes(stdout), bytes(stderr), str(exc)) diff --git a/reme/utils/local_backend.py b/reme/utils/local_backend.py index 4feed857..8e67573b 100644 --- a/reme/utils/local_backend.py +++ b/reme/utils/local_backend.py @@ -193,8 +193,9 @@ class LocalProcessBackend(LocalBackend): return output_limit_result(bytes(stdout_buffer), bytes(stderr_buffer), str(exc)) except TimeoutError: if raise_errors: - stdout, stderr = communication.result() if not communication.cancelled() else (b"", b"") - raise ProcessTimeout(ExecResult(exit_code=process.returncode, stdout=stdout, stderr=stderr)) from None + raise ProcessTimeout( + ExecResult(exit_code=process.returncode, stdout=bytes(stdout_buffer), stderr=bytes(stderr_buffer)), + ) from None return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") return ExecResult(exit_code=process.returncode, stdout=stdout, stderr=stderr) diff --git a/tests/integration/test_home_wheel.py b/tests/integration/test_home_wheel.py new file mode 100644 index 00000000..eb87d742 --- /dev/null +++ b/tests/integration/test_home_wheel.py @@ -0,0 +1,97 @@ +"""Offline installed-wheel smoke test; set REME_TEST_WHEEL_DIR to built release wheels.""" + +import os +import subprocess +import sys +import venv +from pathlib import Path + +import pytest + + +@pytest.mark.skipif(sys.platform != "darwin", reason="Real sandbox smoke test requires macOS") +def test_home_wheel_with_minimum_core(tmp_path): + """Built home and minimum-core wheels enforce policy outside the source tree.""" + directory = os.environ.get("REME_TEST_WHEEL_DIR") + if not directory: + pytest.skip("Set REME_TEST_WHEEL_DIR to a directory containing built core and home wheels") + wheels = Path(directory).resolve() + core = list(wheels.glob("reme_ai-0.4.2.2-*.whl")) + home = list(wheels.glob("reme_home-*.whl")) + assert len(core) == len(home) == 1, "Provide exactly one minimum-core wheel and one home wheel" + environment = tmp_path / "venv" + # Reuse third-party test dependencies, but install both packages into a fresh + # venv. No network or dependency upgrades are needed for this artifact test. + venv.EnvBuilder(with_pip=True, system_site_packages=True).create(environment) + interpreter = str(environment / "bin" / "python") + env = {key: value for key, value in os.environ.items() if key != "PYTHONPATH"} + subprocess.run( + [ + interpreter, + "-m", + "pip", + "install", + "--no-index", + "--no-deps", + "--force-reinstall", + str(core[0]), + str(home[0]), + ], + cwd=tmp_path, + env=env, + check=True, + capture_output=True, + text=True, + timeout=60, + ) + script = """ +import asyncio +import os +import shlex +import sys +from importlib.metadata import distribution, version +from pathlib import Path +from packaging.requirements import Requirement +import reme +import reme_home +from reme.components.application_context import ApplicationContext +from reme_home.auto_bash import AutoBashStep + +assert version('reme-ai') == '0.4.2.2' +for module in (reme, reme_home): + assert Path(module.__file__).is_relative_to(sys.prefix), module.__file__ +dist = distribution('reme-home') +requirement = next(Requirement(item) for item in dist.requires if Requirement(item).name == 'reme-ai') +assert requirement.extras == {'as'} +assert '0.4.2.1' not in requirement.specifier +assert '0.4.2.2' in requirement.specifier +next(point for point in dist.entry_points if point.group == 'reme.plugins' and point.name == 'home').load() + +async def main(): + workspace = Path('workspace').resolve() + workspace.mkdir() + outside = Path('outside.txt').resolve() + outside.write_text('private') + step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(workspace))) + denied = await step(command='true', is_read_only=False, writable_paths=[str(workspace.parent)]) + assert not denied.success and 'must be configured on the Step' in denied.answer + denied = await step(command=f'echo changed > {shlex.quote(str(outside))}', is_read_only=False) + assert not denied.success, denied.answer + assert outside.read_text() == 'private' + os.environ['REME_WHEEL_TEST_SECRET'] = 'must-not-inherit' + command = shlex.join([sys.executable, '-c', "import os; print(os.getenv('REME_WHEEL_TEST_SECRET', 'clean'))"]) + result = await step(command=command, is_read_only=False) + assert result.success and result.answer.strip() == 'clean', result.answer + +asyncio.run(main()) +""" + result = subprocess.run( + [interpreter, "-c", script], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/integration/test_seatbelt_backend.py b/tests/integration/test_seatbelt_backend.py index 8a66379b..792be347 100644 --- a/tests/integration/test_seatbelt_backend.py +++ b/tests/integration/test_seatbelt_backend.py @@ -2,6 +2,8 @@ import asyncio import os +import shlex +import subprocess import sys import sysconfig import tempfile @@ -17,6 +19,64 @@ from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend pytestmark = pytest.mark.skipif(sys.platform != "darwin", reason="Seatbelt requires macOS") +@pytest.mark.asyncio +@pytest.mark.parametrize("direct", [False, True], ids=["auto-launch", "direct"]) +async def test_seatbelt_apple_git_preserves_isolation(tmp_path, monkeypatch, direct): + """Resolve Apple's Git to real tools, with isolated files and temporary storage.""" + root = tmp_path.resolve() + workspace = root / "workspace" + workspace.mkdir() + subprocess.run(["/usr/bin/git", "init", "-q", str(workspace)], check=True, capture_output=True) + (workspace / "note.txt").write_text("hello", encoding="utf-8") + outside = root / "outside.txt" + outside.write_text("private", encoding="utf-8") + monkeypatch.setenv("REME_SANDBOX_TEST_SECRET", "must-not-inherit") + monkeypatch.setenv("DEVELOPER_DIR", str(root / "untrusted-developer")) + backend = PathSeatbeltBackend(cwd=workspace, read_only_paths=[], writable_paths=[]) + if direct: + + async def missing(_domain): + return False + + monkeypatch.setattr(backend, "_has_gui_launchd", missing) + + for args, expected in ((["status", "--short"], "?? note.txt"), (["rev-parse", "--show-toplevel"], str(workspace))): + for command in (["/usr/bin/git", *args], ["git", *args], ["/bin/sh", "-c", shlex.join(["git", *args])]): + result = await backend.exec_shell(command, timeout=10) + assert result.exit_code == 0, result.stderr.decode(errors="replace") + assert result.stderr == b"" + assert expected in result.stdout.decode() + for args in (["--version"], ["--exec-path"], ["hash-object", "note.txt"]): + result = await backend.exec_shell(["/usr/bin/git", *args], timeout=10) + assert result.exit_code == 0, result.stderr.decode(errors="replace") + assert result.stderr == b"" + assert result.stdout.strip() + + code = f""" +import os +from pathlib import Path +assert os.environ['HOME'] == {str(workspace)!r} +assert 'REME_SANDBOX_TEST_SECRET' not in os.environ +assert 'DEVELOPER_DIR' not in os.environ +temporary = Path(os.environ['TMPDIR']) +(temporary / 'cache').write_text('temporary') +for path, operation in [({str(outside)!r}, 'read'), ({str(outside)!r}, 'write'), ('note.txt', 'write')]: + try: + Path(path).read_text() if operation == 'read' else Path(path).write_text('denied') + except PermissionError: + pass + else: + raise AssertionError(f'{{operation}} unexpectedly allowed: {{path}}') +print(temporary) +""" + result = await backend.exec_shell([sys.executable, "-c", code], timeout=10) + assert result.exit_code == 0, result.stderr.decode(errors="replace") + assert not Path(result.stdout.decode().strip()).exists() + assert outside.read_text(encoding="utf-8") == "private" + assert (workspace / "note.txt").read_text(encoding="utf-8") == "hello" + assert os.environ["REME_SANDBOX_TEST_SECRET"] == "must-not-inherit" + + @pytest.mark.asyncio @pytest.mark.parametrize( "command", diff --git a/tests/unit/test_common_steps.py b/tests/unit/test_common_steps.py index b63c5f69..02c8a870 100644 --- a/tests/unit/test_common_steps.py +++ b/tests/unit/test_common_steps.py @@ -17,11 +17,21 @@ from reme.steps.common.health_check import _file_graph_status from reme.steps.common.llm_demo import LLMDemoStep from reme.steps.common.python_execute import PythonExecuteStep from reme.steps.index import traverse as traverse_mod +from reme.utils.local_backend import LocalProcessBackend warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") +@pytest.fixture(autouse=True) +def python_process_backend(monkeypatch): + """Test executor responses with real pipes; sandbox policy is tested separately.""" + monkeypatch.setattr( + "reme.steps.common.python_execute.file_sandbox_for_step", + lambda _step: LocalProcessBackend(), + ) + + def _node( path: str, links: list[tuple[str, str | None]] | None = None, diff --git a/tests/unit/test_file_sandbox_backend.py b/tests/unit/test_file_sandbox_backend.py index 0b46bde6..c261e7ea 100644 --- a/tests/unit/test_file_sandbox_backend.py +++ b/tests/unit/test_file_sandbox_backend.py @@ -11,7 +11,7 @@ from agentscope.tool import ExecResult from reme.components.application_context import ApplicationContext from reme.steps.common.python_execute import PythonExecuteStep -from reme.utils.local_backend import LocalProcessBackend +from reme.utils.local_backend import LocalProcessBackend, ProcessTimeout from reme.utils.file_sandbox import create_file_sandbox_backend from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend @@ -38,18 +38,16 @@ def test_platform_selection_and_memory_defaults(tmp_path, monkeypatch): @pytest.mark.asyncio -async def test_unsupported_platform_requires_backend_only_for_path_policy(tmp_path, monkeypatch): - """An unsupported platform runs ordinary Python but rejects a requested sandbox.""" +async def test_unsupported_platform_never_runs_python_on_host(tmp_path, monkeypatch): + """Default execution fails closed when the platform has no sandbox.""" monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", "unsupported") step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - plain = await step(code="print('host')") - assert plain.success is True - assert plain.answer == "host\n" - - restricted = await step(code="print('blocked')", read_only_paths=[]) + marker = tmp_path / "host-execution" + restricted = await step(code=f"from pathlib import Path; Path({str(marker)!r}).touch()") assert restricted.success is False assert restricted.answer == "No file sandbox backend for unsupported" + assert not marker.exists() def test_seatbelt_profile_limits_file_access(tmp_path): @@ -67,6 +65,178 @@ def test_seatbelt_profile_limits_file_access(tmp_path): assert "network*" not in profile +@pytest.mark.parametrize("redirected", [False, True]) +@pytest.mark.parametrize( + "runtime", + ["/Library/Developer/CommandLineTools", "/Applications/Xcode.app/Contents/Developer"], +) +def test_seatbelt_developer_runtime_grants_are_read_only(tmp_path, monkeypatch, redirected, runtime): + """Standard developer tools are readable, but symlinks cannot expose other trees.""" + original_isdir = os.path.isdir + original_realpath = os.path.realpath + monkeypatch.setattr(os.path, "isdir", lambda path: path == runtime or original_isdir(path)) + monkeypatch.setattr( + os.path, + "realpath", + lambda path: (str(tmp_path) if redirected else runtime) if path == runtime else original_realpath(path), + ) + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + profile = backend._profile(str(tmp_path / "temporary")) # pylint: disable=protected-access + + assert (f'(allow file-read* (subpath "{runtime}"))' in profile) is not redirected + assert f'(allow file-write* (subpath "{runtime}"))' not in profile + for path in ("/Library", "/var", "/private/var"): + assert f'(allow file-read* (subpath "{path}"))' not in profile + assert f'(allow file-write* (subpath "{path}"))' not in profile + for path in ("/etc", "/var", "/var/select/developer_dir", "/private/var/db/xcode_select_link"): + assert f'(allow file-read* (literal "{path}"))' in profile + assert '(allow file-read-metadata (literal "/private/var/select"))' in profile + + +@pytest.mark.parametrize("redirected", [False, True]) +def test_seatbelt_developer_runtime_grants_are_narrow(tmp_path, monkeypatch, redirected): + """Xcode's sibling runtime is readable, without trusting symlinked installations.""" + directories = { + "/Library/Developer/CommandLineTools", + "/Applications/Xcode.app/Contents/Developer", + "/Applications/Xcode.app/Contents/Frameworks", + "/Applications/Xcode.app/Contents/SharedFrameworks", + } + runtime = directories | { + "/Applications/Xcode.app/Contents/Info.plist", + "/Applications/Xcode.app/Contents/version.plist", + } + exists, isdir, realpath = os.path.exists, os.path.isdir, os.path.realpath + monkeypatch.setattr(os.path, "exists", lambda path: path in runtime or exists(path)) + monkeypatch.setattr(os.path, "isdir", lambda path: path in directories or isdir(path)) + # Model both canonical and redirected installations independently of the + # host's Xcode.app symlink (common on GitHub's macOS runners). + monkeypatch.setattr( + os.path, + "realpath", + lambda path: (str(tmp_path / "untrusted") if redirected else path) if path in runtime else realpath(path), + ) + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + profile = backend._profile(str(tmp_path / "temporary")) # pylint: disable=protected-access + + for path in runtime: + rule = f'(subpath "{path}")' if path in directories else f'(literal "{path}")' + assert (f"(allow file-read* {rule})" in profile) is not redirected + assert f"(allow file-write* {rule})" not in profile + for parent in ("/Applications", "/Applications/Xcode.app/Contents", "/Library/Developer", "/private/var/db"): + assert f'(allow file-read* (subpath "{parent}"))' not in profile + assert f'(allow file-write* (subpath "{parent}"))' not in profile + assert '(allow file-read* (literal "/private/var/db/xcode_select_link"))' in profile + + +def test_seatbelt_environment_drops_host_values(tmp_path, monkeypatch): + """The environment is an allowlist, including for future secret and hook names.""" + values = { + "REME_SANDBOX_TEST_SECRET": "secret", + "PYTHONPATH": str(tmp_path), + "PYTHONSTARTUP": str(tmp_path / "hook.py"), + "DYLD_INSERT_LIBRARIES": str(tmp_path / "hook.dylib"), + "DEVELOPER_DIR": str(tmp_path / "developer"), + "PATH": str(tmp_path / "untrusted"), + "GIT_TRACE_FUTURE": str(tmp_path / "trace"), + } + for name, value in values.items(): + monkeypatch.setenv(name, value) + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + environment = backend._environment(str(tmp_path / "temporary")) # pylint: disable=protected-access + assert environment == { + "HOME": str(tmp_path), + "TMPDIR": str(tmp_path / "temporary"), + "PATH": f"{sys.prefix}/bin:/usr/local/bin:/usr/bin:/bin", + } + assert all(os.environ[name] == value for name, value in values.items()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "developer,usable", + [ + ("/Library/Developer/CommandLineTools", True), + ("/Applications/Xcode.app/Contents/Developer", True), + ("/Applications/Other.app/Contents/Developer", False), + ("/tmp/developer", False), + ("relative/developer", False), + ("", False), + ], +) +@pytest.mark.parametrize("redirected", [False, True]) +async def test_seatbelt_resolves_only_trusted_selected_git(tmp_path, monkeypatch, developer, usable, redirected): + """Host selection is sanitized, and only an unredirected runtime can replace Apple's shim.""" + from unittest.mock import AsyncMock + + monkeypatch.setenv("DEVELOPER_DIR", "/untrusted/developer") + monkeypatch.setenv("DYLD_INSERT_LIBRARIES", "/untrusted/loader") + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + process = object() + spawn = AsyncMock(return_value=process) + communicate = AsyncMock(return_value=ExecResult(exit_code=0, stdout=f"{developer}\n".encode(), stderr=b"")) + monkeypatch.setattr(LocalProcessBackend, "_spawn_process", spawn) + monkeypatch.setattr(LocalProcessBackend, "_communicate", communicate) + git = f"{developer}/usr/bin/git" + realpath = os.path.realpath + monkeypatch.setattr(os.path, "realpath", lambda path: "/untrusted/git" if redirected and path == git else path) + monkeypatch.setattr(os.path, "isfile", lambda path: path == git) + monkeypatch.setattr(os, "access", lambda path, _mode: path == git) + temporary = str(tmp_path / "temporary") + command = ["/usr/bin/git", "status", "--short"] + + prepared, environment = await backend._prepare_command(command, temporary) # pylint: disable=protected-access + + spawn.assert_awaited_once_with( + ["/usr/bin/xcode-select", "--print-path"], + cwd="/", + env={}, + stdin=asyncio.subprocess.DEVNULL, + ) + communicate.assert_awaited_once_with(process, timeout=5) + selected = usable and not redirected + assert prepared == ([git, "status", "--short"] if selected else command) + expected_bin = f"{developer}/usr/bin:" if selected else "" + assert environment == { + "HOME": str(tmp_path), + "TMPDIR": temporary, + "PATH": f"{sys.prefix}/bin:/usr/local/bin:{expected_bin}/usr/bin:/bin", + } + assert command == ["/usr/bin/git", "status", "--short"] + # A shell's script and quoting remain intact; it finds Git through PATH. + shell = ["/bin/sh", "-c", "git rev-parse --show-toplevel"] + prepared, _ = await backend._prepare_command(shell, temporary) # pylint: disable=protected-access + assert prepared == shell + monkeypatch.setattr(os.path, "realpath", realpath) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["missing", "timeout", "failed", "not-executable", "not-file"]) +async def test_seatbelt_git_selection_failure_preserves_sandbox(tmp_path, monkeypatch, failure): + """Unavailable tools never turn into arbitrary host execution or broader path grants.""" + from unittest.mock import AsyncMock + + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + spawn = AsyncMock(side_effect=FileNotFoundError() if failure == "missing" else None) + developer = "/Library/Developer/CommandLineTools" + result = ExecResult( + exit_code=0 if failure in {"not-executable", "not-file"} else -1, + stdout=f"{developer}\n".encode(), + stderr=b"timed out" if failure == "timeout" else b"", + ) + monkeypatch.setattr(LocalProcessBackend, "_spawn_process", spawn) + monkeypatch.setattr(LocalProcessBackend, "_communicate", AsyncMock(return_value=result)) + monkeypatch.setattr(os.path, "isfile", lambda _: failure != "not-file") + monkeypatch.setattr(os, "access", lambda *_: failure != "not-executable") + command = ["/usr/bin/git", "status", "--short"] + temporary = str(tmp_path / "temporary") + + prepared, environment = await backend._prepare_command(command, temporary) # pylint: disable=protected-access + + assert prepared == command + assert environment == backend._environment(temporary) # pylint: disable=protected-access + + def test_sandbox_reads_virtual_environment_base_runtime(tmp_path, monkeypatch): """A virtual environment grants its real base interpreter without widening writes.""" workspace = tmp_path / "workspace" @@ -157,6 +327,8 @@ async def test_sandbox_file_io_round_trip(tmp_path, monkeypatch, data): async def test_seatbelt_launch_failure_cleans_up_without_host_fallback(tmp_path, monkeypatch): """A failed sandbox launcher never starts the requested command directly.""" monkeypatch.setattr(sys, "platform", "darwin") + monkeypatch.setenv("REME_TEST_SECRET", "must-not-be-inherited") + monkeypatch.setenv("PYTHONPATH", "/untrusted/python") monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.shutil.which", lambda _: "/usr/bin/sandbox-exec") seen = {} @@ -183,30 +355,110 @@ async def test_seatbelt_launch_failure_cleans_up_without_host_fallback(tmp_path, assert seen["arguments"][0] == "bootstrap" assert seen["job"]["AbandonProcessGroup"] is False assert seen["job"]["ProgramArguments"][-2:] == ["/bin/echo", "unsafe"] + assert "REME_TEST_SECRET" not in seen["job"]["EnvironmentVariables"] + assert "PYTHONPATH" not in seen["job"]["EnvironmentVariables"] + argv = seen["job"]["ProgramArguments"] + assert argv[argv.index("/usr/bin/env") + 1] == "-i" + assert not any("must-not-be-inherited" in argument for argument in argv) assert not Path(seen["arguments"][-1]).exists() @pytest.mark.asyncio async def test_python_job_selects_seatbelt_on_macos(tmp_path, monkeypatch): - """Path injection reaches the macOS backend with no implicit memory cap.""" + """Default execution selects macOS isolation with a read-only workspace.""" monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", "darwin") seen = {} - async def fake_exec(self, command, *, cwd=None, timeout=None): + async def fake_run_job(self, command, cwd, *, data, timeout, raise_errors=False): + assert data is None + assert raise_errors is True seen.update(backend=self, command=command, cwd=cwd, timeout=timeout) return ExecResult(exit_code=0, stdout=b"ok\n", stderr=b"") - monkeypatch.setattr(PathSeatbeltBackend, "exec_shell", fake_exec) + monkeypatch.setattr(PathSeatbeltBackend, "_run_job", fake_run_job) step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(code="print('ok')", read_only_paths=[]) + response = await step(code="print('ok')") assert response.success is True assert isinstance(seen["backend"], PathSeatbeltBackend) assert seen["command"] == [sys.executable, "-c", "print('ok')"] + assert seen["cwd"] == str(tmp_path) + assert seen["timeout"] == 60 + profile = seen["backend"]._profile(str(tmp_path / "temporary")) # pylint: disable=protected-access + assert f'(allow file-read* (subpath "{tmp_path}"))' in profile + assert f'(allow file-write* (subpath "{tmp_path}"))' not in profile @pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["success", "failure", "timeout", "cancel", "error", "kickstart_error", "write"]) -async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode): +@pytest.mark.parametrize("platform", ["linux", "darwin"]) +@pytest.mark.parametrize("grant", ["default", "workspace", "output"]) +async def test_python_workspace_grants(tmp_path, monkeypatch, platform, grant): + """Backend completion keeps defaults read-only and accepts explicit write grants.""" + monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", platform) + output = tmp_path / "output" + output.mkdir() + writable = [] if grant == "default" else [str(tmp_path if grant == "workspace" else output)] + backend_class = PathBubbleBackend if platform == "linux" else PathSeatbeltBackend + + async def inspect_exec(backend, _command, **kwargs): + assert kwargs["raise_errors"] is True + assert backend._writable_paths == writable # pylint: disable=protected-access + read_only = backend._read_only_paths # pylint: disable=protected-access + assert (str(tmp_path) in read_only) is (grant != "workspace") + return ExecResult(exit_code=0, stdout=b"ok\n", stderr=b"") + + monkeypatch.setattr(backend_class, "exec_shell", inspect_exec) + step = PythonExecuteStep( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + **({"writable_paths": writable} if writable else {}), + ) + response = await step(code="print('ok')") + assert response.success is True + assert response.answer == "ok\n" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("platform", ["linux", "darwin"]) +@pytest.mark.parametrize("name", ["read_only_paths", "writable_paths"]) +async def test_python_rejects_request_path_grants(tmp_path, monkeypatch, platform, name): + """Neither platform accepts request-level path authorization.""" + monkeypatch.setattr("reme.utils.file_sandbox.sys.platform", platform) + step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + response = await step(code="print('blocked')", **{name: [str(tmp_path)]}) + assert response.success is False + assert response.answer == f"{name} must be configured on the Step, not supplied by the caller" + + +@pytest.mark.asyncio +async def test_seatbelt_direct_execution_does_not_inherit_service_environment(tmp_path, monkeypatch): + """The non-launchd path strips credentials and interpreter injection settings too.""" + monkeypatch.setenv("REME_TEST_SECRET", "must-not-be-inherited") + monkeypatch.setenv("PYTHONPATH", "/untrusted/python") + monkeypatch.setenv("DYLD_INSERT_LIBRARIES", "/untrusted/library") + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + seen = {} + + async def spawn(_command, **kwargs): + seen.update(kwargs) + raise RuntimeError("stop before spawning") + + monkeypatch.setattr(backend, "_spawn_process", spawn) + with pytest.raises(RuntimeError, match="stop before spawning"): + await backend._start_process( # pylint: disable=protected-access + [sys.executable, "-c", "print('ok')"], + str(tmp_path), + ) + assert set(seen["env"]) == {"HOME", "TMPDIR", "PATH"} + assert seen["env"]["HOME"] == str(tmp_path) + assert not Path(seen["env"]["TMPDIR"]).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raise_errors", [False, True]) +@pytest.mark.parametrize( + "mode", + ["success", "failure", "timeout", "timeout_status", "cancel", "error", "kickstart_error", "write"], +) +async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode, raise_errors): """Every launchd outcome unregisters the job and removes its control files.""" from unittest.mock import AsyncMock, Mock @@ -235,13 +487,15 @@ async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode): # Output is collected after launchd has reaped the process group. (control / "stdout").write_bytes(b"output") (control / "stderr").write_bytes(b"diagnostic") + if mode == "timeout_status": + (control / "status").write_text("137\n", encoding="ascii") return "12345" async def wait_status(*_args): started.set() if mode == "error": raise RuntimeError("status failed") - if mode in ("timeout", "cancel"): + if mode in ("timeout", "timeout_status", "cancel"): await asyncio.Event().wait() return 7 if mode == "failure" else 0 @@ -250,7 +504,11 @@ async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode): operation = ( backend.write_file(str(tmp_path / "note.bin"), b"binary\x00\xff") if mode == "write" - else backend.exec_shell(["true"], timeout=0.01 if mode == "timeout" else None) + else backend.exec_shell( + ["true"], + timeout=0.01 if mode.startswith("timeout") else None, + raise_errors=raise_errors, + ) ) task = asyncio.create_task(operation) if mode == "cancel": @@ -261,23 +519,49 @@ async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode): elif mode in ("error", "kickstart_error"): with pytest.raises(RuntimeError, match="failed"): await asyncio.wait_for(task, timeout=1) + elif mode.startswith("timeout") and raise_errors: + with pytest.raises(ProcessTimeout) as caught: + await asyncio.wait_for(task, timeout=1) + assert caught.value.result == ExecResult( + exit_code=137 if mode == "timeout_status" else -1, + stdout=b"output", + stderr=b"diagnostic", + ) else: result = await asyncio.wait_for(task, timeout=1) if mode == "write": assert result is None - elif mode == "timeout": + elif mode.startswith("timeout"): assert result == ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") else: assert result.exit_code == (7 if mode == "failure" else 0) assert result.stdout == b"output" and result.stderr == b"diagnostic" assert calls == ["bootstrap", "kickstart", "bootout"] assert not control.exists() - assert killed.call_count == int(mode in ("timeout", "cancel", "error")) + assert killed.call_count == int(mode in ("timeout", "timeout_status", "cancel", "error")) @pytest.mark.asyncio -async def test_python_timeout_preserves_exit_status_and_stderr(tmp_path): - """The shared executor retains the existing unsandboxed Python metadata.""" +@pytest.mark.parametrize("backend_class", [PathBubbleBackend, PathSeatbeltBackend]) +async def test_python_timeout_preserves_exit_status_and_stderr(tmp_path, monkeypatch, backend_class): + """Real pipes exercise sandbox timeout propagation without OS isolation requirements.""" + from unittest.mock import AsyncMock + + backend = backend_class(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + + async def start(command, cwd, *, stdin): + return await LocalProcessBackend()._start_process(command, cwd, stdin=stdin) # pylint: disable=protected-access + + monkeypatch.setattr(backend, "_start_process", start) + if backend_class is PathSeatbeltBackend: + monkeypatch.setattr(sys, "platform", "darwin") + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.shutil.which", lambda _: "sandbox-exec") + monkeypatch.setattr(backend, "_has_gui_launchd", AsyncMock(return_value=False)) + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.warnings.warn", lambda *a, **kw: None) + monkeypatch.setattr( + "reme.steps.common.python_execute.file_sandbox_for_step", + lambda _step: backend, + ) step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) response = await step( code="import sys, time; print('diagnostic', file=sys.stderr, flush=True); time.sleep(60)", diff --git a/tests/unit/test_local_backend.py b/tests/unit/test_local_backend.py index d6cb10c9..84892dcc 100644 --- a/tests/unit/test_local_backend.py +++ b/tests/unit/test_local_backend.py @@ -2,12 +2,40 @@ import asyncio import os +from types import SimpleNamespace import anyio import psutil import pytest -from reme.utils.local_backend import LocalProcessBackend, finish_cleanup +from reme.utils.local_backend import LocalProcessBackend, ProcessTimeout, finish_cleanup + + +@pytest.mark.asyncio +async def test_timeout_retains_output_when_pipe_draining_is_cancelled(monkeypatch): + """Bounded cleanup must retain bytes captured before an inherited pipe stalls.""" + from unittest.mock import AsyncMock, Mock + + backend = LocalProcessBackend() + process = SimpleNamespace(returncode=-9, _transport=Mock(), wait=AsyncMock(return_value=-9)) + + async def collect(_process, _data, stdout, stderr, _exceeded): + stdout.extend(b"partial output") + stderr.extend(b"diagnostic") + await asyncio.Event().wait() + + monkeypatch.setattr(backend, "_collect_output", collect) + monkeypatch.setattr(backend, "_kill_process", lambda _process: None) + monkeypatch.setattr("reme.utils.local_backend._CLEANUP_TIMEOUT", 0.01) + with pytest.raises(ProcessTimeout) as caught: + await asyncio.wait_for( + backend._communicate(process, timeout=0.02, raise_errors=True), # pylint: disable=protected-access + timeout=1, + ) + assert caught.value.result.exit_code == -9 + assert caught.value.result.stdout == b"partial output" + assert caught.value.result.stderr == b"diagnostic" + process._transport.close.assert_called_once() # pylint: disable=protected-access @pytest.mark.asyncio diff --git a/tests/unit/test_package_versions.py b/tests/unit/test_package_versions.py index 559e62a9..ac8530ae 100644 --- a/tests/unit/test_package_versions.py +++ b/tests/unit/test_package_versions.py @@ -27,15 +27,16 @@ bump_version = _load_script("bump_version") package_studio = _load_script("package_studio") -def test_home_requires_first_sandbox_release() -> None: - """Home must reject published ReMe wheels that predate its sandbox imports.""" +def test_home_requires_complete_sandbox_fix() -> None: + """Home must reject cores without request path rejection and environment isolation.""" config = tomllib.loads((REPOSITORY / "plugins" / "home" / "pyproject.toml").read_text(encoding="utf-8")) requirement = next(Requirement(value) for value in config["project"]["dependencies"] if value.startswith("reme-ai")) assert requirement.extras == {"as"} assert Version("0.4.1.14") not in requirement.specifier assert Version("0.4.2.0") not in requirement.specifier - assert Version("0.4.2.1") in requirement.specifier + assert Version("0.4.2.1") not in requirement.specifier + assert Version("0.4.2.2") in requirement.specifier assert Version(bump_version.read_version()) in requirement.specifier diff --git a/tests/unit/test_path_bubble_backend.py b/tests/unit/test_path_bubble_backend.py index 03695788..db3a33de 100644 --- a/tests/unit/test_path_bubble_backend.py +++ b/tests/unit/test_path_bubble_backend.py @@ -14,6 +14,8 @@ from agentscope.tool import ExecResult from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend from reme.components.application_context import ApplicationContext +from reme.components.job import BaseJob +from reme.components.service import HttpService, MCPService from reme.steps.common.python_execute import PythonExecuteStep @@ -274,17 +276,19 @@ def test_path_bubble_backend_rejects_replaced_mount_source(tmp_path): @pytest.mark.asyncio @pytest.mark.parametrize("memory_options", [{}, {"memory_limit_mib": 0}]) -async def test_python_execute_uses_path_backend_when_paths_supplied(tmp_path, monkeypatch, memory_options): - """Python execution uses the same backend and timeout units.""" +async def test_python_execute_defaults_to_read_only_sandbox(tmp_path, monkeypatch, memory_options): + """Default Python execution has no host write grants and uses the sandbox.""" seen = {} - async def fake_exec(self, command, *, cwd=None, timeout=None): + async def fake_run_job(self, command, cwd, *, data, timeout, raise_errors=False): + assert data is None + assert raise_errors is True seen.update(command=command, cwd=cwd, timeout=timeout, backend=self) return ExecResult(exit_code=0, stdout=b"ok\n", stderr=b"") - monkeypatch.setattr(PathBubbleBackend, "exec_shell", fake_exec) + monkeypatch.setattr(PathBubbleBackend, "_run_job", fake_run_job) step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(code="print('ok')", read_only_paths=[], writable_paths=[str(tmp_path)], **memory_options) + response = await step(code="print('ok')", **memory_options) assert response.success is True assert response.answer == "ok\n" @@ -293,23 +297,27 @@ async def test_python_execute_uses_path_backend_when_paths_supplied(tmp_path, mo assert seen["timeout"] == 60 expected_limit = memory_options.get("memory_limit_mib", 4096) assert seen["backend"]._memory_limit_mib == expected_limit # pylint: disable=protected-access + argv = seen["backend"]._bwrap_argv(seen["command"], str(tmp_path)) # pylint: disable=protected-access + assert ("--ro-bind", str(tmp_path), str(tmp_path)) in _mounts(argv) + assert not any(mode == "--bind" for mode, _source, _target in _mounts(argv)) + assert "--clearenv" in argv and "--unshare-all" in argv @pytest.mark.asyncio -@pytest.mark.parametrize("limit", [0, 1, -1]) -async def test_python_execute_rejects_memory_limit_without_path_policy(tmp_path, limit): - """A requested memory policy must never be silently ignored.""" +@pytest.mark.parametrize("policy", [{"writable_paths": ["/tmp"]}, {"read_only_paths": []}, {"writable_paths": None}]) +async def test_python_execute_rejects_caller_path_grants(tmp_path, policy): + """A caller cannot widen or disable the server's sandbox path policy.""" step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(code="print('should not run')", memory_limit_mib=limit) + response = await step(code="print('should not run')", **policy) assert response.success is False - assert response.answer == "memory_limit_mib requires read_only_paths or writable_paths to enable the sandbox" + assert "must be configured on the Step, not supplied by the caller" in response.answer @pytest.mark.asyncio async def test_step_rejects_invalid_mount_parameters(tmp_path): """Malformed path injection cannot silently disable sandboxing.""" - step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) - response = await step(code="print('unsafe')", writable_paths="/tmp") + step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path)), writable_paths="/tmp") + response = await step(code="print('unsafe')") assert response.success is False assert "writable_paths must be a list" in response.answer @@ -324,7 +332,7 @@ async def test_sandbox_start_failure_does_not_run_on_host(tmp_path, monkeypatch) monkeypatch.setattr(PathBubbleBackend, "_start_process", unavailable) app_context = ApplicationContext(workspace_dir=str(tmp_path)) python_step = PythonExecuteStep(app_context=app_context) - python_result = await python_step(code="print('host')", read_only_paths=[]) + python_result = await python_step(code="print('host')") assert python_result.success is False assert "bwrap unavailable" in python_result.answer @@ -333,7 +341,39 @@ async def test_sandbox_start_failure_does_not_run_on_host(tmp_path, monkeypatch) def test_job_hides_injected_path_policy(): """Path policy reaches the Step through context without becoming a tool argument.""" config_path = Path(__file__).parents[2] / "reme" / "config" / "benchmark.yaml" - parameters = yaml.safe_load(config_path.read_text(encoding="utf-8"))["jobs"]["python_execute"]["parameters"] + job = yaml.safe_load(config_path.read_text(encoding="utf-8"))["jobs"]["python_execute"] + assert job["enable_serve"] is False + parameters = job["parameters"] assert "read_only_paths" not in parameters["properties"] assert "writable_paths" not in parameters["properties"] assert "default" not in parameters["properties"]["memory_limit_mib"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("service_class", [HttpService, MCPService]) +async def test_benchmark_python_job_is_private_but_locally_callable(tmp_path, monkeypatch, service_class): + """Disabling serving preserves the benchmark's in-process job execution.""" + from unittest.mock import Mock + + config_path = Path(__file__).parents[2] / "reme" / "config" / "benchmark.yaml" + config = yaml.safe_load(config_path.read_text(encoding="utf-8"))["jobs"]["python_execute"] + context = ApplicationContext(workspace_dir=str(tmp_path)) + job = BaseJob(name="python_execute", app_context=context, **config) + context.jobs["python_execute"] = job + service = service_class() + service.add_job = Mock(return_value=True) + service.add_jobs(SimpleNamespace(context=context)) + service.add_job.assert_not_called() + + monkeypatch.setattr( + PathBubbleBackend, + "exec_shell", + AsyncMock(return_value=ExecResult(exit_code=0, stdout=b"local\n", stderr=b"")), + ) + await job.start() + try: + response = await job(code="print('local')") + assert response.success is True + assert response.answer == "local\n" + finally: + await job.close() diff --git a/tests/unit/test_process_output_limits.py b/tests/unit/test_process_output_limits.py index b14f4f6f..989dce42 100644 --- a/tests/unit/test_process_output_limits.py +++ b/tests/unit/test_process_output_limits.py @@ -60,8 +60,12 @@ async def test_continuous_output_stops_process(tmp_path, fd): @pytest.mark.asyncio -async def test_python_step_reports_output_overflow(tmp_path): +async def test_python_step_reports_output_overflow(tmp_path, monkeypatch): """Python callers receive bounded partial output and failure metadata.""" + monkeypatch.setattr( + "reme.steps.common.python_execute.file_sandbox_for_step", + lambda _step: LocalProcessBackend(), + ) step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) response = await step(code=f"import os;os.write(1, b'x' * {OUTPUT_LIMIT_BYTES + 1})") assert response.success is False diff --git a/tests/unit/test_python_execute_timeout.py b/tests/unit/test_python_execute_timeout.py new file mode 100644 index 00000000..88df253b --- /dev/null +++ b/tests/unit/test_python_execute_timeout.py @@ -0,0 +1,96 @@ +"""Timeout bounds at the Python executor and HTTP request boundaries.""" + +from unittest.mock import AsyncMock + +import pytest +from agentscope.tool import ExecResult +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from reme.components.application_context import ApplicationContext +from reme.components.job.base_job import BaseJob +from reme.components.service.http_service import HttpService +from reme.steps.common.python_execute import PythonExecuteStep + + +@pytest.fixture(name="backend") +def mock_backend(monkeypatch): + """Capture deadlines without executing untrusted code or waiting for them.""" + instance = AsyncMock() + instance.exec_shell.return_value = ExecResult(exit_code=0, stdout=b"ok\n", stderr=b"") + monkeypatch.setattr("reme.steps.common.python_execute.file_sandbox_for_step", lambda _step: instance) + return instance + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key", ["python_timeout", "timeout"]) +@pytest.mark.parametrize( + ("requested", "expected"), + [ + (0.25, 0.25), + ("12.5", 12.5), + (86400, 86400), + (86401, 86400), + ("inf", 86400), + (float("inf"), 86400), + ("1e9999", 86400), + (10**400, 86400), + ], +) +async def test_timeout_is_capped_before_execution(backend, key, requested, expected): + """Both supported request keys use a finite deadline and report the effective value.""" + response = await PythonExecuteStep()(code="pass", **{key: requested}) + assert response.success + assert backend.exec_shell.await_args.kwargs["timeout"] == expected + assert response.metadata["python_timeout"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested", ["nan", float("nan"), "-inf", 0, -1, -(10**400), True, None, "bad"]) +async def test_invalid_timeout_never_executes(backend, requested): + """Invalid inputs fail before constructing or invoking an execution backend.""" + response = await PythonExecuteStep()(code="pass", python_timeout=requested) + assert not response.success + assert response.answer == "python_timeout must be a positive number" + backend.exec_shell.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("maximum", "expected"), [(30, 30), (120, 60), (100000, 60)]) +async def test_default_timeout_respects_server_cap(backend, maximum, expected): + """An omitted timeout retains 60 seconds unless the server requires less.""" + response = await PythonExecuteStep(max_python_timeout=maximum)(code="pass") + assert response.success + assert backend.exec_shell.await_args.kwargs["timeout"] == expected + assert response.metadata["python_timeout"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("maximum", [0, -1, "nan", "-inf", None, True]) +async def test_invalid_server_cap_fails_closed(backend, maximum): + """A misconfigured server limit cannot silently become unlimited.""" + response = await PythonExecuteStep(max_python_timeout=maximum)(code="pass") + assert not response.success + assert response.answer == "max_python_timeout must be a positive number" + backend.exec_shell.assert_not_awaited() + + +@pytest.mark.parametrize(("maximum", "expected"), [(86400, 86400), (30, 30)]) +@pytest.mark.parametrize("key", ["python_timeout", "timeout"]) +def test_http_request_cannot_raise_server_timeout_cap(tmp_path, backend, maximum, expected, key): + """Exercise the real HTTP job route, including extra fields that bypass job schemas.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + job = BaseJob(name="python_review", app_context=context) + job.step_specs = [(PythonExecuteStep, {"app_context": context, "max_python_timeout": maximum})] + service = HttpService() + service.service = FastAPI() + service._add_json_job(job) # pylint: disable=protected-access + with TestClient(service.service) as client: + response = client.post( + "/python_review", + json={"code": "pass", key: "inf", "max_python_timeout": "inf"}, + ) + assert response.status_code == 200 + assert response.json()["success"] + assert response.json()["metadata"]["python_timeout"] == expected + assert backend.exec_shell.await_args.kwargs["timeout"] == expected From d323356e6f11c0faf4bf30fae7a2edc30e18fbb9 Mon Sep 17 00:00:00 2001 From: cangen185-sudo Date: Wed, 4 Nov 2026 15:25:11 +0800 Subject: [PATCH 3/5] fix(file-io): refuse edits after incomplete bounded reads (#663) * fix(file-io): refuse edits after incomplete bounded reads * style(file-io): keep complete-read signature formatter-stable * docs(file-io): clarify bounded-read safety guarantee --------- Co-authored-by: cangen185-sudo <234582933+cangen185-sudo@users.noreply.github.com> --- reme/steps/file_io/_file_io.py | 17 ++++- reme/steps/file_io/edit.py | 3 +- tests/unit/test_edit_read_limits.py | 102 ++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_edit_read_limits.py diff --git a/reme/steps/file_io/_file_io.py b/reme/steps/file_io/_file_io.py index 22e945b0..33ff1876 100644 --- a/reme/steps/file_io/_file_io.py +++ b/reme/steps/file_io/_file_io.py @@ -106,15 +106,28 @@ def _decode_known_file(data: bytes, file_extension: str) -> tuple[str, str]: # --------------------------------------------------------------------------- -async def read_file_safe(file_path, max_bytes: int = MAX_FILE_READ_BYTES) -> tuple[str, str]: +async def read_file_safe( + file_path, + max_bytes: int = MAX_FILE_READ_BYTES, + *, + require_complete: bool = False, +) -> tuple[str, str]: """Read file in byte mode and decode using extension-aware strategy. - Returns ``(text, encoding)``. + Returns ``(text, encoding)``. With ``require_complete``, reject files above + the byte limit or reads that leave unread bytes at the EOF check. This only + prevents returning a truncated bounded read for write-back; it does not + support oversized-file editing, provide a stable snapshot against concurrent + changes, or make subsequent writes atomic. """ stat = await aiofiles.os.stat(str(file_path)) + if require_complete and stat.st_size > max_bytes: + raise ValueError(f"file size {stat.st_size} exceeds read limit of {max_bytes} bytes") read_size = min(stat.st_size, max_bytes) async with aiofiles.open(str(file_path), "rb") as f: data = await f.read(read_size) + if require_complete and await f.read(1): + raise ValueError("file changed during read; refusing to return incomplete contents") return _decode_known_file(data, Path(file_path).suffix) diff --git a/reme/steps/file_io/edit.py b/reme/steps/file_io/edit.py index 31a0fd76..62ec1a64 100644 --- a/reme/steps/file_io/edit.py +++ b/reme/steps/file_io/edit.py @@ -9,6 +9,7 @@ from ._base import _FileIOStep from ._file_io import get_path_lock, read_file_safe, write_file_safe from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ...components import R +from ...constants import MAX_FILE_READ_BYTES def _replace_normalized_newlines(body: str, old: str, new: str) -> tuple[str, int]: @@ -100,7 +101,7 @@ class EditStep(_FileIOStep): return None try: - raw_text, encoding = await read_file_safe(target) + raw_text, encoding = await read_file_safe(target, max_bytes=MAX_FILE_READ_BYTES, require_complete=True) except Exception as e: # pylint: disable=broad-except self._fail(f"read failed: {e}", path=str(target)) return None diff --git a/tests/unit/test_edit_read_limits.py b/tests/unit/test_edit_read_limits.py new file mode 100644 index 00000000..6543b3d4 --- /dev/null +++ b/tests/unit/test_edit_read_limits.py @@ -0,0 +1,102 @@ +"""Edits must not write back a prefix produced by a bounded file read.""" + +import asyncio +from functools import partial +from types import SimpleNamespace + +import pytest + +from reme.components import ApplicationContext +from reme.steps.file_io import _file_io, edit as edit_module + + +def _small_edit_limit(monkeypatch): + """Exercise the real bounded reader with small files, without allocating 200 MiB.""" + monkeypatch.setattr(edit_module, "MAX_FILE_READ_BYTES", 64, raising=False) + monkeypatch.setattr(edit_module, "read_file_safe", partial(_file_io.read_file_safe, max_bytes=64)) + + +@pytest.mark.parametrize("suffix", [".md", ".txt"]) +@pytest.mark.parametrize("encoding", ["utf-8", "utf-16"]) +def test_edit_rejects_truncated_read_without_writing(tmp_path, monkeypatch, suffix, encoding): + """An occurrence in the prefix must not cause the unread tail to be discarded.""" + _small_edit_limit(monkeypatch) + path = tmp_path / ("note" + suffix) + source = ("target" + "x" * 80 + "TAIL").encode(encoding) + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert path.read_bytes() == source + assert step.context.response.success is False + assert "limit" in step.context.response.answer + + asyncio.run(run()) + + +@pytest.mark.parametrize("suffix", [".md", ".txt"]) +@pytest.mark.parametrize("encoding,filler", [("utf-8", 58), ("utf-16", 25)]) +def test_edit_at_read_limit_preserves_encoding(tmp_path, monkeypatch, suffix, encoding, filler): + """The exact byte limit remains editable, with the existing write rules.""" + _small_edit_limit(monkeypatch) + text = "target" + "x" * filler + source = text.encode(encoding) + assert len(source) == 64 + path = tmp_path / ("note" + suffix) + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert step.context.response.success is True + expected = text.replace("target", "edited") + ("\n" if suffix == ".md" else "") + assert path.read_bytes() == expected.encode(encoding) + + asyncio.run(run()) + + +def test_edit_rejects_growth_after_stat(tmp_path, monkeypatch): + """A stale small stat cannot turn a larger actual read into a destructive edit.""" + _small_edit_limit(monkeypatch) + path = tmp_path / "note.txt" + source = b"target" + b"x" * 80 + b"TAIL" + path.write_bytes(source) + app = ApplicationContext(workspace_dir=str(tmp_path)) + + async def stale_stat(_path): + return SimpleNamespace(st_size=8) + + monkeypatch.setattr(_file_io.aiofiles.os, "stat", stale_stat) + + async def run(): + step = edit_module.EditStep(app_context=app) + await step(path=path.name, old="target", new="edited") + assert step.context.response.success is False + assert "incomplete contents" in step.context.response.answer + assert path.read_bytes() == source + + asyncio.run(run()) + + +@pytest.mark.parametrize("size", [0, 63, 64, 65]) +def test_read_helper_complete_mode_and_default_truncation(tmp_path, size): + """Ordinary limited reads keep their behavior; complete reads reject only overflow.""" + path = tmp_path / "note.txt" + source = b"x" * size + path.write_bytes(source) + + async def run(): + text, _ = await _file_io.read_file_safe(path, max_bytes=64) + assert text == "x" * min(size, 64) + if size > 64: + with pytest.raises(ValueError, match="limit"): + await _file_io.read_file_safe(path, max_bytes=64, require_complete=True) + else: + text, _ = await _file_io.read_file_safe(path, max_bytes=64, require_complete=True) + assert text == "x" * size + assert path.read_bytes() == source + + asyncio.run(run()) From fe812c916febbeffe85bc64df44a95ed192309b2 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:51:14 +0800 Subject: [PATCH 4/5] feat: add DashScope web search job (#710) --- docs/en/configuration.md | 15 ++ docs/zh/configuration.md | 15 ++ example.env | 3 + pyproject.toml | 1 + reme/config/default.yaml | 20 +++ reme/steps/common/__init__.py | 2 + reme/steps/common/dashscope_search.py | 71 +++++++++ tests/unit/test_dashscope_search_step.py | 182 +++++++++++++++++++++++ 8 files changed, 309 insertions(+) create mode 100644 reme/steps/common/dashscope_search.py create mode 100644 tests/unit/test_dashscope_search_step.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index fedf7ac0..8b4a2e78 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -113,6 +113,21 @@ Built-in registrations include `openai`, `anthropic`, `dashscope`, `deepseek`, ` File operations, BM25 search, wikilink traversal, and `proactive_read` do not require an LLM. Evolution workflows such as `auto_memory`, `auto_resource`, `auto_dream`, and proactive refresh do. +## Web search + +With `reme-ai[core]` installed, `dashscope_search` calls `qwen3.8-flash` through the DashScope SDK, independently of LLM settings. +Set the key and start the service in one terminal: + +```bash +export DASHSCOPE_API_KEY=sk-xxx +reme start config=default +``` + +From another terminal, call `reme dashscope_search query="Latest Qwen models"`. +A missing key does not prevent service startup; calls to this Job return a failure. +The result's `search_info.search_results` preserves provider source fields; `content` is separate model-generated text. +Page bodies and snippets are not guaranteed. See `reme/config/default.yaml` for the model, search strategy, and timeout. + ## Decision Model Configuration `as_decision_model` is a provider-independent interface for probabilistic decisions. The current `agentscope` adapter diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 6cec4614..7053a74c 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -117,6 +117,21 @@ components: 基础文件操作、BM25 检索、wikilink 遍历和 `proactive_read` 不需要 LLM。`auto_memory`、`auto_resource`、 `auto_dream` 和 proactive refresh 等演化流程需要可用 LLM。 +## 联网搜索 + +安装 `reme-ai[core]` 后，默认 `dashscope_search` Job 使用 DashScope SDK 调用 `qwen3.8-flash`，独立于上述 LLM 配置。 +在一个终端设置 Key 并启动服务： + +```bash +export DASHSCOPE_API_KEY=sk-xxx +reme start config=default +``` + +在另一个终端调用 `reme dashscope_search query="Qwen 最新模型"`。 +未配置 Key 不影响服务启动，但此 Job 调用会返回失败。 +结果的 `search_info.search_results` 保留供应商返回的来源字段，`content` 是单独的模型生成正文； +网页正文和摘要不保证提供。模型、搜索策略和超时配置见 `reme/config/default.yaml`。 + ## 决策模型配置 `as_decision_model` 是不依赖供应商的概率决策模型接口；当前的 `agentscope` 适配器通过 AgentScope `2.0.11` 调用 TypeSafe API。 diff --git a/example.env b/example.env index b8e9e3d9..fa78dd5a 100644 --- a/example.env +++ b/example.env @@ -3,6 +3,9 @@ EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 LLM_API_KEY=sk-xxx LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +# Optional dashscope_search job (qwen3.8-flash web search). +# DASHSCOPE_API_KEY=sk-xxx + # Optional decision models: install reme-ai[core,decision]. # Setting these variables alone does not enable an as_decision_model component. # DECISION_MODEL_API_KEY=your-typesafe-api-key diff --git a/pyproject.toml b/pyproject.toml index e508f170..b00a962f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ web = [ ] core = [ "reme-ai[as]", + "dashscope>=1.26.0", "pillow>=13.0.0", "claude-agent-sdk>=0.2.174", "openai-codex>=0.163.0", diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 32233ce4..d2d481ba 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -502,6 +502,26 @@ jobs: expand_links: true max_links_per_direction: 10 + dashscope_search: + backend: base + description: >- + Search the web with qwen3.8-flash; requires DASHSCOPE_API_KEY. + Returns raw search_info sources and separate model-generated content; page bodies and snippets are not guaranteed. + parameters: + type: object + properties: + query: + type: string + description: "web search query" + required: + - query + steps: + - backend: dashscope_search_step + model: qwen3.8-flash + api_key: ${DASHSCOPE_API_KEY:-} + search_strategy: turbo + timeout: 60 + node_search: backend: base description: "Digest node recall — given a candidate abstraction's name+description, surface existing digest nodes similar enough to either dedup against or link to as related." diff --git a/reme/steps/common/__init__.py b/reme/steps/common/__init__.py index 18d2c7b0..86240fe4 100644 --- a/reme/steps/common/__init__.py +++ b/reme/steps/common/__init__.py @@ -3,6 +3,7 @@ from .add import AddStep from .app_config import AppConfigStep from .chat import ChatStep +from .dashscope_search import DashScopeSearchStep from .decision_demo import DecisionDemoStep from .demo import DemoEchoStep1, DemoEchoStep2 from .health_check import HealthCheckStep @@ -19,6 +20,7 @@ __all__ = [ "AddStep", "AppConfigStep", "ChatStep", + "DashScopeSearchStep", "DecisionDemoStep", "DemoEchoStep1", "DemoEchoStep2", diff --git a/reme/steps/common/dashscope_search.py b/reme/steps/common/dashscope_search.py new file mode 100644 index 00000000..230aac83 --- /dev/null +++ b/reme/steps/common/dashscope_search.py @@ -0,0 +1,71 @@ +"""Web search through DashScope, preserving provider search results.""" + +import asyncio +import os +from contextlib import aclosing +from http import HTTPStatus + +from ..base_step import BaseStep +from ...components import R + + +@R.register("dashscope_search_step") +class DashScopeSearchStep(BaseStep): + """Return raw search sources and separate model-generated content.""" + + async def execute(self): + assert self.context is not None + response = self.context.response + query = self.context.get("query") + query = "" if query is None else str(query).strip() + api_key = (self.kwargs.get("api_key") or os.getenv("DASHSCOPE_API_KEY", "")).strip() + if not query or not api_key: + response.success = False + response.answer = "query is required" if not query else "DASHSCOPE_API_KEY is required" + return response + + from aiohttp import ClientSession + from dashscope import AioMultiModalConversation + + model = self.kwargs.get("model", "qwen3.8-flash") + search_info, sources, content, usage = {}, {}, [], {} + request_id = "" + async with asyncio.timeout(self.kwargs.get("timeout", 60)), ClientSession(trust_env=True) as session: + stream = await AioMultiModalConversation.call( + api_key=api_key, + session=session, + model=model, + messages=[{"role": "user", "content": [{"text": query}]}], + enable_search=True, + enable_thinking=False, + stream=True, + incremental_output=True, + search_options={ + "forced_search": True, + "enable_source": True, + "search_strategy": self.kwargs.get("search_strategy", "turbo"), + }, + ) + async with aclosing(stream): + async for chunk in stream: + if chunk.status_code != HTTPStatus.OK: + raise RuntimeError(f"DashScope {chunk.code}: {chunk.message} (request_id={chunk.request_id})") + request_id = chunk.request_id or request_id + if any((chunk.usage or {}).values()): + usage = dict(chunk.usage) + output = chunk.output or {} + info = output.get("search_info") or {} + search_info.update({key: value for key, value in info.items() if key != "search_results"}) + for item in info.get("search_results") or []: + sources.setdefault(item["url"], item) + for choice in output.get("choices") or []: + text = (choice.get("message") or {}).get("content") or [] + if isinstance(text, str): + content.append(text) + else: + content.extend(part["text"] for part in text if "text" in part) + + search_info["search_results"] = list(sources.values()) + response.answer = {"search_info": search_info, "content": "".join(content)} + response.metadata.update(model=model, request_id=request_id, usage=usage) + return response diff --git a/tests/unit/test_dashscope_search_step.py b/tests/unit/test_dashscope_search_step.py new file mode 100644 index 00000000..04bc0c20 --- /dev/null +++ b/tests/unit/test_dashscope_search_step.py @@ -0,0 +1,182 @@ +"""DashScope web search contracts, with no provider requests.""" + +import asyncio +from unittest.mock import AsyncMock + +import pytest +from dashscope import AioMultiModalConversation +from dashscope.api_entities.dashscope_response import DashScopeAPIResponse, MultiModalConversationResponse + +from reme.application import Application +from reme.components.job import BaseJob +from reme.config.config_parser import _load_config +from reme.steps.common.dashscope_search import DashScopeSearchStep + + +@pytest.fixture(name="provider") +def mock_provider(monkeypatch): + """Supply SDK frames and record stream cleanup.""" + frames, closed = [], [] + + async def stream(): + try: + for frame in frames: + yield frame + finally: + closed.append(True) + + call = AsyncMock(side_effect=lambda **_kwargs: stream()) + monkeypatch.setattr(AioMultiModalConversation, "call", call) + monkeypatch.setenv("DASHSCOPE_API_KEY", "test-key") + return frames, closed, call + + +def chunk(output=None, **kwargs): + """Build an SDK response, including source-only frames.""" + return MultiModalConversationResponse.from_api_response( + DashScopeAPIResponse( + status_code=200, + request_id="request-1", + output=output or {}, + **kwargs, + ), + ) + + +@pytest.mark.asyncio +async def test_default_job_starts_without_key_and_rejects_call(monkeypatch, tmp_path): + """Default config loads and the registered job starts without credentials.""" + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + call = AsyncMock(side_effect=AssertionError("must not contact provider")) + monkeypatch.setattr(AioMultiModalConversation, "call", call) + config = _load_config("default.yaml")["jobs"]["dashscope_search"] + app = Application( + workspace_dir=str(tmp_path), + service={"backend": "cli"}, + jobs={"dashscope_search": config}, + enable_logo=False, + log_to_file=False, + ) + await app.start() + try: + response = await app.run_job("dashscope_search", query="latest Qwen models") + assert not response.success + assert response.answer == "DASHSCOPE_API_KEY is required" + assert config["parameters"]["required"] == ["query"] + call.assert_not_awaited() + finally: + await app.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("query", [None, "", " "]) +async def test_empty_query_does_not_call_provider(provider, query): + """Reject missing or blank queries before contacting DashScope.""" + response = await DashScopeSearchStep()(query=query) + assert not response.success + assert response.answer == "query is required" + provider[2].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_stream_preserves_sources_and_separates_generated_content(provider): + """Source-only, repeated, empty, and later source frames retain original fields and order.""" + frames, closed, call = provider + first = {"index": 1, "title": "One", "url": "https://example.com/1", "unknown": {"raw": True}} + second = {"index": 2, "title": "Two", "url": "https://example.com/2", "snippet": "provider snippet"} + frames.extend( + [ + chunk({"search_info": {"search_results": [first], "provider": "web"}}), + chunk( + { + "search_info": {"search_results": [dict(first, index=9), second]}, + "choices": [{"message": {"role": "assistant", "content": [{"text": "Generated "}]}}], + }, + ), + chunk( + { + "search_info": {"search_results": []}, + "choices": [{"message": {"role": "assistant", "content": "content"}}], + }, + usage={"input_tokens": 12, "output_tokens": 2, "plugins": {"search": {"count": 1}}}, + ), + chunk(), + ], + ) + response = await DashScopeSearchStep(api_key="configured-key")(query=0) + + assert response.success + assert response.answer == { + "search_info": {"search_results": [first, second], "provider": "web"}, + "content": "Generated content", + } + assert response.metadata == { + "model": "qwen3.8-flash", + "request_id": "request-1", + "usage": {"input_tokens": 12, "output_tokens": 2, "characters": 0, "plugins": {"search": {"count": 1}}}, + } + assert closed == [True] + session = call.await_args.kwargs["session"] + assert session.closed + call.assert_awaited_once_with( + api_key="configured-key", + session=session, + model="qwen3.8-flash", + messages=[{"role": "user", "content": [{"text": "0"}]}], + enable_search=True, + enable_thinking=False, + stream=True, + incremental_output=True, + search_options={"forced_search": True, "enable_source": True, "search_strategy": "turbo"}, + ) + + +@pytest.mark.asyncio +async def test_empty_results_remain_empty(provider): + """An ordinary response with no search hits does not invent sources.""" + provider[0].append(chunk({"search_info": {"search_results": []}})) + response = await DashScopeSearchStep()(query="no hits") + assert response.success + assert response.answer == {"search_info": {"search_results": []}, "content": ""} + + +@pytest.mark.asyncio +async def test_provider_error_after_sources_fails_job_and_closes_stream(provider): + """A failed stream cannot be published as a successful partial search.""" + frames, closed, _ = provider + frames.extend( + [ + chunk({"search_info": {"search_results": [{"url": "https://example.com"}]}}), + MultiModalConversationResponse( + status_code=429, + code="Throttling", + message="Try later", + request_id="failed-request", + ), + ], + ) + job = BaseJob() + job.step_specs = [(DashScopeSearchStep, {})] + response = await job(query="query") + assert not response.success + assert "Throttling: Try later" in response.answer + assert "failed-request" in response.answer + assert closed == [True] + + +@pytest.mark.asyncio +async def test_timeout_closes_stream(monkeypatch): + """Bound the entire stream, including waiting for another frame.""" + closed = [] + + async def stream(): + try: + yield chunk() + await asyncio.Event().wait() + finally: + closed.append(True) + + monkeypatch.setattr(AioMultiModalConversation, "call", AsyncMock(return_value=stream())) + with pytest.raises(TimeoutError): + await DashScopeSearchStep(api_key="test-key", timeout=0.01)(query="query") + assert closed == [True] From c6049c3fc18da6c8f52e28ebd7211002b81d8044 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:23:13 +0800 Subject: [PATCH 5/5] refactor(codex)!: share app-server processes across agent wrappers (#709) * refactor(codex)!: share app-server and reuse application job tools * test(codex): use FastMCP input_schema without compatibility aliases * feat(codex): configure native model catalog and preserve streamed metadata * fix(codex): default to unattended tools and preserve retry streams * fix(codex): use an explicit Responses provider for custom endpoints * fix(codex): store runtime state in workspace .codex * feat(config): add minimal Studio workspace config with Codex chat agent - Introduce proactive.yaml defining a workspace with HTTP backend and web enabled - Add background job for Markdown index update watching daily and digest directories - Define multiple base and stream backend jobs for chat, config, version, status, health, graph, reindex, and file operations - Configure Codex app server with API key auth and detailed model catalog including qwen3.8-max - Set up default agent wrapper using Codex backend with high effort and read-only sandbox - Establish default tokenizer (regex), keyword index (bm25), file graph and chunker for markdown files - Configure file storage backend with local implementation and indexing components * fix(codex): settle cancelled dispatch and isolate TLS callbacks * test(codex): check cancellation barriers across Python versions --- docs/en/framework.md | 39 +- docs/zh/framework.md | 34 +- reme/components/__init__.py | 2 + .../agent_wrapper/codex_agent_wrapper.py | 551 +--- .../agent_wrapper/codex_mcp_server.py | 85 - reme/components/client/http_client.py | 5 +- reme/components/codex_app_server.py | 421 +++ reme/components/service/http_service.py | 14 +- reme/components/service/mcp_tools.py | 24 +- reme/components/service/scoped_mcp.py | 99 + reme/config/default.yaml | 68 +- reme/config/demo.yaml | 56 +- reme/config/proactive.yaml | 266 ++ reme/enumeration/component_enum.py | 2 + reme/steps/common/stream_llm_demo.py | 4 +- tests/integration/test_codex_agent_wrapper.py | 284 +- tests/integration/test_codex_demo.py | 176 ++ .../test_codex_wrapper_capabilities.py | 139 + tests/unit/test_codex_agent_wrapper.py | 2453 +++++++---------- tests/unit/test_common_steps.py | 32 +- tests/unit/test_http_client.py | 33 + 21 files changed, 2622 insertions(+), 2165 deletions(-) delete mode 100644 reme/components/agent_wrapper/codex_mcp_server.py create mode 100644 reme/components/codex_app_server.py create mode 100644 reme/components/service/scoped_mcp.py create mode 100644 reme/config/proactive.yaml create mode 100644 tests/integration/test_codex_demo.py create mode 100644 tests/integration/test_codex_wrapper_capabilities.py diff --git a/docs/en/framework.md b/docs/en/framework.md index 8b6fc4ca..fb6757fa 100644 --- a/docs/en/framework.md +++ b/docs/en/framework.md @@ -75,6 +75,7 @@ reme/ as_llm/, as_embedding/ # model wrappers as_decision_model/ # optional decision model adapters agent_wrapper/ # AgentScope / Claude Code / Codex wrappers + codex_app_server.py # shared Codex process and conversation locks steps/ base_step.py # BaseStep, Ref, dispatch_steps common/ # version, help, health_check, status, chat @@ -547,7 +548,8 @@ Current default components in `reme/config/default.yaml`: | `as_llm` | `default` | `${LLM_BACKEND:-openai}` | LLM model wrapper. | | `agent_wrapper` | `default` | `agentscope` | AgentScope wrapper. | | `agent_wrapper` | `claude_code` | `claude_code` | Claude Code wrapper. | -| `agent_wrapper` | `codex/codex_oauth` | `codex` | Codex wrappers for API-key and OAuth authentication. | +| `codex_app_server` | `default/oauth` | `sdk` | Shared Codex processes for API-key and OAuth authentication. | +| `agent_wrapper` | `codex/codex_oauth` | `codex` | Agent configuration referencing a named Codex app-server. | | `file_graph` | `default` | `local` | Wikilink graph. | | `file_catalog` | `default/resource/digest/dream` | `local` | File-change checkpoints. | | `file_chunker` | `markdown` | `markdown` | Markdown AST chunking. | @@ -555,6 +557,41 @@ Current default components in `reme/config/default.yaml`: | `keyword_index` | `default` | `bm25` | BM25 keyword index. | | `file_store` | `default` | `local` | Combines file_graph and keyword_index; defaults to `embedding_store: ""`. | +Codex wrappers reference `app_server: default` or `app_server: oauth`. Multiple wrappers can share one server while +keeping their own models, instructions and conversations. Omit `session_id` to start a new thread; pass a returned +`session_id` to continue it. Different threads can execute concurrently; operations on one thread are serialized. + +The default is `approval_mode: deny_all` with `sandbox: full-access`: the SDK maps this mode to +`approval_policy: never`, so available tools run without human or model approval. Requests that need additional +permissions are rejected rather than prompting. Set `approval_mode: auto_review` to opt into the configured +automatic reviewer. Tool allowlists and injected argument constraints continue to apply. + +With an API key and a custom `base_url`, the shared server configures a dedicated Responses provider that uses +HTTP/SSE by default. Set `supports_websockets: true` on `codex_app_server` only for endpoints supporting the Responses +WebSocket transport. Explicit `config_overrides` retain precedence over generated provider defaults. + +API-key defaults use `qwen3.8-max` for main conversations, code reviews, subagents and native memory consolidation, +and `qwen3.8-flash` for automatic approval review and native memory extraction. The model catalog, main requests and +default subagents select `xhigh` reasoning. The catalog exposes only `xhigh`, so native approval reviews also use +that level rather than the SDK's internal `low` default. Main-model overrides use `model` +on the wrapper; the remaining native settings are explicit in `codex_app_server.default.config_overrides` and +`model_catalog`. ReMe writes this native Codex catalog as rebuildable metadata in the server's private home. +`auto_review_model_override` selects the approval model separately from the main model; changing only `model` does +not change that reviewer. Native conversation compaction uses the thread's model. OAuth uses account-supported +models and the separate `CODEX_OAUTH_MODEL_NAME` override. See [demo configuration](../../reme/config/demo.yaml) for +the complete API-key setup using `LLM_API_KEY` and `LLM_BASE_URL`. + +Migration: move `auth_mode`, `api_key`, `base_url`, `codex_home`, `codex_bin`, `config_overrides` and `experimental_api` +from wrappers to `components.codex_app_server.` (`backend: sdk`). The default Codex home is +`{workspace_dir}/.codex`; set `codex_home` to the previous directory to retain its login and saved threads. +Configure separate `codex_home` values when servers need independent account or session storage. +Remove `mcp_config` and `release_threads`. Job tools run against the current Application through a +session-scoped MCP endpoint, reusing its plain HTTP listener when available. TLS, CLI and embedded applications use +a lifecycle-managed loopback listener. Tool selection and injected arguments remain bound to each loaded thread; an unloaded +thread can be resumed with a new binding. Internal jobs are available only through the explicitly selected tools, +without changing their public service exposure. Closing a wrapper leaves the shared server running; application +shutdown closes it. A disconnected request fails without replaying tools; later requests can reconnect and resume. + Note that the `search` Step configuration contains `vector_weight`, but `file_store.default.embedding_store` is empty by default. Vector retrieval is available only when the runtime configuration enables an embedding store. diff --git a/docs/zh/framework.md b/docs/zh/framework.md index ad11342e..1910b1b1 100644 --- a/docs/zh/framework.md +++ b/docs/zh/framework.md @@ -70,6 +70,7 @@ reme/ as_llm/, as_embedding/ # 模型封装 as_decision_model/ # 可选的决策模型适配器 agent_wrapper/ # AgentScope / Claude Code / Codex wrapper + codex_app_server.py # 共享 Codex 进程与会话锁 steps/ base_step.py # BaseStep、Ref、dispatch_steps common/ # version、help、health_check、status、chat @@ -528,7 +529,8 @@ flowchart LR | `as_decision_model` | `default` | 默认未配置；示例为 `agentscope` | 可选概率决策；通过继承 default 的 decision_demo 配置演示 | | `agent_wrapper` | `default` | `agentscope` | AgentScope wrapper | | `agent_wrapper` | `claude_code` | `claude_code` | Claude Code wrapper | -| `agent_wrapper` | `codex/codex_oauth` | `codex` | 分别使用 API key 与 OAuth 的 Codex wrapper | +| `codex_app_server` | `default/oauth` | `sdk` | 分别使用 API key 与 OAuth 的共享 Codex 进程 | +| `agent_wrapper` | `codex/codex_oauth` | `codex` | 引用具名 Codex app-server 的 agent 配置 | | `file_graph` | `default` | `local` | wikilink 图谱 | | `file_catalog` | `default/resource/digest/dream` | `local` | 文件变更 checkpoint | | `file_chunker` | `markdown` | `markdown` | Markdown AST 分块 | @@ -536,6 +538,36 @@ flowchart LR | `keyword_index` | `default` | `bm25` | BM25 关键词索引 | | `file_store` | `default` | `local` | 组合 file_graph、keyword_index；默认 `embedding_store: ""` | +Codex wrapper 通过 `app_server: default` 或 `app_server: oauth` 引用共享组件。多个 wrapper 可以共享同一个进程， +并分别配置模型、指令和会话。不传 `session_id` 时创建新 thread；传回已有 `session_id` 时续聊。 +不同 thread 可以并发，同一 thread 的操作串行。 + +默认使用 `approval_mode: deny_all` 和 `sandbox: full-access`；SDK 将该模式映射为 `approval_policy: never`， +可用工具直接执行，不经过人工或模型审批。需要新增权限的请求直接拒绝，不弹出审批。 +如需自动审批，可显式设置 `approval_mode: auto_review`。工具白名单与注入参数约束继续生效。 + +使用 API key 和自定义 `base_url` 时，共享组件会配置独立的 Responses provider，默认直接使用 HTTP/SSE。 +仅当端点支持 Responses WebSocket 传输时，才在 `codex_app_server` 上设置 `supports_websockets: true`。 +显式 `config_overrides` 优先于自动生成的 provider 默认值。 + +API key 模式默认使用 `qwen3.8-max` 处理主对话、代码审查、子任务和原生记忆整合，使用 `qwen3.8-flash` +处理自动审批和原生记忆提取。模型目录、主请求和默认子任务的 reasoning 均设为 `xhigh`。 +模型目录仅开放 `xhigh` 档位，使原生自动审批也使用该档位，避免继续采用 SDK 内部的 `low` 默认值。 +主模型由 wrapper 的 `model` 配置，其余原生配置显式写在 +`codex_app_server.default.config_overrides` 和 `model_catalog` 中；ReMe 将原生 Codex 模型目录写为共享组件私有目录内的 +可重建元数据。`auto_review_model_override` 单独选择审批模型，仅修改主模型不会改变审批模型。原生会话压缩沿用 +thread 的模型。OAuth 使用账号支持的模型，通过独立的 `CODEX_OAUTH_MODEL_NAME` 覆盖。 +使用 `LLM_API_KEY`、`LLM_BASE_URL` 的完整 API key 配置见 [demo 配置](../../reme/config/demo.yaml)。 + +迁移：将 wrapper 上的 `auth_mode`、`api_key`、`base_url`、`codex_home`、`codex_bin`、`config_overrides` 和 +`experimental_api` 移到 `components.codex_app_server.<名称>`，并使用 `backend: sdk`。默认 Codex 目录为 +`{workspace_dir}/.codex`；如需继续使用原有登录和会话，请将 `codex_home` 指向旧目录。 +多个服务组件需要独立账号或会话存储时，分别显式配置 `codex_home`。 +删除 `mcp_config` 和 `release_threads`。job 工具通过受限 MCP 入口调用当前 Application 中的原有 job，明文 HTTP 模式复用现有 +监听器，TLS、CLI 和嵌入式模式使用随组件关闭的本地监听器。工具白名单及隐藏参数绑定到已加载的 thread；thread 卸载后可在 +恢复时重新绑定。内部 job 仅对显式选择的工具开放，不改变其公共服务可见性。关闭 wrapper 不关闭共享进程，由 Application +统一关闭。连接断开时当前请求报错且不重放工具，后续请求可以重连并恢复会话。 + 注意：`search` step 的配置含 `vector_weight`，但默认 `file_store.default.embedding_store` 为空，因此实际是否有向量检索取决于运行配置是否启用 embedding store。 diff --git a/reme/components/__init__.py b/reme/components/__init__.py index 7ffbda4c..a4620583 100644 --- a/reme/components/__init__.py +++ b/reme/components/__init__.py @@ -4,6 +4,7 @@ from . import agent_wrapper from . import as_llm from . import as_decision_model from . import client +from .codex_app_server import CodexAppServer from . import as_embedding from . import embedding_store from . import file_catalog @@ -27,6 +28,7 @@ __all__ = [ "BaseComponent", "ComponentMixin", "ComponentRegistry", + "CodexAppServer", "R", "create_application_registry", "PromptHandler", diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index b744452d..d7d5afb5 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -1,140 +1,74 @@ -"""Codex Python SDK backend for the unified agent wrapper.""" +"""Codex agent sessions backed by an application-owned app-server.""" from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator -from contextlib import suppress -from dataclasses import dataclass +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager from functools import partial -import hashlib import json import os -from pathlib import Path -import sys -import tempfile from typing import Any, TYPE_CHECKING from .base_agent_wrapper import BaseAgentWrapper +from ..codex_app_server import CodexAppServer from ..component_registry import R +from ..job import BackgroundJob, StreamJob from ...enumeration import ChunkEnum from ...schema import StreamChunk, TokenUsage if TYPE_CHECKING: - from openai_codex import AsyncCodex, AsyncThread, CodexConfig, RunInput + from openai_codex import AsyncThread, RunInput from openai_codex.types import Notification -else: - # Keep the optional Codex SDK out of ReMe's package import path. Tests may - # also replace this value before the SDK is loaded. - AsyncCodex: Any = None - - -def _get_async_codex_class(): - """Return the optional Codex client class, importing its SDK on first use.""" - global AsyncCodex # pylint: disable=global-statement - if AsyncCodex is None: - from openai_codex import AsyncCodex as AsyncCodexClass - - AsyncCodex = AsyncCodexClass - return AsyncCodex - - -@dataclass(frozen=True) -class _CodexAuthConfig: - """Resolved authentication settings for one Codex app-server.""" - - mode: str - api_key: str = "" - base_url: str = "" - - -@dataclass(frozen=True) -class _CodexLaunchConfig: - """Options fixed for the lifetime of one Codex app-server.""" - - auth_mode: str - api_key: str - base_url: str - codex_bin: str | None - config_overrides: tuple[str, ...] - experimental_api: bool @R.register("codex") class CodexAgentWrapper(BaseAgentWrapper): - """Agent wrapper backed by the Codex Python SDK.""" + """Configure agents while sharing transport, authentication and session locks.""" SDK_PACKAGE = "openai-codex" - _CLIENT_OPTION_NAMES = frozenset( + _SERVER_OPTIONS = frozenset( { - "api_key", "auth_mode", + "api_key", "base_url", + "supports_websockets", "codex_bin", "codex_home", "config_overrides", - "cwd", + "model_catalog", "experimental_api", "launch_args_override", + "mcp_config", + "release_threads", }, ) + def __init__(self, app_server: str = "default", **kwargs) -> None: + self._reject_server_options(kwargs) + super().__init__(**kwargs) + self.app_server = self.bind(app_server, CodexAppServer, default_factory=CodexAppServer, optional=False) + self._active = 0 + self._idle = asyncio.Event() + self._idle.set() + + @classmethod + def _reject_server_options(cls, kwargs: dict[str, Any]) -> None: + invalid = sorted(cls._SERVER_OPTIONS.intersection(kwargs)) + if invalid: + raise TypeError(f"Configure Codex process options on codex_app_server: {', '.join(invalid)}") + + def _reply_kwargs(self, kwargs: dict[str, Any], *, stream: bool = False) -> dict[str, Any]: + self._reject_server_options(kwargs) + invalid = sorted({"app_server", "cwd"}.intersection(kwargs)) + if invalid: + raise TypeError(f"Configure Codex wrapper options on the component: {', '.join(invalid)}") + return self._merged_stream_kwargs(kwargs) if stream else self._merged_kwargs(kwargs) + @staticmethod def _codex_usage(usage: Any) -> TokenUsage: - """Normalize Codex's full-turn input/output usage snapshot.""" return TokenUsage.from_provider(usage) - # pylint: disable=too-many-arguments - def __init__( - self, - mcp_config: str | None = None, - codex_home: str | Path | None = None, - *, - auth_mode: str = "auto", - api_key: str = "", - base_url: str = "", - codex_bin: str | None = None, - config_overrides: list[str] | tuple[str, ...] | None = None, - experimental_api: bool = True, - release_threads: bool = True, - **kwargs, - ) -> None: - if "launch_args_override" in kwargs: - raise TypeError("launch_args_override is not supported; configure codex_bin instead") - super().__init__(**kwargs) - self.mcp_config = mcp_config - self._codex_home = codex_home - self._launch_config = _CodexLaunchConfig( - auth_mode=auth_mode, - api_key=api_key, - base_url=base_url, - codex_bin=codex_bin, - config_overrides=tuple(config_overrides or ()), - experimental_api=experimental_api, - ) - self._codex: AsyncCodex | None = None - self._release_threads = release_threads - self._turn_lock = asyncio.Lock() - self._mcp_snapshot_path: Path | None = None - self._thread_tool_contexts: dict[str, str] = {} - - @property - def session_path(self) -> Path: - """Directory used for Codex state and persisted threads.""" - if self._codex_home: - path = Path(self._codex_home).expanduser() - return path if path.is_absolute() else self.workspace_path / path - if self.app_context is None: - return self.workspace_path / "mem_session" / "codex" - return self.workspace_path / self.app_context.app_config.mem_session_dir / "codex" - - @property - def wrapper_session_path(self) -> Path: - """ReMe-owned session data, kept separate from a shared OAuth CODEX_HOME.""" - if self.app_context is None: - return self.workspace_path / "mem_session" / "codex" - return self.workspace_path / self.app_context.app_config.mem_session_dir / "codex" - def _ensure_skills(self, skills: list[str] | str | None) -> None: """Expose selected project skills through Codex's repo-level directory.""" sources = self._resolve_project_skills(skills) @@ -154,178 +88,25 @@ class CodexAgentWrapper(BaseAgentWrapper): relative_source = os.path.relpath(source, target.parent) target.symlink_to(relative_source, target_is_directory=True) - def _explicit_mcp_config(self, kwargs: dict[str, Any]) -> str | None: - value = kwargs.get("mcp_config") if "mcp_config" in kwargs else self.mcp_config - if value is None: - return None - source = Path(str(value)).expanduser() - if source.suffix in {".yaml", ".yml", ".json"}: - if not source.is_absolute(): - source = self.workspace_path / source - return str(source.absolute()) - return str(value) - - def _effective_config_snapshot(self) -> Path: - """Create one private snapshot that remains valid for the client lifetime.""" - if self._mcp_snapshot_path is not None: - return self._mcp_snapshot_path - if self.app_context is None: - raise RuntimeError("Cannot snapshot MCP config without an app_context") - snapshot_dir = self.wrapper_session_path / "reme-mcp" - snapshot_dir.mkdir(parents=True, exist_ok=True) - fd, raw_path = tempfile.mkstemp(prefix="config-", suffix=".json", dir=snapshot_dir) - try: - if hasattr(os, "fchmod"): - os.fchmod(fd, 0o600) - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump(self.app_context.app_config.model_dump(mode="json"), stream) - except BaseException: - with suppress(OSError): - os.close(fd) - Path(raw_path).unlink(missing_ok=True) - raise - self._mcp_snapshot_path = Path(raw_path) - return self._mcp_snapshot_path - - def _mcp_config_source(self, kwargs: dict[str, Any]) -> str: - return self._explicit_mcp_config(kwargs) or str(self._effective_config_snapshot()) - - @staticmethod - def _resolve_auth_config(auth_mode: str, api_key: str = "", base_url: str = "") -> _CodexAuthConfig: - """Resolve login from wrapper options. - - The Codex child process still inherits the parent process environment. - """ - requested_mode = str(auth_mode or "auto").lower() - if requested_mode not in {"auto", "api_key", "oauth"}: - raise ValueError("auth_mode must be one of: auto, api_key, oauth") - if requested_mode == "oauth": - return _CodexAuthConfig(mode="oauth") - - api_key = api_key if isinstance(api_key, str) else "" - if requested_mode == "api_key" and not api_key: - raise ValueError("auth_mode='api_key' requires a non-empty API key") - if not api_key: - return _CodexAuthConfig(mode="oauth") - - base_url = base_url if isinstance(base_url, str) else "" - return _CodexAuthConfig(mode="api_key", api_key=api_key, base_url=base_url) - - def _build_client_config(self, auth: _CodexAuthConfig) -> CodexConfig: - from openai_codex import CodexConfig - - env = dict(self.subprocess_environment) - self.session_path.mkdir(parents=True, exist_ok=True) - env["CODEX_HOME"] = str(self.session_path) - - overrides = list(self._launch_config.config_overrides) - if auth.base_url: - overrides.append(f"openai_base_url={json.dumps(auth.base_url)}") - login_method = "api" if auth.mode == "api_key" else "chatgpt" - overrides.append(f"forced_login_method={json.dumps(login_method)}") - return CodexConfig( - codex_bin=self._launch_config.codex_bin, - config_overrides=tuple(overrides), - cwd=str(self.cwd), - env=env, - client_name="reme", - client_title="ReMe", - experimental_api=self._launch_config.experimental_api, - ) - - @classmethod - def _reject_client_options(cls, kwargs: dict[str, Any]) -> None: - invalid = sorted(cls._CLIENT_OPTION_NAMES.intersection(kwargs)) - if invalid: - names = ", ".join(invalid) - raise TypeError(f"Codex client options must be configured on the wrapper: {names}") - - def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: - from ..job import BackgroundJob, StreamJob - - job_names = list(dict.fromkeys(kwargs.get("job_tools") or [])) - if not job_names: - return None - jobs = self._resolve_job_tools(job_names) - unsupported = [job.name for job in jobs if isinstance(job, (BackgroundJob, StreamJob))] - if unsupported: - raise TypeError(f"Codex job_tools must be non-stream request jobs: {', '.join(unsupported)}") - - config_source = self._mcp_config_source(kwargs) - args = [ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - config_source, - "--workspace", - str(self.workspace_path), - ] - for name in job_names: - args.extend(["--job", name]) - args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) - if injected_job_kwargs := dict(kwargs.get("injected_job_kwargs") or {}): - args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) - return { - "command": sys.executable, - "args": args, - "cwd": str(self.project_path), - "required": True, - "enabled_tools": job_names, - "startup_timeout_sec": kwargs.get("mcp_startup_timeout", 30), - "tool_timeout_sec": kwargs.get("mcp_tool_timeout", 300), - } - - def _thread_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: - config = dict(kwargs.get("config") or {}) - if proxy_environment := self.command_proxy_environment: - shell_environment_policy = dict(config.get("shell_environment_policy") or {}) - environment = dict(shell_environment_policy.get("set") or {}) - environment.update(proxy_environment) - shell_environment_policy["set"] = environment - config["shell_environment_policy"] = shell_environment_policy - if server := self._mcp_server_config(kwargs): - servers = dict(config.get("mcp_servers") or {}) - server_key = hashlib.sha256(json.dumps(server, sort_keys=True).encode()).hexdigest()[:12] - servers[f"reme_jobs_{server_key}"] = server - config["mcp_servers"] = servers - return config or None - @staticmethod def _enum(enum_cls: Any, value: Any, default: Any = None) -> Any: if value is None: return default return value if isinstance(value, enum_cls) else enum_cls(value) - async def _open_thread(self, codex: AsyncCodex, kwargs: dict[str, Any]) -> AsyncThread: + def _thread_options(self, kwargs: dict[str, Any]) -> dict[str, Any]: from openai_codex import ApprovalMode, Sandbox from openai_codex.types import Personality, ThreadSource, ThreadStartSource - resume = kwargs.get("resume") or "" - session_id = kwargs.get("session_id") or "" - if resume and session_id and resume != session_id: - raise ValueError("resume and session_id must identify the same Codex thread") - thread_id = resume or session_id - fork_session = bool(kwargs.get("fork_session", False)) - if fork_session and not thread_id: - raise ValueError("fork_session=True requires resume or session_id") - if self._release_threads and self._thread_tool_contexts: - await self._prune_thread_tool_contexts(codex) - requested_tool_context = json.dumps( - { - "tool_context_id": str(kwargs.get("tool_context_id") or ""), - "injected_job_kwargs": dict(kwargs.get("injected_job_kwargs") or {}), - }, - sort_keys=True, - default=str, - ) - if not fork_session and thread_id in self._thread_tool_contexts: - if requested_tool_context != self._thread_tool_contexts[thread_id]: - raise ValueError("tool_context_id and injected_job_kwargs cannot change when resuming a Codex thread") - - common = { - "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.auto_review), + config = dict(kwargs.get("config") or {}) + if proxy_environment := self.command_proxy_environment: + policy = dict(config.get("shell_environment_policy") or {}) + policy["set"] = {**policy.get("set", {}), **proxy_environment} + config["shell_environment_policy"] = policy + options = { + "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.deny_all), "base_instructions": kwargs.get("base_instructions"), - "config": self._thread_config(kwargs), + "config": config, "cwd": str(self.cwd), "developer_instructions": kwargs.get("system_prompt"), "model": kwargs.get("model"), @@ -333,122 +114,64 @@ class CodexAgentWrapper(BaseAgentWrapper): "sandbox": self._enum(Sandbox, kwargs.get("sandbox"), Sandbox.full_access), "service_tier": kwargs.get("service_tier"), } - personality = self._enum(Personality, kwargs.get("personality")) - thread_source = self._enum(ThreadSource, kwargs.get("thread_source")) - if fork_session: - thread = await codex.thread_fork( - thread_id, + if kwargs.get("fork_session"): + options.update( ephemeral=kwargs.get("ephemeral"), - thread_source=thread_source, - **common, + thread_source=self._enum(ThreadSource, kwargs.get("thread_source")), ) - elif thread_id: - thread = await codex.thread_resume(thread_id, personality=personality, **common) else: - thread = await codex.thread_start( - ephemeral=kwargs.get("ephemeral"), - personality=personality, - service_name=kwargs.get("service_name"), - session_start_source=self._enum(ThreadStartSource, kwargs.get("session_start_source")), - thread_source=thread_source, - **common, - ) - self._thread_tool_contexts[thread.id] = requested_tool_context - return thread + options["personality"] = self._enum(Personality, kwargs.get("personality")) + if not (kwargs.get("resume") or kwargs.get("session_id")): + options.update( + ephemeral=kwargs.get("ephemeral"), + service_name=kwargs.get("service_name"), + session_start_source=self._enum(ThreadStartSource, kwargs.get("session_start_source")), + thread_source=self._enum(ThreadSource, kwargs.get("thread_source")), + ) + return options - def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: - from openai_codex import ApprovalMode, Sandbox - from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary - - return { - "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), - "cwd": str(self.cwd), - "effort": self._enum(ReasoningEffort, kwargs.get("effort")), - "model": kwargs.get("model"), - "output_schema": kwargs.get("output_schema"), - "personality": self._enum(Personality, kwargs.get("personality")), - "sandbox": self._enum(Sandbox, kwargs.get("sandbox")), - "service_tier": kwargs.get("service_tier"), - "summary": self._enum(ReasoningSummary, kwargs.get("summary")), - } - - async def _get_codex(self) -> AsyncCodex: - """Lazily start the component-owned client from its fixed launch configuration.""" - if self._codex is not None: - return self._codex - auth = self._resolve_auth_config( - self._launch_config.auth_mode, - self._launch_config.api_key, - self._launch_config.base_url, - ) - codex = _get_async_codex_class()(self._build_client_config(auth)) + @asynccontextmanager + async def _session(self, kwargs: dict[str, Any], *, compact: bool = False) -> AsyncIterator[AsyncThread]: + resume, session_id = kwargs.get("resume") or "", kwargs.get("session_id") or "" + if resume and session_id and resume != session_id: + raise ValueError("resume and session_id must identify the same Codex thread") + thread_id = resume or session_id + fork = bool(kwargs.get("fork_session")) + if fork and not thread_id: + raise ValueError("fork_session=True requires resume or session_id") + jobs = None if compact else self._resolve_job_tools(list(dict.fromkeys(kwargs.get("job_tools") or []))) + if any(isinstance(job, (BackgroundJob, StreamJob)) for job in jobs or []): + raise TypeError("Codex job_tools must be non-stream request jobs") + if any(not getattr(job, "enabled", True) for job in jobs or []): + raise ValueError("Codex job_tools must be enabled") + self._ensure_skills(kwargs.get("skills")) + await self.start() + self._active += 1 + self._idle.clear() try: - if auth.mode == "api_key": - await codex.login_api_key(auth.api_key) - else: - account = await codex.account() - if account.account is None: - raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") - except BaseException: - await codex.close() - raise - self._codex = codex - return codex + async with self.app_server.thread( + thread_id, + fork=fork, + options={} if compact else self._thread_options(kwargs), + jobs=jobs, + injected=self._resolve_injected_job_kwargs(kwargs), + mcp_startup_timeout=kwargs.get("mcp_startup_timeout", 30), + mcp_tool_timeout=kwargs.get("mcp_tool_timeout", 300), + ) as thread: + yield thread + finally: + self._active -= 1 + if not self._active: + self._idle.set() async def _close(self) -> None: - """Close the persistent app-server and remove its private config snapshot.""" - async with self._turn_lock: - codex, self._codex = self._codex, None - self._thread_tool_contexts.clear() - try: - if codex is not None: - await codex.close() - finally: - if self._mcp_snapshot_path is not None: - self._mcp_snapshot_path.unlink(missing_ok=True) - self._mcp_snapshot_path = None + await self._idle.wait() - async def _prune_thread_tool_contexts(self, codex: AsyncCodex) -> None: - """Forget tool constraints only after the server confirms a thread is unloaded.""" + async def _interrupt(self, turn) -> None: try: - from openai_codex.generated.v2_all import ThreadLoadedListResponse - - loaded_ids: set[str] = set() - cursor = None - while True: - page = await codex._client.request( # pylint: disable=protected-access - "thread/loaded/list", - {"cursor": cursor}, - response_model=ThreadLoadedListResponse, - ) - loaded_ids.update(page.data) - cursor = page.next_cursor - if not cursor: - break - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to list loaded Codex threads: {exc}") - return - self._thread_tool_contexts = { - thread_id: context for thread_id, context in self._thread_tool_contexts.items() if thread_id in loaded_ids - } - - async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None: - """Unsubscribe so the app-server can unload the thread once it is idle.""" - try: - from openai_codex.generated.v2_all import ThreadUnsubscribeResponse - - # The SDK exposes no public unsubscribe; request() is public but - # reached through the private client attribute on SDK 0.162.x. - await codex._client.request( # pylint: disable=protected-access - "thread/unsubscribe", - {"threadId": thread_id}, - response_model=ThreadUnsubscribeResponse, - ) - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to unsubscribe Codex thread {thread_id}: {exc}") - # Unsubscribe is not an unload acknowledgement. In particular, a - # cancelled turn can still be running with its original MCP constraints. - # Keep the guard until thread/loaded/list confirms the thread is gone. + await turn.interrupt() + except Exception as exc: + self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") @staticmethod def _serialize(value: Any) -> Any: @@ -457,20 +180,15 @@ class CodexAgentWrapper(BaseAgentWrapper): return to_jsonable_python(value, by_alias=True) async def reply(self, inputs: RunInput, **kwargs) -> dict: - """Run one Codex turn and return its final response.""" - self._reject_client_options(kwargs) - kwargs = self._merged_kwargs(kwargs) - self._ensure_skills(kwargs.get("skills")) - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await self._open_thread(codex, kwargs) + """Run an isolated or explicitly resumed conversation on the shared server.""" + kwargs = self._reply_kwargs(kwargs) + async with self._session(kwargs) as thread: + turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) try: - result = await thread.run(inputs, **self._turn_kwargs(kwargs)) - finally: - if self._release_threads: - await self._release_thread(codex, thread.id) - + result = await turn.run() + except BaseException: + await self._interrupt(turn) + raise final_response = result.final_response or "" response = { "session_id": thread.id, @@ -491,16 +209,25 @@ class CodexAgentWrapper(BaseAgentWrapper): return response async def compact_session(self, session_id: str) -> None: - """Start native compaction of a Codex thread.""" - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await codex.thread_resume(session_id) - try: - await thread.compact() - finally: - if self._release_threads: - await self._release_thread(codex, thread.id) + """Hold the conversation lock until native compaction has finished.""" + async with self._session({"session_id": session_id}, compact=True) as thread: + await self.app_server.compact(thread) + + def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + from openai_codex import ApprovalMode, Sandbox + from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary + + return { + "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), + "cwd": str(self.cwd), + "effort": self._enum(ReasoningEffort, kwargs.get("effort")), + "model": kwargs.get("model"), + "output_schema": kwargs.get("output_schema"), + "personality": self._enum(Personality, kwargs.get("personality")), + "sandbox": self._enum(Sandbox, kwargs.get("sandbox")), + "service_tier": kwargs.get("service_tier"), + "summary": self._enum(ReasoningSummary, kwargs.get("summary")), + } @classmethod # pylint: disable=too-many-return-statements @@ -619,31 +346,19 @@ class CodexAgentWrapper(BaseAgentWrapper): return [] async def reply_stream(self, inputs: RunInput, **kwargs) -> AsyncGenerator[StreamChunk, None]: - """Stream Codex app-server notifications as unified chunks.""" - self._reject_client_options(kwargs) - kwargs = self._merged_stream_kwargs(kwargs) - self._ensure_skills(kwargs.get("skills")) - await self.start() - async with self._turn_lock: - codex = await self._get_codex() - thread = await self._open_thread(codex, kwargs) + """Stream notifications routed by the SDK to this conversation's turn.""" + kwargs = self._reply_kwargs(kwargs, stream=True) + async with self._session(kwargs) as thread: + turn = await self.app_server.turn(thread, inputs, **self._turn_kwargs(kwargs)) + stream = turn.stream() + completed = False try: - turn = await thread.turn(inputs, **self._turn_kwargs(kwargs)) - stream = turn.stream() - completed = False - try: - async for event in stream: - if event.method == "turn/completed": - completed = True - for chunk in self._event_to_chunks(event, thread.id): - yield chunk - finally: - if not completed: - try: - await turn.interrupt() - except Exception as exc: # pylint: disable=broad-exception-caught - self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") - await stream.aclose() + async for event in stream: + if event.method == "turn/completed": + completed = True + for chunk in self._event_to_chunks(event, thread.id): + yield chunk finally: - if self._release_threads: - await self._release_thread(codex, thread.id) + if not completed: + await self._interrupt(turn) + await stream.aclose() diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py deleted file mode 100644 index 9ed86f5b..00000000 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ /dev/null @@ -1,85 +0,0 @@ -"""FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" - -import argparse -import json -from pathlib import Path -from typing import Any - -from ...config import resolve_app_config -from ...reme import ReMe - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Expose selected ReMe jobs over FastMCP STDIO") - parser.add_argument("--config", default="default", help="ReMe config name or file path") - parser.add_argument("--workspace", required=True, help="ReMe workspace directory") - parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") - parser.add_argument("--tool-context-id", default="", help="Context id injected into every job call") - parser.add_argument( - "--injected-job-kwargs", - default="", - help="JSON object of server-owned kwargs injected into every job call", - ) - return parser.parse_args() - - -def _prepare_config( - config: dict[str, Any], - job_names: list[str], - tool_context_id: str = "", - injected_job_kwargs: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Configure the dedicated child Application to serve selected jobs over MCP STDIO.""" - selected = set(job_names) - jobs: dict[str, dict[str, Any]] = {} - for name, raw_job_config in (config.get("jobs") or {}).items(): - job_config = dict(raw_job_config) - if job_config.get("backend") in {"background", "cron"}: - continue - if name in selected: - # An explicit Codex job_tools selection has always overridden enable_serve. - job_config["enable_serve"] = True - jobs[name] = job_config - - missing = sorted(selected.difference(jobs)) - if missing: - raise KeyError(f"Codex job tools not found or not request jobs: {', '.join(missing)}") - - service: dict[str, Any] = { - "backend": "mcp", - "transport": "stdio", - "jobs": job_names, - "tool_error_on_failure": True, - } - injected = dict(injected_job_kwargs or {}) - if tool_context_id: - injected["tool_context_id"] = tool_context_id - if injected: - service["injected_job_kwargs"] = injected - - prepared = dict(config) - prepared["jobs"] = jobs - prepared["service"] = service - return prepared - - -def main() -> None: - """Load ReMe and serve the requested jobs over STDIO.""" - args = _parse_args() - config = resolve_app_config( - config=args.config, - workspace_dir=str(Path(args.workspace).absolute()), - enable_logo=False, - log_to_console=False, - log_to_file=False, - log_config=False, - ) - injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else None - if injected_job_kwargs is not None and not isinstance(injected_job_kwargs, dict): - raise TypeError("--injected-job-kwargs must be a JSON object") - config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) - ReMe(**config).run_app() - - -if __name__ == "__main__": - main() diff --git a/reme/components/client/http_client.py b/reme/components/client/http_client.py index ef4da0a9..b304e1bf 100644 --- a/reme/components/client/http_client.py +++ b/reme/components/client/http_client.py @@ -60,9 +60,8 @@ class HttpClient(BaseClient): except json.JSONDecodeError: continue chunk = StreamChunk(**data) - if chunk.chunk_type == ChunkEnum.ERROR: - # Surface server-side errors as exceptions so callers don't - # mistake error chunks for valid content. + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True: + # Recoverable notifications remain visible; terminal errors stop the request. raise RuntimeError(str(chunk.chunk)) if chunk.done: return diff --git a/reme/components/codex_app_server.py b/reme/components/codex_app_server.py new file mode 100644 index 00000000..da465865 --- /dev/null +++ b/reme/components/codex_app_server.py @@ -0,0 +1,421 @@ +"""Application-owned Codex process, concurrent sessions, and local job tools.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field +import json +from pathlib import Path +import sys +from typing import Any, TYPE_CHECKING + +from .base_component import BaseComponent +from .component_registry import R +from .service.scoped_mcp import ScopedMCP +from ..constants import normalize_connect_host +from ..enumeration import ComponentEnum +from ..utils.async_utils import finish_cleanup + +if TYPE_CHECKING: + from openai_codex import AsyncCodex, AsyncThread, CodexConfig + + +@dataclass +class _ThreadState: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + users: int = 0 + revision: int = 0 + binding: str | None = None + token: str | None = None + completion: asyncio.Future[None] | None = None + operation_started: bool = False + + +@R.register("sdk") +class CodexAppServer(BaseComponent): + """Share one SDK client across wrappers with the same account and launch settings.""" + + component_type = ComponentEnum.CODEX_APP_SERVER + + def __init__( + self, + auth_mode: str = "auto", + api_key: str = "", + base_url: str = "", + supports_websockets: bool = False, + codex_home: str | Path | None = None, + codex_bin: str | None = None, + config_overrides: list[str] | tuple[str, ...] = (), + model_catalog: dict[str, Any] | None = None, + cwd: str | Path | None = None, + experimental_api: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if auth_mode not in {"auto", "api_key", "oauth"}: + raise ValueError("auth_mode must be one of: auto, api_key, oauth") + self._auth_mode = auth_mode + self._api_key = api_key + self._base_url = base_url + self._supports_websockets = supports_websockets + self._codex_bin = codex_bin + self._config_overrides = tuple(config_overrides) + self._model_catalog = model_catalog + self._cwd = cwd + self._experimental_api = experimental_api + self._codex_home = codex_home + self._codex: AsyncCodex | None = None + self._notifications: asyncio.Task | None = None + self._client_lock = asyncio.Lock() + self._threads: dict[str, _ThreadState] = {} + self._active = 0 + self._idle = asyncio.Event() + self._idle.set() + self._failed = False + self._http_url: str | None = None + self.jobs_mcp = ScopedMCP() + + @property + def session_path(self) -> Path: + """Persistent Codex state under the workspace, unless an explicit home is configured.""" + if self._codex_home: + path = Path(self._codex_home).expanduser() + return path if path.is_absolute() else self.workspace_path / path + return self.workspace_path / ".codex" + + def _client_config(self) -> CodexConfig: + from openai_codex import CodexConfig + + environment = dict(self.app_context.app_config.environment) if self.app_context else {} + environment["CODEX_HOME"] = str(self.session_path) + overrides = list(self._config_overrides) + api_login = self._auth_mode != "oauth" and bool(self._api_key) + if self._auth_mode == "api_key" and not api_login: + raise ValueError("auth_mode='api_key' requires a non-empty API key") + self.session_path.mkdir(parents=True, exist_ok=True) + if self._model_catalog is not None: + if any(option.partition("=")[0].strip() == "model_catalog_json" for option in overrides): + raise ValueError("Configure either model_catalog or config_overrides.model_catalog_json, not both") + catalog_path = self.session_path / "reme-model-catalog.json" + catalog_path.write_text(json.dumps(self._model_catalog, ensure_ascii=False), encoding="utf-8") + overrides.append(f"model_catalog_json={json.dumps(str(catalog_path))}") + if api_login and self._base_url: + # Custom endpoints must not inherit the built-in OpenAI provider's WebSocket capability. + overrides = [ + 'model_provider="reme"', + 'model_providers.reme.name="ReMe API"', + f"model_providers.reme.base_url={json.dumps(self._base_url)}", + 'model_providers.reme.wire_api="responses"', + "model_providers.reme.requires_openai_auth=true", + f"model_providers.reme.supports_websockets={json.dumps(self._supports_websockets)}", + *overrides, + ] + overrides.append(f"forced_login_method={json.dumps('api' if api_login else 'chatgpt')}") + cwd = Path(self._cwd).expanduser() if self._cwd else self.workspace_path + if not cwd.is_absolute(): + cwd = self.workspace_path / cwd + return CodexConfig( + codex_bin=self._codex_bin, + config_overrides=tuple(overrides), + cwd=str(cwd), + env=environment, + client_name="reme", + client_title="ReMe", + experimental_api=self._experimental_api, + ) + + async def _get_client(self) -> AsyncCodex: + from openai_codex import AsyncCodex + + async with self._client_lock: + if self._failed: + if self._active > 1: + raise RuntimeError("Codex disconnected; wait for its active requests to finish before reconnecting") + await self._close_client() + self._threads.clear() + self.jobs_mcp.scopes.clear() + self._failed = False + if self._codex is None: + codex = AsyncCodex(self._client_config()) + try: + if self._auth_mode != "oauth" and self._api_key: + await codex.login_api_key(self._api_key) + elif (await codex.account()).account is None: + raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}") + except BaseException: + await codex.close() + raise + self._codex = codex + self._notifications = asyncio.create_task(self._watch_notifications(codex)) + return self._codex + + async def _watch_notifications(self, codex: AsyncCodex) -> None: + """Route SDK status events to the operation that owns each conversation.""" + try: + while self._codex is codex: + event = await codex._client.next_notification() # pylint: disable=protected-access + if event.method != "thread/status/changed": + continue + state = self._threads.get(event.payload.thread_id) + if state is None or state.completion is None or state.completion.done(): + continue + if event.payload.status.root.type == "active": + state.operation_started = True + elif state.operation_started: + state.completion.set_result(None) + except Exception as exc: + if self._codex is codex: + self._failed = True + for state in self._threads.values(): + if state.completion is not None and not state.completion.done(): + state.completion.set_exception(exc) + + async def _close_client(self) -> None: + codex, self._codex = self._codex, None + try: + if codex is not None: + await codex.close() + finally: + if self._notifications is not None: + self._notifications.cancel() + with suppress(asyncio.CancelledError): + await self._notifications + self._notifications = None + + async def _mcp_url(self) -> str: + from .service.http_service import HttpService + + service = self.app_context.service if self.app_context else None + if ( + isinstance(service, HttpService) + and service.service is not None + and any(getattr(route, "endpoint", None) is self.jobs_mcp for route in service.service.routes) + ): + if self._http_url is None: + host = normalize_connect_host(service.host) + host = "::1" if host == "::" else host + # Background jobs can run before uvicorn finishes binding the main listener. + async with asyncio.timeout(30): + while True: + try: + _, writer = await asyncio.open_connection(host, service.port) + except OSError: + await asyncio.sleep(0.01) + else: + writer.close() + await writer.wait_closed() + break + authority = f"[{host}]" if ":" in host else host + self._http_url = f"http://{authority}:{service.port}{self.jobs_mcp.path}" + return self._http_url + return await self.jobs_mcp.listen() + + async def _prune_threads(self, codex: AsyncCodex) -> None: + """Revoke idle tool bindings only after the server confirms their threads unloaded.""" + from openai_codex.generated.v2_all import ThreadLoadedListResponse + + candidates = {key: (state, state.revision) for key, state in self._threads.items() if not state.users} + if not candidates: + return + try: + loaded: set[str] = set() + cursor = None + while True: + page = await codex._client.request( # pylint: disable=protected-access + "thread/loaded/list", + {"cursor": cursor}, + response_model=ThreadLoadedListResponse, + ) + loaded.update(page.data) + cursor = page.next_cursor + if not cursor: + break + except Exception as exc: + self.logger.warning(f"Failed to list loaded Codex threads: {exc}") + return + for key, (state, revision) in candidates.items(): + if key not in loaded and not state.users and state.revision == revision and self._threads.get(key) is state: + self._threads.pop(key) + self.jobs_mcp.scopes.pop(state.token, None) + + async def wait_idle(self, thread: AsyncThread) -> None: + """Keep the thread lock through cancellation or asynchronous compaction.""" + state = self._threads[thread.id] + if state.completion is not None: + await asyncio.shield(state.completion) + state.completion = None + async with asyncio.timeout(30): + while (await thread.read()).thread.status.root.type == "active": + await asyncio.sleep(0.05) + + @asynccontextmanager + async def _operation(self, thread: AsyncThread) -> AsyncIterator[None]: + """Track a shielded dispatch; clear its barrier if the SDK rejects the request.""" + state = self._threads[thread.id] + state.operation_started = False + state.completion = asyncio.get_running_loop().create_future() + state.completion.add_done_callback(lambda done: done.exception() if not done.cancelled() else None) + try: + yield + except Exception: + state.completion.cancel() + state.completion = None + raise + + async def turn(self, thread: AsyncThread, inputs: Any, **kwargs): + """Start a normal SDK turn with completion tracking owned by this server.""" + turn = None + + async def dispatch(): + nonlocal turn + async with self._operation(thread): + turn = await thread.turn(inputs, **kwargs) + + try: + await finish_cleanup(dispatch()) + except asyncio.CancelledError: + if turn is not None: + try: + await finish_cleanup(turn.interrupt()) + except asyncio.CancelledError: + pass + except Exception as exc: + self.logger.warning(f"Failed to interrupt Codex turn {turn.id}: {exc}") + raise + return turn + + async def compact(self, thread: AsyncThread) -> None: + """Enqueue native compaction; the thread lease waits for its completion event.""" + + async def dispatch(): + async with self._operation(thread): + await thread.compact() + + await finish_cleanup(dispatch()) + + @asynccontextmanager + async def thread( + self, + thread_id: str = "", + *, + fork: bool = False, + options: dict[str, Any] | None = None, + jobs: list[Any] | None = None, + injected: dict[str, Any] | None = None, + mcp_startup_timeout: float = 30, + mcp_tool_timeout: float = 300, + ) -> AsyncIterator[AsyncThread]: + """Serialize one conversation while allowing independent conversations to run concurrently.""" + from openai_codex import TransportClosedError + + await self.start() + self._active += 1 + self._idle.clear() + state = None + created_token = None + thread = None + codex = None + try: + codex = await self._get_client() + await self._prune_threads(codex) + state = self._threads.setdefault(thread_id, _ThreadState()) if thread_id and not fork else _ThreadState() + state.users += 1 + state.revision += 1 + async with state.lock: + options = dict(options or {}) + if jobs is not None: + injected = json.loads(json.dumps(injected or {})) + binding = json.dumps({"jobs": [job.name for job in jobs], "injected": injected}, sort_keys=True) + if state.binding is not None and state.binding != binding: + raise ValueError( + "job_tools, tool_context_id and injected_job_kwargs cannot change on a loaded thread", + ) + state.binding = binding + config = dict(options.get("config") or {}) + servers = dict(config.get("mcp_servers") or {}) + if "reme_jobs" in servers: + raise ValueError("mcp_servers.reme_jobs is reserved for ReMe job tools") + # Codex validates the transport even when a server is disabled. + servers["reme_jobs"] = {"url": "http://127.0.0.1", "enabled": False} + if jobs: + url = await self._mcp_url() + if state.token is None: + state.token = self.jobs_mcp.register(jobs, injected) + created_token = state.token + servers["reme_jobs"] = { + "url": url, + "http_headers": {"Authorization": f"Bearer {state.token}"}, + "required": True, + "enabled_tools": [job.name for job in jobs], + "startup_timeout_sec": mcp_startup_timeout, + "tool_timeout_sec": mcp_tool_timeout, + } + config["mcp_servers"] = servers + options["config"] = config + if fork: + source = self._threads.setdefault(thread_id, _ThreadState()) + source.users += 1 + source.revision += 1 + try: + async with source.lock: + thread = await codex.thread_fork(thread_id, **options) + finally: + source.users -= 1 + elif thread_id: + thread = await codex.thread_resume(thread_id, **options) + else: + thread = await codex.thread_start(**options) + self._threads[thread.id] = state + if thread_id and not fork: + await self.wait_idle(thread) + try: + yield thread + finally: + # Interrupt acknowledgement alone does not mean that tool execution stopped. + error = sys.exception() + try: + async with asyncio.timeout(30 if error is not None else None): + await self.wait_idle(thread) + except Exception as exc: + if isinstance(exc, TransportClosedError) and self._codex is codex: + self._failed = True + if error is None: + raise + self.logger.warning(f"Failed to wait for Codex thread {thread.id}: {exc}") + else: + from openai_codex.generated.v2_all import ThreadUnsubscribeResponse + + try: + await codex._client.request( # pylint: disable=protected-access + "thread/unsubscribe", + {"threadId": thread.id}, + response_model=ThreadUnsubscribeResponse, + ) + except Exception as exc: + self.logger.warning(f"Failed to unsubscribe Codex thread {thread.id}: {exc}") + except TransportClosedError: + if codex is not None and self._codex is codex: + self._failed = True + raise + finally: + if state is not None: + state.users -= 1 + if thread is None and created_token is not None: + self.jobs_mcp.scopes.pop(created_token, None) + state.token = None + self._active -= 1 + if not self._active: + self._idle.set() + + async def _close(self) -> None: + await self._idle.wait() + try: + await self._close_client() + finally: + self._codex = None + self._failed = False + self._http_url = None + self._threads.clear() + await self.jobs_mcp.close() diff --git a/reme/components/service/http_service.py b/reme/components/service/http_service.py index 1d90ab2f..5863852a 100644 --- a/reme/components/service/http_service.py +++ b/reme/components/service/http_service.py @@ -67,10 +67,15 @@ class HttpService(BaseService): def build_service(self, app: "Application") -> None: """Create one FastAPI app containing JSON/SSE and optional MCP routes.""" + from fastmcp.utilities.lifespan import combine_lifespans + lifespan = self._lifespan(app, self.host, self.port) + servers = self.app_context.components.get("codex_app_server", {}) if self.app_context else {} + # TLS listeners may use certificates that cannot be verified through a loopback address. + if self.kwargs.get("ssl_certfile") or self.kwargs.get("ssl_keyfile"): + servers = {} if self.mcp_enabled: from fastmcp import FastMCP - from fastmcp.utilities.lifespan import combine_lifespans self.mcp_server = FastMCP(name=app.config.app_name) self.mcp_app = self.mcp_server.http_app( @@ -84,6 +89,13 @@ class HttpService(BaseService): title=app.config.app_name, lifespan=lifespan, ) + for server in servers.values(): + gateway = server.jobs_mcp + self.service.router.lifespan_context = combine_lifespans( + gateway.app.lifespan, + self.service.router.lifespan_context, + ) + self.service.router.routes.append(Route(gateway.path, endpoint=gateway, methods=["GET", "POST", "DELETE"])) cors_origins = ["*"] self.service.add_middleware( CORSMiddleware, # type: ignore[arg-type] diff --git a/reme/components/service/mcp_tools.py b/reme/components/service/mcp_tools.py index 77209f2c..e29bd323 100644 --- a/reme/components/service/mcp_tools.py +++ b/reme/components/service/mcp_tools.py @@ -13,11 +13,20 @@ def add_mcp_job( tool_error_on_failure: bool, ) -> bool: """Register a non-stream job as an MCP tool on ``server``.""" + tool = make_mcp_job_tool(job, injected_job_kwargs=injected_job_kwargs, tool_error_on_failure=tool_error_on_failure) + if tool is None: + return False + server.add_tool(tool) + return True + + +def make_mcp_job_tool(job: BaseJob, *, injected_job_kwargs: dict[str, Any], tool_error_on_failure: bool) -> Any: + """Build a tool bound to an existing job and server-owned arguments.""" from fastmcp.exceptions import ToolError from fastmcp.tools import FunctionTool if isinstance(job, StreamJob): - return False + return None async def execute_tool(**kwargs): conflicts = sorted(injected_job_kwargs.keys() & kwargs.keys()) @@ -43,12 +52,9 @@ def add_mcp_job( if "required" in parameters: parameters["required"] = [name for name in parameters["required"] if name not in injected_names] - server.add_tool( - FunctionTool( - name=job.name, - description=job.description, - fn=execute_tool, - parameters=parameters, - ), + return FunctionTool( + name=job.name, + description=job.description, + fn=execute_tool, + parameters=parameters, ) - return True diff --git a/reme/components/service/scoped_mcp.py b/reme/components/service/scoped_mcp.py new file mode 100644 index 00000000..c5c1e681 --- /dev/null +++ b/reme/components/service/scoped_mcp.py @@ -0,0 +1,99 @@ +"""Session-scoped MCP tools backed by the current Application's jobs.""" + +import asyncio +from contextlib import contextmanager +import secrets +import socket +from typing import Any + +from .mcp_tools import make_mcp_job_tool + + +class ScopedMCP: + """Serve isolated tool catalogs on one authenticated, application-owned endpoint.""" + + def __init__(self) -> None: + self.path = f"/_reme/codex/{secrets.token_hex(12)}" + self.scopes: dict[str, list[Any]] = {} + self._app: Any = None + self._server: Any = None + self._task: asyncio.Task | None = None + self._socket: socket.socket | None = None + + @property + def app(self) -> Any: + """Build one MCP transport; resolve tools separately for every HTTP request.""" + if self._app is None: + from fastmcp import FastMCP + from fastmcp.server.dependencies import get_http_headers + from fastmcp.server.providers import Provider + + scopes = self.scopes + + class JobProvider(Provider): + """Select the immutable tool binding carried by the request's bearer token.""" + + async def _list_tools(self): + authorization = get_http_headers(include={"authorization"}).get("authorization", "") + return scopes.get(authorization.removeprefix("Bearer "), []) + + self._app = FastMCP("ReMe agent tools", providers=[JobProvider()], tasks=False).http_app( + path=self.path, + stateless_http=True, + ) + return self._app + + def register(self, jobs: list[Any], injected: dict[str, Any]) -> str: + """Bind selected jobs and hidden arguments without cloning the Application.""" + token = secrets.token_urlsafe(32) + self.scopes[token] = [ + make_mcp_job_tool(job, injected_job_kwargs=injected, tool_error_on_failure=True) for job in jobs + ] + return token + + async def __call__(self, scope, receive, send): + if scope["type"] == "http": + authorization = dict(scope["headers"]).get(b"authorization", b"").decode("latin-1") + if not authorization.startswith("Bearer ") or authorization[7:] not in self.scopes: + from starlette.responses import Response + + await Response(status_code=403)(scope, receive, send) + return + await self.app(scope, receive, send) # pylint: disable=not-callable + + async def listen(self) -> str: + """Start a loopback listener for embedded/CLI hosts that have no HTTP service.""" + if self._server is None: + import uvicorn + + class LocalServer(uvicorn.Server): + """Leave signal handling to the containing ReMe application.""" + + @contextmanager + def capture_signals(self): + """Keep the caller's signal handlers intact.""" + yield + + self._socket = socket.socket() + self._socket.bind(("127.0.0.1", 0)) + self._server = LocalServer(uvicorn.Config(self, log_config=None, access_log=False, ws="none")) + self._task = asyncio.create_task(self._server.serve(sockets=[self._socket])) + while not self._server.started: + if self._task.done(): + await self._task + raise RuntimeError("Codex tool endpoint stopped before startup") + await asyncio.sleep(0.01) + return f"http://127.0.0.1:{self._socket.getsockname()[1]}{self.path}" + + async def close(self) -> None: + """Stop the optional listener and revoke all tool bindings.""" + try: + if self._server is not None: + self._server.should_exit = True # pylint: disable=attribute-defined-outside-init + try: + await self._task + finally: + self._socket.close() + self._server = self._task = self._socket = None + finally: + self.scopes.clear() diff --git a/reme/config/default.yaml b/reme/config/default.yaml index d2d481ba..677adf8c 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -930,6 +930,58 @@ components: max_tokens: 65536 thinking_enable: false + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${CODEX_API_KEY:-} + base_url: ${CODEX_BASE_URL:-} + codex_home: .codex + supports_websockets: false + config_overrides: + - 'model_reasoning_effort="xhigh"' + - 'review_model="qwen3.8-max"' + - 'agents.default_subagent_model="qwen3.8-max"' + - 'agents.default_subagent_reasoning_effort="xhigh"' + - 'memories.extract_model="qwen3.8-flash"' + - 'memories.consolidation_model="qwen3.8-max"' + model_catalog: + models: + - &codex_model_defaults + slug: &codex_simple_model qwen3.8-flash + display_name: Qwen3.8 Flash + description: Simple tasks and automatic approval review + auto_review_model_override: *codex_simple_model + default_reasoning_level: xhigh + multi_agent_reasoning_effort: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 1 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + - <<: *codex_model_defaults + slug: &codex_complex_model qwen3.8-max + display_name: Qwen3.8 Max + description: Main conversations and complex tasks + default_reasoning_level: xhigh + priority: 0 + oauth: + backend: sdk + auth_mode: oauth + codex_home: ${CODEX_HOME:-~/.codex} + agent_wrapper: default: backend: agentscope @@ -951,18 +1003,16 @@ components: permission_mode: bypassPermissions codex: backend: codex - auth_mode: api_key - model: ${CODEX_MODEL_NAME:-} - api_key: ${CODEX_API_KEY:-} - base_url: ${CODEX_BASE_URL:-} - approval_mode: auto_review + app_server: default + model: ${CODEX_MODEL_NAME:-qwen3.8-max} + effort: xhigh + approval_mode: deny_all sandbox: full-access codex_oauth: backend: codex - auth_mode: oauth - model: ${CODEX_MODEL_NAME:-} - codex_home: ${CODEX_HOME:-~/.codex} - approval_mode: auto_review + app_server: oauth + model: ${CODEX_OAUTH_MODEL_NAME:-} + approval_mode: deny_all sandbox: full-access file_graph: diff --git a/reme/config/demo.yaml b/reme/config/demo.yaml index b883ce83..dd66cafb 100644 --- a/reme/config/demo.yaml +++ b/reme/config/demo.yaml @@ -115,6 +115,54 @@ components: max_tokens: 65536 thinking_enable: false + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + codex_home: .codex + supports_websockets: false + config_overrides: + - 'model_reasoning_effort="xhigh"' + - 'review_model="qwen3.8-max"' + - 'agents.default_subagent_model="qwen3.8-max"' + - 'agents.default_subagent_reasoning_effort="xhigh"' + - 'memories.extract_model="qwen3.8-flash"' + - 'memories.consolidation_model="qwen3.8-max"' + model_catalog: + models: + - &codex_model_defaults + slug: &codex_simple_model qwen3.8-flash + display_name: Qwen3.8 Flash + description: Simple tasks and automatic approval review + auto_review_model_override: *codex_simple_model + default_reasoning_level: xhigh + multi_agent_reasoning_effort: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 1 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + - <<: *codex_model_defaults + slug: &codex_complex_model qwen3.8-max + display_name: Qwen3.8 Max + description: Main conversations and complex tasks + default_reasoning_level: xhigh + priority: 0 + agent_wrapper: as: backend: agentscope @@ -140,10 +188,10 @@ components: permission_mode: bypassPermissions codex: backend: codex - auth_mode: oauth - model: ${CODEX_MODEL_NAME:-} - codex_home: ${CODEX_HOME:-~/.codex} + app_server: default + model: *codex_complex_model + effort: xhigh job_tools: - add - approval_mode: auto_review + approval_mode: deny_all sandbox: full-access diff --git a/reme/config/proactive.yaml b/reme/config/proactive.yaml new file mode 100644 index 00000000..7f189229 --- /dev/null +++ b/reme/config/proactive.yaml @@ -0,0 +1,266 @@ +# Minimal Studio workspace with one Codex chat agent using API-key authentication. +# Set LLM_API_KEY / LLM_BASE_URL (and optionally CODEX_MODEL_NAME) for chat. +# Custom LLM_BASE_URL endpoints must support the Responses API. +# Start: reme start config=proactive +# Open: http://127.0.0.1:2333 + +service: + backend: http + web_enabled: true + mcp_enabled: false + +jobs: + # Keep Markdown search and graph data in sync with the workspace files. + index_update_loop: + backend: background + watch_dirs: [daily_dir, digest_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 + monitor_type: file_store + dispatch_steps: + - backend: update_index_step + persist: false + + chat: + backend: stream + description: "Stream a read-only agent conversation over the ReMe workspace." + parameters: + type: object + properties: + query: {type: string, description: "user message"} + session_id: {type: string, description: "agent session to resume"} + system_prompt: {type: string, description: "optional system prompt override"} + required: [query] + steps: + - backend: chat_step + agent_wrapper: default + + app_config: + backend: base + description: "Return the effective application config with secrets redacted." + steps: + - backend: app_config_step + + version: + backend: base + description: "Return the ReMe version." + steps: + - backend: version_step + + status: + backend: base + description: "Report component memory estimates and process RSS." + steps: + - backend: status_step + + health_check: + backend: base + description: "Return a concise health-check snapshot of ReMe components." + steps: + - backend: health_check_step + + graph_snapshot: + backend: base + description: "Return the category-rooted digest wikilink graph with daily-note leaves." + steps: + - backend: graph_snapshot_step + + reindex: + backend: base + description: "Rebuild the local BM25 index from already-ingested chunks." + parameters: + type: object + properties: + scope: {type: string, enum: [all, bm25], default: all} + steps: + - backend: reindex_step + + list: + backend: base + description: "List files under a workspace path." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative directory; empty = root", default: ""} + recursive: {type: boolean, default: false} + sort_by: {type: string, enum: [mtime]} + extensions: {type: array, items: {type: string}} + limit: {type: integer, default: 100} + steps: + - backend: list_step + + load: + backend: base + description: "Load a complete text file for the local editor without agent-output truncation." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + required: [path] + steps: + - backend: load_step + max_bytes: 5242880 + + save: + backend: base + description: "Save text verbatim, optionally rejecting external changes since the file was opened." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + content: {type: string, description: "complete file content, including frontmatter"} + expected_mtime: {type: string, description: "optional ISO mtime returned by load/stat"} + required: [path, content] + steps: + - backend: save_step + + # ChatStep exposes these read-only jobs as the agent's tools. + search: + backend: base + description: "Search indexed daily and digest Markdown files with BM25." + parameters: + type: object + properties: + query: {type: string, description: "search query"} + limit: {type: integer, default: 5} + min_score: {type: number, default: 0.0} + start_date: {type: string, description: "optional inclusive YYYY-MM-DD start date"} + end_date: {type: string, description: "optional inclusive YYYY-MM-DD end date"} + required: [query] + steps: + - backend: search_step + + read: + backend: base + description: "Read a Markdown file under the workspace." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative Markdown path"} + start_line: {type: integer, description: "first line (1-based, inclusive)"} + end_line: {type: integer, description: "last line (1-based, inclusive)"} + required: [path] + steps: + - backend: read_step + + read_image: + backend: base + description: "Read an image file as base64." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative image path"} + required: [path] + steps: + - backend: read_image_step + max_bytes: 5242880 + + frontmatter_read: + backend: base + description: "Read a file's frontmatter." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative Markdown path"} + required: [path] + steps: + - backend: frontmatter_read_step + + stat: + backend: base + description: "Stat a workspace path (size, mtime, exists, is_dir, is_file)." + parameters: + type: object + properties: + path: {type: string, description: "workspace-relative path"} + required: [path] + steps: + - backend: stat_step + + traverse: + backend: base + description: "Return a bounded wikilink graph with frontend-ready nodes and directed edges." + parameters: + type: object + properties: + path: + oneOf: + - {type: string} + - {type: array, items: {type: string}, minItems: 1} + description: "workspace-relative path or paths used as graph seeds" + depth: {type: integer, minimum: 0, default: 1} + direction: {type: string, enum: [forward, backward, both], default: both} + required: [path] + steps: + - backend: traverse_step + +components: + codex_app_server: + default: + backend: sdk + auth_mode: api_key + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + codex_home: .codex + model_catalog: + models: + - slug: qwen3.8-max + display_name: Qwen3.8 Max + description: Studio chat + default_reasoning_level: xhigh + supported_reasoning_levels: + - effort: xhigh + description: Maximum reasoning + context_window: 983616 + input_modalities: [text, image] + supports_parallel_tool_calls: false + supports_reasoning_summaries: false + support_verbosity: false + shell_type: default + visibility: list + supported_in_api: true + priority: 0 + base_instructions: "" + experimental_supported_tools: [] + truncation_policy: + mode: bytes + limit: 10000 + + agent_wrapper: + default: + backend: codex + app_server: default + model: ${CODEX_MODEL_NAME:-qwen3.8-max} + effort: xhigh + approval_mode: deny_all + sandbox: read-only + + tokenizer: + default: + backend: regex + + keyword_index: + default: + backend: bm25 + tokenizer: default + + file_graph: + default: + backend: local + + file_chunker: + markdown: + backend: markdown + supported_extensions: [md] + + file_store: + default: + backend: local + embedding_store: "" + keyword_index: default + file_graph: default diff --git a/reme/enumeration/component_enum.py b/reme/enumeration/component_enum.py index 8b317d61..37e03962 100644 --- a/reme/enumeration/component_enum.py +++ b/reme/enumeration/component_enum.py @@ -40,6 +40,8 @@ class ComponentEnum(str, Enum): AGENT_WRAPPER = "agent_wrapper" + CODEX_APP_SERVER = "codex_app_server" + MAIL = "mail" OUTBOUND_PROXY = "outbound_proxy" diff --git a/reme/steps/common/stream_llm_demo.py b/reme/steps/common/stream_llm_demo.py index 4fca8e71..aa554acf 100644 --- a/reme/steps/common/stream_llm_demo.py +++ b/reme/steps/common/stream_llm_demo.py @@ -10,7 +10,7 @@ class StreamLLMDemoStep(BaseStep): """Drive an Agent powered by the ``agent_wrapper`` component with streaming output. When streaming is enabled on the context, text/thinking/tool events are - pushed chunk-by-chunk via ``self.context.add_stream_string``. + pushed chunk-by-chunk via ``self.context.add_stream_chunk``. When streaming is not enabled, falls back to non-streaming reply. Inputs (from RuntimeContext): @@ -63,7 +63,7 @@ class StreamLLMDemoStep(BaseStep): text_parts: list[str] = [] async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): - await self.context.add_stream_string(chunk.chunk, chunk.chunk_type) + await self.context.add_stream_chunk(chunk) if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): text_parts.append(chunk.chunk) diff --git a/tests/integration/test_codex_agent_wrapper.py b/tests/integration/test_codex_agent_wrapper.py index 2cc45a11..3cef6b4d 100644 --- a/tests/integration/test_codex_agent_wrapper.py +++ b/tests/integration/test_codex_agent_wrapper.py @@ -1,23 +1,17 @@ -"""Opt-in live coverage for the Codex app-server wrapper. - -Run with ``REME_CODEX_INTEGRATION=1``. These tests consume the caller's active -API-key or Codex OAuth account and are intentionally excluded from normal CI. -""" +"""Opt-in live Codex SDK coverage; set REME_CODEX_INTEGRATION=1 to use credentials.""" # pylint: disable=protected-access import asyncio import os from pathlib import Path -import subprocess -from types import SimpleNamespace import pytest from pydantic import BaseModel -from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper -from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import ApplicationConfig, Response +from reme.application import Application +from reme.enumeration import ChunkEnum +from reme.steps.base_step import BaseStep pytestmark = pytest.mark.skipif( os.getenv("REME_CODEX_INTEGRATION") != "1", @@ -29,213 +23,123 @@ class _StructuredResult(BaseModel): marker: str -class _CustomJob: - name = "only_custom" - description = "Return the fixed marker CUSTOM_JOB_OK." - parameters = {"type": "object", "properties": {}, "additionalProperties": False} +class _CustomStep(BaseStep): + """Exercise a parent-only registration and the actual parent metadata.""" - async def __call__(self, **_kwargs): - return Response(answer="CUSTOM_JOB_OK") - - -class _DraftJob: - def __init__(self, name: str, parameters: dict): - self.name = name - self.description = f"Live tool-context contract job: {name}" - self.parameters = parameters - - -def _child_pids(root_pid: int) -> set[int]: - result = subprocess.run( - ["ps", "-axo", "pid=,ppid="], - check=True, - capture_output=True, - text=True, - ) - children: dict[int, list[int]] = {} - for line in result.stdout.splitlines(): - pid, parent = (int(value) for value in line.split()) - children.setdefault(parent, []).append(pid) - found: set[int] = set() - pending = list(children.get(root_pid, [])) - while pending: - pid = pending.pop() - if pid not in found: - found.add(pid) - pending.extend(children.get(pid, [])) - return found + async def execute(self): + self.app_context.metadata["live_custom_calls"] = self.app_context.metadata.get("live_custom_calls", 0) + 1 + self.context.response.answer = "CUSTOM_JOB_OK" @pytest.mark.asyncio -async def test_live_codex_reply_stream_tools_skills_resume_fork_approval_and_close(tmp_path): - """Exercise the live Codex wrapper contract when explicitly enabled.""" +async def test_live_shared_codex_sessions_tools_streaming_and_close(tmp_path): + """Use one real SDK process and execute every job inside the containing Application.""" workspace = tmp_path / "workspace" workspace.mkdir() - project_skill = Path(__file__).resolve().parents[2] / "skills" / "reme_memory" skills_root = workspace / "skills" skills_root.mkdir() + project_skill = Path(__file__).resolve().parents[2] / "skills" / "reme_memory" (skills_root / "reme_memory").symlink_to(project_skill, target_is_directory=True) - - app_config = ApplicationConfig( + api_key = base_url = "" + for prefix in ("CODEX", "OPENAI", "LLM"): + if api_key := os.getenv(f"{prefix}_API_KEY", ""): + base_url = os.getenv(f"{prefix}_BASE_URL", "") + break + codex_home = os.getenv("REME_CODEX_HOME") + if not api_key and codex_home is None: + codex_home = str(Path.home() / ".codex") + app = Application( workspace_dir=str(workspace), - mem_session_dir="sessions", enable_logo=False, log_to_console=False, log_to_file=False, service={"backend": "mcp"}, + components={ + "codex_app_server": { + "default": { + "backend": "sdk", + "auth_mode": "api_key" if api_key else "oauth", + "api_key": api_key, + "base_url": base_url, + "codex_home": codex_home, + }, + }, + "agent_wrapper": { + name: {"backend": "codex", "model": os.getenv("CODEX_MODEL_NAME")} for name in ("first", "second") + }, + }, jobs={ "only_custom": { "backend": "base", - "description": _CustomJob.description, - "parameters": _CustomJob.parameters, - "steps": [], + "description": "Return CUSTOM_JOB_OK.", + "enable_serve": False, + "steps": [{"backend": "live_custom_step"}], }, - "add_live_draft": { + "add_draft": { "backend": "base", "description": "Append text to the current tool context.", - "parameters": { - "type": "object", - "properties": {"text": {"type": "string"}}, - "required": ["text"], - }, + "parameters": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}, "steps": [{"backend": "add_draft_step"}], }, - "read_live_draft": { + "read_draft": { "backend": "base", "description": "Read text accumulated in the current tool context.", "parameters": {"type": "object", "properties": {}}, "steps": [{"backend": "read_all_draft_step"}], }, }, - components={ComponentEnum.AS_LLM: {}}, ) - context = SimpleNamespace( - app_config=app_config, - jobs={ - "only_custom": _CustomJob(), - "add_live_draft": _DraftJob( - "add_live_draft", - { - "type": "object", - "properties": {"text": {"type": "string"}}, - "required": ["text"], - }, - ), - "read_live_draft": _DraftJob( - "read_live_draft", - {"type": "object", "properties": {}}, - ), - }, - ) - codex_home = os.getenv("REME_CODEX_HOME") - if codex_home is None and not any(os.getenv(name) for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY")): - codex_home = str(Path.home() / ".codex") - wrapper = CodexAgentWrapper(app_context=context, codex_home=codex_home) - - await wrapper.start() - first = await wrapper.reply( - "Call add_live_draft with text STATE_OK, then reply FIRST_TURN_OK.", - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert "FIRST_TURN_OK" in first["last_message"] - - resumed = await wrapper.reply( - "Call read_live_draft. Include its exact result and RESUME_OK in your answer.", - resume=first["session_id"], - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert resumed["session_id"] == first["session_id"] - assert "RESUME_OK" in resumed["last_message"] - assert "STATE_OK" in resumed["last_message"] - - isolated = await wrapper.reply( - "Call read_live_draft. If it is empty, reply ISOLATED_OK.", - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="other-context", - ) - assert "ISOLATED_OK" in isolated["last_message"] - assert "STATE_OK" not in isolated["last_message"] - - resumed_after_isolated_context = await wrapper.reply( - "Call read_live_draft again and include its exact result.", - resume=first["session_id"], - job_tools=["add_live_draft", "read_live_draft"], - tool_context_id="live-context", - ) - assert "STATE_OK" in resumed_after_isolated_context["last_message"] - - forked = await wrapper.reply( - "Reply with exactly FORK_OK.", - resume=first["session_id"], - fork_session=True, - tool_context_id="fork-context", - ) - assert forked["session_id"] != first["session_id"] - assert "FORK_OK" in forked["last_message"] - - structured = await wrapper.reply( - "Return marker STRUCTURED_OK.", - output_schema=_StructuredResult, - ) - assert structured["structured_output"] == {"marker": "STRUCTURED_OK"} - - tool_result = await wrapper.reply( - "Call the only_custom tool once, then include its result in your answer.", - job_tools=["only_custom"], - tool_context_id="tool-context", - ) - assert "CUSTOM_JOB_OK" in tool_result["last_message"] - - skill_result = await wrapper.reply( - "Use the reme_memory skill. Read only its instructions and reply SKILL_OK; do not run its scripts or CLI.", - skills=["reme_memory"], - ) - assert "SKILL_OK" in skill_result["last_message"] - assert (workspace / ".agents" / "skills" / "reme_memory").is_symlink() - - outside_target = tmp_path / "approval-target.txt" - approval_chunks = [ - chunk - async for chunk in wrapper.reply_stream( - f"Try to write the word approved to {outside_target} using a shell command.", - approval_mode="auto_review", - sandbox="workspace-write", + app.context.registry.add("live_custom_step", _CustomStep, owner=__name__) + server = app.context.components["codex_app_server"]["default"] + async with app: + first, second = (app.context.components["agent_wrapper"][name] for name in ("first", "second")) + original = await first.reply( + "Call add_draft with text STATE_OK, then reply FIRST_OK.", + job_tools=["add_draft", "read_draft"], + tool_context_id="one", ) - ] - assert approval_chunks[0].chunk_type == ChunkEnum.REPLY_START - assert approval_chunks[-1].chunk_type == ChunkEnum.REPLY_END - assert any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in approval_chunks) - - codex_proc = wrapper._codex._client._sync._proc - assert codex_proc is not None and codex_proc.poll() is None - child_pids = _child_pids(codex_proc.pid) - await wrapper.close() - await asyncio.sleep(0.2) - - assert codex_proc.poll() is not None - running_pids = { - int(line) - for line in subprocess.run( - ["ps", "-axo", "pid="], - check=True, - capture_output=True, - text=True, - ).stdout.split() - } - assert child_pids.isdisjoint(running_pids) - assert wrapper._codex is None - assert wrapper._mcp_snapshot_path is None - - second_context = SimpleNamespace(app_config=app_config, jobs=context.jobs.copy()) - second_wrapper = CodexAgentWrapper(app_context=second_context, codex_home=codex_home) - await second_wrapper.start() - isolated_application = await second_wrapper.reply( - "Call read_live_draft. If it is empty, reply NEW_APPLICATION_OK.", - job_tools=["read_live_draft"], - tool_context_id="live-context", - ) - await second_wrapper.close() - assert "NEW_APPLICATION_OK" in isolated_application["last_message"] - assert "STATE_OK" not in isolated_application["last_message"] + assert "FIRST_OK" in original["last_message"] + resumed = await second.reply( + "Call read_draft and include its exact result.", + session_id=original["session_id"], + job_tools=["add_draft", "read_draft"], + tool_context_id="one", + ) + assert resumed["session_id"] == original["session_id"] + assert "STATE_OK" in resumed["last_message"] + isolated = await second.reply( + "Call read_draft. If it is empty, reply ISOLATED_OK.", + job_tools=["read_draft"], + tool_context_id="two", + ) + assert isolated["session_id"] != original["session_id"] + assert "ISOLATED_OK" in isolated["last_message"] + assert "STATE_OK" not in isolated["last_message"] + forked = await first.reply("Reply FORK_OK.", resume=original["session_id"], fork_session=True) + assert forked["session_id"] != original["session_id"] + assert "FORK_OK" in forked["last_message"] + custom = await first.reply("Call only_custom and include its result.", job_tools=["only_custom"]) + assert "CUSTOM_JOB_OK" in custom["last_message"] + assert app.context.metadata["live_custom_calls"] >= 1 + structured = await first.reply("Return marker STRUCTURED_OK.", output_schema=_StructuredResult) + assert structured["structured_output"] == {"marker": "STRUCTURED_OK"} + await first.compact_session(structured["session_id"]) + skill = await first.reply( + "Read the reme_memory skill instructions and reply SKILL_OK. Do not execute its scripts.", + skills=["reme_memory"], + ) + assert "SKILL_OK" in skill["last_message"] + assert (workspace / ".agents/skills/reme_memory").is_symlink() + concurrent = await asyncio.gather(first.reply("Reply ONE."), second.reply("Reply TWO.")) + assert concurrent[0]["session_id"] != concurrent[1]["session_id"] + chunks = [chunk async for chunk in second.reply_stream("Reply STREAM_OK.")] + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.REPLY_END + process = server._codex._client._sync._proc + assert process.poll() is None + await first.close() + assert process.poll() is None + assert (await second.reply("Reply STILL_OK."))["result"] + assert process.poll() is not None + assert server._codex is None + assert server.jobs_mcp._task is None diff --git a/tests/integration/test_codex_demo.py b/tests/integration/test_codex_demo.py new file mode 100644 index 00000000..2755ca62 --- /dev/null +++ b/tests/integration/test_codex_demo.py @@ -0,0 +1,176 @@ +"""Opt-in live coverage of both demo Steps using Codex API-key authentication.""" + +# pylint: disable=redefined-outer-name + +import asyncio +from collections import Counter +import json +import os +from pathlib import Path + +import pytest +import pytest_asyncio + +from reme.application import Application +from reme.config import resolve_app_config +from reme.enumeration import ChunkEnum +from reme.steps.common.stream_llm_demo import StreamLLMDemoStep +from reme.utils import global_counter_get, parse_env_file + +pytestmark = pytest.mark.skipif( + os.getenv("REME_CODEX_INTEGRATION") != "1", + reason="set REME_CODEX_INTEGRATION=1 to run live Codex tests", +) + + +@pytest_asyncio.fixture +async def demo_app(tmp_path, monkeypatch, request): + """Load the requested LLM credentials while keeping all runtime state temporary.""" + env_path = Path(__file__).resolve().parents[2] / ".env" + environment = {**(parse_env_file(env_path) if env_path.is_file() else {}), **os.environ} + for name in ("LLM_API_KEY", "LLM_BASE_URL"): + assert environment.get(name), f"{name} is required for the live Codex demo tests" + monkeypatch.setenv(name, environment[name]) + config = resolve_app_config( + config="demo", + log_config=False, + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + ) + approval_mode = getattr(request, "param", "deny_all") + config["components"]["agent_wrapper"]["codex"]["approval_mode"] = approval_mode + server_config = config["components"]["codex_app_server"]["default"] + assert server_config["auth_mode"] == "api_key" + assert server_config["api_key"] == environment["LLM_API_KEY"] + assert server_config["base_url"] == environment["LLM_BASE_URL"] + app = Application(**config) + server = app.context.components["codex_app_server"]["default"] + async with app: + yield app + models = set() + reasoning = set() + for path in (server.session_path / "sessions").rglob("*.jsonl"): + for line in path.read_text(encoding="utf-8").splitlines(): + record = json.loads(line) + payload = record.get("payload", {}) + if record.get("type") == "event_msg" and payload.get("type") == "thread_settings_applied": + settings = payload["thread_settings"] + models.add(settings["model"]) + reasoning.add((settings["model"], settings.get("reasoning_effort"))) + print(f"\nCodex models observed: {sorted(models)}") + print(f"Codex reasoning observed: {sorted(reasoning)}") + assert ("qwen3.8-max", "xhigh") in reasoning + assert {effort for _, effort in reasoning} == {"xhigh"} + assert models == ( + {"qwen3.8-max", "qwen3.8-flash"} + if approval_mode == "auto_review" and global_counter_get(app.context.metadata, ["__job_counter", "add"]) + else {"qwen3.8-max"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True], ids=["non-stream", "stream"]) +@pytest.mark.parametrize("use_tool", [False, True], ids=["chat", "add-tool"]) +async def test_live_codex_demo_chat_and_add_tool(demo_app, streaming, use_tool): + """Verify text delivery and actual ReMe add execution through each demo Job.""" + query = ( + "Use the add tool to calculate 21 + 21 and return only its result." + if use_tool + else "Write the integers from 1 through 20 in order, separated by spaces." + ) + sys_prompt = "Use the add tool for addition. Otherwise follow the requested output format." + previous_calls = global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) + async with asyncio.timeout(90): + if streaming: + chunks = [ + chunk + async for chunk in demo_app.run_stream_job( + "stream_llm_demo", + query=query, + sys_prompt=sys_prompt, + ) + ] + errors = [ + chunk.chunk + for chunk in chunks + if chunk.chunk_type == ChunkEnum.ERROR and chunk.metadata.get("will_retry") is not True + ] + assert not errors, errors + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.DONE + assert chunks[-1].done + assert not any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert any(chunk.chunk_type == ChunkEnum.REPLY_END for chunk in chunks) + assert chunks[-2].metadata.get("status") == "completed" + content = [chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk] + text = "".join(content).strip() + counts = Counter(chunk.chunk_type.value for chunk in chunks) + print(f"\nstream tool={use_tool}: chunks={dict(counts)}, answer={text!r}") + if use_tool: + assert counts[ChunkEnum.TOOL_CALL.value] > 0 + assert counts[ChunkEnum.TOOL_RESULT.value] > 0 + else: + response = await demo_app.run_job("llm_demo", query=query, sys_prompt=sys_prompt) + assert response.success, response.answer + text = response.answer.strip() + assert response.metadata["response"] == text + print(f"\nnon-stream tool={use_tool}: answer={text!r}") + assert text + if use_tool: + assert "42" in text + assert global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) > previous_calls + else: + assert text.split() == [str(number) for number in range(1, 21)] + + +@pytest.mark.asyncio +async def test_live_codex_stream_demo_without_stream_queue(demo_app): + """Without a queue, StreamLLMDemoStep uses the non-streaming wrapper reply.""" + step = StreamLLMDemoStep(app_context=demo_app.context, agent_wrapper="codex") + async with asyncio.timeout(90): + response = await step(query="Reply exactly FALLBACK_OK.") + assert response.success, response.answer + assert response.answer.strip() == "FALLBACK_OK" + + +@pytest.mark.asyncio +async def test_live_codex_native_shell_tool_needs_no_approval(demo_app): + """The default policy also executes a native shell tool without an approval request.""" + wrapper = demo_app.context.components["agent_wrapper"]["codex"] + async with asyncio.timeout(90): + chunks = [ + chunk + async for chunk in wrapper.reply_stream( + "Run the shell command printf NATIVE_TOOL_OK, then return only the command output.", + ) + ] + assert any( + chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_name == "commandExecution" for chunk in chunks + ) + assert any( + chunk.chunk_type == ChunkEnum.TOOL_RESULT and chunk.tool_call_name == "commandExecution" for chunk in chunks + ) + assert not any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert "NATIVE_TOOL_OK" in "".join( + chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("demo_app", ["auto_review"], indirect=True) +async def test_live_codex_automatic_review_remains_available(demo_app): + """Explicitly selecting auto_review still invokes the configured Flash reviewer.""" + async with asyncio.timeout(90): + chunks = [ + chunk + async for chunk in demo_app.run_stream_job( + "stream_llm_demo", + query="Call the add tool to calculate 21 + 21, then return only its result.", + ) + ] + assert any(chunk.chunk_type == ChunkEnum.APPROVAL for chunk in chunks) + assert global_counter_get(demo_app.context.metadata, ["__job_counter", "add"]) > 0 + assert "42" in "".join(chunk.chunk for chunk in chunks if chunk.chunk_type == ChunkEnum.CONTENT) diff --git a/tests/integration/test_codex_wrapper_capabilities.py b/tests/integration/test_codex_wrapper_capabilities.py new file mode 100644 index 00000000..f529c76a --- /dev/null +++ b/tests/integration/test_codex_wrapper_capabilities.py @@ -0,0 +1,139 @@ +"""Opt-in live checks for skills, additional instructions, and ReMe job tools.""" + +# pylint: disable=redefined-outer-name + +import os +from pathlib import Path +import secrets + +import pytest +import pytest_asyncio + +from reme.application import Application +from reme.steps.base_step import BaseStep + +pytestmark = pytest.mark.skipif( + os.getenv("REME_CODEX_INTEGRATION") != "1", + reason="set REME_CODEX_INTEGRATION=1 to run live Codex tests", +) + + +class _ProbeStep(BaseStep): + """Record execution in the containing Application rather than a child process.""" + + async def execute(self): + """Capture the tool input and return a deterministic result.""" + text = self.context.data["text"] + self.app_context.metadata.setdefault("codex_probe_calls", []).append( + {"text": text, "tool_context_id": self.context.data.get("tool_context_id")}, + ) + self.context.response.answer = f"JOB_OK:{text}" + + +@pytest_asyncio.fixture +async def wrapper(tmp_path): + """Use configured credentials with an isolated workspace and real SDK transport.""" + api_key = base_url = "" + for prefix in ("CODEX", "OPENAI", "LLM"): + if api_key := os.getenv(f"{prefix}_API_KEY", ""): + base_url = os.getenv(f"{prefix}_BASE_URL", "") + break + codex_home = os.getenv("REME_CODEX_HOME") + if not api_key and codex_home is None: + codex_home = str(Path.home() / ".codex") + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "mcp"}, + components={ + "codex_app_server": { + "default": { + "backend": "sdk", + "auth_mode": "api_key" if api_key else "oauth", + "api_key": api_key, + "base_url": base_url, + "codex_home": codex_home, + }, + }, + "agent_wrapper": {"default": {"backend": "codex", "model": os.getenv("CODEX_MODEL_NAME")}}, + }, + jobs={ + "capability_probe_job": { + "backend": "base", + "description": "Return the supplied text prefixed with JOB_OK:.", + "enable_serve": False, + "parameters": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + }, + "steps": [{"backend": "codex_capability_probe_step"}], + }, + }, + ) + app.context.registry.add("codex_capability_probe_step", _ProbeStep, owner=__name__) + async with app: + yield app.context.components["agent_wrapper"]["default"] + + +@pytest.mark.asyncio +async def test_live_codex_add_skill_reads_instructions(wrapper, tmp_path): + """The expected answer exists only inside the newly added skill file.""" + marker = secrets.token_hex(16) + source = tmp_path / "skills" / "capability_probe" + source.mkdir(parents=True) + (source / "SKILL.md").write_text( + "---\nname: capability_probe\ndescription: Return the verification token when explicitly invoked.\n---\n\n" + f"When invoked, reply with exactly this token and no other text: {marker}\n", + encoding="utf-8", + ) + wrapper.add_skills(["capability_probe"]) + result = await wrapper.reply("$capability_probe") + assert result["result"].strip() == marker + assert (tmp_path / ".agents/skills/capability_probe").resolve() == source.resolve() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resume", + [ + False, + pytest.param( + True, + marks=pytest.mark.xfail( + reason="Codex ignores developer instruction overrides on an already loaded thread", + strict=True, + raises=AssertionError, + ), + ), + ], +) +async def test_live_codex_extend_system_prompt(wrapper, resume): + """Check additional instructions on new sessions and expose the loaded-session limitation.""" + first_marker, added_marker = secrets.token_hex(16), secrets.token_hex(16) + wrapper.set_system_prompt(f"For input ORIGINAL_PROMPT_PROBE, reply exactly {first_marker}.") + original = await wrapper.reply("ORIGINAL_PROMPT_PROBE") + assert original["result"].strip() == first_marker + wrapper.set_system_prompt( + wrapper.kwargs["system_prompt"] + f"\nFor input ADDED_PROMPT_PROBE, reply exactly {added_marker}.", + ) + options = {"session_id": original["session_id"]} if resume else {} + extended = await wrapper.reply("ADDED_PROMPT_PROBE", **options) + assert (extended["session_id"] == original["session_id"]) == resume + assert extended["result"].strip() == added_marker + + +@pytest.mark.asyncio +async def test_live_codex_add_job_tool_calls_current_application(wrapper): + """A private ReMe job executes through MCP with its hidden context binding.""" + marker = secrets.token_hex(16) + wrapper.add_job_tools(["capability_probe_job"]) + result = await wrapper.reply( + f"Call capability_probe_job with text {marker}, then return only its exact result.", + tool_context_id="codex_capability_probe", + ) + calls = wrapper.app_context.metadata["codex_probe_calls"] + assert calls == [{"text": marker, "tool_context_id": "codex_capability_probe"}] + assert result["result"].strip() == f"JOB_OK:{marker}" diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index 6f436d52..f4f86937 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -1,340 +1,1052 @@ -"""Unit tests for the Codex agent wrapper and its FastMCP bridge.""" +"""Shared Codex runtime, isolated conversations, and same-Application MCP tools.""" -# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,too-many-lines +# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access,redefined-outer-name import asyncio from dataclasses import dataclass import json -import os from pathlib import Path -import stat -import sys +import socket from types import SimpleNamespace +from unittest.mock import AsyncMock +import httpx import pytest +import pytest_asyncio +from fastmcp import Client +from fastmcp.client import StreamableHttpTransport +from fastmcp.exceptions import ToolError from openai_codex.generated.v2_all import TokenUsageBreakdown from pydantic import BaseModel +from reme.application import Application +from reme.components.application_context import ApplicationContext from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper -from reme.components.agent_wrapper.codex_mcp_server import _prepare_config -from reme.components.job import BackgroundJob -from reme.components.outbound_proxy import FixedHttpOutboundProxy +from reme.components.codex_app_server import CodexAppServer +from reme.components.job import BackgroundJob, StreamJob +from reme.components.service.scoped_mcp import ScopedMCP from reme.config import resolve_app_config from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import ApplicationConfig, Response - -# A fresh bridge imports ReMe before replying to initialize; keep this separate -# from the 10-second budget for requests to an already initialized server. -STDIO_INIT_TIMEOUT = 30 +from reme.schema import Response class _Job: - def __init__(self, name="search"): - self.name = name - self.description = "Search memory" - self.parameters = { - "type": "object", - "properties": {"query": {"type": "string"}}, - "required": ["query"], - } + name = "search" + description = "Search memory" + enabled = True + enable_serve = False + parameters = { + "type": "object", + "properties": {"query": {"type": "string"}, "_allowed_paths": {"type": "array"}}, + "required": ["query"], + } + + def __init__(self): self.calls = [] async def __call__(self, **kwargs): self.calls.append(kwargs) - return Response(answer=f"found:{kwargs['query']}") + return Response(answer={"query": kwargs["query"], "calls": len(self.calls)}) + + +@dataclass +class _Result: + final_response: str + usage: object = None + + +class _Turn: + def __init__(self, thread, inputs): + self.thread = thread + self.inputs = inputs + self.id = f"turn-{len(thread.inputs)}" + + async def run(self): + try: + if self.thread.codex.run_hook: + await self.thread.codex.run_hook(self.thread, self.inputs) + return _Result(self.inputs) + finally: + self.thread.busy = False + + async def interrupt(self): + self.thread.interrupts += 1 + self.thread.busy = False + + async def stream(self): + try: + yield SimpleNamespace(method="turn/started", payload=SimpleNamespace(turn=SimpleNamespace(id=self.id))) + yield SimpleNamespace( + method="item/agentMessage/delta", + payload=SimpleNamespace(item_id="item", delta=self.inputs), + ) + self.thread.busy = False + yield SimpleNamespace( + method="turn/completed", + payload=SimpleNamespace( + turn=SimpleNamespace( + id=self.id, + status=SimpleNamespace(value="completed"), + duration_ms=1, + error=None, + ), + ), + ) + finally: + self.thread.stream_closed = True + + +class _Thread: + def __init__(self, codex, thread_id): + self.codex, self.id = codex, thread_id + self.inputs = [] + self.options = [] + self._busy = False + self.interrupts = 0 + self.stream_closed = False + + @property + def busy(self): + return self._busy + + @busy.setter + def busy(self, value): + self._busy = value + self.codex.notify_status(self.id, "active" if value else "idle") + + async def turn(self, inputs, **kwargs): + assert not self.busy, "overlapping turns on one thread" + self.inputs.append(inputs) + self.options.append(kwargs) + self.busy = True + return _Turn(self, inputs) + + async def read(self): + return SimpleNamespace( + thread=SimpleNamespace( + status=SimpleNamespace(root=SimpleNamespace(type="active" if self.busy else "idle")), + ), + ) + + async def compact(self): + self.busy = True + if self.codex.compact_hook: + await self.codex.compact_hook(self) + else: + self.busy = False + + +class _Codex: + instances = [] + + def __init__(self, config): + self.config = config + self.threads = {} + self.loaded = set() + self.unsubscribed = [] + self.closed = 0 + self.logins = [] + self.opened = [] + self.run_hook = None + self.compact_hook = None + self.notifications = asyncio.Queue() + self._client = SimpleNamespace(request=self.request, next_notification=self.next_notification) + self.instances.append(self) + + async def account(self): + await asyncio.sleep(0) + return SimpleNamespace(account=object()) + + async def login_api_key(self, key): + self.logins.append(key) + + async def close(self): + from openai_codex import TransportClosedError + + self.closed += 1 + self.notifications.put_nowait(TransportClosedError("closed")) + + def notify_status(self, thread_id, status): + self.notifications.put_nowait( + SimpleNamespace( + method="thread/status/changed", + payload=SimpleNamespace(thread_id=thread_id, status=SimpleNamespace(root=SimpleNamespace(type=status))), + ), + ) + + async def next_notification(self): + notification = await self.notifications.get() + if isinstance(notification, Exception): + raise notification + return notification + + async def thread_start(self, **options): + thread = _Thread(self, f"thread-{len(self.threads) + 1}") + self.threads[thread.id] = thread + self.loaded.add(thread.id) + self.opened.append(("start", thread.id, options)) + return thread + + async def thread_resume(self, thread_id, **options): + self.loaded.add(thread_id) + self.opened.append(("resume", thread_id, options)) + return self.threads[thread_id] + + async def thread_fork(self, thread_id, **options): + forked = await self.thread_start(**options) + self.opened[-1] = ("fork", thread_id, options) + return forked + + async def request(self, method, params, **_kwargs): + if method == "thread/loaded/list": + return SimpleNamespace(data=list(self.loaded), next_cursor=None) + assert method == "thread/unsubscribe" + self.unsubscribed.append(params["threadId"]) + return SimpleNamespace(status="unsubscribed") def _wrapper(tmp_path, **kwargs): + context = ApplicationContext(workspace_dir=str(tmp_path)) job = _Job() - config = SimpleNamespace( - workspace_dir=str(tmp_path), - mem_session_dir="mem_session", - environment={}, - components={ComponentEnum.AS_LLM: {}}, - model_dump=lambda **_kwargs: { - "workspace_dir": str(tmp_path), - "enable_logo": False, - "log_to_console": False, - "log_to_file": False, - "jobs": {}, - "components": {}, - }, + context.jobs[job.name] = job + server = CodexAppServer( + name="default", + app_context=context, + auth_mode="oauth", + codex_home=kwargs.pop("codex_home", None), ) - context = SimpleNamespace(app_config=config, components={}, jobs={job.name: job}) + context.components = {ComponentEnum.CODEX_APP_SERVER: {"default": server}} return CodexAgentWrapper(app_context=context, **kwargs), job -def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - - config = wrapper._mcp_server_config( # pylint: disable=protected-access - {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, - ) - - assert config["command"] - assert config["enabled_tools"] == ["search"] - assert config["args"].count("--job") == 1 - assert config["args"][config["args"].index("--job") + 1] == "search" - assert "reme.components.agent_wrapper.codex_mcp_server" in config["args"] - assert config["args"][config["args"].index("--config") + 1] == str(tmp_path / "custom.yaml") - assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" - - -def test_mcp_config_serializes_injected_job_kwargs(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - - config = wrapper._mcp_server_config( # pylint: disable=protected-access - { - "job_tools": ["search"], - "tool_context_id": "ctx-1", - "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]}, - }, - ) - - raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] - assert json.loads(raw) == {"_allowed_paths": ["daily/2025-06-01/note.md"]} - - plain = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access - assert "--injected-job-kwargs" not in plain["args"] - - -def test_prepare_config_merges_injected_job_kwargs_with_tool_context(): - prepared = _prepare_config( - {"jobs": {"selected": {"backend": "base"}}}, - ["selected"], - "ctx-1", - {"_allowed_paths": ["daily/2025-06-01/note.md"]}, - ) - - assert prepared["service"]["injected_job_kwargs"] == { - "_allowed_paths": ["daily/2025-06-01/note.md"], - "tool_context_id": "ctx-1", - } - - -def test_thread_config_preserves_other_mcp_servers(tmp_path): - wrapper, _job = _wrapper(tmp_path) - config = wrapper._thread_config( # pylint: disable=protected-access - { - "job_tools": ["search"], - "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}}, - }, - ) - - assert "docs" in config["mcp_servers"] - assert len(config["mcp_servers"]) == 2 - assert next(name for name in config["mcp_servers"] if name != "docs").startswith("reme_jobs_") - - -@pytest.mark.asyncio -async def test_thread_config_injects_proxy_only_into_codex_shell_policy(tmp_path): - wrapper, _job = _wrapper(tmp_path) - proxy = FixedHttpOutboundProxy(url="http://127.0.0.1:18080") - await proxy.start() - wrapper.app_context.components = {ComponentEnum.OUTBOUND_PROXY: {"default": proxy}} - await wrapper.start() - - config = wrapper._thread_config( # pylint: disable=protected-access - { - "config": { - "shell_environment_policy": { - "inherit": "core", - "set": { - "CUSTOM": "preserved", - "HTTP_PROXY": "http://user-proxy.example:8080", - }, - }, - }, - }, - ) - environment = config["shell_environment_policy"]["set"] - - assert config["shell_environment_policy"]["inherit"] == "core" - assert environment["CUSTOM"] == "preserved" - for key in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"): - assert environment[key] == proxy.http_url - - auth = wrapper._resolve_auth_config("oauth") # pylint: disable=protected-access - client_config = wrapper._build_client_config(auth) # pylint: disable=protected-access - assert "HTTP_PROXY" not in client_config.env - assert "HTTPS_PROXY" not in client_config.env - - await wrapper.close() - await proxy.close() - - -def test_mcp_config_rejects_background_jobs(tmp_path): - wrapper, _job = _wrapper(tmp_path) - wrapper.app_context.jobs["watch"] = BackgroundJob(name="watch", app_context=wrapper.app_context) - - with pytest.raises(TypeError, match="non-stream request jobs"): - wrapper._mcp_server_config({"job_tools": ["watch"]}) - - -def test_prepare_config_reuses_selected_stdio_mcp_service(): - prepared = _prepare_config( - { - "service": {"backend": "http"}, - "jobs": { - "selected": {"backend": "base", "enable_serve": False}, - "helper": {"backend": "base"}, - "watch": {"backend": "background"}, - }, - }, - ["selected"], - "ctx-1", - ) - - assert prepared["service"] == { - "backend": "mcp", - "transport": "stdio", - "jobs": ["selected"], - "tool_error_on_failure": True, - "injected_job_kwargs": {"tool_context_id": "ctx-1"}, - } - assert set(prepared["jobs"]) == {"selected", "helper"} - assert prepared["jobs"]["selected"]["enable_serve"] is True - - -def test_prepare_config_rejects_missing_or_background_selected_jobs(): - with pytest.raises(KeyError, match="missing"): - _prepare_config({"jobs": {}}, ["missing"]) - with pytest.raises(KeyError, match="watch"): - _prepare_config({"jobs": {"watch": {"backend": "background"}}}, ["watch"]) - - -@pytest.mark.asyncio -async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - config_path = tmp_path / "bridge.json" - config_path.write_text( - json.dumps( - { - "service": {"backend": "mcp"}, - "workspace_dir": str(tmp_path / "workspace"), - "jobs": { - "empty": { - "backend": "base", - "description": "Return an empty response", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - }, - ), - encoding="utf-8", - ) - transport = StdioTransport( - command=sys.executable, - args=[ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(config_path), - "--workspace", - str(tmp_path / "workspace"), - "--job", - "empty", - ], - cwd=str(Path(__file__).resolve().parents[2]), - keep_alive=False, - ) - - async with Client(transport, timeout=None, init_timeout=STDIO_INIT_TIMEOUT) as client: - tools = await asyncio.wait_for(client.list_tools(), timeout=10) - - assert [tool.name for tool in tools] == ["empty"] - - -@pytest.mark.asyncio -async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): - config_path = tmp_path / "bridge.json" - config_path.write_text( - json.dumps( - { - "service": {"backend": "mcp"}, - "workspace_dir": str(tmp_path / "workspace"), - "jobs": { - "empty": { - "backend": "base", - "description": "Empty", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - }, - ), - encoding="utf-8", - ) - proc = await asyncio.create_subprocess_exec( - sys.executable, - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(config_path), - "--workspace", - str(tmp_path / "workspace"), - "--job", - "empty", - cwd=str(Path(__file__).resolve().parents[2]), - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - request = { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2025-06-18", - "capabilities": {}, - "clientInfo": {"name": "test", "version": "1"}, - }, - } +@pytest_asyncio.fixture +async def system(tmp_path, monkeypatch): + _Codex.instances = [] + monkeypatch.setattr("openai_codex.AsyncCodex", _Codex) + first, job = _wrapper(tmp_path, name="first") + context = first.app_context + second = CodexAgentWrapper(name="second", app_context=context) + server = context.components[ComponentEnum.CODEX_APP_SERVER]["default"] + monkeypatch.setattr(server, "_mcp_url", AsyncMock(return_value="http://127.0.0.1/mcp")) try: - assert proc.stdin is not None and proc.stdout is not None and proc.stderr is not None - proc.stdin.write((json.dumps(request) + "\n").encode()) - await proc.stdin.drain() - first_line = await asyncio.wait_for(proc.stdout.readline(), timeout=STDIO_INIT_TIMEOUT) - message = json.loads(first_line) - assert message["jsonrpc"] == "2.0" - assert message["id"] == 1 + yield first, second, server, job finally: - if proc.returncode is None: - proc.terminate() + await first.close() + await second.close() + await server.close() + + +@pytest.mark.asyncio +async def test_wrappers_share_one_lazy_client_and_keep_independent_sessions(system): + first, second, server, _ = system + await first.start() + assert not _Codex.instances + results = await asyncio.gather(first.reply("first"), second.reply("second")) + assert len(_Codex.instances) == 1 + assert results[0]["session_id"] != results[1]["session_id"] + assert first.app_server is second.app_server is server + client = _Codex.instances[0] + await first.close() + assert client.closed == 0 + assert (await second.reply("still running"))["result"] == "still running" + await server.close() + assert client.closed == 1 + + +@pytest.mark.asyncio +async def test_different_threads_run_concurrently_and_same_thread_serializes_across_wrappers(system): + first, second, server, _ = system + thread_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + entered, release = asyncio.Event(), asyncio.Event() + observed = [] + + async def run_hook(thread, inputs): + observed.append((thread.id, inputs)) + if inputs == "hold": + entered.set() + await release.wait() + + client.run_hook = run_hook + async with asyncio.TaskGroup() as group: + held = group.create_task(first.reply("hold", session_id=thread_id)) + await asyncio.wait_for(entered.wait(), 2) + queued = group.create_task(second.reply("same", session_id=thread_id)) try: - await asyncio.wait_for(proc.wait(), timeout=10) - except asyncio.TimeoutError: - proc.kill() - await proc.wait() - stdout = first_line + await proc.stdout.read() - stderr = (await proc.stderr.read()).decode() - assert b"Loading config" not in stdout - assert b"INFO" not in stdout - assert b"WARNING" not in stdout - assert b"2026-" not in stdout - assert "Failed to parse JSONRPC message" not in stderr - assert "Invalid JSON" not in stderr + independent = await asyncio.wait_for(second.reply("other"), 2) + assert independent["session_id"] != thread_id + assert not queued.done() + assert (thread_id, "same") not in observed + finally: + release.set() + assert held.result()["session_id"] == queued.result()["session_id"] == thread_id + assert len(server._threads) == 2 -def test_sdk_responds_to_interactive_approval_server_requests(): - from openai_codex.client import CodexClient +@pytest.mark.asyncio +@pytest.mark.parametrize("approval_mode", [None, "auto_review"]) +async def test_resume_fork_and_sdk_options_remain_explicit(system, approval_mode): + from openai_codex import ApprovalMode, Sandbox + from openai_codex.types import Personality, ReasoningEffort - client = CodexClient() - assert client._default_approval_handler( - "item/commandExecution/requestApproval", - {}, - ) == { # pylint: disable=protected-access - "decision": "accept", - } - assert client._default_approval_handler( - "item/fileChange/requestApproval", - {}, - ) == { # pylint: disable=protected-access - "decision": "accept", + first, _, _, _ = system + original = await first.reply( + "one", + model="custom", + system_prompt="instructions", + personality="friendly", + approval_mode=approval_mode, + ) + thread_id = original["session_id"] + resumed = await first.reply("two", resume=thread_id, effort="high") + forked = await first.reply("three", session_id=thread_id, fork_session=True) + assert resumed["session_id"] == thread_id + assert forked["session_id"] != thread_id + client = _Codex.instances[0] + options = client.opened[0][2] + assert options["model"] == "custom" + assert options["developer_instructions"] == "instructions" + assert options["approval_mode"] == (ApprovalMode.auto_review if approval_mode else ApprovalMode.deny_all) + assert options["sandbox"] == Sandbox.full_access + assert options["personality"] == Personality.friendly + assert client.threads[thread_id].options[-1]["effort"] == ReasoningEffort.high + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changes", + [ + {"tool_context_id": "other"}, + {"job_tools": []}, + {"injected_job_kwargs": {"_allowed_paths": ["other.md"]}}, + ], +) +async def test_loaded_thread_constraints_cannot_be_changed_by_another_wrapper(system, changes): + first, second, server, _ = system + binding = { + "job_tools": ["search"], + "tool_context_id": "one", + "injected_job_kwargs": {"_allowed_paths": ["one.md"]}, } + original = await first.reply("one", **binding) + thread_id = original["session_id"] + token = server._threads[thread_id].token + with pytest.raises(ValueError, match="cannot change"): + await second.reply("two", session_id=thread_id, **{**binding, **changes}) + assert token in server.jobs_mcp.scopes + await second.reply("three", session_id=thread_id, **binding) + + +@pytest.mark.asyncio +async def test_unloaded_thread_rebinds_tools_and_revokes_old_token(system): + first, _, server, _ = system + original = await first.reply("one", job_tools=["search"], tool_context_id="one") + thread_id = original["session_id"] + token = server._threads[thread_id].token + _Codex.instances[0].loaded.clear() + await first.reply("two", session_id=thread_id, job_tools=["search"], tool_context_id="two") + assert token not in server.jobs_mcp.scopes + assert server._threads[thread_id].token != token + + +@pytest.mark.asyncio +async def test_cancelling_reply_interrupts_only_its_turn_and_keeps_shared_server(system): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + entered = asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await asyncio.Event().wait() + + client = _Codex.instances[0] + client.run_hook = hold + task = asyncio.create_task(first.reply("cancel", session_id=thread_id)) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + client.run_hook = None + assert client.threads[thread_id].interrupts == 1 + assert client.closed == 0 + assert (await second.reply("okay"))["result"] == "okay" + assert not client.threads[thread_id].busy + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["turn", "compact"]) +@pytest.mark.parametrize("fail", [False, True]) +async def test_cancelled_dispatch_settles_before_resuming_the_thread(system, monkeypatch, operation, fail): + first, second, server, _ = system + thread_id = (await first.reply("initial"))["session_id"] + thread = _Codex.instances[0].threads[thread_id] + original = getattr(thread, operation) + entered, release = asyncio.Event(), asyncio.Event() + + async def dispatch(*args, **kwargs): + entered.set() + await release.wait() + if fail: + raise ValueError("dispatch rejected") + return await original(*args, **kwargs) + + monkeypatch.setattr(thread, operation, dispatch) + request = ( + first.compact_session(thread_id) if operation == "compact" else first.reply("cancel", session_id=thread_id) + ) + task = asyncio.create_task(request) + queued = None + try: + await asyncio.wait_for(entered.wait(), 2) + task.cancel("original cancellation") + await asyncio.sleep(0) + task.cancel("cancel again") + await asyncio.sleep(0) + queued = asyncio.create_task(second.reply("later", session_id=thread_id)) + await asyncio.sleep(0.01) + assert not queued.done() + monkeypatch.setattr(thread, operation, original) + release.set() + with pytest.raises(asyncio.CancelledError, match="original cancellation"): + await asyncio.wait_for(task, 2) + assert (await asyncio.wait_for(queued, 2))["result"] == "later" + assert thread.interrupts == int(operation == "turn" and not fail) + await asyncio.wait_for(server.close(), 2) + finally: + release.set() + task.cancel() + if queued is not None: + queued.cancel() + await asyncio.gather(task, *([queued] if queued is not None else []), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_stream_close_interrupts_and_routes_chunks_per_session(system): + first, second, _, _ = system + stream = first.reply_stream("early") + start = await anext(stream) + thread_id = start.session_id + await stream.aclose() + thread = _Codex.instances[0].threads[thread_id] + assert thread.interrupts == 1 + assert thread.stream_closed + chunks = [chunk async for chunk in second.reply_stream("complete")] + assert chunks[0].chunk_type == ChunkEnum.REPLY_START + assert chunks[-1].chunk_type == ChunkEnum.REPLY_END + assert all(chunk.session_id != thread_id for chunk in chunks) + + +@pytest.mark.asyncio +async def test_compaction_holds_thread_lock_until_it_finishes(system): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + entered, release = asyncio.Event(), asyncio.Event() + + async def compact_hook(thread): + entered.set() + + async def finish(): + await release.wait() + thread.busy = False + + asyncio.create_task(finish()) + + _Codex.instances[0].compact_hook = compact_hook + async with asyncio.TaskGroup() as group: + group.create_task(first.compact_session(thread_id)) + await asyncio.wait_for(entered.wait(), 2) + queued = group.create_task(second.reply("later", session_id=thread_id)) + try: + await asyncio.sleep(0.02) + assert not queued.done() + finally: + release.set() + assert queued.result()["session_id"] == thread_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_compaction_acknowledgement_does_not_release_the_thread_before_it_starts(system, monkeypatch, cancelled): + first, second, _, _ = system + thread_id = (await first.reply("initial"))["session_id"] + client = _Codex.instances[0] + thread = client.threads[thread_id] + acknowledged = asyncio.Event() + + async def compact(): + acknowledged.set() + + monkeypatch.setattr(thread, "compact", compact) + async with asyncio.TaskGroup() as group: + compacting = group.create_task(first.compact_session(thread_id)) + await asyncio.wait_for(acknowledged.wait(), 2) + if cancelled: + compacting.cancel() + queued = group.create_task(second.reply("later", session_id=thread_id)) + try: + await asyncio.sleep(0.01) + # Cancellation may end the caller during cleanup; its barrier must still block the next turn. + if not cancelled: + assert not compacting.done() + assert not queued.done() + client.notify_status(thread_id, "idle") + await asyncio.sleep(0) + if not cancelled: + assert not compacting.done() + assert not queued.done() + finally: + client.notify_status(thread_id, "active") + client.notify_status(thread_id, "idle") + if cancelled: + assert compacting.cancelled() + assert queued.result()["session_id"] == thread_id + + +@pytest.mark.asyncio +async def test_close_waits_for_active_request_without_closing_other_wrappers(system): + first, second, server, _ = system + await first.reply("initial") + entered, release = asyncio.Event(), asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await release.wait() + + client = _Codex.instances[0] + client.run_hook = hold + async with asyncio.TaskGroup() as group: + group.create_task(first.reply("held")) + await asyncio.wait_for(entered.wait(), 2) + closing = group.create_task(first.close()) + try: + await asyncio.sleep(0) + assert not closing.done() + assert client.closed == 0 + finally: + release.set() + client.run_hook = None + await second.reply("still works") + assert server.is_started + + +class _StructuredModel(BaseModel): + ok: bool + + +@pytest.mark.asyncio +async def test_structured_reply_and_native_inputs(system): + from openai_codex import LocalImageInput, TextInput + + first, _, _, _ = system + result = await first.reply('{"ok": true}', output_schema=_StructuredModel) + assert result["structured_output"] == {"ok": True} + inputs = [TextInput("describe"), LocalImageInput("image.png")] + await first.reply(inputs) + assert _Codex.instances[0].threads["thread-2"].inputs[0] is inputs + with pytest.raises(NotImplementedError): + await anext(first.reply_stream("no", output_schema=_StructuredModel)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kwargs", [{"resume": "a", "session_id": "b"}, {"fork_session": True}]) +async def test_invalid_session_selection_fails_without_starting_client(system, kwargs): + first, _, _, _ = system + with pytest.raises(ValueError): + await first.reply("invalid", **kwargs) + assert not _Codex.instances + + +@pytest.mark.asyncio +@pytest.mark.parametrize("job_class", [BackgroundJob, StreamJob]) +async def test_unsupported_job_types_are_rejected(system, job_class): + first, _, _, _ = system + first.app_context.jobs["unsupported"] = job_class(name="unsupported") + with pytest.raises(TypeError, match="non-stream"): + await first.reply("invalid", job_tools=["unsupported"]) + + +@pytest.mark.parametrize("name", sorted(CodexAgentWrapper._SERVER_OPTIONS)) +def test_process_options_must_be_configured_on_server(name): + with pytest.raises(TypeError, match="codex_app_server"): + CodexAgentWrapper(**{name: "invalid"}) + + +def test_server_config_authentication_environment_and_home(tmp_path): + context = ApplicationContext(workspace_dir=str(tmp_path), environment={"CUSTOM": "preserved"}) + server = CodexAppServer( + name="one", + app_context=context, + auth_mode="api_key", + api_key="key", + base_url="https://proxy.test/v1", + ) + config = server._client_config() + assert config.env["CUSTOM"] == "preserved" + assert config.env["CODEX_HOME"] == str(tmp_path / ".codex") + assert 'forced_login_method="api"' in config.config_overrides + assert 'model_provider="reme"' in config.config_overrides + assert 'model_providers.reme.base_url="https://proxy.test/v1"' in config.config_overrides + assert "model_providers.reme.requires_openai_auth=true" in config.config_overrides + assert "model_providers.reme.supports_websockets=false" in config.config_overrides + oauth = CodexAppServer(app_context=context, auth_mode="oauth", api_key="ignored", base_url="ignored") + assert CodexAppServer(codex_home="~/.codex").session_path == Path.home() / ".codex" + assert CodexAppServer(app_context=context, codex_home="other-home").session_path == tmp_path / "other-home" + assert oauth._client_config().config_overrides == ('forced_login_method="chatgpt"',) + with pytest.raises(ValueError, match="requires a non-empty"): + CodexAppServer(app_context=context, auth_mode="api_key")._client_config() + + +def test_custom_provider_websocket_support_and_explicit_overrides_are_preserved(tmp_path): + server = CodexAppServer( + app_context=ApplicationContext(workspace_dir=str(tmp_path)), + auth_mode="api_key", + api_key="key", + base_url="https://proxy.test/v1", + supports_websockets=True, + config_overrides=['model_provider="custom"'], + ) + overrides = server._client_config().config_overrides + assert "model_providers.reme.supports_websockets=true" in overrides + assert overrides.index('model_provider="custom"') > overrides.index('model_provider="reme"') + + +def test_server_materializes_model_catalog_inside_its_private_home(tmp_path): + context = ApplicationContext(workspace_dir=str(tmp_path)) + catalog = {"models": [{"slug": "main", "auto_review_model_override": "reviewer"}, {"slug": "reviewer"}]} + server = CodexAppServer(name="catalog", app_context=context, model_catalog=catalog) + config = server._client_config() + path = server.session_path / "reme-model-catalog.json" + assert json.loads(path.read_text(encoding="utf-8")) == catalog + assert f"model_catalog_json={json.dumps(str(path))}" in config.config_overrides + assert path.is_relative_to(tmp_path) + with pytest.raises(ValueError, match="not both"): + CodexAppServer( + app_context=context, + model_catalog=catalog, + config_overrides=['model_catalog_json="external.json"'], + )._client_config() + + +@pytest.mark.asyncio +async def test_client_initialization_failure_is_closed_and_retryable(system, monkeypatch): + first, _, server, _ = system + monkeypatch.setattr(_Codex, "account", AsyncMock(side_effect=RuntimeError("login failed"))) + with pytest.raises(RuntimeError, match="login failed"): + await first.reply("fails") + assert _Codex.instances[0].closed == 1 + assert server._codex is None + assert server._active == 0 + + +def test_default_and_demo_configuration_use_named_shared_servers(monkeypatch): + monkeypatch.delenv("CODEX_HOME", raising=False) + monkeypatch.delenv("CODEX_MODEL_NAME", raising=False) + monkeypatch.setenv("LLM_API_KEY", "demo-api-key") + monkeypatch.setenv("LLM_BASE_URL", "https://demo.invalid/v1") + for name in ("default", "demo"): + config = resolve_app_config(config=name, log_config=False) + wrappers = config["components"]["agent_wrapper"] + assert wrappers["codex"]["app_server"] == "default" + assert "auth_mode" not in wrappers["codex"] + servers = config["components"]["codex_app_server"] + assert servers["default"]["auth_mode"] == "api_key" + assert servers["default"]["codex_home"] == ".codex" + assert servers["default"]["supports_websockets"] is False + assert wrappers["codex"]["model"] == "qwen3.8-max" + assert wrappers["codex"]["approval_mode"] == "deny_all" + assert wrappers["codex"]["effort"] == "xhigh" + models = servers["default"]["model_catalog"]["models"] + assert {model["slug"] for model in models} == {"qwen3.8-max", "qwen3.8-flash"} + assert all(model["auto_review_model_override"] == "qwen3.8-flash" for model in models) + assert all(model["default_reasoning_level"] == "xhigh" for model in models) + assert all(model["multi_agent_reasoning_effort"] == "xhigh" for model in models) + assert all([level["effort"] for level in model["supported_reasoning_levels"]] == ["xhigh"] for model in models) + assert set(servers["default"]["config_overrides"]) == { + 'model_reasoning_effort="xhigh"', + 'review_model="qwen3.8-max"', + 'agents.default_subagent_model="qwen3.8-max"', + 'agents.default_subagent_reasoning_effort="xhigh"', + 'memories.extract_model="qwen3.8-flash"', + 'memories.consolidation_model="qwen3.8-max"', + } + if name == "demo": + assert servers["default"]["api_key"] == "demo-api-key" + assert servers["default"]["base_url"] == "https://demo.invalid/v1" + else: + assert servers["oauth"]["codex_home"] == "~/.codex" + + +@pytest.mark.asyncio +async def test_scoped_mcp_calls_original_jobs_and_isolates_catalog_arguments_and_failures(): + gateway = ScopedMCP() + job = _Job() + first = gateway.register([job], {"tool_context_id": "one", "_allowed_paths": ["one.md"]}) + second = gateway.register([job], {"tool_context_id": "two", "_allowed_paths": ["two.md"]}) + empty = gateway.register([], {}) + try: + url = await gateway.listen() + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {first}"})) as one: + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {second}"})) as two: + tools = await one.list_tools() + assert [tool.name for tool in tools] == ["search"] + assert "_allowed_paths" not in tools[0].input_schema["properties"] + await one.call_tool("search", {"query": "a"}) + await two.call_tool("search", {"query": "b"}) + await one.call_tool("search", {"query": "c"}) + with pytest.raises(ToolError, match="injected"): + await one.call_tool("search", {"query": "bad", "_allowed_paths": ["escape"]}) + assert [call["tool_context_id"] for call in job.calls] == ["one", "two", "one"] + assert job.calls[1]["_allowed_paths"] == ["two.md"] + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {empty}"})) as client: + assert not await client.list_tools() + with pytest.raises(ToolError): + await client.call_tool("search", {"query": "not allowed"}) + async with httpx.AsyncClient() as client: + assert (await client.post(url, json={})).status_code == 403 + assert (await client.post(url, headers={"Authorization": "Bearer unknown"}, json={})).status_code == 403 + finally: + await gateway.close() + assert gateway._task is None + assert not gateway.scopes + + +@pytest.mark.asyncio +async def test_mounted_mcp_transport_lifespan_can_restart(): + gateway = ScopedMCP() + app = gateway.app + for _ in range(2): + async with app.lifespan(app): + assert gateway.app is app + + +@pytest.mark.asyncio +async def test_scoped_mcp_can_restart_its_loopback_listener(): + gateway = ScopedMCP() + job = _Job() + try: + for _ in range(2): + token = gateway.register([job], {}) + url = await gateway.listen() + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {token}"})) as client: + await client.call_tool("search", {"query": "same application"}) + await gateway.close() + assert len(job.calls) == 2 + finally: + await gateway.close() + + +@pytest.mark.asyncio +async def test_http_service_mounts_scoped_tools_without_another_listener(tmp_path, monkeypatch): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False}, + components={ + "codex_app_server": {"default": {"backend": "sdk", "auth_mode": "oauth"}}, + "agent_wrapper": {"one": {"backend": "codex"}, "two": {"backend": "codex"}}, + }, + ) + service = app.context.service + service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + writer = SimpleNamespace(close=lambda: None, wait_closed=AsyncMock()) + monkeypatch.setattr(asyncio, "open_connection", AsyncMock(return_value=(None, writer))) + job = _Job() + token = server.jobs_mcp.register([job], {"tool_context_id": "parent"}) + + def factory(**kwargs): + return httpx.AsyncClient(transport=httpx.ASGITransport(app=service.service), **kwargs) + + async with service.service.router.lifespan_context(service.service): + url = await server._mcp_url() + async with Client( + StreamableHttpTransport( + url, + headers={"Authorization": f"Bearer {token}"}, + httpx_client_factory=factory, + ), + ) as client: + await client.call_tool("search", {"query": "same application"}) + assert job.calls == [{"query": "same application", "tool_context_id": "parent"}] + assert server.jobs_mcp._server is None + assert app.context.components["agent_wrapper"]["one"].app_server is server + assert app.context.components["agent_wrapper"]["two"].app_server is server + assert not server.is_started + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tls_options", [{"ssl_certfile": "server.pem"}, {"ssl_keyfile": "server.key"}]) +async def test_tls_http_service_uses_managed_loopback_for_scoped_tools(tmp_path, monkeypatch, tls_options): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False, **tls_options}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + service = app.context.service + service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + probe = AsyncMock(side_effect=AssertionError("must not probe the TLS listener")) + monkeypatch.setattr(asyncio, "open_connection", probe) + job = _Job() + token = server.jobs_mcp.register([job], {"tool_context_id": "parent"}) + listener_task = None + async with service.service.router.lifespan_context(service.service): + url = await server._mcp_url() + assert url.startswith("http://127.0.0.1:") + assert await server._mcp_url() == url + listener_task = server.jobs_mcp._task + assert listener_task is not None + async with Client(StreamableHttpTransport(url, headers={"Authorization": f"Bearer {token}"})) as client: + await client.call_tool("search", {"query": "TLS host"}) + assert job.calls == [{"query": "TLS host", "tool_context_id": "parent"}] + probe.assert_not_awaited() + assert listener_task.done() + assert server.jobs_mcp._task is None + assert not server.jobs_mcp.scopes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("option", ["cwd", "app_server"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_call_time_wrapper_options_are_rejected_instead_of_silently_ignored(system, option, stream): + first, _, _, _ = system + with pytest.raises(TypeError, match="wrapper options"): + if stream: + await anext(first.reply_stream("invalid", **{option: "ignored"})) + else: + await first.reply("invalid", **{option: "ignored"}) + assert not _Codex.instances + + +@pytest.mark.asyncio +async def test_mounted_tools_wait_for_the_main_listener_before_returning_a_url(tmp_path): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + sock.setblocking(False) + port = sock.getsockname()[1] + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "port": port, "web_enabled": False}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + app.context.service.build_service(app) + server = app.context.components["codex_app_server"]["default"] + lookup = asyncio.create_task(server._mcp_url()) + listener = None + + async def accept(_reader, writer): + writer.close() + await writer.wait_closed() + + try: + await asyncio.sleep(0.02) + assert not lookup.done() + listener = await asyncio.start_server(accept, sock=sock) + url = await asyncio.wait_for(lookup, 2) + assert url == f"http://127.0.0.1:{port}{server.jobs_mcp.path}" + assert server.jobs_mcp._task is None + finally: + lookup.cancel() + await asyncio.gather(lookup, return_exceptions=True) + if listener is not None: + listener.close() + await listener.wait_closed() + sock.close() + + +@pytest.mark.asyncio +async def test_replaced_runtime_uses_loopback_when_its_gateway_is_not_mounted(tmp_path, monkeypatch): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http", "web_enabled": False}, + components={"codex_app_server": {"default": {"backend": "sdk"}}}, + ) + app.context.service.build_service(app) + replacement = CodexAppServer(app_context=app.context) + app.context.service.service.mount("/assets", app=ScopedMCP()) + listen = AsyncMock(return_value="http://127.0.0.1:1234/local") + monkeypatch.setattr(replacement.jobs_mcp, "listen", listen) + assert await replacement._mcp_url() == "http://127.0.0.1:1234/local" + listen.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_cleanup_preserves_the_original_request_error(system, monkeypatch): + first, _, server, _ = system + await first.reply("initial") + client = _Codex.instances[0] + + async def fail(_thread, _inputs): + raise ValueError("original failure") + + client.run_hook = fail + monkeypatch.setattr(server, "wait_idle", AsyncMock(side_effect=RuntimeError("cleanup failure"))) + with pytest.raises(ValueError, match="original failure"): + await first.reply("failed") + + +@pytest.mark.asyncio +async def test_cancelled_waiter_keeps_the_running_threads_tool_binding(system): + first, second, server, _ = system + binding = {"job_tools": ["search"], "tool_context_id": "one"} + thread_id = (await first.reply("initial", **binding))["session_id"] + token = server._threads[thread_id].token + entered, release = asyncio.Event(), asyncio.Event() + + async def hold(_thread, _inputs): + entered.set() + await release.wait() + + client = _Codex.instances[0] + client.run_hook = hold + async with asyncio.TaskGroup() as group: + group.create_task(first.reply("held", session_id=thread_id, **binding)) + await asyncio.wait_for(entered.wait(), 2) + waiting = group.create_task(second.reply("queued", session_id=thread_id, **binding)) + try: + await asyncio.sleep(0) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + assert token in server.jobs_mcp.scopes + finally: + release.set() + client.run_hook = None + await second.reply("still authorized", session_id=thread_id, **binding) + assert "queued" not in client.threads[thread_id].inputs + + +@pytest.mark.asyncio +async def test_disconnected_turn_is_not_replayed_and_next_request_reconnects(system): + from openai_codex import TransportClosedError + + first, _, server, _ = system + await first.reply("initial") + old_client = _Codex.instances[0] + + async def disconnect(_thread, _inputs): + raise TransportClosedError("disconnected") + + old_client.run_hook = disconnect + with pytest.raises(TransportClosedError): + await first.reply("do not replay") + assert server._failed + assert old_client.threads["thread-2"].inputs == ["do not replay"] + assert (await first.reply("new request"))["result"] == "new request" + assert len(_Codex.instances) == 2 + assert old_client.closed == 1 + assert _Codex.instances[1].threads["thread-1"].inputs == ["new request"] + + +@pytest.mark.asyncio +async def test_initial_transport_failure_does_not_poison_client_retry(system, monkeypatch): + from openai_codex import TransportClosedError + + first, _, server, _ = system + monkeypatch.setattr( + _Codex, + "account", + AsyncMock( + side_effect=[ + TransportClosedError("initialize failed"), + SimpleNamespace(account=object()), + ], + ), + ) + with pytest.raises(TransportClosedError): + await first.reply("failed initialization") + assert not server._failed + assert (await first.reply("retry"))["result"] == "retry" + assert _Codex.instances[0].closed == 1 + + +@pytest.mark.asyncio +async def test_thread_pruning_waits_for_complete_pagination_and_preserves_tokens_on_error(system, monkeypatch): + first, _, server, _ = system + one = (await first.reply("one", job_tools=["search"]))["session_id"] + two = (await first.reply("two", job_tools=["search"]))["session_id"] + token = server._threads[two].token + client = _Codex.instances[0] + pages = [SimpleNamespace(data=[one], next_cursor="next"), RuntimeError("incomplete listing")] + request = AsyncMock(side_effect=pages) + monkeypatch.setattr(client._client, "request", request) + await server._prune_threads(client) + assert token in server.jobs_mcp.scopes + request.side_effect = [ + SimpleNamespace(data=[one], next_cursor="next"), + SimpleNamespace(data=[], next_cursor=None), + ] + await server._prune_threads(client) + assert one in server._threads + assert two not in server._threads + assert token not in server.jobs_mcp.scopes + + +@pytest.mark.asyncio +async def test_stale_unload_snapshot_does_not_revoke_a_resumed_threads_tools(system, monkeypatch): + first, second, server, _ = system + binding = {"job_tools": ["search"], "tool_context_id": "one"} + thread_id = (await first.reply("initial", **binding))["session_id"] + state = server._threads[thread_id] + token = state.token + client = _Codex.instances[0] + original_request = client.request + resume_snapshot, stale_snapshot = asyncio.Event(), asyncio.Event() + resume_query, stale_query = asyncio.Event(), asyncio.Event() + queries = 0 + + async def request(method, params, **kwargs): + nonlocal queries + if method == "thread/loaded/list": + queries += 1 + if queries == 1: + resume_snapshot.set() + await resume_query.wait() + return SimpleNamespace(data=[thread_id], next_cursor=None) + if queries == 2: + stale_snapshot.set() + await stale_query.wait() + return SimpleNamespace(data=[], next_cursor=None) + return await original_request(method, params, **kwargs) + + monkeypatch.setattr(client._client, "request", request) + async with asyncio.TaskGroup() as group: + resumed = group.create_task(second.reply("resume", session_id=thread_id, **binding)) + await asyncio.wait_for(resume_snapshot.wait(), 2) + client.loaded.clear() + pruning = group.create_task(server._prune_threads(client)) + await asyncio.wait_for(stale_snapshot.wait(), 2) + try: + resume_query.set() + await asyncio.wait_for(asyncio.shield(resumed), 2) + finally: + resume_query.set() + stale_query.set() + await pruning + assert server._threads.get(thread_id) is state + assert token in server.jobs_mcp.scopes + + +@pytest.mark.asyncio +async def test_tool_arguments_are_snapshotted_before_the_caller_mutates_them(system): + first, _, server, job = system + allowed = ["one.md"] + result = await first.reply("one", job_tools=["search"], injected_job_kwargs={"_allowed_paths": allowed}) + allowed[:] = ["other.md"] + token = server._threads[result["session_id"]].token + await server.jobs_mcp.scopes[token][0].fn(query="read") + assert job.calls[0]["_allowed_paths"] == ["one.md"] def test_event_to_chunks_maps_content_usage_and_completion(): @@ -397,89 +1109,6 @@ def test_event_to_chunks_preserves_new_turn_scoped_notifications(): assert chunk.metadata == {"codex_method": "turn/diff/updated"} -@dataclass -class _TurnResult: - final_response: str - status: str = "completed" - - -@pytest.mark.asyncio -async def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - - class FakeThread: - id = "thread-1" - - async def run(self, inputs, **kwargs): - assert inputs == "answer" - assert kwargs["output_schema"] == {"type": "object"} - return _TurnResult(final_response=json.dumps({"ok": True})) - - class FakeCodex: - def __init__(self, _config): - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, _exc_type, _exc, _tb): - return None - - async def close(self): - return None - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - result = await wrapper.reply("answer", output_schema={"type": "object"}) - await wrapper.close() - - assert result["session_id"] == "thread-1" - assert result["structured_output"] == {"ok": True} - - -@pytest.mark.asyncio -async def test_reply_accepts_latest_sdk_run_input(tmp_path, monkeypatch): - from openai_codex import LocalImageInput, TextInput - - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - inputs = [TextInput("describe this image"), LocalImageInput("image.png")] - observed = {} - - class FakeThread: - id = "thread-1" - - async def run(self, run_input, **_kwargs): - observed["input"] = run_input - return _TurnResult(final_response="done") - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - result = await wrapper.reply(inputs) - await wrapper.close() - - assert observed["input"] is inputs - assert result["last_message"] == "done" - - def test_codex_skills_add_all_without_deleting_existing_content(tmp_path): wrapper, _job = _wrapper(tmp_path) for name in ("reme_memory", "qwenpaw_memory"): @@ -560,603 +1189,6 @@ def test_codex_skills_do_not_modify_codex_home(tmp_path): assert marker.read_text(encoding="utf-8") == "keep" -def test_effective_mcp_config_snapshot_is_private_and_removed_on_close(tmp_path): - wrapper, _job = _wrapper(tmp_path) - config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access - snapshot = Path(config["args"][config["args"].index("--config") + 1]) - assert snapshot.exists() - assert stat.S_IMODE(snapshot.stat().st_mode) == 0o600 - assert json.loads(snapshot.read_text(encoding="utf-8"))["workspace_dir"] == str(tmp_path) - - async def close_started_wrapper(): - await wrapper.start() - await wrapper.close() - - asyncio.run(close_started_wrapper()) - assert not snapshot.exists() - - -def test_effective_mcp_config_snapshot_needs_no_fchmod(tmp_path, monkeypatch): - """A platform without ``os.fchmod`` (Windows) still gets a usable snapshot.""" - monkeypatch.delattr(os, "fchmod", raising=False) - - wrapper, _job = _wrapper(tmp_path) - config = wrapper._mcp_server_config({"job_tools": ["search"]}) - snapshot = Path(config["args"][config["args"].index("--config") + 1]) - - assert snapshot.exists() - assert json.loads(snapshot.read_text(encoding="utf-8"))["workspace_dir"] == str(tmp_path) - - -@pytest.mark.asyncio -async def test_effective_snapshot_exposes_parent_only_custom_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - app_config = ApplicationConfig( - workspace_dir=str(tmp_path), - enable_logo=False, - log_to_console=False, - log_to_file=False, - service={"backend": "mcp"}, - jobs={ - "only_custom": { - "backend": "base", - "description": "Parent-only inline job", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - "referenced_helper": { - "backend": "base", - "description": "A normal job that custom jobs may reference.", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - ) - job = _Job("only_custom") - context = SimpleNamespace(app_config=app_config, components={}, jobs={"only_custom": job}) - wrapper = CodexAgentWrapper(app_context=context) - server_config = wrapper._mcp_server_config({"job_tools": ["only_custom"]}) # pylint: disable=protected-access - transport = StdioTransport( - command=server_config["command"], - args=server_config["args"], - cwd=server_config["cwd"], - keep_alive=False, - ) - - await wrapper.start() - try: - async with Client(transport, timeout=None, init_timeout=STDIO_INIT_TIMEOUT) as client: - tools = await asyncio.wait_for(client.list_tools(), timeout=10) - snapshot = Path(server_config["args"][server_config["args"].index("--config") + 1]) - assert snapshot.exists() - assert set(json.loads(snapshot.read_text(encoding="utf-8"))["jobs"]) == { - "only_custom", - "referenced_helper", - } - finally: - await wrapper.close() - - assert [tool.name for tool in tools] == ["only_custom"] - assert not snapshot.exists() - - -@pytest.mark.asyncio -async def test_open_thread_defaults_to_full_access(tmp_path): - from openai_codex import ApprovalMode, Sandbox - - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeCodex: - async def thread_start(self, **kwargs): - observed.update(kwargs) - return SimpleNamespace(id="thread-1") - - await wrapper._open_thread(FakeCodex(), {}) # pylint: disable=protected-access - - assert observed["approval_mode"] == ApprovalMode.auto_review - assert observed["sandbox"] == Sandbox.full_access - - -@pytest.mark.asyncio -async def test_compact_session_uses_native_thread_operation(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeThread: - id = "thread-1" - - async def compact(self): - observed["compacted"] = True - - async def start(): - observed["started"] = True - - async def resume(session_id): - observed["resume"] = session_id - return FakeThread() - - async def get_codex(): - return SimpleNamespace(thread_resume=resume) - - monkeypatch.setattr(wrapper, "start", start) - monkeypatch.setattr(wrapper, "_get_codex", get_codex) - - await wrapper.compact_session("thread-1") - - assert observed["started"] is True - assert observed["resume"] == "thread-1" - assert observed["compacted"] is True - - -@pytest.mark.asyncio -async def test_open_thread_forwards_latest_sdk_options(tmp_path): - from openai_codex.types import Personality, ThreadSource, ThreadStartSource - - wrapper, _job = _wrapper(tmp_path) - observed = {} - - class FakeCodex: - async def thread_start(self, **kwargs): - observed["start"] = kwargs - return SimpleNamespace(id="thread-1") - - async def thread_resume(self, _thread_id, **kwargs): - observed["resume"] = kwargs - return SimpleNamespace(id="thread-1") - - async def thread_fork(self, _thread_id, **kwargs): - observed["fork"] = kwargs - return SimpleNamespace(id="thread-2") - - codex = FakeCodex() - await wrapper._open_thread( # pylint: disable=protected-access - codex, - { - "personality": "friendly", - "service_name": "reme", - "session_start_source": "startup", - "thread_source": "user", - }, - ) - await wrapper._open_thread( # pylint: disable=protected-access - codex, - {"session_id": "thread-1", "personality": "pragmatic"}, - ) - await wrapper._open_thread( # pylint: disable=protected-access - codex, - {"session_id": "thread-1", "fork_session": True, "thread_source": "subagent"}, - ) - - assert observed["start"]["personality"] == Personality.friendly - assert observed["start"]["service_name"] == "reme" - assert observed["start"]["session_start_source"] == ThreadStartSource.startup - assert observed["start"]["thread_source"] == ThreadSource.user - assert observed["resume"]["personality"] == Personality.pragmatic - assert observed["fork"]["thread_source"] == ThreadSource.subagent - - -@pytest.mark.asyncio -async def test_resume_reuses_tool_context_and_rejects_context_change(tmp_path): - wrapper, _job = _wrapper(tmp_path) - - class FakeCodex: - async def thread_start(self, **_kwargs): - return SimpleNamespace(id="thread-1") - - async def thread_resume(self, _thread_id, **_kwargs): - return SimpleNamespace(id="thread-1") - - await wrapper._open_thread(FakeCodex(), {"tool_context_id": "ctx-a"}) # pylint: disable=protected-access - await wrapper._open_thread( # pylint: disable=protected-access - FakeCodex(), - {"resume": "thread-1", "tool_context_id": "ctx-a"}, - ) - with pytest.raises(ValueError, match="cannot change"): - await wrapper._open_thread( # pylint: disable=protected-access - FakeCodex(), - {"resume": "thread-1", "tool_context_id": "ctx-b"}, - ) - - -class _StructuredModel(BaseModel): - ok: bool - - -@pytest.mark.parametrize("schema", [_StructuredModel(ok=True), str]) -def test_output_schema_rejects_instances_and_arbitrary_classes(tmp_path, schema): - wrapper, _job = _wrapper(tmp_path) - with pytest.raises(TypeError, match="JSON schema dict or BaseModel class"): - wrapper._merged_kwargs({"output_schema": schema}) # pylint: disable=protected-access - - -def test_output_schema_normalizes_model_class_and_preserves_dict(tmp_path): - wrapper, _job = _wrapper(tmp_path) - schema = {"type": "object", "properties": {"ok": {"type": "boolean"}}} - - assert ( - wrapper._merged_kwargs({"output_schema": _StructuredModel})["output_schema"] # pylint: disable=protected-access - == _StructuredModel.model_json_schema() - ) - assert ( - wrapper._merged_kwargs({"output_schema": schema})["output_schema"] is schema - ) # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - clients = [] - observed_schemas = [] - close_count = 0 - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **kwargs): - observed_schemas.append(kwargs["output_schema"]) - return _TurnResult(final_response=json.dumps({"ok": True})) - - class FakeCodex: - def __init__(self, _config): - clients.append(self) - - async def __aenter__(self): - return self - - async def close(self): - nonlocal close_count - close_count += 1 - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - result = await wrapper.reply("first", output_schema=_StructuredModel) - await wrapper.close() - - assert result["structured_output"] == {"ok": True} - assert observed_schemas == [_StructuredModel.model_json_schema()] - assert len(clients) == 1 - assert close_count == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("schema", [{}, _StructuredModel]) -async def test_reply_stream_rejects_output_schema(tmp_path, schema): - wrapper, _job = _wrapper(tmp_path) - - with pytest.raises(NotImplementedError, match="Structured output is not supported"): - await anext(wrapper.reply_stream("answer", output_schema=schema)) - - -@pytest.mark.asyncio -async def test_reply_stream_interrupts_turn_when_consumer_closes_early(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - stream_closed = False - interrupt_count = 0 - - class FakeTurn: - id = "turn-1" - - async def stream(self): - nonlocal stream_closed - try: - yield SimpleNamespace( - method="turn/started", - payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)), - ) - await asyncio.Event().wait() - finally: - stream_closed = True - - async def interrupt(self): - nonlocal interrupt_count - interrupt_count += 1 - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - stream = wrapper.reply_stream("answer") - first = await anext(stream) - assert first.chunk_type == ChunkEnum.REPLY_START - await stream.aclose() - await wrapper.close() - - assert interrupt_count == 1 - assert stream_closed - - -@pytest.mark.asyncio -async def test_reply_stream_does_not_record_token_usage(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - recorded_usages = [] - usage = TokenUsageBreakdown( - cachedInputTokens=0, - inputTokens=3, - outputTokens=5, - reasoningOutputTokens=0, - totalTokens=8, - ) - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="thread/tokenUsage/updated", - payload=SimpleNamespace(token_usage=SimpleNamespace(last=usage)), - ) - yield SimpleNamespace( - method="turn/completed", - payload=SimpleNamespace( - turn=SimpleNamespace( - id=self.id, - status=SimpleNamespace(value="completed"), - duration_ms=1, - error=None, - ), - ), - ) - - async def interrupt(self): - raise AssertionError("Completed turns must not be interrupted") - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - monkeypatch.setattr(wrapper, "_record_token_usage", recorded_usages.append) - - chunks = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - assert any(chunk.chunk_type == ChunkEnum.USAGE for chunk in chunks) - assert not recorded_usages - - -@pytest.mark.asyncio -async def test_close_waits_for_active_turn(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - turn_started = asyncio.Event() - release_turn = asyncio.Event() - client_closed = asyncio.Event() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - turn_started.set() - await release_turn.wait() - return _TurnResult(final_response="done") - - class FakeCodex: - def __init__(self, _config): - pass - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - client_closed.set() - - async def thread_start(self, **_kwargs): - return FakeThread() - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - reply_task = asyncio.create_task(wrapper.reply("answer")) - await turn_started.wait() - close_task = asyncio.create_task(wrapper.close()) - await asyncio.sleep(0) - - assert not client_closed.is_set() - assert not close_task.done() - - release_turn.set() - result = await reply_task - await close_task - - assert result["last_message"] == "done" - assert client_closed.is_set() - - -@pytest.mark.asyncio -async def test_component_start_keeps_optional_client_lazy(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="") - - def fail_if_constructed(_config): - raise AssertionError("Codex client should be lazy") - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", fail_if_constructed) - - await wrapper.start() - await wrapper.close() - - assert wrapper._codex is None - - -@pytest.mark.asyncio -async def test_client_config_is_fixed_for_component_lifetime(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="api_key", api_key="one") - clients = [] - - class FakeCodex: - def __init__(self, _config): - clients.append(self) - - async def login_api_key(self, _api_key): - return None - - async def close(self): - return None - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - assert await wrapper._get_codex() is clients[0] # pylint: disable=protected-access - with pytest.raises(TypeError, match="configured on the wrapper: api_key"): - await wrapper.reply("answer", api_key="two") - await wrapper.close() - - assert len(clients) == 1 - - -@pytest.mark.parametrize( - ("name", "value"), - [ - ("auth_mode", "oauth"), - ("base_url", "https://example.test/v1"), - ("codex_bin", "/tmp/codex"), - ("codex_home", "/tmp/codex-home"), - ("config_overrides", ['model="test"']), - ("cwd", "/tmp"), - ("experimental_api", False), - ("launch_args_override", ["codex", "app-server"]), - ], -) -@pytest.mark.asyncio -async def test_reply_rejects_call_time_client_options(tmp_path, name, value): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - - with pytest.raises(TypeError, match=f"configured on the wrapper: {name}"): - await wrapper.reply("answer", **{name: value}) - - assert wrapper._codex is None - - -def test_constructor_rejects_launch_args_override(tmp_path): - with pytest.raises(TypeError, match="configure codex_bin instead"): - _wrapper(tmp_path, launch_args_override=["codex", "app-server"]) - - -def test_oauth_mode_ignores_api_credentials_and_forces_chatgpt(tmp_path, monkeypatch): - wrapper, _job = _wrapper( - tmp_path, - auth_mode="oauth", - api_key="explicit-key", - base_url="https://explicit.example.test/v1", - ) - wrapper.app_context.app_config.environment = {"TOOL_ENV": "preserved"} - monkeypatch.setenv("CODEX_API_KEY", "ambient-key") - monkeypatch.setenv("CODEX_BASE_URL", "https://ambient.example.test/v1") - - auth = wrapper._resolve_auth_config( # pylint: disable=protected-access - "oauth", - "explicit-key", - "https://explicit.example.test/v1", - ) - config = wrapper._build_client_config(auth) # pylint: disable=protected-access - - assert auth.mode == "oauth" - assert auth.api_key == "" - assert auth.base_url == "" - assert config.env["TOOL_ENV"] == "preserved" - assert "CODEX_HOME" not in wrapper.app_context.app_config.environment - assert 'forced_login_method="chatgpt"' in config.config_overrides - assert not any(value.startswith("openai_base_url=") for value in config.config_overrides) - - -@pytest.mark.asyncio -async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch): - wrapper, _job = _wrapper( - tmp_path, - auth_mode="api_key", - api_key="explicit-key", - base_url="https://proxy.example.test/v1", - ) - observed = {} - - class FakeCodex: - def __init__(self, config): - observed["config"] = config - - async def login_api_key(self, api_key): - observed["api_key"] = api_key - - async def close(self): - observed["closed"] = True - - monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) - - await wrapper.start() - await wrapper._get_codex() # pylint: disable=protected-access - await wrapper.close() - - config = observed["config"] - assert observed["api_key"] == "explicit-key" - assert 'openai_base_url="https://proxy.example.test/v1"' in config.config_overrides - assert 'forced_login_method="api"' in config.config_overrides - assert observed["closed"] is True - - -def test_auth_selection_uses_explicit_wrapper_options(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - wrapper.app_context.app_config.components[ComponentEnum.AS_LLM]["default"] = SimpleNamespace( - credential={"api_key": "default-key", "base_url": "https://default.example.test/v1"}, - ) - for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY"): - monkeypatch.setenv(name, "ambient-key") - for name in ("CODEX_BASE_URL", "OPENAI_BASE_URL", "LLM_BASE_URL"): - monkeypatch.setenv(name, "https://ambient.example.test/v1") - - auth = wrapper._resolve_auth_config( # pylint: disable=protected-access - "api_key", - "configured-key", - ) - - assert auth.mode == "api_key" - assert auth.api_key == "configured-key" - assert auth.base_url == "" - - with pytest.raises(ValueError, match="requires a non-empty API key"): - wrapper._resolve_auth_config( # pylint: disable=protected-access - "api_key", - ) - - @pytest.mark.parametrize("review_status", ["approved", "denied"]) def test_event_to_chunks_maps_approval_started_and_completed(review_status): action = {"type": "futureApprovalAction", "value": "preserved"} @@ -1192,464 +1224,3 @@ def test_event_to_chunks_maps_approval_started_and_completed(review_status): assert completed_chunk.metadata["status"] == "completed" assert completed_chunk.metadata["review"]["status"] == review_status assert completed_chunk.metadata["decision_source"] == {"type": "guardian"} - - -def test_codex_home_expands_user_directory(tmp_path): - wrapper, _job = _wrapper(tmp_path, codex_home="~/.codex") - assert wrapper.session_path == Path.home() / ".codex" - - -def test_named_default_mcp_config_remains_supported(tmp_path): - wrapper, _job = _wrapper(tmp_path, mcp_config="default") - assert wrapper._mcp_config_source({}) == "default" # pylint: disable=protected-access - - -def test_default_config_provides_codex_oauth_wrapper(monkeypatch): - monkeypatch.delenv("CODEX_HOME", raising=False) - config = resolve_app_config(log_config=False) - oauth = config["components"]["agent_wrapper"]["codex_oauth"] - codex = config["components"]["agent_wrapper"]["codex"] - assert oauth["backend"] == "codex" - assert oauth["auth_mode"] == "oauth" - assert oauth["codex_home"] == "~/.codex" - assert oauth["sandbox"] == "full-access" - assert "api_key" not in oauth - assert codex["auth_mode"] == "api_key" - assert codex["sandbox"] == "full-access" - - -class _RecordingClient: - def __init__(self, fail: bool = False): - self.requests = [] - self.fail = fail - self.loaded_threads = ["thread-1"] - - async def request(self, method, params, *, response_model): - self.requests.append({"method": method, "params": params, "response_model": response_model}) - if self.fail: - raise RuntimeError("unsubscribe failed") - if method == "thread/loaded/list": - return SimpleNamespace(data=self.loaded_threads, next_cursor=None) - return SimpleNamespace(status="unsubscribed") - - -def _assert_unsubscribed(client, thread_id="thread-1"): - from openai_codex.generated.v2_all import ThreadUnsubscribeResponse - - assert client.requests == [ - { - "method": "thread/unsubscribe", - "params": {"threadId": thread_id}, - "response_model": ThreadUnsubscribeResponse, - }, - ] - - -def _fake_codex(client, thread): - class FakeCodex: - def __init__(self, _config): - self._client = client - - async def account(self): - return SimpleNamespace(account=SimpleNamespace()) - - async def close(self): - return None - - async def thread_start(self, **_kwargs): - return thread - - async def thread_resume(self, _thread_id, **_kwargs): - return thread - - return FakeCodex - - -@pytest.mark.asyncio -async def test_reply_releases_thread_after_turn(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - await wrapper.close() - - assert result["last_message"] == "done" - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_releases_thread_when_turn_fails(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - raise ValueError("turn failed") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(ValueError, match="turn failed"): - await wrapper.reply("answer") - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_ignores_unsubscribe_failure(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient(fail=True) - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - - assert result["last_message"] == "done" - assert len(client.requests) == 1 - # The thread is still loaded, so its tool-context guard must survive. - assert "thread-1" in wrapper._thread_tool_contexts # pylint: disable=protected-access - - await wrapper.close() - - -@pytest.mark.asyncio -async def test_failed_release_keeps_tool_context_guard_on_resume(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient(fail=True) - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer", tool_context_id="ctx-a") - assert result["session_id"] == "thread-1" - assert len(client.requests) == 1 - - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - await wrapper.close() - - -@pytest.mark.asyncio -async def test_release_threads_disabled_keeps_threads_loaded(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth", release_threads=False) - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - result = await wrapper.reply("answer") - - assert result["last_message"] == "done" - assert not client.requests - assert "thread-1" in wrapper._thread_tool_contexts # pylint: disable=protected-access - - await wrapper.close() - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_on_completion(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="turn/completed", - payload=SimpleNamespace( - turn=SimpleNamespace( - id="turn-1", - status=SimpleNamespace(value="completed"), - duration_ms=1, - error=None, - ), - ), - ) - - async def interrupt(self): - raise AssertionError("Completed turns must not be interrupted") - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - chunks = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - assert chunks[-1].chunk_type == ChunkEnum.REPLY_END - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_on_early_close(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - interrupt_count = 0 - - class FakeTurn: - id = "turn-1" - - async def stream(self): - yield SimpleNamespace( - method="turn/started", - payload=SimpleNamespace(turn=SimpleNamespace(id=self.id)), - ) - await asyncio.Event().wait() - - async def interrupt(self): - nonlocal interrupt_count - interrupt_count += 1 - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - return FakeTurn() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - stream = wrapper.reply_stream("answer") - first = await anext(stream) - assert first.chunk_type == ChunkEnum.REPLY_START - await stream.aclose() - await wrapper.close() - - assert interrupt_count == 1 - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_compact_session_releases_thread(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path) - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def compact(self): - return None - - async def start(): - return None - - async def resume(_thread_id): - return FakeThread() - - async def get_codex(): - return SimpleNamespace(_client=client, thread_resume=resume) - - monkeypatch.setattr(wrapper, "start", start) - monkeypatch.setattr(wrapper, "_get_codex", get_codex) - - await wrapper.compact_session("thread-1") - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_when_turn_start_fails(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - raise RuntimeError("turn/start failed") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(RuntimeError, match="turn/start failed"): - _ = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -async def test_reply_stream_releases_thread_when_turn_start_is_cancelled(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def turn(self, _inputs, **_kwargs): - raise asyncio.CancelledError("cancelled during turn/start") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - - with pytest.raises(asyncio.CancelledError): - _ = [chunk async for chunk in wrapper.reply_stream("answer")] - await wrapper.close() - - _assert_unsubscribed(client) - assert not wrapper._thread_tool_contexts # pylint: disable=protected-access - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("original", "changed"), - [ - ({"tool_context_id": "ctx-a"}, {"tool_context_id": "ctx-b"}), - ({"injected_job_kwargs": {"tenant": "a"}}, {"injected_job_kwargs": {"tenant": "b"}}), - ], -) -async def test_cancelled_reply_keeps_loaded_thread_constraints(tmp_path, monkeypatch, original, changed): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - entered = asyncio.Event() - - class FakeThread: - id = "thread-1" - calls = 0 - - async def run(self, _inputs, **_kwargs): - self.calls += 1 - if self.calls > 1: - return _TurnResult(final_response="old MCP constraints") - entered.set() - await asyncio.Event().wait() - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - task = asyncio.create_task(wrapper.reply("answer", **original)) - try: - await asyncio.wait_for(entered.wait(), timeout=5) - finally: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - try: - _assert_unsubscribed(client) - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", **changed) - finally: - await wrapper.close() - - -@pytest.mark.asyncio -async def test_unloaded_thread_can_resume_with_new_constraints(tmp_path, monkeypatch): - wrapper, _job = _wrapper(tmp_path, auth_mode="oauth") - client = _RecordingClient() - - class FakeThread: - id = "thread-1" - - async def run(self, _inputs, **_kwargs): - return _TurnResult(final_response="done") - - monkeypatch.setattr( - "reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", - _fake_codex(client, FakeThread()), - ) - try: - await wrapper.reply("answer", tool_context_id="ctx-a") - assert "thread-1" in wrapper._thread_tool_contexts - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-a") - with pytest.raises(ValueError, match="cannot change"): - await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - # Unsubscribe schedules delayed unload; only a later loaded-list query - # establishes that a resume will create a fresh runtime and MCP config. - client.loaded_threads.clear() - result = await wrapper.reply("answer", session_id="thread-1", tool_context_id="ctx-b") - assert result["last_message"] == "done" - assert json.loads(wrapper._thread_tool_contexts["thread-1"])["tool_context_id"] == "ctx-b" - finally: - await wrapper.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("fail_second_page", [False, True]) -async def test_prune_thread_constraints_reads_all_pages_before_removing_entries(tmp_path, fail_second_page): - from openai_codex.generated.v2_all import ThreadLoadedListResponse - - wrapper, _job = _wrapper(tmp_path) - wrapper._thread_tool_contexts = {"first": "a", "last": "b", "unloaded": "c"} - cursors = [] - - class Client: - async def request(self, method, params, *, response_model): - assert method == "thread/loaded/list" - assert response_model is ThreadLoadedListResponse - cursors.append(params["cursor"]) - if params["cursor"] is None: - return ThreadLoadedListResponse(data=["first"], nextCursor="page-2") - if fail_second_page: - raise RuntimeError("loaded-list unavailable") - return ThreadLoadedListResponse(data=["last"]) - - await wrapper._prune_thread_tool_contexts(SimpleNamespace(_client=Client())) - - assert cursors == [None, "page-2"] - expected = {"first": "a", "last": "b"} - if fail_second_page: - expected["unloaded"] = "c" - assert wrapper._thread_tool_contexts == expected diff --git a/tests/unit/test_common_steps.py b/tests/unit/test_common_steps.py index 02c8a870..9738a36f 100644 --- a/tests/unit/test_common_steps.py +++ b/tests/unit/test_common_steps.py @@ -3,6 +3,7 @@ # pylint: disable=protected-access from contextlib import chdir as _temp_chdir +import asyncio import tempfile import warnings @@ -11,11 +12,13 @@ import pytest from reme.components.agent_wrapper import BaseAgentWrapper from reme.components.application_context import ApplicationContext from reme.components.file_store import LocalFileStore -from reme.schema import FileFrontMatter, FileLink, FileNode, TraverseGraph +from reme.enumeration import ChunkEnum +from reme.schema import FileFrontMatter, FileLink, FileNode, StreamChunk, TraverseGraph from reme.steps.common.add import AddStep from reme.steps.common.health_check import _file_graph_status from reme.steps.common.llm_demo import LLMDemoStep from reme.steps.common.python_execute import PythonExecuteStep +from reme.steps.common.stream_llm_demo import StreamLLMDemoStep from reme.steps.index import traverse as traverse_mod from reme.utils.local_backend import LocalProcessBackend @@ -147,11 +150,17 @@ class _FakeAgentWrapper(BaseAgentWrapper): def __init__(self): super().__init__() self.last_kwargs = None + self.stream_chunks = [] async def reply(self, inputs, **kwargs) -> dict: self.last_kwargs = kwargs return {"result": "ok"} + async def reply_stream(self, inputs, **kwargs): + self.last_kwargs = kwargs + for chunk in self.stream_chunks: + yield chunk + @pytest.mark.asyncio async def test_llm_demo_always_registers_add_tool(): @@ -165,6 +174,27 @@ async def test_llm_demo_always_registers_add_tool(): assert "job_tools" not in resp.metadata +@pytest.mark.asyncio +async def test_stream_llm_demo_preserves_backend_chunk_metadata(): + """Forward retry flags, tool identities, usage, and completion status without rebuilding chunks.""" + wrapper = _FakeAgentWrapper() + wrapper.stream_chunks = [ + StreamChunk(chunk_type=ChunkEnum.ERROR, chunk="retrying", metadata={"will_retry": True}), + StreamChunk(chunk_type=ChunkEnum.TOOL_CALL, tool_call_id="call-1", tool_call_name="add", chunk={"a": 1}), + StreamChunk(chunk_type=ChunkEnum.CONTENT, session_id="session-1", block_id="text-1", chunk="hello"), + StreamChunk(chunk_type=ChunkEnum.USAGE, input_tokens=3, output_tokens=1), + StreamChunk(chunk_type=ChunkEnum.REPLY_END, metadata={"status": "completed"}), + ] + queue = asyncio.Queue() + response = await StreamLLMDemoStep()(query="hello", stream_queue=queue, agent_wrapper=wrapper) + forwarded = [queue.get_nowait() for _ in wrapper.stream_chunks] + assert all(actual is expected for actual, expected in zip(forwarded, wrapper.stream_chunks)) + assert queue.empty() + assert response.success + assert response.answer == "hello" + assert response.metadata["session_id"] == "session-1" + + def test_file_graph_health_reports_neo4j_cached_counts(): """Neo4j file graph health should not be reported as an empty local graph.""" diff --git a/tests/unit/test_http_client.py b/tests/unit/test_http_client.py index bd506cb0..7b477f46 100644 --- a/tests/unit/test_http_client.py +++ b/tests/unit/test_http_client.py @@ -2,8 +2,14 @@ # pylint: disable=protected-access +import json + +import httpx +import pytest + from reme.components.client.http_client import HttpClient from reme.constants import REME_DEFAULT_HOST, REME_DEFAULT_PORT +from reme.enumeration import ChunkEnum def test_default_client_uses_loopback_address(): @@ -34,3 +40,30 @@ def test_format_for_display_shows_metadata_when_requested(): text = '{"answer":"0.4.0.7","success":true,"metadata":{"version":"0.4.0.7"}}' assert client._format_for_display(text) == '0.4.0.7\n✅ {"version": "0.4.0.7"}' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("will_retry", [True, False, None, "false"]) +async def test_stream_retry_notifications_continue_but_terminal_errors_raise(will_retry): + """A temporary Codex reconnect must not prevent CLI callers from receiving the answer.""" + events = [ + {"chunk_type": "error", "chunk": "upstream error", "metadata": {"will_retry": will_retry}}, + {"chunk_type": "content", "chunk": "recovered answer"}, + {"chunk_type": "done", "done": True}, + ] + body = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + + def handle(_request): + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + client = HttpClient() + client.client = httpx.AsyncClient(base_url=client.base_url, transport=httpx.MockTransport(handle)) + async with client: + if will_retry is True: + chunks = [chunk async for chunk in client.stream_chunks("stream_llm_demo")] + assert [chunk.chunk_type for chunk in chunks] == [ChunkEnum.ERROR, ChunkEnum.CONTENT] + assert chunks[0].metadata == {"will_retry": True} + assert chunks[-1].chunk == "recovered answer" + else: + with pytest.raises(RuntimeError, match="upstream error"): + _ = [chunk async for chunk in client.stream_chunks("stream_llm_demo")]