diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index fc042eb5..7289aa16 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import collections import contextlib import functools import glob @@ -62,6 +63,10 @@ def _runs_in_container(cmd: list[str]) -> bool: _REAP_SECONDS = 5 +# How long a finished command's output is still read, for a job it left in the background that holds it open; and +# how much of each stream is kept, the end of it: above the largest output core reads whole, run_code's 50 MiB result. +_OUTPUT_GRACE_SECONDS = 1 +_OUTPUT_TAIL_BYTES = 64 * 1024 * 1024 # Every command a local sandbox runs carries this, set to the sandbox's id, and so does whatever the command starts, so # teardown can find a process left running on this machine, such as the agent a host-mode install started. @@ -282,6 +287,7 @@ async def exec(self, *command: str) -> Any: logger.debug(f"LocalSandbox exec: {' '.join(cmd)}") process = await asyncio.create_subprocess_exec( *cmd, + stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, env={**os.environ, _SANDBOX_ENV: self.sandbox_id}, @@ -289,8 +295,10 @@ async def exec(self, *command: str) -> Any: return process async def exec_with_output(self, *args: str) -> tuple[int, str, str]: - """Like the base, except that cancelling it kills the command and every process it started. This host - is the sandbox's VM, so nothing else would.""" + """Like the base, but as a remote exec session behaves. Cancelling it kills the command and every process it + started; this host is the sandbox's VM, so nothing else would. The call ends with the command, even if a job + it left in the background still holds its output, which is then closed. Each stream keeps its last + ``_OUTPUT_TAIL_BYTES``, and bytes that aren't UTF-8 become replacement characters.""" spawning = asyncio.ensure_future(self.exec(*args)) try: process = await asyncio.shield(spawning) @@ -298,13 +306,26 @@ async def exec_with_output(self, *args: str) -> tuple[int, str, str]: with contextlib.suppress(Exception): await _stop(await spawning) raise + stdout, stderr = _Tail(), _Tail() + reading = [asyncio.ensure_future(tail.drain(stream)) + for tail, stream in ((stdout, process.stdout), (stderr, process.stderr))] try: - stdout, stderr = await asyncio.gather(process.stdout.read(), process.stderr.read()) - exit_code = await process.wait() + exit_code = await _exit_code(process) + await asyncio.wait(reading, timeout=_OUTPUT_GRACE_SECONDS) except asyncio.CancelledError: await _stop(process) raise - return exit_code, stdout.decode(), stderr.decode() + finally: + for reader in reading: + reader.cancel() + # Our ends of its pipes, which a background job may hold. asyncio has no public way to close them. + process._transport.close() + output = stdout.text(), stderr.text() + for name, tail in (("stdout", stdout), ("stderr", stderr)): + if tail.dropped: # not the command itself: it can carry secrets + logger.warning("Kept the last %d bytes of a command's %s on sandbox %s; %d before them were dropped", + _OUTPUT_TAIL_BYTES, name, self.sandbox_id, tail.dropped) + return exit_code, *output async def _write_unsigned_object(self, object_store: ObjectStore, object_url: str, vm_path: str) -> None: """This host is the VM, so the store writes the object in place, with no exec transport.""" @@ -319,12 +340,50 @@ async def wait_for_vm(self) -> None: pass +class _Tail: + """The end of a stream as it is read: its last ``_OUTPUT_TAIL_BYTES``, and how many bytes before them were + dropped. Chunks, because cutting the front off one growing buffer keeps the memory it grew to.""" + + def __init__(self) -> None: + self.chunks: collections.deque[bytes] = collections.deque() + self.size = self.dropped = 0 + + async def drain(self, stream: asyncio.StreamReader) -> None: + while chunk := await stream.read(1 << 16): + self.chunks.append(chunk) + self.size += len(chunk) + while self.size - len(self.chunks[0]) >= _OUTPUT_TAIL_BYTES: + first = self.chunks.popleft() + self.size -= len(first) + self.dropped += len(first) + + def text(self) -> str: + data = b"".join(self.chunks) + if (excess := len(data) - _OUTPUT_TAIL_BYTES) > 0: + data = data[excess:] + self.dropped += excess + return data.decode(errors="replace") + + +async def _exit_code(process: asyncio.subprocess.Process) -> int: + """``process``'s exit code, as soon as it exits. ``Process.wait()`` also waits for its pipes to close, which a job + it left in the background can hold open, so the exit is checked for meanwhile.""" + waiting = asyncio.ensure_future(process.wait()) + try: + while process.returncode is None: + await asyncio.wait([waiting], timeout=0.05) + return process.returncode + finally: + waiting.cancel() + + async def _stop(process: asyncio.subprocess.Process) -> None: - """SIGKILL ``process`` and every process it started, then reap it, waiting a few seconds at most for one - that escaped the kill to let go of its output.""" + """SIGKILL ``process`` and every process it started, wait a few seconds at most for it to exit, then close our + ends of its pipes, which a job that escaped the kill may still hold.""" _kill_tree(process.pid) with contextlib.suppress(TimeoutError): - await asyncio.wait_for(process.wait(), _REAP_SECONDS) + await asyncio.wait_for(_exit_code(process), _REAP_SECONDS) + process._transport.close() def _kill_tree(root: int) -> None: diff --git a/tst/unit/providers/sandbox_providers/local_sandbox_test.py b/tst/unit/providers/sandbox_providers/local_sandbox_test.py index f5f92eb0..48118519 100644 --- a/tst/unit/providers/sandbox_providers/local_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/local_sandbox_test.py @@ -1,8 +1,10 @@ import asyncio +import contextlib import os import platform import shutil import signal +import subprocess import tempfile from pathlib import Path from types import SimpleNamespace @@ -18,6 +20,7 @@ from agent_env.store import LocalFilesystemObjectStore from agent_env.store.object_store.local.tls import local_ca from agent_env.store.routing import LocalRunObjectStore +from agent_env.task_step.task_steps.run_code import RunCodeTaskStep from agent_env.a2a_agent import a2a_agent as a2a_agent_module from agent_env.a2a_agent.a2a_agent import A2AAgent from tst.unit.store.fakes import FakeObjectStore @@ -675,8 +678,11 @@ def run(args, **kwargs): class _CapturedProcess: + returncode = 0 + def __init__(self): self.stdout = self.stderr = SimpleNamespace(read=self._empty) + self._transport = SimpleNamespace(close=lambda: None) @staticmethod async def _empty(): @@ -824,6 +830,75 @@ async def slow_spawn(*args, **kwargs): os.kill(int(pid_file.read_text()), 0) +@pytest.mark.asyncio +async def test_a_command_reads_nothing_from_this_process(monkeypatch, tmp_path: Path): + seen = {} + + async def spawn(*argv, **kwargs): + seen.update(kwargs) + return _CapturedProcess() + + monkeypatch.setattr(ls.asyncio, "create_subprocess_exec", spawn) + await LocalSandbox(work_dir=tmp_path).exec("cat") + + assert seen["stdin"] is asyncio.subprocess.DEVNULL + + +def _state(pid: int) -> str: + """``ps``'s state for ``pid``: empty once it is gone, ``Z`` while a parent has yet to reap it.""" + return subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip() + + +@pytest.mark.asyncio +async def test_the_call_ends_with_the_command_though_a_job_it_left_holds_its_output(tmp_path: Path): + started = asyncio.get_running_loop().time() + exit_code, stdout, _ = await LocalSandbox(work_dir=tmp_path).exec_with_output("bash", "-c", "sleep 30 & echo $!") + job = int(stdout) + try: + assert exit_code == 0 + assert asyncio.get_running_loop().time() - started < 5 + finally: + os.kill(job, signal.SIGKILL) + + +@pytest.mark.asyncio +async def test_a_job_left_writing_to_the_commands_output_is_cut_off_with_it(tmp_path: Path): + await LocalSandbox(work_dir=tmp_path).exec_with_output("bash", "-c", "yes & echo $! > /app/job.pid") + job = int((tmp_path / "job.pid").read_text()) + try: + for _ in range(100): + if _state(job)[:1] in ("", "Z"): + break + await asyncio.sleep(0.05) + assert _state(job)[:1] in ("", "Z") + finally: + with contextlib.suppress(ProcessLookupError): + os.kill(job, signal.SIGKILL) + + +@pytest.mark.asyncio +async def test_output_that_isnt_utf8_is_kept_with_replacement_characters(tmp_path: Path): + exit_code, stdout, _ = await LocalSandbox(work_dir=tmp_path).exec_with_output("bash", "-c", r"printf 'ok \xff'") + + assert (exit_code, stdout) == (0, "ok \ufffd") + + +def test_a_run_code_result_at_its_limit_is_kept_whole(): + assert ls._OUTPUT_TAIL_BYTES > RunCodeTaskStep.MAX_OUTPUT_BYTES + + +@pytest.mark.asyncio +async def test_only_the_end_of_a_long_output_is_kept(monkeypatch, tmp_path: Path, caplog): + monkeypatch.setattr(ls, "_OUTPUT_TAIL_BYTES", 1000) + + _, stdout, _ = await LocalSandbox(work_dir=tmp_path).exec_with_output( + "bash", "-c", "head -c 200000 /dev/zero | tr '\\0' x; echo END") + + assert len(stdout) == 1000 and stdout.endswith("xEND\n") + assert "199004 before them were dropped" in caplog.text + assert "head -c" not in caplog.text # the command is left out: it can carry secrets + + @pytest.mark.asyncio async def test_a_command_carries_its_sandboxs_id(tmp_path: Path): sandbox = LocalSandbox(work_dir=tmp_path)