diff --git a/nanobot/agent/tools/exec_session.py b/nanobot/agent/tools/exec_session.py index e4fdc17b..6a38d30b 100644 --- a/nanobot/agent/tools/exec_session.py +++ b/nanobot/agent/tools/exec_session.py @@ -128,7 +128,15 @@ class _ExecSession: ) -> _SessionPoll: self.last_access = time.monotonic() if yield_time_ms > 0 and self.process.returncode is None: - await asyncio.sleep(min(yield_time_ms, MAX_YIELD_MS) / 1000) + wait_s = min(yield_time_ms, MAX_YIELD_MS) / 1000 + remaining_s = self.deadline - time.monotonic() + if remaining_s <= 0: + wait_s = 0 + else: + wait_s = min(wait_s, remaining_s) + if wait_s > 0: + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(self.process.wait(), timeout=wait_s) if self.process.returncode is None and time.monotonic() >= self.deadline: self._timed_out = True diff --git a/tests/tools/test_exec_session_tools.py b/tests/tools/test_exec_session_tools.py index e120b236..f1fc3eed 100644 --- a/tests/tools/test_exec_session_tools.py +++ b/tests/tools/test_exec_session_tools.py @@ -5,6 +5,7 @@ import re import shlex import subprocess import sys +import time from nanobot.agent.tools.exec_session import ( ExecSessionManager, @@ -76,6 +77,23 @@ def test_exec_returns_completed_session_output_when_yield_time_ms_is_used(tmp_pa assert "session_id:" not in result +def test_exec_session_yield_returns_when_process_finishes_early(tmp_path): + async def run() -> tuple[str, float]: + manager = ExecSessionManager() + tool = ExecTool(working_dir=str(tmp_path), timeout=5, session_manager=manager) + command = _python_command("import time; time.sleep(0.1); print('done')") + started = time.monotonic() + result = await tool.execute(command=command, yield_time_ms=1200) + return result, time.monotonic() - started + + result, elapsed = asyncio.run(run()) + + assert "done" in result + assert "Exit code: 0" in result + assert "session_id:" not in result + assert elapsed < 1.0 + + def test_exec_session_accepts_max_output_tokens_alias(tmp_path): async def run() -> str: manager = ExecSessionManager()