mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-07 21:08:34 +03:00
fix(exec): bound buffered session output
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user