Fix four worker review findings from #14

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)

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
 

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):

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)

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:

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