Make a local sandbox's commands read no input, end with the command, and keep a bounded tail of valid text

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,29 +287,45 @@ 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},
         )
         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)
         except asyncio.CancelledError:
             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:

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)