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 292fe6c4..e10debfa 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -79,6 +79,61 @@ 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. Codex wrappers select an application-owned `codex_app_server` through `app_server`. +Its scoped HTTP tools execute the original Application's jobs with the resolved workspace and preserve Job failures +as tool errors. HTTP services can host the endpoint; embedded and one-shot CLI hosts use a managed loopback listener. +This works for both single and multiple workspaces. Following the shared-server migration, move process and authentication +settings to `components.codex_app_server.` and remove `mcp_config` and `release_threads`. + +Once a Codex thread operation enters the SDK, cancellation or an unknown RPC outcome retains that thread's +workspace ownership across restarts: cancelling the caller does not prove that submitted work stopped. +A provisional claim can be released before SDK entry or after an explicit request-shape, method, or parameter +rejection, provided no other operation still uses or has confirmed it. A successful resume confirms ownership even +if a subsequent turn or compaction fails. History, status, and proactive append jobs use the same workspace binding; +application-authored session records live under the selected workspace's `session_dir`. + +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/en/operations.md b/docs/en/operations.md index ab1ba8f5..bab324d2 100644 --- a/docs/en/operations.md +++ b/docs/en/operations.md @@ -22,8 +22,8 @@ reme app_config - `find_reme` confirms the actual host, port, and PID; - `version` verifies that the CLI reaches the service; -- `health_check` reports component health; -- `status` estimates stateful component memory and process RSS; +- `health_check` reports shared and selected-workspace component health and labels the workspace; +- `status` estimates shared and selected-workspace component memory, labels the workspace, and reports process-wide RSS; - `app_config` returns the effective configuration with secrets redacted. ## Job execution status diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 1bdffbfd..5bfd4a66 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -84,6 +84,54 @@ 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`）传播，作为模型无法覆盖的服务端约束。 +Codex wrapper 通过 `app_server` 选择 Application 持有的共享 `codex_app_server`。其 scoped HTTP 工具在原有 +Application 中执行 Job，保留已解析的 workspace 绑定，并将失败的 Job 作为 tool error 返回。 +HTTP 服务可挂载该端点；嵌入式和 one-shot CLI 宿主使用受生命周期管理的 loopback listener，单、多 workspace 均支持。 +按照共享服务迁移约定，将进程和认证配置移到 `components.codex_app_server.`，并移除 `mcp_config` 和 `release_threads`。 + +Codex thread 操作进入 SDK 后，若调用被取消或 RPC 结果不明，会保留 thread 的 workspace 归属，重启后仍然有效； +取消调用方并不能证明已提交的操作停止执行。临时归属仅在进入 SDK 前，或收到明确的请求格式、方法或参数拒绝后， +且没有其他操作仍在使用或已确认该归属时，才可以释放。成功 resume 即确认归属，后续 turn 或 compact 失败不会撤销它。 +会话历史、状态和主动消息 Job 使用相同的 workspace 绑定；Application 写入的会话记录保存在所选 workspace 的 `session_dir` 下。 + +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/docs/zh/operations.md b/docs/zh/operations.md index 248f5506..a87303ad 100644 --- a/docs/zh/operations.md +++ b/docs/zh/operations.md @@ -21,8 +21,8 @@ reme app_config - `find_reme`：服务是否存在，以及实际 host、port、PID； - `version`：CLI 能否成功访问服务； -- `health_check`：组件健康状态； -- `status`：状态组件的内存估算和进程 RSS； +- `health_check`：共享组件与所选工作区组件的健康状态，并标明工作区； +- `status`：共享组件和所选工作区状态组件的内存估算，标明工作区，并报告整个进程的 RSS； - `app_config`：隐藏密钥后的实际生效配置。 ## Job 执行状态 diff --git a/reme/application.py b/reme/application.py index 2d81c258..2c0db79f 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: @@ -258,11 +306,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): @@ -271,6 +318,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, @@ -281,9 +340,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 @@ -316,8 +373,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}") @@ -333,6 +405,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(): @@ -342,14 +415,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) @@ -374,6 +445,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 bf6fe58f..55908a82 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -158,7 +158,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( @@ -166,9 +166,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: @@ -186,9 +191,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), @@ -224,10 +230,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: @@ -235,15 +246,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 @@ -260,7 +271,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 "" @@ -275,7 +286,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}") @@ -293,8 +304,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.""" @@ -311,7 +322,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", []) @@ -323,8 +334,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, @@ -415,7 +440,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 ------------------------------------------- @@ -533,7 +558,7 @@ class AsAgentWrapper(BaseAgentWrapper): finally: if not finishing and model_usages: self._record_token_usage(TokenUsage.combine(model_usages)) - await self._dump_state(agent.state) + await self._dump_state(agent.state, kwargs) terminal.chunk = {key: value for key, value in result.items() if key != "events"} terminal.status = result["status"] terminal.metadata["error"] = result["error"] diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index f195de38..62469079 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -7,7 +7,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 AgentReply, SessionMessage, SessionSnapshot, SessionStatus, 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 @@ -277,11 +319,11 @@ class BaseAgentWrapper(BaseComponent): """Request compaction of one persisted agent session.""" raise NotImplementedError(f"{type(self).__name__} does not support session compaction") - async def read_session(self, session_id: str) -> SessionSnapshot: + async def read_session(self, session_id: str, *, workspace: str | None = None) -> SessionSnapshot: """Read display history without waiting for the session to become idle.""" raise NotImplementedError(f"{type(self).__name__} does not support session reading") - async def get_session_status(self, session_id: str) -> SessionStatus: + async def get_session_status(self, session_id: str, *, workspace: str | None = None) -> SessionStatus: """Read activity without resuming or waiting for the session.""" raise NotImplementedError(f"{type(self).__name__} does not support session status") @@ -292,6 +334,7 @@ class BaseAgentWrapper(BaseComponent): *, message_id: str, timeout: float = 30, + workspace: str | None = None, ) -> SessionMessage: """Append an idempotent proactive assistant message after the active turn finishes.""" raise NotImplementedError(f"{type(self).__name__} does not support session appending") diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 250a36a0..ec8d2955 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -44,10 +44,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.""" @@ -84,11 +89,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): @@ -188,9 +198,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) @@ -204,7 +215,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 d9b91c80..f6e86fd8 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -123,7 +123,7 @@ class CodexAgentWrapper(BaseAgentWrapper): "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.deny_all), "base_instructions": kwargs.get("base_instructions"), "config": config, - "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"), @@ -168,6 +168,7 @@ class CodexAgentWrapper(BaseAgentWrapper): async with self.app_server.thread( thread_id, fork=fork, + workspace=kwargs.get("workspace"), options={} if compact else self._thread_options(kwargs), jobs=jobs, injected=self._resolve_injected_job_kwargs(kwargs), @@ -256,20 +257,27 @@ class CodexAgentWrapper(BaseAgentWrapper): raw=self._serialize(result), ) - async def compact_session(self, session_id: str) -> None: + def _session_workspace(self, workspace: str | None) -> str | None: + """Resolve an explicit alias or the wrapper's configured/default binding.""" + return self._merged_kwargs({} if workspace is None else {"workspace": workspace}).get("workspace") + + async def compact_session(self, session_id: str, workspace: str | None = None) -> None: """Hold the conversation lock until native compaction has finished.""" - async with self._session({"session_id": session_id}, compact=True) as thread: + async with self._session( + {"session_id": session_id, "workspace": self._session_workspace(workspace)}, + compact=True, + ) as thread: await self.app_server.compact(thread) - async def read_session(self, session_id: str) -> SessionSnapshot: + async def read_session(self, session_id: str, *, workspace: str | None = None) -> SessionSnapshot: """Read history without acquiring the generation lease.""" await self.start() - return await self.app_server.read_session(session_id) + return await self.app_server.read_session(session_id, workspace=self._session_workspace(workspace)) - async def get_session_status(self, session_id: str) -> SessionStatus: + async def get_session_status(self, session_id: str, *, workspace: str | None = None) -> SessionStatus: """Report native thread activity, including approval/input waits.""" await self.start() - return await self.app_server.get_session_status(session_id) + return await self.app_server.get_session_status(session_id, workspace=self._session_workspace(workspace)) async def append_session( self, @@ -278,10 +286,17 @@ class CodexAgentWrapper(BaseAgentWrapper): *, message_id: str, timeout: float = 30, + workspace: str | None = None, ) -> SessionMessage: """Append a proactive message through the shared session owner.""" await self.start() - return await self.app_server.append_session(session_id, content, message_id=message_id, timeout=timeout) + return await self.app_server.append_session( + session_id, + content, + message_id=message_id, + timeout=timeout, + workspace=self._session_workspace(workspace), + ) def _turn_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: from openai_codex import ApprovalMode, Sandbox @@ -289,7 +304,7 @@ class CodexAgentWrapper(BaseAgentWrapper): return { "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode")), - "cwd": str(self.cwd), + "cwd": str(self._call_cwd(kwargs)), "effort": self._enum(ReasoningEffort, kwargs.get("effort")), "model": kwargs.get("model"), "output_schema": kwargs.get("output_schema"), diff --git a/reme/components/application_context.py b/reme/components/application_context.py index e5059c44..27b1bf8f 100644 --- a/reme/components/application_context.py +++ b/reme/components/application_context.py @@ -1,11 +1,13 @@ """Application context: shared state container for components, jobs, and service.""" +import asyncio +from collections.abc import Iterator, MutableMapping from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from threading import Lock from typing import TYPE_CHECKING, Any, get_args -from ..schema import ApplicationConfig, JobExecution, JobExecutionStatus +from ..schema import ApplicationConfig, JobExecution, JobExecutionStatus, WorkspaceContext from .component_registry import ComponentRegistry, create_application_registry if TYPE_CHECKING: @@ -14,33 +16,97 @@ if TYPE_CHECKING: from .service import BaseService -class ApplicationContext: - """Shared state container holding parsed config and wired components. +class _LegacyComponentView(MutableMapping): + """Compatibility mapping combining shared components with default-workspace instances.""" - Application performs the wiring; this class stores the results and provides - synchronized access to in-memory job execution snapshots. - """ + 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 + self.codex_thread_ownership_lock = asyncio.Lock() + self.metadata: dict[str, Any] = {} self.job_history: dict[str, list[JobExecution]] = {} self._job_counts: dict[str, int] = {} self._job_lock = Lock() self._job_execution_id = 0 - # Application-lifetime shared state. Values remain available across Job and Step - # invocations while this Application is running, so Jobs may keep cross-call state here. - # This is in-memory state, not durable storage; use workspace files or a store when state - # must survive an Application restart. - self.metadata: dict[str, Any] = {} + @property + def 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 components_for(self, workspace: WorkspaceContext) -> dict[str, dict[str, "BaseComponent"]]: + """Snapshot shared and bound-workspace groups without consulting the legacy view.""" + if self.workspaces.get(workspace.alias) is not workspace: + raise ValueError(f"Workspace {workspace.alias!r} is not bound to this Application") + groups = {ctype: dict(group) for ctype, group in self.shared_components.items()} + for ctype, group in workspace.components.items(): + groups.setdefault(ctype, {}).update(group) + return groups + + def get_component(self, ctype: str, name: str, alias: str | None = None) -> "BaseComponent | None": + """Compatibility resolver for external callers that select a workspace by alias.""" + alias = alias or next(iter(self.workspaces)) + 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) def start_job_execution(self, name: str, *, track_execution: bool = True) -> JobExecution | None: """Count every invocation and optionally append a running execution.""" 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/codex_app_server.py b/reme/components/codex_app_server.py index 91fa7532..aac7259b 100644 --- a/reme/components/codex_app_server.py +++ b/reme/components/codex_app_server.py @@ -7,13 +7,15 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field import json +import os +import tempfile import hashlib import math from pathlib import Path import sys from typing import Any, TYPE_CHECKING -from .base_component import BaseComponent +from .base_component import BaseComponent, _context_workspaces from .component_registry import R from .service.scoped_mcp import ScopedMCP from ..constants import normalize_connect_host @@ -90,45 +92,59 @@ class CodexAppServer(BaseComponent): return path if path.is_absolute() else self.workspace_path / path return self.workspace_path / ".codex" - def _message_path(self, session_id: str) -> Path: + def _message_path(self, session_id: str, workspace: str | None = None) -> Path: from .agent_wrapper.codex_session_history import validate_identifier validate_identifier(session_id) session_dir = self.app_context.app_config.session_dir if self.app_context else "session" namespace = hashlib.sha256(str(self.session_path.resolve()).encode()).hexdigest()[:16] - path = (self.workspace_path / session_dir / "proactive" / namespace / f"{session_id}.jsonl").resolve() - if not path.is_relative_to(self.workspace_path.resolve()): + root = self.workspace_path + if workspaces := _context_workspaces(self.app_context): + root = workspaces[self._require_thread_alias(str(workspace or ""))].workspace_path + path = (root / session_dir / "proactive" / namespace / f"{session_id}.jsonl").resolve() + if not path.is_relative_to(root.resolve()): raise ValueError("Session messages must stay inside the workspace") return path @asynccontextmanager - async def _session_client(self) -> AsyncIterator[AsyncCodex]: - """Keep status/history RPCs under lifecycle ownership without taking the turn lock.""" + async def _session_client(self, session_id: str, workspace: str | None = None) -> AsyncIterator[AsyncCodex]: + """Guard read-only RPCs without taking the generation lease.""" + alias = self._require_thread_alias(str(workspace or "")) await self.start() self._active += 1 self._idle.clear() + reserved = False + sdk_entered = False try: - yield await self._get_client() + reserved = await self._record_thread_ownership(alias, session_id, provisional=True) + codex = await self._get_client() + sdk_entered = True + yield codex + await self._record_thread_ownership(alias, session_id) + except BaseException as exc: + if reserved: + await self._settle_failed_thread_claim(alias, session_id, exc, sdk_entered=sdk_entered) + raise finally: self._active -= 1 if not self._active: self._idle.set() - async def get_session_status(self, session_id: str) -> SessionStatus: + async def get_session_status(self, session_id: str, *, workspace: str | None = None) -> SessionStatus: """Inspect a thread without resuming it or waiting for an active turn.""" from .agent_wrapper.codex_session_history import session_status - self._message_path(session_id) - async with self._session_client() as codex: + self._message_path(session_id, workspace) + async with self._session_client(session_id, workspace) as codex: thread = (await codex._client.thread_read(session_id)).thread # pylint: disable=protected-access return session_status(thread) - async def read_session(self, session_id: str) -> SessionSnapshot: + async def read_session(self, session_id: str, *, workspace: str | None = None) -> SessionSnapshot: """Project public native history plus application-authored proactive messages.""" from .agent_wrapper import codex_session_history as history - path = self._message_path(session_id) - async with self._session_client() as codex: + path = self._message_path(session_id, workspace) + async with self._session_client(session_id, workspace) as codex: response = await codex._client.thread_read( # pylint: disable=protected-access session_id, include_turns=True, @@ -149,12 +165,13 @@ class CodexAppServer(BaseComponent): *, message_id: str, timeout: float = 30, + workspace: str | None = None, ) -> SessionMessage: """Wait for the thread lease, then inject and publish one retry-safe assistant message.""" from openai_codex.generated.v2_all import ThreadInjectItemsResponse from .agent_wrapper import codex_session_history as history - path = self._message_path(session_id) + path = self._message_path(session_id, workspace) history.validate_identifier(message_id) if not isinstance(content, str) or not content.strip() or len(content) > 100_000: raise ValueError("content must contain 1-100000 characters") @@ -162,7 +179,7 @@ class CodexAppServer(BaseComponent): raise ValueError("timeout must be positive and finite") async with asyncio.timeout(timeout): - async with self.thread(session_id) as loaded: + async with self.thread(session_id, workspace=workspace) as loaded: thread = (await loaded.read(include_turns=True)).thread if history.session_status(thread).status != "idle": raise RuntimeError("Session must be idle before appending") @@ -442,11 +459,11 @@ class CodexAppServer(BaseComponent): self._threads[thread.id].token_usage = None await finish_cleanup(dispatch()) - async def _fork_messages(self, source_id: str, thread: AsyncThread) -> None: + async def _fork_messages(self, source_id: str, thread: AsyncThread, workspace: str | None = None) -> None: """Copy inherited proactive provenance while the source conversation is locked.""" from .agent_wrapper import codex_session_history as history - records = await asyncio.to_thread(history.read_records, self._message_path(source_id)) + records = await asyncio.to_thread(history.read_records, self._message_path(source_id, workspace)) if not records: return accepted = set() @@ -461,7 +478,7 @@ class CodexAppServer(BaseComponent): if record["injected"] or record["message"]["id"] in accepted ] if inherited: - await complete_in_thread(history.write_records, self._message_path(thread.id), inherited) + await complete_in_thread(history.write_records, self._message_path(thread.id, workspace), inherited) @asynccontextmanager async def thread( @@ -469,6 +486,7 @@ class CodexAppServer(BaseComponent): thread_id: str = "", *, fork: bool = False, + workspace: str | None = None, options: dict[str, Any] | None = None, jobs: list[Any] | None = None, injected: dict[str, Any] | None = None, @@ -478,6 +496,9 @@ class CodexAppServer(BaseComponent): """Serialize one conversation while allowing independent conversations to run concurrently.""" from openai_codex import TransportClosedError + alias = self._require_thread_alias(str(workspace or "")) + reserved = False + sdk_entered = False await self.start() self._active += 1 self._idle.clear() @@ -486,6 +507,8 @@ class CodexAppServer(BaseComponent): thread = None codex = None try: + if thread_id: + reserved = await self._record_thread_ownership(alias, thread_id, provisional=True) codex = await self._get_client() await self._prune_threads(codex) state = self._threads.setdefault(thread_id, _ThreadState()) if thread_id and not fork else _ThreadState() @@ -522,6 +545,7 @@ class CodexAppServer(BaseComponent): } config["mcp_servers"] = servers options["config"] = config + sdk_entered = True if fork: source = self._threads.setdefault(thread_id, _ThreadState()) source.users += 1 @@ -530,7 +554,7 @@ class CodexAppServer(BaseComponent): async with source.lock: state.token_usage = await self.token_usage(thread_id) thread = await codex.thread_fork(thread_id, **options) - await finish_cleanup(self._fork_messages(thread_id, thread)) + await finish_cleanup(self._fork_messages(thread_id, thread, alias)) finally: source.users -= 1 elif thread_id: @@ -538,6 +562,10 @@ class CodexAppServer(BaseComponent): else: thread = await codex.thread_start(**options) state.token_usage = TokenUsage() + if thread_id: + await self._record_thread_ownership(alias, thread_id) + reserved = False + await self._record_thread_ownership(alias, thread.id) self._threads[thread.id] = state if thread_id and not fork: await self.wait_idle(thread) @@ -570,8 +598,10 @@ class CodexAppServer(BaseComponent): ) except Exception as exc: self.logger.warning(f"Failed to unsubscribe Codex thread {thread.id}: {exc}") - except TransportClosedError: - if codex is not None and self._codex is codex: + except BaseException as exc: + if reserved: + await self._settle_failed_thread_claim(alias, thread_id, exc, sdk_entered=sdk_entered) + if isinstance(exc, TransportClosedError) and codex is not None and self._codex is codex: self._failed = True raise finally: @@ -584,6 +614,156 @@ class CodexAppServer(BaseComponent): if not self._active: self._idle.set() + 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, *, provisional: bool = False) -> bool: + """Reserve or confirm ownership under the application-wide lock. + + Provisional callers receive True while their reservation needs rollback + on failure. Successful SDK operations confirm the claim for all wrappers. + Persisted claims loaded after restart are always treated as confirmed. + """ + workspaces = _context_workspaces(self.app_context) + if not workspaces: + 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: + claims = self._codex_thread_state(alias).provisional_claims + if not provisional: + claims.pop(thread_id, None) + elif thread_id in claims: + claims[thread_id] += 1 + return True + return False + raise ValueError( + f"Codex thread {thread_id} belongs to workspace {owner_alias!r}; " + 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) + async with state.lock: + if path is not None: + self._write_owned_thread_ids(path, {*ids, thread_id}) + ids.add(thread_id) + if provisional: + state.provisional_claims[thread_id] = 1 + return True + + async def _forget_thread_ownership(self, alias: str, thread_id: str) -> None: + """Release a failed reservation without revoking confirmed or active ownership.""" + workspaces = _context_workspaces(self.app_context) + if not workspaces: + return + alias = self._normalize_thread_alias(alias) + async with self.app_context.codex_thread_ownership_lock: + state = self._codex_thread_state(alias) + count = state.provisional_claims.get(thread_id, 0) + if not count: + return + if count > 1: + state.provisional_claims[thread_id] -= 1 + return + ids = await self._owned_thread_ids(alias) + if thread_id not in ids: + return + path = self._thread_ownership_path(alias) + async with state.lock: + updated = ids - {thread_id} + if path is not None: + self._write_owned_thread_ids(path, updated) + ids.discard(thread_id) + state.provisional_claims.pop(thread_id, None) + + async def _settle_failed_thread_claim( + self, + alias: str, + thread_id: str, + error: BaseException, + *, + sdk_entered: bool, + ) -> None: + """Roll back only when no thread RPC ran or the server rejected its input. + + The SDK offloads RPCs to worker threads. Cancellation, transport errors, + and other unknown outcomes cannot establish that submitted work stopped. + Confirm those claims so later rollback and Context recreation stay safe. + """ + if sdk_entered: + from openai_codex.errors import InvalidParamsError, InvalidRequestError, MethodNotFoundError + + if not isinstance(error, (InvalidParamsError, InvalidRequestError, MethodNotFoundError)): + await self._record_thread_ownership(alias, thread_id) + return + await self._forget_thread_ownership(alias, thread_id) + async def _close(self) -> None: await self._idle.wait() try: 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 29af2cee..24a25213 100644 --- a/reme/components/file_catalog/local_file_catalog.py +++ b/reme/components/file_catalog/local_file_catalog.py @@ -20,6 +20,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 aa2fe5d3..98f8d60f 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -131,7 +131,9 @@ class BackgroundJob(BaseJob): async def _run_once(self) -> None: """Track the actual body, including custom __call__ implementations and restarts.""" - context = RuntimeContext(stop_event=self._stop_event, **self.kwargs) + merged = dict(self.kwargs) + workspace = self._bind_workspace(merged) + context = RuntimeContext(stop_event=self._stop_event, workspace=workspace, **merged) token = _supervised_job.set(self) try: with self._execution(context): @@ -144,7 +146,8 @@ class BackgroundJob(BaseJob): async def __call__(self, **kwargs) -> Response: """Default body: run steps in order; errors propagate to supervisor.""" 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) if _supervised_job.get() is self: # Consume once so a nested call of this same job gets its own record. _supervised_job.set(None) diff --git a/reme/components/job/base_job.py b/reme/components/job/base_job.py index 08165431..d741001d 100644 --- a/reme/components/job/base_job.py +++ b/reme/components/job/base_job.py @@ -5,7 +5,7 @@ from contextlib import contextmanager from typing import TYPE_CHECKING from ..application_context import ApplicationContext -from ..base_component import BaseComponent +from ..base_component import BaseComponent, _context_workspaces from ..component_registry import R from ..runtime_context import RuntimeContext, _current_job_execution from ...enumeration import ComponentEnum @@ -109,6 +109,25 @@ 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 + @contextmanager def _execution(self, context: RuntimeContext): """Count one call and optionally track it, preserving the caller's error handling.""" @@ -151,7 +170,8 @@ class BaseJob(BaseComponent): async def __call__(self, **kwargs) -> Response: """Run all steps in order, capturing any failure into the response.""" merged = {**self.kwargs, **kwargs} - context = RuntimeContext(**merged) + workspace = self._bind_workspace(merged) + context = RuntimeContext(workspace=workspace, **merged) with self._execution(context): try: await self._run_steps(context) diff --git a/reme/components/job/cron_job.py b/reme/components/job/cron_job.py index 5c8a49eb..180b5f1e 100644 --- a/reme/components/job/cron_job.py +++ b/reme/components/job/cron_job.py @@ -37,7 +37,9 @@ class CronJob(BackgroundJob): return max(0.0, delay) async def _execute_steps(self) -> Response: - context = RuntimeContext(stop_event=self._stop_event, **self.kwargs) + merged = dict(self.kwargs) + workspace = self._bind_workspace(merged) + context = RuntimeContext(stop_event=self._stop_event, workspace=workspace, **merged) with self._execution(context): return await self._run_steps(context) diff --git a/reme/components/job/stream_job.py b/reme/components/job/stream_job.py index 662575b2..866adde0 100644 --- a/reme/components/job/stream_job.py +++ b/reme/components/job/stream_job.py @@ -13,7 +13,8 @@ class StreamJob(BaseJob): async def __call__(self, **kwargs) -> None: """Run steps; emit failures as ERROR chunks, then a terminal DONE marker.""" merged = {**self.kwargs, **kwargs} - context = RuntimeContext(**merged) + workspace = self._bind_workspace(merged) + context = RuntimeContext(workspace=workspace, **merged) with self._execution(context): try: await self._run_steps(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 b90abe8e..59266352 100644 --- a/reme/components/runtime_context.py +++ b/reme/components/runtime_context.py @@ -1,6 +1,7 @@ """Per-request runtime context shared across steps and jobs.""" import asyncio +from typing import TYPE_CHECKING from contextvars import ContextVar from ..enumeration import ChunkEnum @@ -8,12 +9,17 @@ from ..schema import JobExecution, Response, StreamChunk _current_job_execution: ContextVar[JobExecution | None] = ContextVar("reme_job_execution", default=None) +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__( @@ -21,11 +27,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 self.job_execution: JobExecution | None = _current_job_execution.get() self.step_error: Exception | None = None 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 de9e5717..32af19e7 100644 --- a/reme/schema/__init__.py +++ b/reme/schema/__init__.py @@ -28,6 +28,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__ = [ "AgentReply", @@ -68,4 +69,5 @@ __all__ = [ "TraverseGraph", "TraverseGraphEdge", "TraverseGraphNode", + "WorkspaceContext", ] diff --git a/reme/schema/application_config.py b/reme/schema/application_config.py index cb6a55f9..a7d94e97 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.""" @@ -41,9 +45,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``. @@ -67,6 +80,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..3394b3b4 --- /dev/null +++ b/reme/schema/workspace.py @@ -0,0 +1,95 @@ +"""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 + # Active operations on unconfirmed claims; protected by the application ownership lock. + provisional_claims: dict[str, int] = field(default_factory=dict) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +@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 4357103a..cd718385 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())) execution = self.context.job_execution previous = None failed = False @@ -217,12 +231,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 5bbf22cc..682a1d9c 100644 --- a/reme/steps/common/chat.py +++ b/reme/steps/common/chat.py @@ -29,6 +29,7 @@ this chat intentionally provides read-only tools. Reply in the user's language." "system_prompt": str(self.context.get("system_prompt") or self.DEFAULT_SYSTEM_PROMPT), "job_tools": self.kwargs.get("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/common/health_check.py b/reme/steps/common/health_check.py index d1a8b2d1..0a247ac9 100644 --- a/reme/steps/common/health_check.py +++ b/reme/steps/common/health_check.py @@ -170,13 +170,14 @@ def _is_healthy(ctype: ComponentEnum, status: dict) -> bool: return True -def _collect_components(app_context) -> tuple[dict, bool]: - """Walk every registered component type and produce {type: {name: status}}, plus overall flag.""" +def _collect_components(app_context, workspace=None) -> tuple[dict, bool]: + """Collect shared and selected-workspace health, never default-workspace substitutes.""" components: dict = {} healthy = True + selected_components = app_context.components_for(workspace or app_context.workspace) for ctype, handler in _HANDLERS.items(): bucket = {} - for name, comp in app_context.components.get(ctype, {}).items(): + for name, comp in selected_components.get(ctype, {}).items(): status = handler(comp) bucket[name] = status if not _is_healthy(ctype, status): @@ -193,15 +194,17 @@ class HealthCheckStep(BaseStep): assert self.context is not None if self.app_context is not None: - components, healthy = _collect_components(self.app_context) + components, healthy = _collect_components(self.app_context, self.context.workspace) else: components, healthy = {}, True - health = {"version": __version__, "healthy": healthy, "components": components} + workspace = self.context.workspace.alias if self.context.workspace is not None else None + health = {"version": __version__, "workspace": workspace, "healthy": healthy, "components": components} self.logger.info(f"[{self.name}] health collected: {health}") emoji = "✅" if healthy else "❌" label = "healthy" if healthy else "unhealthy" - self.context.response.answer = f"{emoji} ReMe v{__version__} - {label}" + scope = f" (workspace: {workspace})" if workspace is not None else "" + self.context.response.answer = f"{emoji} ReMe v{__version__} - {label}{scope}" self.context.response.metadata["health"] = health return self.context.response diff --git a/reme/steps/common/session.py b/reme/steps/common/session.py index 8431ff8d..4a395448 100644 --- a/reme/steps/common/session.py +++ b/reme/steps/common/session.py @@ -2,6 +2,7 @@ from ..base_step import BaseStep from ...components import R +from ...components.base_component import _context_workspaces class _SessionStep(BaseStep): @@ -23,6 +24,8 @@ class _SessionStep(BaseStep): "message_id": self.context.get("message_id"), "timeout": float(self.context.get("wait_timeout", 30)), } + if (workspaces := _context_workspaces(self.app_context)) and len(workspaces) > 1: + kwargs["workspace"] = self._bound_workspace_alias() try: result = await getattr(self.agent_wrapper, self.operation)(session_id, **kwargs) except NotImplementedError: diff --git a/reme/steps/common/status.py b/reme/steps/common/status.py index b2cdcc3e..5258ee69 100644 --- a/reme/steps/common/status.py +++ b/reme/steps/common/status.py @@ -90,14 +90,16 @@ def _component_size(obj: object) -> int: return walk(obj) -def _collect_memory(app_context) -> dict: - """Collect per-component estimates, their sum, and process RSS.""" +def _collect_memory(app_context, workspace=None) -> dict: + """Estimate shared and selected-workspace components, plus process-wide RSS.""" components: dict[str, dict[str, dict[str, int | str]]] = {} total = 0 if app_context is not None: + workspace = workspace or app_context.workspace + selected_components = app_context.components_for(workspace) for component_type in sorted(_TRACKED_COMPONENT_TYPES, key=lambda item: item.value): group = {} - for name, component in sorted(app_context.components.get(component_type, {}).items()): + for name, component in sorted(selected_components.get(component_type, {}).items()): # _component_size keeps an independent seen set so each component # remains understandable in isolation. Consequently, a non-component # object shared by multiple components may be included more than once @@ -110,6 +112,7 @@ def _collect_memory(app_context) -> dict: rss = psutil.Process().memory_info().rss return { + "workspace": workspace.alias if workspace is not None else None, "components": components, "components_total_bytes": total, "components_total": _format_bytes(total), @@ -121,6 +124,8 @@ def _collect_memory(app_context) -> dict: def _format_status(memory: dict) -> str: """Build the human-readable CLI response.""" lines = ["Memory (estimated component object size)"] + if memory.get("workspace") is not None: + lines.append(f" Workspace {memory['workspace']}") for component_type, group in memory["components"].items(): for name, usage in group.items(): lines.append(f" {component_type}:{name} {usage['human']}") @@ -140,7 +145,7 @@ class StatusStep(BaseStep): async def execute(self): assert self.context is not None - memory = _collect_memory(self.app_context) + memory = _collect_memory(self.app_context, self.context.workspace) self.context.response.answer = _format_status(memory) self.context.response.metadata["status"] = {"memory": memory} return self.context.response diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index 3f90a112..52107b9a 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -676,6 +676,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(), ) text = agent_reply_result_text(result) 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 028b7ee1..c23a558c 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(), ) text = agent_reply_result_text(result) diff --git a/reme/steps/evolve/base_auto_resource.py b/reme/steps/evolve/base_auto_resource.py index 8a1ab07e..cecdb36c 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 14f6c13a..bf1c71b0 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(), ) raw_result = agent_reply_result_text(result) except Exception as e: # noqa: BLE001 diff --git a/reme/steps/evolve/proactive/topics.py b/reme/steps/evolve/proactive/topics.py index 9d03fe70..4c59931d 100644 --- a/reme/steps/evolve/proactive/topics.py +++ b/reme/steps/evolve/proactive/topics.py @@ -20,6 +20,7 @@ import re 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 @@ -337,7 +338,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: @@ -355,7 +363,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/codex_workspace_test_support.py b/tests/unit/codex_workspace_test_support.py new file mode 100644 index 00000000..9a07fe9c --- /dev/null +++ b/tests/unit/codex_workspace_test_support.py @@ -0,0 +1,25 @@ +"""Shared-server SDK doubles for workspace ownership and tool-boundary tests.""" + +# pylint: disable=protected-access + +from reme.components.codex_app_server import CodexAppServer +from reme.enumeration import ComponentEnum + +from .test_codex_agent_wrapper import _Codex, _Thread + + +class WorkspaceCodex(_Codex): + """Permit first use of an external thread while retaining real server lifecycle hooks.""" + + async def thread_resume(self, thread_id, **options): + if thread_id not in self.threads: + self.threads[thread_id] = _Thread(self, thread_id) + return await super().thread_resume(thread_id, **options) + + +def workspace_server(context, monkeypatch): + """Register the application-owned server without contacting a provider.""" + monkeypatch.setattr("openai_codex.AsyncCodex", WorkspaceCodex) + server = CodexAppServer(name="default", app_context=context, auth_mode="oauth") + context.shared_components.setdefault(ComponentEnum.CODEX_APP_SERVER, {})["default"] = server + return server diff --git a/tests/unit/test_agent_reply.py b/tests/unit/test_agent_reply.py index 8a5a6124..f3acff1b 100644 --- a/tests/unit/test_agent_reply.py +++ b/tests/unit/test_agent_reply.py @@ -142,7 +142,11 @@ async def test_agentscope_propagates_cancellation_even_if_sdk_returns_a_message( task.cancel() with pytest.raises(asyncio.CancelledError): await task - wrapper._dump_state.assert_awaited_once_with(agent.state) # pylint: disable=protected-access + dump_state = wrapper._dump_state # pylint: disable=protected-access + dump_state.assert_awaited_once_with( + agent.state, + {"workspace": "default", "_workspace_explicit": False}, + ) counters = global_counter_get_all(wrapper.app_context.metadata, ["__token_counter", wrapper.name]) assert counters["children"]["total_tokens"]["value"] == 12 diff --git a/tests/unit/test_agentscope_workspace_stream.py b/tests/unit/test_agentscope_workspace_stream.py new file mode 100644 index 00000000..31368095 --- /dev/null +++ b/tests/unit/test_agentscope_workspace_stream.py @@ -0,0 +1,114 @@ +"""Real workspace persistence through AgentScope's shared streaming/collection path.""" + +# pylint: disable=protected-access + +import asyncio +from types import SimpleNamespace +from uuid import uuid4 + +from agentscope.event import ModelCallEndEvent, ReplyEndEvent +from agentscope.message import AssistantMsg, Usage +from agentscope.permission import PermissionMode +import pytest + +from reme.components.agent_wrapper.as_agent_wrapper import AsAgentWrapper +from reme.components.application_context import ApplicationContext +from reme.utils import global_counter_get_all + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api", "outcome"), + [(api, outcome) for api in ("reply", "stream") for outcome in ("complete", "error", "cancel")] + + [("stream", "close")], +) +async def test_reply_persists_only_bound_workspace_on_all_exits(tmp_path, monkeypatch, api, outcome): + """Public return modes save resumable beta state on completion, error, cancellation and close.""" + config = {"workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}} + context = ApplicationContext(**config) + wrapper = AsAgentWrapper(name="research", app_context=context, as_llm="") + session_id = str(uuid4()) + entered, closed = asyncio.Event(), asyncio.Event() + saved = [] + dump_state = wrapper._dump_state + + async def persist(state, kwargs): + saved.append(kwargs["workspace"]) + await dump_state(state, kwargs) + + async def build_agent(inputs, **kwargs): + state = await wrapper._load_state(kwargs, PermissionMode.DEFAULT) + state.summary = f"beta-{outcome}" + message = AssistantMsg( + name="agent", + content="beta answer", + finished_reason="completed", + usage=Usage(input_tokens=3, output_tokens=2), + ) + state.context.append(message) + + async def native_stream(_inputs, *, yield_final_msg): + assert yield_final_msg + try: + yield ModelCallEndEvent(reply_id=message.id, input_tokens=3, output_tokens=2) + entered.set() + if outcome == "error": + raise RuntimeError("native stream failed") + if outcome in {"cancel", "close"}: + await asyncio.Event().wait() + yield ReplyEndEvent(session_id=session_id, reply_id=message.id, finished_reason="completed") + yield message + finally: + closed.set() + + return SimpleNamespace(state=state, reply_stream=native_stream), inputs + + monkeypatch.setattr(wrapper, "_build_agent", build_agent) + monkeypatch.setattr(wrapper, "_dump_state", persist) + + async def consume(): + kwargs = {"session_id": session_id, "workspace": "beta"} + if api == "reply": + return await wrapper.reply("hello", **kwargs) + chunks = [chunk async for chunk in wrapper.reply_stream("hello", **kwargs)] + return chunks[-1].chunk + + task = None + try: + if outcome == "close": + stream = wrapper.reply_stream("hello", session_id=session_id, workspace="beta") + await anext(stream) + await stream.aclose() + elif outcome == "cancel": + task = asyncio.create_task(consume()) + await asyncio.wait_for(entered.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + elif outcome == "error": + with pytest.raises(RuntimeError, match="native stream failed"): + await consume() + else: + result = await consume() + assert result["result"] == "beta answer" + assert result["status"] == "completed" + assert result["usage"]["total_tokens"] == 5 + + assert closed.is_set() + assert saved == ["beta"] + beta_file = tmp_path / "beta" / "mem_session" / "agentscope" / f"{session_id}.jsonl" + assert beta_file.is_file() + assert not (tmp_path / "alpha" / "mem_session" / "agentscope" / f"{session_id}.jsonl").exists() + # Recreate the application context to verify the durable result, not a live object cache. + restarted = AsAgentWrapper(app_context=ApplicationContext(**config), as_llm="") + state = await restarted._load_state({"resume": session_id, "workspace": "beta"}, PermissionMode.DEFAULT) + assert state.summary == f"beta-{outcome}" + with pytest.raises(FileNotFoundError): + await restarted._load_state({"resume": session_id, "workspace": "alpha"}, PermissionMode.DEFAULT) + counters = global_counter_get_all(context.metadata, ["__token_counter", "research"]) + assert counters["children"]["total_tokens"]["value"] == 5 + finally: + if task is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await wrapper.close() 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 e806ff9b..fa663b43 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 [] @@ -237,7 +237,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 e65071d2..b42009c1 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_compaction_ownership.py b/tests/unit/test_codex_compaction_ownership.py new file mode 100644 index 00000000..c80f4f63 --- /dev/null +++ b/tests/unit/test_codex_compaction_ownership.py @@ -0,0 +1,372 @@ +"""Workspace ownership survives cancellation across the shared SDK process boundary.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +import json +import threading + +import pytest +from openai_codex import AsyncCodex +from openai_codex.errors import InvalidParamsError, TransportClosedError + +from reme.components import ApplicationContext +from reme.components.agent_wrapper.codex_agent_wrapper import CodexAgentWrapper +from reme.components.codex_app_server import CodexAppServer +from reme.enumeration import ComponentEnum + +from .codex_workspace_test_support import workspace_server + + +def _config(tmp_path): + return {"workspace_dir": {"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}} + + +def _owned_on_disk(tmp_path, alias): + path = tmp_path / alias / "mem_session" / "codex" / "owned_threads.json" + return set(json.loads(path.read_text(encoding="utf-8"))["thread_ids"]) if path.is_file() else set() + + +async def _assert_isolated(tmp_path, context): + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {"T"} + assert _owned_on_disk(tmp_path, "alpha") == {"T"} + for checked_context in (context, ApplicationContext(**_config(tmp_path))): + # Rejection must happen before a new server initializes any SDK client. + server = CodexAppServer(app_context=checked_context) + checked_context.shared_components.setdefault(ComponentEnum.CODEX_APP_SERVER, {})["default"] = server + beta = CodexAgentWrapper(name="beta", app_context=checked_context) + try: + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await beta.reply("hello", resume="T", workspace="beta") + assert server._codex is None + finally: + await beta.close() + await server.close() + assert not _owned_on_disk(tmp_path, "beta") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.parametrize("operation", ["reply", "compact"]) +async def test_failed_compactor_keeps_another_wrappers_success(tmp_path, monkeypatch, cancel, operation): + context = ApplicationContext(**_config(tmp_path)) + server = workspace_server(context, monkeypatch) + compactor = CodexAgentWrapper(name="compactor", app_context=context) + worker = CodexAgentWrapper(name="worker", app_context=context) + original_get_client = server._get_client + entered, finish = asyncio.Event(), asyncio.Event() + calls = 0 + + async def get_client(): + nonlocal calls + calls += 1 + if calls == 1: + entered.set() + await finish.wait() + raise RuntimeError("initialization failed") + return await original_get_client() + + monkeypatch.setattr(server, "_get_client", get_client) + task = asyncio.create_task(compactor.compact_session("T", workspace="alpha")) + try: + await asyncio.wait_for(entered.wait(), 3) + if operation == "reply": + assert (await worker.reply("hello", resume="T", workspace="alpha"))["result"] == "hello" + else: + await worker.compact_session("T", workspace="alpha") + if cancel: + task.cancel() + else: + finish.set() + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + await asyncio.wait_for(task, 3) + await _assert_isolated(tmp_path, context) + finally: + finish.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await compactor.close() + await worker.close() + await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_last_predispatch_failure_releases_unused_claim(tmp_path, monkeypatch, cancel): + context = ApplicationContext(**_config(tmp_path)) + server = workspace_server(context, monkeypatch) + alpha = CodexAgentWrapper(app_context=context) + entered = asyncio.Event() + + async def get_client(): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(server, "_get_client", get_client) + task = asyncio.create_task(alpha.compact_session("T", workspace="alpha")) + try: + await asyncio.wait_for(entered.wait(), 3) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not _owned_on_disk(tmp_path, "alpha") + assert not context.workspaces["alpha"].codex_thread_state.owned_thread_ids + assert not context.workspaces["alpha"].codex_thread_state.provisional_claims + if not cancel: + + async def reject_client(): + raise RuntimeError("initialization failed") + + monkeypatch.setattr(server, "_get_client", reject_client) + with pytest.raises(RuntimeError, match="initialization failed"): + await alpha.compact_session("T", workspace="alpha") + assert not _owned_on_disk(tmp_path, "alpha") + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await alpha.close() + await server.close() + + +def _resume_response(thread_id, cwd): + """Valid wire response parsed by the real synchronous SDK client.""" + return { + "approvalPolicy": "never", + "approvalsReviewer": "user", + "cwd": str(cwd), + "model": "test", + "modelProvider": "test", + "sandbox": {"type": "dangerFullAccess"}, + "thread": { + "cliVersion": "0.163.0", + "createdAt": 0, + "cwd": str(cwd), + "ephemeral": False, + "id": thread_id, + "modelProvider": "test", + "preview": "", + "sessionId": thread_id, + "source": "appServer", + "status": {"type": "idle"}, + "turns": [], + "updatedAt": 0, + }, + } + + +def _native_server(context, sdk, monkeypatch): + """Use the real SDK clients with controlled RPC and notification transports.""" + from types import SimpleNamespace + + server = CodexAppServer(name="default", app_context=context) + context.shared_components.setdefault(ComponentEnum.CODEX_APP_SERVER, {})["default"] = server + sdk._initialized = True + server._codex = sdk + notifications = asyncio.Queue() + monkeypatch.setattr(sdk._client, "next_notification", notifications.get) + server._notifications = asyncio.create_task(server._watch_notifications(sdk)) + loop = asyncio.get_running_loop() + + def notify(status): + event = SimpleNamespace( + method="thread/status/changed", + payload=SimpleNamespace(thread_id="T", status=SimpleNamespace(root=SimpleNamespace(type=status))), + ) + loop.call_soon_threadsafe(notifications.put_nowait, event) + + return server, notify + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("operation", "blocked_method"), + [("compact", "thread/compact/start"), ("compact", "thread/resume"), ("reply", "thread/resume")], +) +@pytest.mark.parametrize("outcome", ["success", "rejected", "transport_closed"]) +async def test_cancelled_sdk_worker_preserves_persisted_ownership( + tmp_path, + monkeypatch, + operation, + blocked_method, + outcome, +): + """The real SDK's thread-offloaded RPC stays owned while its waiter is cancelled.""" + context = ApplicationContext(**_config(tmp_path)) + sdk = AsyncCodex() + server, notify = _native_server(context, sdk, monkeypatch) + alpha = CodexAgentWrapper(name="alpha", app_context=context) + entered, release, finished = threading.Event(), threading.Event(), threading.Event() + ops = [] + + def request(method, params=None): + ops.append(method) + if method == blocked_method: + if method == "thread/compact/start": + notify("active") + entered.set() + try: + if not release.wait(10): + raise TimeoutError("transport not released") + if outcome == "rejected": + raise InvalidParamsError(-32602, "rejected") + if outcome == "transport_closed": + raise TransportClosedError("outcome unknown") + if method == "thread/compact/start": + notify("idle") + finally: + finished.set() + if method == "thread/resume": + return _resume_response(params["threadId"], tmp_path / "alpha") + if method == "thread/read": + return {"thread": _resume_response("T", tmp_path / "alpha")["thread"]} + if method == "thread/compact/start": + return {} + if method == "thread/unsubscribe": + return {"status": "unsubscribed"} + raise AssertionError(f"Unexpected RPC: {method}") + + monkeypatch.setattr(sdk._client._sync, "_request_raw", request) + call = ( + alpha.compact_session("T", workspace="alpha") + if operation == "compact" + else alpha.reply( + "hello", + resume="T", + workspace="alpha", + ) + ) + task = asyncio.create_task(call) + try: + assert await asyncio.to_thread(entered.wait, 3) + task.cancel() + await asyncio.sleep(0) # Deliver cancellation while the transport remains blocked. + assert not finished.is_set() + await _assert_isolated(tmp_path, context) + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 3) + assert await asyncio.to_thread(finished.wait, 3) + await _assert_isolated(tmp_path, context) + if blocked_method == "thread/compact/start": + assert "thread/unsubscribe" in ops # Default cleanup cannot revoke ownership. + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await asyncio.to_thread(finished.wait, 3) + await alpha.close() + await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["thread/resume", "thread/compact/start"]) +@pytest.mark.parametrize("error_type", [TransportClosedError, TimeoutError, RuntimeError]) +async def test_unknown_sdk_failure_retains_ownership(tmp_path, monkeypatch, method, error_type): + context = ApplicationContext(**_config(tmp_path)) + sdk = AsyncCodex() + server, _ = _native_server(context, sdk, monkeypatch) + alpha = CodexAgentWrapper(app_context=context) + + def request(rpc_method, params=None): + if rpc_method == method: + raise error_type("outcome unknown") + if rpc_method == "thread/resume": + return _resume_response(params["threadId"], tmp_path / "alpha") + if rpc_method == "thread/read": + return {"thread": _resume_response("T", tmp_path / "alpha")["thread"]} + if rpc_method == "thread/unsubscribe": + return {"status": "unsubscribed"} + raise AssertionError(f"Unexpected RPC: {rpc_method}") + + monkeypatch.setattr(sdk._client._sync, "_request_raw", request) + try: + with pytest.raises(error_type, match="outcome unknown"): + await alpha.compact_session("T", workspace="alpha") + await _assert_isolated(tmp_path, context) + finally: + await alpha.close() + await server.close() + + +@pytest.mark.asyncio +async def test_explicit_resume_rejection_releases_unused_claim(tmp_path, monkeypatch): + context = ApplicationContext(**_config(tmp_path)) + sdk = AsyncCodex() + server, _ = _native_server(context, sdk, monkeypatch) + alpha = CodexAgentWrapper(app_context=context) + + def reject(*_args, **_kwargs): + raise InvalidParamsError(-32602, "invalid thread") + + monkeypatch.setattr(sdk._client._sync, "_request_raw", reject) + try: + with pytest.raises(InvalidParamsError): + await alpha.compact_session("T", workspace="alpha") + assert not _owned_on_disk(tmp_path, "alpha") + assert not context.workspaces["alpha"].codex_thread_state.owned_thread_ids + finally: + await alpha.close() + await server.close() + + +@pytest.mark.asyncio +async def test_failed_first_reservation_preserves_pending_second_operation(tmp_path, monkeypatch): + context = ApplicationContext(**_config(tmp_path)) + server = workspace_server(context, monkeypatch) + wrappers = [CodexAgentWrapper(name=name, app_context=context) for name in ("first", "second")] + entered = [asyncio.Event(), asyncio.Event()] + finish = [asyncio.Event(), asyncio.Event()] + calls = 0 + + async def get_client(): + nonlocal calls + index = calls + calls += 1 + entered[index].set() + await finish[index].wait() + raise RuntimeError("initialization failed") + + monkeypatch.setattr(server, "_get_client", get_client) + tasks = [] + try: + for index, wrapper in enumerate(wrappers): + tasks.append(asyncio.create_task(wrapper.compact_session("T", workspace="alpha"))) + await asyncio.wait_for(entered[index].wait(), 3) + tasks[0].cancel() + with pytest.raises(asyncio.CancelledError): + await tasks[0] + await _assert_isolated(tmp_path, context) + finish[1].set() + with pytest.raises(RuntimeError): + await tasks[1] + assert not _owned_on_disk(tmp_path, "alpha") + assert not context.workspaces["alpha"].codex_thread_state.provisional_claims + finally: + for event in finish: + event.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + for wrapper in wrappers: + await wrapper.close() + await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["reply", "compact"]) +async def test_unknown_workspace_fails_before_sdk_initialization(tmp_path, monkeypatch, operation): + context = ApplicationContext(**_config(tmp_path)) + server = workspace_server(context, monkeypatch) + wrapper = CodexAgentWrapper(app_context=context) + try: + with pytest.raises(ValueError, match="Unknown workspace 'typo'"): + if operation == "reply": + await wrapper.reply("invalid", resume="T", workspace="typo") + else: + await wrapper.compact_session("T", workspace="typo") + assert server._codex is None + assert not _owned_on_disk(tmp_path, "alpha") + assert not _owned_on_disk(tmp_path, "beta") + finally: + await wrapper.close() + await server.close() diff --git a/tests/unit/test_health_check_workspace.py b/tests/unit/test_health_check_workspace.py new file mode 100644 index 00000000..bae65e25 --- /dev/null +++ b/tests/unit/test_health_check_workspace.py @@ -0,0 +1,103 @@ +"""Health and memory reports must describe the request's actual workspace.""" + +import pytest + +from reme.application import Application +from reme.components.application_context import ApplicationContext +from reme.components.base_component import BaseComponent +from reme.enumeration import ComponentEnum +from reme.schema import FileNode +from reme.steps.common.health_check import HealthCheckStep +from reme.steps.common.status import _collect_memory + + +@pytest.mark.asyncio +async def test_health_job_reports_beta_failure_without_affecting_alpha(tmp_path): + """Exercise real graphs and the public job before and after only beta stops.""" + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "local"}}}, + jobs={"health_check": {"backend": "base", "steps": [{"backend": "health_check_step"}]}}, + ) + async with app: + alpha = app.context.workspaces["alpha"].get_component(ComponentEnum.FILE_GRAPH, "default") + beta = app.context.workspaces["beta"].get_component(ComponentEnum.FILE_GRAPH, "default") + await beta.upsert_nodes([FileNode(path="beta.md", st_mtime=1)]) + alpha_before = (await app.run_job("health_check", workspace="alpha")).metadata["health"] + response = await app.run_job("health_check", workspace="beta") + before = response.metadata["health"] + assert before["components"]["file_graph"]["default"]["n_nodes"] == 1 + assert before["workspace"] == "beta" and before["healthy"] is True + assert "beta" in response.answer + assert before["components"]["file_graph"]["default"]["is_started"] is True + await beta.close() + response = await app.run_job("health_check", workspace="beta") + assert response.success is True # The diagnostic job succeeded; its health result is negative. + after = response.metadata["health"] + assert after["workspace"] == "beta" and after["healthy"] is False + assert after["components"]["file_graph"]["default"]["is_started"] is False + assert after["components"]["file_graph"]["default"]["n_nodes"] == 1 + assert alpha.is_started is True and beta.is_started is False + assert (await app.run_job("health_check", workspace="alpha")).metadata["health"] == alpha_before + assert (await app.run_job("health_check")).metadata["health"] == alpha_before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("named", [False, True]) +async def test_health_single_workspace_labels_default_and_preserves_payload(tmp_path, named): + """String and one-entry mapping reports identify their default workspace.""" + context = ApplicationContext(workspace_dir={"named": str(tmp_path)} if named else str(tmp_path)) + alias = context.workspace.alias + step = HealthCheckStep(app_context=context) + response = await step() + assert response.metadata["health"]["workspace"] == alias + assert response.metadata["health"]["healthy"] is True + assert set(response.metadata["health"]["components"]) == { + "embedding_store", + "file_graph", + "file_store", + "keyword_index", + } + + +@pytest.mark.asyncio +async def test_reports_include_shared_and_selected_components_without_mutating_registries(tmp_path): + """Shared and scoped names of the same type both remain visible; alpha never leaks.""" + context = ApplicationContext(workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}) + shared, alpha, beta = (BaseComponent() for _ in range(3)) + context.shared_components[ComponentEnum.FILE_GRAPH] = {"shared": shared} + context.workspaces["alpha"].components[ComponentEnum.FILE_GRAPH] = {"alpha": alpha} + context.workspaces["beta"].components[ComponentEnum.FILE_GRAPH] = {"beta": beta} + await shared.start() + await alpha.start() + await beta.start() + try: + from reme.components.runtime_context import RuntimeContext + + step = HealthCheckStep(app_context=context) + response = await step(RuntimeContext(workspace=context.workspaces["beta"])) + assert set(response.metadata["health"]["components"]["file_graph"]) == {"shared", "beta"} + memory = _collect_memory(context, context.workspaces["beta"]) + assert set(memory["components"]["file_graph"]) == {"shared", "beta"} + await shared.close() + response = await step(RuntimeContext(workspace=context.workspaces["beta"])) + assert response.metadata["health"]["healthy"] is False + assert set(context.shared_components[ComponentEnum.FILE_GRAPH]) == {"shared"} + assert set(context.workspaces["alpha"].components[ComponentEnum.FILE_GRAPH]) == {"alpha"} + assert set(context.workspaces["beta"].components[ComponentEnum.FILE_GRAPH]) == {"beta"} + finally: + await shared.close() + await alpha.close() + await beta.close() + + +def test_component_report_view_rejects_foreign_workspace(tmp_path): + """A workspace object from another application cannot select this application's resources.""" + context = ApplicationContext(workspace_dir=str(tmp_path / "current")) + foreign = ApplicationContext(workspace_dir=str(tmp_path / "foreign")) + with pytest.raises(ValueError, match="not bound to this Application"): + context.components_for(foreign.workspace) 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_status_step.py b/tests/unit/test_status_step.py index 1951259e..2ca44899 100644 --- a/tests/unit/test_status_step.py +++ b/tests/unit/test_status_step.py @@ -2,6 +2,8 @@ import asyncio +from reme.application import Application +from reme.schema import FileNode from reme.components.application_context import ApplicationContext from reme.components.base_component import BaseComponent, Dependency from reme.enumeration import ComponentEnum @@ -108,3 +110,45 @@ def test_status_step_returns_human_summary_and_exact_metadata(tmp_path): assert "file_store:default" in response.answer assert "Storage" not in response.answer assert set(response.metadata["status"]) == {"memory"} + + +def test_status_job_attributes_growth_to_selected_workspace(tmp_path): + """The real job includes beta's graph growth without charging it to alpha.""" + + async def check(): + app = Application( + workspace_dir={"alpha": str(tmp_path / "alpha"), "beta": str(tmp_path / "beta")}, + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "cli"}, + components={"file_graph": {"default": {"backend": "local"}}}, + jobs={"status": {"backend": "base", "steps": [{"backend": "status_step"}]}}, + ) + await app.start() + try: + + async def memory(alias): + response = await app.run_job("status", workspace=alias) + assert response.success + report = response.metadata["status"]["memory"] + assert report["workspace"] == alias + assert f"Workspace {alias}" in response.answer + return report + + alpha_before = await memory("alpha") + beta_before = await memory("beta") + graph = app.context.workspaces["beta"].get_component(ComponentEnum.FILE_GRAPH, "default") + await graph.upsert_nodes( + [FileNode(path="beta.md", st_mtime=1, front_matter={"description": "x" * 2_000_000})], + ) + beta_after = await memory("beta") + alpha_after = await memory("alpha") + assert beta_after["components"]["file_graph"]["default"]["bytes"] == _component_size(graph) + assert beta_after["components_total_bytes"] > beta_before["components_total_bytes"] + 1_900_000 + assert alpha_before["components"] == alpha_after["components"] + assert alpha_before["components_total_bytes"] == alpha_after["components_total_bytes"] + finally: + await app.close() + + asyncio.run(check()) diff --git a/tests/unit/test_token_usage.py b/tests/unit/test_token_usage.py index 3b6d8ca9..e4e9c206 100644 --- a/tests/unit/test_token_usage.py +++ b/tests/unit/test_token_usage.py @@ -151,7 +151,7 @@ async def test_agentscope_records_usage_in_both_return_modes(tmp_path, monkeypat async def build_agent(inputs, **_kwargs): return agent, inputs - async def dump_state(_state): + async def dump_state(_state, _kwargs=None): return None monkeypatch.setattr(wrapper, "_build_agent", build_agent) 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..43411586 --- /dev/null +++ b/tests/unit/test_workspace_phase3.py @@ -0,0 +1,570 @@ +"""Phase 3: nested propagation, session namespacing, and background binding.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio +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 + +from .codex_workspace_test_support import workspace_server + + +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 + + +# ----- 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={ + "codex_app_server": {"default": {"backend": "sdk", "auth_mode": "oauth"}}, + "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 ---------------------------------------------------------- + + +@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") + + +# ----- 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")) + + +# ----- Shared Codex server with real workspace jobs --------------------------- + + +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": "cli"}, + components={ + "codex_app_server": {"default": {"backend": "sdk", "auth_mode": "oauth"}}, + "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"}], + }, + }, + ) + + +@pytest.mark.asyncio +async def test_codex_shared_server_preserves_workspace_guards_and_cwd(tmp_path, monkeypatch): + context = _context(tmp_path) + server = workspace_server(context, monkeypatch) + first = CodexAgentWrapper(name="first", app_context=context) + second = CodexAgentWrapper(name="second", app_context=context) + try: + alpha = await first.reply("alpha", tool_context_id="ctx-alpha") + beta = await second.reply("beta", workspace="beta", tool_context_id="ctx-beta") + sdk = server._codex + assert alpha["session_id"] != beta["session_id"] + assert sdk.opened[0][2]["cwd"] == str(first.project_path) + assert sdk.opened[1][2]["cwd"] == str(tmp_path / "beta") + assert sdk.threads[beta["session_id"]].options[0]["cwd"] == str(tmp_path / "beta") + with pytest.raises(ValueError, match="cannot change"): + await second.reply("changed", resume=alpha["session_id"], tool_context_id="ctx-beta") + for operation in ("resume", "fork", "compact"): + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + if operation == "compact": + await second.compact_session(alpha["session_id"], workspace="beta") + else: + await second.reply( + "wrong workspace", + resume=alpha["session_id"], + workspace="beta", + fork_session=operation == "fork", + ) + assert context.workspaces["alpha"].codex_thread_state.owned_thread_ids == {alpha["session_id"]} + assert context.workspaces["beta"].codex_thread_state.owned_thread_ids == {beta["session_id"]} + finally: + await first.close() + await second.close() + await server.close() + + +@pytest.mark.asyncio +async def test_codex_ownership_survives_shared_server_and_context_restart(tmp_path, monkeypatch): + context = _context(tmp_path) + server = workspace_server(context, monkeypatch) + wrapper = CodexAgentWrapper(app_context=context) + try: + await wrapper.reply("external", resume="external", workspace="alpha") + await wrapper.compact_session("compact", workspace="alpha") + finally: + await wrapper.close() + await server.close() + restarted = _context(tmp_path) + server = workspace_server(restarted, monkeypatch) + wrapper = CodexAgentWrapper(app_context=restarted) + try: + for thread_id in ("external", "compact"): + with pytest.raises(ValueError, match="belongs to workspace 'alpha'"): + await wrapper.reply("wrong workspace", resume=thread_id, workspace="beta") + assert server._codex is None + assert (await wrapper.reply("correct", resume="external", workspace="alpha"))["result"] == "correct" + finally: + await wrapper.close() + await server.close() + + +@pytest.mark.asyncio +async def test_codex_concurrent_first_resume_has_one_persisted_owner(tmp_path, monkeypatch): + context = _context(tmp_path) + server = workspace_server(context, monkeypatch) + wrappers = [CodexAgentWrapper(name=name, app_context=context) for name in ("first", "second")] + try: + results = await asyncio.gather( + wrappers[0].reply("alpha", resume="T", workspace="alpha"), + wrappers[1].reply("beta", resume="T", workspace="beta"), + return_exceptions=True, + ) + assert sum(isinstance(result, dict) for result in results) == 1 + assert sum(isinstance(result, ValueError) for result in results) == 1 + owners = [ + alias for alias, ws in context.workspaces.items() if "T" in (ws.codex_thread_state.owned_thread_ids or ()) + ] + assert len(owners) == 1 + assert len(server._codex.opened) == 1 + other = "beta" if owners[0] == "alpha" else "alpha" + restarted = _context(tmp_path) + new_server = workspace_server(restarted, monkeypatch) + new_wrapper = CodexAgentWrapper(app_context=restarted) + try: + with pytest.raises(ValueError, match="belongs to workspace"): + await new_wrapper.reply("wrong", resume="T", workspace=other) + assert new_server._codex is None + finally: + await new_wrapper.close() + await new_server.close() + finally: + for wrapper in wrappers: + await wrapper.close() + await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["multi", "single_string", "single_mapping"]) +@pytest.mark.parametrize("explicit", [False, True]) +async def test_codex_scoped_mcp_calls_original_workspace_jobs(tmp_path, monkeypatch, mode, explicit): + """Managed HTTP tools preserve binding, parent state and failures for CLI hosts too.""" + from fastmcp import Client + from fastmcp.client import StreamableHttpTransport + + from .codex_workspace_test_support import WorkspaceCodex + + monkeypatch.setattr("openai_codex.AsyncCodex", WorkspaceCodex) + workspace_dir = ( + None + if mode == "multi" + else (str(tmp_path / "single") if mode == "single_string" else {"named": str(tmp_path / "single")}) + ) + app = _codex_test_app(tmp_path, workspace_dir=workspace_dir) + async with app: + server = app.context.shared_components[ComponentEnum.CODEX_APP_SERVER]["default"] + wrapper = CodexAgentWrapper(app_context=app.context) + alias = "beta" if mode == "multi" and explicit else next(iter(app.context.workspaces)) + root = app.context.workspaces[alias].workspace_path + await app.run_job("write_note", path="daily/resident.md", content="parent-state", workspace=alias) + sdk = await server._get_client() + seen = [] + + async def tools_from_turn(_thread, _inputs): + config = sdk.opened[-1][2]["config"]["mcp_servers"]["reme_jobs"] + transport = StreamableHttpTransport(url=config["url"], headers=config["http_headers"]) + async with Client(transport) as client: + assert {tool.name for tool in await client.list_tools()} == {"read_note", "write_note"} + read = await client.call_tool("read_note", {"path": "daily/resident.md"}) + assert "parent-state" in read.content[0].text + with pytest.raises(Exception): + await client.call_tool("read_note", {"path": "daily/resident.md", "workspace": "other"}) + await client.call_tool("write_note", {"path": "daily/tool.md", "content": "bound-write"}) + failure = await client.call_tool( + "write_note", + {"path": "../escape.md", "content": "must-not-write"}, + raise_on_error=False, + ) + assert failure.is_error + seen.append(read) + + sdk.run_hook = tools_from_turn + try: + kwargs = {"workspace": alias} if explicit else {} + await wrapper.reply("use tools", job_tools=["read_note", "write_note"], **kwargs) + assert seen + assert (root / "daily/tool.md").read_text(encoding="utf-8") == "bound-write\n" + assert not (root.parent / "escape.md").exists() + for other_alias, workspace in app.context.workspaces.items(): + if other_alias != alias: + assert not (workspace.workspace_path / "daily/tool.md").exists() + assert not await workspace.get_component(ComponentEnum.FILE_STORE, "default").get_nodes() + finally: + await wrapper.close() + + +@pytest.mark.asyncio +async def test_codex_session_history_and_appends_follow_workspace_binding(tmp_path, monkeypatch): + """New session endpoints cannot read or inject another workspace's thread.""" + from unittest.mock import AsyncMock + + from reme.components.codex_app_server import CodexAppServer + from reme.steps.common.session import SessionAppendStep, SessionReadStep, SessionStatusStep + + from .test_agent_session_api import NativeSession + + context = _context(tmp_path) + server = CodexAppServer(app_context=context) + context.shared_components[ComponentEnum.CODEX_APP_SERVER] = {"default": server} + native = NativeSession(server) + monkeypatch.setattr(server, "_get_client", AsyncMock(return_value=native.client)) + wrapper = CodexAgentWrapper(app_context=context) + configured_beta = CodexAgentWrapper(name="beta", app_context=context, workspace="beta") + session_id = native.thread.id + try: + for step_type in (SessionAppendStep, SessionReadStep, SessionStatusStep): + step = step_type(app_context=context, agent_wrapper=wrapper) + runtime = RuntimeContext( + workspace=context.workspaces["beta"], + session_id=session_id, + content="beta proactive message", + message_id="beta-message", + ) + await step(runtime) + assert runtime.response.success + beta_path = server._message_path(session_id, "beta") + assert beta_path.is_relative_to(tmp_path / "beta") and beta_path.is_file() + assert not server._message_path(session_id, "alpha").exists() + assert (await configured_beta.read_session(session_id)).session_id == session_id + fork_path = server.session_path / "sessions" / "fork.jsonl" + fork_path.write_text("") + fork_thread = SimpleNamespace( + id="fork-beta", + path=str(fork_path), + turns=native.thread.turns, + status=native.thread.status, + ) + native.threads[fork_thread.id] = fork_thread + fork_handle = SimpleNamespace( + id=fork_thread.id, + read=AsyncMock(return_value=SimpleNamespace(thread=fork_thread)), + ) + native.client.thread_fork = AsyncMock(return_value=fork_handle) + async with server.thread(session_id, fork=True, workspace="beta") as forked: + assert forked.id == "fork-beta" + fork_snapshot = await server.read_session("fork-beta", workspace="beta") + assert any(message.content == "beta proactive message" for message in fork_snapshot.messages) + assert server._message_path("fork-beta", "beta").is_file() + assert not server._message_path("fork-beta", "alpha").exists() + reads = native.client._client.thread_read.await_count + for operation in ("read_session", "get_session_status", "append_session"): + kwargs = {"content": "wrong", "message_id": "wrong"} if operation == "append_session" else {} + with pytest.raises(ValueError, match="belongs to workspace 'beta'"): + await getattr(wrapper, operation)(session_id, **kwargs) + assert native.client._client.thread_read.await_count == reads + assert len(native.injected) == 1 + finally: + await wrapper.close() + await configured_beta.close() + await server.close()