From 6365280250ce72e5c3b633482b6a0d8321f03792 Mon Sep 17 00:00:00 2001 From: celanwang Date: Tue, 20 Oct 2026 16:25:19 +0800 Subject: [PATCH 01/21] feat(config): support named multi-workspace workspace_dir Phase 1 of the multi-workspace plan: configuration union, validation, and a per-workspace path layer, with zero behavior change for existing single-workspace setups. workspace_dir now accepts an ordered alias-to-path mapping in addition to a plain string. Both forms normalize to an ordered workspaces mapping whose first entry is the default; the string form becomes {"default": path}. Aliases are taken verbatim and must be non-empty, and resolved directories must be distinct and must not overlap (no parent/child pairs). workspace_dir keeps its existing meaning as the default workspace's absolute path, so every current reader is unaffected, and the derived mapping is excluded from serialization. ApplicationContext builds one frozen WorkspaceContext per alias holding the workspace root plus the six configured storage directories, and Application creates that layout for every configured workspace at startup. Request routing and component instantiation are unchanged and remain single-workspace for later phases. Refs #632 --- docs/en/configuration.md | 14 ++ docs/zh/configuration.md | 12 ++ reme/application.py | 20 +-- reme/components/application_context.py | 9 +- reme/schema/__init__.py | 2 + reme/schema/application_config.py | 51 ++++++- reme/schema/workspace.py | 61 +++++++++ tests/unit/test_multi_workspace_config.py | 159 ++++++++++++++++++++++ 8 files changed, 310 insertions(+), 18 deletions(-) create mode 100644 reme/schema/workspace.py create mode 100644 tests/unit/test_multi_workspace_config.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index fedf7ac0..7b550ee5 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -78,6 +78,20 @@ Keep secrets in `.env` or the process environment, never in committed configurat `session_dir` must remain workspace-relative. +`workspace_dir` also accepts an ordered mapping of alias to path for hosts that serve several workspaces from one +configuration, for example in a YAML config file: + +```yaml +workspace_dir: + main: /data/reme/main + scratch: /data/reme/scratch +``` + +The first entry is the default workspace; every alias gets the same subdirectory layout (`metadata_dir`, +`session_dir`, and the others above) materialized under its own root at startup. Aliases must be non-empty, and +resolved directories must be distinct and must not contain one another. A single string keeps the existing +behavior and is normalized internally to a `default` alias. + ## Enabling jobs `jobs..enabled` defaults to `true` for all Job types. Disabled jobs retain their configuration but do not diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 6cec4614..2609275b 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -83,6 +83,18 @@ base_url: ${LLM_BASE_URL:-https://example.com/v1} `session_dir` 必须保持 workspace-relative。其他 workspace 子目录也应使用清晰、稳定的相对名称。 +`workspace_dir` 也接受"别名到路径"的有序映射,便于一个配置服务多个 workspace,例如在 YAML 配置文件中: + +```yaml +workspace_dir: + main: /data/reme/main + scratch: /data/reme/scratch +``` + +第一个 entry 是默认 workspace;每个别名在启动时都会在自己的根目录下创建相同的子目录布局 +(`metadata_dir`、`session_dir` 等上述目录)。别名必须非空,且解析后的目录必须互不相同、不能互相包含。 +单个字符串保持原有行为,内部归一化为 `default` 别名。 + ## Job 启停 `jobs..enabled` 默认为 `true`,适用于所有 Job 类型。禁用的 Job 保留配置,但不启动、不暴露接口, diff --git a/reme/application.py b/reme/application.py index 5bbe3d14..405ac022 100644 --- a/reme/application.py +++ b/reme/application.py @@ -5,7 +5,6 @@ import heapq from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from inspect import getattr_static -from pathlib import Path from typing import Any, AsyncGenerator, TypeVar from . import __version__ @@ -54,20 +53,11 @@ class Application(BaseComponent): # ----- Wiring (called once during __init__) -------------------------- def _setup_workspace_directories(self) -> None: - """Ensure the workspace root and configured subdirectories exist on disk.""" - cfg = self.config - workspace_path = Path(cfg.workspace_dir).absolute() - workspace_path.mkdir(parents=True, exist_ok=True) - for subdir in [ - cfg.metadata_dir, - cfg.session_dir, - cfg.mem_session_dir, - cfg.resource_dir, - cfg.daily_dir, - cfg.digest_dir, - ]: - if subdir: - (workspace_path / subdir).mkdir(parents=True, exist_ok=True) + """Ensure every configured workspace root and its configured subdirectories exist on disk.""" + for workspace in self.context.workspaces.values(): + workspace.workspace_path.mkdir(parents=True, exist_ok=True) + for subdir in workspace.storage_dirs(): + subdir.mkdir(parents=True, exist_ok=True) def _init_service(self) -> None: """Instantiate the single service backend declared in config.service.""" diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 0e0c393f..dca5dd67 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -3,7 +3,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any -from ..schema import ApplicationConfig +from ..schema import ApplicationConfig, WorkspaceContext from .component_registry import ComponentRegistry, create_application_registry if TYPE_CHECKING: @@ -25,6 +25,13 @@ class ApplicationContext: self.app_config: ApplicationConfig = ApplicationConfig(**kwargs) self.registry = registry or create_application_registry() + # Ordered per-workspace path views; the first entry is the default workspace. + self.workspaces: dict[str, WorkspaceContext] = { + alias: WorkspaceContext.from_config(alias, path, self.app_config) + for alias, path in self.app_config.workspaces.items() + } + self.workspace: WorkspaceContext = next(iter(self.workspaces.values())) + # Populated by Application during initialization. self.service: "BaseService | None" = None self.components: dict[str, dict[str, "BaseComponent"]] = {} diff --git a/reme/schema/__init__.py b/reme/schema/__init__.py index 187ee50b..d163ab6f 100644 --- a/reme/schema/__init__.py +++ b/reme/schema/__init__.py @@ -25,6 +25,7 @@ from .response import Response from .stream_chunk import StreamChunk from .token_usage import TokenUsage from .traverse_graph import TraverseGraph, TraverseGraphEdge, TraverseGraphNode +from .workspace import WorkspaceContext __all__ = [ "CallDecision", @@ -57,4 +58,5 @@ __all__ = [ "TraverseGraph", "TraverseGraphEdge", "TraverseGraphNode", + "WorkspaceContext", ] diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index dc12737e..1a19d4ef 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -1,13 +1,17 @@ """Application configuration models.""" +from collections.abc import Mapping +from itertools import combinations import os from pathlib import Path from pathlib import PurePosixPath, PureWindowsPath -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from ..enumeration import component_type_name +DEFAULT_WORKSPACE_ALIAS = "default" + class ComponentConfig(BaseModel): """Base config for a component; extra fields allowed for backend-specific options.""" @@ -37,9 +41,18 @@ class ApplicationConfig(BaseModel): ) workspace_dir: str = Field( default=".reme", - description="Workspace root directory for runtime files", + description=( + "Workspace root directory for runtime files. Accepts a single path or an ordered mapping of " + "alias to path; the mapping form normalizes to the first entry here and exposes all entries " + "through the derived workspaces mapping." + ), validate_default=True, ) + workspaces: dict[str, str] = Field( + default_factory=dict, + exclude=True, + description="Ordered alias-to-absolute-path workspaces derived from workspace_dir", + ) metadata_dir: str = Field(default="metadata", description="Subdirectory for ReMe persistent state") session_dir: str = Field(default="session", description="Subdirectory for persisted agent sessions") # dialog_dir was removed; standard transcripts are always derived as ``{session_dir}/dialog``. @@ -62,6 +75,40 @@ class ApplicationConfig(BaseModel): description="Component registry keyed by type then name", ) + @model_validator(mode="before") + @classmethod + def normalize_workspace_mapping(cls, data): + """Resolve a mapping workspace_dir, keeping its first entry as the default workspace.""" + if not isinstance(data, dict): + return data + raw = data.get("workspace_dir") + if not isinstance(raw, Mapping): + return data + if not raw: + raise ValueError("workspace_dir mapping must contain at least one workspace") + resolved: dict[str, str] = {} + for alias, path in raw.items(): + alias = str(alias).strip() + if not alias: + raise ValueError("workspace_dir aliases must be non-empty") + resolved[alias] = str(Path(str(path)).expanduser().resolve(strict=False)) + directories = list(resolved.values()) + if len(set(directories)) != len(directories): + raise ValueError(f"workspace_dir aliases must point to distinct directories: {resolved}") + for first, second in combinations(map(Path, directories), 2): + if first.is_relative_to(second) or second.is_relative_to(first): + raise ValueError( + f"workspace_dir directories must not overlap: {str(first)!r} and {str(second)!r}", + ) + return {**data, "workspace_dir": directories[0], "workspaces": resolved} + + @model_validator(mode="after") + def derive_default_workspace(self): + """Normalize every workspace_dir form to an ordered mapping; the first entry is the default.""" + if not self.workspaces: + self.workspaces = {DEFAULT_WORKSPACE_ALIAS: self.workspace_dir} + return self + @field_validator("components", mode="before") @classmethod def normalize_component_types(cls, value): diff --git a/reme/schema/workspace.py b/reme/schema/workspace.py new file mode 100644 index 00000000..921b7e67 --- /dev/null +++ b/reme/schema/workspace.py @@ -0,0 +1,61 @@ +"""Per-workspace path views for multi-workspace applications.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .application_config import ApplicationConfig + + +@dataclass(frozen=True) +class WorkspaceContext: + """Resolved paths for one named workspace of an application. + + Built from ApplicationConfig when the application context is created. The + configured subdirectory names stay workspace-relative, matching the layout + a single-workspace application derives from the same config. + """ + + alias: str + workspace_path: Path + metadata_path: Path | None + session_path: Path | None + mem_session_path: Path | None + resource_path: Path | None + daily_path: Path | None + digest_path: Path | None + + @classmethod + def from_config(cls, alias: str, workspace: str, config: "ApplicationConfig") -> "WorkspaceContext": + """Resolve one alias of the config's workspaces mapping into absolute paths.""" + root = Path(workspace).absolute() + + def child(subdir: str) -> Path | None: + return root / subdir if subdir else None + + return cls( + alias=alias, + workspace_path=root, + metadata_path=child(config.metadata_dir), + session_path=child(config.session_dir), + mem_session_path=child(config.mem_session_dir), + resource_path=child(config.resource_dir), + daily_path=child(config.daily_dir), + digest_path=child(config.digest_dir), + ) + + def storage_dirs(self) -> list[Path]: + """Subdirectories the application materializes at startup, in configuration order.""" + return [ + path + for path in ( + self.metadata_path, + self.session_path, + self.mem_session_path, + self.resource_path, + self.daily_path, + self.digest_path, + ) + if path is not None + ] diff --git a/tests/unit/test_multi_workspace_config.py b/tests/unit/test_multi_workspace_config.py new file mode 100644 index 00000000..4d88f248 --- /dev/null +++ b/tests/unit/test_multi_workspace_config.py @@ -0,0 +1,159 @@ +"""Multi-workspace workspace_dir configuration contracts (Phase 1).""" + +# pylint: disable=protected-access,missing-function-docstring + +from pathlib import Path + +import pytest +from pydantic import ValidationError +import yaml + +from reme.application import Application +from reme.components import ApplicationContext +from reme.config import resolve_app_config +from reme.schema import ApplicationConfig, WorkspaceContext + + +def test_string_workspace_normalizes_to_default_alias(tmp_path): + config = ApplicationConfig(workspace_dir=str(tmp_path / "ws")) + + assert config.workspace_dir == str(tmp_path / "ws") + assert config.workspaces == {"default": str(tmp_path / "ws")} + assert "workspaces" not in config.model_dump(mode="json") + assert "workspaces" not in config.model_dump() + + +def test_mapping_workspace_keeps_order_and_first_entry_as_default(tmp_path): + config = ApplicationConfig( + workspace_dir={ + "beta": str(tmp_path / "b"), + "alpha": str(tmp_path / "a"), + }, + ) + + assert list(config.workspaces) == ["beta", "alpha"] + assert config.workspace_dir == str(tmp_path / "b") + assert config.workspaces["beta"] == config.workspace_dir + + +def test_mapping_workspace_expands_home_and_dedupes_by_resolved_path(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + config = ApplicationConfig(workspace_dir={"home": "~/ws", "other": str(tmp_path / "ws2")}) + + assert config.workspaces["home"] == str(tmp_path / "ws") + with pytest.raises(ValidationError, match="distinct directories"): + ApplicationConfig(workspace_dir={"one": "~/ws", "two": str(tmp_path / "ws")}) + + +def test_empty_mapping_rejected(): + with pytest.raises(ValidationError, match="at least one workspace"): + ApplicationConfig(workspace_dir={}) + + +def test_empty_alias_rejected(tmp_path): + with pytest.raises(ValidationError, match="aliases must be non-empty"): + ApplicationConfig(workspace_dir={" ": str(tmp_path / "ws")}) + + +def test_overlapping_directories_rejected(tmp_path): + parent = tmp_path / "parent" + child = parent / "child" + child.mkdir(parents=True) + + with pytest.raises(ValidationError, match="must not overlap"): + ApplicationConfig(workspace_dir={"parent": str(parent), "child": str(child)}) + with pytest.raises(ValidationError, match="must not overlap"): + ApplicationConfig(workspace_dir={"child": str(child), "parent": str(parent)}) + + sibling = tmp_path / "parent-sibling" + config = ApplicationConfig(workspace_dir={"parent": str(parent), "sibling": str(sibling)}) + assert list(config.workspaces) == ["parent", "sibling"] + + +def test_application_context_exposes_ordered_workspace_contexts(tmp_path): + context = ApplicationContext( + workspace_dir={"main": str(tmp_path / "main"), "scratch": str(tmp_path / "scratch")}, + ) + + assert list(context.workspaces) == ["main", "scratch"] + assert context.workspace is context.workspaces["main"] + main = context.workspaces["main"] + assert main.alias == "main" + assert main.workspace_path == tmp_path / "main" + assert main.metadata_path == tmp_path / "main" / "metadata" + assert main.session_path == tmp_path / "main" / "session" + assert main.mem_session_path == tmp_path / "main" / "mem_session" + assert main.resource_path == tmp_path / "main" / "resource" + assert main.daily_path == tmp_path / "main" / "daily" + assert main.digest_path == tmp_path / "main" / "digest" + + +def test_workspace_context_skips_empty_subdirs(tmp_path): + context = ApplicationContext(workspace_dir=str(tmp_path / "ws"), digest_dir="") + workspace = context.workspace + + assert workspace.digest_path is None + assert tmp_path / "ws" / "digest" not in workspace.storage_dirs() + assert len(workspace.storage_dirs()) == 5 + + +def test_string_workspace_layout_unchanged(tmp_path): + app = Application( + workspace_dir=str(tmp_path / "ws"), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + ) + + root = tmp_path / "ws" + assert root.is_dir() + for subdir in ("metadata", "session", "mem_session", "resource", "daily", "digest"): + assert (root / subdir).is_dir() + assert list(app.context.workspaces) == ["default"] + assert app.context.workspace.workspace_path == root + + +def test_application_creates_directories_for_every_alias(tmp_path): + Application( + workspace_dir={"main": str(tmp_path / "main"), "scratch": str(tmp_path / "scratch")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + ) + + for alias in ("main", "scratch"): + root = tmp_path / alias + assert root.is_dir() + for subdir in ("metadata", "session", "mem_session", "resource", "daily", "digest"): + assert (root / subdir).is_dir() + + +def test_yaml_config_loads_mapping_workspace(tmp_path): + config_path = tmp_path / "multi.yaml" + config_path.write_text( + yaml.safe_dump( + { + "workspace_dir": {"main": str(tmp_path / "m"), "scratch": str(tmp_path / "s")}, + "enable_logo": False, + }, + ), + encoding="utf-8", + ) + + resolved = resolve_app_config(config=str(config_path), log_config=False) + config = ApplicationConfig(**resolved) + + assert list(config.workspaces) == ["main", "scratch"] + assert config.workspace_dir == str(tmp_path / "m") + + +def test_default_yaml_string_workspace_still_loads(monkeypatch): + monkeypatch.delenv("CODEX_HOME", raising=False) + resolved = resolve_app_config(log_config=False) + config = ApplicationConfig(**resolved) + + assert isinstance(config.workspace_dir, str) + assert list(config.workspaces) == ["default"] + assert config.workspaces["default"] == str(Path(".reme").resolve()) From 66d45ac86341a30daeba5409e3969208179c1312 Mon Sep 17 00:00:00 2001 From: celanwang Date: Wed, 21 Oct 2026 13:04:57 +0800 Subject: [PATCH 02/21] fix(config): reject workspace alias collisions and validate mapping values Tighten workspace_dir mapping validation found in review of the Phase 1 multi-workspace change: - Reject aliases that collide after normalization (e.g. "alice" and " alice ") instead of silently dropping the earlier workspace, and reject non-string aliases such as YAML integer keys. - Validate each mapping value is a non-empty string before path normalization, so None, numbers, and blank entries fail with a clear per-alias error instead of becoming literal directories. - Reject any explicitly supplied workspaces field: it is derived from workspace_dir only, so an inconsistent input mapping can no longer steer directory creation away from the configured default workspace. Refs #632 --- reme/schema/application_config.py | 17 +++++--- tests/unit/test_multi_workspace_config.py | 47 ++++++++++++++++++++++- 2 files changed, 58 insertions(+), 6 deletions(-) diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index 1a19d4ef..6894c346 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -81,17 +81,25 @@ class ApplicationConfig(BaseModel): """Resolve a mapping workspace_dir, keeping its first entry as the default workspace.""" if not isinstance(data, dict): return data + if "workspaces" in data: + raise ValueError("workspaces is derived from workspace_dir; do not supply it") raw = data.get("workspace_dir") if not isinstance(raw, Mapping): return data if not raw: raise ValueError("workspace_dir mapping must contain at least one workspace") resolved: dict[str, str] = {} - for alias, path in raw.items(): - alias = str(alias).strip() + for key, path in raw.items(): + if not isinstance(key, str): + raise ValueError(f"workspace_dir aliases must be strings: {key!r}") + alias = key.strip() if not alias: raise ValueError("workspace_dir aliases must be non-empty") - resolved[alias] = str(Path(str(path)).expanduser().resolve(strict=False)) + if alias in resolved: + raise ValueError(f"workspace_dir aliases collide after normalization: {alias!r}") + if not isinstance(path, str) or not path.strip(): + raise ValueError(f"workspace_dir[{key!r}] must be a non-empty path string: {path!r}") + resolved[alias] = str(Path(path).expanduser().resolve(strict=False)) directories = list(resolved.values()) if len(set(directories)) != len(directories): raise ValueError(f"workspace_dir aliases must point to distinct directories: {resolved}") @@ -105,8 +113,7 @@ class ApplicationConfig(BaseModel): @model_validator(mode="after") def derive_default_workspace(self): """Normalize every workspace_dir form to an ordered mapping; the first entry is the default.""" - if not self.workspaces: - self.workspaces = {DEFAULT_WORKSPACE_ALIAS: self.workspace_dir} + self.workspaces = dict(self.workspaces) or {DEFAULT_WORKSPACE_ALIAS: self.workspace_dir} return self @field_validator("components", mode="before") diff --git a/tests/unit/test_multi_workspace_config.py b/tests/unit/test_multi_workspace_config.py index 4d88f248..e1b2199e 100644 --- a/tests/unit/test_multi_workspace_config.py +++ b/tests/unit/test_multi_workspace_config.py @@ -11,7 +11,7 @@ import yaml from reme.application import Application from reme.components import ApplicationContext from reme.config import resolve_app_config -from reme.schema import ApplicationConfig, WorkspaceContext +from reme.schema import ApplicationConfig def test_string_workspace_normalizes_to_default_alias(tmp_path): @@ -157,3 +157,48 @@ def test_default_yaml_string_workspace_still_loads(monkeypatch): assert isinstance(config.workspace_dir, str) assert list(config.workspaces) == ["default"] assert config.workspaces["default"] == str(Path(".reme").resolve()) + + +def test_alias_normalization_collision_rejected(tmp_path): + with pytest.raises(ValidationError, match="aliases collide"): + ApplicationConfig(workspace_dir={"alice": str(tmp_path / "a"), " alice ": str(tmp_path / "b")}) + + +def test_non_string_alias_rejected(tmp_path): + with pytest.raises(ValidationError, match="aliases must be strings"): + ApplicationConfig(workspace_dir={123: str(tmp_path / "a")}) + + +@pytest.mark.parametrize("bad_value", [None, 123, "", " "]) +def test_invalid_mapping_value_rejected(bad_value): + with pytest.raises(ValidationError, match="non-empty path string"): + ApplicationConfig(workspace_dir={"alice": bad_value}) + + +def test_supplied_workspaces_rejected_with_string_form(tmp_path): + with pytest.raises(ValidationError, match="derived from workspace_dir"): + ApplicationConfig(workspace_dir=str(tmp_path / "expected"), workspaces={"": str(tmp_path / "other")}) + + +def test_supplied_workspaces_rejected_with_mapping_form(tmp_path): + with pytest.raises(ValidationError, match="derived from workspace_dir"): + ApplicationConfig( + workspace_dir={"main": str(tmp_path / "main")}, + workspaces={"other": str(tmp_path / "other")}, + ) + + +def test_default_workspace_matches_directory_creation_source(tmp_path): + """The default workspace_dir and the materialized directories share one derivation.""" + app = Application( + workspace_dir={"main": str(tmp_path / "main"), "scratch": str(tmp_path / "scratch")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + ) + + assert app.config.workspace_dir == str(tmp_path / "main") + assert app.config.workspace_dir == app.config.workspaces["main"] + assert app.context.workspace.workspace_path == tmp_path / "main" + assert (tmp_path / "main" / "daily").is_dir() From e1dfb8f56de8986c2c6008e808b0184923549485 Mon Sep 17 00:00:00 2001 From: celanwang Date: Wed, 21 Oct 2026 14:09:05 +0800 Subject: [PATCH 03/21] feat(workspace): bind requests and components to the selected workspace Phase 2 of the multi-workspace plan. Job calls accept an optional workspace alias; the whole call binds to that workspace while single workspace (string) configurations behave exactly as before. - RuntimeContext carries the resolved WorkspaceContext; BaseJob and StreamJob consume the workspace kwarg at the execution boundary so it never reaches step payloads, and unknown aliases fail fast. - ComponentMixin.workspace_path is binding-aware: workspace-scoped component instances resolve their own alias, steps follow the current call binding, and everything else falls back to the default workspace, so existing path call sites need no changes. - Workspace-scoped components (local file stores including faiss/zvec, local embedding store, bm25 index, local/nx file graphs, local tag index, local file catalog) are instantiated once per alias; the default alias instance stays registered under the plain components mapping for backward-compatible lookups. Dependency binding and the topological start/reverse-close ordering resolve same-alias instances first, and Ref resolution in steps picks the bound alias instance, failing instead of falling back when one is missing. - External stateful backends opt out via supports_multi_workspace; Application rejects them in multi-workspace configurations (neo4j file graph is the first). - replace_component rejects scoped components when multiple workspaces are configured instead of silently diverging alias instances. Refs #632 --- docs/en/configuration.md | 7 + docs/zh/configuration.md | 6 + reme/application.py | 120 ++++++-- reme/components/application_context.py | 19 ++ reme/components/base_component.py | 56 +++- .../embedding_store/local_embedding_store.py | 2 + .../file_catalog/local_file_catalog.py | 2 + .../components/file_graph/local_file_graph.py | 2 + .../components/file_graph/neo4j_file_graph.py | 2 + reme/components/file_graph/nx_file_graph.py | 2 + .../components/file_store/local_file_store.py | 2 + reme/components/job/base_job.py | 24 +- reme/components/job/stream_job.py | 3 +- reme/components/keyword_index/bm25_index.py | 2 + reme/components/runtime_context.py | 8 + reme/components/tag_index/local_tag_index.py | 2 + reme/steps/base_step.py | 14 +- tests/unit/test_workspace_binding.py | 277 ++++++++++++++++++ 18 files changed, 518 insertions(+), 32 deletions(-) create mode 100644 tests/unit/test_workspace_binding.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 7b550ee5..7b55a54d 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -92,6 +92,13 @@ The first entry is the default workspace; every alias gets the same subdirectory resolved directories must be distinct and must not contain one another. A single string keeps the existing behavior and is normalized internally to a `default` alias. +Job calls accept an optional `workspace` field naming the alias to bind, for example +`reme read_note path=daily/note.md workspace=scratch`; omitting it binds the default workspace, and an unknown +alias fails fast. The binding is scoped to that call: workspace-aware state (file stores, indexes, graphs, +catalogs) resolves to per-alias instances, while shared components (models, tokenizers) stay single-instance. +External stateful backends that cannot isolate workspaces (currently the `neo4j` file graph) are rejected at +startup in multi-workspace configurations. + ## Enabling jobs `jobs..enabled` defaults to `true` for all Job types. Disabled jobs retain their configuration but do not diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 2609275b..14f0ff95 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -95,6 +95,12 @@ workspace_dir: (`metadata_dir`、`session_dir` 等上述目录)。别名必须非空,且解析后的目录必须互不相同、不能互相包含。 单个字符串保持原有行为,内部归一化为 `default` 别名。 +Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 +`reme read_note path=daily/note.md workspace=scratch`;缺省绑定默认 workspace,未知别名会直接报错。 +绑定仅在单次调用内生效:workspace 相关的状态组件(file store、索引、graph、catalog)按别名解析到各自实例, +共享组件(模型、tokenizer 等)保持单实例。无法隔离 workspace 的外部有状态后端(目前是 `neo4j` file graph) +在多 workspace 配置下会在启动时被拒绝。 + ## Job 启停 `jobs..enabled` 默认为 `true`,适用于所有 Job 类型。禁用的 Job 保留配置,但不启动、不暴露接口, diff --git a/reme/application.py b/reme/application.py index 405ac022..31ae9491 100644 --- a/reme/application.py +++ b/reme/application.py @@ -17,7 +17,8 @@ from .schema import ComponentConfig, Response, StreamChunk from .utils import execute_stream_task, print_logo, get_logger T = TypeVar("T", bound=BaseComponent) -_NodeKey = tuple[str, str] +# (component type, component name, workspace alias); "" marks shared instances. +_NodeKey = tuple[str, str, str] _UNSET = object() @@ -69,17 +70,40 @@ class Application(BaseComponent): ) def _init_components(self) -> None: - """Instantiate every component declared under config.components.""" + """Instantiate declared components; workspace-scoped ones once per alias.""" + multi_workspace = len(self.context.workspaces) > 1 for ctype, group in self.config.components.items(): self.context.components[ctype] = {} for name, cfg in group.items(): - self.context.components[ctype][name] = self._instantiate( - ctype, - cfg, - label=f"Component '{name}'", - expected_type=BaseComponent, - name=name, - ) + label = f"Component '{name}'" + backend_cls = self._backend_class(ctype, cfg, label) + if multi_workspace and not backend_cls.supports_multi_workspace: + raise ValueError( + f"{label} backend '{cfg.backend}' does not support multi-workspace configurations", + ) + if backend_cls.workspace_scoped: + instances = { + alias: self._instantiate( + ctype, + cfg, + label=label, + expected_type=BaseComponent, + name=name, + workspace_alias=alias, + ) + for alias in self.context.workspaces + } + # The default alias instance doubles as the plain lookup entry. + self.context.components[ctype][name] = instances[next(iter(instances))] + self.context.workspace_components.setdefault(ctype, {})[name] = instances + else: + self.context.components[ctype][name] = self._instantiate( + ctype, + cfg, + label=label, + expected_type=BaseComponent, + name=name, + ) def _init_jobs(self) -> None: """Instantiate every job declared under config.jobs.""" @@ -92,6 +116,15 @@ class Application(BaseComponent): name=name, ) + def _backend_class(self, ctype: ComponentType, cfg: ComponentConfig, label: str) -> type[T]: + """Resolve cfg.backend to a registered class without instantiating it.""" + if not cfg.backend: + raise ValueError(f"{label} is missing the required 'backend' field") + backend_cls = self.context.registry.get(ctype, cfg.backend) + if backend_cls is None: + raise ValueError(f"Unregistered backend '{cfg.backend}' for {label}") + return backend_cls + def _instantiate( self, ctype: ComponentType, @@ -100,6 +133,7 @@ class Application(BaseComponent): label: str, expected_type: type[T], name: str | None = None, + workspace_alias: str | None = None, ) -> T: """Resolve cfg.backend through the registry and construct the instance. @@ -108,17 +142,16 @@ class Application(BaseComponent): registered under the wrong component type. `name` is forwarded to the constructor for named components/jobs; leave it None for the service, which is keyed solely by type. + `workspace_alias` binds workspace-scoped instances to one workspace. """ - if not cfg.backend: - raise ValueError(f"{label} is missing the required 'backend' field") - backend_cls = self.context.registry.get(ctype, cfg.backend) - if backend_cls is None: - raise ValueError(f"Unregistered backend '{cfg.backend}' for {label}") + backend_cls = self._backend_class(ctype, cfg, label) params = cfg.model_dump() params["app_context"] = self.context if name is not None: params.setdefault("name", name) + if workspace_alias is not None: + params["workspace_alias"] = workspace_alias instance = backend_cls(**params) if not isinstance(instance, expected_type): got, want = type(instance).__name__, expected_type.__name__ @@ -127,16 +160,30 @@ class Application(BaseComponent): # ----- Dependency ordering ------------------------------------------ + def _component_nodes( + self, + replacement: tuple[_NodeKey, BaseComponent] | None = None, + ) -> dict[_NodeKey, BaseComponent]: + """One node per component instance; scoped components contribute one node per alias.""" + nodes: dict[_NodeKey, BaseComponent] = {} + for ctype, group in self.context.components.items(): + for name, comp in group.items(): + scoped = self.context.workspace_components.get(ctype, {}).get(name) + if scoped: + for alias, instance in scoped.items(): + nodes[(ctype, name, alias)] = instance + else: + nodes[(ctype, name, "")] = comp + if replacement is not None: + nodes[replacement[0]] = replacement[1] + return nodes + def _topological_order( self, replacement: tuple[_NodeKey, BaseComponent] | None = None, ) -> list[BaseComponent]: """Return components in dependency order via Kahn's algorithm; raise on missing dep or cycle.""" - nodes: dict[_NodeKey, BaseComponent] = { - (ctype, name): comp for ctype, group in self.context.components.items() for name, comp in group.items() - } - if replacement is not None: - nodes[replacement[0]] = replacement[1] + nodes = self._component_nodes(replacement) in_degree, dependents = self._build_dependency_graph(nodes) ready = [k for k, d in in_degree.items() if d == 0] @@ -159,13 +206,30 @@ class Application(BaseComponent): def _build_dependency_graph( nodes: dict[_NodeKey, BaseComponent], ) -> tuple[dict[_NodeKey, int], dict[_NodeKey, list[_NodeKey]]]: - """Compute in-degree and adjacency lists; raise if a required dep is missing.""" + """Compute in-degree and adjacency lists; raise if a required dep is missing. + + Scoped instances depend on the same-alias instance of a scoped + dependency, falling back to shared ("" alias) instances, then to the + dependency's default-alias instance for shared dependents. + """ in_degree: dict[_NodeKey, int] = dict.fromkeys(nodes, 0) dependents: dict[_NodeKey, list[_NodeKey]] = {k: [] for k in nodes} for key, comp in nodes.items(): for dep in comp.dependencies: - dep_key = (dep.ctype, dep.name) - if dep_key in nodes: + dep_key = next( + ( + candidate + for candidate in ( + (dep.ctype, dep.name, key[2]), + (dep.ctype, dep.name, ""), + ) + if candidate in nodes + ), + None, + ) + if dep_key is None: + dep_key = next((k for k in nodes if k[:2] == (dep.ctype, dep.name)), None) + if dep_key is not None: dependents[dep_key].append(key) in_degree[key] += 1 elif not dep.optional: @@ -295,6 +359,13 @@ class Application(BaseComponent): raise KeyError(f"Component '{name}' not found in {component_type}") old_component = group[name] + scoped_instances = self.context.workspace_components.get(component_type, {}).get(name) + if scoped_instances is not None and len(scoped_instances) > 1: + raise NotImplementedError( + f"replace_component does not support workspace-scoped component " + f"'{component_type}:{name}' in multi-workspace configurations", + ) + alias = next(iter(scoped_instances)) if scoped_instances is not None else None replacement_config = ( config.model_copy(deep=True) if isinstance(config, ComponentConfig) @@ -306,6 +377,7 @@ class Application(BaseComponent): label=f"Component '{name}'", expected_type=BaseComponent, name=name, + workspace_alias=alias, ) # Validate fields without evaluating lazy properties before injection. for key, value in (runtime_updates or {}).items(): @@ -315,7 +387,7 @@ class Application(BaseComponent): ) setattr(replacement, key, value) - node_key = (component_type, name) + node_key = (component_type, name, alias or "") replacement_order = Application._topological_order(self, replacement=(node_key, replacement)) was_started = old_component.is_started @@ -347,6 +419,8 @@ class Application(BaseComponent): # Commit the new generation synchronously so observers cannot see # a context with only some dependency references updated. group[name] = replacement + if scoped_instances is not None: + scoped_instances[alias] = replacement config_group[name] = replacement_config for consumer, attr in consumers: consumer.__dict__[attr] = replacement diff --git a/reme/components/application_context.py b/reme/components/application_context.py index dca5dd67..0aff4d1a 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -35,6 +35,10 @@ class ApplicationContext: # Populated by Application during initialization. self.service: "BaseService | None" = None self.components: dict[str, dict[str, "BaseComponent"]] = {} + # Per-alias instances of workspace-scoped components, keyed + # ctype -> name -> alias. The default alias instance is also the one + # registered under ``components`` for backward-compatible lookups. + self.workspace_components: dict[str, dict[str, dict[str, "BaseComponent"]]] = {} self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None @@ -43,3 +47,18 @@ class ApplicationContext: # This is in-memory state, not durable storage; use workspace files or a store when state # must survive an Application restart. self.metadata: dict[str, Any] = {} + + def get_component(self, ctype: str, name: str, alias: str | None = None) -> "BaseComponent | None": + """Resolve a component, honoring workspace bindings for scoped components. + + Returns None when the component is not configured at all. Raises + KeyError when a workspace-scoped component has no instance for the + requested alias — a missing scoped resource must fail, never fall back. + """ + scoped = self.workspace_components.get(ctype, {}).get(name) + if scoped is not None: + alias = alias or next(iter(self.workspaces)) + if alias in scoped: + return scoped[alias] + raise KeyError(f"Component '{name}' in {ctype} has no instance for workspace alias {alias!r}") + return self.components.get(ctype, {}).get(name) diff --git a/reme/components/base_component.py b/reme/components/base_component.py index 74ef4633..51f32717 100644 --- a/reme/components/base_component.py +++ b/reme/components/base_component.py @@ -14,21 +14,44 @@ if TYPE_CHECKING: T = TypeVar("T", bound="BaseComponent") +def _context_workspaces(app_context) -> dict | None: + """Return the workspace mapping of a real ApplicationContext, else None. + + Duck-typing is not enough here: MagicMock-style test doubles answer every + attribute access, so the workspace layer engages only for the real type. + """ + from .application_context import ApplicationContext # lazy: registry imports this module + + if isinstance(app_context, ApplicationContext): + return app_context.workspaces + return None + + class ComponentMixin: """Shared state for components and steps: identity, config, workspace paths.""" component_type = ComponentEnum.BASE + # Workspace-scoped components keep per-workspace state (indexes, graphs, + # catalogs, chunk stores) and get one instance per configured alias. + workspace_scoped: bool = False + # Stateful external backends that cannot isolate workspaces must opt out + # explicitly; Application rejects them in multi-workspace configurations. + supports_multi_workspace: bool = True def __init__( self, name: str | None = None, backend: str = "", app_context: "ApplicationContext | None" = None, + workspace_alias: str | None = None, **kwargs, ) -> None: self.name: str = name or self.__class__.__name__ self.backend: str = backend self.app_context: "ApplicationContext | None" = app_context + # Bound alias for workspace-scoped component instances; None for steps, + # jobs, and shared components. + self.workspace_alias: str | None = workspace_alias self.kwargs: dict = dict(kwargs) logger = get_logger() @@ -36,10 +59,25 @@ class ComponentMixin: @property def workspace_path(self) -> Path: - """Absolute workspace root directory (cwd when no app_context is attached).""" + """Workspace root: the bound alias for scoped components, the current + call's workspace binding for steps, otherwise the default workspace.""" if self.app_context is None: return Path.cwd() - return Path(self.app_context.app_config.workspace_dir).absolute() + workspaces = _context_workspaces(self.app_context) + if workspaces is None: + # Test doubles and standalone-style contexts without the workspace layer. + return Path(self.app_context.app_config.workspace_dir).absolute() + alias = self.workspace_alias + if alias is None: + context = getattr(self, "context", None) + bound = getattr(context, "workspace", None) if context is not None else None + alias = bound.alias if bound is not None else None + if alias is None: + return next(iter(workspaces.values())).workspace_path + workspace = workspaces.get(alias) + if workspace is None: + raise KeyError(f"Unknown workspace alias {alias!r}") + return workspace.workspace_path def to_workspace_relative(self, path: str | Path) -> str: """Convert ``path`` to a POSIX workspace path; keep outside paths absolute.""" @@ -183,8 +221,18 @@ class BaseComponent(ComponentMixin, ABC): setattr(self, attr, None) def _resolve_from_context(self, attr: str, dep: Dependency) -> None: - """Context-bound mode: look up the component from ``app_context.components``.""" - target = self.app_context.components.get(dep.ctype, {}).get(dep.name) + """Context-bound mode: look up the component from ``app_context.components``. + + Workspace-scoped dependencies resolve to the same-alias instance as the + dependent; a missing alias instance is an error, never a fallback. + """ + scoped = getattr(self.app_context, "workspace_components", {}).get(dep.ctype, {}).get(dep.name) + if scoped is not None and self.workspace_alias is not None: + target = scoped.get(self.workspace_alias) + if target is None: + raise ValueError(f"{dep.ctype} '{dep.name}' has no instance for workspace '{self.workspace_alias}'.") + else: + target = self.app_context.components.get(dep.ctype, {}).get(dep.name) if target is not None: setattr(self, attr, target) elif dep.optional: diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index e49bae9f..b5465a82 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -47,6 +47,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): Delegates actual embedding computation to a bound ``as_embedding`` component. """ + workspace_scoped = True + def __init__( self, as_embedding: str = "default", diff --git a/reme/components/file_catalog/local_file_catalog.py b/reme/components/file_catalog/local_file_catalog.py index 1cfe1199..b058e790 100644 --- a/reme/components/file_catalog/local_file_catalog.py +++ b/reme/components/file_catalog/local_file_catalog.py @@ -12,6 +12,8 @@ from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst class LocalFileCatalog(BaseFileCatalog): """Dict-backed catalog persisted as JSONL.""" + workspace_scoped = True + def __init__(self, encoding: str = "utf-8", **kwargs): super().__init__(**kwargs) self.encoding = encoding diff --git a/reme/components/file_graph/local_file_graph.py b/reme/components/file_graph/local_file_graph.py index f47c75e4..2adc8067 100644 --- a/reme/components/file_graph/local_file_graph.py +++ b/reme/components/file_graph/local_file_graph.py @@ -13,6 +13,8 @@ from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst class LocalFileGraph(BaseFileGraph): """Dict-backed file graph; uses FileLink.target_path for adjacency.""" + workspace_scoped = True + def __init__(self, **kwargs): super().__init__(**kwargs) self._nodes: dict[str, FileNode] = {} diff --git a/reme/components/file_graph/neo4j_file_graph.py b/reme/components/file_graph/neo4j_file_graph.py index 5d2881b4..c66eb7af 100644 --- a/reme/components/file_graph/neo4j_file_graph.py +++ b/reme/components/file_graph/neo4j_file_graph.py @@ -68,6 +68,8 @@ class Neo4jFileGraph(BaseFileGraph): database: target db name (default ``neo4j``) """ + supports_multi_workspace = False + def __init__( self, uri: str = "bolt://localhost:7687", diff --git a/reme/components/file_graph/nx_file_graph.py b/reme/components/file_graph/nx_file_graph.py index 27257925..b521d877 100644 --- a/reme/components/file_graph/nx_file_graph.py +++ b/reme/components/file_graph/nx_file_graph.py @@ -16,6 +16,8 @@ class NxFileGraph(BaseFileGraph): Real node carries ``node`` attr; virtual (dangling target) does not. """ + workspace_scoped = True + def __init__(self, **kwargs): super().__init__(**kwargs) try: diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 7cee6d85..25c78a8b 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -86,6 +86,8 @@ class LocalFileStore(BaseFileStore): at least one of embedding / keyword must be present. """ + workspace_scoped = True + def __init__( self, embedding_store: str = "default", diff --git a/reme/components/job/base_job.py b/reme/components/job/base_job.py index eb8314bc..98b2c050 100644 --- a/reme/components/job/base_job.py +++ b/reme/components/job/base_job.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING -from ..base_component import BaseComponent +from ..base_component import BaseComponent, _context_workspaces from ..component_registry import R from ..runtime_context import RuntimeContext from ...enumeration import ComponentEnum @@ -90,12 +90,32 @@ class BaseJob(BaseComponent): if not self.enabled: raise ValueError(f"Job '{self.name}' is disabled") + def _bind_workspace(self, merged: dict): + """Consume the optional ``workspace`` alias and resolve the bound workspace. + + The alias never reaches step payloads. An unknown alias fails fast, + matching disabled-job and missing-job behavior at this boundary. + """ + alias = merged.pop("workspace", None) + workspaces = _context_workspaces(self.app_context) + if workspaces is None: + if alias is not None: + raise ValueError(f"workspace={alias!r} requires an application with configured workspaces") + return None + if alias is None: + return next(iter(workspaces.values())) + workspace = workspaces.get(alias) + if workspace is None: + raise ValueError(f"Unknown workspace {alias!r}; configured: {sorted(workspaces)}") + return workspace + async def __call__(self, **kwargs) -> Response: """Run all steps in order, capturing any failure into the response.""" self.check_enabled() self._record_call() merged = {**self.kwargs, **kwargs} - context = RuntimeContext(**merged) + workspace = self._bind_workspace(merged) + context = RuntimeContext(workspace=workspace, **merged) try: for step in self._build_steps(): await step(context) diff --git a/reme/components/job/stream_job.py b/reme/components/job/stream_job.py index b947a12a..f7c042ad 100644 --- a/reme/components/job/stream_job.py +++ b/reme/components/job/stream_job.py @@ -15,7 +15,8 @@ class StreamJob(BaseJob): self.check_enabled() self._record_call() merged = {**self.kwargs, **kwargs} - context = RuntimeContext(**merged) + workspace = self._bind_workspace(merged) + context = RuntimeContext(workspace=workspace, **merged) try: for step in self._build_steps(): await step(context) diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 06daa870..49c381ac 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/reme/components/keyword_index/bm25_index.py @@ -37,6 +37,8 @@ from ...utils.async_utils import complete_in_thread class BM25Index(BaseKeywordIndex): """BM25 inverted index with lazy deletion and on-disk persistence.""" + workspace_scoped = True + def __init__(self, k1: float = 1.5, b: float = 0.75, index_version: str = "v1", **kwargs): super().__init__(**kwargs) self.k1 = k1 diff --git a/reme/components/runtime_context.py b/reme/components/runtime_context.py index 355b94ce..456b0bee 100644 --- a/reme/components/runtime_context.py +++ b/reme/components/runtime_context.py @@ -1,16 +1,22 @@ """Per-request runtime context shared across steps and jobs.""" import asyncio +from typing import TYPE_CHECKING from ..enumeration import ChunkEnum from ..schema import Response, StreamChunk +if TYPE_CHECKING: + from ..schema import WorkspaceContext + class RuntimeContext: """Scratch space for a single execution. Holds the response object, an optional stream queue, a stop event, and a free-form data dict accessible via mapping-style operators (``ctx[key]``). + ``workspace`` binds the whole execution to one configured workspace; job + entry points resolve the request's workspace alias into it. """ def __init__( @@ -18,11 +24,13 @@ class RuntimeContext: response: Response | None = None, stream_queue: asyncio.Queue | None = None, stop_event: asyncio.Event | None = None, + workspace: "WorkspaceContext | None" = None, **kwargs, ): self.response: Response = response or Response() self.stream_queue: asyncio.Queue | None = stream_queue self.stop_event: asyncio.Event | None = stop_event + self.workspace: "WorkspaceContext | None" = workspace self.data: dict = kwargs # ----- Data dict access ---------------------------------------------- diff --git a/reme/components/tag_index/local_tag_index.py b/reme/components/tag_index/local_tag_index.py index 7392d668..43ada6e1 100644 --- a/reme/components/tag_index/local_tag_index.py +++ b/reme/components/tag_index/local_tag_index.py @@ -12,6 +12,8 @@ from ...schema import FileFrontMatter, FileNode class LocalTagIndex(BaseTagIndex): """Maintain bidirectional path/tag relationships without separate source I/O.""" + workspace_scoped = True + reserved_tag_keys = frozenset(FileFrontMatter.model_fields) | { "kind", "session_id", diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index b2042237..ad9e7b4b 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -7,7 +7,7 @@ from typing import Any, TYPE_CHECKING from agentscope.model import ChatModelBase from ..components.agent_wrapper.base_agent_wrapper import BaseAgentWrapper -from ..components.base_component import ComponentMixin +from ..components.base_component import ComponentMixin, _context_workspaces from ..components.component_registry import R from ..components.file_catalog import BaseFileCatalog from ..components.file_store import BaseFileStore @@ -87,7 +87,17 @@ class Ref: if self.optional: return None raise RuntimeError(f"app_context is not set when resolving '{self.key}'") - comp = obj.app_context.components.get(self.comp_enum, {}).get(name) + if _context_workspaces(obj.app_context) is not None: + context = obj.context + bound = getattr(context, "workspace", None) if context is not None else None + comp = obj.app_context.get_component( + self.comp_enum, + name, + alias=bound.alias if bound is not None else None, + ) + else: + # Test doubles and standalone-style contexts without the workspace layer. + comp = obj.app_context.components.get(self.comp_enum, {}).get(name) if comp is None: if self.optional: return None diff --git a/tests/unit/test_workspace_binding.py b/tests/unit/test_workspace_binding.py new file mode 100644 index 00000000..07bd83f5 --- /dev/null +++ b/tests/unit/test_workspace_binding.py @@ -0,0 +1,277 @@ +"""Phase 2: requests and components bind to the selected workspace.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +from pathlib import Path + +import pytest + +from reme.application import Application +from reme.components import ApplicationContext +from reme.components.base_component import BaseComponent +from reme.components.job import BaseJob, StreamJob +from reme.components.runtime_context import RuntimeContext +from reme.components.component_registry import R +from reme.enumeration import ChunkEnum, ComponentEnum +from reme.schema import FileChunk, FileNode +from reme.steps.base_step import BaseStep + + +def _app(tmp_path, **overrides): + config = { + "workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "cli"}, + "components": { + "tokenizer": {"default": {"backend": "regex"}}, + "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, + "keyword_index": {"default": {"backend": "bm25"}}, + "file_graph": {"default": {"backend": "local"}}, + }, + "jobs": { + "write_note": { + "backend": "base", + "description": "write a note", + "parameters": {"type": "object", "properties": {}}, + "steps": [{"backend": "write_step"}], + }, + "read_note": { + "backend": "base", + "description": "read a note", + "parameters": {"type": "object", "properties": {}}, + "steps": [{"backend": "read_step"}], + }, + }, + **overrides, + } + return Application(**config) + + +class _WorkspaceReporterStep(BaseStep): + """Report the bound workspace alias as answer and stream chunk.""" + + async def execute(self): + bound = self.context.workspace + alias = bound.alias if bound is not None else "" + assert "workspace" not in self.context.data, "workspace alias must not leak into step payloads" + self.context.response.answer = alias + if self.context.stream: + await self.context.add_stream_string(alias, ChunkEnum.CONTENT) + + +def _reporter_context(tmp_path, job_cls): + registry = R.copy() + registry.register(_WorkspaceReporterStep, "workspace_reporter_step") + context = ApplicationContext( + registry=registry, + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + job = job_cls( + name="report", + app_context=context, + parameters={"type": "object", "properties": {}}, + steps=[{"backend": "workspace_reporter_step"}], + ) + return job + + +# ----- Selection semantics --------------------------------------------------- + + +@pytest.mark.asyncio +async def test_default_binding_uses_first_workspace(tmp_path): + app = _app(tmp_path) + async with app: + response = await app.run_job("write_note", path="daily/note.md", content="default-target") + + assert response.success, response.answer + assert "default-target" in (tmp_path / "alpha" / "daily" / "note.md").read_text(encoding="utf-8") + assert not (tmp_path / "beta" / "daily" / "note.md").exists() + + +@pytest.mark.asyncio +async def test_unknown_workspace_alias_rejected(tmp_path): + app = _app(tmp_path) + async with app: + with pytest.raises(ValueError, match="Unknown workspace 'nope'"): + await app.run_job("write_note", path="daily/note.md", content="x", workspace="nope") + + +@pytest.mark.asyncio +async def test_unknown_workspace_alias_rejected_for_stream_job(tmp_path): + job = _reporter_context(tmp_path, StreamJob) + await job.start() + with pytest.raises(ValueError, match="Unknown workspace 'nope'"): + async for _chunk in _stream(job, workspace="nope"): + pass + await job.close() + + +async def _stream(job, **kwargs): + from reme.utils import execute_stream_task + + queue = asyncio.Queue() + task = asyncio.create_task(job(stream_queue=queue, **kwargs)) + async for chunk in execute_stream_task(stream_queue=queue, task=task, output_format="chunk"): + yield chunk + + +# ----- Concurrent and streaming isolation ------------------------------------ + + +@pytest.mark.asyncio +async def test_concurrent_same_relative_path_stays_isolated(tmp_path): + app = _app(tmp_path) + async with app: + await asyncio.gather( + *( + app.run_job("write_note", path="daily/shared.md", content=f"{alias}-{i}", workspace=alias) + for i in range(5) + for alias in ("alpha", "beta") + ) + ) + alpha = (tmp_path / "alpha" / "daily" / "shared.md").read_text(encoding="utf-8") + beta = (tmp_path / "beta" / "daily" / "shared.md").read_text(encoding="utf-8") + assert "alpha-4" in alpha and "beta" not in alpha + assert "beta-4" in beta and "alpha" not in beta + + +@pytest.mark.asyncio +async def test_stream_chunks_carry_caller_workspace(tmp_path): + job = _reporter_context(tmp_path, StreamJob) + await job.start() + + async def collect(alias): + return [chunk async for chunk in _stream(job, workspace=alias)] + + chunks_alpha, chunks_beta = await asyncio.gather(collect("alpha"), collect("beta")) + content_alpha = [c.chunk for c in chunks_alpha if c.chunk_type == ChunkEnum.CONTENT] + content_beta = [c.chunk for c in chunks_beta if c.chunk_type == ChunkEnum.CONTENT] + assert content_alpha == ["alpha"] + assert content_beta == ["beta"] + assert chunks_alpha[-1].done and chunks_beta[-1].done + await job.close() + + +# ----- Component binding ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_workspace_scoped_components_instantiate_per_alias(tmp_path): + app = _app(tmp_path) + async with app: + store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + store_beta = app.context.get_component(ComponentEnum.FILE_STORE, "default", "beta") + + assert store_alpha is not store_beta + assert app.context.components[ComponentEnum.FILE_STORE]["default"] is store_alpha # default alias entry + assert store_alpha.workspace_path == tmp_path / "alpha" + assert store_beta.workspace_path == tmp_path / "beta" + assert store_alpha.chunks_path.parent == tmp_path / "alpha" / "metadata" / "file_store" + + # Same-alias dependency wiring: each store binds its own alias's index/graph. + index_alpha = app.context.get_component(ComponentEnum.KEYWORD_INDEX, "default", "alpha") + graph_beta = app.context.get_component(ComponentEnum.FILE_GRAPH, "default", "beta") + assert store_alpha.keyword_index is index_alpha + assert store_beta.file_graph is graph_beta + + with pytest.raises(KeyError, match="no instance for workspace alias 'nope'"): + app.context.get_component(ComponentEnum.FILE_STORE, "default", "nope") + + +@pytest.mark.asyncio +async def test_store_chunks_and_search_stay_in_own_workspace(tmp_path): + app = _app(tmp_path) + async with app: + for alias, text in (("alpha", "apple pie"), ("beta", "banana bread")): + store = app.context.get_component(ComponentEnum.FILE_STORE, "default", alias) + node = FileNode(path=f"daily/{alias}.md", st_mtime=0.0) + chunk = FileChunk(path=f"daily/{alias}.md", text=text, start_line=0, end_line=0) + await store.upsert([(node, [chunk])]) + + store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + store_beta = app.context.get_component(ComponentEnum.FILE_STORE, "default", "beta") + hits_alpha = await store_alpha.keyword_search("apple", limit=5, search_filter={}) + hits_beta = await store_beta.keyword_search("apple", limit=5, search_filter={}) + assert [c.path for c in hits_alpha] == ["daily/alpha.md"] + assert hits_beta == [] + + +@pytest.mark.asyncio +async def test_store_state_persists_per_alias_across_restart(tmp_path): + app = _app(tmp_path) + async with app: + store = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + await store.upsert([(FileNode(path="daily/a.md", st_mtime=0.0), [FileChunk(path="daily/a.md", text="alpha data", start_line=0, end_line=0)])]) + + rebuilt = _app(tmp_path) + async with rebuilt: + store_alpha = rebuilt.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + store_beta = rebuilt.context.get_component(ComponentEnum.FILE_STORE, "default", "beta") + nodes_alpha = await store_alpha.get_nodes() + nodes_beta = await store_beta.get_nodes() + assert [n.path for n in nodes_alpha] == ["daily/a.md"] + assert nodes_beta == [] + + +# ----- External backends and single-workspace parity ------------------------- + + +@pytest.mark.asyncio +async def test_external_stateful_backend_rejected_in_multi_workspace(tmp_path): + class FakeExternalGraph(BaseComponent): + component_type = ComponentEnum.FILE_GRAPH + supports_multi_workspace = False + + with R.preserve(allow_mutation=True): + R.register(FakeExternalGraph, "fake_external") + with pytest.raises(ValueError, match="does not support multi-workspace"): + _app( + tmp_path, + components={ + "file_graph": {"default": {"backend": "fake_external"}}, + }, + ) + + # Single-workspace configurations may still use the same backend. + with R.preserve(allow_mutation=True): + R.register(FakeExternalGraph, "fake_external") + app = Application( + workspace_dir=str(tmp_path / "single"), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "fake_external"}}}, + ) + async with app: + assert app.context.components[ComponentEnum.FILE_GRAPH]["default"] is not None + + +@pytest.mark.asyncio +async def test_single_workspace_string_config_binds_default(tmp_path): + app = Application( + workspace_dir=str(tmp_path / "only"), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + jobs={ + "write_note": { + "backend": "base", + "parameters": {"type": "object", "properties": {}}, + "steps": [{"backend": "write_step"}], + }, + }, + ) + async with app: + response = await app.run_job("write_note", path="daily/note.md", content="plain") + assert response.success, response.answer + assert (tmp_path / "only" / "daily" / "note.md").exists() + # The explicit default alias resolves to the same workspace. + response = await app.run_job("write_note", path="daily/other.md", content="plain", workspace="default") + assert response.success + assert (tmp_path / "only" / "daily" / "other.md").exists() From 08a35acf21a1232391d3e9e8d64e789324eb933d Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 22 Oct 2026 11:44:53 +0800 Subject: [PATCH 04/21] fix(workspace): resolve scoped components through binding and fix replace_component enumeration Review follow-ups on Phase 2 of the multi-workspace change: - ProactiveTopicsStep._digest_titles read the digest file catalog through the plain components table, always landing on the default workspace's instance; under a non-default binding it compared candidates against the wrong workspace's digest titles. Resolve it through the bound workspace instead. A full audit of direct components-table reads found this to be the only scoped-component access; the remaining readers target shared components or reporting. - replace_component enumerated consumers and shutdown tracking from the plain components table only, missing non-default alias instances of workspace-scoped components: their dependency references were never rewired and they were double-counted in _started_components. Enumerate every component instance, deduplicated by identity. - test_workspace_binding.py lint: drop unused imports, add the missing class docstring, and commit hook reformatting. Refs #632 --- reme/application.py | 20 ++++-- reme/steps/evolve/proactive/topics.py | 12 +++- tests/unit/test_workspace_binding.py | 91 +++++++++++++++++++++++++-- 3 files changed, 111 insertions(+), 12 deletions(-) diff --git a/reme/application.py b/reme/application.py index 31ae9491..c49b8c21 100644 --- a/reme/application.py +++ b/reme/application.py @@ -308,6 +308,18 @@ class Application(BaseComponent): setattr(component, key, value) return component + def _all_component_instances(self) -> list[BaseComponent]: + """Every component instance: plain entries plus all alias instances of scoped components.""" + instances: dict[int, BaseComponent] = {} + for group in self.context.components.values(): + for component in group.values(): + instances[id(component)] = component + for group in self.context.workspace_components.values(): + for scoped in group.values(): + for instance in scoped.values(): + instances[id(instance)] = instance + return list(instances.values()) + def _replacement_shutdown_order( self, old_component: BaseComponent, @@ -318,9 +330,7 @@ class Application(BaseComponent): started_ids = {id(component) for component in self._started_components} started_ids.discard(id(old_component)) started_ids.add(id(replacement)) - component_ids = { - id(component) for components in self.context.components.values() for component in components.values() - } + component_ids = {id(component) for component in self._all_component_instances()} non_components = [ component for component in self._started_components @@ -392,9 +402,7 @@ class Application(BaseComponent): was_started = old_component.is_started consumers: list[tuple[BaseComponent, str]] = [] - candidates: list[BaseComponent] = [ - component for components in self.context.components.values() for component in components.values() - ] + candidates: list[BaseComponent] = self._all_component_instances() candidates.extend(self.context.jobs.values()) if self.context.service is not None: candidates.append(self.context.service) diff --git a/reme/steps/evolve/proactive/topics.py b/reme/steps/evolve/proactive/topics.py index 9bf5df75..a22156cc 100644 --- a/reme/steps/evolve/proactive/topics.py +++ b/reme/steps/evolve/proactive/topics.py @@ -21,6 +21,7 @@ import time from ...base_step import BaseStep from .._evolve import passthrough_response from ....components import R +from ....components.base_component import _context_workspaces from ....enumeration import ComponentEnum from ....schema import ProactiveState, ProactiveStateFile, ProactiveTopic from ....schema.proactive import clamp_confidence @@ -340,7 +341,16 @@ class ProactiveTopicsStep(BaseStep): async def _digest_titles(self) -> list[str]: if self.app_context is None or self.digest_compare_limit <= 0: return [] - catalog = self.app_context.components.get(ComponentEnum.FILE_CATALOG, {}).get("digest") + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if _context_workspaces(self.app_context) is not None: + catalog = self.app_context.get_component( + ComponentEnum.FILE_CATALOG, + "digest", + alias=bound.alias if bound is not None else None, + ) + else: + # Test doubles and standalone-style contexts without the workspace layer. + catalog = self.app_context.components.get(ComponentEnum.FILE_CATALOG, {}).get("digest") if catalog is None: return [] try: diff --git a/tests/unit/test_workspace_binding.py b/tests/unit/test_workspace_binding.py index 07bd83f5..97d39e09 100644 --- a/tests/unit/test_workspace_binding.py +++ b/tests/unit/test_workspace_binding.py @@ -3,15 +3,13 @@ # pylint: disable=protected-access,missing-function-docstring import asyncio -from pathlib import Path import pytest from reme.application import Application from reme.components import ApplicationContext from reme.components.base_component import BaseComponent -from reme.components.job import BaseJob, StreamJob -from reme.components.runtime_context import RuntimeContext +from reme.components.job import StreamJob from reme.components.component_registry import R from reme.enumeration import ChunkEnum, ComponentEnum from reme.schema import FileChunk, FileNode @@ -131,7 +129,7 @@ async def test_concurrent_same_relative_path_stays_isolated(tmp_path): app.run_job("write_note", path="daily/shared.md", content=f"{alias}-{i}", workspace=alias) for i in range(5) for alias in ("alpha", "beta") - ) + ), ) alpha = (tmp_path / "alpha" / "daily" / "shared.md").read_text(encoding="utf-8") beta = (tmp_path / "beta" / "daily" / "shared.md").read_text(encoding="utf-8") @@ -205,7 +203,14 @@ async def test_store_state_persists_per_alias_across_restart(tmp_path): app = _app(tmp_path) async with app: store = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") - await store.upsert([(FileNode(path="daily/a.md", st_mtime=0.0), [FileChunk(path="daily/a.md", text="alpha data", start_line=0, end_line=0)])]) + await store.upsert( + [ + ( + FileNode(path="daily/a.md", st_mtime=0.0), + [FileChunk(path="daily/a.md", text="alpha data", start_line=0, end_line=0)], + ), + ], + ) rebuilt = _app(tmp_path) async with rebuilt: @@ -223,6 +228,8 @@ async def test_store_state_persists_per_alias_across_restart(tmp_path): @pytest.mark.asyncio async def test_external_stateful_backend_rejected_in_multi_workspace(tmp_path): class FakeExternalGraph(BaseComponent): + """Stand-in for a stateful external backend that cannot isolate workspaces.""" + component_type = ComponentEnum.FILE_GRAPH supports_multi_workspace = False @@ -275,3 +282,77 @@ async def test_single_workspace_string_config_binds_default(tmp_path): response = await app.run_job("write_note", path="daily/other.md", content="plain", workspace="default") assert response.success assert (tmp_path / "only" / "daily" / "other.md").exists() + + +# ----- Review follow-ups: binding-aware catalog reads and replace_component -- + + +@pytest.mark.asyncio +async def test_digest_titles_follow_workspace_binding(tmp_path): + """The digest dedup source must read the bound workspace's catalog.""" + from reme.components.runtime_context import RuntimeContext + from reme.steps.evolve.proactive.topics import ProactiveTopicsStep + + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_catalog": {"digest": {"backend": "local"}}}, + ) + async with app: + catalog_alpha = app.context.get_component(ComponentEnum.FILE_CATALOG, "digest", "alpha") + catalog_beta = app.context.get_component(ComponentEnum.FILE_CATALOG, "digest", "beta") + await catalog_alpha.upsert([FileNode(path="digest/2026-01-01/alpha-only.md", st_mtime=1.0)]) + await catalog_beta.upsert([FileNode(path="digest/2026-01-01/beta-only.md", st_mtime=1.0)]) + + step = ProactiveTopicsStep(app_context=app.context, digest_compare_limit=10) + step.context = RuntimeContext(workspace=app.context.workspaces["beta"]) + # Beta binding sees beta's digests only: an alpha-only title must not + # mark a beta candidate as a known duplicate. + assert await step._digest_titles() == ["beta-only"] + + step.context = RuntimeContext() + assert await step._digest_titles() == ["alpha-only"] + + +@pytest.mark.asyncio +async def test_replace_shared_component_rewires_every_alias(tmp_path): + """Replacing a shared component rewires scoped consumers of every alias once.""" + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={ + "tokenizer": {"default": {"backend": "regex"}}, + "keyword_index": {"default": {"backend": "bm25"}}, + }, + ) + async with app: + old = app.context.components[ComponentEnum.TOKENIZER]["default"] + close_calls = 0 + original_close = old.close + + async def counting_close(): + nonlocal close_calls + close_calls += 1 + await original_close() + + old.close = counting_close + + index_alpha = app.context.get_component(ComponentEnum.KEYWORD_INDEX, "default", "alpha") + index_beta = app.context.get_component(ComponentEnum.KEYWORD_INDEX, "default", "beta") + assert index_alpha.tokenizer is old + assert index_beta.tokenizer is old + + replacement = await app.replace_component(ComponentEnum.TOKENIZER, "default", config={"backend": "regex"}) + + assert replacement is not old + assert index_alpha.tokenizer is replacement + assert index_beta.tokenizer is replacement + assert close_calls == 1 + started_ids = [id(component) for component in app._started_components] + assert len(started_ids) == len(set(started_ids)) From 9c85de395b1ba82f454ac2d68d268a559fefb78b Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 22 Oct 2026 11:50:53 +0800 Subject: [PATCH 05/21] feat(workspace): namespace agent sessions and bind background jobs Phase 3 of the multi-workspace plan, completing the approved contract: - Nested propagation: step run_job calls inherit the caller's workspace binding (explicit selection still wins), and agent job_tools carry the binding through the server-owned injected_job_kwargs channel in the Claude Code and AgentScope wrappers, so models can never switch workspaces. The MCP tool boundary forwards injected bindings and rejects caller conflicts, and the Codex job bridge launches against the bound workspace path. - Session namespacing: Codex thread tool-context guards key on workspace alias plus thread id (the shared turn lock intentionally serializes turns across workspaces), resource agent session UUIDs mix in the bound alias, and tool-context state (search dedup, drafts, search call budgets) is namespaced per non-default workspace while the default keeps its legacy keys. - Background and cron jobs bind through their configuration: a job may declare a workspace alias (defaulting to the first entry, unknown aliases fail fast), and extra workspaces are watched via explicit per-alias job variants rather than automatic fan-out. - The Claude Code transcript source stays global (a Claude Code-side limitation, now documented); token accounting remains application-level. Refs #632 --- docs/en/configuration.md | 20 ++ docs/zh/configuration.md | 19 ++ .../agent_wrapper/as_agent_wrapper.py | 17 +- .../agent_wrapper/base_agent_wrapper.py | 7 +- .../agent_wrapper/cc_agent_wrapper.py | 14 +- .../agent_wrapper/codex_agent_wrapper.py | 30 +- reme/components/job/background_job.py | 3 +- reme/components/job/cron_job.py | 4 +- reme/steps/base_step.py | 35 ++- reme/steps/common/chat.py | 1 + reme/steps/evolve/auto_memory.py | 1 + reme/steps/evolve/auto_tag.py | 1 + reme/steps/evolve/base_auto_resource.py | 12 +- reme/steps/evolve/dream/extract.py | 2 + reme/steps/evolve/dream/integrate.py | 1 + reme/steps/index/_tool_context.py | 2 +- reme/steps/index/search.py | 2 +- tests/unit/test_workspace_phase3.py | 269 ++++++++++++++++++ 18 files changed, 419 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_workspace_phase3.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 7b55a54d..d681c786 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -99,6 +99,26 @@ catalogs) resolves to per-alias instances, while shared components (models, toke External stateful backends that cannot isolate workspaces (currently the `neo4j` file graph) are rejected at startup in multi-workspace configurations. +The binding propagates through nested job calls and agent tool calls (`job_tools`) as a server-owned constraint +that models cannot override. Background and cron jobs bind through their job configuration, so extra workspaces +are watched by declaring explicit per-alias variants: + +```yaml +jobs: + index_update_loop: + backend: background + # watches the default workspace; workspace defaults to the first alias + steps: [...] + index_update_loop_scratch: + backend: background + workspace: scratch + steps: [...] +``` + +Two scoping notes: the Claude Code integration reads transcripts from the global `~/.claude/projects` directory +(a Claude Code-side limitation — only ReMe's own copies are workspace-scoped), and token usage accounting remains +application-level across workspaces. + ## Enabling jobs `jobs..enabled` defaults to `true` for all Job types. Disabled jobs retain their configuration but do not diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 14f0ff95..d2422797 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -101,6 +101,25 @@ Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 共享组件(模型、tokenizer 等)保持单实例。无法隔离 workspace 的外部有状态后端(目前是 `neo4j` file graph) 在多 workspace 配置下会在启动时被拒绝。 +绑定会经嵌套 Job 调用和 agent 工具调用(`job_tools`)传播,作为模型无法覆盖的服务端约束。 +Background 和 cron Job 通过 Job 配置绑定别名,额外 workspace 的 watch 通过显式配置按别名的变体实现: + +```yaml +jobs: + index_update_loop: + backend: background + # 监控默认 workspace;workspace 缺省为第一个别名 + steps: [...] + index_update_loop_scratch: + backend: background + workspace: scratch + steps: [...] +``` + +两点作用域说明:Claude Code 集成从全局 `~/.claude/projects` 目录读取 transcript +(这是 Claude Code 源端的固有限制,ReMe 自己的副本按 workspace 隔离);token 用量统计保持 Application 级, +不随 workspace 拆分。 + ## Job 启停 `jobs..enabled` 默认为 `true`,适用于所有 Job 类型。禁用的 Job 保留配置,但不启动、不暴露接口, diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 935c39e0..5e95f78f 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -165,9 +165,14 @@ class AsAgentWrapper(BaseAgentWrapper): job: "BaseJob", tool_context_id: str | None = None, injected_job_kwargs: dict[str, Any] | None = None, + workspace: str | None = None, ) -> FunctionTool: injected = cls._resolve_injected_job_kwargs( - {"tool_context_id": tool_context_id, "injected_job_kwargs": injected_job_kwargs}, + { + "tool_context_id": tool_context_id, + "injected_job_kwargs": injected_job_kwargs, + "workspace": workspace, + }, ) async def run_job(**kwargs) -> ToolChunk: @@ -318,7 +323,15 @@ class AsAgentWrapper(BaseAgentWrapper): builtin_tools = "all" tools: list[ToolBase] = [] tools.extend(self._builtin_tools(builtin_tools, sequential_tool_calls=sequential_tool_calls)) - tools.extend(self._make_tool(job, tool_context_id, kwargs.get("injected_job_kwargs")) for job in resolved_jobs) + tools.extend( + self._make_tool( + job, + tool_context_id, + kwargs.get("injected_job_kwargs"), + workspace=kwargs.get("workspace"), + ) + for job in resolved_jobs + ) toolkit = kwargs.get("toolkit") or Toolkit( tools=tools, skills_or_loaders=skills, diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 69f1ace6..167c1ec3 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -163,12 +163,15 @@ class BaseAgentWrapper(BaseComponent): ``injected_job_kwargs`` values are enforced constraints added by the wrapper after receiving the model's tool arguments; the model can - neither see nor override them. ``tool_context_id`` keeps its dedicated - option but is carried through the same mechanism. + neither see nor override them. ``tool_context_id`` and the workspace + binding keep their dedicated options but are carried through the same + mechanism, so the model can never switch workspaces. """ injected = dict(kwargs.get("injected_job_kwargs") or {}) if tool_context_id := kwargs.get("tool_context_id"): injected["tool_context_id"] = tool_context_id + if workspace := kwargs.get("workspace"): + injected["workspace"] = workspace return injected @staticmethod diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 068175df..f2e460a4 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -83,11 +83,16 @@ class CcAgentWrapper(BaseAgentWrapper): job: "BaseJob", tool_context_id: str | None = None, injected_job_kwargs: dict[str, Any] | None = None, + workspace: str | None = None, ): from claude_agent_sdk import SdkMcpTool injected = cls._resolve_injected_job_kwargs( - {"tool_context_id": tool_context_id, "injected_job_kwargs": injected_job_kwargs}, + { + "tool_context_id": tool_context_id, + "injected_job_kwargs": injected_job_kwargs, + "workspace": workspace, + }, ) async def run_job(args): @@ -203,7 +208,12 @@ class CcAgentWrapper(BaseAgentWrapper): if self.MCP_SERVER_NAME in opts.mcp_servers: raise ValueError(f"mcp_servers already contains reserved server name {self.MCP_SERVER_NAME!r}") sdk_tools = [ - self._make_tool(job, kwargs.get("tool_context_id"), kwargs.get("injected_job_kwargs")) + self._make_tool( + job, + kwargs.get("tool_context_id"), + kwargs.get("injected_job_kwargs"), + workspace=kwargs.get("workspace"), + ) for job in resolved_jobs ] opts.mcp_servers[self.MCP_SERVER_NAME] = create_sdk_mcp_server( diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index b744452d..ada1103e 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -16,6 +16,7 @@ import tempfile from typing import Any, TYPE_CHECKING from .base_agent_wrapper import BaseAgentWrapper +from ..base_component import _context_workspaces from ..component_registry import R from ...enumeration import ChunkEnum from ...schema import StreamChunk, TokenUsage @@ -240,6 +241,17 @@ class CodexAgentWrapper(BaseAgentWrapper): names = ", ".join(invalid) raise TypeError(f"Codex client options must be configured on the wrapper: {names}") + def _call_workspace_path(self, kwargs: dict[str, Any]) -> Path: + """Resolve the workspace binding forwarded by the calling step to a path.""" + alias = str(kwargs.get("workspace") or "") + if not alias: + return self.workspace_path + workspaces = _context_workspaces(self.app_context) + if workspaces is None or alias not in workspaces: + configured = sorted(workspaces) if workspaces else [] + raise ValueError(f"Unknown workspace {alias!r}; configured: {configured}") + return workspaces[alias].workspace_path + def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: from ..job import BackgroundJob, StreamJob @@ -258,12 +270,16 @@ class CodexAgentWrapper(BaseAgentWrapper): "--config", config_source, "--workspace", - str(self.workspace_path), + str(self._call_workspace_path(kwargs)), ] 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 {}): + injected_job_kwargs = dict(kwargs.get("injected_job_kwargs") or {}) + if workspace := kwargs.get("workspace"): + # The bridge injects the binding into nested job calls; models cannot override it. + injected_job_kwargs["workspace"] = workspace + if injected_job_kwargs: args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) return { "command": sys.executable, @@ -300,6 +316,7 @@ class CodexAgentWrapper(BaseAgentWrapper): from openai_codex import ApprovalMode, Sandbox from openai_codex.types import Personality, ThreadSource, ThreadStartSource + alias = str(kwargs.get("workspace") or "") resume = kwargs.get("resume") or "" session_id = kwargs.get("session_id") or "" if resume and session_id and resume != session_id: @@ -318,8 +335,9 @@ class CodexAgentWrapper(BaseAgentWrapper): 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]: + guard_key = f"{alias}/{thread_id}" if alias else thread_id + if not fork_session and guard_key in self._thread_tool_contexts: + if requested_tool_context != self._thread_tool_contexts[guard_key]: raise ValueError("tool_context_id and injected_job_kwargs cannot change when resuming a Codex thread") common = { @@ -353,7 +371,7 @@ class CodexAgentWrapper(BaseAgentWrapper): thread_source=thread_source, **common, ) - self._thread_tool_contexts[thread.id] = requested_tool_context + self._thread_tool_contexts[f"{alias}/{thread.id}" if alias else thread.id] = requested_tool_context return thread def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: @@ -429,7 +447,7 @@ class CodexAgentWrapper(BaseAgentWrapper): 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 + key: context for key, context in self._thread_tool_contexts.items() if key.rsplit("/", 1)[-1] in loaded_ids } async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None: diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index ec8406e6..9c50b40a 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -126,7 +126,8 @@ class BackgroundJob(BaseJob): self.check_enabled() self._record_call() merged = {**self.kwargs, **kwargs} - context = RuntimeContext(stop_event=self._stop_event, **merged) + workspace = self._bind_workspace(merged) + context = RuntimeContext(stop_event=self._stop_event, workspace=workspace, **merged) for step in self._build_steps(): await step(context) return context.response diff --git a/reme/components/job/cron_job.py b/reme/components/job/cron_job.py index 71948597..709a49ee 100644 --- a/reme/components/job/cron_job.py +++ b/reme/components/job/cron_job.py @@ -35,7 +35,9 @@ class CronJob(BackgroundJob): return max(0.0, (nxt - now).total_seconds()) async def _execute_steps(self) -> Response: - context = RuntimeContext(**self.kwargs) + merged = dict(self.kwargs) + workspace = self._bind_workspace(merged) + context = RuntimeContext(workspace=workspace, **merged) for step in self._build_steps(): await step(context) return context.response diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index ad9e7b4b..f54b465a 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -199,12 +199,45 @@ class BaseStep(ComponentMixin, ABC): return self.app_context.jobs.get(name) async def run_job(self, name: str, /, **kwargs) -> Response: - """Execute a job by name and kwargs, return the final response.""" + """Execute a job by name and kwargs, return the final response. + + Nested calls inherit the current call's workspace binding unless the + caller selects a workspace explicitly. + """ job: "BaseJob | None" = self.get_job(name) if job is None: raise RuntimeError(f"Job {name} not found") + if alias := self._bound_workspace_alias(): + kwargs.setdefault("workspace", alias) return await job(**kwargs) + def _bound_workspace_alias(self) -> str: + """Alias of the current call's workspace binding. + + Returns "" when unbound or bound to the default workspace: the default + needs no forwarding because every resolution already lands there, and + single-workspace behavior stays byte-identical. + """ + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if bound is None: + return "" + workspaces = _context_workspaces(self.app_context) + if workspaces and bound.alias == next(iter(workspaces)): + return "" + return bound.alias + + def _workspace_forward_kwargs(self) -> dict: + """Forward the current call's workspace binding to agent wrapper calls.""" + alias = self._bound_workspace_alias() + return {"workspace": alias} if alias else {} + + def _scoped_state_key(self, key: str) -> str: + """Namespace a shared in-memory state key with the bound workspace alias.""" + alias = self._bound_workspace_alias() + if not alias or not key: + return key + return f"{alias}/{key}" + def _resolve_dispatch_step(self, raw: _DispatchStep): """Resolve a dispatch step spec to (step class, init params).""" if isinstance(raw, str): diff --git a/reme/steps/common/chat.py b/reme/steps/common/chat.py index 17feda80..a7eaf062 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -50,6 +50,7 @@ this chat intentionally provides read-only tools. Reply in the user's language." "system_prompt": self._system_prompt(), "job_tools": self.READ_ONLY_TOOLS, "builtin_tools": [], + **self._workspace_forward_kwargs(), } if session_id := self.context.get("session_id"): wrapper_kwargs["resume"] = str(session_id) diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 431785c3..6c15f7de 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -630,6 +630,7 @@ class AutoMemoryStep(BaseStep): system_prompt=self.prompt_format("system_prompt"), job_tools=self.create_tools if created else self.update_tools, **reply_kwargs, + **self._workspace_forward_kwargs(), ) self.logger.info(f"[{self.name}] agent done path={note_path} has_result={bool(result.get('result'))}") diff --git a/reme/steps/evolve/auto_tag.py b/reme/steps/evolve/auto_tag.py index 823d59cb..8053050f 100644 --- a/reme/steps/evolve/auto_tag.py +++ b/reme/steps/evolve/auto_tag.py @@ -119,6 +119,7 @@ class AutoTagStep(BaseStep): ), job_tools=self.tools, injected_job_kwargs=tool_context, + **self._workspace_forward_kwargs(), ) path = self.workspace_path / target.path diff --git a/reme/steps/evolve/base_auto_resource.py b/reme/steps/evolve/base_auto_resource.py index ff7a9275..fa1d74ad 100644 --- a/reme/steps/evolve/base_auto_resource.py +++ b/reme/steps/evolve/base_auto_resource.py @@ -74,9 +74,9 @@ def _compute_note_stem(filename: str) -> str: return PurePosixPath(filename).stem -def _compute_agent_session_id(path: str) -> str: - """Return a stable UUID session id for agent backends.""" - return str(uuid.uuid5(uuid.NAMESPACE_URL, path)) +def _compute_agent_session_id(path: str, workspace_alias: str = "") -> str: + """Return a stable UUID session id for agent backends, namespaced per workspace.""" + return str(uuid.uuid5(uuid.NAMESPACE_URL, f"{workspace_alias}:{path}" if workspace_alias else path)) def _parse_resource_path(file_path: str, resource_dir: str) -> tuple[str, str]: @@ -401,9 +401,13 @@ class BaseAutoResourceStep(BaseStep): "system_prompt": self.prompt_format("system_prompt"), "job_tools": self.create_tools if state.created else self.update_tools, **(reply_kwargs or {}), + **self._workspace_forward_kwargs(), } if "session_id" not in agent_kwargs: - agent_kwargs["session_id"] = _compute_agent_session_id(file_path) + agent_kwargs["session_id"] = _compute_agent_session_id( + file_path, + workspace_alias=agent_kwargs.get("workspace", ""), + ) # Consume the resource-only option before forwarding kwargs to the wrapper. if agent_kwargs.pop("scope_note_tools", False): # Bind the final allocated path last, never a path supplied by the agent. diff --git a/reme/steps/evolve/dream/extract.py b/reme/steps/evolve/dream/extract.py index ba4fd4c1..bb99b4f6 100644 --- a/reme/steps/evolve/dream/extract.py +++ b/reme/steps/evolve/dream/extract.py @@ -302,6 +302,7 @@ class DreamExtractStep(BaseStep): user_message, system_prompt=system_prompt, job_tools=list(_TOOLS), + **self._workspace_forward_kwargs(), ) self.logger.info(f"[{self.name}] agent done files={len(batch)} has_result={bool(result.get('result'))}") raw_result = agent_reply_result_text(result) @@ -324,6 +325,7 @@ class DreamExtractStep(BaseStep): user_message + retry_feedback, system_prompt=system_prompt, job_tools=list(_TOOLS), + **self._workspace_forward_kwargs(), ) raw_result = agent_reply_result_text(result) state.extract_summary = "\n".join(filter(None, (state.extract_summary, raw_result))) diff --git a/reme/steps/evolve/dream/integrate.py b/reme/steps/evolve/dream/integrate.py index f3946318..9bcdd94b 100644 --- a/reme/steps/evolve/dream/integrate.py +++ b/reme/steps/evolve/dream/integrate.py @@ -145,6 +145,7 @@ class DreamIntegrateStep(BaseStep): bucket=bucket, ), job_tools=list(_TOOLS), + **self._workspace_forward_kwargs(), ) except Exception as e: # noqa: BLE001 after = _snapshot_digest(workspace, digest_dir) diff --git a/reme/steps/index/_tool_context.py b/reme/steps/index/_tool_context.py index 4b14a6d5..993f47f1 100644 --- a/reme/steps/index/_tool_context.py +++ b/reme/steps/index/_tool_context.py @@ -20,4 +20,4 @@ class _ToolContextMixin: contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) else: contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) - return contexts.setdefault(tool_context_id, {}) + return contexts.setdefault(self._scoped_state_key(tool_context_id), {}) diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index f4e20b3c..eff8222c 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -238,7 +238,7 @@ class SearchStep(_ToolContextMixin, BaseStep): if maximum < 1: raise ValueError("max_search_calls must be positive") budgets = self.app_context.metadata.setdefault("__search_call_budgets", {}) - budget = budgets.setdefault(tool_context_id, {"count": 0, "lock": asyncio.Lock()}) + budget = budgets.setdefault(self._scoped_state_key(tool_context_id), {"count": 0, "lock": asyncio.Lock()}) async with budget["lock"]: if budget["count"] >= maximum: self.logger.warning( diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py new file mode 100644 index 00000000..8ecd1538 --- /dev/null +++ b/tests/unit/test_workspace_phase3.py @@ -0,0 +1,269 @@ +"""Phase 3: nested propagation, session namespacing, and background binding.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio + +import pytest + +from reme.application import Application +from reme.components import ApplicationContext +from reme.components.agent_wrapper.base_agent_wrapper import BaseAgentWrapper +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper +from reme.components.component_registry import R +from reme.components.job import BackgroundJob, BaseJob, CronJob +from reme.components.runtime_context import RuntimeContext +from reme.steps.base_step import BaseStep +from reme.steps.evolve.base_auto_resource import _compute_agent_session_id +from reme.steps.index.draft import AddDraftStep + + +class _AliasReporterStep(BaseStep): + """Answer with the bound workspace alias.""" + + async def execute(self): + bound = self.context.workspace + self.context.response.answer = bound.alias if bound is not None else "" + + +class _NestedCallerStep(BaseStep): + """Call the report job as a nested job and relay its answer.""" + + async def execute(self): + response = await self.run_job("report") + self.context.response.answer = f"parent->{response.answer}" + + +def _registry(): + registry = R.copy() + registry.register(_AliasReporterStep, "alias_reporter_step") + registry.register(_NestedCallerStep, "nested_caller_step") + return registry + + +def _context(tmp_path): + return ApplicationContext( + registry=_registry(), + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + + +def _job(context, name, backend="alias_reporter_step", job_cls=BaseJob, **kwargs): + return job_cls( + name=name, + app_context=context, + parameters={"type": "object", "properties": {}}, + steps=[{"backend": backend}], + **kwargs, + ) + + +# ----- Nested propagation ---------------------------------------------------- + + +@pytest.mark.asyncio +async def test_nested_run_job_inherits_workspace_binding(tmp_path): + context = _context(tmp_path) + parent = _job(context, "parent", backend="nested_caller_step") + child = _job(context, "report") + context.jobs = {"parent": parent, "report": child} + await parent.start() + await child.start() + + response = await parent(workspace="beta") + assert response.answer == "parent->beta" + # Default workspace stays the implicit binding for unscoped calls. + response = await parent() + assert response.answer == "parent->alpha" + # Explicit nested selection wins over inheritance. + caller = _NestedCallerStep(app_context=context) + caller.context = RuntimeContext(workspace=context.workspaces["alpha"]) + response = await caller.run_job("report", workspace="beta") + assert response.answer == "beta" + + +def test_workspace_forwarded_through_injected_channel_only(): + # The calling step's binding joins the server-owned injected kwargs... + injected = BaseAgentWrapper._resolve_injected_job_kwargs( + {"injected_job_kwargs": {"_allowed_paths": ["daily/x.md"]}, "workspace": "beta"}, + ) + assert injected == {"_allowed_paths": ["daily/x.md"], "workspace": "beta"} + # ...and model-supplied tool arguments can never override it. + with pytest.raises(ValueError, match="cannot be provided by the model"): + BaseAgentWrapper._merge_injected_job_kwargs({"workspace": "alpha"}, injected) + + +# ----- Session/tool-context namespacing -------------------------------------- + + +def test_agent_session_uuid_namespaced_per_workspace(): + path = "resource/2026-01-01/report.pdf" + legacy = _compute_agent_session_id(path) + alpha = _compute_agent_session_id(path, workspace_alias="alpha") + beta = _compute_agent_session_id(path, workspace_alias="beta") + # Default/empty alias keeps the historical UUID for backward compatibility. + assert _compute_agent_session_id(path, workspace_alias="") == legacy + assert len({legacy, alpha, beta}) == 3 + + +@pytest.mark.asyncio +async def test_tool_context_state_namespaced_per_workspace(tmp_path): + context = _context(tmp_path) + step = AddDraftStep(app_context=context) + for alias in ("alpha", "beta"): + step.context = RuntimeContext(workspace=context.workspaces[alias], tool_context_id="ctx-1", text=f"draft-{alias}") + await step.execute() + + stores = context.metadata["tool_contexts"] + # The default workspace keeps the legacy plain key; other aliases are namespaced. + assert sorted(stores) == ["beta/ctx-1", "ctx-1"] + assert stores["ctx-1"]["drafts"] == ["draft-alpha"] + assert stores["beta/ctx-1"]["drafts"] == ["draft-beta"] + + +@pytest.mark.asyncio +async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): + from types import SimpleNamespace + + wrapper = CodexAgentWrapper( + app_context=ApplicationContext( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ), + ) + + 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") + + codex = FakeCodex() + await wrapper._open_thread(codex, {"tool_context_id": "ctx-a"}) + # Same thread id under a different workspace is a different guard entry. + await wrapper._open_thread(codex, {"session_id": "thread-1", "workspace": "beta", "tool_context_id": "ctx-b"}) + # Same workspace resume with changed constraints still raises. + with pytest.raises(ValueError, match="cannot change"): + await wrapper._open_thread(codex, {"session_id": "thread-1", "tool_context_id": "ctx-b"}) + + assert sorted(wrapper._thread_tool_contexts) == ["beta/thread-1", "thread-1"] + + +# ----- Background/cron binding ----------------------------------------------- + + +@pytest.mark.asyncio +async def test_background_job_binds_configured_workspace(tmp_path): + context = _context(tmp_path) + job = _job(context, "bg", job_cls=BackgroundJob, workspace="beta") + job.step_specs = [job._resolve_step({"backend": "alias_reporter_step"})] + assert (await job()).answer == "beta" + + default_job = _job(context, "bg-default", job_cls=BackgroundJob) + default_job.step_specs = [default_job._resolve_step({"backend": "alias_reporter_step"})] + assert (await default_job()).answer == "alpha" + + with pytest.raises(ValueError, match="Unknown workspace 'nope'"): + await _job(context, "bg-bad", job_cls=BackgroundJob, workspace="nope")() + + +@pytest.mark.asyncio +async def test_cron_job_binds_configured_workspace(tmp_path): + context = _context(tmp_path) + job = _job(context, "cron", job_cls=CronJob, cron="0 0 * * *", workspace="beta") + job.step_specs = [job._resolve_step({"backend": "alias_reporter_step"})] + assert (await job._execute_steps()).answer == "beta" + + default_job = _job(context, "cron-default", job_cls=CronJob, cron="0 0 * * *") + default_job.step_specs = [default_job._resolve_step({"backend": "alias_reporter_step"})] + assert (await default_job._execute_steps()).answer == "alpha" + + +@pytest.mark.asyncio +async def test_per_alias_watch_variant_jobs_write_own_workspace(tmp_path): + """Two watch-style variants with explicit aliases act on their own workspace.""" + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={ + "tokenizer": {"default": {"backend": "regex"}}, + "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, + "keyword_index": {"default": {"backend": "bm25"}}, + "file_graph": {"default": {"backend": "local"}}, + }, + jobs={ + "watch_alpha": { + "backend": "background", + "workspace": "alpha", + "enable_serve": False, + "parameters": {"type": "object", "properties": {}}, + "path": "daily/marker.md", + "content": "alpha-watch", + "steps": [{"backend": "write_step"}], + }, + "watch_beta": { + "backend": "background", + "workspace": "beta", + "enable_serve": False, + "parameters": {"type": "object", "properties": {}}, + "path": "daily/marker.md", + "content": "beta-watch", + "steps": [{"backend": "write_step"}], + }, + }, + ) + async with app: + await asyncio.sleep(0.5) + + alpha = (tmp_path / "alpha" / "daily" / "marker.md").read_text(encoding="utf-8") + beta = (tmp_path / "beta" / "daily" / "marker.md").read_text(encoding="utf-8") + assert "alpha-watch" in alpha and "beta-watch" not in alpha + assert "beta-watch" in beta and "alpha-watch" not in beta + + +# ----- MCP boundary and Codex bridge ----------------------------------------- + + +@pytest.mark.asyncio +async def test_mcp_tool_boundary_forwards_injected_workspace(tmp_path): + from types import SimpleNamespace + + from reme.components.service.mcp_tools import add_mcp_job + + context = _context(tmp_path) + job = _job(context, "report") + job.step_specs = [job._resolve_step({"backend": "alias_reporter_step"})] + + tools = [] + server = SimpleNamespace(add_tool=tools.append) + assert add_mcp_job(server, job, injected_job_kwargs={"workspace": "beta"}, tool_error_on_failure=False) + + # The server-injected binding reaches the job boundary and wins over callers. + assert await tools[0].fn() == "beta" + with pytest.raises(Exception, match="cannot be provided by the caller"): + await tools[0].fn(workspace="alpha") + + +@pytest.mark.asyncio +async def test_codex_bridge_uses_bound_workspace_path(tmp_path): + import json + + context = _context(tmp_path) + context.jobs = {"search_job": _job(context, "search_job")} + wrapper = CodexAgentWrapper(app_context=context) + + config = wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "beta"}) + assert config["args"][config["args"].index("--workspace") + 1] == str(tmp_path / "beta") + raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] + assert json.loads(raw) == {"workspace": "beta"} + + default_config = wrapper._mcp_server_config({"job_tools": ["search_job"]}) + assert default_config["args"][default_config["args"].index("--workspace") + 1] == str(tmp_path / "alpha") + + with pytest.raises(ValueError, match="Unknown workspace 'nope'"): + wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "nope"}) + + await wrapper.close() From 0ede9664f430538d4b25269757748eb296fca4d5 Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 22 Oct 2026 12:27:19 +0800 Subject: [PATCH 06/21] fix(workspace): carry explicit bindings through sessions, bridges, and nested jobs Review follow-ups closing four isolation gaps: - Agent session state follows the call's workspace binding: the AS wrapper loads/dumps session state and cleans expired sessions under the bound workspace's mem_session dir, the CC wrapper builds its CLAUDE_CONFIG_DIR and cwd from the bound workspace, and Codex thread cwd follows the binding (CODEX_HOME stays shared by design; thread guards are already alias-namespaced). Same session id in two workspaces no longer sees each other's state. - The Codex job bridge passes --workspace-alias alongside the config snapshot, which now carries the full workspace mapping; the child application keeps every workspace and injects the server-owned alias into nested job calls. Single-workspace snapshots keep the legacy --workspace path override. - Execution bindings are explicit end to end: the resolved alias (including the default) is forwarded through nested run_job calls, so a child job's own configured workspace cannot silently override the caller's binding. Storage keys stay decoupled: the default workspace keeps its legacy plain state keys, while other aliases use unambiguous (alias, key) tuples that cannot collide with ids containing separators. Refs #632 --- .../agent_wrapper/as_agent_wrapper.py | 52 ++-- .../agent_wrapper/base_agent_wrapper.py | 21 +- .../agent_wrapper/cc_agent_wrapper.py | 18 +- .../agent_wrapper/codex_agent_wrapper.py | 29 +- .../agent_wrapper/codex_mcp_server.py | 23 +- reme/steps/base_step.py | 44 ++-- tests/unit/test_base_agent_wrapper.py | 4 +- tests/unit/test_cc_agent_wrapper.py | 2 +- tests/unit/test_token_usage.py | 2 +- tests/unit/test_workspace_phase3.py | 247 +++++++++++++++++- 10 files changed, 372 insertions(+), 70 deletions(-) diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 5e95f78f..a9c34de8 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -157,7 +157,7 @@ class AsAgentWrapper(BaseAgentWrapper): super().__init__(**kwargs) self.as_llm = self.bind(as_llm, BaseAsLLM, optional=False) self.session_retention_days = int(session_retention_days) - self._session_cleanup_done = False + self._session_cleanup_done: set[str] = set() @classmethod def _make_tool( @@ -190,9 +190,10 @@ class AsAgentWrapper(BaseAgentWrapper): names: list[str] | str | bool | None = "all", *, sequential_tool_calls: bool = False, + call_kwargs: dict[str, Any] | None = None, ) -> list[ToolBase]: - """Return selected AgentScope built-in tools rooted at ``self.cwd``.""" - cwd = str(self.cwd) + """Return selected AgentScope built-in tools rooted at this call's workspace.""" + cwd = str(self._call_cwd(call_kwargs or {})) backend = WorkspaceBackend(cwd, self.bash_environment) factories = { "bash": lambda: BypassAnalysisBash(cwd=cwd, backend=backend), @@ -228,10 +229,15 @@ class AsAgentWrapper(BaseAgentWrapper): @property def session_path(self) -> Path: - """Directory used for persisted AgentScope sessions.""" + """Directory used for persisted AgentScope sessions in the default workspace.""" + return self._session_path_for({}) + + def _session_path_for(self, kwargs: dict[str, Any]) -> Path: + """Session directory under the workspace this call is bound to.""" + root = self._call_workspace_path(kwargs) if self.app_context is None: - return self.workspace_path / "mem_session" / "agentscope" - return self.workspace_path / self.app_context.app_config.mem_session_dir / "agentscope" + return root / "mem_session" / "agentscope" + return root / self.app_context.app_config.mem_session_dir / "agentscope" @staticmethod def _validate_session_id(session_id: str, field: str = "session_id") -> str: @@ -239,15 +245,15 @@ class AsAgentWrapper(BaseAgentWrapper): raise ValueError(f"{field} must be a valid UUID: {session_id!r}") return session_id.lower() - def _cleanup_expired_sessions(self) -> None: + def _cleanup_expired_sessions(self, session_path: Path) -> None: """Delete persisted session files older than ``session_retention_days``.""" - if self._session_cleanup_done or self.session_retention_days <= 0: - self._session_cleanup_done = True + key = str(session_path) + if key in self._session_cleanup_done or self.session_retention_days <= 0: + self._session_cleanup_done.add(key) return - session_path = self.session_path if not session_path.is_dir(): - self._session_cleanup_done = True + self._session_cleanup_done.add(key) return cutoff = time.time() - self.session_retention_days * 24 * 60 * 60 @@ -264,7 +270,7 @@ class AsAgentWrapper(BaseAgentWrapper): self.logger.info( f"Cleaned {removed} AgentScope session(s) older than {self.session_retention_days} day(s)", ) - self._session_cleanup_done = True + self._session_cleanup_done.add(key) async def _load_state(self, kwargs: dict[str, Any], perm_mode: PermissionMode) -> AgentState: resume = kwargs.get("resume") or "" @@ -279,7 +285,7 @@ class AsAgentWrapper(BaseAgentWrapper): raise ValueError("session_id cannot be used with resume unless fork_session=True") if resume: - handler = AsStateHandler.for_session(self.session_path, resume) + handler = AsStateHandler.for_session(self._session_path_for(kwargs), resume) state = await handler.load_or_none() if state is None: raise FileNotFoundError(f"AgentScope session not found: {resume}") @@ -297,8 +303,8 @@ class AsAgentWrapper(BaseAgentWrapper): return AgentState(session_id=session_id or str(uuid4()), permission_context=PermissionContext(mode=perm_mode)) - async def _dump_state(self, state: AgentState) -> None: - await AsStateHandler.for_session(self.session_path, state.session_id).dump(state) + async def _dump_state(self, state: AgentState, kwargs: dict[str, Any]) -> None: + await AsStateHandler.for_session(self._session_path_for(kwargs), state.session_id).dump(state) def _resolve_skills(self, skills: list[str] | str | None) -> list[str]: """Resolve configured skill names to AgentScope local skill directories.""" @@ -310,7 +316,7 @@ class AsAgentWrapper(BaseAgentWrapper): if model is None: raise ValueError("AsAgentWrapper requires a bound as_llm component with a valid model.") - self._cleanup_expired_sessions() + self._cleanup_expired_sessions(self._session_path_for(kwargs)) system_prompt = kwargs.get("system_prompt", "You are a helpful assistant.") job_tools: list[str] = kwargs.get("job_tools", []) @@ -322,7 +328,13 @@ class AsAgentWrapper(BaseAgentWrapper): if "builtin_tools" not in kwargs and bool(kwargs.get("use_builtin_tools", False)): builtin_tools = "all" tools: list[ToolBase] = [] - tools.extend(self._builtin_tools(builtin_tools, sequential_tool_calls=sequential_tool_calls)) + tools.extend( + self._builtin_tools( + builtin_tools, + sequential_tool_calls=sequential_tool_calls, + call_kwargs=kwargs, + ), + ) tools.extend( self._make_tool( job, @@ -363,7 +375,7 @@ class AsAgentWrapper(BaseAgentWrapper): await agent.observe(inputs) last_msg = await agent.reply() usage = self._agentscope_usage(last_msg.usage) if last_msg.usage is not None else None - await self._dump_state(agent.state) + await self._dump_state(agent.state, kwargs) result = { "session_id": agent.state.session_id, @@ -404,7 +416,7 @@ class AsAgentWrapper(BaseAgentWrapper): agent, _ = await self._build_agent(None, **kwargs) config = {**(kwargs.get("context_config") or {}), "trigger_ratio": 1e-9} await agent.compress_context(ContextConfig(**config)) - await self._dump_state(agent.state) + await self._dump_state(agent.state, kwargs) # ----- StreamChunk conversion ------------------------------------------- @@ -509,4 +521,4 @@ class AsAgentWrapper(BaseAgentWrapper): chunk.session_id = chunk.session_id or agent.state.session_id yield chunk - await self._dump_state(agent.state) + await self._dump_state(agent.state, kwargs) diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 167c1ec3..4017b481 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -8,7 +8,7 @@ from typing import Any, ClassVar, TYPE_CHECKING from pydantic import BaseModel -from ..base_component import BaseComponent +from ..base_component import BaseComponent, _context_workspaces from ..outbound_proxy import BaseOutboundProxy from ...enumeration import ChunkEnum, ComponentEnum from ...schema import StreamChunk, TokenUsage @@ -58,6 +58,25 @@ class BaseAgentWrapper(BaseComponent): cwd = Path(self._cwd) return cwd if cwd.is_absolute() else (self.workspace_path / cwd) + def _call_workspace_path(self, kwargs: dict[str, Any]) -> Path: + """Workspace root for this call: the forwarded binding, else the default workspace.""" + alias = str(kwargs.get("workspace") or "") + if not alias: + return self.workspace_path + workspaces = _context_workspaces(self.app_context) + if workspaces is None or alias not in workspaces: + configured = sorted(workspaces) if workspaces else [] + raise ValueError(f"Unknown workspace {alias!r}; configured: {configured}") + return workspaces[alias].workspace_path + + def _call_cwd(self, kwargs: dict[str, Any]) -> Path: + """Working directory for this call: explicit cwd config wins, then the bound workspace.""" + if self._cwd: + return self.cwd + if kwargs.get("workspace"): + return self._call_workspace_path(kwargs) + return self.cwd + def set_system_prompt(self, prompt: str) -> "BaseAgentWrapper": """Set the agent's system prompt. Returns self for chaining.""" self.kwargs["system_prompt"] = prompt diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index f2e460a4..d799fa81 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -43,10 +43,15 @@ class CcAgentWrapper(BaseAgentWrapper): @property def session_path(self) -> Path: - """Directory used for persisted Claude Code sessions.""" + """Directory used for persisted Claude Code sessions in the default workspace.""" + return self._session_path_for({}) + + def _session_path_for(self, kwargs) -> Path: + """Session directory under the workspace this call is bound to.""" + root = self._call_workspace_path(kwargs) if self.app_context is None: - return self.workspace_path / "mem_session" - return self.workspace_path / self.app_context.app_config.mem_session_dir + return root / "mem_session" + return root / self.app_context.app_config.mem_session_dir def _ensure_claude_skill_dir(self, config_dir: Path, skills: list[str] | str) -> None: """Add selected project skills to Claude Code discovery locations.""" @@ -192,9 +197,10 @@ class CcAgentWrapper(BaseAgentWrapper): ) opts.env.update(extra_env_dict) self._add_bash_proxy_hook(opts) - self.session_path.mkdir(parents=True, exist_ok=True) - opts.cwd = opts.cwd or self.cwd - claude_config_dir = self.session_path / "claude_config" + session_path = self._session_path_for(kwargs) + session_path.mkdir(parents=True, exist_ok=True) + opts.cwd = opts.cwd or self._call_cwd(kwargs) + claude_config_dir = session_path / "claude_config" opts.env.setdefault("CLAUDE_CONFIG_DIR", str(claude_config_dir)) if selected_skills is not None: self._ensure_claude_skill_dir(claude_config_dir, selected_skills) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index ada1103e..15e3e420 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -179,7 +179,15 @@ class CodexAgentWrapper(BaseAgentWrapper): 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) + payload = self.app_context.app_config.model_dump(mode="json") + workspaces = _context_workspaces(self.app_context) + if workspaces and len(workspaces) > 1: + # The bridge rebuilds the full multi-workspace config; + # the serialized workspace_dir string alone would lose it. + payload["workspace_dir"] = { + alias: str(workspace.workspace_path) for alias, workspace in workspaces.items() + } + json.dump(payload, stream) except BaseException: with suppress(OSError): os.close(fd) @@ -241,17 +249,6 @@ class CodexAgentWrapper(BaseAgentWrapper): names = ", ".join(invalid) raise TypeError(f"Codex client options must be configured on the wrapper: {names}") - def _call_workspace_path(self, kwargs: dict[str, Any]) -> Path: - """Resolve the workspace binding forwarded by the calling step to a path.""" - alias = str(kwargs.get("workspace") or "") - if not alias: - return self.workspace_path - workspaces = _context_workspaces(self.app_context) - if workspaces is None or alias not in workspaces: - configured = sorted(workspaces) if workspaces else [] - raise ValueError(f"Unknown workspace {alias!r}; configured: {configured}") - return workspaces[alias].workspace_path - def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: from ..job import BackgroundJob, StreamJob @@ -263,6 +260,7 @@ class CodexAgentWrapper(BaseAgentWrapper): if unsupported: raise TypeError(f"Codex job_tools must be non-stream request jobs: {', '.join(unsupported)}") + alias = str(kwargs.get("workspace") or "") config_source = self._mcp_config_source(kwargs) args = [ "-m", @@ -272,6 +270,9 @@ class CodexAgentWrapper(BaseAgentWrapper): "--workspace", str(self._call_workspace_path(kwargs)), ] + if alias: + # The child builds the full multi-workspace config and binds via the alias. + args.extend(["--workspace-alias", alias]) for name in job_names: args.extend(["--job", name]) args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) @@ -344,7 +345,7 @@ class CodexAgentWrapper(BaseAgentWrapper): "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.auto_review), "base_instructions": kwargs.get("base_instructions"), "config": self._thread_config(kwargs), - "cwd": str(self.cwd), + "cwd": str(self._call_cwd(kwargs)), "developer_instructions": kwargs.get("system_prompt"), "model": kwargs.get("model"), "model_provider": kwargs.get("model_provider"), @@ -380,7 +381,7 @@ class CodexAgentWrapper(BaseAgentWrapper): return { "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), - "cwd": str(self.cwd), + "cwd": str(self._call_cwd(kwargs)), "effort": self._enum(ReasoningEffort, kwargs.get("effort")), "model": kwargs.get("model"), "output_schema": kwargs.get("output_schema"), diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index 9ed86f5b..cedd428a 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -13,6 +13,12 @@ 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( + "--workspace-alias", + default="", + help="Alias selecting one workspace of a multi-workspace config snapshot; " + "without it, --workspace overrides the config with a single workspace (legacy behavior)", + ) 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( @@ -66,17 +72,32 @@ def _prepare_config( def main() -> None: """Load ReMe and serve the requested jobs over STDIO.""" args = _parse_args() + overrides: dict[str, Any] = {} + if not args.workspace_alias: + # Legacy single-workspace mode: the CLI path replaces the config value. + overrides["workspace_dir"] = str(Path(args.workspace).absolute()) 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, + **overrides, ) + if args.workspace_alias: + raw_workspaces = config.get("workspace_dir") + if not isinstance(raw_workspaces, dict) or args.workspace_alias not in raw_workspaces: + raise ValueError( + f"--workspace-alias {args.workspace_alias!r} requires a config snapshot " + "carrying the full workspace mapping", + ) 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") + injected_job_kwargs = dict(injected_job_kwargs or {}) + if args.workspace_alias: + # The binding is server-owned; nested job calls inherit it from here. + injected_job_kwargs.setdefault("workspace", args.workspace_alias) config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) ReMe(**config).run_app() diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index f54b465a..8a3cc9c8 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -212,31 +212,43 @@ class BaseStep(ComponentMixin, ABC): return await job(**kwargs) def _bound_workspace_alias(self) -> str: - """Alias of the current call's workspace binding. + """Alias of the current call's workspace binding ("" only when unbound). - Returns "" when unbound or bound to the default workspace: the default - needs no forwarding because every resolution already lands there, and - single-workspace behavior stays byte-identical. + Execution bindings are explicit end to end: the default alias is a + real binding and is forwarded like any other, so a nested job's own + configured workspace cannot silently override the caller's binding. """ bound = getattr(self.context, "workspace", None) if self.context is not None else None - if bound is None: - return "" + return bound.alias if bound is not None else "" + + def _is_default_workspace_alias(self, alias: str) -> bool: + """Whether the alias names the first (default) configured workspace.""" workspaces = _context_workspaces(self.app_context) - if workspaces and bound.alias == next(iter(workspaces)): - return "" - return bound.alias + return bool(workspaces) and alias == next(iter(workspaces)) def _workspace_forward_kwargs(self) -> dict: - """Forward the current call's workspace binding to agent wrapper calls.""" - alias = self._bound_workspace_alias() - return {"workspace": alias} if alias else {} + """Forward the current call's workspace binding to agent wrapper calls. - def _scoped_state_key(self, key: str) -> str: - """Namespace a shared in-memory state key with the bound workspace alias.""" + The default workspace needs no forwarding because wrappers already + resolve there; explicit non-default bindings are always forwarded. + """ alias = self._bound_workspace_alias() - if not alias or not key: + if not alias or self._is_default_workspace_alias(alias): + return {} + return {"workspace": alias} + + def _scoped_state_key(self, key: str): + """Namespace a shared in-memory state key with the bound workspace alias. + + The default workspace keeps its legacy plain key (omitting the binding + and selecting the default alias share state by design); other aliases + use an unambiguous ``(alias, key)`` tuple so ids containing separators + of another workspace cannot collide with it. + """ + alias = self._bound_workspace_alias() + if not alias or not key or self._is_default_workspace_alias(alias): return key - return f"{alias}/{key}" + return (alias, key) def _resolve_dispatch_step(self, raw: _DispatchStep): """Resolve a dispatch step spec to (step class, init params).""" diff --git a/tests/unit/test_base_agent_wrapper.py b/tests/unit/test_base_agent_wrapper.py index 13235156..d1d3baa6 100644 --- a/tests/unit/test_base_agent_wrapper.py +++ b/tests/unit/test_base_agent_wrapper.py @@ -182,7 +182,7 @@ async def test_agentscope_builtin_tools_are_opt_in(tmp_path, monkeypatch, reply_ wrapper.as_llm = SimpleNamespace(model=object()) observed = {} - def builtin_tools(names, *, sequential_tool_calls=False): + def builtin_tools(names, *, sequential_tool_calls=False, **_extra): observed["names"] = names observed["sequential_tool_calls"] = sequential_tool_calls return [] @@ -251,7 +251,7 @@ async def test_agentscope_structured_output_uses_model_tool_choice_policy(tmp_pa """Return the fake AgentScope agent.""" return FakeAgent(), inputs - async def dump_state(_state): + async def dump_state(_state, _kwargs=None): """Skip session persistence.""" return None diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index c123faa7..a0db6ebc 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -379,7 +379,7 @@ async def test_agentscope_compact_session_forces_and_persists_compression(tmp_pa observed["build"] = (inputs, kwargs) return FakeAgent(), inputs - async def dump_state(state): + async def dump_state(state, _kwargs=None): observed["state"] = state monkeypatch.setattr(wrapper, "_build_agent", build_agent) diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index cd505815..519f417f 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -216,7 +216,7 @@ async def test_agentscope_stream_reply_does_not_record_token_usage(tmp_path, mon """Build the minimal stream double.""" return FakeAgent(), inputs - async def dump_state(_state): + async def dump_state(_state, _kwargs=None): """Avoid durable state writes in this accounting test.""" return None diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index 8ecd1538..0e1f3c2a 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -3,6 +3,8 @@ # pylint: disable=protected-access,missing-function-docstring import asyncio +from pathlib import Path +from types import SimpleNamespace import pytest @@ -111,20 +113,23 @@ async def test_tool_context_state_namespaced_per_workspace(tmp_path): context = _context(tmp_path) step = AddDraftStep(app_context=context) for alias in ("alpha", "beta"): - step.context = RuntimeContext(workspace=context.workspaces[alias], tool_context_id="ctx-1", text=f"draft-{alias}") + step.context = RuntimeContext( + workspace=context.workspaces[alias], + tool_context_id="ctx-1", + text=f"draft-{alias}", + ) await step.execute() stores = context.metadata["tool_contexts"] - # The default workspace keeps the legacy plain key; other aliases are namespaced. - assert sorted(stores) == ["beta/ctx-1", "ctx-1"] + # The default workspace keeps the legacy plain key; other aliases use + # unambiguous tuple keys. + assert set(stores) == {"ctx-1", ("beta", "ctx-1")} assert stores["ctx-1"]["drafts"] == ["draft-alpha"] - assert stores["beta/ctx-1"]["drafts"] == ["draft-beta"] + assert stores[("beta", "ctx-1")]["drafts"] == ["draft-beta"] @pytest.mark.asyncio async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): - from types import SimpleNamespace - wrapper = CodexAgentWrapper( app_context=ApplicationContext( workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, @@ -132,6 +137,8 @@ async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): ) class FakeCodex: + """Minimal Codex double returning a fixed thread id.""" + async def thread_start(self, **_kwargs): return SimpleNamespace(id="thread-1") @@ -229,8 +236,6 @@ async def test_per_alias_watch_variant_jobs_write_own_workspace(tmp_path): @pytest.mark.asyncio async def test_mcp_tool_boundary_forwards_injected_workspace(tmp_path): - from types import SimpleNamespace - from reme.components.service.mcp_tools import add_mcp_job context = _context(tmp_path) @@ -267,3 +272,229 @@ async def test_codex_bridge_uses_bound_workspace_path(tmp_path): wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "nope"}) await wrapper.close() + + +# ----- Review follow-ups: explicit bindings, session isolation, key collisions + + +@pytest.mark.asyncio +async def test_nested_call_default_binding_beats_child_config(tmp_path): + """An explicit default binding propagates; the child's configured alias must not win.""" + context = _context(tmp_path) + parent = _job(context, "parent", backend="nested_caller_step") + child = _job(context, "report", workspace="beta") # job-level configured binding + context.jobs = {"parent": parent, "report": child} + await parent.start() + await child.start() + + # Top-level call to the configured child still honors its configuration. + assert (await child()).answer == "beta" + # Nested under an alpha-bound parent, the parent binding wins. + response = await parent(workspace="alpha") + assert response.answer == "parent->alpha" + response = await parent(workspace="beta") + assert response.answer == "parent->beta" + + +@pytest.mark.asyncio +async def test_tool_context_keys_cannot_collide_across_workspaces(tmp_path): + context = _context(tmp_path) + step = AddDraftStep(app_context=context) + # Default-workspace id that spells like another alias's namespaced key. + step.context = RuntimeContext(workspace=context.workspaces["alpha"], tool_context_id="beta/ctx-1", text="d-alpha") + await step.execute() + step.context = RuntimeContext(workspace=context.workspaces["beta"], tool_context_id="ctx-1", text="d-beta") + await step.execute() + + stores = context.metadata["tool_contexts"] + assert stores["beta/ctx-1"]["drafts"] == ["d-alpha"] + assert stores[("beta", "ctx-1")]["drafts"] == ["d-beta"] + + +@pytest.mark.asyncio +async def test_agentscope_sessions_isolated_per_workspace(tmp_path): + from uuid import uuid4 + + from agentscope.permission import PermissionMode + + from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper + + context = _context(tmp_path) + wrapper = AsAgentWrapper(app_context=context, as_llm="") + wrapper.as_llm = SimpleNamespace(model=object()) + + session_id = str(uuid4()) + state = await wrapper._load_state({"session_id": session_id, "workspace": "beta"}, PermissionMode.DEFAULT) + state.summary = "beta summary" + await wrapper._dump_state(state, {"workspace": "beta"}) + + session_file = tmp_path / "beta" / "mem_session" / "agentscope" / f"{session_id}.jsonl" + assert session_file.is_file() + assert not (tmp_path / "alpha" / "mem_session" / "agentscope" / f"{session_id}.jsonl").exists() + + # Same id under another workspace must not read beta's state. + with pytest.raises(FileNotFoundError): + await wrapper._load_state({"resume": session_id, "workspace": "alpha"}, PermissionMode.DEFAULT) + resumed = await wrapper._load_state({"resume": session_id, "workspace": "beta"}, PermissionMode.DEFAULT) + assert resumed.summary == "beta summary" + # Default workspace is untouched by the beta session. + with pytest.raises(FileNotFoundError): + await wrapper._load_state({"resume": session_id}, PermissionMode.DEFAULT) + + +def test_cc_options_follow_bound_workspace(tmp_path): + from reme.components.agent_wrapper.cc_agent_wrapper import CcAgentWrapper + + context = _context(tmp_path) + wrapper = CcAgentWrapper(app_context=context) + + opts = wrapper._build_options("hi", workspace="beta") + assert opts.env["CLAUDE_CONFIG_DIR"].startswith(str(tmp_path / "beta")) + assert opts.cwd == tmp_path / "beta" + + default_opts = wrapper._build_options("hi") + assert default_opts.env["CLAUDE_CONFIG_DIR"].startswith(str(tmp_path / "alpha")) + + +@pytest.mark.asyncio +async def test_codex_thread_cwd_follows_bound_workspace(tmp_path): + context = _context(tmp_path) + wrapper = CodexAgentWrapper(app_context=context) + observed = {} + + class FakeCodex: + """Capture thread_start parameters without an app-server.""" + + async def thread_start(self, **kwargs): + observed.update(kwargs) + return SimpleNamespace(id="thread-1") + + await wrapper._open_thread(FakeCodex(), {"workspace": "beta"}) + assert observed["cwd"] == str(tmp_path / "beta") + + observed.clear() + await wrapper._open_thread(FakeCodex(), {}) + assert observed["cwd"] != str(tmp_path / "beta") + + +# ----- Codex bridge: real child-process execution ---------------------------- + + +@pytest.mark.asyncio +async def test_codex_bridge_executes_non_default_alias(tmp_path): + """The bridge child rebuilds the full workspace mapping and binds via the alias.""" + import json + import sys + + from fastmcp import Client + from fastmcp.client import StdioTransport + + snapshot = tmp_path / "bridge.json" + snapshot.write_text( + json.dumps( + { + "workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "mcp"}, + "components": { + "tokenizer": {"default": {"backend": "regex"}}, + "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, + "keyword_index": {"default": {"backend": "bm25"}}, + "file_graph": {"default": {"backend": "local"}}, + }, + "jobs": { + "write_note": { + "backend": "base", + "description": "write a note", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, + "required": ["path", "content"], + }, + "steps": [{"backend": "write_step"}], + }, + }, + }, + ), + encoding="utf-8", + ) + transport = StdioTransport( + command=sys.executable, + args=[ + "-m", + "reme.components.agent_wrapper.codex_mcp_server", + "--config", + str(snapshot), + "--workspace", + str(tmp_path / "beta"), + "--workspace-alias", + "beta", + "--job", + "write_note", + ], + cwd=str(Path(__file__).resolve().parents[2]), + keep_alive=False, + ) + + async with Client(transport, timeout=None, init_timeout=30) as client: + await asyncio.wait_for( + client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-beta"}), + timeout=30, + ) + + assert "via-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") + assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() + + +@pytest.mark.asyncio +async def test_codex_bridge_rejects_alias_without_mapping_snapshot(tmp_path): + import json + import sys + + from fastmcp import Client + from fastmcp.client import StdioTransport + + snapshot = tmp_path / "bridge-single.json" + snapshot.write_text( + json.dumps( + { + "workspace_dir": str(tmp_path / "only"), + "enable_logo": False, + "log_to_console": False, + "log_to_file": False, + "service": {"backend": "mcp"}, + "jobs": { + "write_note": { + "backend": "base", + "parameters": {"type": "object", "properties": {}}, + "steps": [{"backend": "write_step"}], + }, + }, + }, + ), + encoding="utf-8", + ) + transport = StdioTransport( + command=sys.executable, + args=[ + "-m", + "reme.components.agent_wrapper.codex_mcp_server", + "--config", + str(snapshot), + "--workspace", + str(tmp_path / "only"), + "--workspace-alias", + "beta", + "--job", + "write_note", + ], + cwd=str(Path(__file__).resolve().parents[2]), + keep_alive=False, + ) + + # The child process fails fast with a ValueError; the client sees the closed connection. + with pytest.raises(Exception, match="Connection closed"): + async with Client(transport, timeout=None, init_timeout=30): + pass From 4bc6bdc0350355be762b7ef64a12a14ec0c07d7d Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 22 Oct 2026 17:25:07 +0800 Subject: [PATCH 07/21] refactor(workspace): move scoped state ownership into WorkspaceContext Commit A of the final review round, restructuring context ownership per the approved target shape: ApplicationContext coordinates, each WorkspaceContext owns everything scoped to its workspace, and RuntimeContext carries the resolved binding. - Workspace-scoped component instances move from the ApplicationContext workspace_components table into each WorkspaceContext (which keeps a compatibility view for transitional readers); startup ordering, reverse shutdown, rollback, and replace_component traverse WorkspaceContext-owned instances with identical semantics. - Drafts, search dedup state, search call budgets, and the dream integration lock move from manually namespaced application metadata into WorkspaceContext-owned state and locks, so identical tool_context_ids in different workspaces stay independent without key encoding. Standalone/test contexts keep the legacy metadata buckets. - File chunkers become workspace-scoped: their node paths are workspace-relative, so a shared chunker mis-relativized non-default workspace files; the update steps resolve the bound alias's chunker. - Execution bindings stay explicit end to end: agent wrapper calls now forward the resolved alias for every workspace including the default, closing the hole where a default-bound caller's nested tool call could land on a child job's configured alias. Session UUIDs keep their historical unprefixed form for the default workspace. - Wrapper cwd resolution: a relative configured cwd resolves against the bound workspace; absolute paths keep their explicit meaning. Refs #632 --- reme/application.py | 36 ++-- .../agent_wrapper/base_agent_wrapper.py | 11 +- reme/components/application_context.py | 30 +++- reme/components/base_component.py | 8 +- .../file_chunker/base_file_chunker.py | 7 +- reme/schema/workspace.py | 32 +++- reme/steps/base_step.py | 28 +--- reme/steps/evolve/base_auto_resource.py | 9 +- reme/steps/evolve/dream/integrate.py | 9 +- reme/steps/index/_tool_context.py | 5 +- reme/steps/index/search.py | 8 +- reme/steps/index/update_changes.py | 8 +- tests/unit/test_auto_resource_agent_inputs.py | 4 +- tests/unit/test_workspace_binding.py | 2 +- tests/unit/test_workspace_ownership.py | 154 ++++++++++++++++++ tests/unit/test_workspace_phase3.py | 15 +- 16 files changed, 285 insertions(+), 81 deletions(-) create mode 100644 tests/unit/test_workspace_ownership.py diff --git a/reme/application.py b/reme/application.py index c49b8c21..0b4d4b68 100644 --- a/reme/application.py +++ b/reme/application.py @@ -82,8 +82,8 @@ class Application(BaseComponent): f"{label} backend '{cfg.backend}' does not support multi-workspace configurations", ) if backend_cls.workspace_scoped: - instances = { - alias: self._instantiate( + for alias, workspace in self.context.workspaces.items(): + instance = self._instantiate( ctype, cfg, label=label, @@ -91,11 +91,11 @@ class Application(BaseComponent): name=name, workspace_alias=alias, ) - for alias in self.context.workspaces - } + workspace.components.setdefault(ctype, {})[name] = instance # The default alias instance doubles as the plain lookup entry. - self.context.components[ctype][name] = instances[next(iter(instances))] - self.context.workspace_components.setdefault(ctype, {})[name] = instances + self.context.components[ctype][name] = next(iter(self.context.workspaces.values())).components[ + ctype + ][name] else: self.context.components[ctype][name] = self._instantiate( ctype, @@ -166,13 +166,15 @@ class Application(BaseComponent): ) -> dict[_NodeKey, BaseComponent]: """One node per component instance; scoped components contribute one node per alias.""" nodes: dict[_NodeKey, BaseComponent] = {} + scoped_keys: set[tuple[str, str]] = set() + for alias, workspace in self.context.workspaces.items(): + for ctype, group in workspace.components.items(): + for name, instance in group.items(): + nodes[(ctype, name, alias)] = instance + scoped_keys.add((ctype, name)) for ctype, group in self.context.components.items(): for name, comp in group.items(): - scoped = self.context.workspace_components.get(ctype, {}).get(name) - if scoped: - for alias, instance in scoped.items(): - nodes[(ctype, name, alias)] = instance - else: + if (ctype, name) not in scoped_keys: nodes[(ctype, name, "")] = comp if replacement is not None: nodes[replacement[0]] = replacement[1] @@ -311,13 +313,13 @@ class Application(BaseComponent): def _all_component_instances(self) -> list[BaseComponent]: """Every component instance: plain entries plus all alias instances of scoped components.""" instances: dict[int, BaseComponent] = {} + for workspace in self.context.workspaces.values(): + for group in workspace.components.values(): + for instance in group.values(): + instances[id(instance)] = instance for group in self.context.components.values(): for component in group.values(): instances[id(component)] = component - for group in self.context.workspace_components.values(): - for scoped in group.values(): - for instance in scoped.values(): - instances[id(instance)] = instance return list(instances.values()) def _replacement_shutdown_order( @@ -427,8 +429,8 @@ class Application(BaseComponent): # Commit the new generation synchronously so observers cannot see # a context with only some dependency references updated. group[name] = replacement - if scoped_instances is not None: - scoped_instances[alias] = replacement + if alias is not None: + self.context.workspaces[alias].components[component_type][name] = replacement config_group[name] = replacement_config for consumer, attr in consumers: consumer.__dict__[attr] = replacement diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 4017b481..48ff158f 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -70,9 +70,16 @@ class BaseAgentWrapper(BaseComponent): return workspaces[alias].workspace_path def _call_cwd(self, kwargs: dict[str, Any]) -> Path: - """Working directory for this call: explicit cwd config wins, then the bound workspace.""" + """Working directory for this call: explicit cwd config wins, then the bound workspace. + + A relative configured cwd resolves against the bound workspace; + an absolute path keeps its explicit meaning. + """ if self._cwd: - return self.cwd + cwd = Path(self._cwd) + if cwd.is_absolute(): + return cwd + return self._call_workspace_path(kwargs) / cwd if kwargs.get("workspace"): return self._call_workspace_path(kwargs) return self.cwd diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 0aff4d1a..e52475b6 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -35,10 +35,6 @@ class ApplicationContext: # Populated by Application during initialization. self.service: "BaseService | None" = None self.components: dict[str, dict[str, "BaseComponent"]] = {} - # Per-alias instances of workspace-scoped components, keyed - # ctype -> name -> alias. The default alias instance is also the one - # registered under ``components`` for backward-compatible lookups. - self.workspace_components: dict[str, dict[str, dict[str, "BaseComponent"]]] = {} self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None @@ -48,6 +44,20 @@ class ApplicationContext: # must survive an Application restart. self.metadata: dict[str, Any] = {} + @property + def workspace_components(self) -> dict[str, dict[str, dict[str, "BaseComponent"]]]: + """Compatibility view of scoped instances, ctype -> name -> alias. + + The authoritative owner of these instances is each WorkspaceContext; + this view exists for transitional readers. + """ + grouped: dict[str, dict[str, dict[str, "BaseComponent"]]] = {} + for alias, workspace in self.workspaces.items(): + for ctype, group in workspace.components.items(): + for name, instance in group.items(): + grouped.setdefault(ctype, {}).setdefault(name, {})[alias] = instance + return grouped + def get_component(self, ctype: str, name: str, alias: str | None = None) -> "BaseComponent | None": """Resolve a component, honoring workspace bindings for scoped components. @@ -55,10 +65,12 @@ class ApplicationContext: KeyError when a workspace-scoped component has no instance for the requested alias — a missing scoped resource must fail, never fall back. """ - scoped = self.workspace_components.get(ctype, {}).get(name) - if scoped is not None: - alias = alias or next(iter(self.workspaces)) - if alias in scoped: - return scoped[alias] + alias = alias or next(iter(self.workspaces)) + workspace = self.workspaces.get(alias) + if workspace is None: + raise KeyError(f"Unknown workspace alias {alias!r}") + if name in workspace.components.get(ctype, {}): + return workspace.components[ctype][name] + if any(name in ws.components.get(ctype, {}) for ws in self.workspaces.values()): raise KeyError(f"Component '{name}' in {ctype} has no instance for workspace alias {alias!r}") return self.components.get(ctype, {}).get(name) diff --git a/reme/components/base_component.py b/reme/components/base_component.py index 51f32717..1f704bfb 100644 --- a/reme/components/base_component.py +++ b/reme/components/base_component.py @@ -226,9 +226,11 @@ class BaseComponent(ComponentMixin, ABC): Workspace-scoped dependencies resolve to the same-alias instance as the dependent; a missing alias instance is an error, never a fallback. """ - scoped = getattr(self.app_context, "workspace_components", {}).get(dep.ctype, {}).get(dep.name) - if scoped is not None and self.workspace_alias is not None: - target = scoped.get(self.workspace_alias) + workspaces = getattr(self.app_context, "workspaces", None) + workspace = workspaces.get(self.workspace_alias) if workspaces and self.workspace_alias is not None else None + scoped_group = workspace.components.get(dep.ctype) if workspace is not None else None + if scoped_group is not None: + target = scoped_group.get(dep.name) if target is None: raise ValueError(f"{dep.ctype} '{dep.name}' has no instance for workspace '{self.workspace_alias}'.") else: diff --git a/reme/components/file_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 49312989..5633fa5d 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -9,9 +9,14 @@ from ...schema import FileChunk, FileNode class BaseFileChunker(BaseComponent): - """Abstract base for file chunkers. Subclasses implement `chunk`.""" + """Abstract base for file chunkers. Subclasses implement `chunk`. + + Chunkers are workspace-scoped: node paths are workspace-relative, so each + alias gets its own instance rooted at that workspace. + """ component_type = ComponentEnum.FILE_CHUNKER + workspace_scoped = True def __init__(self, supported_extensions: list[str] | None = None, **kwargs): super().__init__(**kwargs) diff --git a/reme/schema/workspace.py b/reme/schema/workspace.py index 921b7e67..f30cb92b 100644 --- a/reme/schema/workspace.py +++ b/reme/schema/workspace.py @@ -1,8 +1,9 @@ -"""Per-workspace path views for multi-workspace applications.""" +"""Per-workspace path views and owned state for multi-workspace applications.""" -from dataclasses import dataclass +import asyncio +from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .application_config import ApplicationConfig @@ -10,11 +11,14 @@ if TYPE_CHECKING: @dataclass(frozen=True) class WorkspaceContext: - """Resolved paths for one named workspace of an application. + """Owns everything scoped to one workspace of an application. - Built from ApplicationConfig when the application context is created. The - configured subdirectory names stay workspace-relative, matching the layout - a single-workspace application derives from the same config. + Built from ApplicationConfig when the application context is created. + Beyond the resolved paths, the workspace owns its stateful component + instances (stores, graphs, indexes, catalogs), in-memory execution state + (drafts, search dedup, call budgets), session ownership records, and + workspace-scoped locks. Shared, stateless clients stay on the + ApplicationContext; per-call bindings reference this object. """ alias: str @@ -25,6 +29,16 @@ class WorkspaceContext: resource_path: Path | None daily_path: Path | None digest_path: Path | None + # Workspace-scoped component instances keyed by component type then name. + # Values are BaseComponent instances; typed as Any to keep this module + # free of component imports (schema must not depend on components). + components: dict[str, dict[str, Any]] = field(default_factory=dict) + # Drafts and search dedup state, keyed by tool_context_id. + tool_contexts: dict[str, Any] = field(default_factory=dict) + # Search call budgets keyed by tool_context_id. + search_call_budgets: dict[str, Any] = field(default_factory=dict) + # Workspace-scoped named asyncio locks (e.g. dream integration). + locks: dict[str, asyncio.Lock] = field(default_factory=dict) @classmethod def from_config(cls, alias: str, workspace: str, config: "ApplicationConfig") -> "WorkspaceContext": @@ -59,3 +73,7 @@ class WorkspaceContext: ) if path is not None ] + + def lock(self, name: str) -> asyncio.Lock: + """Return the workspace-scoped lock for ``name``, created lazily.""" + return self.locks.setdefault(name, asyncio.Lock()) diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index 8a3cc9c8..15702fad 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -221,34 +221,14 @@ class BaseStep(ComponentMixin, ABC): bound = getattr(self.context, "workspace", None) if self.context is not None else None return bound.alias if bound is not None else "" - def _is_default_workspace_alias(self, alias: str) -> bool: - """Whether the alias names the first (default) configured workspace.""" - workspaces = _context_workspaces(self.app_context) - return bool(workspaces) and alias == next(iter(workspaces)) - def _workspace_forward_kwargs(self) -> dict: - """Forward the current call's workspace binding to agent wrapper calls. + """Forward the current call's resolved workspace binding to agent wrapper calls. - The default workspace needs no forwarding because wrappers already - resolve there; explicit non-default bindings are always forwarded. + Every resolved binding is forwarded, including the default alias, so + agent tool paths execute against the caller's workspace. """ alias = self._bound_workspace_alias() - if not alias or self._is_default_workspace_alias(alias): - return {} - return {"workspace": alias} - - def _scoped_state_key(self, key: str): - """Namespace a shared in-memory state key with the bound workspace alias. - - The default workspace keeps its legacy plain key (omitting the binding - and selecting the default alias share state by design); other aliases - use an unambiguous ``(alias, key)`` tuple so ids containing separators - of another workspace cannot collide with it. - """ - alias = self._bound_workspace_alias() - if not alias or not key or self._is_default_workspace_alias(alias): - return key - return (alias, key) + return {"workspace": alias} if alias else {} def _resolve_dispatch_step(self, raw: _DispatchStep): """Resolve a dispatch step spec to (step class, init params).""" diff --git a/reme/steps/evolve/base_auto_resource.py b/reme/steps/evolve/base_auto_resource.py index fa1d74ad..9f9791c8 100644 --- a/reme/steps/evolve/base_auto_resource.py +++ b/reme/steps/evolve/base_auto_resource.py @@ -15,6 +15,7 @@ import frontmatter from agentscope.message import DataBlock, TextBlock, UserMsg from watchfiles import Change +from ...components.base_component import _context_workspaces from ...components.runtime_context import RuntimeContext from ..base_step import BaseStep from ..file_io import refresh_day_index, validate_filename_component @@ -404,9 +405,15 @@ class BaseAutoResourceStep(BaseStep): **self._workspace_forward_kwargs(), } if "session_id" not in agent_kwargs: + # Storage compatibility: the default workspace keeps its historical + # (unprefixed) session UUIDs; other aliases are namespaced. + session_alias = agent_kwargs.get("workspace", "") + workspaces = _context_workspaces(self.app_context) + if workspaces and session_alias == next(iter(workspaces)): + session_alias = "" agent_kwargs["session_id"] = _compute_agent_session_id( file_path, - workspace_alias=agent_kwargs.get("workspace", ""), + workspace_alias=session_alias, ) # Consume the resource-only option before forwarding kwargs to the wrapper. if agent_kwargs.pop("scope_note_tools", False): diff --git a/reme/steps/evolve/dream/integrate.py b/reme/steps/evolve/dream/integrate.py index 9bcdd94b..441dc7e6 100644 --- a/reme/steps/evolve/dream/integrate.py +++ b/reme/steps/evolve/dream/integrate.py @@ -100,9 +100,16 @@ class DreamIntegrateStep(BaseStep): return self._finish(state, not state.failed_units, answer) def _integration_lock(self) -> asyncio.Lock: - """Return the Application-wide lock used by AutoDream integration.""" + """Return the workspace-scoped lock used by AutoDream integration. + + Bound calls lock per workspace so integrations in different workspaces + never block each other; standalone contexts keep the app-wide lock. + """ if self.app_context is None: raise RuntimeError("dream_integrate_step requires app_context") + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if bound is not None: + return bound.lock("dream_integration") key = "dream_integration_lock" lock = self.app_context.metadata.get(key) if lock is None: diff --git a/reme/steps/index/_tool_context.py b/reme/steps/index/_tool_context.py index 993f47f1..0d9aba3d 100644 --- a/reme/steps/index/_tool_context.py +++ b/reme/steps/index/_tool_context.py @@ -16,8 +16,11 @@ class _ToolContextMixin: kwargs: dict[str, Any] def _tool_context_store(self, tool_context_id: str) -> dict: + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if bound is not None: + return bound.tool_contexts.setdefault(tool_context_id, {}) if self.app_context is not None: contexts = self.app_context.metadata.setdefault(self.TOOL_CONTEXTS_KEY, {}) else: contexts = self.kwargs.setdefault(self.TOOL_CONTEXTS_KEY, {}) - return contexts.setdefault(self._scoped_state_key(tool_context_id), {}) + return contexts.setdefault(tool_context_id, {}) diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index eff8222c..9a5b11b4 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -237,8 +237,12 @@ class SearchStep(_ToolContextMixin, BaseStep): maximum = int(max_search_calls) if maximum < 1: raise ValueError("max_search_calls must be positive") - budgets = self.app_context.metadata.setdefault("__search_call_budgets", {}) - budget = budgets.setdefault(self._scoped_state_key(tool_context_id), {"count": 0, "lock": asyncio.Lock()}) + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if bound is not None: + budgets = bound.search_call_budgets + else: + budgets = self.app_context.metadata.setdefault("__search_call_budgets", {}) + budget = budgets.setdefault(tool_context_id, {"count": 0, "lock": asyncio.Lock()}) async with budget["lock"]: if budget["count"] >= maximum: self.logger.warning( diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 6ce461f0..7afa9fed 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -364,8 +364,12 @@ class UpdateIndexStep(ChangeApplyStep): return await chunker.chunk(path) def _resolve_chunker(self, path: Path) -> BaseFileChunker: - """Resolve a file chunker for a given path.""" - chunkers: dict[str, BaseFileChunker] = self.app_context.components[ComponentEnum.FILE_CHUNKER] + """Resolve a file chunker for a given path, bound to this call's workspace.""" + bound = getattr(self.context, "workspace", None) if self.context is not None else None + if bound is not None: + chunkers: dict[str, BaseFileChunker] = bound.components.get(ComponentEnum.FILE_CHUNKER, {}) + else: + chunkers = self.app_context.components[ComponentEnum.FILE_CHUNKER] suffix = path.suffix.lstrip(".").lower() for candidate in chunkers.values(): if suffix and suffix in {ext.lower().lstrip(".") for ext in candidate.supported_extensions}: diff --git a/tests/unit/test_auto_resource_agent_inputs.py b/tests/unit/test_auto_resource_agent_inputs.py index f1f07f04..7a8b906d 100644 --- a/tests/unit/test_auto_resource_agent_inputs.py +++ b/tests/unit/test_auto_resource_agent_inputs.py @@ -563,7 +563,9 @@ async def test_configured_wrapper_and_job_image_options( if suffix == "txt": assert response.answer == "ok" assert isinstance(reply.call_args.args[0], str) - assert set(reply.call_args.kwargs) == {"system_prompt", "job_tools", "session_id"} + assert set(reply.call_args.kwargs) == {"system_prompt", "job_tools", "session_id", "workspace"} + # Execution bindings are forwarded everywhere, including the default alias. + assert reply.call_args.kwargs["workspace"] == "default" elif expected_action == "added": note = frontmatter.load(tmp_path / result["path"]) assert note["source_resource"] == f"[[resource/2026-01-01/photo.{suffix}]]" diff --git a/tests/unit/test_workspace_binding.py b/tests/unit/test_workspace_binding.py index 97d39e09..f7122489 100644 --- a/tests/unit/test_workspace_binding.py +++ b/tests/unit/test_workspace_binding.py @@ -176,7 +176,7 @@ async def test_workspace_scoped_components_instantiate_per_alias(tmp_path): assert store_alpha.keyword_index is index_alpha assert store_beta.file_graph is graph_beta - with pytest.raises(KeyError, match="no instance for workspace alias 'nope'"): + with pytest.raises(KeyError, match="Unknown workspace alias 'nope'"): app.context.get_component(ComponentEnum.FILE_STORE, "default", "nope") diff --git a/tests/unit/test_workspace_ownership.py b/tests/unit/test_workspace_ownership.py new file mode 100644 index 00000000..d388673a --- /dev/null +++ b/tests/unit/test_workspace_ownership.py @@ -0,0 +1,154 @@ +"""Ownership restructure: WorkspaceContext owns state; bindings stay explicit.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio + +import pytest + +from reme.application import Application +from reme.components import ApplicationContext +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper +from reme.components.component_registry import R +from reme.components.job import BaseJob +from reme.components.runtime_context import RuntimeContext +from reme.enumeration import ComponentEnum +from reme.steps.base_step import BaseStep +from reme.steps.evolve.dream.integrate import DreamIntegrateStep + + +class _AliasReporterStep(BaseStep): + """Answer with the bound workspace alias.""" + + async def execute(self): + bound = self.context.workspace + self.context.response.answer = bound.alias if bound is not None else "" + + +def _context(tmp_path): + registry = R.copy() + registry.register(_AliasReporterStep, "alias_reporter_step") + return ApplicationContext( + registry=registry, + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + + +def _reporter_job(context, name, **kwargs): + return BaseJob( + name=name, + app_context=context, + parameters={"type": "object", "properties": {}}, + steps=[{"backend": "alias_reporter_step"}], + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_default_binding_reaches_real_agent_tool_call(tmp_path): + """A default-bound caller keeps its binding through a real FunctionTool call.""" + context = _context(tmp_path) + job = _reporter_job(context, "report", workspace="beta") # child configured for beta + job.step_specs = [job._resolve_step({"backend": "alias_reporter_step"})] + + # The wrapper tool boundary injects the caller's resolved (default) binding. + tool = AsAgentWrapper._make_tool(job, None, None, workspace="alpha") + chunk = await tool._func() + text = "".join(block.text for block in chunk.content if hasattr(block, "text")) + assert text == "alpha" + + # Model-supplied workspace arguments are rejected at the same boundary. + with pytest.raises(ValueError, match="cannot be provided by the model"): + await tool._func(workspace="beta") + + +def test_call_cwd_resolves_relative_against_bound_workspace(tmp_path): + context = _context(tmp_path) + absolute = tmp_path / "elsewhere" + + wrapper = CodexAgentWrapper(app_context=context, cwd="relative/sub") + assert wrapper._call_cwd({"workspace": "beta"}) == tmp_path / "beta" / "relative" / "sub" + assert wrapper._call_cwd({}) == tmp_path / "alpha" / "relative" / "sub" + + absolute_wrapper = CodexAgentWrapper(app_context=context, cwd=str(absolute)) + assert absolute_wrapper._call_cwd({"workspace": "beta"}) == absolute + + default_wrapper = CodexAgentWrapper(app_context=context) + assert default_wrapper._call_cwd({"workspace": "beta"}) == tmp_path / "beta" + assert default_wrapper._call_cwd({}) != tmp_path / "beta" # project path, not a workspace + + +def test_dream_integration_lock_is_workspace_scoped(tmp_path): + context = _context(tmp_path) + step = DreamIntegrateStep(app_context=context) + + step.context = RuntimeContext(workspace=context.workspaces["alpha"]) + lock_alpha = step._integration_lock() + step.context = RuntimeContext(workspace=context.workspaces["beta"]) + lock_beta = step._integration_lock() + step.context = RuntimeContext(workspace=context.workspaces["alpha"]) + assert step._integration_lock() is lock_alpha + assert lock_alpha is not lock_beta + + +# ----- Real watch loops ------------------------------------------------------- + +_WATCH_COMPONENTS = { + "tokenizer": {"default": {"backend": "regex"}}, + "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, + "keyword_index": {"default": {"backend": "bm25"}}, + "file_graph": {"default": {"backend": "local"}}, + "file_chunker": {"markdown": {"backend": "markdown", "supported_extensions": ["md"]}}, +} + + +def _watch_job(alias): + return { + "backend": "background", + "workspace": alias, + "enable_serve": False, + "watch_dirs": ["daily_dir"], + "watch_suffixes": ["md"], + "parameters": {"type": "object", "properties": {}}, + "steps": [ + { + "backend": "watch_changes_step", + "debounce": 200, + "step": 100, + "poll_delay_ms": 200, + "dispatch_steps": [{"backend": "update_index_step", "persist": False}], + }, + ], + } + + +@pytest.mark.asyncio +async def test_real_watch_loops_index_own_workspace_and_clean_up(tmp_path): + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components=_WATCH_COMPONENTS, + jobs={"watch_alpha": _watch_job("alpha"), "watch_beta": _watch_job("beta")}, + ) + note = tmp_path / "beta" / "daily" / "watch-me.md" + note.parent.mkdir(parents=True, exist_ok=True) + async with app: + store_beta = app.context.get_component(ComponentEnum.FILE_STORE, "default", "beta") + store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + for i in range(30): # up to ~15s for watcher startup + debounce + dispatch + note.write_text(f"# beta watch\nreal watch event content v{i}\n", encoding="utf-8") + await asyncio.sleep(0.5) + if await store_beta.get_nodes(): + break + nodes_beta = await store_beta.get_nodes() + nodes_alpha = await store_alpha.get_nodes() + + assert [n.path for n in nodes_beta] == ["daily/watch-me.md"] + assert nodes_alpha == [] + # Background loops stopped with the application; nothing keeps running. + for job in app.context.jobs.values(): + assert job._task is None or job._task.done() diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index 0e1f3c2a..a33f5399 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -120,12 +120,10 @@ async def test_tool_context_state_namespaced_per_workspace(tmp_path): ) await step.execute() - stores = context.metadata["tool_contexts"] - # The default workspace keeps the legacy plain key; other aliases use - # unambiguous tuple keys. - assert set(stores) == {"ctx-1", ("beta", "ctx-1")} - assert stores["ctx-1"]["drafts"] == ["draft-alpha"] - assert stores[("beta", "ctx-1")]["drafts"] == ["draft-beta"] + # State is owned by each WorkspaceContext; identical ids stay independent. + assert context.workspaces["alpha"].tool_contexts["ctx-1"]["drafts"] == ["draft-alpha"] + assert context.workspaces["beta"].tool_contexts["ctx-1"]["drafts"] == ["draft-beta"] + assert "tool_contexts" not in context.metadata @pytest.mark.asyncio @@ -306,9 +304,8 @@ async def test_tool_context_keys_cannot_collide_across_workspaces(tmp_path): step.context = RuntimeContext(workspace=context.workspaces["beta"], tool_context_id="ctx-1", text="d-beta") await step.execute() - stores = context.metadata["tool_contexts"] - assert stores["beta/ctx-1"]["drafts"] == ["d-alpha"] - assert stores[("beta", "ctx-1")]["drafts"] == ["d-beta"] + assert context.workspaces["alpha"].tool_contexts["beta/ctx-1"]["drafts"] == ["d-alpha"] + assert context.workspaces["beta"].tool_contexts["ctx-1"]["drafts"] == ["d-beta"] @pytest.mark.asyncio From 440b85b5c49cb34f805b909139d14e2b3739f711 Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 22 Oct 2026 18:00:15 +0800 Subject: [PATCH 08/21] feat(workspace): serve Codex job tools from a parent-owned loopback endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Commit B of the final review round, restructuring the Codex job bridge and codex thread ownership: - The bridge subprocess no longer builds a child Application per thread. CodexAgentWrapper owns a LoopbackJobServer: a minimal HTTP/JSON endpoint on 127.0.0.1 with an ephemeral port and bearer token, started lazily on first job_tools use and closed with the wrapper. codex_mcp_server is now a thin FastMCP forwarder that fetches job metadata from the parent and forwards calls — with the server-owned workspace alias, tool context id, and injected kwargs — back to the parent for execution, so workspace stores have exactly one persistent owner and bridge calls see resident parent state. The private config snapshot machinery is gone; mcp_config stays accepted with a one-time warning. - Codex thread ownership is persisted per workspace under mem_session/codex/owned_threads.json: resume, fork, and compact of a thread recorded under another alias fail explicitly, including after a wrapper restart, while unrecorded (pre-existing or external) threads pass through to the SDK. The alias-namespaced tool-context guard is unchanged. Refs #632 --- .../agent_wrapper/codex_agent_wrapper.py | 167 +++++++----- .../agent_wrapper/codex_mcp_server.py | 163 ++++++----- .../agent_wrapper/loopback_job_server.py | 148 ++++++++++ tests/unit/test_codex_agent_wrapper.py | 174 ++++-------- tests/unit/test_workspace_phase3.py | 257 ++++++++++-------- 5 files changed, 539 insertions(+), 370 deletions(-) create mode 100644 reme/components/agent_wrapper/loopback_job_server.py diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 15e3e420..f40c678d 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -4,7 +4,6 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator -from contextlib import suppress from dataclasses import dataclass from functools import partial import hashlib @@ -12,10 +11,10 @@ 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 .loopback_job_server import LoopbackJobServer from ..base_component import _context_workspaces from ..component_registry import R from ...enumeration import ChunkEnum @@ -116,8 +115,10 @@ class CodexAgentWrapper(BaseAgentWrapper): 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] = {} + self._thread_owners: dict[str, set[str]] = {} + self._job_server: LoopbackJobServer | None = None + self._mcp_config_warned = False @property def session_path(self) -> Path: @@ -129,13 +130,6 @@ class CodexAgentWrapper(BaseAgentWrapper): 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) @@ -155,49 +149,15 @@ 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: + def _mcp_config_source_unused(self, kwargs: dict[str, Any]) -> None: + """Warn once that job_tools now execute in the parent via the loopback endpoint.""" 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: - payload = self.app_context.app_config.model_dump(mode="json") - workspaces = _context_workspaces(self.app_context) - if workspaces and len(workspaces) > 1: - # The bridge rebuilds the full multi-workspace config; - # the serialized workspace_dir string alone would lose it. - payload["workspace_dir"] = { - alias: str(workspace.workspace_path) for alias, workspace in workspaces.items() - } - json.dump(payload, 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()) + if value is not None and not self._mcp_config_warned: + self._mcp_config_warned = True + self.logger.warning( + "mcp_config is ignored for job_tools: Codex job tools execute in the parent " + "process via the wrapper-owned loopback endpoint", + ) @staticmethod def _resolve_auth_config(auth_mode: str, api_key: str = "", base_url: str = "") -> _CodexAuthConfig: @@ -249,9 +209,19 @@ class CodexAgentWrapper(BaseAgentWrapper): names = ", ".join(invalid) raise TypeError(f"Codex client options must be configured on the wrapper: {names}") + async def _ensure_job_server(self) -> LoopbackJobServer: + """Start the wrapper-owned loopback job endpoint on first bridge use.""" + if self._job_server is None: + if self.app_context is None: + raise RuntimeError("Codex job_tools require an app_context") + self._job_server = LoopbackJobServer(self.app_context) + await self._job_server.start() + return self._job_server + def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: from ..job import BackgroundJob, StreamJob + self._mcp_config_source_unused(kwargs) job_names = list(dict.fromkeys(kwargs.get("job_tools") or [])) if not job_names: return None @@ -259,19 +229,19 @@ class CodexAgentWrapper(BaseAgentWrapper): 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)}") + if self._job_server is None: + raise RuntimeError("Codex job endpoint is not started; call _ensure_job_server first") alias = str(kwargs.get("workspace") or "") - config_source = self._mcp_config_source(kwargs) args = [ "-m", "reme.components.agent_wrapper.codex_mcp_server", - "--config", - config_source, - "--workspace", - str(self._call_workspace_path(kwargs)), + "--endpoint", + self._job_server.url, + "--token", + self._job_server.token, ] if alias: - # The child builds the full multi-workspace config and binds via the alias. args.extend(["--workspace-alias", alias]) for name in job_names: args.extend(["--job", name]) @@ -326,6 +296,11 @@ class CodexAgentWrapper(BaseAgentWrapper): 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 thread_id: + # Cross-workspace thread access is rejected before touching the SDK. + await self._check_thread_ownership(thread_id, alias) + if kwargs.get("job_tools"): + await self._ensure_job_server() if self._release_threads and self._thread_tool_contexts: await self._prune_thread_tool_contexts(codex) requested_tool_context = json.dumps( @@ -372,9 +347,72 @@ class CodexAgentWrapper(BaseAgentWrapper): thread_source=thread_source, **common, ) + if fork_session or not thread_id: + # Threads created under this binding are owned by this workspace. + self._record_thread_ownership(alias, thread.id) self._thread_tool_contexts[f"{alias}/{thread.id}" if alias else thread.id] = requested_tool_context return thread + # ----- Thread ownership (per-workspace, persisted) --------------------- + + def _normalize_thread_alias(self, alias: str) -> str: + """Map a call alias to a concrete workspace alias ("" means the default).""" + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return alias + return alias if alias in workspaces else next(iter(workspaces)) + + def _thread_ownership_path(self, alias: str) -> Path | None: + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return None + workspace = workspaces.get(alias) + if workspace is None or workspace.mem_session_path is None: + return None + return workspace.mem_session_path / "codex" / "owned_threads.json" + + def _owned_thread_ids(self, alias: str) -> set[str]: + """Load (once) the persisted ids of threads this workspace created.""" + if alias in self._thread_owners: + return self._thread_owners[alias] + ids: set[str] = set() + path = self._thread_ownership_path(alias) + if path is not None and path.is_file(): + try: + raw = json.loads(path.read_text(encoding="utf-8")) + ids = {str(thread_id) for thread_id in raw.get("thread_ids", [])} + except (OSError, ValueError, AttributeError) as exc: + self.logger.warning(f"Ignoring unreadable Codex thread ownership file {path}: {exc}") + self._thread_owners[alias] = ids + return ids + + def _record_thread_ownership(self, alias: str, thread_id: str) -> None: + alias = self._normalize_thread_alias(alias) + path = self._thread_ownership_path(alias) + ids = self._owned_thread_ids(alias) + if path is None or thread_id in ids: + return + ids.add(thread_id) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"thread_ids": sorted(ids)}, indent=1), encoding="utf-8") + + async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: + """Reject resume/fork/compact of a thread owned by another workspace. + + Threads recorded nowhere predate ownership tracking or came from other + tools; they pass through and let the SDK decide. + """ + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return + alias = self._normalize_thread_alias(alias) + for owner_alias in workspaces: + if owner_alias != alias and thread_id in self._owned_thread_ids(owner_alias): + raise ValueError( + f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " + f"it cannot be used from workspace {alias!r}", + ) + 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 @@ -415,7 +453,7 @@ class CodexAgentWrapper(BaseAgentWrapper): return codex async def _close(self) -> None: - """Close the persistent app-server and remove its private config snapshot.""" + """Close the persistent app-server and the wrapper-owned job endpoint.""" async with self._turn_lock: codex, self._codex = self._codex, None self._thread_tool_contexts.clear() @@ -423,9 +461,9 @@ class CodexAgentWrapper(BaseAgentWrapper): 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 + if self._job_server is not None: + job_server, self._job_server = self._job_server, None + await job_server.close() async def _prune_thread_tool_contexts(self, codex: AsyncCodex) -> None: """Forget tool constraints only after the server confirms a thread is unloaded.""" @@ -509,10 +547,11 @@ class CodexAgentWrapper(BaseAgentWrapper): raise ValueError("Codex returned invalid JSON for the requested output_schema") from exc return response - async def compact_session(self, session_id: str) -> None: - """Start native compaction of a Codex thread.""" + async def compact_session(self, session_id: str, workspace: str | None = None) -> None: + """Start native compaction of a Codex thread owned by the bound workspace.""" await self.start() async with self._turn_lock: + await self._check_thread_ownership(session_id, str(workspace or "")) codex = await self._get_codex() thread = await codex.thread_resume(session_id) try: diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index cedd428a..dd0d6ad9 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -1,105 +1,118 @@ -"""FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" +"""Thin FastMCP STDIO bridge forwarding Codex job tools to the parent process. + +This subprocess never builds a ReMe Application or any workspace store. It +asks the parent-owned loopback endpoint (LoopbackJobServer) for job metadata +and forwards every tool call — together with the server-owned workspace alias, +tool context id, and injected kwargs — back to the parent for execution, so +all workspace state has exactly one owner. +""" import argparse import json -from pathlib import Path from typing import Any -from ...config import resolve_app_config -from ...reme import ReMe +import httpx 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( - "--workspace-alias", - default="", - help="Alias selecting one workspace of a multi-workspace config snapshot; " - "without it, --workspace overrides the config with a single workspace (legacy behavior)", - ) + parser = argparse.ArgumentParser(description="Forward selected ReMe jobs from the parent loopback endpoint") + parser.add_argument("--endpoint", required=True, help="Parent loopback endpoint base URL") + parser.add_argument("--token", required=True, help="Bearer token for the loopback endpoint") parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") + parser.add_argument("--workspace-alias", default="", help="Workspace alias injected into every job call") 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", ) + parser.add_argument("--tool-timeout", type=float, default=300, help="Per-call timeout in seconds") 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 +class _EndpointClient: + """Minimal client for the parent's loopback job server.""" - missing = sorted(selected.difference(jobs)) - if missing: - raise KeyError(f"Codex job tools not found or not request jobs: {', '.join(missing)}") + def __init__(self, endpoint: str, token: str, timeout: float) -> None: + self._client = httpx.Client( + base_url=endpoint, + headers={"Authorization": f"Bearer {token}"}, + timeout=timeout, + ) - 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 + def list_jobs(self) -> list[dict[str, Any]]: + """Fetch the servable job metadata from the parent endpoint.""" + response = self._client.get("/jobs") + response.raise_for_status() + return response.json()["jobs"] - prepared = dict(config) - prepared["jobs"] = jobs - prepared["service"] = service - return prepared + def call(self, job: str, kwargs: dict[str, Any], injected_job_kwargs: dict[str, Any] | None = None) -> str: + """Forward one tool call to the parent endpoint and unwrap its answer.""" + response = self._client.post( + "/call", + json={"job": job, "kwargs": kwargs, "injected_job_kwargs": injected_job_kwargs or {}}, + ) + if response.status_code != 200: + try: + detail = response.json().get("error") or response.text + except ValueError: + detail = response.text + raise RuntimeError(str(detail)) + payload = response.json() + if not payload.get("success"): + raise RuntimeError(str(payload.get("answer") or payload)) + return str(payload.get("answer") or "") + + def close(self) -> None: + """Close the underlying HTTP client.""" + self._client.close() def main() -> None: - """Load ReMe and serve the requested jobs over STDIO.""" + """Fetch job metadata from the parent and serve forwarding tools over STDIO.""" + from fastmcp import FastMCP + from fastmcp.exceptions import ToolError + from fastmcp.tools import FunctionTool + args = _parse_args() - overrides: dict[str, Any] = {} - if not args.workspace_alias: - # Legacy single-workspace mode: the CLI path replaces the config value. - overrides["workspace_dir"] = str(Path(args.workspace).absolute()) - config = resolve_app_config( - config=args.config, - enable_logo=False, - log_to_console=False, - log_to_file=False, - log_config=False, - **overrides, - ) - if args.workspace_alias: - raw_workspaces = config.get("workspace_dir") - if not isinstance(raw_workspaces, dict) or args.workspace_alias not in raw_workspaces: - raise ValueError( - f"--workspace-alias {args.workspace_alias!r} requires a config snapshot " - "carrying the full workspace mapping", - ) - 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): + injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else {} + if not isinstance(injected_job_kwargs, dict): raise TypeError("--injected-job-kwargs must be a JSON object") - injected_job_kwargs = dict(injected_job_kwargs or {}) + injected = dict(injected_job_kwargs) if args.workspace_alias: - # The binding is server-owned; nested job calls inherit it from here. - injected_job_kwargs.setdefault("workspace", args.workspace_alias) - config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) - ReMe(**config).run_app() + injected["workspace"] = args.workspace_alias + if args.tool_context_id: + injected["tool_context_id"] = args.tool_context_id + + client = _EndpointClient(args.endpoint, args.token, args.tool_timeout) + metadata = {entry["name"]: entry for entry in client.list_jobs()} + missing = sorted(set(args.jobs) - set(metadata)) + if missing: + raise KeyError(f"Codex job tools not found on the parent endpoint: {', '.join(missing)}") + + server = FastMCP("reme-codex-bridge") + for name in args.jobs: + entry = metadata[name] + + def forward(_client=client, _name=name, _injected=None, **kwargs): + try: + return _client.call(_name, kwargs, _injected or injected) + except Exception as exc: # noqa: BLE001 - surface every failure as a tool error + raise ToolError(str(exc)) from exc + + server.add_tool( + FunctionTool( + name=name, + description=entry.get("description") or "", + parameters=entry.get("parameters") or {"type": "object", "properties": {}}, + fn=forward, + ), + ) + + try: + server.run(transport="stdio") + finally: + client.close() if __name__ == "__main__": diff --git a/reme/components/agent_wrapper/loopback_job_server.py b/reme/components/agent_wrapper/loopback_job_server.py new file mode 100644 index 00000000..fa242336 --- /dev/null +++ b/reme/components/agent_wrapper/loopback_job_server.py @@ -0,0 +1,148 @@ +"""Wrapper-owned loopback job endpoint for agent-bridge subprocesses. + +The Codex job bridge used to spawn a full child Application per thread; every +thread then held its own workspace stores and became a second persistence +writer. This endpoint keeps job execution in the parent process instead: a +minimal HTTP/JSON server on 127.0.0.1 with an ephemeral port and a bearer +token, owned by one CodexAgentWrapper instance for its lifetime. Bridge +subprocesses are thin FastMCP forwarders that call this endpoint. +""" + +import asyncio +import json +import secrets +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + from ..application_context import ApplicationContext + from ..job import BaseJob + +_REASON = { + 200: "OK", + 400: "Bad Request", + 401: "Unauthorized", + 404: "Not Found", + 500: "Internal Server Error", +} + + +class LoopbackJobServer: + """Serve this application's request jobs to trusted local subprocesses.""" + + def __init__(self, app_context: "ApplicationContext") -> None: + self._app_context = app_context + self.token = secrets.token_urlsafe(24) + self._server: asyncio.Server | None = None + + @property + def is_running(self) -> bool: + """Whether the endpoint is currently bound.""" + return self._server is not None + + @property + def url(self) -> str: + """Base URL of the bound endpoint.""" + if self._server is None: + raise RuntimeError("Loopback job server is not started") + port = self._server.sockets[0].getsockname()[1] + return f"http://127.0.0.1:{port}" + + async def start(self) -> None: + """Bind the endpoint; idempotent so lazy bridge creation can retry.""" + if self._server is None: + self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0) + + async def close(self) -> None: + """Stop the endpoint and release the port.""" + server, self._server = self._server, None + if server is not None: + server.close() + await server.wait_closed() + + # ----- HTTP handling --------------------------------------------------- + + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=30) + request_line, *header_lines = head.decode("latin-1").split("\r\n") + method, target, _version = request_line.split(" ", 2) + headers = {} + for line in header_lines: + if ":" in line: + key, value = line.split(":", 1) + headers[key.strip().lower()] = value.strip() + length = int(headers.get("content-length", 0) or 0) + body = await reader.readexactly(length) if length else b"" + status, payload = await self._dispatch(method, target, headers, body) + except Exception as exc: # pylint: disable=broad-exception-caught + status, payload = 500, {"error": f"{type(exc).__name__}: {exc}"} + data = json.dumps(payload).encode() + writer.write( + f"HTTP/1.1 {status} {_REASON.get(status, 'OK')}\r\n" + "content-type: application/json\r\n" + f"content-length: {len(data)}\r\n" + "connection: close\r\n\r\n".encode() + data, + ) + try: + await writer.drain() + finally: + writer.close() + + async def _dispatch( + self, + method: str, + target: str, + headers: dict[str, str], + body: bytes, + ) -> tuple[int, dict[str, Any]]: + if headers.get("authorization") != f"Bearer {self.token}": + return 401, {"error": "unauthorized"} + if method == "GET" and target == "/jobs": + return 200, {"jobs": self._job_metadata()} + if method == "POST" and target == "/call": + try: + request = json.loads(body or b"{}") + except json.JSONDecodeError as exc: + return 400, {"error": f"invalid JSON body: {exc}"} + return await self._call(request) + return 404, {"error": "not found"} + + # ----- Job surface ------------------------------------------------------- + + def _request_jobs(self) -> dict[str, "BaseJob"]: + """Return the application's non-stream, non-background jobs by name.""" + from ..job import BackgroundJob, StreamJob + + return { + name: job for name, job in self._app_context.jobs.items() if not isinstance(job, (StreamJob, BackgroundJob)) + } + + def _job_metadata(self) -> list[dict[str, Any]]: + """Describe every servable job for bridge tool registration.""" + return [ + {"name": name, "description": job.description, "parameters": job.parameters} + for name, job in sorted(self._request_jobs().items()) + ] + + async def _call(self, request: dict[str, Any]) -> tuple[int, dict[str, Any]]: + name = str(request.get("job") or "") + job = self._request_jobs().get(name) + if job is None: + return 404, {"error": f"job {name!r} not found or not a request job"} + + kwargs = dict(request.get("kwargs") or {}) + injected = dict(request.get("injected_job_kwargs") or {}) + conflicts = sorted(injected.keys() & kwargs.keys()) + if conflicts: + return 400, {"error": f"{', '.join(conflicts)} injected by the server and cannot be provided"} + kwargs.update(injected) + if workspace := request.get("workspace"): + kwargs["workspace"] = workspace + if tool_context_id := request.get("tool_context_id"): + kwargs["tool_context_id"] = tool_context_id + + try: + response = await job(**kwargs) + except ValueError as exc: + return 400, {"error": str(exc)} + return 200, {"success": response.success, "answer": response.answer} diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index 6f436d52..b0da6d96 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -7,8 +7,6 @@ from dataclasses import dataclass import json import os from pathlib import Path -import stat -import sys from types import SimpleNamespace import pytest @@ -16,12 +14,11 @@ from openai_codex.generated.v2_all import TokenUsageBreakdown from pydantic import BaseModel 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.config import resolve_app_config from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import ApplicationConfig, Response +from reme.schema import 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. @@ -64,24 +61,32 @@ def _wrapper(tmp_path, **kwargs): return CodexAgentWrapper(app_context=context, **kwargs), job -def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): +@pytest.mark.asyncio +async def test_mcp_config_uses_loopback_bridge_and_selected_jobs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") + await wrapper._ensure_job_server() # pylint: disable=protected-access config = wrapper._mcp_server_config( # pylint: disable=protected-access {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, ) + await wrapper.close() 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("--endpoint") + 1].startswith("http://127.0.0.1:") + assert config["args"][config["args"].index("--token") + 1] assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" + # mcp_config is accepted for compatibility but no longer builds a child app. + assert "--config" not in config["args"] -def test_mcp_config_serializes_injected_job_kwargs(tmp_path): +@pytest.mark.asyncio +async def test_mcp_config_serializes_injected_job_kwargs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") + await wrapper._ensure_job_server() # pylint: disable=protected-access config = wrapper._mcp_server_config( # pylint: disable=protected-access { @@ -90,6 +95,7 @@ def test_mcp_config_serializes_injected_job_kwargs(tmp_path): "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]}, }, ) + await wrapper.close() raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] assert json.loads(raw) == {"_allowed_paths": ["daily/2025-06-01/note.md"]} @@ -98,28 +104,17 @@ def test_mcp_config_serializes_injected_job_kwargs(tmp_path): 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): +@pytest.mark.asyncio +async def test_thread_config_preserves_other_mcp_servers(tmp_path): wrapper, _job = _wrapper(tmp_path) + await wrapper._ensure_job_server() # pylint: disable=protected-access config = wrapper._thread_config( # pylint: disable=protected-access { "job_tools": ["search"], "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}}, }, ) + await wrapper.close() assert "docs" in config["mcp_servers"] assert len(config["mcp_servers"]) == 2 @@ -171,114 +166,69 @@ def test_mcp_config_rejects_background_jobs(tmp_path): 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", - ) +@pytest.mark.asyncio +async def test_loopback_endpoint_serves_parent_jobs(tmp_path): + import httpx - 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 + wrapper, job = _wrapper(tmp_path) + await wrapper.start() + server = await wrapper._ensure_job_server() # pylint: disable=protected-access + base = server.url + headers = {"Authorization": f"Bearer {server.token}"} + try: + async with httpx.AsyncClient(base_url=base, headers=headers, timeout=10) as http: + listed = (await http.get("/jobs")).json()["jobs"] + assert [entry["name"] for entry in listed] == ["search"] + called = await http.post("/call", json={"job": "search", "kwargs": {"query": "memo"}}) + assert called.json() == {"success": True, "answer": "found:memo"} + # The job executed in the parent process, not in any child. + assert job.calls == [{"query": "memo"}] -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"]) + missing = await http.post("/call", json={"job": "missing", "kwargs": {}}) + assert missing.status_code == 404 + + unauthorized = await http.get("/jobs", headers={"Authorization": "Bearer wrong"}) + assert unauthorized.status_code == 401 + finally: + await wrapper.close() + assert server.is_running is False @pytest.mark.asyncio -async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path): +async def test_loopback_bridge_lists_and_calls_parent_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", - ) + wrapper, job = _wrapper(tmp_path) + await wrapper._ensure_job_server() # pylint: disable=protected-access + server_config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access 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]), + command=server_config["command"], + args=server_config["args"], + cwd=server_config["cwd"], 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] == ["search"] + result = await asyncio.wait_for(client.call_tool("search", {"query": "via-bridge"}), timeout=10) + await wrapper.close() - assert [tool.name for tool in tools] == ["empty"] + assert "found:via-bridge" in str(result.content[0].text) + assert job.calls == [{"query": "via-bridge"}] @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", - ) + wrapper, _job = _wrapper(tmp_path) + await wrapper._ensure_job_server() # pylint: disable=protected-access + server_config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access 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]), + server_config["command"], + *server_config["args"], + cwd=server_config["cwd"], stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, @@ -309,6 +259,7 @@ async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): except asyncio.TimeoutError: proc.kill() await proc.wait() + await wrapper.close() stdout = first_line + await proc.stdout.read() stderr = (await proc.stderr.read()).decode() assert b"Loading config" not in stdout @@ -1199,11 +1150,6 @@ def test_codex_home_expands_user_directory(tmp_path): 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) diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index a33f5399..81bc8a72 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -3,7 +3,6 @@ # pylint: disable=protected-access,missing-function-docstring import asyncio -from pathlib import Path from types import SimpleNamespace import pytest @@ -15,6 +14,7 @@ from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper from reme.components.component_registry import R from reme.components.job import BackgroundJob, BaseJob, CronJob from reme.components.runtime_context import RuntimeContext +from reme.enumeration import ComponentEnum from reme.steps.base_step import BaseStep from reme.steps.evolve.base_auto_resource import _compute_agent_session_id from reme.steps.index.draft import AddDraftStep @@ -135,23 +135,67 @@ async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): ) class FakeCodex: - """Minimal Codex double returning a fixed thread id.""" + """Minimal Codex double echoing the requested thread id.""" + + def __init__(self): + self.next_id = "thread-1" + + async def thread_start(self, **_kwargs): + return SimpleNamespace(id=self.next_id) + + async def thread_resume(self, thread_id, **_kwargs): + return SimpleNamespace(id=thread_id) + + codex = FakeCodex() + await wrapper._open_thread(codex, {"tool_context_id": "ctx-a"}) + # A different thread under another workspace gets its own guard entry. + codex.next_id = "thread-2" + await wrapper._open_thread(codex, {"workspace": "beta", "tool_context_id": "ctx-b"}) + # Same workspace resume with changed constraints still raises. + with pytest.raises(ValueError, match="cannot change"): + await wrapper._open_thread(codex, {"session_id": "thread-1", "tool_context_id": "ctx-b"}) + # A thread owned by one workspace cannot be resumed from another. + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await wrapper._open_thread(codex, {"session_id": "thread-1", "workspace": "beta"}) + + assert sorted(wrapper._thread_tool_contexts) == ["beta/thread-2", "thread-1"] + + +@pytest.mark.asyncio +async def test_codex_thread_ownership_persists_across_restart(tmp_path): + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + + class FakeCodex: + """Minimal Codex double echoing the requested thread id.""" async def thread_start(self, **_kwargs): return SimpleNamespace(id="thread-1") - async def thread_resume(self, _thread_id, **_kwargs): - return SimpleNamespace(id="thread-1") + async def thread_resume(self, thread_id, **_kwargs): + return SimpleNamespace(id=thread_id) + + async def thread_fork(self, thread_id, **_kwargs): + return SimpleNamespace(id=f"{thread_id}-fork") codex = FakeCodex() - await wrapper._open_thread(codex, {"tool_context_id": "ctx-a"}) - # Same thread id under a different workspace is a different guard entry. - await wrapper._open_thread(codex, {"session_id": "thread-1", "workspace": "beta", "tool_context_id": "ctx-b"}) - # Same workspace resume with changed constraints still raises. - with pytest.raises(ValueError, match="cannot change"): - await wrapper._open_thread(codex, {"session_id": "thread-1", "tool_context_id": "ctx-b"}) + wrapper = CodexAgentWrapper(app_context=context) + await wrapper._open_thread(codex, {"workspace": "alpha"}) + ownership = tmp_path / "alpha" / "mem_session" / "codex" / "owned_threads.json" + assert ownership.is_file() - assert sorted(wrapper._thread_tool_contexts) == ["beta/thread-1", "thread-1"] + # A fresh wrapper (restart) reloads the ownership records and still rejects. + restarted = CodexAgentWrapper(app_context=context) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await restarted._open_thread(codex, {"session_id": "thread-1", "workspace": "beta"}) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await restarted._open_thread(codex, {"session_id": "thread-1", "workspace": "beta", "fork_session": True}) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await restarted.compact_session("thread-1", workspace="beta") + # The owning workspace keeps working normally. + thread = await restarted._open_thread(codex, {"session_id": "thread-1", "workspace": "alpha"}) + assert thread.id == "thread-1" + # Threads recorded nowhere pass through and let the SDK decide. + await restarted._open_thread(codex, {"session_id": "thread-elsewhere", "workspace": "beta"}) # ----- Background/cron binding ----------------------------------------------- @@ -251,25 +295,27 @@ async def test_mcp_tool_boundary_forwards_injected_workspace(tmp_path): @pytest.mark.asyncio -async def test_codex_bridge_uses_bound_workspace_path(tmp_path): +async def test_codex_bridge_uses_loopback_endpoint(tmp_path): import json context = _context(tmp_path) context.jobs = {"search_job": _job(context, "search_job")} wrapper = CodexAgentWrapper(app_context=context) + await wrapper.start() + await wrapper._ensure_job_server() # pylint: disable=protected-access config = wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "beta"}) - assert config["args"][config["args"].index("--workspace") + 1] == str(tmp_path / "beta") + assert "--workspace-alias" in config["args"] + assert config["args"][config["args"].index("--workspace-alias") + 1] == "beta" raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] assert json.loads(raw) == {"workspace": "beta"} + assert "--config" not in config["args"] and "--workspace" not in config["args"] default_config = wrapper._mcp_server_config({"job_tools": ["search_job"]}) - assert default_config["args"][default_config["args"].index("--workspace") + 1] == str(tmp_path / "alpha") - - with pytest.raises(ValueError, match="Unknown workspace 'nope'"): - wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "nope"}) + assert "--workspace-alias" not in default_config["args"] await wrapper.close() + assert wrapper._job_server is None # ----- Review follow-ups: explicit bindings, session isolation, key collisions @@ -378,120 +424,97 @@ async def test_codex_thread_cwd_follows_bound_workspace(tmp_path): @pytest.mark.asyncio -async def test_codex_bridge_executes_non_default_alias(tmp_path): - """The bridge child rebuilds the full workspace mapping and binds via the alias.""" - import json - import sys - +async def test_codex_bridge_reads_resident_parent_state(tmp_path): + """Bridge calls execute against the parent's live stores; there is no second owner.""" from fastmcp import Client from fastmcp.client import StdioTransport - snapshot = tmp_path / "bridge.json" - snapshot.write_text( - json.dumps( - { - "workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, - "enable_logo": False, - "log_to_console": False, - "log_to_file": False, - "service": {"backend": "mcp"}, - "components": { - "tokenizer": {"default": {"backend": "regex"}}, - "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, - "keyword_index": {"default": {"backend": "bm25"}}, - "file_graph": {"default": {"backend": "local"}}, - }, - "jobs": { - "write_note": { - "backend": "base", - "description": "write a note", - "parameters": { - "type": "object", - "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, - "required": ["path", "content"], - }, - "steps": [{"backend": "write_step"}], - }, + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={ + "tokenizer": {"default": {"backend": "regex"}}, + "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, + "keyword_index": {"default": {"backend": "bm25"}}, + "file_graph": {"default": {"backend": "local"}}, + }, + jobs={ + "write_note": { + "backend": "base", + "description": "write a note", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, + "required": ["path", "content"], }, + "steps": [{"backend": "write_step"}], }, - ), - encoding="utf-8", + "read_note": { + "backend": "base", + "description": "read a note", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + "steps": [{"backend": "read_step"}], + }, + }, ) - transport = StdioTransport( - command=sys.executable, - args=[ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(snapshot), - "--workspace", - str(tmp_path / "beta"), - "--workspace-alias", - "beta", - "--job", - "write_note", - ], - cwd=str(Path(__file__).resolve().parents[2]), - keep_alive=False, - ) - - async with Client(transport, timeout=None, init_timeout=30) as client: - await asyncio.wait_for( - client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-beta"}), - timeout=30, + wrapper = CodexAgentWrapper(app_context=app.context) + async with app: + await wrapper.start() + # Index a note directly in the parent's beta store. + await app.run_job("write_note", path="daily/resident.md", content="parent-resident-beta", workspace="beta") + await wrapper._ensure_job_server() # pylint: disable=protected-access + server_config = wrapper._mcp_server_config( # pylint: disable=protected-access + {"job_tools": ["read_note", "write_note"], "workspace": "beta"}, ) - - assert "via-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") - assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() + transport = StdioTransport( + command=server_config["command"], + args=server_config["args"], + cwd=server_config["cwd"], + keep_alive=False, + ) + async with Client(transport, timeout=None, init_timeout=30) as client: + read = await asyncio.wait_for(client.call_tool("read_note", {"path": "daily/resident.md"}), timeout=30) + await asyncio.wait_for( + client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-bridge-beta"}), + timeout=30, + ) + # The bridge saw the parent's resident beta state and wrote through it. + assert "parent-resident-beta" in str(read.content[0].text) + assert "via-bridge-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") + assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() + store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + assert all("bridge" not in node.path for node in await store_alpha.get_nodes()) + await wrapper.close() @pytest.mark.asyncio -async def test_codex_bridge_rejects_alias_without_mapping_snapshot(tmp_path): - import json - import sys - +async def test_codex_bridge_unknown_alias_fails_at_execution(tmp_path): + """An unknown workspace alias surfaces as a tool error from the parent boundary.""" from fastmcp import Client from fastmcp.client import StdioTransport - snapshot = tmp_path / "bridge-single.json" - snapshot.write_text( - json.dumps( - { - "workspace_dir": str(tmp_path / "only"), - "enable_logo": False, - "log_to_console": False, - "log_to_file": False, - "service": {"backend": "mcp"}, - "jobs": { - "write_note": { - "backend": "base", - "parameters": {"type": "object", "properties": {}}, - "steps": [{"backend": "write_step"}], - }, - }, - }, - ), - encoding="utf-8", + context = _context(tmp_path) + context.jobs = {"search_job": _job(context, "search_job")} + wrapper = CodexAgentWrapper(app_context=context, project_path=str(tmp_path)) + await wrapper.start() + await wrapper._ensure_job_server() # pylint: disable=protected-access + server_config = wrapper._mcp_server_config( # pylint: disable=protected-access + {"job_tools": ["search_job"], "workspace": "nope"}, ) transport = StdioTransport( - command=sys.executable, - args=[ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(snapshot), - "--workspace", - str(tmp_path / "only"), - "--workspace-alias", - "beta", - "--job", - "write_note", - ], - cwd=str(Path(__file__).resolve().parents[2]), + command=server_config["command"], + args=server_config["args"], + cwd=server_config["cwd"], keep_alive=False, ) - - # The child process fails fast with a ValueError; the client sees the closed connection. - with pytest.raises(Exception, match="Connection closed"): - async with Client(transport, timeout=None, init_timeout=30): - pass + async with Client(transport, timeout=None, init_timeout=30) as client: + with pytest.raises(Exception, match="Unknown workspace 'nope'"): + await asyncio.wait_for(client.call_tool("search_job", {}), timeout=30) + await wrapper.close() From f0afcfd13a54d2e3aace845395abc7157c139133 Mon Sep 17 00:00:00 2001 From: celanwang Date: Mon, 26 Oct 2026 18:56:30 +0800 Subject: [PATCH 09/21] refactor(workspace): consolidate ownership and Codex integration --- AGENTS.md | 13 +- docs/en/configuration.md | 9 +- docs/zh/configuration.md | 4 + reme/application.py | 45 ++-- .../agent_wrapper/codex_agent_wrapper.py | 188 +++++++++------ .../agent_wrapper/codex_mcp_server.py | 224 +++++++++++------- .../agent_wrapper/loopback_job_server.py | 148 ------------ reme/components/application_context.py | 88 ++++--- reme/components/base_component.py | 26 +- reme/schema/workspace.py | 36 +-- reme/steps/base_step.py | 16 +- reme/steps/evolve/proactive/topics.py | 10 +- reme/steps/evolve/proactive/utils.py | 2 +- reme/steps/index/update_changes.py | 5 +- tests/unit/test_codex_agent_wrapper.py | 95 ++------ tests/unit/test_workspace_binding.py | 10 +- tests/unit/test_workspace_phase3.py | 176 ++++++++------ 17 files changed, 548 insertions(+), 547 deletions(-) delete mode 100644 reme/components/agent_wrapper/loopback_job_server.py diff --git a/AGENTS.md b/AGENTS.md index bb53cc0a..3836471e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -132,8 +132,12 @@ Treat Step instances as invocation-scoped: - `self.context.response.answer`, `success`, and `metadata` are request-scoped output. Because the same response travels through the Step chain, later Steps may consume metadata produced earlier, but it is not application-lifetime or durable storage. -- `self.app_context.metadata` holds in-memory state shared across Job/Step invocations for the life of one - `Application`, such as counters, tool-context state, session maps, or locks. +- `self.app_context.metadata` holds application-lifetime operational state shared across Job/Step invocations, + such as counters. `ApplicationContext.shared_components` owns genuinely shared components; its `components` + compatibility view also exposes default-workspace instances to legacy callers. +- Each `WorkspaceContext` owns its paths, stateful components, tool contexts, search budgets, agent session state, + and locks. `RuntimeContext.workspace` is the explicit binding used by workspace-bound execution; internal Steps + must not resolve a default workspace repeatedly or use the compatibility component view. - Workspace files or a dedicated Component/store hold durable state that must survive restart. Use narrow, namespaced keys in `app_context.metadata` and protect shared mutable values against concurrent access. The @@ -144,8 +148,9 @@ dedicated lifecycle, locking, or persistence, promote it to a typed context fiel Additional Step contracts: -- `Ref` dependencies resolve in this order: Step kwargs, current `RuntimeContext`, then the named application component. - The value is cached only on the current Step instance and cleared before each call. +- `Ref` dependencies resolve in this order: Step kwargs, current `RuntimeContext`, then the bound workspace for scoped + components or the application shared-component registry for shared components. The value is cached only on the + current Step instance and cleared before each call. - `input_mapping` and `output_mapping` copy keys within `RuntimeContext.data`; missing sources are ignored. - Dispatched Steps receive the current `RuntimeContext`, so their data and response are shared. - Base jobs convert uncaught Step errors into `Response(success=False)`; stream jobs emit an error chunk and always a diff --git a/docs/en/configuration.md b/docs/en/configuration.md index d681c786..5bb68d2b 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -100,8 +100,13 @@ External stateful backends that cannot isolate workspaces (currently the `neo4j` startup in multi-workspace configurations. The binding propagates through nested job calls and agent tool calls (`job_tools`) as a server-owned constraint -that models cannot override. Background and cron jobs bind through their job configuration, so extra workspaces -are watched by declaring explicit per-alias variants: +that models cannot override. For Codex with multiple workspaces, `job_tools` requires the resident HTTP service +with MCP enabled, and the selected jobs must be exposed by that service. Its STDIO bridge forwards to this service +without creating another Application. The multi-workspace one-shot CLI combination is not supported; `mcp_config` +retains its legacy meaning only for single-workspace Codex job tools. + +Background and cron jobs bind through their job configuration, so extra workspaces are watched by declaring +explicit per-alias variants: ```yaml jobs: diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index d2422797..019fbd72 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -102,6 +102,10 @@ Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 在多 workspace 配置下会在启动时被拒绝。 绑定会经嵌套 Job 调用和 agent 工具调用(`job_tools`)传播,作为模型无法覆盖的服务端约束。 +多 workspace 下的 Codex `job_tools` 需要启用 MCP 的常驻 HTTP 服务,且所选 Job 必须由该服务暴露; +其 STDIO bridge 转发到该服务,不会创建另一个 Application。多 workspace 与 one-shot CLI 的组合暂不支持; +`mcp_config` 仅在单 workspace 的 Codex job tools 中保留原有语义。 + Background 和 cron Job 通过 Job 配置绑定别名,额外 workspace 的 watch 通过显式配置按别名的变体实现: ```yaml diff --git a/reme/application.py b/reme/application.py index 0b4d4b68..76c6e88e 100644 --- a/reme/application.py +++ b/reme/application.py @@ -73,7 +73,6 @@ class Application(BaseComponent): """Instantiate declared components; workspace-scoped ones once per alias.""" multi_workspace = len(self.context.workspaces) > 1 for ctype, group in self.config.components.items(): - self.context.components[ctype] = {} for name, cfg in group.items(): label = f"Component '{name}'" backend_cls = self._backend_class(ctype, cfg, label) @@ -92,12 +91,8 @@ class Application(BaseComponent): workspace_alias=alias, ) workspace.components.setdefault(ctype, {})[name] = instance - # The default alias instance doubles as the plain lookup entry. - self.context.components[ctype][name] = next(iter(self.context.workspaces.values())).components[ - ctype - ][name] else: - self.context.components[ctype][name] = self._instantiate( + self.context.shared_components.setdefault(ctype, {})[name] = self._instantiate( ctype, cfg, label=label, @@ -166,16 +161,13 @@ class Application(BaseComponent): ) -> dict[_NodeKey, BaseComponent]: """One node per component instance; scoped components contribute one node per alias.""" nodes: dict[_NodeKey, BaseComponent] = {} - scoped_keys: set[tuple[str, str]] = set() for alias, workspace in self.context.workspaces.items(): for ctype, group in workspace.components.items(): for name, instance in group.items(): nodes[(ctype, name, alias)] = instance - scoped_keys.add((ctype, name)) - for ctype, group in self.context.components.items(): + for ctype, group in self.context.shared_components.items(): for name, comp in group.items(): - if (ctype, name) not in scoped_keys: - nodes[(ctype, name, "")] = comp + nodes[(ctype, name, "")] = comp if replacement is not None: nodes[replacement[0]] = replacement[1] return nodes @@ -297,11 +289,10 @@ class Application(BaseComponent): """Update an existing component by type/name; never creates missing components.""" async with self._component_mutation_lock: component_type = component_type_name(component_enum) - group = self.context.components.get(component_type) - if not group or name not in group: + component = self.context.get_component(component_type, name) + if component is None: raise KeyError(f"Component '{name}' not found in {component_type}") - component = group[name] # Validate fields without evaluating lazy properties before injection. for key in kwargs: if getattr_static(component, key, _UNSET) is _UNSET and not hasattr(component, key): @@ -317,7 +308,7 @@ class Application(BaseComponent): for group in workspace.components.values(): for instance in group.values(): instances[id(instance)] = instance - for group in self.context.components.values(): + for group in self.context.shared_components.values(): for component in group.values(): instances[id(component)] = component return list(instances.values()) @@ -365,19 +356,27 @@ class Application(BaseComponent): """ async with self._component_mutation_lock: component_type = component_type_name(component_enum) - group = self.context.components.get(component_type) config_group = self.config.components.get(component_type) - if not group or name not in group or config_group is None: - raise KeyError(f"Component '{name}' not found in {component_type}") - - old_component = group[name] - scoped_instances = self.context.workspace_components.get(component_type, {}).get(name) - if scoped_instances is not None and len(scoped_instances) > 1: + scoped_instances = { + alias: workspace.components[component_type][name] + for alias, workspace in self.context.workspaces.items() + if name in workspace.components.get(component_type, {}) + } + if len(scoped_instances) > 1: raise NotImplementedError( f"replace_component does not support workspace-scoped component " f"'{component_type}:{name}' in multi-workspace configurations", ) - alias = next(iter(scoped_instances)) if scoped_instances is not None else None + alias = next(iter(scoped_instances), None) + group = ( + self.context.workspaces[alias].components.get(component_type) + if alias is not None + else self.context.shared_components.get(component_type) + ) + if not group or name not in group or config_group is None: + raise KeyError(f"Component '{name}' not found in {component_type}") + + old_component = group[name] replacement_config = ( config.model_copy(deep=True) if isinstance(config, ComponentConfig) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index f40c678d..078a8500 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator +from contextlib import suppress from dataclasses import dataclass from functools import partial import hashlib @@ -11,14 +12,16 @@ 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 .loopback_job_server import LoopbackJobServer from ..base_component import _context_workspaces from ..component_registry import R +from ...constants import normalize_connect_host from ...enumeration import ChunkEnum from ...schema import StreamChunk, TokenUsage +from ...schema.workspace import WorkspaceAgentState if TYPE_CHECKING: from openai_codex import AsyncCodex, AsyncThread, CodexConfig, RunInput @@ -116,9 +119,7 @@ class CodexAgentWrapper(BaseAgentWrapper): self._release_threads = release_threads self._turn_lock = asyncio.Lock() self._thread_tool_contexts: dict[str, str] = {} - self._thread_owners: dict[str, set[str]] = {} - self._job_server: LoopbackJobServer | None = None - self._mcp_config_warned = False + self._mcp_snapshot_path: Path | None = None @property def session_path(self) -> Path: @@ -149,15 +150,47 @@ class CodexAgentWrapper(BaseAgentWrapper): relative_source = os.path.relpath(source, target.parent) target.symlink_to(relative_source, target_is_directory=True) - def _mcp_config_source_unused(self, kwargs: dict[str, Any]) -> None: - """Warn once that job_tools now execute in the parent via the loopback endpoint.""" + @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 _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 not None and not self._mcp_config_warned: - self._mcp_config_warned = True - self.logger.warning( - "mcp_config is ignored for job_tools: Codex job tools execute in the parent " - "process via the wrapper-owned loopback endpoint", - ) + 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: + 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: @@ -209,19 +242,32 @@ class CodexAgentWrapper(BaseAgentWrapper): names = ", ".join(invalid) raise TypeError(f"Codex client options must be configured on the wrapper: {names}") - async def _ensure_job_server(self) -> LoopbackJobServer: - """Start the wrapper-owned loopback job endpoint on first bridge use.""" - if self._job_server is None: - if self.app_context is None: - raise RuntimeError("Codex job_tools require an app_context") - self._job_server = LoopbackJobServer(self.app_context) - await self._job_server.start() - return self._job_server + def _resident_mcp_endpoint(self) -> str: + from ..service import HttpService, MCPService + + service = self.app_context.service if self.app_context is not None else None + if isinstance(service, HttpService) and service.mcp_enabled: + return f"http://{normalize_connect_host(service.host)}:{service.port}{service.mcp_path}" + if isinstance(service, MCPService) and service.transport == "streamable-http": + return f"http://{normalize_connect_host(service.host)}:{service.port}/mcp" + raise RuntimeError( + "Multi-workspace Codex job_tools require a resident HTTP MCP service; " + "one-shot CLI and other service backends are not supported for this combination", + ) + + @staticmethod + def _validate_resident_jobs(service, jobs) -> None: + unavailable = [ + job.name + for job in jobs + if not job.enabled or not job.enable_serve or (service.jobs is not None and job.name not in service.jobs) + ] + if unavailable: + raise ValueError(f"Codex job_tools must be exposed by the resident MCP service: {', '.join(unavailable)}") def _mcp_server_config(self, kwargs: dict[str, Any]) -> dict[str, Any] | None: from ..job import BackgroundJob, StreamJob - self._mcp_config_source_unused(kwargs) job_names = list(dict.fromkeys(kwargs.get("job_tools") or [])) if not job_names: return None @@ -229,26 +275,29 @@ class CodexAgentWrapper(BaseAgentWrapper): 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)}") - if self._job_server is None: - raise RuntimeError("Codex job endpoint is not started; call _ensure_job_server first") - alias = str(kwargs.get("workspace") or "") - args = [ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--endpoint", - self._job_server.url, - "--token", - self._job_server.token, - ] - if alias: - args.extend(["--workspace-alias", alias]) + workspaces = _context_workspaces(self.app_context) + multi_workspace = workspaces is not None and len(workspaces) > 1 + if multi_workspace: + if self._explicit_mcp_config(kwargs) is not None: + raise ValueError("mcp_config is only supported for single-workspace Codex job_tools") + endpoint = self._resident_mcp_endpoint() + self._validate_resident_jobs(self.app_context.service, jobs) + args = ["-m", "reme.components.agent_wrapper.codex_mcp_server", "--endpoint", endpoint] + else: + args = [ + "-m", + "reme.components.agent_wrapper.codex_mcp_server", + "--config", + self._mcp_config_source(kwargs), + "--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 "")]) injected_job_kwargs = dict(kwargs.get("injected_job_kwargs") or {}) - if workspace := kwargs.get("workspace"): - # The bridge injects the binding into nested job calls; models cannot override it. + if multi_workspace and (workspace := kwargs.get("workspace")): injected_job_kwargs["workspace"] = workspace if injected_job_kwargs: args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) @@ -288,6 +337,7 @@ class CodexAgentWrapper(BaseAgentWrapper): from openai_codex.types import Personality, ThreadSource, ThreadStartSource alias = str(kwargs.get("workspace") or "") + tool_contexts = self._tool_contexts(alias) resume = kwargs.get("resume") or "" session_id = kwargs.get("session_id") or "" if resume and session_id and resume != session_id: @@ -297,12 +347,9 @@ class CodexAgentWrapper(BaseAgentWrapper): if fork_session and not thread_id: raise ValueError("fork_session=True requires resume or session_id") if thread_id: - # Cross-workspace thread access is rejected before touching the SDK. await self._check_thread_ownership(thread_id, alias) - if kwargs.get("job_tools"): - await self._ensure_job_server() - if self._release_threads and self._thread_tool_contexts: - await self._prune_thread_tool_contexts(codex) + if self._release_threads and tool_contexts: + await self._prune_thread_tool_contexts(codex, alias) requested_tool_context = json.dumps( { "tool_context_id": str(kwargs.get("tool_context_id") or ""), @@ -311,9 +358,8 @@ class CodexAgentWrapper(BaseAgentWrapper): sort_keys=True, default=str, ) - guard_key = f"{alias}/{thread_id}" if alias else thread_id - if not fork_session and guard_key in self._thread_tool_contexts: - if requested_tool_context != self._thread_tool_contexts[guard_key]: + if not fork_session and thread_id in tool_contexts: + if requested_tool_context != tool_contexts[thread_id]: raise ValueError("tool_context_id and injected_job_kwargs cannot change when resuming a Codex thread") common = { @@ -348,19 +394,28 @@ class CodexAgentWrapper(BaseAgentWrapper): **common, ) if fork_session or not thread_id: - # Threads created under this binding are owned by this workspace. self._record_thread_ownership(alias, thread.id) - self._thread_tool_contexts[f"{alias}/{thread.id}" if alias else thread.id] = requested_tool_context + tool_contexts[thread.id] = requested_tool_context return thread # ----- Thread ownership (per-workspace, persisted) --------------------- - def _normalize_thread_alias(self, alias: str) -> str: - """Map a call alias to a concrete workspace alias ("" means the default).""" + def _agent_state(self, alias: str) -> WorkspaceAgentState | None: workspaces = _context_workspaces(self.app_context) if not workspaces: - return alias - return alias if alias in workspaces else next(iter(workspaces)) + return None + return workspaces[self._normalize_thread_alias(alias)].agent_state(self.name) + + def _tool_contexts(self, alias: str) -> dict[str, str]: + state = self._agent_state(alias) + return state.thread_tool_contexts if state is not None else self._thread_tool_contexts + + def _normalize_thread_alias(self, alias: str) -> str: + """Map an omitted alias to the default while preserving explicit aliases.""" + workspaces = _context_workspaces(self.app_context) + if workspaces and not alias: + return next(iter(workspaces)) + return alias def _thread_ownership_path(self, alias: str) -> Path | None: workspaces = _context_workspaces(self.app_context) @@ -372,9 +427,10 @@ class CodexAgentWrapper(BaseAgentWrapper): return workspace.mem_session_path / "codex" / "owned_threads.json" def _owned_thread_ids(self, alias: str) -> set[str]: - """Load (once) the persisted ids of threads this workspace created.""" - if alias in self._thread_owners: - return self._thread_owners[alias] + """Load the persisted ids into the owning workspace's agent state.""" + state = self._agent_state(alias) + if state is not None and state.owned_thread_ids is not None: + return state.owned_thread_ids ids: set[str] = set() path = self._thread_ownership_path(alias) if path is not None and path.is_file(): @@ -383,7 +439,8 @@ class CodexAgentWrapper(BaseAgentWrapper): ids = {str(thread_id) for thread_id in raw.get("thread_ids", [])} except (OSError, ValueError, AttributeError) as exc: self.logger.warning(f"Ignoring unreadable Codex thread ownership file {path}: {exc}") - self._thread_owners[alias] = ids + if state is not None: + state.owned_thread_ids = ids return ids def _record_thread_ownership(self, alias: str, thread_id: str) -> None: @@ -453,20 +510,19 @@ class CodexAgentWrapper(BaseAgentWrapper): return codex async def _close(self) -> None: - """Close the persistent app-server and the wrapper-owned job endpoint.""" + """Close the persistent app-server and clear session execution guards.""" 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._job_server is not None: - job_server, self._job_server = self._job_server, None - await job_server.close() + for workspace in (_context_workspaces(self.app_context) or {}).values(): + if state := workspace.agent_states.get(self.name): + state.thread_tool_contexts.clear() + if codex is not None: + await codex.close() - async def _prune_thread_tool_contexts(self, codex: AsyncCodex) -> None: + async def _prune_thread_tool_contexts(self, codex: AsyncCodex, alias: str = "") -> None: """Forget tool constraints only after the server confirms a thread is unloaded.""" + tool_contexts = self._tool_contexts(alias) try: from openai_codex.generated.v2_all import ThreadLoadedListResponse @@ -485,9 +541,9 @@ class CodexAgentWrapper(BaseAgentWrapper): 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 = { - key: context for key, context in self._thread_tool_contexts.items() if key.rsplit("/", 1)[-1] in loaded_ids - } + retained = {thread_id: context for thread_id, context in tool_contexts.items() if thread_id in loaded_ids} + tool_contexts.clear() + tool_contexts.update(retained) async def _release_thread(self, codex: AsyncCodex, thread_id: str) -> None: """Unsubscribe so the app-server can unload the thread once it is idle.""" diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index dd0d6ad9..1cba8a9b 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -1,25 +1,27 @@ -"""Thin FastMCP STDIO bridge forwarding Codex job tools to the parent process. +"""Expose ReMe jobs to Codex over FastMCP STDIO. -This subprocess never builds a ReMe Application or any workspace store. It -asks the parent-owned loopback endpoint (LoopbackJobServer) for job metadata -and forwards every tool call — together with the server-owned workspace alias, -tool context id, and injected kwargs — back to the parent for execution, so -all workspace state has exactly one owner. +Multi-workspace callers forward to the resident MCP service without creating +an Application. Single-workspace callers retain the legacy child-service path. """ import argparse +import asyncio import json +from pathlib import Path from typing import Any -import httpx +from .base_agent_wrapper import BaseAgentWrapper +from ..client import MCPClient +from ...config import resolve_app_config +from ...reme import ReMe def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Forward selected ReMe jobs from the parent loopback endpoint") - parser.add_argument("--endpoint", required=True, help="Parent loopback endpoint base URL") - parser.add_argument("--token", required=True, help="Bearer token for the loopback endpoint") + 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", help="Single-workspace directory for the legacy child service") + parser.add_argument("--endpoint", help="Resident streamable-HTTP MCP endpoint for forwarding") parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") - parser.add_argument("--workspace-alias", default="", help="Workspace alias injected into every job call") parser.add_argument("--tool-context-id", default="", help="Context id injected into every job call") parser.add_argument( "--injected-job-kwargs", @@ -30,89 +32,145 @@ def _parse_args() -> argparse.Namespace: return parser.parse_args() -class _EndpointClient: - """Minimal client for the parent's loopback job server.""" +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 legacy single-workspace child Application for 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: + job_config["enable_serve"] = True + jobs[name] = job_config - def __init__(self, endpoint: str, token: str, timeout: float) -> None: - self._client = httpx.Client( - base_url=endpoint, - headers={"Authorization": f"Bearer {token}"}, - timeout=timeout, - ) + missing = sorted(selected.difference(jobs)) + if missing: + raise KeyError(f"Codex job tools not found or not request jobs: {', '.join(missing)}") - def list_jobs(self) -> list[dict[str, Any]]: - """Fetch the servable job metadata from the parent endpoint.""" - response = self._client.get("/jobs") - response.raise_for_status() - return response.json()["jobs"] + 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 - def call(self, job: str, kwargs: dict[str, Any], injected_job_kwargs: dict[str, Any] | None = None) -> str: - """Forward one tool call to the parent endpoint and unwrap its answer.""" - response = self._client.post( - "/call", - json={"job": job, "kwargs": kwargs, "injected_job_kwargs": injected_job_kwargs or {}}, - ) - if response.status_code != 200: - try: - detail = response.json().get("error") or response.text - except ValueError: - detail = response.text - raise RuntimeError(str(detail)) - payload = response.json() - if not payload.get("success"): - raise RuntimeError(str(payload.get("answer") or payload)) - return str(payload.get("answer") or "") - - def close(self) -> None: - """Close the underlying HTTP client.""" - self._client.close() + prepared = dict(config) + prepared["jobs"] = jobs + prepared["service"] = service + return prepared -def main() -> None: - """Fetch job metadata from the parent and serve forwarding tools over STDIO.""" +class _ResidentMCPClient: + """Reuse the standard ReMe MCP client for resident tool discovery and calls.""" + + def __init__(self, endpoint: str, timeout: float) -> None: + from fastmcp.client import StreamableHttpTransport + + self._client = MCPClient(transport=StreamableHttpTransport(url=endpoint), timeout=timeout) + + async def start(self) -> None: + """Connect to the resident service.""" + await self._client.start() + + async def list_jobs(self) -> list[dict[str, Any]]: + """Return the resident service's tool descriptors.""" + return await self._client.list_actions() + + async def call(self, job: str, kwargs: dict[str, Any], injected: dict[str, Any]) -> str: + """Call one resident tool with server-owned arguments enforced.""" + # pylint: disable=protected-access + payload = BaseAgentWrapper._merge_injected_job_kwargs(kwargs, injected) + async for text in self._client(action=job, **payload): + return text + return "" + + async def close(self) -> None: + """Close the resident connection.""" + await self._client.close() + + +def _injected_kwargs(args: argparse.Namespace) -> dict[str, Any]: + injected = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else {} + if not isinstance(injected, dict): + raise TypeError("--injected-job-kwargs must be a JSON object") + injected = dict(injected) + if args.tool_context_id: + injected["tool_context_id"] = args.tool_context_id + return injected + + +async def _run_forwarder(args: argparse.Namespace, injected: dict[str, Any]) -> None: from fastmcp import FastMCP from fastmcp.exceptions import ToolError from fastmcp.tools import FunctionTool - args = _parse_args() - injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else {} - if not isinstance(injected_job_kwargs, dict): - raise TypeError("--injected-job-kwargs must be a JSON object") - injected = dict(injected_job_kwargs) - if args.workspace_alias: - injected["workspace"] = args.workspace_alias - if args.tool_context_id: - injected["tool_context_id"] = args.tool_context_id - - client = _EndpointClient(args.endpoint, args.token, args.tool_timeout) - metadata = {entry["name"]: entry for entry in client.list_jobs()} - missing = sorted(set(args.jobs) - set(metadata)) - if missing: - raise KeyError(f"Codex job tools not found on the parent endpoint: {', '.join(missing)}") - - server = FastMCP("reme-codex-bridge") - for name in args.jobs: - entry = metadata[name] - - def forward(_client=client, _name=name, _injected=None, **kwargs): - try: - return _client.call(_name, kwargs, _injected or injected) - except Exception as exc: # noqa: BLE001 - surface every failure as a tool error - raise ToolError(str(exc)) from exc - - server.add_tool( - FunctionTool( - name=name, - description=entry.get("description") or "", - parameters=entry.get("parameters") or {"type": "object", "properties": {}}, - fn=forward, - ), - ) - + client = _ResidentMCPClient(args.endpoint, args.tool_timeout) + await client.start() try: - server.run(transport="stdio") + metadata = {entry["name"]: entry for entry in await client.list_jobs()} + missing = sorted(set(args.jobs) - set(metadata)) + if missing: + raise KeyError(f"Codex job tools not found on the resident MCP service: {', '.join(missing)}") + + server = FastMCP("reme-codex-bridge") + for name in args.jobs: + entry = metadata[name] + + def make_forwarder(job_name: str): + async def forward(**kwargs): + try: + return await client.call(job_name, kwargs, injected) + except Exception as exc: # noqa: BLE001 - surface every failure as a tool error + raise ToolError(str(exc)) from exc + + return forward + + parameters = entry.get("inputSchema") or entry.get("parameters") + # pylint: disable=protected-access + parameters = BaseAgentWrapper._strip_injected_parameters(parameters, injected) + server.add_tool( + FunctionTool( + name=name, + description=entry.get("description") or "", + parameters=parameters or {"type": "object", "properties": {}}, + fn=make_forwarder(name), + ), + ) + await server.run_async(transport="stdio") finally: - client.close() + await client.close() + + +def main() -> None: + """Forward to a resident service or start the legacy single-workspace child service.""" + args = _parse_args() + injected = _injected_kwargs(args) + if args.endpoint: + asyncio.run(_run_forwarder(args, injected)) + return + if not args.workspace: + raise ValueError("--workspace is required when --endpoint is not supplied") + 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, + ) + config = _prepare_config(config, args.jobs, args.tool_context_id, injected) + ReMe(**config).run_app() if __name__ == "__main__": diff --git a/reme/components/agent_wrapper/loopback_job_server.py b/reme/components/agent_wrapper/loopback_job_server.py deleted file mode 100644 index fa242336..00000000 --- a/reme/components/agent_wrapper/loopback_job_server.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Wrapper-owned loopback job endpoint for agent-bridge subprocesses. - -The Codex job bridge used to spawn a full child Application per thread; every -thread then held its own workspace stores and became a second persistence -writer. This endpoint keeps job execution in the parent process instead: a -minimal HTTP/JSON server on 127.0.0.1 with an ephemeral port and a bearer -token, owned by one CodexAgentWrapper instance for its lifetime. Bridge -subprocesses are thin FastMCP forwarders that call this endpoint. -""" - -import asyncio -import json -import secrets -from typing import Any, TYPE_CHECKING - -if TYPE_CHECKING: - from ..application_context import ApplicationContext - from ..job import BaseJob - -_REASON = { - 200: "OK", - 400: "Bad Request", - 401: "Unauthorized", - 404: "Not Found", - 500: "Internal Server Error", -} - - -class LoopbackJobServer: - """Serve this application's request jobs to trusted local subprocesses.""" - - def __init__(self, app_context: "ApplicationContext") -> None: - self._app_context = app_context - self.token = secrets.token_urlsafe(24) - self._server: asyncio.Server | None = None - - @property - def is_running(self) -> bool: - """Whether the endpoint is currently bound.""" - return self._server is not None - - @property - def url(self) -> str: - """Base URL of the bound endpoint.""" - if self._server is None: - raise RuntimeError("Loopback job server is not started") - port = self._server.sockets[0].getsockname()[1] - return f"http://127.0.0.1:{port}" - - async def start(self) -> None: - """Bind the endpoint; idempotent so lazy bridge creation can retry.""" - if self._server is None: - self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0) - - async def close(self) -> None: - """Stop the endpoint and release the port.""" - server, self._server = self._server, None - if server is not None: - server.close() - await server.wait_closed() - - # ----- HTTP handling --------------------------------------------------- - - async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - try: - head = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=30) - request_line, *header_lines = head.decode("latin-1").split("\r\n") - method, target, _version = request_line.split(" ", 2) - headers = {} - for line in header_lines: - if ":" in line: - key, value = line.split(":", 1) - headers[key.strip().lower()] = value.strip() - length = int(headers.get("content-length", 0) or 0) - body = await reader.readexactly(length) if length else b"" - status, payload = await self._dispatch(method, target, headers, body) - except Exception as exc: # pylint: disable=broad-exception-caught - status, payload = 500, {"error": f"{type(exc).__name__}: {exc}"} - data = json.dumps(payload).encode() - writer.write( - f"HTTP/1.1 {status} {_REASON.get(status, 'OK')}\r\n" - "content-type: application/json\r\n" - f"content-length: {len(data)}\r\n" - "connection: close\r\n\r\n".encode() + data, - ) - try: - await writer.drain() - finally: - writer.close() - - async def _dispatch( - self, - method: str, - target: str, - headers: dict[str, str], - body: bytes, - ) -> tuple[int, dict[str, Any]]: - if headers.get("authorization") != f"Bearer {self.token}": - return 401, {"error": "unauthorized"} - if method == "GET" and target == "/jobs": - return 200, {"jobs": self._job_metadata()} - if method == "POST" and target == "/call": - try: - request = json.loads(body or b"{}") - except json.JSONDecodeError as exc: - return 400, {"error": f"invalid JSON body: {exc}"} - return await self._call(request) - return 404, {"error": "not found"} - - # ----- Job surface ------------------------------------------------------- - - def _request_jobs(self) -> dict[str, "BaseJob"]: - """Return the application's non-stream, non-background jobs by name.""" - from ..job import BackgroundJob, StreamJob - - return { - name: job for name, job in self._app_context.jobs.items() if not isinstance(job, (StreamJob, BackgroundJob)) - } - - def _job_metadata(self) -> list[dict[str, Any]]: - """Describe every servable job for bridge tool registration.""" - return [ - {"name": name, "description": job.description, "parameters": job.parameters} - for name, job in sorted(self._request_jobs().items()) - ] - - async def _call(self, request: dict[str, Any]) -> tuple[int, dict[str, Any]]: - name = str(request.get("job") or "") - job = self._request_jobs().get(name) - if job is None: - return 404, {"error": f"job {name!r} not found or not a request job"} - - kwargs = dict(request.get("kwargs") or {}) - injected = dict(request.get("injected_job_kwargs") or {}) - conflicts = sorted(injected.keys() & kwargs.keys()) - if conflicts: - return 400, {"error": f"{', '.join(conflicts)} injected by the server and cannot be provided"} - kwargs.update(injected) - if workspace := request.get("workspace"): - kwargs["workspace"] = workspace - if tool_context_id := request.get("tool_context_id"): - kwargs["tool_context_id"] = tool_context_id - - try: - response = await job(**kwargs) - except ValueError as exc: - return 400, {"error": str(exc)} - return 200, {"success": response.success, "answer": response.answer} diff --git a/reme/components/application_context.py b/reme/components/application_context.py index e52475b6..e2b617be 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -1,5 +1,6 @@ """Application context: shared state container for components, jobs, and service.""" +from collections.abc import Iterator, MutableMapping from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any @@ -12,65 +13,80 @@ if TYPE_CHECKING: from .service import BaseService -class ApplicationContext: - """Passive state container holding parsed config and wired components. +class _LegacyComponentView(MutableMapping): + """Compatibility mapping combining shared components with default-workspace instances.""" - The Application class performs the actual wiring (registry lookups and - component instantiation); this class only stores the results so that - components, jobs, and the service can find each other at runtime. - """ + def __init__(self, context: "ApplicationContext") -> None: + self._context = context + + def _default_components(self) -> dict: + return next(iter(self._context.workspaces.values())).components + + def __getitem__(self, key: str) -> dict: + if key in self._context.shared_components: + return self._context.shared_components[key] + return self._default_components()[key] + + def __setitem__(self, key: str, value: dict) -> None: + if value and all(getattr(component, "workspace_scoped", False) is True for component in value.values()): + self._default_components()[key] = value + self._context.shared_components.pop(key, None) + else: + self._context.shared_components[key] = value + self._default_components().pop(key, None) + + def __delitem__(self, key: str) -> None: + if key in self._context.shared_components: + del self._context.shared_components[key] + else: + del self._default_components()[key] + + def __iter__(self) -> Iterator[str]: + yield from self._context.shared_components + yield from (key for key in self._default_components() if key not in self._context.shared_components) + + def __len__(self) -> int: + return len(set(self._context.shared_components) | set(self._default_components())) + + +class ApplicationContext: + """Application-wide configuration, shared resources, and workspace registry.""" def __init__(self, *, registry: ComponentRegistry | None = None, **kwargs): - # Parse raw kwargs into a typed, validated config object. self.app_config: ApplicationConfig = ApplicationConfig(**kwargs) self.registry = registry or create_application_registry() - - # Ordered per-workspace path views; the first entry is the default workspace. self.workspaces: dict[str, WorkspaceContext] = { alias: WorkspaceContext.from_config(alias, path, self.app_config) for alias, path in self.app_config.workspaces.items() } - self.workspace: WorkspaceContext = next(iter(self.workspaces.values())) - # Populated by Application during initialization. + self.workspace: WorkspaceContext = next(iter(self.workspaces.values())) self.service: "BaseService | None" = None - self.components: dict[str, dict[str, "BaseComponent"]] = {} + self.shared_components: dict[str, dict[str, "BaseComponent"]] = {} + self._components = _LegacyComponentView(self) self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None - - # Application-lifetime shared state. Values remain available across Job and Step - # invocations while this Application is running, so Jobs may keep cross-call state here. - # This is in-memory state, not durable storage; use workspace files or a store when state - # must survive an Application restart. self.metadata: dict[str, Any] = {} @property - def workspace_components(self) -> dict[str, dict[str, dict[str, "BaseComponent"]]]: - """Compatibility view of scoped instances, ctype -> name -> alias. + def components(self) -> MutableMapping: + """Legacy component lookup; internal execution uses explicit shared or workspace ownership.""" + return self._components - The authoritative owner of these instances is each WorkspaceContext; - this view exists for transitional readers. - """ - grouped: dict[str, dict[str, dict[str, "BaseComponent"]]] = {} - for alias, workspace in self.workspaces.items(): - for ctype, group in workspace.components.items(): - for name, instance in group.items(): - grouped.setdefault(ctype, {}).setdefault(name, {})[alias] = instance - return grouped + @components.setter + def components(self, value: dict) -> None: + self.shared_components = {} + for ctype, group in value.items(): + self._components[ctype] = group def get_component(self, ctype: str, name: str, alias: str | None = None) -> "BaseComponent | None": - """Resolve a component, honoring workspace bindings for scoped components. - - Returns None when the component is not configured at all. Raises - KeyError when a workspace-scoped component has no instance for the - requested alias — a missing scoped resource must fail, never fall back. - """ + """Compatibility resolver for external callers that select a workspace by alias.""" alias = alias or next(iter(self.workspaces)) workspace = self.workspaces.get(alias) if workspace is None: raise KeyError(f"Unknown workspace alias {alias!r}") if name in workspace.components.get(ctype, {}): - return workspace.components[ctype][name] + return workspace.get_component(ctype, name) if any(name in ws.components.get(ctype, {}) for ws in self.workspaces.values()): raise KeyError(f"Component '{name}' in {ctype} has no instance for workspace alias {alias!r}") - return self.components.get(ctype, {}).get(name) + return self.shared_components.get(ctype, {}).get(name) diff --git a/reme/components/base_component.py b/reme/components/base_component.py index 1f704bfb..f4075639 100644 --- a/reme/components/base_component.py +++ b/reme/components/base_component.py @@ -221,20 +221,24 @@ class BaseComponent(ComponentMixin, ABC): setattr(self, attr, None) def _resolve_from_context(self, attr: str, dep: Dependency) -> None: - """Context-bound mode: look up the component from ``app_context.components``. - - Workspace-scoped dependencies resolve to the same-alias instance as the - dependent; a missing alias instance is an error, never a fallback. - """ + """Resolve scoped dependencies from their workspace and shared ones from the application.""" workspaces = getattr(self.app_context, "workspaces", None) workspace = workspaces.get(self.workspace_alias) if workspaces and self.workspace_alias is not None else None - scoped_group = workspace.components.get(dep.ctype) if workspace is not None else None - if scoped_group is not None: - target = scoped_group.get(dep.name) - if target is None: - raise ValueError(f"{dep.ctype} '{dep.name}' has no instance for workspace '{self.workspace_alias}'.") + if workspace is not None and dep.name in workspace.components.get(dep.ctype, {}): + target = workspace.get_component(dep.ctype, dep.name) + elif workspace is not None and any( + dep.name in candidate.components.get(dep.ctype, {}) for candidate in workspaces.values() + ): + raise ValueError(f"{dep.ctype} '{dep.name}' has no instance for workspace '{self.workspace_alias}'.") else: - target = self.app_context.components.get(dep.ctype, {}).get(dep.name) + shared = getattr(self.app_context, "shared_components", None) + if shared is None: + shared = self.app_context.components + target = shared.get(dep.ctype, {}).get(dep.name) + if target is None and workspaces: + default_workspace = next(iter(workspaces.values())) + if dep.name in default_workspace.components.get(dep.ctype, {}): + target = default_workspace.get_component(dep.ctype, dep.name) if target is not None: setattr(self, attr, target) elif dep.optional: diff --git a/reme/schema/workspace.py b/reme/schema/workspace.py index f30cb92b..30f3070c 100644 --- a/reme/schema/workspace.py +++ b/reme/schema/workspace.py @@ -9,17 +9,17 @@ if TYPE_CHECKING: from .application_config import ApplicationConfig +@dataclass +class WorkspaceAgentState: + """Session execution state owned by one agent wrapper within one workspace.""" + + thread_tool_contexts: dict[str, str] = field(default_factory=dict) + owned_thread_ids: set[str] | None = None + + @dataclass(frozen=True) class WorkspaceContext: - """Owns everything scoped to one workspace of an application. - - Built from ApplicationConfig when the application context is created. - Beyond the resolved paths, the workspace owns its stateful component - instances (stores, graphs, indexes, catalogs), in-memory execution state - (drafts, search dedup, call budgets), session ownership records, and - workspace-scoped locks. Shared, stateless clients stay on the - ApplicationContext; per-call bindings reference this object. - """ + """Owns the paths, stateful components, execution state, and locks of one workspace.""" alias: str workspace_path: Path @@ -29,17 +29,16 @@ class WorkspaceContext: resource_path: Path | None daily_path: Path | None digest_path: Path | None - # Workspace-scoped component instances keyed by component type then name. - # Values are BaseComponent instances; typed as Any to keep this module - # free of component imports (schema must not depend on components). components: dict[str, dict[str, Any]] = field(default_factory=dict) - # Drafts and search dedup state, keyed by tool_context_id. tool_contexts: dict[str, Any] = field(default_factory=dict) - # Search call budgets keyed by tool_context_id. search_call_budgets: dict[str, Any] = field(default_factory=dict) - # Workspace-scoped named asyncio locks (e.g. dream integration). + agent_states: dict[str, WorkspaceAgentState] = field(default_factory=dict) locks: dict[str, asyncio.Lock] = field(default_factory=dict) + def agent_state(self, wrapper_name: str) -> WorkspaceAgentState: + """Return the named wrapper's state within this workspace.""" + return self.agent_states.setdefault(wrapper_name, WorkspaceAgentState()) + @classmethod def from_config(cls, alias: str, workspace: str, config: "ApplicationConfig") -> "WorkspaceContext": """Resolve one alias of the config's workspaces mapping into absolute paths.""" @@ -59,6 +58,13 @@ class WorkspaceContext: digest_path=child(config.digest_dir), ) + def get_component(self, ctype: str, name: str) -> Any: + """Return a component owned by this workspace, never a default or shared substitute.""" + try: + return self.components[ctype][name] + except KeyError as exc: + raise KeyError(f"Component '{name}' in {ctype} has no instance for workspace '{self.alias}'") from exc + def storage_dirs(self) -> list[Path]: """Subdirectories the application materializes at startup, in configuration order.""" return [ diff --git a/reme/steps/base_step.py b/reme/steps/base_step.py index 15702fad..c70beb97 100644 --- a/reme/steps/base_step.py +++ b/reme/steps/base_step.py @@ -87,14 +87,16 @@ class Ref: if self.optional: return None raise RuntimeError(f"app_context is not set when resolving '{self.key}'") - if _context_workspaces(obj.app_context) is not None: + if (workspaces := _context_workspaces(obj.app_context)) is not None: context = obj.context bound = getattr(context, "workspace", None) if context is not None else None - comp = obj.app_context.get_component( - self.comp_enum, - name, - alias=bound.alias if bound is not None else None, - ) + is_scoped = any(name in workspace.components.get(self.comp_enum, {}) for workspace in workspaces.values()) + if is_scoped: + if bound is None: + raise RuntimeError(f"A workspace binding is required when resolving '{self.key}'") + comp = bound.get_component(self.comp_enum, name) + else: + comp = obj.app_context.shared_components.get(self.comp_enum, {}).get(name) else: # Test doubles and standalone-style contexts without the workspace layer. comp = obj.app_context.components.get(self.comp_enum, {}).get(name) @@ -159,6 +161,8 @@ class BaseStep(ComponentMixin, ABC): del self.__dict__[key] self.context = RuntimeContext.from_context(context, **kwargs) assert self.context is not None + if (workspaces := _context_workspaces(self.app_context)) is not None and self.context.workspace is None: + self.context.workspace = next(iter(workspaces.values())) if self.input_mapping: self.context.apply_mapping(self.input_mapping) result = await self.execute() diff --git a/reme/steps/evolve/proactive/topics.py b/reme/steps/evolve/proactive/topics.py index a22156cc..d49caa9b 100644 --- a/reme/steps/evolve/proactive/topics.py +++ b/reme/steps/evolve/proactive/topics.py @@ -343,11 +343,9 @@ class ProactiveTopicsStep(BaseStep): return [] bound = getattr(self.context, "workspace", None) if self.context is not None else None if _context_workspaces(self.app_context) is not None: - catalog = self.app_context.get_component( - ComponentEnum.FILE_CATALOG, - "digest", - alias=bound.alias if bound is not None else None, - ) + if bound is None: + raise RuntimeError("A workspace binding is required to resolve the digest catalog") + catalog = bound.components.get(ComponentEnum.FILE_CATALOG, {}).get("digest") else: # Test doubles and standalone-style contexts without the workspace layer. catalog = self.app_context.components.get(ComponentEnum.FILE_CATALOG, {}).get("digest") @@ -368,7 +366,7 @@ class ProactiveTopicsStep(BaseStep): name = self.kwargs.get("as_embedding", "default") if self.app_context is None: return None - return self.app_context.components.get(ComponentEnum.AS_EMBEDDING, {}).get(name) + return self.app_context.shared_components.get(ComponentEnum.AS_EMBEDDING, {}).get(name) async def _semantic_verdict( self, diff --git a/reme/steps/evolve/proactive/utils.py b/reme/steps/evolve/proactive/utils.py index a4a5b0ef..8839cc33 100644 --- a/reme/steps/evolve/proactive/utils.py +++ b/reme/steps/evolve/proactive/utils.py @@ -681,7 +681,7 @@ def resolve_agent_wrapper(step): if wrapper is not None: return wrapper if step.app_context is not None: - fallback = step.app_context.components.get(ComponentEnum.AGENT_WRAPPER, {}).get("default") + fallback = step.app_context.shared_components.get(ComponentEnum.AGENT_WRAPPER, {}).get("default") if fallback is not None: configured = step.kwargs.get("agent_wrapper", "default") step.logger.warning(f"[{step.name}] agent_wrapper '{configured}' missing; using default") diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 7afa9fed..04c39206 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -12,6 +12,7 @@ from watchfiles import Change from ._change_batch import bucket_changes from ..base_step import BaseStep from ...components import R +from ...components.base_component import _context_workspaces from ...components.file_chunker import BaseFileChunker from ...enumeration import ComponentEnum from ...schema import FileChunk, FileNode @@ -368,8 +369,10 @@ class UpdateIndexStep(ChangeApplyStep): bound = getattr(self.context, "workspace", None) if self.context is not None else None if bound is not None: chunkers: dict[str, BaseFileChunker] = bound.components.get(ComponentEnum.FILE_CHUNKER, {}) - else: + elif _context_workspaces(self.app_context) is None: chunkers = self.app_context.components[ComponentEnum.FILE_CHUNKER] + else: + raise RuntimeError("A workspace binding is required to resolve file chunkers") suffix = path.suffix.lstrip(".").lower() for candidate in chunkers.values(): if suffix and suffix in {ext.lower().lstrip(".") for ext in candidate.supported_extensions}: diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index b0da6d96..c498fc41 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -61,60 +61,51 @@ def _wrapper(tmp_path, **kwargs): return CodexAgentWrapper(app_context=context, **kwargs), job -@pytest.mark.asyncio -async def test_mcp_config_uses_loopback_bridge_and_selected_jobs(tmp_path): +def test_mcp_config_uses_legacy_child_bridge_and_selected_jobs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - await wrapper._ensure_job_server() # pylint: disable=protected-access - config = wrapper._mcp_server_config( # pylint: disable=protected-access + config = wrapper._mcp_server_config( {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, ) - await wrapper.close() 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("--endpoint") + 1].startswith("http://127.0.0.1:") - assert config["args"][config["args"].index("--token") + 1] + assert config["args"][config["args"].index("--config") + 1] == str(tmp_path / "custom.yaml") + assert config["args"][config["args"].index("--workspace") + 1] == str(tmp_path) assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" - # mcp_config is accepted for compatibility but no longer builds a child app. - assert "--config" not in config["args"] + assert "--endpoint" not in config["args"] + assert "--token" not in config["args"] -@pytest.mark.asyncio -async def test_mcp_config_serializes_injected_job_kwargs(tmp_path): +def test_mcp_config_serializes_injected_job_kwargs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - await wrapper._ensure_job_server() # pylint: disable=protected-access - config = wrapper._mcp_server_config( # pylint: disable=protected-access + config = wrapper._mcp_server_config( { "job_tools": ["search"], "tool_context_id": "ctx-1", "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]}, }, ) - await wrapper.close() 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 + plain = wrapper._mcp_server_config({"job_tools": ["search"]}) assert "--injected-job-kwargs" not in plain["args"] -@pytest.mark.asyncio -async def test_thread_config_preserves_other_mcp_servers(tmp_path): +def test_thread_config_preserves_other_mcp_servers(tmp_path): wrapper, _job = _wrapper(tmp_path) - await wrapper._ensure_job_server() # pylint: disable=protected-access - config = wrapper._thread_config( # pylint: disable=protected-access + config = wrapper._thread_config( { "job_tools": ["search"], "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}}, }, ) - await wrapper.close() assert "docs" in config["mcp_servers"] assert len(config["mcp_servers"]) == 2 @@ -166,65 +157,15 @@ def test_mcp_config_rejects_background_jobs(tmp_path): wrapper._mcp_server_config({"job_tools": ["watch"]}) -@pytest.mark.asyncio -async def test_loopback_endpoint_serves_parent_jobs(tmp_path): - import httpx - - wrapper, job = _wrapper(tmp_path) - await wrapper.start() - server = await wrapper._ensure_job_server() # pylint: disable=protected-access - base = server.url - headers = {"Authorization": f"Bearer {server.token}"} - try: - async with httpx.AsyncClient(base_url=base, headers=headers, timeout=10) as http: - listed = (await http.get("/jobs")).json()["jobs"] - assert [entry["name"] for entry in listed] == ["search"] - - called = await http.post("/call", json={"job": "search", "kwargs": {"query": "memo"}}) - assert called.json() == {"success": True, "answer": "found:memo"} - # The job executed in the parent process, not in any child. - assert job.calls == [{"query": "memo"}] - - missing = await http.post("/call", json={"job": "missing", "kwargs": {}}) - assert missing.status_code == 404 - - unauthorized = await http.get("/jobs", headers={"Authorization": "Bearer wrong"}) - assert unauthorized.status_code == 401 - finally: - await wrapper.close() - assert server.is_running is False - - -@pytest.mark.asyncio -async def test_loopback_bridge_lists_and_calls_parent_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - wrapper, job = _wrapper(tmp_path) - await wrapper._ensure_job_server() # pylint: disable=protected-access - server_config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access - transport = StdioTransport( - command=server_config["command"], - args=server_config["args"], - cwd=server_config["cwd"], - 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] == ["search"] - result = await asyncio.wait_for(client.call_tool("search", {"query": "via-bridge"}), timeout=10) - await wrapper.close() - - assert "found:via-bridge" in str(result.content[0].text) - assert job.calls == [{"query": "via-bridge"}] - - @pytest.mark.asyncio async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): - wrapper, _job = _wrapper(tmp_path) - await wrapper._ensure_job_server() # pylint: disable=protected-access - server_config = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps({"jobs": {"search": {"backend": "base", "steps": []}}}), + encoding="utf-8", + ) + wrapper, _job = _wrapper(tmp_path, mcp_config=str(config_path)) + server_config = wrapper._mcp_server_config({"job_tools": ["search"]}) proc = await asyncio.create_subprocess_exec( server_config["command"], *server_config["args"], diff --git a/tests/unit/test_workspace_binding.py b/tests/unit/test_workspace_binding.py index f7122489..008b78d5 100644 --- a/tests/unit/test_workspace_binding.py +++ b/tests/unit/test_workspace_binding.py @@ -165,7 +165,12 @@ async def test_workspace_scoped_components_instantiate_per_alias(tmp_path): store_beta = app.context.get_component(ComponentEnum.FILE_STORE, "default", "beta") assert store_alpha is not store_beta - assert app.context.components[ComponentEnum.FILE_STORE]["default"] is store_alpha # default alias entry + assert app.context.components[ComponentEnum.FILE_STORE]["default"] is store_alpha + assert ComponentEnum.FILE_STORE not in app.context.shared_components + assert app.context.workspaces["alpha"].get_component(ComponentEnum.FILE_STORE, "default") is store_alpha + assert app.context.workspaces["beta"].get_component(ComponentEnum.FILE_STORE, "default") is store_beta + with pytest.raises(KeyError, match="no instance for workspace 'alpha'"): + app.context.workspaces["alpha"].get_component(ComponentEnum.FILE_STORE, "missing") assert store_alpha.workspace_path == tmp_path / "alpha" assert store_beta.workspace_path == tmp_path / "beta" assert store_alpha.chunks_path.parent == tmp_path / "alpha" / "metadata" / "file_store" @@ -314,7 +319,8 @@ async def test_digest_titles_follow_workspace_binding(tmp_path): assert await step._digest_titles() == ["beta-only"] step.context = RuntimeContext() - assert await step._digest_titles() == ["alpha-only"] + with pytest.raises(RuntimeError, match="workspace binding is required"): + await step._digest_titles() @pytest.mark.asyncio diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index 81bc8a72..db00f606 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -3,6 +3,8 @@ # pylint: disable=protected-access,missing-function-docstring import asyncio +import json +import os from types import SimpleNamespace import pytest @@ -128,11 +130,8 @@ async def test_tool_context_state_namespaced_per_workspace(tmp_path): @pytest.mark.asyncio async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): - wrapper = CodexAgentWrapper( - app_context=ApplicationContext( - workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, - ), - ) + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + wrapper = CodexAgentWrapper(app_context=context) class FakeCodex: """Minimal Codex double echoing the requested thread id.""" @@ -158,7 +157,13 @@ async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): await wrapper._open_thread(codex, {"session_id": "thread-1", "workspace": "beta"}) - assert sorted(wrapper._thread_tool_contexts) == ["beta/thread-2", "thread-1"] + alpha_state = context.workspaces["alpha"].agent_state(wrapper.name) + beta_state = context.workspaces["beta"].agent_state(wrapper.name) + assert set(alpha_state.thread_tool_contexts) == {"thread-1"} + assert set(beta_state.thread_tool_contexts) == {"thread-2"} + assert not wrapper._thread_tool_contexts + assert alpha_state.owned_thread_ids == {"thread-1"} + assert beta_state.owned_thread_ids == {"thread-2"} @pytest.mark.asyncio @@ -295,27 +300,24 @@ async def test_mcp_tool_boundary_forwards_injected_workspace(tmp_path): @pytest.mark.asyncio -async def test_codex_bridge_uses_loopback_endpoint(tmp_path): - import json +async def test_codex_bridge_uses_resident_mcp_endpoint(tmp_path): + from reme.components.service import HttpService context = _context(tmp_path) + context.service = HttpService(app_context=context) context.jobs = {"search_job": _job(context, "search_job")} wrapper = CodexAgentWrapper(app_context=context) - await wrapper.start() - await wrapper._ensure_job_server() # pylint: disable=protected-access config = wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "beta"}) - assert "--workspace-alias" in config["args"] - assert config["args"][config["args"].index("--workspace-alias") + 1] == "beta" + assert config["args"][config["args"].index("--endpoint") + 1] == "http://127.0.0.1:2333/mcp" raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] assert json.loads(raw) == {"workspace": "beta"} assert "--config" not in config["args"] and "--workspace" not in config["args"] + assert "--token" not in config["args"] default_config = wrapper._mcp_server_config({"job_tools": ["search_job"]}) - assert "--workspace-alias" not in default_config["args"] - + assert "--injected-job-kwargs" not in default_config["args"] await wrapper.close() - assert wrapper._job_server is None # ----- Review follow-ups: explicit bindings, session isolation, key collisions @@ -423,18 +425,13 @@ async def test_codex_thread_cwd_follows_bound_workspace(tmp_path): # ----- Codex bridge: real child-process execution ---------------------------- -@pytest.mark.asyncio -async def test_codex_bridge_reads_resident_parent_state(tmp_path): - """Bridge calls execute against the parent's live stores; there is no second owner.""" - from fastmcp import Client - from fastmcp.client import StdioTransport - - app = Application( +def _codex_test_app(tmp_path): + return Application( workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, enable_logo=False, log_to_console=False, log_to_file=False, - service={"backend": "cli"}, + service={"backend": "http", "host": "127.0.0.1", "port": 0, "web_enabled": False}, components={ "tokenizer": {"default": {"backend": "regex"}}, "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, @@ -464,57 +461,104 @@ async def test_codex_bridge_reads_resident_parent_state(tmp_path): }, }, ) + + +async def _start_resident_http(app): + import uvicorn + + service = app.context.service + service.build_service(app) + service.add_jobs(app) + config = uvicorn.Config(service.service, host=service.host, port=0, log_level="error", lifespan="on") + server = uvicorn.Server(config) + task = asyncio.create_task(server.serve()) + for _ in range(100): + if server.started: + break + await asyncio.sleep(0.05) + assert server.started + service.port = server.servers[0].sockets[0].getsockname()[1] + return server, task + + +async def _stop_resident_http(server, task): + server.should_exit = True + await asyncio.wait_for(task, timeout=30) + + +@pytest.mark.asyncio +async def test_codex_bridge_reads_resident_parent_state(tmp_path): + """The stdio bridge reuses the resident MCP service and its live workspace state.""" + from fastmcp import Client + from fastmcp.client import StdioTransport + + app = _codex_test_app(tmp_path) wrapper = CodexAgentWrapper(app_context=app.context) async with app: - await wrapper.start() - # Index a note directly in the parent's beta store. - await app.run_job("write_note", path="daily/resident.md", content="parent-resident-beta", workspace="beta") - await wrapper._ensure_job_server() # pylint: disable=protected-access - server_config = wrapper._mcp_server_config( # pylint: disable=protected-access - {"job_tools": ["read_note", "write_note"], "workspace": "beta"}, - ) - transport = StdioTransport( - command=server_config["command"], - args=server_config["args"], - cwd=server_config["cwd"], - keep_alive=False, - ) - async with Client(transport, timeout=None, init_timeout=30) as client: - read = await asyncio.wait_for(client.call_tool("read_note", {"path": "daily/resident.md"}), timeout=30) - await asyncio.wait_for( - client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-bridge-beta"}), - timeout=30, + server, task = await _start_resident_http(app) + try: + await app.run_job("write_note", path="daily/resident.md", content="parent-resident-beta", workspace="beta") + from reme.components.client import MCPClient + + resident_client = MCPClient(transport="streamable-http", host="127.0.0.1", port=app.context.service.port) + await resident_client.start() + assert {entry["name"] for entry in await resident_client.list_actions()} == {"read_note", "write_note"} + await resident_client.close() + server_config = wrapper._mcp_server_config( + {"job_tools": ["read_note", "write_note"], "workspace": "beta"}, ) - # The bridge saw the parent's resident beta state and wrote through it. - assert "parent-resident-beta" in str(read.content[0].text) - assert "via-bridge-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") - assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() - store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") - assert all("bridge" not in node.path for node in await store_alpha.get_nodes()) + transport = StdioTransport( + command=server_config["command"], + args=server_config["args"], + cwd=server_config["cwd"], + keep_alive=False, + env=dict(os.environ), + ) + async with Client(transport, timeout=None, init_timeout=30) as client: + tools = await client.list_tools() + assert {tool.name for tool in tools} == {"read_note", "write_note"} + read = await asyncio.wait_for(client.call_tool("read_note", {"path": "daily/resident.md"}), timeout=30) + with pytest.raises(Exception): + await client.call_tool("read_note", {"path": "daily/resident.md", "workspace": "alpha"}) + await asyncio.wait_for( + client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-bridge-beta"}), + timeout=30, + ) + assert "parent-resident-beta" in str(read.content[0].text) + assert "via-bridge-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") + assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() + store_alpha = app.context.get_component(ComponentEnum.FILE_STORE, "default", "alpha") + assert all("bridge" not in node.path for node in await store_alpha.get_nodes()) + + invalid_config = wrapper._mcp_server_config({"job_tools": ["read_note"], "workspace": "nope"}) + invalid_transport = StdioTransport( + command=invalid_config["command"], + args=invalid_config["args"], + cwd=invalid_config["cwd"], + keep_alive=False, + env=dict(os.environ), + ) + async with Client(invalid_transport, timeout=None, init_timeout=30) as client: + with pytest.raises(Exception, match="Unknown workspace 'nope'"): + await asyncio.wait_for(client.call_tool("read_note", {"path": "daily/resident.md"}), timeout=30) + finally: + await _stop_resident_http(server, task) await wrapper.close() @pytest.mark.asyncio -async def test_codex_bridge_unknown_alias_fails_at_execution(tmp_path): - """An unknown workspace alias surfaces as a tool error from the parent boundary.""" - from fastmcp import Client - from fastmcp.client import StdioTransport +async def test_codex_multi_workspace_rejects_unsupported_job_transport(tmp_path): + from reme.components.service import CliService context = _context(tmp_path) + context.service = CliService(app_context=context) context.jobs = {"search_job": _job(context, "search_job")} - wrapper = CodexAgentWrapper(app_context=context, project_path=str(tmp_path)) - await wrapper.start() - await wrapper._ensure_job_server() # pylint: disable=protected-access - server_config = wrapper._mcp_server_config( # pylint: disable=protected-access - {"job_tools": ["search_job"], "workspace": "nope"}, - ) - transport = StdioTransport( - command=server_config["command"], - args=server_config["args"], - cwd=server_config["cwd"], - keep_alive=False, - ) - async with Client(transport, timeout=None, init_timeout=30) as client: - with pytest.raises(Exception, match="Unknown workspace 'nope'"): - await asyncio.wait_for(client.call_tool("search_job", {}), timeout=30) + wrapper = CodexAgentWrapper(app_context=context) + + with pytest.raises(RuntimeError, match="resident HTTP MCP service"): + wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "beta"}) + with pytest.raises(ValueError, match="only supported for single-workspace"): + wrapper._mcp_server_config( + {"job_tools": ["search_job"], "workspace": "beta", "mcp_config": "custom.yaml"}, + ) await wrapper.close() From f9a957beeed9bd036e249f29e2ca7ec19093114f Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 29 Oct 2026 14:06:02 +0800 Subject: [PATCH 10/21] fix(workspace): validate derived storage ownership --- reme/schema/application_config.py | 17 +++++++ tests/unit/test_multi_workspace_config.py | 58 ++++++++++++++++++++++- 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index 6894c346..8d7bcbd6 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -114,6 +114,23 @@ class ApplicationConfig(BaseModel): def derive_default_workspace(self): """Normalize every workspace_dir form to an ordered mapping; the first entry is the default.""" self.workspaces = dict(self.workspaces) or {DEFAULT_WORKSPACE_ALIAS: self.workspace_dir} + if len(self.workspaces) > 1: + for alias, workspace in self.workspaces.items(): + root = Path(workspace).resolve(strict=False) + for field_name in ( + "metadata_dir", + "session_dir", + "mem_session_dir", + "resource_dir", + "daily_dir", + "digest_dir", + ): + subdir = getattr(self, field_name) + if subdir and not (root / subdir).resolve(strict=False).is_relative_to(root): + raise ValueError( + f"{field_name} for workspace {alias!r} must stay inside the workspace " + "in multi-workspace configurations", + ) return self @field_validator("components", mode="before") diff --git a/tests/unit/test_multi_workspace_config.py b/tests/unit/test_multi_workspace_config.py index e1b2199e..cd61b7fb 100644 --- a/tests/unit/test_multi_workspace_config.py +++ b/tests/unit/test_multi_workspace_config.py @@ -11,7 +11,8 @@ import yaml from reme.application import Application from reme.components import ApplicationContext from reme.config import resolve_app_config -from reme.schema import ApplicationConfig +from reme.enumeration import ComponentEnum +from reme.schema import ApplicationConfig, FileNode def test_string_workspace_normalizes_to_default_alias(tmp_path): @@ -202,3 +203,58 @@ def test_default_workspace_matches_directory_creation_source(tmp_path): assert app.config.workspace_dir == app.config.workspaces["main"] assert app.context.workspace.workspace_path == tmp_path / "main" assert (tmp_path / "main" / "daily").is_dir() + + +@pytest.mark.parametrize( + "field_name", + ["metadata_dir", "session_dir", "mem_session_dir", "resource_dir", "daily_dir", "digest_dir"], +) +def test_multi_workspace_rejects_derived_directory_outside_root(tmp_path, field_name): + with pytest.raises(ValidationError, match="workspace"): + ApplicationConfig( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + **{field_name: str(tmp_path / "shared")}, + ) + + +def test_multi_workspace_rejects_derived_directory_escape(tmp_path): + with pytest.raises(ValidationError, match="must stay inside the workspace"): + ApplicationConfig( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + metadata_dir="../shared", + ) + + +def test_single_workspace_keeps_absolute_derived_directory(tmp_path): + external = tmp_path / "external-metadata" + config = ApplicationConfig(workspace_dir=str(tmp_path / "single"), metadata_dir=str(external)) + context = ApplicationContext(**config.model_dump()) + + assert context.workspace.metadata_path == external + + +def _graph_app(tmp_path): + return Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "local"}}}, + ) + + +@pytest.mark.asyncio +async def test_workspace_graph_contents_stay_isolated_across_restart(tmp_path): + app = _graph_app(tmp_path) + async with app: + for alias, path in (("alpha", "daily/alpha-private.md"), ("beta", "daily/beta-note.md")): + graph = app.context.workspaces[alias].get_component(ComponentEnum.FILE_GRAPH, "default") + await graph.upsert_nodes([FileNode(path=path, st_mtime=1.0)]) + + reopened = _graph_app(tmp_path) + async with reopened: + alpha = reopened.context.workspaces["alpha"].get_component(ComponentEnum.FILE_GRAPH, "default") + beta = reopened.context.workspaces["beta"].get_component(ComponentEnum.FILE_GRAPH, "default") + assert [node.path for node in await alpha.get_nodes()] == ["daily/alpha-private.md"] + assert [node.path for node in await beta.get_nodes()] == ["daily/beta-note.md"] From 1b4e5198dc07a8d135909a383dc898d87be6f651 Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 29 Oct 2026 14:10:01 +0800 Subject: [PATCH 11/21] fix(codex): unify workspace thread ownership --- .../agent_wrapper/codex_agent_wrapper.py | 67 ++++++++++++------- reme/schema/workspace.py | 8 +++ tests/unit/test_workspace_phase3.py | 31 ++++++++- 3 files changed, 80 insertions(+), 26 deletions(-) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 078a8500..bcbd41a4 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -394,7 +394,7 @@ class CodexAgentWrapper(BaseAgentWrapper): **common, ) if fork_session or not thread_id: - self._record_thread_ownership(alias, thread.id) + await self._record_thread_ownership(alias, thread.id) tool_contexts[thread.id] = requested_tool_context return thread @@ -426,32 +426,51 @@ class CodexAgentWrapper(BaseAgentWrapper): return None return workspace.mem_session_path / "codex" / "owned_threads.json" - def _owned_thread_ids(self, alias: str) -> set[str]: - """Load the persisted ids into the owning workspace's agent state.""" - state = self._agent_state(alias) - if state is not None and state.owned_thread_ids is not None: - return state.owned_thread_ids - ids: set[str] = set() - path = self._thread_ownership_path(alias) - if path is not None and path.is_file(): - try: - raw = json.loads(path.read_text(encoding="utf-8")) - ids = {str(thread_id) for thread_id in raw.get("thread_ids", [])} - except (OSError, ValueError, AttributeError) as exc: - self.logger.warning(f"Ignoring unreadable Codex thread ownership file {path}: {exc}") - if state is not None: - state.owned_thread_ids = ids - return ids + def _codex_thread_state(self, alias: str): + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return None + return workspaces[self._normalize_thread_alias(alias)].codex_thread_state - def _record_thread_ownership(self, alias: str, thread_id: str) -> None: + async def _owned_thread_ids(self, alias: str) -> set[str]: + """Load ownership through the single authoritative state of each workspace.""" + state = self._codex_thread_state(alias) + if state is None: + return set() + async with state.lock: + if state.owned_thread_ids is not None: + return state.owned_thread_ids + ids: set[str] = set() + path = self._thread_ownership_path(alias) + if path is not None and path.is_file(): + try: + raw = json.loads(path.read_text(encoding="utf-8")) + ids = {str(thread_id) for thread_id in raw.get("thread_ids", [])} + except (OSError, ValueError, AttributeError) as exc: + self.logger.warning(f"Ignoring unreadable Codex thread ownership file {path}: {exc}") + state.owned_thread_ids = ids + return ids + + async def _record_thread_ownership(self, alias: str, thread_id: str) -> None: alias = self._normalize_thread_alias(alias) + state = self._codex_thread_state(alias) path = self._thread_ownership_path(alias) - ids = self._owned_thread_ids(alias) - if path is None or thread_id in ids: + if state is None or path is None: return - ids.add(thread_id) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps({"thread_ids": sorted(ids)}, indent=1), encoding="utf-8") + ids = await self._owned_thread_ids(alias) + async with state.lock: + if thread_id in ids: + return + updated = {*ids, thread_id} + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump({"thread_ids": sorted(updated)}, stream, indent=1) + os.replace(raw_path, path) + finally: + Path(raw_path).unlink(missing_ok=True) + ids.add(thread_id) async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -464,7 +483,7 @@ class CodexAgentWrapper(BaseAgentWrapper): return alias = self._normalize_thread_alias(alias) for owner_alias in workspaces: - if owner_alias != alias and thread_id in self._owned_thread_ids(owner_alias): + if owner_alias != alias and thread_id in await self._owned_thread_ids(owner_alias): raise ValueError( f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " f"it cannot be used from workspace {alias!r}", diff --git a/reme/schema/workspace.py b/reme/schema/workspace.py index 30f3070c..96276cb3 100644 --- a/reme/schema/workspace.py +++ b/reme/schema/workspace.py @@ -14,7 +14,14 @@ class WorkspaceAgentState: """Session execution state owned by one agent wrapper within one workspace.""" thread_tool_contexts: dict[str, str] = field(default_factory=dict) + + +@dataclass +class WorkspaceCodexThreadState: + """Authoritative Codex thread ownership shared by every wrapper in a workspace.""" + owned_thread_ids: set[str] | None = None + lock: asyncio.Lock = field(default_factory=asyncio.Lock) @dataclass(frozen=True) @@ -33,6 +40,7 @@ class WorkspaceContext: tool_contexts: dict[str, Any] = field(default_factory=dict) search_call_budgets: dict[str, Any] = field(default_factory=dict) agent_states: dict[str, WorkspaceAgentState] = field(default_factory=dict) + codex_thread_state: WorkspaceCodexThreadState = field(default_factory=WorkspaceCodexThreadState) locks: dict[str, asyncio.Lock] = field(default_factory=dict) def agent_state(self, wrapper_name: str) -> WorkspaceAgentState: diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index db00f606..938cfc27 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -162,8 +162,8 @@ async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): assert set(alpha_state.thread_tool_contexts) == {"thread-1"} assert set(beta_state.thread_tool_contexts) == {"thread-2"} assert not wrapper._thread_tool_contexts - assert alpha_state.owned_thread_ids == {"thread-1"} - assert beta_state.owned_thread_ids == {"thread-2"} + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"thread-1"} + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"thread-2"} @pytest.mark.asyncio @@ -203,6 +203,33 @@ async def test_codex_thread_ownership_persists_across_restart(tmp_path): await restarted._open_thread(codex, {"session_id": "thread-elsewhere", "workspace": "beta"}) +@pytest.mark.asyncio +async def test_codex_thread_ownership_is_shared_and_atomic_across_wrappers(tmp_path): + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + wrapper_a = CodexAgentWrapper(name="codex-a", app_context=context) + wrapper_b = CodexAgentWrapper(name="codex-b", app_context=context) + + assert await wrapper_b._owned_thread_ids("alpha") == set() + await wrapper_a._record_thread_ownership("alpha", "alpha-thread") + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await wrapper_b._check_thread_ownership("alpha-thread", "beta") + + await wrapper_b._record_thread_ownership("alpha", "second-alpha-thread") + ownership = tmp_path / "alpha" / "mem_session" / "codex" / "owned_threads.json" + assert set(json.loads(ownership.read_text(encoding="utf-8"))["thread_ids"]) == { + "alpha-thread", + "second-alpha-thread", + } + assert not list(ownership.parent.glob("*.tmp")) + + restarted_context = ApplicationContext( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + restarted = CodexAgentWrapper(name="codex-restarted", app_context=restarted_context) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await restarted._check_thread_ownership("alpha-thread", "beta") + + # ----- Background/cron binding ----------------------------------------------- From 9954d63bc5ebb7341f22463e48c67862c4fd24ae Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 29 Oct 2026 14:15:31 +0800 Subject: [PATCH 12/21] fix(codex): preserve workspace and tool failures --- docs/en/configuration.md | 7 ++- docs/zh/configuration.md | 7 ++- .../agent_wrapper/codex_agent_wrapper.py | 7 ++- .../agent_wrapper/codex_mcp_server.py | 5 +- tests/unit/test_workspace_phase3.py | 57 +++++++++++++++++-- 5 files changed, 70 insertions(+), 13 deletions(-) diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 5bb68d2b..5b65bdfd 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -101,9 +101,10 @@ startup in multi-workspace configurations. The binding propagates through nested job calls and agent tool calls (`job_tools`) as a server-owned constraint that models cannot override. For Codex with multiple workspaces, `job_tools` requires the resident HTTP service -with MCP enabled, and the selected jobs must be exposed by that service. Its STDIO bridge forwards to this service -without creating another Application. The multi-workspace one-shot CLI combination is not supported; `mcp_config` -retains its legacy meaning only for single-workspace Codex job tools. +with MCP enabled and `tool_error_on_failure: true`; the selected jobs must be exposed by that service. Its STDIO +bridge forwards to this service without creating another Application and preserves failed Jobs as tool errors. +The multi-workspace one-shot CLI combination is not supported; `mcp_config` retains its legacy meaning only for +single-workspace Codex job tools. Background and cron jobs bind through their job configuration, so extra workspaces are watched by declaring explicit per-alias variants: diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 019fbd72..3351af3e 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -102,9 +102,10 @@ Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 在多 workspace 配置下会在启动时被拒绝。 绑定会经嵌套 Job 调用和 agent 工具调用(`job_tools`)传播,作为模型无法覆盖的服务端约束。 -多 workspace 下的 Codex `job_tools` 需要启用 MCP 的常驻 HTTP 服务,且所选 Job 必须由该服务暴露; -其 STDIO bridge 转发到该服务,不会创建另一个 Application。多 workspace 与 one-shot CLI 的组合暂不支持; -`mcp_config` 仅在单 workspace 的 Codex job tools 中保留原有语义。 +多 workspace 下的 Codex `job_tools` 需要启用 MCP 且配置 `tool_error_on_failure: true` 的常驻 HTTP 服务, +且所选 Job 必须由该服务暴露;其 STDIO bridge 转发到该服务,不会创建另一个 Application,并会将失败的 +Job 保留为 tool error。多 workspace 与 one-shot CLI 的组合暂不支持;`mcp_config` 仅在单 workspace 的 +Codex job tools 中保留原有语义。 Background 和 cron Job 通过 Job 配置绑定别名,额外 workspace 的 watch 通过显式配置按别名的变体实现: diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index bcbd41a4..90b691dd 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -257,6 +257,8 @@ class CodexAgentWrapper(BaseAgentWrapper): @staticmethod def _validate_resident_jobs(service, jobs) -> None: + if not service.tool_error_on_failure: + raise ValueError("Multi-workspace Codex job_tools require resident MCP tool_error_on_failure=true") unavailable = [ job.name for job in jobs @@ -293,11 +295,14 @@ class CodexAgentWrapper(BaseAgentWrapper): "--workspace", str(self.workspace_path), ] + if workspaces: + alias = str(kwargs.get("workspace") or next(iter(workspaces))) + args.extend(["--workspace-alias", alias]) for name in job_names: args.extend(["--job", name]) args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) injected_job_kwargs = dict(kwargs.get("injected_job_kwargs") or {}) - if multi_workspace and (workspace := kwargs.get("workspace")): + if workspaces and (workspace := kwargs.get("workspace")): injected_job_kwargs["workspace"] = workspace if injected_job_kwargs: args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index 1cba8a9b..1ad168ca 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -20,6 +20,7 @@ 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", help="Single-workspace directory for the legacy child service") + parser.add_argument("--workspace-alias", help="Configured alias retained by the legacy child service") parser.add_argument("--endpoint", help="Resident streamable-HTTP MCP endpoint for forwarding") 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") @@ -161,9 +162,11 @@ def main() -> None: return if not args.workspace: raise ValueError("--workspace is required when --endpoint is not supplied") + workspace = str(Path(args.workspace).absolute()) + workspace_dir = {args.workspace_alias: workspace} if args.workspace_alias else workspace config = resolve_app_config( config=args.config, - workspace_dir=str(Path(args.workspace).absolute()), + workspace_dir=workspace_dir, enable_logo=False, log_to_console=False, log_to_file=False, diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index 938cfc27..9e746547 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -331,7 +331,7 @@ async def test_codex_bridge_uses_resident_mcp_endpoint(tmp_path): from reme.components.service import HttpService context = _context(tmp_path) - context.service = HttpService(app_context=context) + context.service = HttpService(app_context=context, tool_error_on_failure=True) context.jobs = {"search_job": _job(context, "search_job")} wrapper = CodexAgentWrapper(app_context=context) @@ -452,13 +452,21 @@ async def test_codex_thread_cwd_follows_bound_workspace(tmp_path): # ----- Codex bridge: real child-process execution ---------------------------- -def _codex_test_app(tmp_path): +def _codex_test_app(tmp_path, workspace_dir=None): + if workspace_dir is None: + workspace_dir = {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")} return Application( - workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + workspace_dir=workspace_dir, enable_logo=False, log_to_console=False, log_to_file=False, - service={"backend": "http", "host": "127.0.0.1", "port": 0, "web_enabled": False}, + service={ + "backend": "http", + "host": "127.0.0.1", + "port": 0, + "web_enabled": False, + "tool_error_on_failure": True, + }, components={ "tokenizer": {"default": {"backend": "regex"}}, "file_store": {"default": {"backend": "local", "embedding_store": "", "tag_index": ""}}, @@ -551,6 +559,13 @@ async def test_codex_bridge_reads_resident_parent_state(tmp_path): client.call_tool("write_note", {"path": "daily/bridge.md", "content": "via-bridge-beta"}), timeout=30, ) + failure = await client.call_tool( + "write_note", + {"path": "../escape.md", "content": "must-not-write"}, + raise_on_error=False, + ) + assert failure.is_error + assert not (tmp_path / "escape.md").exists() assert "parent-resident-beta" in str(read.content[0].text) assert "via-bridge-beta" in (tmp_path / "beta" / "daily" / "bridge.md").read_text(encoding="utf-8") assert not (tmp_path / "alpha" / "daily" / "bridge.md").exists() @@ -573,9 +588,38 @@ async def test_codex_bridge_reads_resident_parent_state(tmp_path): await wrapper.close() +@pytest.mark.asyncio +@pytest.mark.parametrize("named_alias", [False, True]) +async def test_codex_legacy_bridge_preserves_single_workspace_binding(tmp_path, named_alias): + from fastmcp import Client + from fastmcp.client import StdioTransport + + alias = "alice" if named_alias else "default" + root = tmp_path / alias + workspace_dir = {alias: str(root)} if named_alias else str(root) + app = _codex_test_app(tmp_path, workspace_dir=workspace_dir) + wrapper = CodexAgentWrapper(app_context=app.context) + server_config = wrapper._mcp_server_config({"job_tools": ["write_note"], "workspace": alias}) + transport = StdioTransport( + command=server_config["command"], + args=server_config["args"], + cwd=server_config["cwd"], + keep_alive=False, + env=dict(os.environ), + ) + + async with Client(transport, timeout=None, init_timeout=30) as client: + await asyncio.wait_for( + client.call_tool("write_note", {"path": "daily/legacy.md", "content": alias}), + timeout=30, + ) + assert alias in (root / "daily" / "legacy.md").read_text(encoding="utf-8") + await wrapper.close() + + @pytest.mark.asyncio async def test_codex_multi_workspace_rejects_unsupported_job_transport(tmp_path): - from reme.components.service import CliService + from reme.components.service import CliService, HttpService context = _context(tmp_path) context.service = CliService(app_context=context) @@ -588,4 +632,7 @@ async def test_codex_multi_workspace_rejects_unsupported_job_transport(tmp_path) wrapper._mcp_server_config( {"job_tools": ["search_job"], "workspace": "beta", "mcp_config": "custom.yaml"}, ) + context.service = HttpService(app_context=context, tool_error_on_failure=False) + with pytest.raises(ValueError, match="tool_error_on_failure=true"): + wrapper._mcp_server_config({"job_tools": ["search_job"], "workspace": "beta"}) await wrapper.close() From abecd22813ed529033ea0dd389b1d15fee0eb61d Mon Sep 17 00:00:00 2001 From: celanwang Date: Fri, 30 Oct 2026 15:57:51 +0800 Subject: [PATCH 13/21] fix(codex): retain snapshot lifecycle compatibility --- reme/components/agent_wrapper/codex_agent_wrapper.py | 12 +++++++++--- tests/unit/test_codex_agent_wrapper.py | 3 ++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 90b691dd..35aabc18 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -178,7 +178,8 @@ class CodexAgentWrapper(BaseAgentWrapper): snapshot_dir.mkdir(parents=True, exist_ok=True) fd, raw_path = tempfile.mkstemp(prefix="config-", suffix=".json", dir=snapshot_dir) try: - os.fchmod(fd, 0o600) + 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: @@ -541,8 +542,13 @@ class CodexAgentWrapper(BaseAgentWrapper): for workspace in (_context_workspaces(self.app_context) or {}).values(): if state := workspace.agent_states.get(self.name): state.thread_tool_contexts.clear() - if codex is not None: - await codex.close() + 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 async def _prune_thread_tool_contexts(self, codex: AsyncCodex, alias: str = "") -> None: """Forget tool constraints only after the server confirms a thread is unloaded.""" diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index c498fc41..dbd6f294 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -7,6 +7,7 @@ from dataclasses import dataclass import json import os from pathlib import Path +import stat from types import SimpleNamespace import pytest @@ -18,7 +19,7 @@ from reme.components.job import BackgroundJob from reme.components.outbound_proxy import FixedHttpOutboundProxy from reme.config import resolve_app_config from reme.enumeration import ChunkEnum, ComponentEnum -from reme.schema import Response +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. From 637909ec7adf04a871b3b1d764a006e1d821a0e7 Mon Sep 17 00:00:00 2001 From: celanwang Date: Fri, 30 Oct 2026 16:04:51 +0800 Subject: [PATCH 14/21] fix(codex): claim ownership for resumed threads --- docs/en/configuration.md | 3 +- docs/zh/configuration.md | 3 +- .../agent_wrapper/codex_agent_wrapper.py | 53 +++++++---- reme/components/application_context.py | 2 + reme/schema/application_config.py | 2 + tests/unit/test_multi_workspace_config.py | 11 +++ tests/unit/test_workspace_phase3.py | 93 ++++++++++++++++++- 7 files changed, 143 insertions(+), 24 deletions(-) diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 5b65bdfd..4be39008 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -90,7 +90,8 @@ workspace_dir: The first entry is the default workspace; every alias gets the same subdirectory layout (`metadata_dir`, `session_dir`, and the others above) materialized under its own root at startup. Aliases must be non-empty, and resolved directories must be distinct and must not contain one another. A single string keeps the existing -behavior and is normalized internally to a `default` alias. +behavior and is normalized internally to a `default` alias. Multi-workspace configurations also require a non-empty +`mem_session_dir` so agent thread ownership remains enforceable after restart. Job calls accept an optional `workspace` field naming the alias to bind, for example `reme read_note path=daily/note.md workspace=scratch`; omitting it binds the default workspace, and an unknown diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 3351af3e..a556cd43 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -93,7 +93,8 @@ workspace_dir: 第一个 entry 是默认 workspace;每个别名在启动时都会在自己的根目录下创建相同的子目录布局 (`metadata_dir`、`session_dir` 等上述目录)。别名必须非空,且解析后的目录必须互不相同、不能互相包含。 -单个字符串保持原有行为,内部归一化为 `default` 别名。 +单个字符串保持原有行为,内部归一化为 `default` 别名。多 workspace 配置还要求 `mem_session_dir` 非空, +以确保重启后仍能强制执行 agent thread 的归属隔离。 Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 `reme read_note path=daily/note.md workspace=scratch`;缺省绑定默认 workspace,未知别名会直接报错。 diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 35aabc18..d90eb73e 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -399,8 +399,7 @@ class CodexAgentWrapper(BaseAgentWrapper): thread_source=thread_source, **common, ) - if fork_session or not thread_id: - await self._record_thread_ownership(alias, thread.id) + await self._record_thread_ownership(alias, thread.id) tool_contexts[thread.id] = requested_tool_context return thread @@ -458,25 +457,40 @@ class CodexAgentWrapper(BaseAgentWrapper): return ids async def _record_thread_ownership(self, alias: str, thread_id: str) -> None: - alias = self._normalize_thread_alias(alias) - state = self._codex_thread_state(alias) - path = self._thread_ownership_path(alias) - if state is None or path is None: + workspaces = _context_workspaces(self.app_context) + if not workspaces: return - ids = await self._owned_thread_ids(alias) - async with state.lock: - if thread_id in ids: + alias = self._normalize_thread_alias(alias) + async with self.app_context.codex_thread_ownership_lock: + for owner_alias in workspaces: + if thread_id not in await self._owned_thread_ids(owner_alias): + continue + if owner_alias == alias: + return + raise ValueError( + f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " + f"it cannot be used from workspace {alias!r}", + ) + + state = self._codex_thread_state(alias) + ids = await self._owned_thread_ids(alias) + path = self._thread_ownership_path(alias) + if path is None: + ids.add(thread_id) return - updated = {*ids, thread_id} - path.parent.mkdir(parents=True, exist_ok=True) - fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) - try: - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump({"thread_ids": sorted(updated)}, stream, indent=1) - os.replace(raw_path, path) - finally: - Path(raw_path).unlink(missing_ok=True) - ids.add(thread_id) + async with state.lock: + if thread_id in ids: + return + updated = {*ids, thread_id} + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump({"thread_ids": sorted(updated)}, stream, indent=1) + os.replace(raw_path, path) + finally: + Path(raw_path).unlink(missing_ok=True) + ids.add(thread_id) async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -642,6 +656,7 @@ class CodexAgentWrapper(BaseAgentWrapper): thread = await codex.thread_resume(session_id) try: await thread.compact() + await self._record_thread_ownership(str(workspace or ""), thread.id) finally: if self._release_threads: await self._release_thread(codex, thread.id) diff --git a/reme/components/application_context.py b/reme/components/application_context.py index e2b617be..a5f8ad1e 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -1,5 +1,6 @@ """Application context: shared state container for components, jobs, and service.""" +import asyncio from collections.abc import Iterator, MutableMapping from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any @@ -66,6 +67,7 @@ class ApplicationContext: self._components = _LegacyComponentView(self) self.jobs: dict[str, "BaseJob"] = {} self.thread_pool: ThreadPoolExecutor | None = None + self.codex_thread_ownership_lock = asyncio.Lock() self.metadata: dict[str, Any] = {} @property diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index 8d7bcbd6..512664ca 100644 --- a/reme/schema/application_config.py +++ b/reme/schema/application_config.py @@ -115,6 +115,8 @@ class ApplicationConfig(BaseModel): """Normalize every workspace_dir form to an ordered mapping; the first entry is the default.""" self.workspaces = dict(self.workspaces) or {DEFAULT_WORKSPACE_ALIAS: self.workspace_dir} if len(self.workspaces) > 1: + if not self.mem_session_dir.strip(): + raise ValueError("mem_session_dir must be non-empty in multi-workspace configurations") for alias, workspace in self.workspaces.items(): root = Path(workspace).resolve(strict=False) for field_name in ( diff --git a/tests/unit/test_multi_workspace_config.py b/tests/unit/test_multi_workspace_config.py index cd61b7fb..7adeabe9 100644 --- a/tests/unit/test_multi_workspace_config.py +++ b/tests/unit/test_multi_workspace_config.py @@ -225,6 +225,17 @@ def test_multi_workspace_rejects_derived_directory_escape(tmp_path): ) +def test_multi_workspace_requires_session_ownership_storage(tmp_path): + with pytest.raises(ValidationError, match="mem_session_dir must be non-empty"): + ApplicationConfig( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + mem_session_dir="", + ) + + config = ApplicationConfig(workspace_dir=str(tmp_path / "single"), mem_session_dir="") + assert config.mem_session_dir == "" + + def test_single_workspace_keeps_absolute_derived_directory(tmp_path): external = tmp_path / "external-metadata" config = ApplicationConfig(workspace_dir=str(tmp_path / "single"), metadata_dir=str(external)) diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index 9e746547..b91a2527 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -230,6 +230,89 @@ async def test_codex_thread_ownership_is_shared_and_atomic_across_wrappers(tmp_p await restarted._check_thread_ownership("alpha-thread", "beta") +@pytest.mark.asyncio +async def test_codex_resume_and_compact_claim_external_threads(tmp_path, monkeypatch): + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + wrapper_a = CodexAgentWrapper(name="codex-a", app_context=context, release_threads=False) + wrapper_b = CodexAgentWrapper(name="codex-b", app_context=context, release_threads=False) + + class FakeCodex: + """Resume an existing external thread.""" + + async def thread_resume(self, thread_id, **_kwargs): + return SimpleNamespace(id=thread_id) + + await wrapper_a._open_thread(FakeCodex(), {"session_id": "external-thread", "workspace": "alpha"}) + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"external-thread"} + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await wrapper_b._open_thread(FakeCodex(), {"session_id": "external-thread", "workspace": "beta"}) + + class CompactThread: + """Minimal native compaction double.""" + + id = "compact-thread" + + async def compact(self): + return None + + compact_codex = SimpleNamespace(thread_resume=lambda _thread_id: None) + + async def resume_compact(_thread_id): + return CompactThread() + + compact_codex.thread_resume = resume_compact + + async def get_compact_codex(): + return compact_codex + + monkeypatch.setattr(wrapper_a, "_get_codex", get_compact_codex) + await wrapper_a.compact_session("compact-thread", workspace="alpha") + assert "compact-thread" in context.workspaces["alpha"].codex_thread_state.owned_thread_ids + + restarted_context = ApplicationContext( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + restarted = CodexAgentWrapper(name="codex-restarted", app_context=restarted_context) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await restarted._open_thread(FakeCodex(), {"session_id": "external-thread", "workspace": "beta"}) + + +@pytest.mark.asyncio +async def test_codex_concurrent_first_resume_has_one_owner(tmp_path): + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + wrapper_a = CodexAgentWrapper(name="codex-a", app_context=context) + wrapper_b = CodexAgentWrapper(name="codex-b", app_context=context) + + class FakeCodex: + """Yield between resume and ownership registration to force interleaving.""" + + async def thread_resume(self, thread_id, **_kwargs): + await asyncio.sleep(0) + return SimpleNamespace(id=thread_id) + + results = await asyncio.gather( + wrapper_a._open_thread(FakeCodex(), {"session_id": "shared-external", "workspace": "alpha"}), + wrapper_b._open_thread(FakeCodex(), {"session_id": "shared-external", "workspace": "beta"}), + return_exceptions=True, + ) + assert len([result for result in results if not isinstance(result, Exception)]) == 1 + assert len([result for result in results if isinstance(result, ValueError)]) == 1 + + owners = [ + alias + for alias, workspace in context.workspaces.items() + if "shared-external" in (workspace.codex_thread_state.owned_thread_ids or set()) + ] + assert len(owners) == 1 + other = "beta" if owners[0] == "alpha" else "alpha" + restarted_context = ApplicationContext( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + ) + restarted = CodexAgentWrapper(name="codex-restarted", app_context=restarted_context) + with pytest.raises(ValueError, match="belongs to workspace"): + await restarted._open_thread(FakeCodex(), {"session_id": "shared-external", "workspace": other}) + + # ----- Background/cron binding ----------------------------------------------- @@ -437,15 +520,19 @@ async def test_codex_thread_cwd_follows_bound_workspace(tmp_path): class FakeCodex: """Capture thread_start parameters without an app-server.""" + next_id = 1 + async def thread_start(self, **kwargs): observed.update(kwargs) - return SimpleNamespace(id="thread-1") + self.next_id += 1 + return SimpleNamespace(id=f"thread-{self.next_id - 1}") - await wrapper._open_thread(FakeCodex(), {"workspace": "beta"}) + codex = FakeCodex() + await wrapper._open_thread(codex, {"workspace": "beta"}) assert observed["cwd"] == str(tmp_path / "beta") observed.clear() - await wrapper._open_thread(FakeCodex(), {}) + await wrapper._open_thread(codex, {}) assert observed["cwd"] != str(tmp_path / "beta") From b5d0bbbc2064f2b6c48da0010c450331d4e28d6d Mon Sep 17 00:00:00 2001 From: celanwang Date: Fri, 30 Oct 2026 16:09:50 +0800 Subject: [PATCH 15/21] fix(agent): resolve default workspace at wrapper entry --- .../agent_wrapper/base_agent_wrapper.py | 2 + tests/unit/test_cc_agent_wrapper.py | 2 +- tests/unit/test_workspace_ownership.py | 12 +++ tests/unit/test_workspace_phase3.py | 75 +++++++++++++++++++ 4 files changed, 90 insertions(+), 1 deletion(-) diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 48ff158f..150f3d53 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -242,6 +242,8 @@ class BaseAgentWrapper(BaseComponent): def _merged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: """Merge component defaults with call-time kwargs; call-time values win.""" merged = {**self.kwargs, **kwargs} + if (workspaces := _context_workspaces(self.app_context)) and not merged.get("workspace"): + merged["workspace"] = next(iter(workspaces)) if "output_schema" in merged: merged["output_schema"] = self._normalize_output_schema(merged["output_schema"]) return merged diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index a0db6ebc..8fdbefbb 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -387,7 +387,7 @@ async def test_agentscope_compact_session_forces_and_persists_compression(tmp_pa await wrapper.compact_session("session-1") - assert observed["build"] == (None, {"resume": "session-1"}) + assert observed["build"] == (None, {"resume": "session-1", "workspace": "default"}) assert observed["config"].trigger_ratio == 1e-9 assert observed["state"] is FakeAgent.state diff --git a/tests/unit/test_workspace_ownership.py b/tests/unit/test_workspace_ownership.py index d388673a..beaa0440 100644 --- a/tests/unit/test_workspace_ownership.py +++ b/tests/unit/test_workspace_ownership.py @@ -9,6 +9,7 @@ import pytest from reme.application import Application from reme.components import ApplicationContext from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.agent_wrapper.cc_agent_wrapper import CcAgentWrapper from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper from reme.components.component_registry import R from reme.components.job import BaseJob @@ -63,6 +64,17 @@ async def test_default_binding_reaches_real_agent_tool_call(tmp_path): await tool._func(workspace="beta") +@pytest.mark.parametrize("wrapper_cls", [AsAgentWrapper, CcAgentWrapper, CodexAgentWrapper]) +def test_agent_entry_resolves_default_workspace_for_job_tools(tmp_path, wrapper_cls): + context = _context(tmp_path) + wrapper = wrapper_cls(app_context=context) + + merged = wrapper._merged_kwargs({}) + assert merged["workspace"] == "alpha" + assert wrapper._resolve_injected_job_kwargs(merged)["workspace"] == "alpha" + assert wrapper._merged_kwargs({"workspace": "beta"})["workspace"] == "beta" + + def test_call_cwd_resolves_relative_against_bound_workspace(tmp_path): context = _context(tmp_path) absolute = tmp_path / "elsewhere" diff --git a/tests/unit/test_workspace_phase3.py b/tests/unit/test_workspace_phase3.py index b91a2527..38ef75cb 100644 --- a/tests/unit/test_workspace_phase3.py +++ b/tests/unit/test_workspace_phase3.py @@ -5,6 +5,7 @@ import asyncio import json import os +from dataclasses import dataclass from types import SimpleNamespace import pytest @@ -675,6 +676,80 @@ async def test_codex_bridge_reads_resident_parent_state(tmp_path): await wrapper.close() +@pytest.mark.asyncio +async def test_codex_reply_defaults_workspace_for_real_job_tools(tmp_path, monkeypatch): + from fastmcp import Client + from fastmcp.client import StdioTransport + + @dataclass + class TurnResult: + """Serializable fake turn result.""" + + final_response: str + usage: object = None + + class FakeThread: + """Execute the configured job tool through the real STDIO bridge.""" + + def __init__(self, thread_id, config): + self.id = thread_id + self.config = config + + async def run(self, _inputs, **_kwargs): + server = next(iter(self.config["mcp_servers"].values())) + transport = StdioTransport( + command=server["command"], + args=server["args"], + cwd=server["cwd"], + keep_alive=False, + env=dict(os.environ), + ) + arguments = {"path": f"daily/reply-{self.id}.md", "content": self.id} + if self.id == "thread-override": + arguments["workspace"] = "beta" + async with Client(transport, timeout=None, init_timeout=30) as client: + await client.call_tool("write_note", arguments) + return TurnResult(final_response="done") + + class FakeCodex: + """Create threads with deterministic IDs.""" + + def __init__(self): + self.next_id = 0 + + async def thread_start(self, **kwargs): + self.next_id += 1 + thread_id = "thread-override" if self.next_id == 3 else f"thread-{self.next_id}" + return FakeThread(thread_id, kwargs["config"]) + + app = _codex_test_app(tmp_path) + app.context.jobs["write_note"].kwargs["workspace"] = "beta" + wrapper = CodexAgentWrapper(app_context=app.context, release_threads=False) + codex = FakeCodex() + + async def get_codex(): + return codex + + monkeypatch.setattr(wrapper, "_get_codex", get_codex) + async with app: + server, task = await _start_resident_http(app) + try: + await wrapper.reply("default", job_tools=["write_note"]) + assert (tmp_path / "alpha" / "daily" / "reply-thread-1.md").is_file() + assert not (tmp_path / "beta" / "daily" / "reply-thread-1.md").exists() + + await wrapper.reply("explicit", job_tools=["write_note"], workspace="alpha") + assert (tmp_path / "alpha" / "daily" / "reply-thread-2.md").is_file() + assert not (tmp_path / "beta" / "daily" / "reply-thread-2.md").exists() + + with pytest.raises(Exception): + await wrapper.reply("override", job_tools=["write_note"]) + assert not (tmp_path / "beta" / "daily" / "reply-thread-override.md").exists() + finally: + await _stop_resident_http(server, task) + await wrapper.close() + + @pytest.mark.asyncio @pytest.mark.parametrize("named_alias", [False, True]) async def test_codex_legacy_bridge_preserves_single_workspace_binding(tmp_path, named_alias): From ffda20c45148f27a6cae7744c82d864762978743 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 2 Nov 2026 12:00:06 +0000 Subject: [PATCH 16/21] fix(codex): claim compaction ownership before SDK side effects Validate the workspace alias and settle the cross-workspace claim before thread resume or compact, so a losing request cannot compact. Drop a claim made for this call when the SDK rejects the operation, and leave an existing claim in place. Co-authored-by: celanwang --- .../agent_wrapper/codex_agent_wrapper.py | 100 ++++++++--- tests/unit/test_codex_compaction_ownership.py | 160 ++++++++++++++++++ 2 files changed, 240 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_codex_compaction_ownership.py diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index d90eb73e..7fd43ee4 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -456,17 +456,46 @@ class CodexAgentWrapper(BaseAgentWrapper): state.owned_thread_ids = ids return ids - async def _record_thread_ownership(self, alias: str, thread_id: str) -> None: + def _require_thread_alias(self, alias: str) -> str: + """Resolve an omitted alias and reject names that are not configured.""" workspaces = _context_workspaces(self.app_context) if not workspaces: - return - alias = self._normalize_thread_alias(alias) + return alias + resolved = self._normalize_thread_alias(alias) + if resolved not in workspaces: + configured = sorted(workspaces) + raise ValueError(f"Unknown workspace {resolved!r}; configured: {configured}") + return resolved + + @staticmethod + def _write_owned_thread_ids(path: Path, thread_ids: set[str]) -> None: + """Atomically replace one workspace's ownership file.""" + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump({"thread_ids": sorted(thread_ids)}, stream, indent=1) + os.replace(raw_path, path) + finally: + Path(raw_path).unlink(missing_ok=True) + + async def _record_thread_ownership(self, alias: str, thread_id: str) -> bool: + """Claim ``thread_id`` for ``alias``. + + Returns True when this call newly records the thread. An unknown alias, + or a thread already owned by another workspace, raises ValueError before + the claim is stored. + """ + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return False + alias = self._require_thread_alias(alias) async with self.app_context.codex_thread_ownership_lock: for owner_alias in workspaces: if thread_id not in await self._owned_thread_ids(owner_alias): continue if owner_alias == alias: - return + return False raise ValueError( f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " f"it cannot be used from workspace {alias!r}", @@ -476,21 +505,40 @@ class CodexAgentWrapper(BaseAgentWrapper): ids = await self._owned_thread_ids(alias) path = self._thread_ownership_path(alias) if path is None: + if thread_id in ids: + return False ids.add(thread_id) - return + return True async with state.lock: if thread_id in ids: - return + return False updated = {*ids, thread_id} - path.parent.mkdir(parents=True, exist_ok=True) - fd, raw_path = tempfile.mkstemp(prefix=".owned_threads-", suffix=".tmp", dir=path.parent) - try: - with os.fdopen(fd, "w", encoding="utf-8") as stream: - json.dump({"thread_ids": sorted(updated)}, stream, indent=1) - os.replace(raw_path, path) - finally: - Path(raw_path).unlink(missing_ok=True) + self._write_owned_thread_ids(path, updated) ids.add(thread_id) + return True + + async def _forget_thread_ownership(self, alias: str, thread_id: str) -> None: + """Drop a claim reserved for an SDK operation that did not succeed.""" + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return + alias = self._normalize_thread_alias(alias) + async with self.app_context.codex_thread_ownership_lock: + ids = await self._owned_thread_ids(alias) + if thread_id not in ids: + return + path = self._thread_ownership_path(alias) + if path is None: + ids.discard(thread_id) + return + state = self._codex_thread_state(alias) + async with state.lock: + if thread_id not in ids: + return + updated = set(ids) + updated.discard(thread_id) + self._write_owned_thread_ids(path, updated) + ids.discard(thread_id) async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -648,17 +696,29 @@ class CodexAgentWrapper(BaseAgentWrapper): return response async def compact_session(self, session_id: str, workspace: str | None = None) -> None: - """Start native compaction of a Codex thread owned by the bound workspace.""" + """Compact a Codex thread after its workspace wins ownership. + + The alias is validated and the cross-workspace claim is settled before + any SDK call, so a losing request cannot compact. A claim made for this + call is removed when the SDK rejects the operation; an existing claim + is left in place. Thread release still runs once a thread was resumed. + """ await self.start() + alias = str(workspace or "") async with self._turn_lock: - await self._check_thread_ownership(session_id, str(workspace or "")) - codex = await self._get_codex() - thread = await codex.thread_resume(session_id) + newly_claimed = await self._record_thread_ownership(alias, session_id) + codex = None + thread = None try: + codex = await self._get_codex() + thread = await codex.thread_resume(session_id) await thread.compact() - await self._record_thread_ownership(str(workspace or ""), thread.id) + except BaseException: + if newly_claimed: + await self._forget_thread_ownership(alias, session_id) + raise finally: - if self._release_threads: + if thread is not None and codex is not None and self._release_threads: await self._release_thread(codex, thread.id) @classmethod diff --git a/tests/unit/test_codex_compaction_ownership.py b/tests/unit/test_codex_compaction_ownership.py new file mode 100644 index 00000000..47562045 --- /dev/null +++ b/tests/unit/test_codex_compaction_ownership.py @@ -0,0 +1,160 @@ +"""Compaction claims a Codex thread before any SDK side effect.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from reme.components import ApplicationContext +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper + + +def _owned_on_disk(tmp_path, alias) -> set[str]: + path = tmp_path / alias / "mem_session" / "codex" / "owned_threads.json" + if not path.is_file(): + return set() + return set(json.loads(path.read_text(encoding="utf-8"))["thread_ids"]) + + +@pytest.mark.asyncio +async def test_compact_arbitrates_before_sdk_side_effects(tmp_path, monkeypatch): + """Compaction claims ownership before resume/compact, including under concurrency.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + alpha = CodexAgentWrapper(name="codex-alpha", app_context=context, release_threads=False) + beta = CodexAgentWrapper(name="codex-beta", app_context=context, release_threads=False) + sdk_ops: list[str] = [] + entered_sdk = asyncio.Event() + release_sdk = asyncio.Event() + + class PausingThread: + """Resume pauses inside the SDK so a competing resume can try to claim.""" + + id = "shared-thread" + + async def compact(self): + sdk_ops.append("compact") + + async def alpha_resume(_thread_id): + sdk_ops.append("alpha-resume") + entered_sdk.set() + await release_sdk.wait() + return PausingThread() + + async def beta_resume(thread_id, **_kwargs): + sdk_ops.append("beta-resume") + return SimpleNamespace(id=thread_id) + + async def alpha_codex(): + return SimpleNamespace(thread_resume=alpha_resume) + + monkeypatch.setattr(alpha, "_get_codex", alpha_codex) + compact_task = asyncio.create_task(alpha.compact_session("shared-thread", workspace="alpha")) + await asyncio.wait_for(entered_sdk.wait(), timeout=2) + + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta._open_thread( + SimpleNamespace(thread_resume=beta_resume), + {"session_id": "shared-thread", "workspace": "beta"}, + ) + + release_sdk.set() + await asyncio.wait_for(compact_task, timeout=2) + assert sdk_ops == ["alpha-resume", "compact"] + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"shared-thread"} + assert "shared-thread" not in (context.workspaces["beta"].codex_thread_state.owned_thread_ids or set()) + assert _owned_on_disk(tmp_path, "alpha") == {"shared-thread"} + assert "shared-thread" not in _owned_on_disk(tmp_path, "beta") + + # An alias that was never configured fails before either SDK method. + unknown_ops: list[str] = [] + + async def unknown_resume(_thread_id): + unknown_ops.append("resume") + return PausingThread() + + async def unknown_codex(): + return SimpleNamespace(thread_resume=unknown_resume) + + monkeypatch.setattr(alpha, "_get_codex", unknown_codex) + with pytest.raises(ValueError, match="Unknown workspace 'typo'"): + await alpha.compact_session("unrecorded-thread", workspace="typo") + assert not unknown_ops + assert "unrecorded-thread" not in (context.workspaces["alpha"].codex_thread_state.owned_thread_ids or set()) + + # Once beta owns a thread, alpha's compaction does not reach the SDK. + await beta._open_thread( + SimpleNamespace(thread_resume=beta_resume), + {"session_id": "beta-thread", "workspace": "beta"}, + ) + loser_ops: list[str] = [] + + async def loser_resume(_thread_id): + loser_ops.append("resume") + return PausingThread() + + async def loser_codex(): + return SimpleNamespace(thread_resume=loser_resume) + + monkeypatch.setattr(alpha, "_get_codex", loser_codex) + with pytest.raises(ValueError, match="belongs to workspace 'beta'"): + await alpha.compact_session("beta-thread", workspace="alpha") + assert not loser_ops + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"beta-thread"} + + +@pytest.mark.asyncio +async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): + """A failed resume or compact releases a claim this call just made.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + alpha = CodexAgentWrapper(name="codex-alpha", app_context=context) + beta = CodexAgentWrapper(name="codex-beta", app_context=context, release_threads=False) + released: list[str] = [] + + class RecordingClient: + """Capture thread ids passed to the SDK unsubscribe request.""" + + async def request(self, method, params, *, response_model): + released.append(params["threadId"]) + assert method == "thread/unsubscribe" + assert response_model is not None + return SimpleNamespace() + + class RejectingThread: + """Resume succeeds and compaction is rejected by the SDK.""" + + def __init__(self, thread_id): + self.id = thread_id + + async def compact(self): + raise RuntimeError("sdk rejected") + + async def resume(thread_id): + if thread_id == "resume-fails": + raise RuntimeError("resume rejected") + return RejectingThread(thread_id) + + async def get_codex(): + return SimpleNamespace(_client=RecordingClient(), thread_resume=resume) + + monkeypatch.setattr(alpha, "_get_codex", get_codex) + + with pytest.raises(RuntimeError, match="resume rejected"): + await alpha.compact_session("resume-fails", workspace="alpha") + with pytest.raises(RuntimeError, match="sdk rejected"): + await alpha.compact_session("new-thread", workspace="alpha") + await alpha._record_thread_ownership("alpha", "kept-thread") + with pytest.raises(RuntimeError, match="sdk rejected"): + await alpha.compact_session("kept-thread", workspace="alpha") + + owned = context.workspaces["alpha"].codex_thread_state.owned_thread_ids or set() + assert "resume-fails" not in owned + assert "new-thread" not in owned + assert owned == {"kept-thread"} + assert _owned_on_disk(tmp_path, "alpha") == {"kept-thread"} + assert released == ["new-thread", "kept-thread"] + + await beta._record_thread_ownership("beta", "new-thread") + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"new-thread"} From f2c5b16ae366938e7e60e8d6b7af1789310d653b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 2 Nov 2026 12:00:06 +0000 Subject: [PATCH 17/21] fix(agent): keep project-root cwd when the default workspace is filled in Job tools still receive the default alias. That fill-in is marked so cwd resolution can ignore it and keep the single-workspace project root at the SDK boundary. Co-authored-by: celanwang --- .../agent_wrapper/base_agent_wrapper.py | 23 ++++++-- tests/unit/test_cc_agent_wrapper.py | 6 +- tests/unit/test_codex_agent_wrapper.py | 58 +++++++++++++++++++ tests/unit/test_workspace_ownership.py | 2 + 4 files changed, 82 insertions(+), 7 deletions(-) diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index 150f3d53..4d110099 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -70,17 +70,20 @@ class BaseAgentWrapper(BaseComponent): return workspaces[alias].workspace_path def _call_cwd(self, kwargs: dict[str, Any]) -> Path: - """Working directory for this call: explicit cwd config wins, then the bound workspace. + """Working directory for this call. - A relative configured cwd resolves against the bound workspace; - an absolute path keeps its explicit meaning. + An explicit cwd wins. A relative one resolves against the bound + workspace; an absolute path keeps its explicit meaning. Otherwise an + explicitly selected workspace selects that workspace root. The default + alias injected for job tools is not an explicit selection, so a + single-workspace call keeps the project root. """ if self._cwd: cwd = Path(self._cwd) if cwd.is_absolute(): return cwd return self._call_workspace_path(kwargs) / cwd - if kwargs.get("workspace"): + if kwargs.get("workspace") and kwargs.get("_workspace_explicit", True): return self._call_workspace_path(kwargs) return self.cwd @@ -240,10 +243,18 @@ class BaseAgentWrapper(BaseComponent): return resolved def _merged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: - """Merge component defaults with call-time kwargs; call-time values win.""" + """Merge component defaults with call-time kwargs; call-time values win. + + An omitted workspace becomes the default alias so job tools keep an + explicit binding. ``_workspace_explicit`` records whether the caller or + the component config selected that alias; cwd resolution uses it to + keep the historical project root when the alias was only filled in. + """ merged = {**self.kwargs, **kwargs} - if (workspaces := _context_workspaces(self.app_context)) and not merged.get("workspace"): + explicit_workspace = bool(merged.get("workspace")) + if (workspaces := _context_workspaces(self.app_context)) and not explicit_workspace: merged["workspace"] = next(iter(workspaces)) + merged["_workspace_explicit"] = explicit_workspace if "output_schema" in merged: merged["output_schema"] = self._normalize_output_schema(merged["output_schema"]) return merged diff --git a/tests/unit/test_cc_agent_wrapper.py b/tests/unit/test_cc_agent_wrapper.py index 8fdbefbb..a2b28ae1 100644 --- a/tests/unit/test_cc_agent_wrapper.py +++ b/tests/unit/test_cc_agent_wrapper.py @@ -387,7 +387,11 @@ async def test_agentscope_compact_session_forces_and_persists_compression(tmp_pa await wrapper.compact_session("session-1") - assert observed["build"] == (None, {"resume": "session-1", "workspace": "default"}) + build_inputs, build_kwargs = observed["build"] + assert build_inputs is None + # Default-alias fill-in is marked so cwd resolution can ignore it. + assert build_kwargs.pop("_workspace_explicit") is False + assert build_kwargs == {"resume": "session-1", "workspace": "default"} assert observed["config"].trigger_ratio == 1e-9 assert observed["state"] is FakeAgent.state diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index dbd6f294..4874b4e5 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -14,6 +14,7 @@ import pytest from openai_codex.generated.v2_all import TokenUsageBreakdown from pydantic import BaseModel +from reme.components import ApplicationContext from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper from reme.components.job import BackgroundJob from reme.components.outbound_proxy import FixedHttpOutboundProxy @@ -721,6 +722,63 @@ async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, mo assert close_count == 1 +@pytest.mark.asyncio +async def test_reply_keeps_project_root_cwd_without_explicit_workspace(tmp_path, monkeypatch): + """A single-workspace reply delivers the project root to the SDK. + + The default alias is still injected for job tools. Naming that alias, or + choosing another workspace, selects that workspace directory instead. + """ + project = tmp_path / "project" + workspace = project / ".reme" + context = ApplicationContext(workspace_dir=str(workspace)) + wrapper = CodexAgentWrapper( + app_context=context, + project_path="..", + auth_mode="oauth", + release_threads=False, + ) + observed = {"thread": [], "turn": []} + + class FakeThread: + id = "thread-1" + + async def run(self, _inputs, **kwargs): + observed["turn"].append(kwargs["cwd"]) + return _TurnResult(final_response="ok") + + 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): + observed["thread"].append(kwargs["cwd"]) + return FakeThread() + + monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.AsyncCodex", FakeCodex) + + merged = wrapper._merged_kwargs({}) # pylint: disable=protected-access + assert merged["workspace"] == "default" + assert wrapper._resolve_injected_job_kwargs(merged)["workspace"] == "default" # pylint: disable=protected-access + + await wrapper.reply("hello") + await wrapper.reply("again", workspace="default") + await wrapper.close() + + project_cwd = str(wrapper.project_path) + workspace_cwd = str(context.workspaces["default"].workspace_path) + assert wrapper.cwd == wrapper.project_path + assert project_cwd != workspace_cwd + assert observed["thread"] == [project_cwd, workspace_cwd] + assert observed["turn"] == [project_cwd, workspace_cwd] + + @pytest.mark.asyncio @pytest.mark.parametrize("schema", [{}, _StructuredModel]) async def test_reply_stream_rejects_output_schema(tmp_path, schema): diff --git a/tests/unit/test_workspace_ownership.py b/tests/unit/test_workspace_ownership.py index beaa0440..2f7ce6c6 100644 --- a/tests/unit/test_workspace_ownership.py +++ b/tests/unit/test_workspace_ownership.py @@ -89,6 +89,8 @@ def test_call_cwd_resolves_relative_against_bound_workspace(tmp_path): default_wrapper = CodexAgentWrapper(app_context=context) assert default_wrapper._call_cwd({"workspace": "beta"}) == tmp_path / "beta" assert default_wrapper._call_cwd({}) != tmp_path / "beta" # project path, not a workspace + # Filling in the default alias for job tools does not switch a relative cwd. + assert wrapper._call_cwd(wrapper._merged_kwargs({})) == tmp_path / "alpha" / "relative" / "sub" def test_dream_integration_lock_is_workspace_scoped(tmp_path): From 5b14319c4ed947aa1daa015f488fd3223a991517 Mon Sep 17 00:00:00 2001 From: celanwang Date: Wed, 4 Nov 2026 21:23:19 +0800 Subject: [PATCH 18/21] fix(codex): preserve ownership adopted by concurrent operations --- .../agent_wrapper/codex_agent_wrapper.py | 69 ++++++----- reme/schema/workspace.py | 2 + tests/unit/test_codex_compaction_ownership.py | 115 ++++++++++++++++++ 3 files changed, 157 insertions(+), 29 deletions(-) diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index 7fd43ee4..f8c95cd7 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -389,7 +389,14 @@ class CodexAgentWrapper(BaseAgentWrapper): **common, ) elif thread_id: - thread = await codex.thread_resume(thread_id, personality=personality, **common) + reserved = await self._record_thread_ownership(alias, thread_id, provisional=True) + try: + thread = await codex.thread_resume(thread_id, personality=personality, **common) + await self._record_thread_ownership(alias, thread.id) + except BaseException: + if reserved: + await self._forget_thread_ownership(alias, thread_id) + raise else: thread = await codex.thread_start( ephemeral=kwargs.get("ephemeral"), @@ -479,12 +486,12 @@ class CodexAgentWrapper(BaseAgentWrapper): finally: Path(raw_path).unlink(missing_ok=True) - async def _record_thread_ownership(self, alias: str, thread_id: str) -> bool: - """Claim ``thread_id`` for ``alias``. + async def _record_thread_ownership(self, alias: str, thread_id: str, *, provisional: bool = False) -> bool: + """Reserve or confirm ownership under the application-wide lock. - Returns True when this call newly records the thread. An unknown alias, - or a thread already owned by another workspace, raises ValueError before - the claim is stored. + Provisional callers receive True while their reservation needs rollback + on failure. Successful SDK operations confirm the claim for all wrappers. + Persisted claims loaded after restart are always treated as confirmed. """ workspaces = _context_workspaces(self.app_context) if not workspaces: @@ -495,6 +502,12 @@ class CodexAgentWrapper(BaseAgentWrapper): if thread_id not in await self._owned_thread_ids(owner_alias): continue if owner_alias == alias: + claims = self._codex_thread_state(alias).provisional_claims + if not provisional: + claims.pop(thread_id, None) + elif thread_id in claims: + claims[thread_id] += 1 + return True return False raise ValueError( f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " @@ -504,41 +517,38 @@ class CodexAgentWrapper(BaseAgentWrapper): state = self._codex_thread_state(alias) ids = await self._owned_thread_ids(alias) path = self._thread_ownership_path(alias) - if path is None: - if thread_id in ids: - return False - ids.add(thread_id) - return True async with state.lock: - if thread_id in ids: - return False - updated = {*ids, thread_id} - self._write_owned_thread_ids(path, updated) + if path is not None: + self._write_owned_thread_ids(path, {*ids, thread_id}) ids.add(thread_id) + if provisional: + state.provisional_claims[thread_id] = 1 return True async def _forget_thread_ownership(self, alias: str, thread_id: str) -> None: - """Drop a claim reserved for an SDK operation that did not succeed.""" + """Release a failed reservation without revoking confirmed or active ownership.""" workspaces = _context_workspaces(self.app_context) if not workspaces: return alias = self._normalize_thread_alias(alias) async with self.app_context.codex_thread_ownership_lock: + state = self._codex_thread_state(alias) + count = state.provisional_claims.get(thread_id, 0) + if not count: + return + if count > 1: + state.provisional_claims[thread_id] -= 1 + return ids = await self._owned_thread_ids(alias) if thread_id not in ids: return path = self._thread_ownership_path(alias) - if path is None: - ids.discard(thread_id) - return - state = self._codex_thread_state(alias) async with state.lock: - if thread_id not in ids: - return - updated = set(ids) - updated.discard(thread_id) - self._write_owned_thread_ids(path, updated) + updated = ids - {thread_id} + if path is not None: + self._write_owned_thread_ids(path, updated) ids.discard(thread_id) + state.provisional_claims.pop(thread_id, None) async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -700,21 +710,22 @@ class CodexAgentWrapper(BaseAgentWrapper): The alias is validated and the cross-workspace claim is settled before any SDK call, so a losing request cannot compact. A claim made for this - call is removed when the SDK rejects the operation; an existing claim - is left in place. Thread release still runs once a thread was resumed. + call is removed on failure only if no other operation still uses or has + confirmed it. Thread release still runs once a thread was resumed. """ await self.start() alias = str(workspace or "") async with self._turn_lock: - newly_claimed = await self._record_thread_ownership(alias, session_id) + reserved = await self._record_thread_ownership(alias, session_id, provisional=True) codex = None thread = None try: codex = await self._get_codex() thread = await codex.thread_resume(session_id) await thread.compact() + await self._record_thread_ownership(alias, session_id) except BaseException: - if newly_claimed: + if reserved: await self._forget_thread_ownership(alias, session_id) raise finally: diff --git a/reme/schema/workspace.py b/reme/schema/workspace.py index 96276cb3..3394b3b4 100644 --- a/reme/schema/workspace.py +++ b/reme/schema/workspace.py @@ -21,6 +21,8 @@ class WorkspaceCodexThreadState: """Authoritative Codex thread ownership shared by every wrapper in a workspace.""" owned_thread_ids: set[str] | None = None + # Active operations on unconfirmed claims; protected by the application ownership lock. + provisional_claims: dict[str, int] = field(default_factory=dict) lock: asyncio.Lock = field(default_factory=asyncio.Lock) diff --git a/tests/unit/test_codex_compaction_ownership.py b/tests/unit/test_codex_compaction_ownership.py index 47562045..9aff5297 100644 --- a/tests/unit/test_codex_compaction_ownership.py +++ b/tests/unit/test_codex_compaction_ownership.py @@ -158,3 +158,118 @@ async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): await beta._record_thread_ownership("beta", "new-thread") assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {"new-thread"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.parametrize("worker_operation", ["reply", "compact"]) +async def test_failed_compactor_keeps_ownership_confirmed_by_another_wrapper( + tmp_path, + monkeypatch, + cancel, + worker_operation, +): + """Failure/cancellation cannot revoke another wrapper's successful use.""" + config = {"workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}} + context = ApplicationContext(**config) + compactor = CodexAgentWrapper(name="compactor", app_context=context, release_threads=False) + worker = CodexAgentWrapper(name="worker", app_context=context, release_threads=False) + entered = asyncio.Event() + finish = asyncio.Event() + + async def failing_client(): + entered.set() + await finish.wait() + raise RuntimeError("client initialization failed") + + async def run(_inputs, **_kwargs): + return SimpleNamespace(final_response="ok") + + async def compact(): + return None + + async def resume(thread_id, **_kwargs): + return SimpleNamespace(id=thread_id, run=run, compact=compact) + + async def working_client(): + return SimpleNamespace(thread_resume=resume) + + monkeypatch.setattr(compactor, "_get_codex", failing_client) + monkeypatch.setattr(worker, "_get_codex", working_client) + monkeypatch.setattr(worker, "_serialize", lambda _result: {}) + task = asyncio.create_task(compactor.compact_session("T", workspace="alpha")) + try: + await asyncio.wait_for(entered.wait(), timeout=2) + if worker_operation == "reply": + result = await worker.reply("hello", resume="T", workspace="alpha") + assert result["result"] == "ok" + else: + await worker.compact_session("T", workspace="alpha") + if cancel: + task.cancel() + else: + finish.set() + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + await asyncio.wait_for(task, timeout=2) + + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"T"} + assert not context.workspaces["alpha"].codex_thread_state.provisional_claims + assert _owned_on_disk(tmp_path, "alpha") == {"T"} + for checked_context in (context, ApplicationContext(**config)): + beta = CodexAgentWrapper(name="beta", app_context=checked_context, release_threads=False) + monkeypatch.setattr(beta, "_get_codex", working_client) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta.reply("hello", resume="T", workspace="beta") + assert not _owned_on_disk(tmp_path, "beta") + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await compactor.close() + await worker.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_failed_compactor_retains_another_pending_reservation(tmp_path, monkeypatch, cancel): + """The last failed reservation cleans up; earlier failures preserve active work.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + wrappers = [CodexAgentWrapper(name=name, app_context=context) for name in ("first", "second")] + entered = [asyncio.Event(), asyncio.Event()] + finish = [asyncio.Event(), asyncio.Event()] + + def client(index): + async def get_client(): + entered[index].set() + await finish[index].wait() + raise RuntimeError("failed") + + return get_client + + tasks = [] + try: + for index, wrapper in enumerate(wrappers): + monkeypatch.setattr(wrapper, "_get_codex", client(index)) + tasks.append(asyncio.create_task(wrapper.compact_session("T", workspace="alpha"))) + await asyncio.wait_for(entered[index].wait(), timeout=2) + if cancel: + tasks[0].cancel() + else: + finish[0].set() + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + await asyncio.wait_for(tasks[0], timeout=2) + assert _owned_on_disk(tmp_path, "alpha") == {"T"} + beta = CodexAgentWrapper(app_context=context) + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta.compact_session("T", workspace="beta") + finish[1].set() + with pytest.raises(RuntimeError, match="failed"): + await asyncio.wait_for(tasks[1], timeout=2) + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == set() + assert not context.workspaces["alpha"].codex_thread_state.provisional_claims + assert not _owned_on_disk(tmp_path, "alpha") + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + for wrapper in wrappers: + await wrapper.close() From 46faffcb9fec23a7d9329e5b8d8725647e8bbbc3 Mon Sep 17 00:00:00 2001 From: celanwang Date: Wed, 4 Nov 2026 21:23:20 +0800 Subject: [PATCH 19/21] fix(status): attribute component memory to the selected workspace --- docs/en/operations.md | 2 +- docs/zh/operations.md | 2 +- reme/steps/common/status.py | 16 +++++++++---- tests/unit/test_status_step.py | 44 ++++++++++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 6 deletions(-) diff --git a/docs/en/operations.md b/docs/en/operations.md index f5e62823..a1f0c6b2 100644 --- a/docs/en/operations.md +++ b/docs/en/operations.md @@ -23,7 +23,7 @@ reme app_config - `find_reme` confirms the actual host, port, and PID; - `version` verifies that the CLI reaches the service; - `health_check` reports component health; -- `status` estimates stateful component memory and process RSS; +- `status` estimates shared and selected-workspace component memory, labels the workspace, and reports process-wide RSS; - `app_config` returns the effective configuration with secrets redacted. ## Logs and common symptoms diff --git a/docs/zh/operations.md b/docs/zh/operations.md index 27571bf9..53ea8ebc 100644 --- a/docs/zh/operations.md +++ b/docs/zh/operations.md @@ -22,7 +22,7 @@ reme app_config - `find_reme`:服务是否存在,以及实际 host、port、PID; - `version`:CLI 能否成功访问服务; - `health_check`:组件健康状态; -- `status`:状态组件的内存估算和进程 RSS; +- `status`:共享组件和所选工作区状态组件的内存估算,标明工作区,并报告整个进程的 RSS; - `app_config`:隐藏密钥后的实际生效配置。 ## 日志 diff --git a/reme/steps/common/status.py b/reme/steps/common/status.py index b2cdcc3e..0d4b5e7d 100644 --- a/reme/steps/common/status.py +++ b/reme/steps/common/status.py @@ -90,14 +90,19 @@ def _component_size(obj: object) -> int: return walk(obj) -def _collect_memory(app_context) -> dict: - """Collect per-component estimates, their sum, and process RSS.""" +def _collect_memory(app_context, workspace=None) -> dict: + """Estimate shared and selected-workspace components, plus process-wide RSS.""" components: dict[str, dict[str, dict[str, int | str]]] = {} total = 0 if app_context is not None: + workspace = workspace or app_context.workspace for component_type in sorted(_TRACKED_COMPONENT_TYPES, key=lambda item: item.value): group = {} - for name, component in sorted(app_context.components.get(component_type, {}).items()): + selected_components = { + **app_context.shared_components.get(component_type, {}), + **workspace.components.get(component_type, {}), + } + for name, component in sorted(selected_components.items()): # _component_size keeps an independent seen set so each component # remains understandable in isolation. Consequently, a non-component # object shared by multiple components may be included more than once @@ -110,6 +115,7 @@ def _collect_memory(app_context) -> dict: rss = psutil.Process().memory_info().rss return { + "workspace": workspace.alias if workspace is not None else None, "components": components, "components_total_bytes": total, "components_total": _format_bytes(total), @@ -121,6 +127,8 @@ def _collect_memory(app_context) -> dict: def _format_status(memory: dict) -> str: """Build the human-readable CLI response.""" lines = ["Memory (estimated component object size)"] + if memory.get("workspace") is not None: + lines.append(f" Workspace {memory['workspace']}") for component_type, group in memory["components"].items(): for name, usage in group.items(): lines.append(f" {component_type}:{name} {usage['human']}") @@ -140,7 +148,7 @@ class StatusStep(BaseStep): async def execute(self): assert self.context is not None - memory = _collect_memory(self.app_context) + memory = _collect_memory(self.app_context, self.context.workspace) self.context.response.answer = _format_status(memory) self.context.response.metadata["status"] = {"memory": memory} return self.context.response diff --git a/tests/unit/test_status_step.py b/tests/unit/test_status_step.py index 1951259e..2ca44899 100644 --- a/tests/unit/test_status_step.py +++ b/tests/unit/test_status_step.py @@ -2,6 +2,8 @@ import asyncio +from reme.application import Application +from reme.schema import FileNode from reme.components.application_context import ApplicationContext from reme.components.base_component import BaseComponent, Dependency from reme.enumeration import ComponentEnum @@ -108,3 +110,45 @@ def test_status_step_returns_human_summary_and_exact_metadata(tmp_path): assert "file_store:default" in response.answer assert "Storage" not in response.answer assert set(response.metadata["status"]) == {"memory"} + + +def test_status_job_attributes_growth_to_selected_workspace(tmp_path): + """The real job includes beta's graph growth without charging it to alpha.""" + + async def check(): + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "local"}}}, + jobs={"status": {"backend": "base", "steps": [{"backend": "status_step"}]}}, + ) + await app.start() + try: + + async def memory(alias): + response = await app.run_job("status", workspace=alias) + assert response.success + report = response.metadata["status"]["memory"] + assert report["workspace"] == alias + assert f"Workspace {alias}" in response.answer + return report + + alpha_before = await memory("alpha") + beta_before = await memory("beta") + graph = app.context.workspaces["beta"].get_component(ComponentEnum.FILE_GRAPH, "default") + await graph.upsert_nodes( + [FileNode(path="beta.md", st_mtime=1, front_matter={"description": "x" * 2_000_000})], + ) + beta_after = await memory("beta") + alpha_after = await memory("alpha") + assert beta_after["components"]["file_graph"]["default"]["bytes"] == _component_size(graph) + assert beta_after["components_total_bytes"] > beta_before["components_total_bytes"] + 1_900_000 + assert alpha_before["components"] == alpha_after["components"] + assert alpha_before["components_total_bytes"] == alpha_after["components_total_bytes"] + finally: + await app.close() + + asyncio.run(check()) From 15f2bdbdc06c3094a36995ee1a33a3a8c7d97922 Mon Sep 17 00:00:00 2001 From: celanwang Date: Thu, 5 Nov 2026 11:43:09 +0800 Subject: [PATCH 20/21] fix(codex): retain thread ownership for uncertain SDK outcomes --- docs/en/configuration.md | 5 + docs/zh/configuration.md | 4 + .../agent_wrapper/codex_agent_wrapper.py | 38 +++- tests/unit/test_codex_compaction_ownership.py | 181 +++++++++++++++++- 4 files changed, 215 insertions(+), 13 deletions(-) diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 4be39008..5cb4c5cb 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -107,6 +107,11 @@ bridge forwards to this service without creating another Application and preserv The multi-workspace one-shot CLI combination is not supported; `mcp_config` retains its legacy meaning only for single-workspace Codex job tools. +Once a Codex thread operation enters the SDK, cancellation or an unknown RPC outcome retains that thread's +workspace ownership across restarts: cancelling the caller does not prove that submitted work stopped. +A provisional claim can be released before SDK entry or after an explicit request-shape, method, or parameter +rejection, provided no other operation still uses or has confirmed it. + Background and cron jobs bind through their job configuration, so extra workspaces are watched by declaring explicit per-alias variants: diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index a556cd43..fbc727d3 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -108,6 +108,10 @@ Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 Job 保留为 tool error。多 workspace 与 one-shot CLI 的组合暂不支持;`mcp_config` 仅在单 workspace 的 Codex job tools 中保留原有语义。 +Codex thread 操作进入 SDK 后,若调用被取消或 RPC 结果不明,会保留 thread 的 workspace 归属,重启后仍然有效; +取消调用方并不能证明已提交的操作停止执行。临时归属仅在进入 SDK 前,或收到明确的请求格式、方法或参数拒绝后, +且没有其他操作仍在使用或已确认该归属时,才可以释放。 + Background 和 cron Job 通过 Job 配置绑定别名,额外 workspace 的 watch 通过显式配置按别名的变体实现: ```yaml diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index f8c95cd7..c6e56b41 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -393,9 +393,9 @@ class CodexAgentWrapper(BaseAgentWrapper): try: thread = await codex.thread_resume(thread_id, personality=personality, **common) await self._record_thread_ownership(alias, thread.id) - except BaseException: + except BaseException as exc: if reserved: - await self._forget_thread_ownership(alias, thread_id) + await self._settle_failed_thread_claim(alias, thread_id, exc, sdk_entered=True) raise else: thread = await codex.thread_start( @@ -550,6 +550,28 @@ class CodexAgentWrapper(BaseAgentWrapper): ids.discard(thread_id) state.provisional_claims.pop(thread_id, None) + async def _settle_failed_thread_claim( + self, + alias: str, + thread_id: str, + error: BaseException, + *, + sdk_entered: bool, + ) -> None: + """Roll back only when no thread RPC ran or the server rejected its input. + + The SDK offloads RPCs to worker threads. Cancellation, transport errors, + and other unknown outcomes cannot establish that submitted work stopped. + Confirm those claims so later rollback and Context recreation stay safe. + """ + if sdk_entered: + from openai_codex.errors import InvalidParamsError, InvalidRequestError, MethodNotFoundError + + if not isinstance(error, (InvalidParamsError, InvalidRequestError, MethodNotFoundError)): + await self._record_thread_ownership(alias, thread_id) + return + await self._forget_thread_ownership(alias, thread_id) + async def _check_thread_ownership(self, thread_id: str, alias: str) -> None: """Reject resume/fork/compact of a thread owned by another workspace. @@ -710,8 +732,10 @@ class CodexAgentWrapper(BaseAgentWrapper): The alias is validated and the cross-workspace claim is settled before any SDK call, so a losing request cannot compact. A claim made for this - call is removed on failure only if no other operation still uses or has - confirmed it. Thread release still runs once a thread was resumed. + call is removed only before SDK entry or on explicit input rejection, + and only if no other operation uses or has confirmed it. Cancellation + after SDK entry retains ownership because the RPC may still execute. + Thread release still runs once a thread was resumed. """ await self.start() alias = str(workspace or "") @@ -719,14 +743,16 @@ class CodexAgentWrapper(BaseAgentWrapper): reserved = await self._record_thread_ownership(alias, session_id, provisional=True) codex = None thread = None + sdk_entered = False try: codex = await self._get_codex() + sdk_entered = True thread = await codex.thread_resume(session_id) await thread.compact() await self._record_thread_ownership(alias, session_id) - except BaseException: + except BaseException as exc: if reserved: - await self._forget_thread_ownership(alias, session_id) + await self._settle_failed_thread_claim(alias, session_id, exc, sdk_entered=sdk_entered) raise finally: if thread is not None and codex is not None and self._release_threads: diff --git a/tests/unit/test_codex_compaction_ownership.py b/tests/unit/test_codex_compaction_ownership.py index 9aff5297..819cd787 100644 --- a/tests/unit/test_codex_compaction_ownership.py +++ b/tests/unit/test_codex_compaction_ownership.py @@ -4,9 +4,12 @@ import asyncio import json +import threading from types import SimpleNamespace import pytest +from openai_codex import AsyncCodex +from openai_codex.errors import InvalidParamsError, TransportClosedError from reme.components import ApplicationContext from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper @@ -106,8 +109,8 @@ async def test_compact_arbitrates_before_sdk_side_effects(tmp_path, monkeypatch) @pytest.mark.asyncio -async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): - """A failed resume or compact releases a claim this call just made.""" +async def test_compact_drops_new_claim_when_sdk_rejects_input(tmp_path, monkeypatch): + """An explicit SDK input rejection releases an otherwise unused provisional claim.""" context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) alpha = CodexAgentWrapper(name="codex-alpha", app_context=context) beta = CodexAgentWrapper(name="codex-beta", app_context=context, release_threads=False) @@ -129,11 +132,11 @@ async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): self.id = thread_id async def compact(self): - raise RuntimeError("sdk rejected") + raise InvalidParamsError(-32602, "sdk rejected") async def resume(thread_id): if thread_id == "resume-fails": - raise RuntimeError("resume rejected") + raise InvalidParamsError(-32602, "resume rejected") return RejectingThread(thread_id) async def get_codex(): @@ -141,12 +144,12 @@ async def test_compact_drops_new_claim_when_sdk_rejects(tmp_path, monkeypatch): monkeypatch.setattr(alpha, "_get_codex", get_codex) - with pytest.raises(RuntimeError, match="resume rejected"): + with pytest.raises(InvalidParamsError, match="resume rejected"): await alpha.compact_session("resume-fails", workspace="alpha") - with pytest.raises(RuntimeError, match="sdk rejected"): + with pytest.raises(InvalidParamsError, match="sdk rejected"): await alpha.compact_session("new-thread", workspace="alpha") await alpha._record_thread_ownership("alpha", "kept-thread") - with pytest.raises(RuntimeError, match="sdk rejected"): + with pytest.raises(InvalidParamsError, match="sdk rejected"): await alpha.compact_session("kept-thread", workspace="alpha") owned = context.workspaces["alpha"].codex_thread_state.owned_thread_ids or set() @@ -273,3 +276,167 @@ async def test_failed_compactor_retains_another_pending_reservation(tmp_path, mo await asyncio.gather(*tasks, return_exceptions=True) for wrapper in wrappers: await wrapper.close() + + +def _resume_response(thread_id, cwd): + """Valid wire response parsed by the real synchronous SDK client.""" + return { + "approvalPolicy": "never", + "approvalsReviewer": "user", + "cwd": str(cwd), + "model": "test", + "modelProvider": "test", + "sandbox": {"type": "dangerFullAccess"}, + "thread": { + "cliVersion": "0.163.0", + "createdAt": 0, + "cwd": str(cwd), + "ephemeral": False, + "id": thread_id, + "modelProvider": "test", + "preview": "", + "sessionId": thread_id, + "source": "appServer", + "status": {"type": "idle"}, + "turns": [], + "updatedAt": 0, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("operation", "blocked_method"), + [("compact", "thread/compact/start"), ("compact", "thread/resume"), ("reply", "thread/resume")], +) +@pytest.mark.parametrize("outcome", ["success", "rejected", "transport_closed"]) +async def test_cancelled_sdk_worker_preserves_persisted_ownership( + tmp_path, + monkeypatch, + operation, + blocked_method, + outcome, +): + """Cancelling the SDK's to_thread waiter leaves its dispatched RPC running.""" + config = {"workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}} + context = ApplicationContext(**config) + alpha = CodexAgentWrapper(name="alpha", app_context=context) # default release_threads=True + sdk = AsyncCodex() + sdk._initialized = True # Replace only startup and the raw RPC transport; never launch a provider. + entered = threading.Event() + release = threading.Event() + finished = threading.Event() + sdk_ops = [] + + def request(method, params=None): + sdk_ops.append(method) + if method == blocked_method: + entered.set() + try: + if not release.wait(timeout=10): + raise TimeoutError("test transport was not released") + if outcome == "rejected": + raise InvalidParamsError(-32602, "rejected") + if outcome == "transport_closed": + raise TransportClosedError("unknown outcome") + finally: + finished.set() + if method == "thread/resume": + return _resume_response(params["threadId"], tmp_path / "alpha") + if method == "thread/compact/start": + return {} + if method == "thread/unsubscribe": + return {"status": "unsubscribed"} + raise AssertionError(f"Unexpected RPC: {method}") + + async def get_sdk(): + return sdk + + monkeypatch.setattr(sdk._client._sync, "_request_raw", request) + monkeypatch.setattr(alpha, "_get_codex", get_sdk) + call = ( + alpha.compact_session("T", workspace="alpha") + if operation == "compact" + else alpha.reply( + "hello", + resume="T", + workspace="alpha", + ) + ) + task = asyncio.create_task(call) + + async def assert_isolated(): + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"T"} + assert _owned_on_disk(tmp_path, "alpha") == {"T"} + for checked_context in (context, ApplicationContext(**config)): + beta = CodexAgentWrapper(name="beta", app_context=checked_context) + monkeypatch.setattr(beta, "_get_codex", get_sdk) + try: + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta.reply("hello", resume="T", workspace="beta") + finally: + await beta.close() + assert not _owned_on_disk(tmp_path, "beta") + + try: + assert await asyncio.to_thread(entered.wait, 3) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + assert not finished.is_set() # The RPC is still executing after cancellation/cleanup. + if blocked_method == "thread/compact/start": + assert sdk_ops == ["thread/resume", "thread/compact/start", "thread/unsubscribe"] + await assert_isolated() + release.set() + assert await asyncio.to_thread(finished.wait, 3) + await assert_isolated() + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await asyncio.to_thread(finished.wait, 3) + await alpha.close() + await sdk.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["thread/resume", "thread/compact/start"]) +@pytest.mark.parametrize("error_type", [TransportClosedError, TimeoutError, RuntimeError]) +async def test_unknown_sdk_failure_retains_ownership(tmp_path, monkeypatch, method, error_type): + """An RPC failure without a rejection response cannot establish non-execution.""" + config = {"workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}} + context = ApplicationContext(**config) + alpha = CodexAgentWrapper(app_context=context) + sdk = AsyncCodex() + sdk._initialized = True + + def request(rpc_method, params=None): + if rpc_method == method: + raise error_type("outcome unknown") + if rpc_method == "thread/resume": + return _resume_response(params["threadId"], tmp_path / "alpha") + if rpc_method == "thread/unsubscribe": + return {"status": "unsubscribed"} + raise AssertionError(f"Unexpected RPC: {rpc_method}") + + async def get_sdk(): + return sdk + + monkeypatch.setattr(sdk._client._sync, "_request_raw", request) + monkeypatch.setattr(alpha, "_get_codex", get_sdk) + try: + with pytest.raises(error_type, match="outcome unknown"): + await alpha.compact_session("T", workspace="alpha") + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"T"} + assert not context.workspaces["alpha"].codex_thread_state.provisional_claims + assert _owned_on_disk(tmp_path, "alpha") == {"T"} + for checked_context in (context, ApplicationContext(**config)): + beta = CodexAgentWrapper(app_context=checked_context) + try: + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta.compact_session("T", workspace="beta") + finally: + await beta.close() + finally: + await alpha.close() + await sdk.close() From bbf7a93d07c8d5376cc6b101fa49f60054017b11 Mon Sep 17 00:00:00 2001 From: celanwang Date: Fri, 6 Nov 2026 11:54:05 +0800 Subject: [PATCH 21/21] fix(health): use explicit workspace component views for runtime reports --- docs/en/operations.md | 2 +- docs/zh/operations.md | 2 +- reme/components/application_context.py | 9 ++ reme/steps/common/health_check.py | 15 ++-- reme/steps/common/status.py | 7 +- tests/unit/test_health_check_workspace.py | 103 ++++++++++++++++++++++ 6 files changed, 125 insertions(+), 13 deletions(-) create mode 100644 tests/unit/test_health_check_workspace.py diff --git a/docs/en/operations.md b/docs/en/operations.md index 90027d6f..bab324d2 100644 --- a/docs/en/operations.md +++ b/docs/en/operations.md @@ -22,7 +22,7 @@ reme app_config - `find_reme` confirms the actual host, port, and PID; - `version` verifies that the CLI reaches the service; -- `health_check` reports component health; +- `health_check` reports shared and selected-workspace component health and labels the workspace; - `status` estimates shared and selected-workspace component memory, labels the workspace, and reports process-wide RSS; - `app_config` returns the effective configuration with secrets redacted. diff --git a/docs/zh/operations.md b/docs/zh/operations.md index 89e8b890..a87303ad 100644 --- a/docs/zh/operations.md +++ b/docs/zh/operations.md @@ -21,7 +21,7 @@ reme app_config - `find_reme`:服务是否存在,以及实际 host、port、PID; - `version`:CLI 能否成功访问服务; -- `health_check`:组件健康状态; +- `health_check`:共享组件与所选工作区组件的健康状态,并标明工作区; - `status`:共享组件和所选工作区状态组件的内存估算,标明工作区,并报告整个进程的 RSS; - `app_config`:隐藏密钥后的实际生效配置。 diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 1cd6cc3d..27b1bf8f 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -87,6 +87,15 @@ class ApplicationContext: for ctype, group in value.items(): self._components[ctype] = group + def components_for(self, workspace: WorkspaceContext) -> dict[str, dict[str, "BaseComponent"]]: + """Snapshot shared and bound-workspace groups without consulting the legacy view.""" + if self.workspaces.get(workspace.alias) is not workspace: + raise ValueError(f"Workspace {workspace.alias!r} is not bound to this Application") + groups = {ctype: dict(group) for ctype, group in self.shared_components.items()} + for ctype, group in workspace.components.items(): + groups.setdefault(ctype, {}).update(group) + return groups + def get_component(self, ctype: str, name: str, alias: str | None = None) -> "BaseComponent | None": """Compatibility resolver for external callers that select a workspace by alias.""" alias = alias or next(iter(self.workspaces)) diff --git a/reme/steps/common/health_check.py b/reme/steps/common/health_check.py index d1a8b2d1..0a247ac9 100644 --- a/reme/steps/common/health_check.py +++ b/reme/steps/common/health_check.py @@ -170,13 +170,14 @@ def _is_healthy(ctype: ComponentEnum, status: dict) -> bool: return True -def _collect_components(app_context) -> tuple[dict, bool]: - """Walk every registered component type and produce {type: {name: status}}, plus overall flag.""" +def _collect_components(app_context, workspace=None) -> tuple[dict, bool]: + """Collect shared and selected-workspace health, never default-workspace substitutes.""" components: dict = {} healthy = True + selected_components = app_context.components_for(workspace or app_context.workspace) for ctype, handler in _HANDLERS.items(): bucket = {} - for name, comp in app_context.components.get(ctype, {}).items(): + for name, comp in selected_components.get(ctype, {}).items(): status = handler(comp) bucket[name] = status if not _is_healthy(ctype, status): @@ -193,15 +194,17 @@ class HealthCheckStep(BaseStep): assert self.context is not None if self.app_context is not None: - components, healthy = _collect_components(self.app_context) + components, healthy = _collect_components(self.app_context, self.context.workspace) else: components, healthy = {}, True - health = {"version": __version__, "healthy": healthy, "components": components} + workspace = self.context.workspace.alias if self.context.workspace is not None else None + health = {"version": __version__, "workspace": workspace, "healthy": healthy, "components": components} self.logger.info(f"[{self.name}] health collected: {health}") emoji = "✅" if healthy else "❌" label = "healthy" if healthy else "unhealthy" - self.context.response.answer = f"{emoji} ReMe v{__version__} - {label}" + scope = f" (workspace: {workspace})" if workspace is not None else "" + self.context.response.answer = f"{emoji} ReMe v{__version__} - {label}{scope}" self.context.response.metadata["health"] = health return self.context.response diff --git a/reme/steps/common/status.py b/reme/steps/common/status.py index 0d4b5e7d..5258ee69 100644 --- a/reme/steps/common/status.py +++ b/reme/steps/common/status.py @@ -96,13 +96,10 @@ def _collect_memory(app_context, workspace=None) -> dict: total = 0 if app_context is not None: workspace = workspace or app_context.workspace + selected_components = app_context.components_for(workspace) for component_type in sorted(_TRACKED_COMPONENT_TYPES, key=lambda item: item.value): group = {} - selected_components = { - **app_context.shared_components.get(component_type, {}), - **workspace.components.get(component_type, {}), - } - for name, component in sorted(selected_components.items()): + for name, component in sorted(selected_components.get(component_type, {}).items()): # _component_size keeps an independent seen set so each component # remains understandable in isolation. Consequently, a non-component # object shared by multiple components may be included more than once diff --git a/tests/unit/test_health_check_workspace.py b/tests/unit/test_health_check_workspace.py new file mode 100644 index 00000000..bae65e25 --- /dev/null +++ b/tests/unit/test_health_check_workspace.py @@ -0,0 +1,103 @@ +"""Health and memory reports must describe the request's actual workspace.""" + +import pytest + +from reme.application import Application +from reme.components.application_context import ApplicationContext +from reme.components.base_component import BaseComponent +from reme.enumeration import ComponentEnum +from reme.schema import FileNode +from reme.steps.common.health_check import HealthCheckStep +from reme.steps.common.status import _collect_memory + + +@pytest.mark.asyncio +async def test_health_job_reports_beta_failure_without_affecting_alpha(tmp_path): + """Exercise real graphs and the public job before and after only beta stops.""" + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "local"}}}, + jobs={"health_check": {"backend": "base", "steps": [{"backend": "health_check_step"}]}}, + ) + async with app: + alpha = app.context.workspaces["alpha"].get_component(ComponentEnum.FILE_GRAPH, "default") + beta = app.context.workspaces["beta"].get_component(ComponentEnum.FILE_GRAPH, "default") + await beta.upsert_nodes([FileNode(path="beta.md", st_mtime=1)]) + alpha_before = (await app.run_job("health_check", workspace="alpha")).metadata["health"] + response = await app.run_job("health_check", workspace="beta") + before = response.metadata["health"] + assert before["components"]["file_graph"]["default"]["n_nodes"] == 1 + assert before["workspace"] == "beta" and before["healthy"] is True + assert "beta" in response.answer + assert before["components"]["file_graph"]["default"]["is_started"] is True + await beta.close() + response = await app.run_job("health_check", workspace="beta") + assert response.success is True # The diagnostic job succeeded; its health result is negative. + after = response.metadata["health"] + assert after["workspace"] == "beta" and after["healthy"] is False + assert after["components"]["file_graph"]["default"]["is_started"] is False + assert after["components"]["file_graph"]["default"]["n_nodes"] == 1 + assert alpha.is_started is True and beta.is_started is False + assert (await app.run_job("health_check", workspace="alpha")).metadata["health"] == alpha_before + assert (await app.run_job("health_check")).metadata["health"] == alpha_before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("named", [False, True]) +async def test_health_single_workspace_labels_default_and_preserves_payload(tmp_path, named): + """String and one-entry mapping reports identify their default workspace.""" + context = ApplicationContext(workspace_dir={"named": str(tmp_path)} if named else str(tmp_path)) + alias = context.workspace.alias + step = HealthCheckStep(app_context=context) + response = await step() + assert response.metadata["health"]["workspace"] == alias + assert response.metadata["health"]["healthy"] is True + assert set(response.metadata["health"]["components"]) == { + "embedding_store", + "file_graph", + "file_store", + "keyword_index", + } + + +@pytest.mark.asyncio +async def test_reports_include_shared_and_selected_components_without_mutating_registries(tmp_path): + """Shared and scoped names of the same type both remain visible; alpha never leaks.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + shared, alpha, beta = (BaseComponent() for _ in range(3)) + context.shared_components[ComponentEnum.FILE_GRAPH] = {"shared": shared} + context.workspaces["alpha"].components[ComponentEnum.FILE_GRAPH] = {"alpha": alpha} + context.workspaces["beta"].components[ComponentEnum.FILE_GRAPH] = {"beta": beta} + await shared.start() + await alpha.start() + await beta.start() + try: + from reme.components.runtime_context import RuntimeContext + + step = HealthCheckStep(app_context=context) + response = await step(RuntimeContext(workspace=context.workspaces["beta"])) + assert set(response.metadata["health"]["components"]["file_graph"]) == {"shared", "beta"} + memory = _collect_memory(context, context.workspaces["beta"]) + assert set(memory["components"]["file_graph"]) == {"shared", "beta"} + await shared.close() + response = await step(RuntimeContext(workspace=context.workspaces["beta"])) + assert response.metadata["health"]["healthy"] is False + assert set(context.shared_components[ComponentEnum.FILE_GRAPH]) == {"shared"} + assert set(context.workspaces["alpha"].components[ComponentEnum.FILE_GRAPH]) == {"alpha"} + assert set(context.workspaces["beta"].components[ComponentEnum.FILE_GRAPH]) == {"beta"} + finally: + await shared.close() + await alpha.close() + await beta.close() + + +def test_component_report_view_rejects_foreign_workspace(tmp_path): + """A workspace object from another application cannot select this application's resources.""" + context = ApplicationContext(workspace_dir=str(tmp_path / "current")) + foreign = ApplicationContext(workspace_dir=str(tmp_path / "foreign")) + with pytest.raises(ValueError, match="not bound to this Application"): + context.components_for(foreign.workspace)