fix(run_code): enforce the timeout in the runner, report full stderr, and probe cores on macOS

src/agent_env/env/envs/multi_env.py

@@ -17,6 +17,9 @@
 # behaviour, or hold it fixed while bisecting a contention problem.
 _LOAD_CONCURRENCY_ENV_VAR = "AGENT_ENV_LOAD_CONCURRENCY"
 
+# `nproc` is GNU coreutils; macOS has only `getconf`.
+_CORE_COUNT_PROBE = "nproc 2>/dev/null || getconf _NPROCESSORS_ONLN"
+
 def _snapshot_after_load_default() -> bool:
     return os.environ.get(_SNAPSHOT_AFTER_LOAD_ENV_VAR, "").strip().lower() in ("1", "true", "yes")
 
@@ -475,7 +478,7 @@ async def _load_concurrency_limit(self, total: int) -> int:
         cores = None
         if self._sandbox is not None:
             try:
-                cores = int((await self._sandbox.exec_script("nproc")).strip())
+                cores = int((await self._sandbox.exec_script(_CORE_COUNT_PROBE)).strip())
             except Exception as e:  # noqa: BLE001 — a probe, not a dependency
                 logger.warning("Could not read sandbox core count (%s); loading all %d services at once", e, total)
         if not cores or cores < 1:

src/agent_env/task_step/task_steps/run_code.py

@@ -38,7 +38,7 @@
 
 _RUN_DIR = "/tmp/agent_env_run_code"
 
-# The bootstrap we run in the sandbox: `python3 runner.py <run_dir> <entrypoint> <script_file>`.
+# The bootstrap we run in the sandbox: `python3 runner.py <run_dir> <entrypoint> <script_file> <timeout_seconds>`.
 _RUNNER_SOURCE = (Path(__file__).parent / "run_code_runner.py").read_text(encoding="utf-8")
 
 # A universe member with one of these names would shadow run_code's own staging files.
@@ -68,8 +68,9 @@ class RunCodeTaskStep(TaskStep):
     # Not an __init__ default (script_file defaults to None) — the universe entry
     # chosen when several .py members exist.
     ENTRY_WHEN_AMBIGUOUS: ClassVar[str] = "run.py"
-    _TIMEOUT_EXIT_CODE: ClassVar[int] = 124  # `timeout` coreutil convention
-    _STDERR_TRUNCATE: ClassVar[int] = 1000  # max stderr chars surfaced in errors
+    _TIMEOUT_EXIT_CODE: ClassVar[int] = 124  # the runner's status at the limit, after GNU timeout's
+    # The end of a failed script's stderr that its error keeps: the error is stored and raised whole.
+    _STDERR_TAIL_BYTES: ClassVar[int] = 1024 * 1024
     # A collect_artifacts universe can hold hundreds of non-.py members; name a few.
     _SKIPPED_MEMBERS_LOGGED: ClassVar[int] = 10
 
