diff --git a/sdk/python/agentfield/harness/_cli.py b/sdk/python/agentfield/harness/_cli.py index 7fc63d548..9b03197d3 100644 --- a/sdk/python/agentfield/harness/_cli.py +++ b/sdk/python/agentfield/harness/_cli.py @@ -4,15 +4,19 @@ import asyncio import json +import logging import os import re import shutil import signal import subprocess -from typing import Any, Dict, List, Optional, Tuple +from collections import deque +from typing import Any, Deque, Dict, List, Optional, Tuple from agentfield.openrouter_attribution import apply_subprocess_env +logger = logging.getLogger("agentfield.harness.cli") + _ANSI_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]") # 300 rather than 120: provider CLIs in JSON mode (e.g. `opencode run @@ -100,9 +104,86 @@ def _resolve_idle_seconds(idle_seconds: Optional[float]) -> Optional[float]: return idle_seconds if idle_seconds and idle_seconds > 0 else None +# Cap on captured bytes PER STREAM (stdout and stderr each). Providers hold +# the joined text plus its parsed JSONL events in memory, so an unbounded +# child stream is buffered several times over — and N concurrent harness +# calls multiply that (the pr-af#65 OOM). Real provider streams are +# completion-boundary events (hundreds of KB), so 16MB is ~50x headroom +# while bounding the pathological case. <= 0 disables the cap. +_DEFAULT_MAX_OUTPUT_BYTES = 16 * 1024 * 1024 +# Keep-head fraction on overflow: the head carries session/model info and the +# first error, but the FINAL events (result text, cumulative usage) live at +# the tail, so the tail gets most of the budget. +_TRUNCATION_HEAD_FRACTION = 0.25 + + +def _resolve_max_output_bytes() -> int: + raw = os.environ.get("AGENTFIELD_HARNESS_MAX_OUTPUT_BYTES") + if raw is not None: + try: + return int(raw) + except ValueError: + return _DEFAULT_MAX_OUTPUT_BYTES + return _DEFAULT_MAX_OUTPUT_BYTES + + +class _BoundedChunks: + """Chunk accumulator with a byte cap that keeps the stream's head and tail. + + Below the cap, ``joined()`` is byte-identical to the stream. Above it, the + first ``head`` budget bytes and the most recent bytes within the remaining + budget are kept, spliced around a marker line. The marker (and any partial + line at the seam) parses as invalid JSON, which ``parse_jsonl`` skips — so + the final result/usage events at the tail stay extractable. + """ + + def __init__(self, max_bytes: int) -> None: + self._max = max_bytes + self._head_budget = ( + int(max_bytes * _TRUNCATION_HEAD_FRACTION) if max_bytes > 0 else 0 + ) + self._tail_budget = max(0, max_bytes - self._head_budget) + self._head: List[bytes] = [] + self._head_bytes = 0 + self._tail: Deque[bytes] = deque() + self._tail_bytes = 0 + self.dropped_bytes = 0 + + def append(self, chunk: bytes) -> None: + if self._max <= 0: + self._head.append(chunk) + return + if self._head_bytes < self._head_budget: + take = min(len(chunk), self._head_budget - self._head_bytes) + self._head.append(chunk[:take]) + self._head_bytes += take + chunk = chunk[take:] + if not chunk: + return + self._tail.append(chunk) + self._tail_bytes += len(chunk) + # Always retain at least the newest chunk, even if it alone exceeds + # the tail budget — the newest bytes are the ones that matter. + while self._tail_bytes > self._tail_budget and len(self._tail) > 1: + evicted = self._tail.popleft() + self._tail_bytes -= len(evicted) + self.dropped_bytes += len(evicted) + + def joined(self) -> bytes: + head = b"".join(self._head) + tail = b"".join(self._tail) + if not self.dropped_bytes: + return head + tail + marker = ( + b"\n[agentfield: output truncated, %d bytes dropped]\n" + % self.dropped_bytes + ) + return head + marker + tail + + async def _drain( stream: Optional[asyncio.StreamReader], - chunks: List[bytes], + chunks: _BoundedChunks, last_activity: List[float], ) -> None: """Read a stream incrementally, recording each chunk and its arrival time.""" @@ -200,6 +281,11 @@ async def run_cli( without it stdin is /dev/null. Providers use this to hand over prompts too large for a command line — on Windows an npm ``.cmd`` shim runs via cmd.exe, which caps the command line at ~8k characters. + + Captured output is bounded per stream (env + ``AGENTFIELD_HARNESS_MAX_OUTPUT_BYTES``, default 16MB; <= 0 disables): + on overflow the head and tail are kept around a truncation marker, so + the final JSONL result/usage events remain parseable. """ merged_env = {**os.environ} if env: @@ -221,8 +307,9 @@ async def run_cli( start_new_session=True, ) - stdout_chunks: List[bytes] = [] - stderr_chunks: List[bytes] = [] + max_output_bytes = _resolve_max_output_bytes() + stdout_chunks = _BoundedChunks(max_output_bytes) + stderr_chunks = _BoundedChunks(max_output_bytes) last_activity = [asyncio.get_event_loop().time()] # Pump all pipes concurrently to avoid a pipe-buffer deadlock (a large @@ -317,9 +404,19 @@ def _kill_group() -> None: returncode = proc.returncode if returncode is None: returncode = fallback_returncode + for name, buf in (("stdout", stdout_chunks), ("stderr", stderr_chunks)): + if buf.dropped_bytes: + logger.warning( + "CLI %s exceeded the %d-byte capture cap; dropped %d bytes " + "from the middle of the stream: %s", + name, + max_output_bytes, + buf.dropped_bytes, + " ".join(cmd[:3]), + ) return ( - b"".join(stdout_chunks).decode("utf-8", errors="replace"), - b"".join(stderr_chunks).decode("utf-8", errors="replace"), + stdout_chunks.joined().decode("utf-8", errors="replace"), + stderr_chunks.joined().decode("utf-8", errors="replace"), returncode if returncode is not None else -1, ) diff --git a/sdk/python/tests/test_harness_cli.py b/sdk/python/tests/test_harness_cli.py index 6d3b12149..4b8acc3b1 100644 --- a/sdk/python/tests/test_harness_cli.py +++ b/sdk/python/tests/test_harness_cli.py @@ -338,3 +338,90 @@ async def test_run_cli_feeds_input_text_via_stdin(): b"a prompt too large for a cmd.exe command line" ) process.stdin.close.assert_called_once() + + +# --------------------------------------------------------------------------- +# Bounded output capture (AGENTFIELD_HARNESS_MAX_OUTPUT_BYTES) +# --------------------------------------------------------------------------- + + +def test_bounded_chunks_is_byte_identical_below_cap(): + from agentfield.harness._cli import _BoundedChunks + + buf = _BoundedChunks(1024) + buf.append(b"hello ") + buf.append(b"world") + assert buf.joined() == b"hello world" + assert buf.dropped_bytes == 0 + + +def test_bounded_chunks_keeps_head_and_tail_above_cap(): + from agentfield.harness._cli import _BoundedChunks + + buf = _BoundedChunks(1000) + for i in range(100): + buf.append(f"chunk-{i:04d}-".encode() + b"x" * 90) + joined = buf.joined() + assert buf.dropped_bytes > 0 + assert joined.startswith(b"chunk-0000-") + assert joined.rstrip().endswith(b"x") + assert b"chunk-0099-" in joined + assert b"[agentfield: output truncated," in joined + assert len(joined) <= 1000 + 200 # cap + marker + one-chunk slack + + +def test_bounded_chunks_retains_newest_chunk_even_when_oversized(): + from agentfield.harness._cli import _BoundedChunks + + buf = _BoundedChunks(100) + buf.append(b"a" * 500) + buf.append(b"FINAL" * 100) + assert b"FINAL" in buf.joined() + + +def test_bounded_chunks_cap_disabled_when_nonpositive(): + from agentfield.harness._cli import _BoundedChunks + + buf = _BoundedChunks(0) + payload = b"y" * 100_000 + buf.append(payload) + buf.append(payload) + assert buf.joined() == payload + payload + assert buf.dropped_bytes == 0 + + +@pytest.mark.asyncio +async def test_run_cli_truncation_preserves_final_jsonl_events(monkeypatch): + monkeypatch.setenv("AGENTFIELD_HARNESS_MAX_OUTPUT_BYTES", "4096") + filler = [(b'{"type":"text","text":"' + b"x" * 200 + b'"}\n') for _ in range(200)] + final = b'{"type":"result","result":"the final answer"}\n' + process = MagicMock() + process.stdout = _stream_reader(filler + [final]) + process.stderr = _stream_reader([]) + process.returncode = 0 + process.wait = AsyncMock(return_value=0) + + with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=process)): + stdout, stderr, returncode = await run_cli(["opencode", "run"], timeout=5) + + assert returncode == 0 + assert "[agentfield: output truncated," in stdout + events = parse_jsonl(stdout) + assert {"type": "result", "result": "the final answer"} in events + assert extract_final_text(events) == "the final answer" + + +@pytest.mark.asyncio +async def test_run_cli_no_truncation_below_cap(monkeypatch): + monkeypatch.delenv("AGENTFIELD_HARNESS_MAX_OUTPUT_BYTES", raising=False) + payload = b'{"type":"result","result":"ok"}\n' + process = MagicMock() + process.stdout = _stream_reader([payload]) + process.stderr = _stream_reader([]) + process.returncode = 0 + process.wait = AsyncMock(return_value=0) + + with patch("asyncio.create_subprocess_exec", AsyncMock(return_value=process)): + stdout, _, _ = await run_cli(["opencode", "run"], timeout=5) + + assert stdout == payload.decode()