diff --git a/control-plane/mergit_worker/models.py b/control-plane/mergit_worker/models.py index b315a813..7ebdfdaa 100644 --- a/control-plane/mergit_worker/models.py +++ b/control-plane/mergit_worker/models.py @@ -5,16 +5,25 @@ from dataclasses import dataclass, field from typing import Any -TRUNCATION_MARKER = b"\n[mergit: output truncated at 1048576 bytes]\n" OUTPUT_CAP_BYTES = 1_048_576 +def truncation_marker(cap: int) -> bytes: + """The line that ends truncated output. It names the cutoff actually applied, so a deployment that sets + `MAX_OUTPUT_BYTES` does not ship a receipt that claims a different one.""" + return f"\n[mergit: output truncated at {cap} bytes]\n".encode() + + +# spec/receipt-v0.md fixes this literal for the default 1 MiB cap. +TRUNCATION_MARKER = truncation_marker(OUTPUT_CAP_BYTES) + + def cap_output(data: bytes, cap: int = OUTPUT_CAP_BYTES) -> bytes: """Cap captured output per spec/receipt-v0.md: keep the first `cap` bytes and end with the marker. The digest of the step covers exactly the returned bytes.""" if len(data) <= cap: return data - return data[:cap] + TRUNCATION_MARKER + return data[:cap] + truncation_marker(cap) @dataclass(frozen=True, slots=True) diff --git a/control-plane/mergit_worker/phi.py b/control-plane/mergit_worker/phi.py index 05118d0d..29e9d54c 100644 --- a/control-plane/mergit_worker/phi.py +++ b/control-plane/mergit_worker/phi.py @@ -13,7 +13,7 @@ import mergit_receipt from .models import CommandResult, cap_output -from .sandbox import SandboxPort +from .sandbox import SandboxError, SandboxPort from .workspace import WorkspaceError @@ -32,7 +32,10 @@ def run_phi(sandbox: SandboxPort, phi: dict[str, Any], max_output_bytes: int, ma sandbox.reset_processes() try: command = sandbox.run(list(phi["command"]), phi["workdir"], int(phi["timeout_s"])) - except WorkspaceError as exc: # a workdir that escapes the checkout + except ( + WorkspaceError, + SandboxError, + ) as exc: # a workdir that escapes the checkout, or a sandbox that failed command = CommandResult(-1, b"", str(exc).encode(), 0, could_not_start=True) inconclusive = command.timed_out or command.could_not_start diff --git a/control-plane/mergit_worker/sandbox.py b/control-plane/mergit_worker/sandbox.py index 9123768d..0650e97b 100644 --- a/control-plane/mergit_worker/sandbox.py +++ b/control-plane/mergit_worker/sandbox.py @@ -9,13 +9,12 @@ from __future__ import annotations import contextlib +import logging import os import threading import time import uuid from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor -from concurrent.futures import TimeoutError as FutureTimeout from dataclasses import dataclass from pathlib import Path from typing import Any, Protocol @@ -40,6 +39,12 @@ } +# How long to wait for an exec stream to end once its container was killed. A healthy daemon ends it at once. +_STREAM_END_AFTER_KILL_S = 30.0 + +log = logging.getLogger(__name__) + + class SandboxError(Exception): """The sandbox could not be created or used.""" @@ -161,14 +166,26 @@ def pump() -> None: except BaseException as exc: # reported to the caller below error.append(exc) - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit(pump) + # A daemon thread, not an executor: an executor's `with` block waits for its thread on exit, so a + # stream the daemon never ends would hang the worker there however the wait above is bounded. + pumper = threading.Thread(target=pump, name="sandbox-exec-pump", daemon=True) + pumper.start() + pumper.join(timeout_s) + if pumper.is_alive(): try: - future.result(timeout=timeout_s) - except FutureTimeout: - self._restart() # kills the process tree; the pump's stream then ends - future.result() - return CommandResult(-1, out.bytes(), err.bytes(), _ms(started), timed_out=True) + self._restart() # kills the process tree; the pump's stream then normally ends + except Exception as exc: # a stuck or refusing daemon: say so, do not hang or crash the worker + raise SandboxError( + f"command timed out and the sandbox could not be restarted: {exc}" + ) from exc + pumper.join(_STREAM_END_AFTER_KILL_S) + if pumper.is_alive(): + # The container is killed and started again, so nothing the command ran survives; only the + # stream reader is stuck. It is abandoned (daemon) and what it captured so far is returned. + log.warning( + "exec stream did not end %.0fs after the sandbox was restarted", _STREAM_END_AFTER_KILL_S + ) + return CommandResult(-1, out.bytes(), err.bytes(), _ms(started), timed_out=True) if error: exc = error[0] if isinstance(exc, docker.errors.APIError): diff --git a/control-plane/mergit_worker/tools.py b/control-plane/mergit_worker/tools.py index 5bf9c7f5..30a8b275 100644 --- a/control-plane/mergit_worker/tools.py +++ b/control-plane/mergit_worker/tools.py @@ -16,7 +16,7 @@ from .llm import ToolCall from .meter import Cost, tool_cost from .models import cap_output -from .sandbox import SandboxPort +from .sandbox import SandboxError, SandboxPort from .workspace import WorkspaceError Recorder = Callable[[str, bytes, bytes, int | None, Cost], int] @@ -115,7 +115,7 @@ def execute(self, call: ToolCall) -> ToolOutcome: step_type, output, exit_status, outcome = handler(self, call.input) is_error, finished = outcome.is_error, outcome.finished shown = outcome.content - except (_BadCall, WorkspaceError, OSError) as exc: + except (_BadCall, WorkspaceError, SandboxError, OSError) as exc: output = str(exc).encode() is_error, finished, shown = True, False, str(exc) wall_ms = int((time.monotonic() - started) * 1000) diff --git a/control-plane/mergit_worker/workspace.py b/control-plane/mergit_worker/workspace.py index 10f85e13..be351712 100644 --- a/control-plane/mergit_worker/workspace.py +++ b/control-plane/mergit_worker/workspace.py @@ -41,6 +41,10 @@ _REF = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/@+-]{0,254}$") _MAX_LIST_ENTRIES = 2000 +# The nested-.git scan visits every entry of the checkout, so an agent-made tree could make every Φ evaluation +# arbitrarily slow. Unlike a listing, it cannot just stop and report what it has: a `.git` planted past the +# cutoff would then go unseen and git would run over it, so reaching the limit refuses (fails closed) instead. +_MAX_SCAN_ENTRIES = 500_000 class WorkspaceError(Exception): @@ -431,30 +435,56 @@ def _walk(self, fd: int, prefix: PurePosixPath, found: list[str]) -> None: os.close(f) -def _find_name(fd: int, name: str) -> str | None: +def _find_name(fd: int, name: str, limit: int | None = None) -> str | None: """The relative path of the first entry called `name` anywhere under `fd`, or None. Iterative, for the same reason `_walk` is: an attacker-made directory chain must not be able to overflow the - stack. Never follows a symlink, so a symlink cannot be used to point this scan outside the tree.""" - stack = [(fd, PurePosixPath())] - opened: list[int] = [] + stack. Never follows a symlink, so a symlink cannot be used to point this scan outside the tree. + + Raises `WorkspaceError` rather than answering None when the tree has more than `limit` entries (default + `_MAX_SCAN_ENTRIES`) or cannot be scanned in full: "not found" is only true if all of it was looked at.""" + cap = _MAX_SCAN_ENTRIES if limit is None else limit + seen = 0 + # One frame per open directory: its fd, its path and the subdirectories still to visit. Only the + # directories on the current path are open, so descriptors grow with the depth of the tree, not its + # size; holding one per directory ran a normal checkout out of them under a 1024 soft limit. + frames: list[tuple[int, PurePosixPath, list[str]]] = [] + cur: tuple[int, PurePosixPath] | None = (fd, PurePosixPath()) try: - while stack: - cur_fd, cur_prefix = stack.pop() + while cur is not None: + cur_fd, cur_prefix = cur + subdirs: list[str] = [] + frames.append((cur_fd, cur_prefix, subdirs)) with os.scandir(cur_fd) as it: - entries = list(it) - for entry in entries: - if entry.name == name: - return str(cur_prefix / entry.name) - if entry.is_symlink(): - continue - if entry.is_dir(follow_symlinks=False): - child = _openat(entry.name, os.O_RDONLY | os.O_DIRECTORY, cur_fd) - opened.append(child) - stack.append((child, cur_prefix / entry.name)) + for entry in it: + seen += 1 + if seen > cap: + raise WorkspaceError( + f"the checkout has more than {cap} entries: refusing to run git over a tree " + f"too large to check for a nested {name}" + ) + if entry.name == name: + return str(cur_prefix / entry.name) + if entry.is_symlink(): + continue + if entry.is_dir(follow_symlinks=False): + subdirs.append(entry.name) + cur = None + while frames: + top_fd, top_prefix, pending = frames[-1] + if pending: + child = pending.pop() + cur = (_openat(child, os.O_RDONLY | os.O_DIRECTORY, top_fd), top_prefix / child) + break + frames.pop() + if top_fd != fd: # the caller owns and closes the fd it passed in + os.close(top_fd) return None + except OSError as exc: # for example out of file descriptors on a pathologically deep tree + raise WorkspaceError(f"cannot scan the checkout for a nested {name}: {exc}") from exc finally: - for f in opened: - os.close(f) + for f, _, _ in frames: + if f != fd: + os.close(f) def _openat(name: str, flags: int, dir_fd: int, mode: int = 0o644) -> int: diff --git a/control-plane/tests/test_harness.py b/control-plane/tests/test_harness.py index 543f6dae..05d5b77d 100644 --- a/control-plane/tests/test_harness.py +++ b/control-plane/tests/test_harness.py @@ -10,7 +10,7 @@ from mergit_worker.harness import HarnessError, run_task from mergit_worker.llm import LLMError -from mergit_worker.models import TRUNCATION_MARKER +from mergit_worker.models import truncation_marker from mergit_worker.tools import ToolLimits from tests.support import ( CHAIN_ID, @@ -246,8 +246,8 @@ def test_tool_output_truncated_with_marker(sandbox, publisher, repo_url): assert cmd_step["type"] == "command" blobs = json.loads(publisher.sealed[-1].trajectory_blobs_json) stored = base64.b64decode(blobs[cmd_step["output_digest"]]) - assert stored.endswith(TRUNCATION_MARKER) - assert len(stored) == cap + len(TRUNCATION_MARKER) + assert stored.endswith(truncation_marker(cap)) + assert len(stored) == cap + len(truncation_marker(cap)) assert mergit_receipt.keccak_hex(stored) == cmd_step["output_digest"] # digest covers stored bytes diff --git a/control-plane/tests/test_review_findings.py b/control-plane/tests/test_review_findings.py new file mode 100644 index 00000000..ebeab794 --- /dev/null +++ b/control-plane/tests/test_review_findings.py @@ -0,0 +1,308 @@ +"""Regression tests for the four findings in issue #31 (left open by the review of #14). + +The sandbox tests drive the real `Sandbox.run` with a fake Docker API, so a daemon that wedges can be +simulated without one. Each test names what it would catch if the fix were removed. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import pytest + +from mergit_worker import sandbox as sandbox_mod +from mergit_worker.llm import ToolCall +from mergit_worker.meter import Cost +from mergit_worker.models import OUTPUT_CAP_BYTES, TRUNCATION_MARKER, cap_output, truncation_marker +from mergit_worker.phi import run_phi +from mergit_worker.sandbox import Sandbox, SandboxConfig, SandboxError +from mergit_worker.tools import ToolBox, ToolLimits +from mergit_worker.workspace import GitLimits, WorkspaceError, _find_name +from tests.support import ProcessSandbox, phi + +# -- 1. a stuck daemon after the restart must not wedge the worker -------------------------------------- + + +class FakeContainer: + id = "c1" + + def __init__(self, on_kill: threading.Event, start_error: Exception | None = None) -> None: + self.calls: list[str] = [] + self._on_kill = on_kill + self._start_error = start_error + + def kill(self) -> None: + self.calls.append("kill") + self._on_kill.set() + + def start(self) -> None: + self.calls.append("start") + if self._start_error: + raise self._start_error + + +class FakeApi: + """exec_start streams one chunk, then blocks until `release` is set (a command still running).""" + + def __init__(self, release: threading.Event, honours_kill: bool) -> None: + self.release = release + self.honours_kill = honours_kill + + def exec_create(self, *_a: Any, **_k: Any) -> dict[str, str]: + return {"Id": "e1"} + + def exec_start(self, *_a: Any, **_k: Any) -> Iterator[tuple[bytes | None, bytes | None]]: + yield b"partial output", None + self.release.wait(60) # a stream that never ends unless the container is killed (or the test ends) + + def exec_inspect(self, _id: str) -> dict[str, int]: + return {"ExitCode": 0} + + +class FakeClient: + def __init__(self, api: FakeApi) -> None: + self.api = api + + +def make_sandbox(tmp_path: Path, *, honours_kill: bool, start_error: Exception | None = None): + release = threading.Event() + container = FakeContainer(release if honours_kill else threading.Event(), start_error) + cfg = SandboxConfig( + image="x", memory="1g", nano_cpus=1, pids_limit=1, runtime=None, max_output_bytes=1000, + workspace_root=tmp_path, git=GitLimits(("file",), 5.0), owner="t", + ) # fmt: skip + api = FakeApi(release, honours_kill) + sbx = Sandbox(object(), container, cfg, FakeClient(api)) # type: ignore[arg-type] + return sbx, container, release + + +def test_a_healthy_kill_ends_the_stream_and_reports_the_timeout(tmp_path: Path) -> None: + sbx, container, _ = make_sandbox(tmp_path, honours_kill=True) + result = sbx.run(["sleep", "99"], ".", 1) + assert result.timed_out and result.exit_code == -1 + assert result.stdout == b"partial output" + assert container.calls == ["kill", "start"] + + +def test_a_stream_that_never_ends_cannot_wedge_the_worker( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Before the fix this line was `future.result()` with no timeout, and the executor's `with` block waited + for the stuck thread again on the way out: the call below never returned.""" + monkeypatch.setattr(sandbox_mod, "_STREAM_END_AFTER_KILL_S", 0.3) + sbx, container, release = make_sandbox(tmp_path, honours_kill=False) + done: list[object] = [] + t = threading.Thread(target=lambda: done.append(sbx.run(["sleep", "99"], ".", 1)), daemon=True) + started = time.monotonic() + t.start() + t.join(10) + release.set() # let the abandoned reader thread finish + assert not t.is_alive(), "Sandbox.run is wedged on a stream the daemon never ended" + assert time.monotonic() - started < 5 + result = done[0] + assert result.timed_out and result.stdout == b"partial output" # type: ignore[attr-defined] + assert container.calls == ["kill", "start"], "the container must still have been killed and restarted" + + +def test_the_sandbox_is_usable_again_after_an_abandoned_stream( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(sandbox_mod, "_STREAM_END_AFTER_KILL_S", 0.2) + sbx, _, release = make_sandbox(tmp_path, honours_kill=False) + sbx.run(["sleep", "99"], ".", 1) + assert sbx._lock.acquire(timeout=1), "the sandbox lock must be released after an abandoned stream" + sbx._lock.release() + release.set() + + +def test_a_restart_that_fails_is_a_sandbox_error_not_a_hang_or_a_raw_docker_error(tmp_path: Path) -> None: + sbx, _, release = make_sandbox(tmp_path, honours_kill=False, start_error=ConnectionError("daemon gone")) + with pytest.raises(SandboxError, match="could not be restarted"): + sbx.run(["sleep", "99"], ".", 1) + release.set() + + +# -- 2. the nested-.git scan is capped, and a cap means refuse, never "none found" ---------------------- + + +def test_the_scan_still_finds_a_nested_git(tmp_path: Path) -> None: + (tmp_path / "a" / "b").mkdir(parents=True) + (tmp_path / "a" / "b" / ".git").write_text("x") + assert _scan(tmp_path) == "a/b/.git" + + +def test_a_clean_tree_under_the_cap_has_no_nested_git(tmp_path: Path) -> None: + for i in range(20): + (tmp_path / f"f{i}").write_text("x") + assert _scan(tmp_path, limit=1000) is None + + +def test_a_tree_over_the_cap_is_refused_not_passed(tmp_path: Path) -> None: + """The dangerous version of this fix returns None at the cap: a `.git` planted after the cutoff would + then be missed and git would run over it. The scan must raise instead.""" + for i in range(50): + (tmp_path / f"f{i:02d}").write_text("x") + with pytest.raises(WorkspaceError, match="more than 10 entries"): + _scan(tmp_path, limit=10) + + +def test_a_git_planted_after_the_cap_is_never_reported_as_absent(tmp_path: Path) -> None: + for i in range(50): + (tmp_path / f"d{i:02d}").mkdir() + (tmp_path / "zzz").mkdir() + (tmp_path / "zzz" / ".git").write_text("gitdir: /hostile\n") + with pytest.raises(WorkspaceError): + _scan(tmp_path, limit=10) # raises; the one outcome it must not have is returning None + + +def test_the_cap_is_on_entries_not_directories_so_one_huge_folder_is_bounded(tmp_path: Path) -> None: + wide = tmp_path / "wide" + wide.mkdir() + for i in range(200): + (wide / f"{i}.txt").write_text("x") + with pytest.raises(WorkspaceError): + _scan(tmp_path, limit=50) + + +def test_a_deep_chain_is_bounded_by_the_same_cap(tmp_path: Path) -> None: + d = tmp_path + for _ in range(60): + d = d / "d" + d.mkdir() + with pytest.raises(WorkspaceError): + _scan(tmp_path, limit=20) + + +@pytest.mark.skipif(not Path("/proc/self/fd").is_dir(), reason="counts open descriptors through /proc") +def test_the_scan_holds_descriptors_for_the_depth_of_the_tree_not_its_size( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Holding one descriptor per directory ran a checkout of about a thousand directories out of them + under the usual 1024 soft limit, so every git step on it failed.""" + import os + + import mergit_worker.workspace as ws + + for i in range(300): + (tmp_path / f"d{i:03d}" / "inner").mkdir(parents=True) + real_openat = ws._openat + baseline = len(os.listdir("/proc/self/fd")) + peak = baseline + + def counting(*a: Any, **k: Any) -> int: + nonlocal peak + opened = real_openat(*a, **k) + peak = max(peak, len(os.listdir("/proc/self/fd"))) + return opened + + monkeypatch.setattr(ws, "_openat", counting) + assert _scan(tmp_path) is None + assert peak - baseline < 10, f"{peak - baseline} descriptors open at once for a tree two levels deep" + + +def test_the_default_cap_is_generous_enough_for_a_real_repository(sandbox: ProcessSandbox) -> None: + assert sandbox.workspace.changed_paths() == [] # a normal checkout is nowhere near the default cap + + +def test_a_scan_the_os_refuses_is_a_workspace_error(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import mergit_worker.workspace as ws + + (tmp_path / "sub").mkdir() + + def boom(*_a: Any, **_k: Any) -> int: + raise OSError(24, "Too many open files") + + monkeypatch.setattr(ws, "_openat", boom) + with pytest.raises(WorkspaceError, match="cannot scan"): + _scan(tmp_path) + + +def test_an_oversized_checkout_makes_phi_inconclusive_not_a_crash( + sandbox: ProcessSandbox, monkeypatch: pytest.MonkeyPatch +) -> None: + import mergit_worker.workspace as ws + + monkeypatch.setattr(ws, "_MAX_SCAN_ENTRIES", 1) + run = run_phi(sandbox, phi(), 1000, 1000) + assert run.holds is None and run.result["exit_code"] == -1 # a person decides, never a guess + + +def _scan(root: Path, limit: int | None = None) -> str | None: + import os + + fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY) + try: + return _find_name(fd, ".git", limit) + finally: + os.close(fd) + + +# -- 3. the truncation marker names the cap that was applied -------------------------------------------- + + +def test_the_default_marker_is_exactly_the_literal_the_spec_fixes() -> None: + assert TRUNCATION_MARKER == b"\n[mergit: output truncated at 1048576 bytes]\n" + assert truncation_marker(OUTPUT_CAP_BYTES) == TRUNCATION_MARKER + + +def test_the_default_cap_ships_the_spec_marker() -> None: + out = cap_output(b"a" * (OUTPUT_CAP_BYTES + 1)) + assert out == b"a" * OUTPUT_CAP_BYTES + TRUNCATION_MARKER + + +@pytest.mark.parametrize("cap", [1, 100, 4096, 2_000_000]) +def test_a_custom_cap_ships_a_marker_that_names_it(cap: int) -> None: + out = cap_output(b"a" * (cap + 1), cap) + assert out == b"a" * cap + f"\n[mergit: output truncated at {cap} bytes]\n".encode() + assert b"1048576" not in out or cap == OUTPUT_CAP_BYTES + + +def test_output_at_the_cap_is_untouched_whatever_the_cap() -> None: + assert cap_output(b"a" * 100, 100) == b"a" * 100 + + +# -- 4. SandboxError is one failed step and one inconclusive Φ, never an aborted task ------------------- + + +class RaisingSandbox(ProcessSandbox): + def run(self, argv: list[str], workdir: str, timeout_s: int) -> Any: + raise SandboxError("argv must be a non-empty list of strings") + + def reset_processes(self) -> None: + pass + + +def test_a_sandbox_error_from_a_tool_call_is_recorded_as_one_failed_step(sandbox: ProcessSandbox) -> None: + rows: list[tuple[str, bytes, bytes, int | None, Cost]] = [] + + def record(step_type: str, inp: bytes, out: bytes, status: int | None, cost: Cost) -> int: + rows.append((step_type, inp, out, status, cost)) + return len(rows) + + raising = RaisingSandbox(sandbox.workspace) + box = ToolBox(raising, record, ToolLimits(5, 1000, 100, 1024)) + outcome = box.execute(ToolCall("t1", "run_command", {"argv": ["true"]})) + assert outcome.is_error and not outcome.finished + assert len(rows) == 1, "a failure is still exactly one receipt step" + assert b"non-empty list" in rows[0][2] + + +def test_a_nul_byte_in_argv_is_a_failed_step_through_the_real_sandbox_check(tmp_path: Path) -> None: + sbx, _, release = make_sandbox(tmp_path, honours_kill=True) + rows: list[Any] = [] + box = ToolBox(sbx, lambda *a: rows.append(a) or len(rows), ToolLimits(5, 1000, 100, 1024)) # type: ignore[arg-type] + outcome = box.execute(ToolCall("t1", "run_command", {"argv": ["echo", "a\x00b"]})) + release.set() + assert outcome.is_error and len(rows) == 1 + + +def test_a_sandbox_error_while_running_phi_is_inconclusive_not_an_aborted_task( + sandbox: ProcessSandbox, +) -> None: + run = run_phi(RaisingSandbox(sandbox.workspace), phi(), 1000, 1000) + assert run.result["exit_code"] == -1 + assert run.holds is None