@@ -325,17 +326,16 @@ async def _stage(
         await target.write_file(f"{run_dir}/runner.py", _RUNNER_SOURCE)
 
     async def _run(self, target: _ExecTarget, run_dir: str, script_file: str) -> None:
-        """Run the script under a timeout."""
+        """Run the script, which the runner stops at the timeout."""
         exit_code, _stdout, stderr = await target.run(
-            "timeout", str(self.timeout_seconds),
-            "python3", f"{run_dir}/runner.py", run_dir, self.entrypoint, script_file,
+            "python3", f"{run_dir}/runner.py", run_dir, self.entrypoint, script_file, str(self.timeout_seconds),
         )
         if exit_code == 0:
             return
         reason = (
             f"timed out after {self.timeout_seconds}s"
             if exit_code == self._TIMEOUT_EXIT_CODE
-            else f"exited {exit_code}: {stderr[:self._STDERR_TRUNCATE]}"
+            else f"exited {exit_code}: {self._stderr_tail(stderr)}"
         )
         raise self._exec_error(reason)
 
@@ -348,14 +348,14 @@ async def _read_output(self, target: _ExecTarget, run_dir: str) -> Any:
         path = f"{run_dir}/output.json"
         exit_code, size_out, stderr = await target.run("wc", "-c", path)
         if exit_code != 0:
-            raise self._exec_error(f"wrote no result file: {stderr[:self._STDERR_TRUNCATE]}")
+            raise self._exec_error(f"wrote no result file: {self._stderr_tail(stderr)}")
         if int(size_out.split()[0]) > self.MAX_OUTPUT_BYTES:
             raise self._validation_error(
                 f"result exceeds {self.MAX_OUTPUT_BYTES}-byte cap"
             )
         exit_code, stdout, stderr = await target.run("cat", path)
         if exit_code != 0:
-            raise self._exec_error(f"could not read result file: {stderr[:self._STDERR_TRUNCATE]}")
+            raise self._exec_error(f"could not read result file: {self._stderr_tail(stderr)}")
         try:
             return json.loads(stdout)
         except json.JSONDecodeError as exc:
@@ -448,6 +448,14 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget:
 
     # --- errors ---
 
+    def _stderr_tail(self, stderr: str) -> str:
+        """``stderr`` within ``_STDERR_TAIL_BYTES``, kept from the end, where the exception is."""
+        data = stderr.encode("utf-8", "replace")
+        if len(data) <= self._STDERR_TAIL_BYTES:
+            return stderr
+        tail = data[-self._STDERR_TAIL_BYTES:].decode("utf-8", "ignore")  # a character cut in two goes
+        return f"[first {len(data) - len(tail.encode('utf-8')):,} bytes of stderr cut]\n{tail}"
+
     def _exec_error(self, msg: str) -> RunCodeExecutionError:
         return RunCodeExecutionError(f"run_code '{self.id}': {msg}")
 

src/agent_env/task_step/task_steps/run_code_runner.py

@@ -3,16 +3,96 @@
 Imports the entry script, calls its entrypoint with the JSON input, and writes the
 JSON result. Reads `run_dir`, `entrypoint` and `script_file` from argv.
 
+Given a fourth argument, `timeout_seconds`, it first supervises: it reruns itself without
+it in a new session and stops that whole process group at the limit, so no target needs
+GNU `timeout`.
+
 Both the entry script's own directory and `run_dir` go on sys.path, so a multi-file
 artifact can import its siblings whether it is laid out flat or under a subdirectory.
-Stdlib only — it runs in the sandbox, which may not have agent_env installed.
+Stdlib only, and Python 3.9 compatible: it runs in the sandbox, which may not have
+agent_env installed, or anything newer than macOS's /usr/bin/python3.
 """
 
+import sys
+
+_TIMEOUT_EXIT_CODE = 124  # GNU timeout's, which run_code reports as a timeout
+_GRACE_SECONDS = 5  # between SIGTERM and SIGKILL
+_POLL_SECONDS = 0.1
+
+
+def _signal_group(pgid, signum):
+    try:
+        os.killpg(pgid, signum)
+    except (ProcessLookupError, PermissionError):  # gone, or (macOS) only zombies are left
+        pass
+
+
+def _supervise(command, timeout_seconds):
+    """Run `command` in a new session; return its exit status, or 124 if it hit the limit.
+
+    At the limit, or on a signal forwarded to it, the group gets that signal (SIGTERM at the
+    limit), then SIGKILL once the child exits or the grace period runs out.
+    """
+    child = None
+    early = []  # signals that came before the child did
+    deadline = time.monotonic() + timeout_seconds
+    kill_at = None
+    timed_out = False
+
+    def forward(signum, _frame):
+        nonlocal kill_at
+        if child is None:
+            early.append(signum)
+        else:
+            _signal_group(child.pid, signum)
+        kill_at = kill_at or time.monotonic() + _GRACE_SECONDS
+
+    # Caught before the child starts, as GNU timeout does, so it gets their defaults: one ignored here (a background
+    # job's SIGINT) would stay ignored in it.
+    for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
+        signal.signal(signum, forward)
+    child = subprocess.Popen(command, start_new_session=True)
+    for signum in early:
+        _signal_group(child.pid, signum)
+    while child.poll() is None:
+        now = time.monotonic()
+        if kill_at is not None and now >= kill_at:
+            break
+        if not timed_out and now >= deadline:
+            timed_out = True
+            _signal_group(child.pid, signal.SIGTERM)
+            kill_at = kill_at or now + _GRACE_SECONDS
+        try:
+            child.wait(_POLL_SECONDS)
+        except subprocess.TimeoutExpired:
+            pass
+    if kill_at is not None:
+        _signal_group(child.pid, signal.SIGKILL)
+    status = child.wait()
+    if timed_out:
+        return _TIMEOUT_EXIT_CODE
+    return status if status >= 0 else 128 - status
+
+
+if len(sys.argv) == 5:
+    # Only the supervisor imports these, and not from the run dir (sys.path[0]): an artifact member named after one
+    # would otherwise run here, outside the limit. The script's own process imports just what it always has.
+    import os  # loaded at startup, so never from the run dir
+
+    # Python puts the realpath there, and the run dir is often reached through a symlink (macOS's /tmp).
+    here = os.path.dirname(os.path.realpath(__file__))
+    sys.path[:] = [entry for entry in sys.path if os.path.realpath(entry or ".") != here]
+    import signal
+    import subprocess
+    import time
+
+    sys.exit(_supervise([sys.executable] + sys.argv[:4], float(sys.argv[4])))
+
+# Below the supervisor, which must not import these.
 import importlib.machinery
 import importlib.util
 import json
 import os
-import sys
 
 run_dir, entrypoint, script_file = sys.argv[1], sys.argv[2], sys.argv[3]
 

tst/unit/env/envs/test_multi_env_load_concurrency.py

@@ -10,11 +10,13 @@
 from __future__ import annotations
 
 import asyncio
+import shutil
+import subprocess
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from agent_env.env.envs.multi_env import MultiEnv
+from agent_env.env.envs.multi_env import _CORE_COUNT_PROBE, MultiEnv
 from agent_env.providers.env_providers import EnvironmentGatewayProvider
 from agent_env.providers.env_state import LocalPostgresStateProvider
 
@@ -48,6 +50,22 @@ async def test_limit_scales_with_cores(cores, total, expected):
     assert await _env(cores)._load_concurrency_limit(total) == expected
 
 
+def _probe(path):
+    out = subprocess.run([shutil.which("bash"), "-c", _CORE_COUNT_PROBE], env={"PATH": path},
+                         capture_output=True, text=True, check=True).stdout
+    return int(out.strip())
+
+
+def test_the_core_count_probe_reads_this_hosts_cores():
+    assert _probe("/usr/bin:/bin") >= 1
+
+
+def test_the_core_count_probe_falls_back_to_getconf_without_nproc(tmp_path):
+    # macOS has no `nproc`; this stands in for it on any host.
+    (tmp_path / "getconf").symlink_to(shutil.which("getconf"))
+    assert _probe(str(tmp_path)) >= 1
+
+
 @pytest.mark.asyncio
 async def test_unreadable_core_count_keeps_todays_behaviour():
     """A probe failure must not quietly change how a load runs."""