| 1 | """Bounded real-agent capture and owned-process-group cleanup.
|
| 2 |
|
| 3 | This controls one POSIX process group, not an OS sandbox. Descendants that
|
| 4 | create independent sessions can escape the group; non-POSIX cleanup only owns
|
| 5 | the direct child. Never describe this helper as process-tree containment.
|
| 6 | """
|
| 7 | from __future__ import annotations
|
| 8 |
|
| 9 | from dataclasses import dataclass
|
| 10 | import math
|
| 11 | import os
|
| 12 | from pathlib import Path
|
| 13 | import subprocess
|
| 14 | import threading
|
| 15 | import time
|
| 16 |
|
| 17 | from oakbench.execution import stop_owned_process
|
| 18 |
|
| 19 | EXECUTION_VERSION = "owned-group-bounded-output-v1"
|
| 20 | DEFAULT_OUTPUT_LIMIT_BYTES = 8 * 1024 * 1024
|
| 21 |
|
| 22 |
|
| 23 | @dataclass(frozen=True)
|
| 24 | class CapturedProcess:
|
| 25 | returncode: int
|
| 26 | timed_out: bool
|
| 27 | output_limited: bool
|
| 28 | elapsed_ms: float
|
| 29 | stdout_observed_bytes: int
|
| 30 | stderr_observed_bytes: int
|
| 31 | stdout_retained_bytes: int
|
| 32 | stderr_retained_bytes: int
|
| 33 | capture_complete: bool
|
| 34 |
|
| 35 |
|
| 36 | def run_captured(command: list[str], cwd: Path, env: dict[str, str],
|
| 37 | stdin_text: str | None, stdout_path: Path, stderr_path: Path,
|
| 38 | timeout_seconds: float,
|
| 39 | output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES) -> CapturedProcess:
|
| 40 | if not math.isfinite(timeout_seconds) or timeout_seconds <= 0:
|
| 41 | raise ValueError("timeout_seconds must be positive and finite")
|
| 42 | if isinstance(output_limit_bytes, bool) or not isinstance(output_limit_bytes, int) or output_limit_bytes < 1:
|
| 43 | raise ValueError("output_limit_bytes must be a positive integer")
|
| 44 | started = time.perf_counter()
|
| 45 | counts = {"stdout": 0, "stderr": 0}
|
| 46 | retained = {"stdout": 0, "stderr": 0}
|
| 47 | count_lock = threading.Lock()
|
| 48 | exceeded = threading.Event()
|
| 49 | capture_errors = []
|
| 50 | readers = []
|
| 51 | writers = []
|
| 52 | proc = None
|
| 53 | timed_out = False
|
| 54 | forced_pipe_closure = False
|
| 55 | with stdout_path.open("wb") as stdout_file, stderr_path.open("wb") as stderr_file:
|
| 56 | try:
|
| 57 | proc = subprocess.Popen(command, cwd=cwd, env=env,
|
| 58 | stdin=subprocess.PIPE if stdin_text is not None else subprocess.DEVNULL,
|
| 59 | stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
| 60 | bufsize=0, start_new_session=os.name == "posix")
|
| 61 |
|
| 62 | def drain(stream, destination, name):
|
| 63 | try:
|
| 64 | while True:
|
| 65 | data = stream.read(65536)
|
| 66 | if not data:
|
| 67 | break
|
| 68 | with count_lock:
|
| 69 | counts[name] += len(data)
|
| 70 | allowance = max(0, output_limit_bytes - sum(retained.values()))
|
| 71 | admitted = data[:allowance]
|
| 72 | destination.write(admitted)
|
| 73 | retained[name] += len(admitted)
|
| 74 | if len(admitted) != len(data):
|
| 75 | exceeded.set()
|
| 76 | except (OSError, ValueError) as exc:
|
| 77 | capture_errors.append(str(exc))
|
| 78 | finally:
|
| 79 | stream.close()
|
| 80 |
|
| 81 | for name, stream, destination in (("stdout", proc.stdout, stdout_file),
|
| 82 | ("stderr", proc.stderr, stderr_file)):
|
| 83 | thread = threading.Thread(target=drain, args=(stream, destination, name), daemon=True)
|
| 84 | readers.append(thread)
|
| 85 | thread.start()
|
| 86 | if stdin_text is not None:
|
| 87 | def supply_input():
|
| 88 | try:
|
| 89 | pending = memoryview(stdin_text.encode("utf-8"))
|
| 90 | while pending:
|
| 91 | written = proc.stdin.write(pending)
|
| 92 | if not written:
|
| 93 | break
|
| 94 | pending = pending[written:]
|
| 95 | except (BrokenPipeError, OSError, ValueError):
|
| 96 | pass
|
| 97 | finally:
|
| 98 | proc.stdin.close()
|
| 99 | writer = threading.Thread(target=supply_input, daemon=True)
|
| 100 | writers.append(writer)
|
| 101 | writer.start()
|
| 102 | deadline = started + timeout_seconds
|
| 103 | while proc.poll() is None:
|
| 104 | if exceeded.is_set():
|
| 105 | break
|
| 106 | remaining = deadline - time.perf_counter()
|
| 107 | if remaining <= 0:
|
| 108 | timed_out = True
|
| 109 | break
|
| 110 | try:
|
| 111 | proc.wait(timeout=min(0.025, remaining))
|
| 112 | except subprocess.TimeoutExpired:
|
| 113 | pass
|
| 114 | finally:
|
| 115 | # Even an exit-zero CLI may leave tools running or pipe handles
|
| 116 | # open. Stop the owned group before oracle inspection on every path,
|
| 117 | # including KeyboardInterrupt and reader/launcher exceptions.
|
| 118 | if proc is not None:
|
| 119 | stop_owned_process(proc)
|
| 120 | for thread in readers + writers:
|
| 121 | thread.join(timeout=1.0)
|
| 122 | forced_pipe_closure = any(thread.is_alive() for thread in readers)
|
| 123 | if proc is not None:
|
| 124 | for stream in (proc.stdout, proc.stderr, proc.stdin):
|
| 125 | if stream is not None and not stream.closed:
|
| 126 | stream.close()
|
| 127 | complete = not capture_errors and not forced_pipe_closure
|
| 128 | output_limited = exceeded.is_set()
|
| 129 | returncode = 124 if timed_out else (125 if output_limited or not complete else proc.returncode)
|
| 130 | return CapturedProcess(returncode, timed_out, output_limited,
|
| 131 | (time.perf_counter() - started) * 1000,
|
| 132 | counts["stdout"], counts["stderr"],
|
| 133 | retained["stdout"], retained["stderr"], complete)
|