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 fedf7ac0..4be39008 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -78,6 +78,54 @@ 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. 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 +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. + +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 `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: + +```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 6cec4614..a556cd43 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -83,6 +83,49 @@ 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` 别名。多 workspace 配置还要求 `mem_session_dir` 非空, +以确保重启后仍能强制执行 agent thread 的归属隔离。 + +Job 调用接受可选的 `workspace` 字段指定绑定的别名,例如 +`reme read_note path=daily/note.md workspace=scratch`;缺省绑定默认 workspace,未知别名会直接报错。 +绑定仅在单次调用内生效:workspace 相关的状态组件(file store、索引、graph、catalog)按别名解析到各自实例, +共享组件(模型、tokenizer 等)保持单实例。无法隔离 workspace 的外部有状态后端(目前是 `neo4j` file graph) +在多 workspace 配置下会在启动时被拒绝。 + +绑定会经嵌套 Job 调用和 agent 工具调用(`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 通过显式配置按别名的变体实现: + +```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/application.py b/reme/application.py index 5bbe3d14..76c6e88e 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__ @@ -18,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() @@ -54,20 +54,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.""" @@ -79,17 +70,35 @@ 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: + for alias, workspace in self.context.workspaces.items(): + instance = self._instantiate( + ctype, + cfg, + label=label, + expected_type=BaseComponent, + name=name, + workspace_alias=alias, + ) + workspace.components.setdefault(ctype, {})[name] = instance + else: + self.context.shared_components.setdefault(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.""" @@ -102,6 +111,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, @@ -110,6 +128,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. @@ -118,17 +137,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__ @@ -137,16 +155,29 @@ 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 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 + for ctype, group in self.context.shared_components.items(): + for name, comp in group.items(): + 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] @@ -169,13 +200,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: @@ -241,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): @@ -254,6 +301,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 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.shared_components.values(): + for component in group.values(): + instances[id(component)] = component + return list(instances.values()) + def _replacement_shutdown_order( self, old_component: BaseComponent, @@ -264,9 +323,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 @@ -299,8 +356,23 @@ 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) + 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), 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}") @@ -316,6 +388,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(): @@ -325,14 +398,12 @@ 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 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) @@ -357,6 +428,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 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/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 935c39e0..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( @@ -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: @@ -185,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), @@ -223,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: @@ -234,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 @@ -259,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 "" @@ -274,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}") @@ -292,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.""" @@ -305,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", []) @@ -317,8 +328,22 @@ 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._make_tool(job, tool_context_id, kwargs.get("injected_job_kwargs")) for job in resolved_jobs) + tools.extend( + self._builtin_tools( + builtin_tools, + sequential_tool_calls=sequential_tool_calls, + call_kwargs=kwargs, + ), + ) + 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, @@ -350,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, @@ -391,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 ------------------------------------------- @@ -496,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 69f1ace6..4d110099 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,35 @@ 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. + + 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") and kwargs.get("_workspace_explicit", True): + 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 @@ -163,12 +192,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 @@ -211,8 +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} + 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/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 068175df..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.""" @@ -83,11 +88,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): @@ -187,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) @@ -203,7 +214,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..7fd43ee4 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -16,9 +16,12 @@ 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 ...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 @@ -115,8 +118,8 @@ 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._mcp_snapshot_path: Path | None = None @property def session_path(self) -> Path: @@ -128,13 +131,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) @@ -154,6 +150,13 @@ class CodexAgentWrapper(BaseAgentWrapper): relative_source = os.path.relpath(source, target.parent) target.symlink_to(relative_source, target_is_directory=True) + @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 None: @@ -240,6 +243,31 @@ class CodexAgentWrapper(BaseAgentWrapper): names = ", ".join(invalid) raise TypeError(f"Codex client options must be configured on the wrapper: {names}") + 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: + 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 + 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 @@ -251,19 +279,33 @@ class CodexAgentWrapper(BaseAgentWrapper): if unsupported: raise TypeError(f"Codex job_tools must be non-stream request jobs: {', '.join(unsupported)}") - config_source = self._mcp_config_source(kwargs) - args = [ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - config_source, - "--workspace", - str(self.workspace_path), - ] + 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), + ] + 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 "")]) - if injected_job_kwargs := dict(kwargs.get("injected_job_kwargs") or {}): + injected_job_kwargs = dict(kwargs.get("injected_job_kwargs") or {}) + 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)]) return { "command": sys.executable, @@ -300,6 +342,8 @@ class CodexAgentWrapper(BaseAgentWrapper): from openai_codex import ApprovalMode, Sandbox 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: @@ -308,8 +352,10 @@ 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 self._release_threads and self._thread_tool_contexts: - await self._prune_thread_tool_contexts(codex) + if thread_id: + await self._check_thread_ownership(thread_id, alias) + 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 ""), @@ -318,15 +364,15 @@ 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]: + 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 = { "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"), @@ -353,16 +399,171 @@ class CodexAgentWrapper(BaseAgentWrapper): thread_source=thread_source, **common, ) - self._thread_tool_contexts[thread.id] = requested_tool_context + await self._record_thread_ownership(alias, thread.id) + tool_contexts[thread.id] = requested_tool_context return thread + # ----- Thread ownership (per-workspace, persisted) --------------------- + + def _agent_state(self, alias: str) -> WorkspaceAgentState | None: + workspaces = _context_workspaces(self.app_context) + if not 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) + 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 _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 + + 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 + + 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 + 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 False + 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: + 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) + 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. + + 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 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}", + ) + def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: from openai_codex import ApprovalMode, Sandbox from openai_codex.types import Personality, ReasoningEffort, ReasoningSummary return { "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), - "cwd": str(self.cwd), + "cwd": str(self._call_cwd(kwargs)), "effort": self._enum(ReasoningEffort, kwargs.get("effort")), "model": kwargs.get("model"), "output_schema": kwargs.get("output_schema"), @@ -396,10 +597,13 @@ 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 clear session execution guards.""" async with self._turn_lock: codex, self._codex = self._codex, None self._thread_tool_contexts.clear() + for workspace in (_context_workspaces(self.app_context) or {}).values(): + if state := workspace.agent_states.get(self.name): + state.thread_tool_contexts.clear() try: if codex is not None: await codex.close() @@ -408,8 +612,9 @@ class CodexAgentWrapper(BaseAgentWrapper): self._mcp_snapshot_path.unlink(missing_ok=True) self._mcp_snapshot_path = None - 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 @@ -428,9 +633,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 = { - thread_id: context for thread_id, context in self._thread_tool_contexts.items() if thread_id 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.""" @@ -490,16 +695,30 @@ 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: + """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: - 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() + 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/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index 9ed86f5b..1ad168ca 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -1,10 +1,17 @@ -"""FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" +"""Expose ReMe jobs to Codex over FastMCP STDIO. + +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 +from .base_agent_wrapper import BaseAgentWrapper +from ..client import MCPClient from ...config import resolve_app_config from ...reme import ReMe @@ -12,7 +19,9 @@ from ...reme import ReMe def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Expose selected ReMe jobs over FastMCP STDIO") parser.add_argument("--config", default="default", help="ReMe config name or file path") - parser.add_argument("--workspace", required=True, help="ReMe workspace directory") + parser.add_argument("--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") parser.add_argument( @@ -20,6 +29,7 @@ def _parse_args() -> argparse.Namespace: 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() @@ -29,7 +39,7 @@ def _prepare_config( 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.""" + """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(): @@ -37,7 +47,6 @@ def _prepare_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 @@ -63,21 +72,107 @@ def _prepare_config( return prepared +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 + + client = _ResidentMCPClient(args.endpoint, args.tool_timeout) + await client.start() + try: + 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: + await client.close() + + def main() -> None: - """Load ReMe and serve the requested jobs over STDIO.""" + """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") + 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, log_config=False, ) - injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else None - if injected_job_kwargs is not None and not isinstance(injected_job_kwargs, dict): - raise TypeError("--injected-job-kwargs must be a JSON object") - config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) + config = _prepare_config(config, args.jobs, args.tool_context_id, injected) ReMe(**config).run_app() diff --git a/reme/components/application_context.py b/reme/components/application_context.py index 0e0c393f..a5f8ad1e 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -1,9 +1,11 @@ """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 -from ..schema import ApplicationConfig +from ..schema import ApplicationConfig, WorkspaceContext from .component_registry import ComponentRegistry, create_application_registry if TYPE_CHECKING: @@ -12,27 +14,81 @@ 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() + self.workspaces: dict[str, WorkspaceContext] = { + alias: WorkspaceContext.from_config(alias, path, self.app_config) + for alias, path in self.app_config.workspaces.items() + } - # 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.codex_thread_ownership_lock = asyncio.Lock() self.metadata: dict[str, Any] = {} + + @property + def components(self) -> MutableMapping: + """Legacy component lookup; internal execution uses explicit shared or workspace ownership.""" + return self._components + + @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": + """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.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.shared_components.get(ctype, {}).get(name) diff --git a/reme/components/base_component.py b/reme/components/base_component.py index 74ef4633..f4075639 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,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``.""" - target = self.app_context.components.get(dep.ctype, {}).get(dep.name) + """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 + 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: + 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/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_chunker/base_file_chunker.py b/reme/components/file_chunker/base_file_chunker.py index 6d9fed61..475a08c7 100644 --- a/reme/components/file_chunker/base_file_chunker.py +++ b/reme/components/file_chunker/base_file_chunker.py @@ -17,9 +17,14 @@ from ...utils.file_utils import check_file_version 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/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/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/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/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/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/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..512664ca 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,66 @@ 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 + 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 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") + 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}") + 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.""" + 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 ( + "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") @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..96276cb3 --- /dev/null +++ b/reme/schema/workspace.py @@ -0,0 +1,93 @@ +"""Per-workspace path views and owned state for multi-workspace applications.""" + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +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) + + +@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) +class WorkspaceContext: + """Owns the paths, stateful components, execution state, and locks of one workspace.""" + + 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 + components: dict[str, dict[str, Any]] = field(default_factory=dict) + 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: + """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.""" + 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 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 [ + 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 + ] + + 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 b2042237..c70beb97 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,19 @@ 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 (workspaces := _context_workspaces(obj.app_context)) is not None: + context = obj.context + bound = getattr(context, "workspace", None) if context 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) if comp is None: if self.optional: return None @@ -149,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() @@ -189,12 +203,37 @@ 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 ("" only when unbound). + + 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 + return bound.alias if bound is not None else "" + + def _workspace_forward_kwargs(self) -> dict: + """Forward the current call's resolved workspace binding to agent wrapper calls. + + Every resolved binding is forwarded, including the default alias, so + agent tool paths execute against the caller's workspace. + """ + alias = self._bound_workspace_alias() + return {"workspace": alias} if alias else {} + 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..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 @@ -74,9 +75,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 +402,19 @@ 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) + # 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=session_alias, + ) # 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..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: @@ -145,6 +152,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/evolve/proactive/topics.py b/reme/steps/evolve/proactive/topics.py index 9bf5df75..d49caa9b 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,14 @@ 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: + 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") if catalog is None: return [] try: @@ -358,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/_tool_context.py b/reme/steps/index/_tool_context.py index 4b14a6d5..0d9aba3d 100644 --- a/reme/steps/index/_tool_context.py +++ b/reme/steps/index/_tool_context.py @@ -16,6 +16,9 @@ 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: diff --git a/reme/steps/index/search.py b/reme/steps/index/search.py index f4e20b3c..9a5b11b4 100644 --- a/reme/steps/index/search.py +++ b/reme/steps/index/search.py @@ -237,7 +237,11 @@ 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", {}) + 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: diff --git a/reme/steps/index/update_changes.py b/reme/steps/index/update_changes.py index 11a74b01..829be626 100644 --- a/reme/steps/index/update_changes.py +++ b/reme/steps/index/update_changes.py @@ -13,6 +13,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 @@ -534,8 +535,14 @@ 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, {}) + 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_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_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..a2b28ae1 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) @@ -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"}) + 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 6f436d52..4874b4e5 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -8,15 +8,14 @@ import json import os from pathlib import Path import stat -import sys from types import SimpleNamespace 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.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 @@ -64,10 +63,10 @@ def _wrapper(tmp_path, **kwargs): return CodexAgentWrapper(app_context=context, **kwargs), job -def test_mcp_config_uses_stdio_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") - config = wrapper._mcp_server_config( # pylint: disable=protected-access + config = wrapper._mcp_server_config( {"job_tools": ["search", "search"], "tool_context_id": "ctx-1"}, ) @@ -77,13 +76,16 @@ def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): 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("--workspace") + 1] == str(tmp_path) assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" + assert "--endpoint" not in config["args"] + assert "--token" not in config["args"] def test_mcp_config_serializes_injected_job_kwargs(tmp_path): wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") - config = wrapper._mcp_server_config( # pylint: disable=protected-access + config = wrapper._mcp_server_config( { "job_tools": ["search"], "tool_context_id": "ctx-1", @@ -94,27 +96,13 @@ def test_mcp_config_serializes_injected_job_kwargs(tmp_path): 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"] -def test_prepare_config_merges_injected_job_kwargs_with_tool_context(): - prepared = _prepare_config( - {"jobs": {"selected": {"backend": "base"}}}, - ["selected"], - "ctx-1", - {"_allowed_paths": ["daily/2025-06-01/note.md"]}, - ) - - assert prepared["service"]["injected_job_kwargs"] == { - "_allowed_paths": ["daily/2025-06-01/note.md"], - "tool_context_id": "ctx-1", - } - - def test_thread_config_preserves_other_mcp_servers(tmp_path): wrapper, _job = _wrapper(tmp_path) - config = wrapper._thread_config( # pylint: disable=protected-access + config = wrapper._thread_config( { "job_tools": ["search"], "config": {"mcp_servers": {"docs": {"url": "https://example.test/mcp"}}}, @@ -171,114 +159,19 @@ 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", - ) - - assert prepared["service"] == { - "backend": "mcp", - "transport": "stdio", - "jobs": ["selected"], - "tool_error_on_failure": True, - "injected_job_kwargs": {"tool_context_id": "ctx-1"}, - } - assert set(prepared["jobs"]) == {"selected", "helper"} - assert prepared["jobs"]["selected"]["enable_serve"] is True - - -def test_prepare_config_rejects_missing_or_background_selected_jobs(): - with pytest.raises(KeyError, match="missing"): - _prepare_config({"jobs": {}}, ["missing"]) - with pytest.raises(KeyError, match="watch"): - _prepare_config({"jobs": {"watch": {"backend": "background"}}}, ["watch"]) - - -@pytest.mark.asyncio -async def test_stdio_bridge_starts_and_lists_selected_job(tmp_path): - from fastmcp import Client - from fastmcp.client import StdioTransport - - config_path = tmp_path / "bridge.json" - config_path.write_text( - json.dumps( - { - "service": {"backend": "mcp"}, - "workspace_dir": str(tmp_path / "workspace"), - "jobs": { - "empty": { - "backend": "base", - "description": "Return an empty response", - "parameters": {"type": "object", "properties": {}}, - "steps": [], - }, - }, - }, - ), - encoding="utf-8", - ) - transport = StdioTransport( - command=sys.executable, - args=[ - "-m", - "reme.components.agent_wrapper.codex_mcp_server", - "--config", - str(config_path), - "--workspace", - str(tmp_path / "workspace"), - "--job", - "empty", - ], - cwd=str(Path(__file__).resolve().parents[2]), - keep_alive=False, - ) - - async with Client(transport, timeout=None, init_timeout=STDIO_INIT_TIMEOUT) as client: - tools = await asyncio.wait_for(client.list_tools(), timeout=10) - - assert [tool.name for tool in tools] == ["empty"] - - @pytest.mark.asyncio async def test_stdio_bridge_stdout_is_protocol_clean(tmp_path): - config_path = tmp_path / "bridge.json" + config_path = tmp_path / "config.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": [], - }, - }, - }, - ), + 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( - 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 +202,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 @@ -828,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): @@ -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_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"} diff --git a/tests/unit/test_multi_workspace_config.py b/tests/unit/test_multi_workspace_config.py new file mode 100644 index 00000000..7adeabe9 --- /dev/null +++ b/tests/unit/test_multi_workspace_config.py @@ -0,0 +1,271 @@ +"""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.enumeration import ComponentEnum +from reme.schema import ApplicationConfig, FileNode + + +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()) + + +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() + + +@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_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)) + 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"] 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_binding.py b/tests/unit/test_workspace_binding.py new file mode 100644 index 00000000..008b78d5 --- /dev/null +++ b/tests/unit/test_workspace_binding.py @@ -0,0 +1,364 @@ +"""Phase 2: requests and components bind to the selected workspace.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio + +import pytest + +from reme.application import Application +from reme.components import ApplicationContext +from reme.components.base_component import BaseComponent +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 +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 + 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" + + # 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="Unknown 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): + """Stand-in for a stateful external backend that cannot isolate workspaces.""" + + 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() + + +# ----- 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() + with pytest.raises(RuntimeError, match="workspace binding is required"): + await step._digest_titles() + + +@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)) diff --git a/tests/unit/test_workspace_ownership.py b/tests/unit/test_workspace_ownership.py new file mode 100644 index 00000000..2f7ce6c6 --- /dev/null +++ b/tests/unit/test_workspace_ownership.py @@ -0,0 +1,168 @@ +"""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.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 +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") + + +@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" + + 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 + # 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): + 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 new file mode 100644 index 00000000..38ef75cb --- /dev/null +++ b/tests/unit/test_workspace_phase3.py @@ -0,0 +1,800 @@ +"""Phase 3: nested propagation, session namespacing, and background binding.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +import json +import os +from dataclasses import dataclass +from types import SimpleNamespace + +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.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 + + +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() + + # 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 +async def test_codex_thread_tool_contexts_namespaced_per_workspace(tmp_path): + 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.""" + + 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"}) + + 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 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 +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_id) + + async def thread_fork(self, thread_id, **_kwargs): + return SimpleNamespace(id=f"{thread_id}-fork") + + codex = FakeCodex() + 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() + + # 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"}) + + +@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") + + +@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 ----------------------------------------------- + + +@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 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_resident_mcp_endpoint(tmp_path): + from reme.components.service import HttpService + + context = _context(tmp_path) + context.service = HttpService(app_context=context, tool_error_on_failure=True) + 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("--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 "--injected-job-kwargs" not in default_config["args"] + 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() + + 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 +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.""" + + next_id = 1 + + async def thread_start(self, **kwargs): + observed.update(kwargs) + self.next_id += 1 + return SimpleNamespace(id=f"thread-{self.next_id - 1}") + + codex = FakeCodex() + await wrapper._open_thread(codex, {"workspace": "beta"}) + assert observed["cwd"] == str(tmp_path / "beta") + + observed.clear() + await wrapper._open_thread(codex, {}) + assert observed["cwd"] != str(tmp_path / "beta") + + +# ----- Codex bridge: real child-process execution ---------------------------- + + +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=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, + "tool_error_on_failure": True, + }, + 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"}], + }, + "read_note": { + "backend": "base", + "description": "read a note", + "parameters": { + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + }, + "steps": [{"backend": "read_step"}], + }, + }, + ) + + +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: + 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"}, + ) + 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, + ) + 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() + 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_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 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, HttpService + + context = _context(tmp_path) + context.service = CliService(app_context=context) + context.jobs = {"search_job": _job(context, "search_job")} + 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"}, + ) + 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()