From f34dcdb09b381658ce7ee30137cd32a118475e46 Mon Sep 17 00:00:00 2001 From: xyf2020 <75460675+xyf2020@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:20:21 +0800 Subject: [PATCH] feat(Step tools): add white/black path prefix permission filtering to read, edit, write (#391) * feat(read): add white/black path prefix permission filtering to ReadStep * feat: add PrefixCheck mixin for path-prefix permission in file I/O steps * feat: add injected_job_kwargs mechanism and refine path-prefix permission * refactor(file_io): consolidate prefix_check into _path module --- .../agent_wrapper/as_agent_wrapper.py | 22 +- .../agent_wrapper/base_agent_wrapper.py | 40 +++ .../agent_wrapper/cc_agent_wrapper.py | 27 +- .../agent_wrapper/codex_agent_wrapper.py | 13 +- .../agent_wrapper/codex_mcp_server.py | 23 +- reme/steps/benchmark/beam/auto_memory.py | 9 +- reme/steps/benchmark/lme/auto_memory.py | 9 +- reme/steps/evolve/auto_memory.py | 8 +- reme/steps/file_io/_path.py | 51 +++- reme/steps/file_io/edit.py | 9 +- reme/steps/file_io/frontmatter_update.py | 12 +- reme/steps/file_io/read.py | 12 +- reme/steps/file_io/write.py | 9 +- tests/unit/test_background_steps.py | 1 + tests/unit/test_codex_agent_wrapper.py | 32 +++ tests/unit/test_injected_job_kwargs.py | 241 ++++++++++++++++ tests/unit/test_path_permission.py | 263 ++++++++++++++++++ 17 files changed, 741 insertions(+), 40 deletions(-) create mode 100644 tests/unit/test_injected_job_kwargs.py create mode 100644 tests/unit/test_path_permission.py diff --git a/reme/components/agent_wrapper/as_agent_wrapper.py b/reme/components/agent_wrapper/as_agent_wrapper.py index 19927171..f1f8e9ec 100644 --- a/reme/components/agent_wrapper/as_agent_wrapper.py +++ b/reme/components/agent_wrapper/as_agent_wrapper.py @@ -156,18 +156,24 @@ class AsAgentWrapper(BaseAgentWrapper): self._session_cleanup_done = False @classmethod - def _make_tool(cls, job: "BaseJob", tool_context_id: str | None = None) -> FunctionTool: + def _make_tool( + cls, + job: "BaseJob", + tool_context_id: str | None = None, + injected_job_kwargs: dict[str, Any] | None = None, + ) -> FunctionTool: + injected = cls._resolve_injected_job_kwargs( + {"tool_context_id": tool_context_id, "injected_job_kwargs": injected_job_kwargs}, + ) + async def run_job(**kwargs) -> ToolChunk: - if tool_context_id: - assert "tool_context_id" not in kwargs, "tool_context_id is injected by agent_wrapper" - kwargs["tool_context_id"] = tool_context_id - response = await job(**kwargs) + response = await job(**cls._merge_injected_job_kwargs(kwargs, injected)) state = ToolResultState.SUCCESS if response.success else ToolResultState.ERROR return ToolChunk(content=[TextBlock(text=str(response.answer))], state=state) tool = FunctionTool(func=run_job, name=job.name, description=job.description, is_concurrency_safe=False) - if job.parameters: - tool.input_schema = job.parameters + if parameters := cls._strip_injected_parameters(job.parameters, injected): + tool.input_schema = parameters return tool def _builtin_tools( @@ -308,7 +314,7 @@ class AsAgentWrapper(BaseAgentWrapper): builtin_tools = [] 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) for job in resolved_jobs) + tools.extend(self._make_tool(job, tool_context_id, kwargs.get("injected_job_kwargs")) for job in resolved_jobs) toolkit = kwargs.get("toolkit") or Toolkit( tools=tools, skills_or_loaders=skills, diff --git a/reme/components/agent_wrapper/base_agent_wrapper.py b/reme/components/agent_wrapper/base_agent_wrapper.py index b571b15f..81960268 100644 --- a/reme/components/agent_wrapper/base_agent_wrapper.py +++ b/reme/components/agent_wrapper/base_agent_wrapper.py @@ -155,6 +155,46 @@ class BaseAgentWrapper(BaseComponent): return schema raise TypeError("output_schema must be a JSON schema dict or BaseModel class") + @staticmethod + def _resolve_injected_job_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + """Collect server-owned kwargs merged into every job tool call. + + ``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. + """ + injected = dict(kwargs.get("injected_job_kwargs") or {}) + if tool_context_id := kwargs.get("tool_context_id"): + injected["tool_context_id"] = tool_context_id + return injected + + @staticmethod + def _merge_injected_job_kwargs(model_kwargs: dict[str, Any], injected: dict[str, Any]) -> dict[str, Any]: + """Merge server-owned kwargs over model tool arguments, rejecting conflicts. + + Silently letting model values win would make injected constraints + bypassable, so any overlap is an explicit error. + """ + if conflicts := sorted(injected.keys() & model_kwargs.keys()): + names = ", ".join(conflicts) + raise ValueError(f"injected tool arguments cannot be provided by the model: {names}") + return {**model_kwargs, **injected} + + @staticmethod + def _strip_injected_parameters(parameters: dict | None, injected: dict[str, Any]) -> dict | None: + """Hide injected keys from the tool parameter schema exposed to the model.""" + if not parameters or not injected: + return parameters + parameters = dict(parameters) + if "properties" in parameters: + parameters["properties"] = { + name: schema for name, schema in parameters["properties"].items() if name not in injected + } + if "required" in parameters: + parameters["required"] = [name for name in parameters["required"] if name not in injected] + return parameters + def _resolve_job_tools(self, job_tools: list[str]) -> list["BaseJob"]: """Resolve job name strings to BaseJob instances via app_context.""" if not job_tools: diff --git a/reme/components/agent_wrapper/cc_agent_wrapper.py b/reme/components/agent_wrapper/cc_agent_wrapper.py index 469e208f..064b832c 100644 --- a/reme/components/agent_wrapper/cc_agent_wrapper.py +++ b/reme/components/agent_wrapper/cc_agent_wrapper.py @@ -73,16 +73,20 @@ class CcAgentWrapper(BaseAgentWrapper): self.logger.warning(f"Failed to link Claude Code skills into {target}: {exc}") @classmethod - def _make_tool(cls, job: "BaseJob", tool_context_id: str | None = None): + def _make_tool( + cls, + job: "BaseJob", + tool_context_id: str | None = None, + injected_job_kwargs: dict[str, Any] | 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}, + ) + async def run_job(args): - call_args = dict(args) - if tool_context_id: - if "tool_context_id" in call_args: - raise ValueError("tool_context_id is injected by agent_wrapper") - call_args["tool_context_id"] = tool_context_id - response = await job(**call_args) + response = await job(**cls._merge_injected_job_kwargs(dict(args), injected)) return { "content": [{"type": "text", "text": str(response.answer)}], "is_error": not response.success, @@ -91,7 +95,7 @@ class CcAgentWrapper(BaseAgentWrapper): return SdkMcpTool( name=job.name, description=job.description, - input_schema=job.parameters, + input_schema=cls._strip_injected_parameters(job.parameters, injected), handler=run_job, ) @@ -142,7 +146,7 @@ class CcAgentWrapper(BaseAgentWrapper): if "setting_sources" not in kwargs and kwargs.get("skills") is None: kwargs["setting_sources"] = [] - skip_keys = {"job_tools", "output_schema", "api_key", "base_url", "credential"} + skip_keys = {"job_tools", "injected_job_kwargs", "output_schema", "api_key", "base_url", "credential"} option_fields = {field.name for field in fields(ClaudeAgentOptions)} option_kwargs = {key: value for key, value in kwargs.items() if key not in skip_keys and key in option_fields} option_kwargs["disallowed_tools"] = list( @@ -193,7 +197,10 @@ class CcAgentWrapper(BaseAgentWrapper): opts.mcp_servers = dict(opts.mcp_servers) 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")) for job in resolved_jobs] + sdk_tools = [ + self._make_tool(job, kwargs.get("tool_context_id"), kwargs.get("injected_job_kwargs")) + for job in resolved_jobs + ] opts.mcp_servers[self.MCP_SERVER_NAME] = create_sdk_mcp_server( name=self.MCP_SERVER_NAME, tools=sdk_tools, diff --git a/reme/components/agent_wrapper/codex_agent_wrapper.py b/reme/components/agent_wrapper/codex_agent_wrapper.py index edef958e..bd97203b 100644 --- a/reme/components/agent_wrapper/codex_agent_wrapper.py +++ b/reme/components/agent_wrapper/codex_agent_wrapper.py @@ -255,6 +255,8 @@ class CodexAgentWrapper(BaseAgentWrapper): for name in job_names: args.extend(["--job", name]) args.extend(["--tool-context-id", str(kwargs.get("tool_context_id") or "")]) + if injected_job_kwargs := dict(kwargs.get("injected_job_kwargs") or {}): + args.extend(["--injected-job-kwargs", json.dumps(injected_job_kwargs, sort_keys=True)]) return { "command": sys.executable, "args": args, @@ -298,10 +300,17 @@ class CodexAgentWrapper(BaseAgentWrapper): fork_session = bool(kwargs.get("fork_session", False)) if fork_session and not thread_id: raise ValueError("fork_session=True requires resume or session_id") - requested_tool_context = str(kwargs.get("tool_context_id") or "") + requested_tool_context = json.dumps( + { + "tool_context_id": str(kwargs.get("tool_context_id") or ""), + "injected_job_kwargs": dict(kwargs.get("injected_job_kwargs") or {}), + }, + sort_keys=True, + default=str, + ) if not fork_session and thread_id in self._thread_tool_contexts: if requested_tool_context != self._thread_tool_contexts[thread_id]: - raise ValueError("tool_context_id cannot change when resuming a Codex thread") + raise ValueError("tool_context_id and injected_job_kwargs cannot change when resuming a Codex thread") common = { "approval_mode": self._enum(ApprovalMode, kwargs.get("approval_mode"), ApprovalMode.auto_review), diff --git a/reme/components/agent_wrapper/codex_mcp_server.py b/reme/components/agent_wrapper/codex_mcp_server.py index 7ddfbe59..9ed86f5b 100644 --- a/reme/components/agent_wrapper/codex_mcp_server.py +++ b/reme/components/agent_wrapper/codex_mcp_server.py @@ -1,6 +1,7 @@ """FastMCP STDIO bridge that exposes selected ReMe jobs to Codex.""" import argparse +import json from pathlib import Path from typing import Any @@ -14,10 +15,20 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--workspace", required=True, help="ReMe workspace directory") parser.add_argument("--job", dest="jobs", action="append", required=True, help="ReMe job name; repeat as needed") parser.add_argument("--tool-context-id", default="", help="Context id injected into every job call") + parser.add_argument( + "--injected-job-kwargs", + default="", + help="JSON object of server-owned kwargs injected into every job call", + ) return parser.parse_args() -def _prepare_config(config: dict[str, Any], job_names: list[str], tool_context_id: str = "") -> dict[str, Any]: +def _prepare_config( + config: dict[str, Any], + job_names: list[str], + tool_context_id: str = "", + injected_job_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: """Configure the dedicated child Application to serve selected jobs over MCP STDIO.""" selected = set(job_names) jobs: dict[str, dict[str, Any]] = {} @@ -40,8 +51,11 @@ def _prepare_config(config: dict[str, Any], job_names: list[str], tool_context_i "jobs": job_names, "tool_error_on_failure": True, } + injected = dict(injected_job_kwargs or {}) if tool_context_id: - service["injected_job_kwargs"] = {"tool_context_id": tool_context_id} + injected["tool_context_id"] = tool_context_id + if injected: + service["injected_job_kwargs"] = injected prepared = dict(config) prepared["jobs"] = jobs @@ -60,7 +74,10 @@ def main() -> None: log_to_file=False, log_config=False, ) - config = _prepare_config(config, args.jobs, args.tool_context_id) + injected_job_kwargs = json.loads(args.injected_job_kwargs) if args.injected_job_kwargs else None + if injected_job_kwargs is not None and not isinstance(injected_job_kwargs, dict): + raise TypeError("--injected-job-kwargs must be a JSON object") + config = _prepare_config(config, args.jobs, args.tool_context_id, injected_job_kwargs) ReMe(**config).run_app() diff --git a/reme/steps/benchmark/beam/auto_memory.py b/reme/steps/benchmark/beam/auto_memory.py index 1a2a9945..b0013dc5 100644 --- a/reme/steps/benchmark/beam/auto_memory.py +++ b/reme/steps/benchmark/beam/auto_memory.py @@ -117,7 +117,11 @@ def _interpolate_timestamps(items: list[dict]) -> list[dict]: @R.register("beam_auto_memory_step") class BeamAutoMemoryStep(AutoMemoryStep): - """AutoMemoryStep variant that interpolates timestamps and pins daily_write to the resolved day.""" + """AutoMemoryStep variant that interpolates timestamps for BEAM sessions. + + Date pinning for ``daily_write`` is handled by the base class through the + agent wrapper's server-owned ``injected_job_kwargs``. + """ def _build_messages(self, raw_messages: list) -> list[Msg]: # Interpolate timestamps: if any message carries created_at, fill in @@ -127,6 +131,3 @@ class BeamAutoMemoryStep(AutoMemoryStep): [item if not isinstance(item, dict) else dict(item) for item in raw_messages], ) return [self._to_msg(item) for item in interpolated] - - def _reply_extra_kwargs(self, day: str) -> dict: - return {"tool_defaults": {"daily_write": {"date": day}}} diff --git a/reme/steps/benchmark/lme/auto_memory.py b/reme/steps/benchmark/lme/auto_memory.py index 433a0d58..a27cb45c 100644 --- a/reme/steps/benchmark/lme/auto_memory.py +++ b/reme/steps/benchmark/lme/auto_memory.py @@ -117,7 +117,11 @@ def _interpolate_timestamps(items: list[dict]) -> list[dict]: @R.register("lme_auto_memory_step") class LmeAutoMemoryStep(AutoMemoryStep): - """AutoMemoryStep variant that interpolates timestamps and pins daily_write to the resolved day.""" + """AutoMemoryStep variant that interpolates timestamps for LongMemEval sessions. + + Date pinning for ``daily_write`` is handled by the base class through the + agent wrapper's server-owned ``injected_job_kwargs``. + """ def _build_messages(self, raw_messages: list) -> list[Msg]: # Interpolate timestamps: if any message carries created_at, fill in @@ -127,6 +131,3 @@ class LmeAutoMemoryStep(AutoMemoryStep): [item if not isinstance(item, dict) else dict(item) for item in raw_messages], ) return [self._to_msg(item) for item in interpolated] - - def _reply_extra_kwargs(self, day: str) -> dict: - return {"tool_defaults": {"daily_write": {"date": day}}} diff --git a/reme/steps/evolve/auto_memory.py b/reme/steps/evolve/auto_memory.py index c6605618..b0ceef88 100644 --- a/reme/steps/evolve/auto_memory.py +++ b/reme/steps/evolve/auto_memory.py @@ -312,11 +312,17 @@ class AutoMemoryStep(BaseStep): ) self.logger.info(f"[{self.name}] agent start path={note_path} template={template_key}") + # Existing-note updates are restricted to the resolved note path. New + # notes retain the upstream ``daily_write`` date behavior, where the + # model supplies the date from the prompt. + reply_kwargs = self._reply_extra_kwargs(day) + if not created: + reply_kwargs["injected_job_kwargs"] = {"_allowed_paths": [note_path]} result = await self.agent_wrapper.reply( user_message, system_prompt=self.prompt_format("system_prompt"), job_tools=self.create_tools if created else self.update_tools, - **self._reply_extra_kwargs(day), + **reply_kwargs, ) self.logger.info(f"[{self.name}] agent done path={note_path} has_result={bool(result.get('result'))}") diff --git a/reme/steps/file_io/_path.py b/reme/steps/file_io/_path.py index 1d87d0b6..e426ac34 100644 --- a/reme/steps/file_io/_path.py +++ b/reme/steps/file_io/_path.py @@ -82,10 +82,17 @@ def resolve_path( return workspace_path.resolve(), None return None, "`path` is required" s = str(raw).strip() + # ``~`` paths refer to the user's home directory rather than the workspace. + # They are deliberately unsupported: do not expand them (which could make a + # home-relative path look like a workspace escape), and report them as a + # missing target to callers. + if s.startswith("~"): + return None, f"file {s!r} does not exist" p = Path(s) workspace = workspace_path.resolve() if p.is_absolute(): - logger.info("absolute path detected, recommending relative paths") + if Path(s).is_absolute(): + logger.info("absolute path detected, recommending relative paths") target = p.resolve() else: for part in p.parts: @@ -98,6 +105,48 @@ def resolve_path( return target, None +def _check_path_permission(workspace_path: Path, target: Path, allowed_paths) -> bool: + """Return whether ``target`` is covered by the optional allowed-path scope. + + ``None`` leaves access unrestricted. Existing files allow that exact path; + existing directories allow their descendants. Missing entries are treated + as directory-like path prefixes, so a scoped ``write`` may create the path + itself or a descendant. Containment is path-component based, never a string + prefix. Entries beginning with ``~`` are ignored. Invalid constraints and + workspace-escaping entries fail closed. + """ + if allowed_paths is None: + return True + if isinstance(allowed_paths, (str, Path)): + allowed_paths = [allowed_paths] + if not isinstance(allowed_paths, (list, tuple)) or not allowed_paths: + logger.warning("invalid _allowed_paths constraint; denying access (fail closed)") + return False + + allowed_files: list[Path] = [] + allowed_dirs: list[Path] = [] + for raw in allowed_paths: + raw_string = str(raw).strip() + if raw_string.startswith("~"): + logger.warning(f"home-relative _allowed_paths entry {raw_string!r}; skipping") + continue + resolved, err = resolve_path(workspace_path, raw_string) + if err or resolved is None: + logger.warning(f"invalid _allowed_paths entry {raw_string!r} ({err}); denying access (fail closed)") + return False + if resolved.is_file(): + allowed_files.append(resolved) + elif resolved.is_dir(): + allowed_dirs.append(resolved) + else: + allowed_dirs.append(resolved) + + resolved_target = target.resolve() + return resolved_target in allowed_files or any( + is_relative_to(resolved_target, directory) for directory in allowed_dirs + ) + + def gate_md(target: Path) -> tuple[Path, bool]: """Markdown gate with compatibility fallback. diff --git a/reme/steps/file_io/edit.py b/reme/steps/file_io/edit.py index 09707e3c..2e412d54 100644 --- a/reme/steps/file_io/edit.py +++ b/reme/steps/file_io/edit.py @@ -4,7 +4,7 @@ import frontmatter import yaml from ._file_io import get_path_lock, read_file_safe, write_file_safe -from ._path import NON_MD_WARNING, gate_md, resolve_path +from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ..base_step import BaseStep from ...components import R @@ -17,6 +17,9 @@ class EditStep(BaseStep): re-emitted unchanged — matches that fall inside front matter are ignored, so a typo in `old` cannot corrupt structured metadata. + Permission: honors the request-scoped ``_allowed_paths`` constraint + injected by the server into the RuntimeContext. + Concurrency: in-process per-path ``asyncio.Lock`` serializes the read-modify-write cycle against the same file (multi-worker / multi- process safety is out of scope).""" @@ -50,6 +53,10 @@ class EditStep(BaseStep): target, is_md = gate_md(target) + if not _check_path_permission(self.workspace_path, target, self.context.get("_allowed_paths")): + self._fail("no permission to edit this file", path=str(target)) + return None + lock = await get_path_lock(target) async with lock: if not target.exists(): diff --git a/reme/steps/file_io/frontmatter_update.py b/reme/steps/file_io/frontmatter_update.py index 40d3644f..3e9cf910 100644 --- a/reme/steps/file_io/frontmatter_update.py +++ b/reme/steps/file_io/frontmatter_update.py @@ -17,14 +17,20 @@ from pathlib import Path import frontmatter from ._file_io import get_path_lock -from ._path import resolve_path +from ._path import _check_path_permission, resolve_path from ..base_step import BaseStep from ...components import R @R.register("frontmatter_update_step") class FrontmatterUpdateStep(BaseStep): - """Set frontmatter keys on a markdown file from a ``metadata`` dict.""" + """Set frontmatter keys on a markdown file from a ``metadata`` dict. + + Permission: honors the request-scoped ``_allowed_paths`` constraint + injected by the server into the RuntimeContext; + without it, restricting read/edit/write alone would still leave + frontmatter of arbitrary workspace Markdown files mutable. + """ async def execute(self): assert self.context is not None @@ -37,6 +43,8 @@ class FrontmatterUpdateStep(BaseStep): target, err = resolve_path(workspace_dir, path) if err or target is None: payload: dict = {"path": path, "error": err or "invalid path"} + elif not _check_path_permission(workspace_dir, target, self.context.get("_allowed_paths")): + payload = {"path": path, "error": "no permission to update this file"} else: lock = await get_path_lock(target) async with lock: diff --git a/reme/steps/file_io/read.py b/reme/steps/file_io/read.py index 310d0db0..56fcecbe 100644 --- a/reme/steps/file_io/read.py +++ b/reme/steps/file_io/read.py @@ -3,7 +3,7 @@ from pathlib import Path from ._file_io import read_file_lines_safe, read_file_safe, truncate_text_output -from ._path import NON_MD_WARNING, gate_md, resolve_path +from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ..base_step import BaseStep from ...components import R from ...constants import DEFAULT_MAX_BYTES, MAX_FILE_READ_BYTES @@ -21,6 +21,9 @@ class ReadStep(BaseStep): neighbors (out/in link targets) with name/description meta, fetched via the file_store. Same rendering as SearchStep. max_neighbors_per_direction (int, default 10): cap per direction. + + Permission: honors the request-scoped ``_allowed_paths`` constraint + injected by the server into the RuntimeContext. """ def _fail(self, message: str, **meta) -> None: @@ -101,6 +104,9 @@ class ReadStep(BaseStep): target = self._resolve_target(raw) if target is None: return None + if not _check_path_permission(self.workspace_path, target, self.context.get("_allowed_paths")): + self._fail("no permission to access this file", path=str(target)) + return None if not self._validate_line_args(start_line, end_line): return None if not self._check_file(target): @@ -137,8 +143,8 @@ class ReadStep(BaseStep): requested_end, max_collect_bytes=DEFAULT_MAX_BYTES * 2, ) - except Exception as e: # pylint: disable=broad-except - self._fail(f"read failed: {e}", path=str(target)) + except Exception as exc: # pylint: disable=broad-except + self._fail(f"read failed: {exc}", path=str(target)) return None if s > total: self._fail(f"start_line {s} exceeds file length ({total} lines)", path=str(target), total_lines=total) diff --git a/reme/steps/file_io/write.py b/reme/steps/file_io/write.py index 7a1a3cb6..44d1425d 100644 --- a/reme/steps/file_io/write.py +++ b/reme/steps/file_io/write.py @@ -3,7 +3,7 @@ import frontmatter from ._file_io import detect_file_encoding, get_path_lock, write_file_safe -from ._path import NON_MD_WARNING, gate_md, resolve_path +from ._path import _check_path_permission, NON_MD_WARNING, gate_md, resolve_path from ..base_step import BaseStep from ...components import R @@ -18,6 +18,9 @@ class WriteStep(BaseStep): ``metadata`` are ignored — only the top-level explicit parameters are honored for those two reserved fields. + Permission: honors the request-scoped ``_allowed_paths`` constraint + injected by the server into the RuntimeContext. + Concurrency: in-process per-path ``asyncio.Lock`` serializes concurrent writes to the same file (multi-worker / multi-process safety is out of scope). @@ -44,6 +47,10 @@ class WriteStep(BaseStep): target, is_md = gate_md(target) + if not _check_path_permission(self.workspace_path, target, self.context.get("_allowed_paths")): + self._fail("no permission to write this file", path=str(target)) + return None + # Non-markdown files have no frontmatter convention: name/description # and metadata are silently dropped and the body is written verbatim. if is_md: diff --git a/tests/unit/test_background_steps.py b/tests/unit/test_background_steps.py index a051b0b8..2d515ae7 100644 --- a/tests/unit/test_background_steps.py +++ b/tests/unit/test_background_steps.py @@ -1755,6 +1755,7 @@ def test_auto_memory_uses_message_day_for_historical_create(): def write_historical_note(inputs, _kwargs): assert f"Today: {historical_day}" in inputs assert f"date={historical_day}" in inputs + assert "injected_job_kwargs" not in _kwargs write_file( cwd / "daily" / historical_day / "memory.md", "---\nname: memory\nsession_id: s1\n" diff --git a/tests/unit/test_codex_agent_wrapper.py b/tests/unit/test_codex_agent_wrapper.py index fa7f4d76..b32a3a13 100644 --- a/tests/unit/test_codex_agent_wrapper.py +++ b/tests/unit/test_codex_agent_wrapper.py @@ -75,6 +75,38 @@ def test_mcp_config_uses_stdio_bridge_and_selected_jobs(tmp_path): assert config["args"][config["args"].index("--tool-context-id") + 1] == "ctx-1" +def test_mcp_config_serializes_injected_job_kwargs(tmp_path): + wrapper, _job = _wrapper(tmp_path, mcp_config="custom.yaml") + + config = wrapper._mcp_server_config( # pylint: disable=protected-access + { + "job_tools": ["search"], + "tool_context_id": "ctx-1", + "injected_job_kwargs": {"_allowed_paths": ["daily/2025-06-01/note.md"]}, + }, + ) + + raw = config["args"][config["args"].index("--injected-job-kwargs") + 1] + assert json.loads(raw) == {"_allowed_paths": ["daily/2025-06-01/note.md"]} + + plain = wrapper._mcp_server_config({"job_tools": ["search"]}) # pylint: disable=protected-access + assert "--injected-job-kwargs" not in plain["args"] + + +def test_prepare_config_merges_injected_job_kwargs_with_tool_context(): + prepared = _prepare_config( + {"jobs": {"selected": {"backend": "base"}}}, + ["selected"], + "ctx-1", + {"_allowed_paths": ["daily/2025-06-01/note.md"]}, + ) + + assert prepared["service"]["injected_job_kwargs"] == { + "_allowed_paths": ["daily/2025-06-01/note.md"], + "tool_context_id": "ctx-1", + } + + def test_thread_config_preserves_other_mcp_servers(tmp_path): wrapper, _job = _wrapper(tmp_path) config = wrapper._thread_config( # pylint: disable=protected-access diff --git a/tests/unit/test_injected_job_kwargs.py b/tests/unit/test_injected_job_kwargs.py new file mode 100644 index 00000000..5f4185fe --- /dev/null +++ b/tests/unit/test_injected_job_kwargs.py @@ -0,0 +1,241 @@ +"""Tests for the server-owned ``injected_job_kwargs`` wrapper mechanism. + +Each agent wrapper merges these kwargs into every job tool call after +receiving the model's arguments, rejects model-supplied conflicts, and hides +the injected keys from the tool schema exposed to the model. +""" + +# pylint: disable=protected-access,missing-function-docstring + + +import pytest + +from reme.components.agent_wrapper import AsAgentWrapper, BaseAgentWrapper, CcAgentWrapper +from reme.components.file_store import LocalFileStore +from reme.schema import Response +from reme.steps.evolve.auto_memory import AutoMemoryStep + + +class _Job: + """Minimal job double recording every call.""" + + def __init__(self, name="write"): + self.name = name + self.description = "Write a note" + self.parameters = { + "type": "object", + "properties": { + "path": {"type": "string"}, + "content": {"type": "string"}, + "date": {"type": "string"}, + }, + "required": ["path", "content", "date"], + } + self.calls: list[dict] = [] + + async def __call__(self, **kwargs): + self.calls.append(kwargs) + return Response(answer="done") + + +# -- base helpers --------------------------------------------------------------- + + +def test_merge_rejects_model_supplied_conflicts(): + injected = {"_allowed_paths": ["daily/a.md"], "date": "2025-06-01"} + merged = BaseAgentWrapper._merge_injected_job_kwargs({"path": "daily/a.md"}, injected) + assert merged == {"path": "daily/a.md", "_allowed_paths": ["daily/a.md"], "date": "2025-06-01"} + + with pytest.raises(ValueError, match="cannot be provided by the model: _allowed_paths, date"): + BaseAgentWrapper._merge_injected_job_kwargs( + {"_allowed_paths": ["anywhere"], "date": "1999-01-01"}, + injected, + ) + + +def test_strip_injected_parameters_hides_keys_from_schema(): + job = _Job() + stripped = BaseAgentWrapper._strip_injected_parameters(job.parameters, {"date": "2025-06-01"}) + assert "date" not in stripped["properties"] + assert stripped["required"] == ["path", "content"] + # Underscore keys never appear in schemas; stripping is a no-op then. + untouched = BaseAgentWrapper._strip_injected_parameters(job.parameters, {"_allowed_paths": ["x"]}) + assert untouched["properties"].keys() == job.parameters["properties"].keys() + # The original job schema is never mutated. + assert "date" in job.parameters["properties"] + + +# -- AgentScope wrapper ----------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_as_tool_injects_kwargs_and_rejects_conflicts(): + job = _Job() + tool = AsAgentWrapper._make_tool(job, "ctx-1", {"_allowed_paths": ["daily/a.md"], "date": "2025-06-01"}) + + assert "date" not in tool.input_schema["properties"] + assert tool.input_schema["required"] == ["path", "content"] + + await tool.call(path="daily/a.md", content="hi") + assert job.calls == [ + { + "path": "daily/a.md", + "content": "hi", + "_allowed_paths": ["daily/a.md"], + "date": "2025-06-01", + "tool_context_id": "ctx-1", + }, + ] + + with pytest.raises(ValueError, match="cannot be provided by the model"): + await tool.call(path="daily/a.md", content="hi", _allowed_paths=["everything"]) + assert len(job.calls) == 1 + + +@pytest.mark.asyncio +async def test_as_tool_without_injection_keeps_original_schema(): + job = _Job() + tool = AsAgentWrapper._make_tool(job) + assert tool.input_schema["required"] == ["path", "content", "date"] + + await tool.call(path="a.md", content="x", date="2025-06-01") + assert job.calls == [{"path": "a.md", "content": "x", "date": "2025-06-01"}] + + +# -- Claude Code wrapper ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cc_tool_injects_kwargs_and_rejects_conflicts(): + job = _Job() + tool = CcAgentWrapper._make_tool(job, "ctx-1", {"_allowed_paths": ["daily/a.md"]}) + + assert "_allowed_paths" not in tool.input_schema["properties"] + + result = await tool.handler({"path": "daily/a.md", "content": "hi", "date": "2025-06-01"}) + assert result["is_error"] is False + assert job.calls == [ + { + "path": "daily/a.md", + "content": "hi", + "date": "2025-06-01", + "_allowed_paths": ["daily/a.md"], + "tool_context_id": "ctx-1", + }, + ] + + with pytest.raises(ValueError, match="cannot be provided by the model"): + await tool.handler({"path": "daily/a.md", "content": "hi", "_allowed_paths": ["everything"]}) + assert len(job.calls) == 1 + + +# -- AutoMemoryStep --------------------------------------------------------------- + + +class _RecordingWrapper(BaseAgentWrapper): + """Agent wrapper double capturing reply() kwargs.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.calls: list[dict] = [] + + async def reply(self, inputs, **kwargs) -> dict: + self.calls.append(kwargs) + return {"session_id": "s-1", "last_message": {}, "result": "ok"} + + +async def _auto_memory_step(wrapper) -> AutoMemoryStep: + store = LocalFileStore(name="t_inject", embedding_store="") # workspace rooted at cwd + await store.start() + return AutoMemoryStep(name="auto_memory", agent_wrapper=wrapper, file_store=store) + + +@pytest.mark.asyncio +async def test_auto_memory_create_uses_model_supplied_date_and_daily_write(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + wrapper = _RecordingWrapper(name="fake") + step = await _auto_memory_step(wrapper) + + async def no_note(_day, _session_id): + return None + + monkeypatch.setattr(step, "_list_session_note", no_note) + + await step( + session_id="sess-1", + messages=[{"name": "user", "role": "user", "content": "hi", "created_at": "2025-06-01T10:00:00"}], + ) + + assert step.context.response.success is True + assert wrapper.calls[0]["job_tools"] == ["daily_write"] + assert "injected_job_kwargs" not in wrapper.calls[0] + + +@pytest.mark.asyncio +async def test_auto_memory_update_scopes_tools_to_exact_note_path(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + wrapper = _RecordingWrapper(name="fake") + step = await _auto_memory_step(wrapper) + session_id = "sess-1" + note_path = "daily/2025-06-01/existing-note.md" + note_file = tmp_path / note_path + note_file.parent.mkdir(parents=True) + note_file.write_text( + f"---\nsession_id: {session_id}\nsource_conversation: '{step._session_link(session_id)}'\n---\nbody\n", + encoding="utf-8", + ) + + async def existing_note(_day, _session_id): + return {"path": note_path, "session_id": session_id} + + async def no_index(_store, _day, _daily_dir): + return {} + + monkeypatch.setattr(step, "_list_session_note", existing_note) + monkeypatch.setattr("reme.steps.evolve.auto_memory.refresh_day_index", no_index) + + await step( + session_id=session_id, + messages=[{"name": "user", "role": "user", "content": "hi again", "created_at": "2025-06-01T11:00:00"}], + ) + + assert step.context.response.success is True + assert wrapper.calls[0]["job_tools"] == ["read", "edit", "frontmatter_update", "write"] + assert wrapper.calls[0]["injected_job_kwargs"] == {"_allowed_paths": [note_path]} + + +def test_auto_memory_keeps_original_tool_names(): + """BEAM/LME configs define only the original jobs; no *_daily variants exist.""" + step = AutoMemoryStep(name="auto_memory") + assert step.create_tools == ["daily_write"] + assert step.update_tools == ["read", "edit", "frontmatter_update", "write"] + + +def test_auto_memory_create_prompts_match_upstream_date_arguments(): + """Auto-memory prompts keep the upstream model-supplied date argument.""" + from pathlib import Path + + prompt_files = ( + Path("reme/steps/evolve/auto_memory.yaml"), + Path("reme/steps/benchmark/beam/auto_memory.yaml"), + Path("reme/steps/benchmark/lme/auto_memory.yaml"), + ) + for prompt_file in prompt_files: + content = prompt_file.read_text(encoding="utf-8") + assert "date={today}" in content or "`date`: {today}" in content or "`date`:{today}" in content + + +def test_configs_define_original_jobs_without_daily_variants(): + from reme.config import resolve_app_config + + for config_name in ("default", "lme", "beam"): + config = resolve_app_config(config=config_name, log_config=False) + jobs = config["jobs"] + for name in ("read", "edit", "write", "frontmatter_update", "daily_write"): + assert name in jobs, f"{config_name} missing job {name}" + for name in ("read_daily", "edit_daily", "write_daily"): + assert name not in jobs, f"{config_name} unexpectedly defines {name}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit/test_path_permission.py b/tests/unit/test_path_permission.py new file mode 100644 index 00000000..e8ab63d8 --- /dev/null +++ b/tests/unit/test_path_permission.py @@ -0,0 +1,263 @@ +"""Tests for the request-scoped ``_allowed_paths`` permission mechanism. + +The constraint is server-owned: AutoMemoryStep passes it to the agent wrapper +via ``injected_job_kwargs``; the wrapper merges it into every job tool call +(rejecting model-supplied conflicts) and BaseJob places it into the +per-invocation RuntimeContext, where the file I/O steps read it. + +Covers: +- ReadStep / EditStep / WriteStep / FrontmatterUpdateStep honoring the scope. +- Exact-file scope (least privilege) vs directory scope, custom daily dirs. +- Fail-closed behavior for invalid injected constraints. +- Wrapper-level injection: merge, conflict rejection, schema hiding + (AgentScope, Claude Code, Codex serialization). +- AutoMemoryStep injecting ``date`` (create) / ``_allowed_paths`` (update) + while keeping the original read/edit/write/frontmatter_update tool names. +""" + +# pylint: disable=protected-access + +import os +import tempfile +from pathlib import Path + +import pytest + +from reme.components.file_store import LocalFileStore +from reme.steps.file_io._path import _check_path_permission +from reme.steps.file_io.edit import EditStep +from reme.steps.file_io.frontmatter_update import FrontmatterUpdateStep +from reme.steps.file_io.read import ReadStep +from reme.steps.file_io.write import WriteStep + + +class temp_chdir: + """Context manager to temporarily chdir into a path and restore on exit.""" + + def __init__(self, path): + self.path = path + self.old = None + + def __enter__(self): + self.old = os.getcwd() + os.chdir(self.path) + return self + + def __exit__(self, *exc): + os.chdir(self.old) + + +def _seed(workspace: Path, rel: str, body: str = "body\n") -> Path: + target = workspace / rel + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(body, encoding="utf-8") + return target + + +async def _make_store() -> LocalFileStore: + store = LocalFileStore(name="t_perm", embedding_store="") + await store.start() + return store + + +async def _run(step_cls, store: LocalFileStore, **kwargs): + """Run a file I/O step; ``_allowed_paths`` rides in like injected job kwargs.""" + step = step_cls(file_store=store) + await step(**kwargs) + return step.context.response + + +NOTE = "daily/2025-06-01/podcast-habits.md" +SIBLING = "daily/2025-06-01/other-note.md" +OUTSIDE = "topics/roadmap.md" + + +@pytest.mark.asyncio +async def test_update_scope_allows_exact_note(): + """All four update tools succeed on the exact note_path they are scoped to.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), NOTE, "---\nname: x\n---\nold text\n") + store = await _make_store() + scope = {"_allowed_paths": [NOTE]} + + resp = await _run(ReadStep, store, path=NOTE, **scope) + assert resp.success is True + assert "old text" in str(resp.answer) + + resp = await _run(EditStep, store, path=NOTE, old="old text", new="new text", **scope) + assert resp.success is True + + resp = await _run(FrontmatterUpdateStep, store, path=NOTE, metadata={"name": "renamed"}, **scope) + assert resp.success is True + + resp = await _run(WriteStep, store, path=NOTE, name="n", description="d", content="rewritten", **scope) + assert resp.success is True + assert "rewritten" in (Path(tmp) / NOTE).read_text(encoding="utf-8") + await store.close() + + +@pytest.mark.asyncio +async def test_update_scope_rejects_sibling_note_in_same_daily_dir(): + """Exact-file scope denies another note in the same daily directory (least privilege).""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), NOTE) + sibling = _seed(Path(tmp), SIBLING, "---\nname: s\n---\nkeep\n") + store = await _make_store() + scope = {"_allowed_paths": [NOTE]} + + for coro in ( + _run(ReadStep, store, path=SIBLING, **scope), + _run(EditStep, store, path=SIBLING, old="keep", new="gone", **scope), + _run(WriteStep, store, path=SIBLING, content="overwrite", **scope), + _run(FrontmatterUpdateStep, store, path=SIBLING, metadata={"name": "hijack"}, **scope), + ): + resp = await coro + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + assert "keep" in sibling.read_text(encoding="utf-8") + await store.close() + + +@pytest.mark.asyncio +async def test_update_scope_rejects_paths_outside_daily_dir(): + """Exact-file scope denies files elsewhere in the workspace.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), NOTE) + outside = _seed(Path(tmp), OUTSIDE, "roadmap\n") + store = await _make_store() + scope = {"_allowed_paths": [NOTE]} + + for coro in ( + _run(ReadStep, store, path=OUTSIDE, **scope), + _run(EditStep, store, path=OUTSIDE, old="roadmap", new="x", **scope), + _run(WriteStep, store, path=OUTSIDE, content="x", **scope), + _run(FrontmatterUpdateStep, store, path=OUTSIDE, metadata={"k": "v"}, **scope), + ): + resp = await coro + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + assert outside.read_text(encoding="utf-8") == "roadmap\n" + await store.close() + + +@pytest.mark.asyncio +async def test_scope_follows_markdown_suffix_gating(): + """A model path without ``.md`` gates to the same file and stays in scope.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), NOTE, "gated\n") + store = await _make_store() + + resp = await _run(ReadStep, store, path=NOTE.removesuffix(".md"), _allowed_paths=[NOTE]) + assert resp.success is True + assert "gated" in str(resp.answer) + await store.close() + + +@pytest.mark.asyncio +async def test_custom_daily_dir_scope_needs_no_config_jobs(): + """A customized daily_dir (e.g. journal/) works because the scope is the note path itself.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + note = "journal/2025-06-01/trip.md" + _seed(Path(tmp), note, "trip\n") + _seed(Path(tmp), "journal/2025-06-01/other.md") + store = await _make_store() + scope = {"_allowed_paths": [note]} + + resp = await _run(ReadStep, store, path=note, **scope) + assert resp.success is True + + resp = await _run(WriteStep, store, path="journal/2025-06-01/other.md", content="x", **scope) + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + await store.close() + + +@pytest.mark.asyncio +async def test_directory_scope_allows_nested_paths(): + """A directory entry acts as a prefix scope with path-component boundaries.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), "daily/2025-06-01/deep/nested.md", "nested\n") + _seed(Path(tmp), "daily-report/leak.md", "leak\n") + store = await _make_store() + scope = {"_allowed_paths": ["daily"]} + + resp = await _run(ReadStep, store, path="daily/2025-06-01/deep/nested.md", **scope) + assert resp.success is True + + resp = await _run(ReadStep, store, path="daily-report/leak.md", **scope) + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + await store.close() + + +@pytest.mark.asyncio +async def test_nonexistent_allowed_paths_allow_component_bounded_descendants(): + """A missing allowed path permits itself and descendants, not string-prefix siblings.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + workspace = Path(tmp) + _seed(workspace, "X/YZ/not-allowed.md", "keep\n") + store = await _make_store() + + scope = {"_allowed_paths": ["X/Y"]} + assert _check_path_permission(workspace, workspace / "X/Y", scope["_allowed_paths"]) + assert _check_path_permission(workspace, workspace / "X/Y/Z", scope["_allowed_paths"]) + assert not _check_path_permission(workspace, workspace / "X/YZ", scope["_allowed_paths"]) + + resp = await _run(WriteStep, store, path="X/Y/nested", content="nested", **scope) + assert resp.success is True + assert (workspace / "X/Y/nested.md").read_text(encoding="utf-8") == "nested\n" + + resp = await _run(ReadStep, store, path="X/YZ/not-allowed.md", **scope) + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + await store.close() + + +@pytest.mark.asyncio +async def test_home_relative_paths_are_not_supported(): + """Home-relative targets are rejected and home-relative scopes grant nothing.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), "journal/kept.md", "kept\n") + store = await _make_store() + + resp = await _run(ReadStep, store, path="journal/kept.md", _allowed_paths=["~/journal/kept.md"]) + assert resp.success is False + assert "no permission" in str(resp.answer).lower() + + resp = await _run(ReadStep, store, path="~/journal/kept.md") + assert resp.success is False + assert "does not exist" in str(resp.answer).lower() + await store.close() + + +@pytest.mark.asyncio +async def test_invalid_injected_constraints_fail_closed(): + """Empty lists, escaping entries, and non-list values all deny access.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), NOTE, "target\n") + store = await _make_store() + + for bad_scope in ([], ["../escape.md"], ["/etc/passwd"], 42, {"path": NOTE}): + resp = await _run(ReadStep, store, path=NOTE, _allowed_paths=bad_scope) + assert resp.success is False, f"scope {bad_scope!r} should fail closed" + assert "no permission" in str(resp.answer).lower() + await store.close() + + +@pytest.mark.asyncio +async def test_no_constraint_allows_all(): + """Without ``_allowed_paths`` the steps impose no additional restriction.""" + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + _seed(Path(tmp), "anywhere/x.md", "free\n") + store = await _make_store() + + resp = await _run(ReadStep, store, path="anywhere/x.md") + assert resp.success is True + + resp = await _run(WriteStep, store, path="anywhere/new.md", content="ok") + assert resp.success is True + await store.close() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])