diff --git a/tests/test_sandbox_resource_limits.py b/tests/test_sandbox_resource_limits.py new file mode 100644 index 0000000..6148922 --- /dev/null +++ b/tests/test_sandbox_resource_limits.py @@ -0,0 +1,39 @@ +"""Resource limits are validated and forwarded without shell interpolation.""" +import pytest + +from tiny_harness._sandbox import SandboxUnavailable, build_bubblewrap_argv + + +@pytest.mark.parametrize("keyword,value", [ + ("cpu_seconds", 0), ("cpu_seconds", -10), ("cpu_seconds", True), + ("memory_mb", 0), ("memory_mb", 16385), ("memory_mb", "1024"), + ("file_mb", -1), ("file_mb", 1025), ("file_mb", 1.5), +]) +def test_resource_caps_reject_invalid_values(tmp_path, monkeypatch, keyword, value): + monkeypatch.setattr("tiny_harness._sandbox.shutil.which", + lambda name, path=None: "/usr/bin/" + name) + args = {"cpu_seconds": 30, "memory_mb": 1024, "file_mb": 32} + args[keyword] = value + with pytest.raises(SandboxUnavailable, match=keyword): + build_bubblewrap_argv("true", workspace=str(tmp_path), cwd=None, **args) + + +def test_resource_caps_are_structured_argv(tmp_path, monkeypatch): + monkeypatch.setattr("tiny_harness._sandbox.shutil.which", + lambda name, path=None: "/usr/bin/" + name) + marker = "echo --as=untrusted; printf 'data'" + argv = build_bubblewrap_argv( + marker, workspace=str(tmp_path), cwd=None, + cpu_seconds=7, memory_mb=256, file_mb=4, + ) + assert "--cpu=7" in argv + assert "--as=268435456" in argv + assert "--fsize=4194304" in argv + assert argv[-3:] == ["/bin/sh", "-c", marker] + + +def test_missing_prlimit_blocks_sandbox(tmp_path, monkeypatch): + monkeypatch.setattr("tiny_harness._sandbox.shutil.which", + lambda name, path=None: "/usr/bin/bwrap" if name == "bwrap" else None) + with pytest.raises(SandboxUnavailable, match="prlimit"): + build_bubblewrap_argv("true", workspace=str(tmp_path), cwd=None) diff --git a/tests/test_shell_bwrap.py b/tests/test_shell_bwrap.py index 07747fe..9acae7a 100644 --- a/tests/test_shell_bwrap.py +++ b/tests/test_shell_bwrap.py @@ -18,7 +18,7 @@ def workspace(tmp_path): def test_sandbox_command_is_structured_and_never_interpolated(workspace, monkeypatch): monkeypatch.setattr("tiny_harness._sandbox.shutil.which", - lambda name, path=None: "/usr/bin/bwrap") + lambda name, path=None: "/usr/bin/" + name) marker = "echo hello; printf 'unsafe?'; echo --ro-bind" argv = build_bubblewrap_argv(marker, workspace=str(workspace), cwd=None) assert argv[0] == "/usr/bin/bwrap" @@ -33,6 +33,10 @@ def test_sandbox_command_is_structured_and_never_interpolated(workspace, monkeyp assert "--share-net" not in argv assert "--ro-bind-try" not in argv assert argv[-3:] == ["/bin/sh", "-c", marker] + assert "--cpu=30" in argv + assert "--as=1073741824" in argv + assert "--fsize=33554432" in argv + assert argv[argv.index("--") + 1] == "/usr/bin/prlimit" assert argv[argv.index("--chdir") + 1] == "/workspace" # The container root is isolated; host / and /etc are not bulk mounted. assert not any( @@ -44,7 +48,7 @@ def test_sandbox_command_is_structured_and_never_interpolated(workspace, monkeyp def test_sandbox_maps_nested_working_directory(workspace, monkeypatch): (workspace / "nested").mkdir() monkeypatch.setattr("tiny_harness._sandbox.shutil.which", - lambda name, path=None: "/usr/bin/bwrap") + lambda name, path=None: "/usr/bin/" + name) argv = build_bubblewrap_argv( "pwd", workspace=str(workspace), cwd=str(workspace / "nested") ) @@ -53,7 +57,7 @@ def test_sandbox_maps_nested_working_directory(workspace, monkeypatch): def test_sandbox_blocks_workspace_escape_and_missing_cwd(workspace, tmp_path, monkeypatch): monkeypatch.setattr("tiny_harness._sandbox.shutil.which", - lambda name, path=None: "/usr/bin/bwrap") + lambda name, path=None: "/usr/bin/" + name) for cwd in (str(tmp_path), str(workspace / "missing")): with pytest.raises(SandboxUnavailable): build_bubblewrap_argv("true", workspace=str(workspace), cwd=cwd) diff --git a/tiny_harness/_config.py b/tiny_harness/_config.py index 011f890..220820b 100644 --- a/tiny_harness/_config.py +++ b/tiny_harness/_config.py @@ -21,6 +21,9 @@ class AgentConfig: no_hitl: bool = False allow_unsafe_shell: bool = False shell_sandbox: str = "off" # 'off' or 'bubblewrap' (Linux only) + sandbox_cpu_seconds: int = 30 + sandbox_memory_mb: int = 1024 + sandbox_file_mb: int = 32 class Prompt: diff --git a/tiny_harness/_sandbox.py b/tiny_harness/_sandbox.py index 22018fb..4400c6a 100644 --- a/tiny_harness/_sandbox.py +++ b/tiny_harness/_sandbox.py @@ -21,7 +21,10 @@ class SandboxUnavailable(RuntimeError): _RUNTIME_FILES = ("/etc/ld.so.cache", "/etc/nsswitch.conf") -def build_bubblewrap_argv(command: str, *, workspace: str, cwd: str | None) -> list[str]: +def build_bubblewrap_argv( + command: str, *, workspace: str, cwd: str | None, + cpu_seconds: int = 30, memory_mb: int = 1024, file_mb: int = 32, +) -> list[str]: """Build an argv array; no untrusted command bytes become bwrap options. User commands run through the shell *inside* the namespace. An unavailable @@ -38,6 +41,18 @@ def build_bubblewrap_argv(command: str, *, workspace: str, cwd: str | None) -> l "bubblewrap executable is unavailable; install bwrap or choose " "shell_sandbox='off' explicitly" ) + prlimit = shutil.which("prlimit", path="/usr/bin:/bin") + if prlimit is None: + raise SandboxUnavailable( + "prlimit is unavailable; cannot enforce requested resource limits" + ) + for label, value, limit in ( + ("cpu_seconds", cpu_seconds, 3600), + ("memory_mb", memory_mb, 16384), + ("file_mb", file_mb, 1024), + ): + if isinstance(value, bool) or not isinstance(value, int) or not 1 <= value <= limit: + raise SandboxUnavailable(f"{label} must be an integer between 1 and {limit}") if not os.path.isdir(workspace): raise SandboxUnavailable("Sandbox workspace must be an existing directory") @@ -79,6 +94,12 @@ def build_bubblewrap_argv(command: str, *, workspace: str, cwd: str | None) -> l "--bind", root, "/workspace", # Only writable host mount. "--chdir", sandbox_cwd, "--", - "/bin/sh", "-c", command, + # RLIMIT_CPU and RLIMIT_AS are PER PROCESS, not aggregate cgroup quotas. + # Every descendant inherits the limits from the shell process. + prlimit, + f"--cpu={cpu_seconds}", + f"--as={memory_mb * 1024 * 1024}", + f"--fsize={file_mb * 1024 * 1024}", + "--", "/bin/sh", "-c", command, )) return argv diff --git a/tiny_harness/cli.py b/tiny_harness/cli.py index aa16598..73cca09 100644 --- a/tiny_harness/cli.py +++ b/tiny_harness/cli.py @@ -19,6 +19,9 @@ def parse_args() -> Namespace: parser.add_argument("--no-hitl", action="store_true", help="Disable human-in-the-loop approval prompts") parser.add_argument("--allow-unsafe-shell", action="store_true", help="Explicitly permit Shell without human approval (dangerous)") parser.add_argument("--shell-sandbox", choices=("off", "bubblewrap"), default="off", help="Optional Linux bwrap Shell sandbox; fails closed if unavailable") + parser.add_argument("--sandbox-cpu-seconds", type=int, default=30, help="Bubblewrap process CPU limit (seconds)") + parser.add_argument("--sandbox-memory-mb", type=int, default=1024, help="Bubblewrap per-process virtual memory limit (MiB)") + parser.add_argument("--sandbox-file-mb", type=int, default=32, help="Bubblewrap file size limit (MiB)") parser.add_argument("--tui", action="store_true", help="Use rich TUI mode (requires 'rich' package)") parser.add_argument("--api-key-env", default=None, help="Env var for API key (defaults to provider-specific variable)") return parser.parse_args() @@ -163,7 +166,7 @@ def main(): from tiny_harness import Agent, Prompt, Config prompt = Prompt("You are a helpful AI assistant.") - config = Config(model=args.model, api_key=api_key, workspace=args.workspace, provider=args.provider, api_base_url=args.api_base_url, max_iterations=args.max_iterations, no_hitl=args.no_hitl, allow_unsafe_shell=args.allow_unsafe_shell, shell_sandbox=args.shell_sandbox) + config = Config(model=args.model, api_key=api_key, workspace=args.workspace, provider=args.provider, api_base_url=args.api_base_url, max_iterations=args.max_iterations, no_hitl=args.no_hitl, allow_unsafe_shell=args.allow_unsafe_shell, shell_sandbox=args.shell_sandbox, sandbox_cpu_seconds=args.sandbox_cpu_seconds, sandbox_memory_mb=args.sandbox_memory_mb, sandbox_file_mb=args.sandbox_file_mb) agent = Agent(prompt=prompt, config=config) if not args.no_hitl: diff --git a/tiny_harness/skills/shell.py b/tiny_harness/skills/shell.py index 8463161..a1a2229 100644 --- a/tiny_harness/skills/shell.py +++ b/tiny_harness/skills/shell.py @@ -92,6 +92,13 @@ def register(agent) -> None: risk_level="dangerous", path_args=("cwd",), ), - partial(run_command_async, structured=True, sandbox_mode=agent._config.shell_sandbox, workspace=agent._guard.guard(".")), + partial( + run_command_async, structured=True, + sandbox_mode=agent._config.shell_sandbox, + workspace=agent._guard.guard("."), + sandbox_cpu_seconds=agent._config.sandbox_cpu_seconds, + sandbox_memory_mb=agent._config.sandbox_memory_mb, + sandbox_file_mb=agent._config.sandbox_file_mb, + ), ) agent._prompt.append(SHELL_PROMPT) diff --git a/tiny_harness/tools/shell.py b/tiny_harness/tools/shell.py index 3c2cb4b..607bdd0 100644 --- a/tiny_harness/tools/shell.py +++ b/tiny_harness/tools/shell.py @@ -76,7 +76,7 @@ def run_command(args: dict) -> str: # The agent registers the async entrypoint below. Unlike the legacy sync helper # it can stop its whole process group when the executor cancels the coroutine. -async def run_command_async(args: dict, *, structured: bool = False, sandbox_mode: str = "off", workspace: str | None = None) -> str | ToolOutcome: +async def run_command_async(args: dict, *, structured: bool = False, sandbox_mode: str = "off", workspace: str | None = None, sandbox_cpu_seconds: int = 30, sandbox_memory_mb: int = 1024, sandbox_file_mb: int = 32) -> str | ToolOutcome: import asyncio import os import signal @@ -103,6 +103,8 @@ async def run_command_async(args: dict, *, structured: bool = False, sandbox_mod # Fail closed: do not use an ordinary shell if bwrap is absent. cmd, use_shell = build_bubblewrap_argv( command, workspace=workspace, cwd=cwd, + cpu_seconds=sandbox_cpu_seconds, memory_mb=sandbox_memory_mb, + file_mb=sandbox_file_mb, ), False cwd = workspace else: