fix(exec): return early when session command exits

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

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()