From 5e67fbf93e3b46927c7edc7ba07b1838c6c8fb58 Mon Sep 17 00:00:00 2001 From: yu-xin-c <2182712990@qq.com> Date: Wed, 29 Jul 2026 01:03:29 +0800 Subject: [PATCH] fix(exec): bound buffered session output --- nanobot/agent/tools/exec_session.py | 119 +++++++++++++++++++------ tests/tools/test_exec_session_tools.py | 84 +++++++++++++++++ 2 files changed, 176 insertions(+), 27 deletions(-) diff --git a/nanobot/agent/tools/exec_session.py b/nanobot/agent/tools/exec_session.py index 1245edfdc..c063bdbaf 100644 --- a/nanobot/agent/tools/exec_session.py +++ b/nanobot/agent/tools/exec_session.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import time import uuid +from collections import deque from contextlib import suppress from dataclasses import dataclass from typing import Any @@ -51,6 +52,66 @@ class ExecSessionInfo: owner_session_key: str | None = None +class _BoundedOutputBuffer: + """Keep the first and most recent characters within a fixed budget.""" + + def __init__(self, max_chars: int) -> None: + self.max_chars = max_chars + self._content = "" + self._tail: deque[str] = deque() + self._tail_chars = 0 + self._total_chars = 0 + self._truncated = False + + @property + def has_output(self) -> bool: + return self._total_chars > 0 + + @property + def retained_chars(self) -> int: + return len(self._content) + self._tail_chars + + def append(self, text: str) -> None: + if not text: + return + self._total_chars += len(text) + if not self._truncated: + combined = self._content + text + if len(combined) <= self.max_chars: + self._content = combined + return + head_chars = self.max_chars // 2 + tail_chars = self.max_chars - head_chars + self._content = combined[:head_chars] + self._tail.append(combined[-tail_chars:]) + self._tail_chars = tail_chars + self._truncated = True + return + + tail_chars = self.max_chars - len(self._content) + self._tail.append(text) + self._tail_chars += len(text) + while self._tail_chars > tail_chars: + excess = self._tail_chars - tail_chars + first = self._tail[0] + if len(first) <= excess: + self._tail.popleft() + self._tail_chars -= len(first) + else: + self._tail[0] = first[excess:] + self._tail_chars -= excess + + def drain(self) -> tuple[str, int]: + output = self._content + "".join(self._tail) + truncated_chars = self._total_chars - len(output) + self._content = "" + self._tail.clear() + self._tail_chars = 0 + self._total_chars = 0 + self._truncated = False + return output, truncated_chars + + class _ExecSession: def __init__( self, @@ -73,30 +134,27 @@ class _ExecSession: # timeout None/0 means no limit; an infinite deadline is never reached. self.deadline = time.monotonic() + timeout if timeout else float("inf") self.last_access = time.monotonic() - self._chunks: list[str] = [] + self._stdout = _BoundedOutputBuffer(MAX_OUTPUT_CHARS) + self._stderr = _BoundedOutputBuffer(MAX_OUTPUT_CHARS) self._lock = asyncio.Lock() self._timed_out = False - self._stdout_task = asyncio.create_task(self._read_stream(process.stdout, "")) - self._stderr_task = asyncio.create_task(self._read_stream(process.stderr, "STDERR:\n")) + self._stdout_task = asyncio.create_task(self._read_stream(process.stdout, self._stdout)) + self._stderr_task = asyncio.create_task(self._read_stream(process.stderr, self._stderr)) async def _read_stream( self, stream: asyncio.StreamReader | None, - prefix: str, + buffer: _BoundedOutputBuffer, ) -> None: if stream is None: return - first = True while True: chunk = await stream.read(4096) if not chunk: break text = chunk.decode("utf-8", errors="replace") - if prefix and first: - text = prefix + text - first = False async with self._lock: - self._chunks.append(text) + buffer.append(text) async def write(self, chars: str) -> str | None: if self.process.returncode is not None: @@ -157,10 +215,14 @@ class _ExecSession: await self._wait_for_buffered_output() async with self._lock: - output = "".join(self._chunks) - self._chunks.clear() + stdout, stdout_truncated = self._stdout.drain() + stderr, stderr_truncated = self._stderr.drain() - output, truncated = _truncate_output(output, max_output_chars) + output_parts = [stdout] if stdout else [] + if stderr: + output_parts.append(f"STDERR:\n{stderr}") + output = "\n".join(output_parts) + output, response_truncated = _truncate_output(output, max_output_chars) return _SessionPoll( output=output, done=self.process.returncode is not None, @@ -169,7 +231,7 @@ class _ExecSession: timed_out=self._timed_out, terminated=terminated, stdin_closed=stdin_closed, - truncated_chars=truncated, + truncated_chars=stdout_truncated + stderr_truncated + response_truncated, ) async def kill(self) -> None: @@ -195,7 +257,7 @@ class _ExecSession: deadline = time.monotonic() + OUTPUT_DRAIN_GRACE_S while time.monotonic() < deadline: async with self._lock: - if self._chunks: + if self._stdout.has_output or self._stderr.has_output: return await asyncio.sleep(0.01) @@ -403,20 +465,16 @@ def clamp_session_int(value: int | None, default: int, minimum: int, maximum: in def _truncate_output(output: str, max_output_chars: int) -> tuple[str, int]: if len(output) <= max_output_chars: return output, 0 - half = max_output_chars // 2 + head_chars = max_output_chars // 2 + tail_chars = max_output_chars - head_chars omitted = len(output) - max_output_chars - return ( - output[:half] - + f"\n\n... ({omitted:,} chars truncated) ...\n\n" - + output[-half:], - omitted, - ) + return output[:head_chars] + output[-tail_chars:], omitted def format_session_poll(session_id: str, poll: _SessionPoll) -> str: parts = [poll.output] if poll.output else [] if poll.truncated_chars: - parts.append(f"(output truncated by {poll.truncated_chars:,} chars)") + parts.append(f"({poll.truncated_chars:,} chars truncated from output)") if poll.timed_out: parts.append("Error: Command timed out; session was terminated.") if poll.terminated and not poll.timed_out: @@ -587,7 +645,9 @@ class WriteStdinTool(Tool): max_output_chars: int, ) -> str: deadline = time.monotonic() + (wait_timeout_ms / 1000) - aggregate: list[str] = [] + aggregate = _BoundedOutputBuffer(max_output_chars) + upstream_truncated = 0 + search_overlap = "" first = True poll: _SessionPoll | None = None @@ -604,15 +664,20 @@ class WriteStdinTool(Tool): owner_session_key=current_request_session_key(), ) first = False + upstream_truncated += poll.truncated_chars if poll.output: aggregate.append(poll.output) - joined = "".join(aggregate) - if wait_for in joined: - poll.output = joined + searchable = search_overlap + poll.output + if wait_for in searchable: + poll.output, aggregate_truncated = aggregate.drain() + poll.truncated_chars = upstream_truncated + aggregate_truncated result = format_session_poll(session_id, poll) return ToolResult.error(result) if poll.timed_out else result + overlap_chars = max(0, len(wait_for) - 1) + search_overlap = searchable[-overlap_chars:] if overlap_chars else "" if poll.done or remaining_ms <= 0: - poll.output = "".join(aggregate) + poll.output, aggregate_truncated = aggregate.drain() + poll.truncated_chars = upstream_truncated + aggregate_truncated result = format_session_poll(session_id, poll) if wait_for not in poll.output: result += f"\nWait target not observed: {wait_for!r}" diff --git a/tests/tools/test_exec_session_tools.py b/tests/tools/test_exec_session_tools.py index 138389c3d..77195e52f 100644 --- a/tests/tools/test_exec_session_tools.py +++ b/tests/tools/test_exec_session_tools.py @@ -19,6 +19,8 @@ from nanobot.agent.tools.exec_session import ( ExecSessionManager, ListExecSessionsTool, WriteStdinTool, + _BoundedOutputBuffer, + _SessionPoll, ) from nanobot.agent.tools.registry import is_tool_error_result from nanobot.agent.tools.shell import ExecTool @@ -143,6 +145,88 @@ def test_exec_session_accepts_max_output_tokens_alias(tmp_path): assert "Exit code: 0" in result +def test_bounded_output_buffer_keeps_head_tail_and_exact_drop_count(): + buffer = _BoundedOutputBuffer(10) + + buffer.append("012345") + buffer.append("6789ABCDEF") + + assert buffer.retained_chars == 10 + assert buffer.drain() == ("01234BCDEF", 6) + assert buffer.retained_chars == 0 + + +def test_exec_session_bounds_unpolled_stdout_and_stderr(tmp_path): + async def run() -> tuple[int, int, str, int]: + manager = ExecSessionManager() + tool = ExecTool(working_dir=str(tmp_path), timeout=5, session_manager=manager) + command = _python_command( + "import sys,time; time.sleep(0.05); " + "sys.stdout.write('OUT_HEAD' + 'o' * 200000 + 'OUT_TAIL'); " + "sys.stderr.write('ERR_HEAD' + 'e' * 200000 + 'ERR_TAIL')" + ) + + initial = await tool.execute( + command=command, + yield_time_ms=0, + max_output_chars=1000, + ) + sid = _session_id(initial) + session = manager._sessions[sid] + await asyncio.wait_for(session.process.wait(), timeout=5) + await asyncio.wait_for( + asyncio.gather(session._stdout_task, session._stderr_task), + timeout=5, + ) + retained_stdout = session._stdout.retained_chars + retained_stderr = session._stderr.retained_chars + poll = await manager.write( + session_id=sid, + chars=None, + close_stdin=False, + terminate=False, + yield_time_ms=0, + max_output_chars=1000, + ) + return retained_stdout, retained_stderr, poll.output, poll.truncated_chars + + retained_stdout, retained_stderr, output, truncated_chars = asyncio.run(run()) + + assert retained_stdout == 50000 + assert retained_stderr == 50000 + assert output.startswith("OUT_HEAD") + assert output.endswith("ERR_TAIL") + assert truncated_chars > 390000 + + +def test_write_stdin_wait_for_keeps_aggregate_within_output_budget(): + async def run() -> str: + manager = SimpleNamespace( + write=AsyncMock(side_effect=[ + _SessionPoll(output="HEAD" + "a" * 596, done=False, exit_code=None), + _SessionPoll(output="b" * 600, done=False, exit_code=None), + _SessionPoll(output="c" * 590 + "TARGET", done=False, exit_code=None), + ]) + ) + tool = WriteStdinTool(manager=manager) + return await tool._wait_for_output( + session_id="session", + chars=None, + close_stdin=False, + terminate=False, + wait_for="TARGET", + wait_timeout_ms=1000, + max_output_chars=1000, + ) + + result = asyncio.run(run()) + + assert result.startswith("HEAD") + assert "TARGET" in result + assert "(796 chars truncated from output)" in result + assert len(result) < 1100 + + def test_exec_one_shot_accepts_max_output_tokens_alias(tmp_path): async def run() -> str: tool = ExecTool(working_dir=str(tmp_path), timeout=5)