fix(agent): defer subagent wait until turn exit

nanobot/agent/loop.py

@@ -112,6 +112,7 @@ if TYPE_CHECKING:
 
 _T = TypeVar("_T")
 _SUBAGENT_PROVIDER_TASK_META = "subagent_provider_task_id"
+_SUBAGENT_TERMINAL_WAIT_SECONDS = 300.0
 
 
 class TurnKind(Enum):
@@ -1000,15 +1001,12 @@ class AgentLoop:
                 )
             self._set_runtime_checkpoint(session, public_payload)
 
-        async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
-            """Drain follow-up messages from the pending queue.
-
-            When no messages are immediately available but sub-agents
-            spawned in this dispatch are still running, blocks until at
-            least one result arrives (or timeout).  This keeps the runner
-            loop alive so subsequent sub-agent completions are consumed
-            in-order rather than dispatched separately.
-            """
+        async def _drain_pending(
+            *,
+            limit: int = _MAX_INJECTIONS_PER_TURN,
+            first_msg: InboundMessage | None = None,
+        ) -> list[dict[str, Any]]:
+            """Drain only messages that are already available."""
             if pending_queue is None:
                 return []
 
@@ -1080,35 +1078,52 @@ class AgentLoop:
                 return row
 
             items: list[dict[str, Any]] = []
+            if first_msg is not None:
+                items.append(await _to_user_message(first_msg))
             while len(items) < limit:
                 try:
                     items.append(await _to_user_message(pending_queue.get_nowait()))
                 except asyncio.QueueEmpty:
                     break
 
-            # Block if nothing drained but sub-agents spawned in this dispatch
-            # are still running.  Keeps the runner loop alive so subsequent
-            # completions are injected in-order rather than dispatched separately.
-            if (not items
-                    and session is not None
-                    and self.subagents.get_running_count_by_session(session.key) > 0):
-                try:
-                    msg = await asyncio.wait_for(pending_queue.get(), timeout=300)
-                except asyncio.TimeoutError:
-                    logger.warning(
-                        "Timeout waiting for sub-agent completion in session {}",
-                        session.key,
-                    )
-                    return items
-                items.append(await _to_user_message(msg))
-                while len(items) < limit:
-                    try:
-                        items.append(await _to_user_message(pending_queue.get_nowait()))
-                    except asyncio.QueueEmpty:
-                        break
-
             return items
 
