python-yama 0.1.0__py3-none-any.whl

yama/tools/bash/fs.py

@@ -0,0 +1,530 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import re
+import shutil
+import signal
+import sys
+import uuid
+from pathlib import Path, PurePosixPath
+from typing import Any, Mapping
+
+from ...mocks import MockConfigurationError, MockInvocationError, ToolMockContext, ToolMockError
+from ...models import sanitize_case_key
+
+_FS_FIELDS = {"cwd", "network", "backend", "files", "bin", "env", "limits"}
+_CONTENT_FIELDS = {"content", "file", "mode"}
+_LIMIT_FIELDS = {"cpu_seconds", "memory_mb", "max_processes", "timeout_seconds"}
+_BACKENDS = {"auto", "macos-sandbox-exec", "none"}
+_ENV_ALLOWLIST = ("PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "USER", "SHELL")
+DEFAULT_TIMEOUT_SECONDS = 30.0
+
+# Live persistent-shell sessions, keyed by id(context.state) -- the same run's
+# ToolMockContext.state dict is reused (mutated in place, never replaced) for
+# the whole run, so its identity is a stable per-run key. The session itself
+# holds a real subprocess handle, which cannot go into context.state: that
+# dict must stay JSON serializable (mocks.py clones it after every call).
+_SESSIONS: dict[int, "ShellSession"] = {}
+
+
+def _resolve_backend(requested: str) -> str:
+    if requested not in _BACKENDS:
+        raise MockConfigurationError(
+            f"mocks.fs.backend must be one of {sorted(_BACKENDS)}: {requested!r}"
+        )
+    if requested == "macos-sandbox-exec" and sys.platform != "darwin":
+        raise MockConfigurationError(
+            "mocks.fs.backend: macos-sandbox-exec requires macOS"
+        )
+    if requested == "auto":
+        if sys.platform == "darwin":
+            return "macos-sandbox-exec"
+        raise MockConfigurationError(
+            "mocks.fs has no sandboxing backend on this platform; set "
+            "mocks.fs.backend: none to run without OS-level isolation"
+        )
+    return requested
+
+
+def _safe_relpath(label: str, value: Any) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise MockConfigurationError(f"{label} path must be a non-empty string")
+    pure = PurePosixPath(value)
+    if pure.is_absolute():
+        raise MockConfigurationError(f"{label} path must be relative: {value}")
+    if any(part == ".." for part in pure.parts):
+        raise MockConfigurationError(f"{label} path cannot contain '..': {value}")
+    return pure.as_posix()
+
+
+def _resolve_content(
+    label: str,
+    value: Any,
+    *,
+    base_dir: Path | None,
+    root: Path | None,
+) -> dict[str, Any]:
+    if isinstance(value, str):
+        return {"content": value, "mode": None}
+    if not isinstance(value, dict):
+        raise MockConfigurationError(f"{label} must be a string or an object")
+    unknown = set(value) - _CONTENT_FIELDS
+    if unknown:
+        raise MockConfigurationError(f"unknown {label} fields: {sorted(unknown)}")
+    if ("content" in value) == ("file" in value):
+        raise MockConfigurationError(f"{label} must set exactly one of content or file")
+    mode = value.get("mode")
+    if mode is not None and not isinstance(mode, str):
+        raise MockConfigurationError(f"{label} mode must be a string")
+    if "content" in value:
+        content = value["content"]
+        if not isinstance(content, str):
+            raise MockConfigurationError(f"{label} content must be a string")
+        return {"content": content, "mode": mode}
+    reference = value["file"]
+    if not isinstance(reference, str) or not reference.strip():
+        raise MockConfigurationError(f"{label} file must be a relative path")
+    if base_dir is None:
+        raise MockConfigurationError(f"{label} file references require an evals directory")
+    path = (base_dir / reference).resolve()
+    boundary = Path(root or base_dir).resolve()
+    if not path.is_relative_to(boundary):
+        raise MockConfigurationError(f"{label} file must be inside the plugin: {path}")
+    if not path.is_file():
+        raise MockConfigurationError(f"{label} file does not exist: {path}")
+    try:
+        content = path.read_text(encoding="utf-8")
+    except OSError as exc:
+        raise MockConfigurationError(f"cannot read {label} file {path}: {exc}") from exc
+    return {"content": content, "mode": mode}
+
+
+def normalize_fs_config(
+    fs: Any,
+    *,
+    base_dir: Path | None = None,
+    root: Path | None = None,
+) -> dict[str, Any]:
+    if not isinstance(fs, dict):
+        raise MockConfigurationError("mocks.fs must be an object")
+    unknown = set(fs) - _FS_FIELDS
+    if unknown:
+        raise MockConfigurationError(f"unknown mocks.fs fields: {sorted(unknown)}")
+
+    cwd_raw = fs.get("cwd", "")
+    cwd = _safe_relpath("mocks.fs.cwd", cwd_raw) if cwd_raw else ""
+
+    network = fs.get("network", False)
+    if not isinstance(network, bool):
+        raise MockConfigurationError("mocks.fs.network must be a boolean")
+
+    backend = _resolve_backend(fs.get("backend", "auto"))
+
+    files_raw = fs.get("files", {})
+    if not isinstance(files_raw, dict):
+        raise MockConfigurationError("mocks.fs.files must be an object")
+    files: dict[str, dict[str, Any]] = {}
+    for key, value in files_raw.items():
+        relpath = _safe_relpath("mocks.fs.files", key)
+        files[relpath] = _resolve_content(
+            f"mocks.fs.files.{key}", value, base_dir=base_dir, root=root
+        )
+
+    bin_raw = fs.get("bin", {})
+    if not isinstance(bin_raw, dict):
+        raise MockConfigurationError("mocks.fs.bin must be an object")
+    bin_files: dict[str, dict[str, Any]] = {}
+    for key, value in bin_raw.items():
+        if not isinstance(key, str) or not key.strip() or "/" in key:
+            raise MockConfigurationError(
+                f"mocks.fs.bin command names must be single path segments: {key!r}"
+            )
+        bin_files[key] = _resolve_content(
+            f"mocks.fs.bin.{key}", value, base_dir=base_dir, root=root
+        )
+
+    env_raw = fs.get("env", {})
+    if not isinstance(env_raw, dict) or not all(
+        isinstance(k, str) and isinstance(v, str) for k, v in env_raw.items()
+    ):
+        raise MockConfigurationError("mocks.fs.env must be an object of string values")
+
+    limits_raw = fs.get("limits", {})
+    if not isinstance(limits_raw, dict):
+        raise MockConfigurationError("mocks.fs.limits must be an object")
+    unknown_limits = set(limits_raw) - _LIMIT_FIELDS
+    if unknown_limits:
+        raise MockConfigurationError(
+            f"unknown mocks.fs.limits fields: {sorted(unknown_limits)}"
+        )
+    limits: dict[str, float] = {}
+    for field_name in _LIMIT_FIELDS:
+        if field_name not in limits_raw:
+            continue
+        value = limits_raw[field_name]
+        if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
+            raise MockConfigurationError(
+                f"mocks.fs.limits.{field_name} must be a positive number"
+            )
+        limits[field_name] = float(value)
+
+    return {
+        "cwd": cwd,
+        "network": network,
+        "backend": backend,
+        "files": files,
+        "bin": bin_files,
+        "env": dict(env_raw),
+        "limits": limits,
+    }
+
+
+def compile_cli_config(cli: Mapping[str, Any]) -> dict[str, Any]:
+    """Compile a normalized mocks.cli tree into a synthetic fs_config.
+
+    Each mocked top-level command becomes a real, PATH-shadowed shell script
+    (materialized the same way as mocks.fs.bin) that delegates matching to
+    yama.tools.bash._cli_shim, which reuses the exact same rule-tree walk and _matches()
+    semantics the old in-process interpreter used -- so behavior (including
+    regex fidelity) is unchanged even though execution now goes through a
+    real shell instead of a hand-rolled tokenizer.
+    """
+    files: dict[str, dict[str, Any]] = {}
+    bin_files: dict[str, dict[str, Any]] = {}
+    for name, node in cli.items():
+        if name in {"passthrough", "passthrough_timeout_seconds"}:
+            continue
+        rules_relpath = f".mockbin-rules/{name}.json"
+        files[rules_relpath] = {
+            "content": json.dumps({"name": name, "node": node}, ensure_ascii=False),
+            "mode": None,
+        }
+        bin_files[name] = {
+            "content": (
+                "#!/bin/sh\n"
+                f'exec "$YAMA_PYTHON" -m yama.tools.bash._cli_shim '
+                f'"$YAMA_WORKSPACE/{rules_relpath}" "$@"\n'
+            ),
+            "mode": None,
+        }
+    return {
+        "cwd": "",
+        "network": False,
+        "backend": "none",
+        "files": files,
+        "bin": bin_files,
+        "env": {},
+        "limits": {},
+        "include_system_path": bool(cli.get("passthrough", False)),
+    }
+
+
+def materialize_workspace(workspace_dir: Path, fs_config: Mapping[str, Any]) -> None:
+    workspace_dir.mkdir(parents=True, exist_ok=True)
+    for relpath, entry in fs_config["files"].items():
+        target = workspace_dir / relpath
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(entry["content"], encoding="utf-8")
+        if entry.get("mode") is not None:
+            target.chmod(int(entry["mode"], 8))
+    bin_files = fs_config["bin"]
+    if bin_files:
+        bin_dir = workspace_dir / ".mockbin"
+        bin_dir.mkdir(parents=True, exist_ok=True)
+        for name, entry in bin_files.items():
+            target = bin_dir / name
+            target.write_text(entry["content"], encoding="utf-8")
+            target.chmod(0o755)
+
+
+def build_macos_profile(workspace_dir: Path, *, network: bool) -> str:
+    home = Path.home()
+    sensitive = (".ssh", ".aws", ".config/gcloud", ".gnupg", "Library/Keychains")
+    deny_reads = "\n".join(
+        f'(deny file-read* (subpath "{home / path}"))' for path in sensitive
+    )
+    network_rule = "" if network else '(deny network*)\n'
+    return f"""(version 1)
+(deny default)
+
+(allow process-fork)
+(allow process-exec)
+(allow signal (target self))
+(allow sysctl-read)
+(allow mach-lookup)
+(allow iokit-open)
+
+(allow file-read*)
+(allow file-write* (subpath "{workspace_dir}"))
+(allow file-write-data (literal "/dev/null") (literal "/dev/tty") (literal "/dev/dtracehelper"))
+
+{deny_reads}
+
+{network_rule}"""
+
+
+def _make_preexec(limits: Mapping[str, float]):
+    def _apply() -> None:
+        os.setsid()
+        try:
+            import resource
+
+            if limits.get("cpu_seconds"):
+                value = int(limits["cpu_seconds"])
+                resource.setrlimit(resource.RLIMIT_CPU, (value, value))
+            if limits.get("memory_mb"):
+                value = int(limits["memory_mb"]) * 1024 * 1024
+                resource.setrlimit(resource.RLIMIT_AS, (value, value))
+            if limits.get("max_processes"):
+                value = int(limits["max_processes"])
+                resource.setrlimit(resource.RLIMIT_NPROC, (value, value))
+        except (ValueError, OSError):
+            pass
+
+    return _apply
+
+
+def _build_env(workspace_dir: Path, fs_config: Mapping[str, Any]) -> dict[str, str]:
+    real_path = os.environ.get("PATH", "/usr/bin:/bin:/usr/sbin:/sbin")
+    bin_dir = workspace_dir / ".mockbin"
+    env = {key: os.environ[key] for key in _ENV_ALLOWLIST if key in os.environ}
+    include_system_path = fs_config.get("include_system_path", True)
+    env["PATH"] = f"{bin_dir}:{real_path}" if include_system_path else str(bin_dir)
+    env.setdefault("HOME", os.environ.get("HOME", str(workspace_dir)))
+    # Exposed so generated .mockbin scripts (mocks.fs.bin as well as compiled
+    # mocks.cli commands) can locate the interpreter running yama, the
+    # workspace root, and -- for rule/node-level `passthrough: true` -- the
+    # real PATH to exec the underlying binary without looping back into their
+    # own shadow script.
+    env["YAMA_WORKSPACE"] = str(workspace_dir)
+    env["YAMA_PYTHON"] = sys.executable
+    env["YAMA_REAL_PATH"] = real_path
+    env.update(fs_config.get("env", {}))
+    return env
+
+
+async def _read_until_marker(
+    stream: asyncio.StreamReader, marker: str
+) -> tuple[str, int | None]:
+    """Scan raw bytes for `<marker> <exit code>\\n`, not readline().
+
+    A command's own output does not always end with a newline (`printf -n`,
+    `echo -n`, ...); if it doesn't, our own marker printf lands glued onto
+    the same line and a readline()-based, "does this line start with the
+    marker" check would never match -- the read then blocks until timeout.
+    Scanning the accumulated buffer for the marker pattern anywhere in it
+    (not just at a line start) makes this correct regardless of what the
+    command's own output looks like.
+    """
+    pattern = re.compile(re.escape(marker).encode("ascii") + rb" (-?\d+)\n")
+    buffer = b""
+    while True:
+        match = pattern.search(buffer)
+        if match is not None:
+            output = buffer[: match.start()]
+            try:
+                code = int(match.group(1))
+            except ValueError:
+                code = None
+            return output.decode("utf-8", errors="replace"), code
+        chunk = await stream.read(4096)
+        if not chunk:
+            return buffer.decode("utf-8", errors="replace"), None
+        buffer += chunk
+
+
+class ShellSession:
+    def __init__(self, workspace_dir: Path, fs_config: Mapping[str, Any]) -> None:
+        self.workspace_dir = workspace_dir
+        self.fs_config = fs_config
+        self.process: asyncio.subprocess.Process | None = None
+        self.closed = False
+
+    def _command_argv(self) -> tuple[list[str], str]:
+        # Resolved against the host's real PATH, not the (possibly
+        # restricted) env we build for the session below -- otherwise a
+        # mocks.cli session with no system PATH in its own env couldn't even
+        # find the `bash` binary to launch in the first place. Passed via
+        # `executable=` rather than as argv[0] itself, so bash's own
+        # self-reported name in messages like "bash: line 3: ..." stays the
+        # short, familiar "bash" instead of an absolute path.
+        bash_path = shutil.which("bash") or "/bin/bash"
+        backend = self.fs_config["backend"]
+        if backend == "macos-sandbox-exec":
+            sandbox_exec_path = shutil.which("sandbox-exec") or "/usr/bin/sandbox-exec"
+            profile_path = self.workspace_dir.parent / "sandbox.sb"
+            profile_path.write_text(
+                build_macos_profile(
+                    self.workspace_dir, network=self.fs_config["network"]
+                ),
+                encoding="utf-8",
+            )
+            return (
+                [
+                    "sandbox-exec",
+                    "-f",
+                    str(profile_path),
+                    bash_path,
+                    "--noprofile",
+                    "--norc",
+                ],
+                sandbox_exec_path,
+            )
+        return (["bash", "--noprofile", "--norc"], bash_path)
+
+    async def start(self) -> None:
+        cwd = self.workspace_dir
+        if self.fs_config["cwd"]:
+            cwd = cwd / self.fs_config["cwd"]
+        cwd.mkdir(parents=True, exist_ok=True)
+        env = _build_env(self.workspace_dir, self.fs_config)
+        argv, executable = self._command_argv()
+        self.process = await asyncio.create_subprocess_exec(
+            *argv,
+            executable=executable,
+            cwd=str(cwd),
+            env=env,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            preexec_fn=_make_preexec(self.fs_config.get("limits", {})),
+        )
+
+    async def run(self, command: str, timeout: float) -> dict[str, Any]:
+        assert self.process is not None
+        marker = f"__yama_fs_{uuid.uuid4().hex}__"
+        payload = (
+            "{\n"
+            f"{command}\n"
+            "} < /dev/null\n"
+            "__yama_ec__=$?\n"
+            f"printf '%s %d\\n' '{marker}' \"$__yama_ec__\"\n"
+            f"printf '%s %d\\n' '{marker}' \"$__yama_ec__\" 1>&2\n"
+        )
+        try:
+            assert self.process.stdin is not None
+            self.process.stdin.write(payload.encode("utf-8"))
+            await self.process.stdin.drain()
+            assert self.process.stdout is not None
+            assert self.process.stderr is not None
+            (out, out_code), (err, err_code) = await asyncio.wait_for(
+                asyncio.gather(
+                    _read_until_marker(self.process.stdout, marker),
+                    _read_until_marker(self.process.stderr, marker),
+                ),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            await self.stop()
+            raise ToolMockError(
+                "BASH_TIMEOUT", f"sandboxed command timed out after {timeout}s"
+            )
+        except (BrokenPipeError, ConnectionResetError) as exc:
+            await self.stop()
+            raise ToolMockError(
+                "BASH_SESSION_DIED", f"sandbox shell session died: {exc}"
+            )
+        if out_code is not None:
+            exit_code = out_code
+        elif err_code is not None:
+            exit_code = err_code
+        elif self.process.returncode is not None:
+            # The shell process itself exited before it got to print either
+            # marker -- e.g. the model's command was malformed enough that
+            # bash raised its own fatal parse error and quit. What it did
+            # manage to write (its own error message) is still real, useful
+            # output; fall back to the real process exit code instead of a
+            # made-up one.
+            exit_code = self.process.returncode
+        else:
+            exit_code = 1
+        if self.process.returncode is not None:
+            # The session is gone either way (see above, or a `exit`/`kill`
+            # inside the command itself) -- mark it closed so the next call
+            # spawns a fresh shell instead of writing to a dead process.
+            await self.stop()
+        return {"stdout": out, "stderr": err, "exit_code": exit_code}
+
+    async def stop(self) -> None:
+        if self.closed:
+            return
+        self.closed = True
+        process = self.process
+        if process is not None and process.returncode is None:
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError):
+                pass
+            try:
+                await asyncio.wait_for(process.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                pass
+
+
+def workspace_dir(context: ToolMockContext) -> Path:
+    """Bash-sandbox workspace directory of the current run.
+
+    Public helper for python mock handlers: the path is a pure function of
+    plugin_root/case_key/run_index, so it can be computed before the run's
+    first `bash` call actually materializes the directory.
+    """
+    return (
+        context.plugin_root
+        / ".yama"
+        / "fs-runs"
+        / sanitize_case_key(context.case_key)
+        / f"run-{context.run_index:03d}"
+    )
+
+
+async def _get_session(
+    context: ToolMockContext, fs_config: Mapping[str, Any]
+) -> ShellSession:
+    session = _SESSIONS.get(id(context.state))
+    if session is None or session.closed:
+        workspace = workspace_dir(context)
+        if not workspace.exists():
+            materialize_workspace(workspace, fs_config)
+        session = ShellSession(workspace, fs_config)
+        await session.start()
+        _SESSIONS[id(context.state)] = session
+        context.state["__fs__"] = {
+            "workspace_dir": str(workspace),
+            "backend": fs_config["backend"],
+        }
+    return session
+
+
+async def run_bash_fs(
+    arguments: Mapping[str, Any],
+    context: ToolMockContext | None,
+    fs_config: Mapping[str, Any],
+) -> dict[str, Any]:
+    if context is None:
+        raise MockInvocationError("mocks.fs requires a ToolMockContext")
+    command = arguments["command"]
+    session = await _get_session(context, fs_config)
+    timeout = fs_config.get("limits", {}).get("timeout_seconds", DEFAULT_TIMEOUT_SECONDS)
+    return await session.run(command, timeout=timeout)
+
+
+async def run_bash_cli(
+    arguments: Mapping[str, Any],
+    context: ToolMockContext | None,
+    cli_config: Mapping[str, Any],
+) -> dict[str, Any]:
+    if context is None:
+        raise MockInvocationError("mocks.cli requires a ToolMockContext")
+    command = arguments["command"]
+    fs_config = compile_cli_config(cli_config)
+    session = await _get_session(context, fs_config)
+    timeout = cli_config.get("passthrough_timeout_seconds", DEFAULT_TIMEOUT_SECONDS)
+    return await session.run(command, timeout=timeout)
+
+
+async def cleanup_session(state: dict[str, Any]) -> None:
+    session = _SESSIONS.pop(id(state), None)
+    if session is not None:
+        await session.stop()

yama/tools/skill/skill.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Mapping
+
+from ..base import Tool
+from ...mocks import ToolMockContext, ToolMockError
+
+
+class SkillTool(Tool):
+    def run(
+        self,
+        arguments: Mapping[str, Any],
+        context: ToolMockContext,
+        config: Mapping[str, Any],
+    ) -> str:
+        name = arguments.get("name")
+        enabled = config.get("enabled", [])
+        if not isinstance(name, str) or name not in enabled:
+            raise ToolMockError(
+                "SKILL_NOT_ENABLED",
+                f"skill {name!r} is not enabled for this case",
+            )
+
+        skills_dir = str(config.get("skills_dir", "skills"))
+        entry = str(config.get("skill_entry", "SKILL.md"))
+        skill_root = (context.plugin_root / skills_dir / name).resolve()
+        requested_file = arguments.get("file", entry)
+        if not isinstance(requested_file, str) or not requested_file.strip():
+            raise ToolMockError(
+                "INVALID_SKILL_FILE", "file must be a non-empty relative path"
+            )
+
+        relative = Path(requested_file)
+        if relative.is_absolute():
+            raise ToolMockError(
+                "INVALID_SKILL_FILE", "file must be relative to the Skill directory"
+            )
+        target = (skill_root / relative).resolve()
+        if not target.is_relative_to(skill_root):
+            raise ToolMockError(
+                "INVALID_SKILL_FILE", "file escapes the selected Skill directory"
+            )
+        if not target.is_file():
+            raise ToolMockError(
+                "SKILL_FILE_NOT_FOUND", f"skill file does not exist: {requested_file}"
+            )
+        try:
+            return target.read_text(encoding="utf-8")
+        except UnicodeDecodeError as exc:
+            raise ToolMockError(
+                "SKILL_FILE_NOT_TEXT",
+                f"skill file is not UTF-8 text: {requested_file}",
+            ) from exc
+
+
+read_skill = SkillTool().run