+        terminal_wait_deadline: float | None = None
+
+        async def _wait_for_pending(
+            *,
+            limit: int = _MAX_INJECTIONS_PER_TURN,
+        ) -> list[dict[str, Any]]:
+            """Wait for a pending result only when the runner is ready to exit."""
+            nonlocal terminal_wait_deadline
+
+            items = await _drain_pending(limit=limit)
+            if (
+                items
+                or pending_queue is None
+                or session is None
+                or self.subagents.get_running_count_by_session(session.key) == 0
+            ):
+                return items
+
+            now = asyncio.get_running_loop().time()
+            if terminal_wait_deadline is None:
+                terminal_wait_deadline = now + _SUBAGENT_TERMINAL_WAIT_SECONDS
+            remaining = terminal_wait_deadline - now
+            if remaining <= 0:
+                return []
+
+            try:
+                msg = await asyncio.wait_for(pending_queue.get(), timeout=remaining)
+            except asyncio.TimeoutError:
+                logger.warning(
+                    "Timeout waiting for sub-agent completion before session {} exits",
+                    session.key,
+                )
+                return []
+
+            return await _drain_pending(limit=limit, first_msg=msg)
+
         active_session_key = session.key if session else session_key
         effective_scope = self.workspace_scopes.for_turn(
             channel=channel,
@@ -1184,6 +1199,7 @@ class AgentLoop:
                 retry_wait_callback=on_retry_wait,
                 checkpoint_callback=_checkpoint,
                 injection_callback=_drain_pending,
+                terminal_injection_callback=_wait_for_pending,
                 # Sustained goals may legitimately exceed NANOBOT_LLM_TIMEOUT_S; idle stall
                 # is still capped by NANOBOT_STREAM_IDLE_TIMEOUT_S in streaming providers.
                 llm_timeout_s=runner_wall_llm_timeout_s(

nanobot/agent/runner.py

@@ -118,6 +118,7 @@ class AgentRunSpec:
     retry_wait_callback: RetryWaitCallback | None = None
     checkpoint_callback: CheckpointCallback | None = None
     injection_callback: InjectionCallback | None = None
+    terminal_injection_callback: InjectionCallback | None = None
     llm_timeout_s: float | None = None
     goal_active_predicate: Callable[[], bool] | None = None
     goal_continue_message: GoalContinueMessage | None = None
@@ -274,6 +275,7 @@ class AgentRunner:
         phase: str = "after error",
         iteration: int | None = None,
         allow_goal_continue: bool = False,
+        wait_at_terminal: bool = False,
     ) -> tuple[bool, int]:
         """Drain pending injections. Returns (should_continue, updated_cycles).
 
@@ -291,6 +293,13 @@ class AgentRunner:
             predicate = spec.goal_active_predicate
             if predicate is not None and predicate():
                 injections = [self._build_goal_continue_message(spec)]
+        if (
+            not injections
+            and wait_at_terminal
+            and injection_cycles < _MAX_INJECTION_CYCLES
+        ):
+            injections = await self._drain_injections(spec, terminal=True)
+            real_injection = bool(injections)
         if not injections:
             return False, injection_cycles
         if real_injection:
@@ -334,7 +343,12 @@ class AgentRunner:
                 custom = None
         return build_goal_continue_message(custom)
 
-    async def _drain_injections(self, spec: AgentRunSpec) -> list[dict[str, Any]]:
+    async def _drain_injections(
+        self,
+        spec: AgentRunSpec,
+        *,
+        terminal: bool = False,
+    ) -> list[dict[str, Any]]:
         """Drain pending user messages via the injection callback.
 
         Returns normalized user messages (capped by
@@ -342,10 +356,15 @@ class AgentRunner:
         nothing to inject. Messages beyond the cap are logged so they
         are not silently lost.
         """
-        if spec.injection_callback is None:
+        callback = (
+            spec.terminal_injection_callback
+            if terminal
+            else spec.injection_callback
+        )
+        if callback is None:
             return []
         try:
-            signature = inspect.signature(spec.injection_callback)
+            signature = inspect.signature(callback)
             accepts_limit = (
                 "limit" in signature.parameters
                 or any(
@@ -354,9 +373,9 @@ class AgentRunner:
                 )
             )
             if accepts_limit:
-                items = await spec.injection_callback(limit=_MAX_INJECTIONS_PER_TURN)
+                items = await callback(limit=_MAX_INJECTIONS_PER_TURN)
             else:
-                items = await spec.injection_callback()
+                items = await callback()
         except Exception:
             logger.exception("injection_callback failed")
             return []
@@ -753,6 +772,11 @@ class AgentRunner:
                 allow_goal_continue=(
                     response.finish_reason not in {"refusal", "content_filter"}
                 ),
+                wait_at_terminal=(
+                    assistant_message is not None
+                    and response.finish_reason
+                    not in {"error", "length", "refusal", "content_filter"}
+                ),
             )
             if should_continue:
                 had_injections = True

tests/agent/test_runner_core.py

@@ -985,6 +985,7 @@ async def test_runner_does_not_auto_continue_goal_after_policy_terminal(
     ))
     tools = MagicMock()
     tools.get_definitions.return_value = []
+    terminal_injection_callback = AsyncMock(return_value=[])
 
     result = await AgentRunner().run(make_run_spec(
         provider,
@@ -994,9 +995,11 @@ async def test_runner_does_not_auto_continue_goal_after_policy_terminal(
         max_iterations=3,
         max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
         goal_active_predicate=lambda: True,
+        terminal_injection_callback=terminal_injection_callback,
     ))
 
     assert provider.chat_with_retry.await_count == 1
+    terminal_injection_callback.assert_not_awaited()
     assert result.final_content == "Request blocked by provider policy."
     assert result.stop_reason == "completed"
 

tests/agent/test_runner_injections.py

@@ -299,6 +299,130 @@ async def test_checkpoint1_injects_after_tool_execution():
     assert len(injected) == 1
 
 
+@pytest.mark.asyncio
+async def test_terminal_wait_does_not_block_next_iteration_after_tools():
+    """Background waits begin only after a no-tool response is ready to finish."""
+    from nanobot.agent.runner import AgentRunner
+    from nanobot.bus.events import InboundMessage
+
+    provider = MagicMock()
+    second_request_started = asyncio.Event()
+    allow_final_response = asyncio.Event()
+    terminal_wait_started = asyncio.Event()
+    release_terminal_result = asyncio.Event()
+    call_count = 0
+    terminal_result_delivered = False
+
+    async def chat_with_retry(*, messages, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return LLMResponse(
+                content="working",
+                tool_calls=[ToolCallRequest(id="c1", name="read_file", arguments={"path": "x"})],
+            )
+        if call_count == 2:
+            second_request_started.set()
+            await allow_final_response.wait()
+            return LLMResponse(content="main work finished", tool_calls=[])
+        return LLMResponse(content="combined final answer", tool_calls=[])
+
+    async def drain_available():
+        return []
+
+    async def wait_at_terminal():
+        nonlocal terminal_result_delivered
+        if terminal_result_delivered:
+            return []
+        terminal_wait_started.set()
+        await release_terminal_result.wait()
+        terminal_result_delivered = True
+        return [
+            InboundMessage(
+                channel="system",
+                sender_id="subagent",
+                chat_id="c",
+                content="background result",
+            )
+        ]
+
+    provider.chat_with_retry = chat_with_retry
+    tools = MagicMock()
+    tools.get_definitions.return_value = []
+    tools.execute = AsyncMock(return_value="file content")
+
+    runner = AgentRunner()
+    run_task = asyncio.create_task(runner.run(make_run_spec(
+        provider,
+        initial_messages=[{"role": "user", "content": "hello"}],
+        tools=tools,
+        model="test-model",
+        max_iterations=5,
+        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        injection_callback=drain_available,
+        terminal_injection_callback=wait_at_terminal,
+    )))
+
+    await asyncio.wait_for(second_request_started.wait(), timeout=1.0)
+    assert not terminal_wait_started.is_set()
+
+    allow_final_response.set()
+    await asyncio.wait_for(terminal_wait_started.wait(), timeout=1.0)
+    assert not run_task.done()
+
+    release_terminal_result.set()
+    result = await asyncio.wait_for(run_task, timeout=1.0)
+
+    assert call_count == 3
+    assert result.had_injections is True
+    assert result.final_content == "combined final answer"
+
+
+@pytest.mark.asyncio
+async def test_goal_continuation_precedes_terminal_wait():
+    """An active sustained goal keeps running without joining background work."""
+    from nanobot.agent.runner import AgentRunner
+
+    provider = MagicMock()
+    provider.chat_with_retry = AsyncMock(side_effect=[
+        LLMResponse(content="goal checkpoint", tool_calls=[]),
+        LLMResponse(content="goal complete", tool_calls=[]),
+    ])
+    tools = MagicMock()
+    tools.get_definitions.return_value = []
+    goal_checks = 0
+    terminal_waits = 0
+
+    def goal_active() -> bool:
+        nonlocal goal_checks
+        goal_checks += 1
+        return goal_checks == 1
+
+    async def drain_available():
+        return []
+
+    async def wait_at_terminal():
+        nonlocal terminal_waits
+        terminal_waits += 1
+        return []
+
+    result = await AgentRunner().run(make_run_spec(
+        provider,
+        initial_messages=[{"role": "user", "content": "complete the goal"}],
+        tools=tools,
+        model="test-model",
+        max_iterations=3,
+        max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
+        injection_callback=drain_available,
+        terminal_injection_callback=wait_at_terminal,
+        goal_active_predicate=goal_active,
+    ))
+
+    assert provider.chat_with_retry.await_count == 2
+    assert terminal_waits == 1
+    assert result.final_content == "goal complete"
+
+
 @pytest.mark.asyncio
 async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
     """After final response, if injections exist, stream_end should get resuming=True."""

tests/agent/tools/test_subagent_tools.py

@@ -470,8 +470,8 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path):
 
 
 @pytest.mark.asyncio
-async def test_drain_pending_blocks_while_subagents_running(tmp_path):
-    """_drain_pending should block when no messages are available but sub-agents are still running."""
+async def test_drain_pending_waits_only_at_terminal_boundary(tmp_path):
+    """Ordinary drains stay non-blocking while terminal drains await subagent results."""
     from nanobot.agent.loop import AgentLoop
     from nanobot.bus.events import InboundMessage
     from nanobot.bus.queue import MessageBus
@@ -486,15 +486,14 @@ async def test_drain_pending_blocks_while_subagents_running(tmp_path):
     pending_queue: asyncio.Queue[InboundMessage] = asyncio.Queue()
     session = Session(key="test:drain-block")
     injection_callback = None
+    terminal_injection_callback = None
 
     # Capture the injection_callback that _run_agent_loop creates
     async def fake_runner_run(spec):
-        nonlocal injection_callback
+        nonlocal injection_callback, terminal_injection_callback
         injection_callback = spec.injection_callback
+        terminal_injection_callback = spec.terminal_injection_callback
 
-        # Simulate: first call to injection_callback should block because
-        # sub-agents are running and no messages are in the queue yet.
-        # We'll resolve this from a concurrent task.
         return SimpleNamespace(
             stop_reason="done",
             final_content="done",
@@ -528,16 +527,19 @@ async def test_drain_pending_blocks_while_subagents_running(tmp_path):
     )
 
     assert injection_callback is not None
+    assert terminal_injection_callback is not None
 
-    # Now test the callback directly
-    # With sub-agents running and an empty queue, it should block
-    drain_task = asyncio.create_task(injection_callback())
+    # Tool-boundary drains must let the runner start its next model iteration.
+    assert await asyncio.wait_for(injection_callback(), timeout=1.0) == []
+
+    # Once the runner is ready to exit, it may wait for a background result.
+    drain_task = asyncio.create_task(terminal_injection_callback())
 
     # Let the task enter the blocking queue wait.
     await asyncio.sleep(0)
 
     # Should still be running (blocked on pending_queue.get())
-    assert not drain_task.done(), "drain should block while sub-agents are running"
+    assert not drain_task.done(), "terminal drain should wait while subagents are running"
 
     # Now put a message in the queue (simulating sub-agent completion)
     await pending_queue.put(InboundMessage(
@@ -577,10 +579,12 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path):
 
     pending_queue: asyncio.Queue = asyncio.Queue()
     injection_callback = None
+    terminal_injection_callback = None
 
     async def fake_runner_run(spec):
-        nonlocal injection_callback
+        nonlocal injection_callback, terminal_injection_callback
         injection_callback = spec.injection_callback
+        terminal_injection_callback = spec.terminal_injection_callback
         return SimpleNamespace(
             stop_reason="done",
             final_content="done",
@@ -605,15 +609,16 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path):
     )
 
     assert injection_callback is not None
+    assert terminal_injection_callback is not None
 
-    # With no sub-agents and empty queue, should return immediately
-    results = await asyncio.wait_for(injection_callback(), timeout=1.0)
-    assert results == []
+    # With no sub-agents and an empty queue, both paths return immediately.
+    assert await asyncio.wait_for(injection_callback(), timeout=1.0) == []
+    assert await asyncio.wait_for(terminal_injection_callback(), timeout=1.0) == []
 
 
 @pytest.mark.asyncio
-async def test_drain_pending_timeout(tmp_path):
-    """_drain_pending should return empty after timeout when sub-agents hang."""
+async def test_terminal_drain_timeout(tmp_path):
+    """The terminal drain should return empty after its shared timeout expires."""
     from nanobot.agent.loop import AgentLoop
     from nanobot.bus.queue import MessageBus
     from nanobot.session.manager import Session
@@ -626,11 +631,11 @@ async def test_drain_pending_timeout(tmp_path):
 
     pending_queue: asyncio.Queue = asyncio.Queue()
     session = Session(key="test:drain-timeout")
-    injection_callback = None
+    terminal_injection_callback = None
 
     async def fake_runner_run(spec):
-        nonlocal injection_callback
-        injection_callback = spec.injection_callback
+        nonlocal terminal_injection_callback
+        terminal_injection_callback = spec.terminal_injection_callback
         return SimpleNamespace(
             stop_reason="done",
             final_content="done",
@@ -662,7 +667,7 @@ async def test_drain_pending_timeout(tmp_path):
         pending_queue=pending_queue,
     )
 
-    assert injection_callback is not None
+    assert terminal_injection_callback is not None
 
     # Patch the timeout path without leaking the queue.get() coroutine.
     async def _timeout(awaitable, timeout):
@@ -670,7 +675,7 @@ async def test_drain_pending_timeout(tmp_path):
         raise asyncio.TimeoutError
 
     with patch("nanobot.agent.loop.asyncio.wait_for", side_effect=_timeout):
-        results = await injection_callback()
+        results = await terminal_injection_callback()
         assert results == []
 
     # Cleanup
@@ -679,3 +684,81 @@ async def test_drain_pending_timeout(tmp_path):
         await hang_task
     except asyncio.CancelledError:
         pass
+
+
+@pytest.mark.asyncio
+async def test_terminal_drain_reuses_one_timeout_budget(tmp_path):
+    """Repeated terminal rendezvous calls share one 300-second deadline."""
+    from nanobot.agent.loop import AgentLoop
+    from nanobot.bus.events import InboundMessage
+    from nanobot.bus.queue import MessageBus
+    from nanobot.session.manager import Session
+
+    loop = AgentLoop(
+        bus=MessageBus(),
+        provider=MagicMock(),
+        workspace=tmp_path,
+        model="test-model",
+    )
+    pending_queue: asyncio.Queue[InboundMessage] = asyncio.Queue()
+    session = Session(key="test:shared-deadline")
+    terminal_injection_callback = None
+
+    async def fake_runner_run(spec):
+        nonlocal terminal_injection_callback
+        terminal_injection_callback = spec.terminal_injection_callback
+        return SimpleNamespace(
+            stop_reason="done",
+            final_content="done",
+            error=None,
+            tool_events=[],
+            messages=[],
+            usage=None,
+            had_injections=False,
+            tools_used=[],
+            provider_state=None,
+        )
+
+    loop.runner.run = AsyncMock(side_effect=fake_runner_run)
+
+    async def _hang_forever():
+        await asyncio.Event().wait()
+
+    hang_task = asyncio.create_task(_hang_forever())
+    loop.subagents._session_tasks.setdefault(session.key, set()).add("sub-deadline-1")
+    loop.subagents._running_tasks["sub-deadline-1"] = hang_task
+
+    await loop._run_agent_loop(
+        [{"role": "user", "content": "test"}],
+        runtime=loop.llm_runtime(),
+        session=session,
+        pending_queue=pending_queue,
+    )
+    assert terminal_injection_callback is not None
+
+    timeouts: list[float] = []
+    clock = MagicMock(side_effect=[10.0, 110.0])
+
+    async def _deliver(awaitable, timeout):
+        awaitable.close()
+        timeouts.append(timeout)
+        return InboundMessage(
+            sender_id="subagent",
+            channel="test",
+            chat_id="c1",
+            content="result",
+        )
+
+    fake_loop = SimpleNamespace(time=clock)
+    with (
+        patch("nanobot.agent.loop.asyncio.get_running_loop", return_value=fake_loop),
+        patch("nanobot.agent.loop.asyncio.wait_for", side_effect=_deliver),
+    ):
+        assert await terminal_injection_callback()
+        assert await terminal_injection_callback()
+
+    assert timeouts == [300.0, 200.0]
+
+    hang_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await hang_task