From 66d9328a003335cfade0459a948c1685b194f2b1 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Tue, 25 Aug 2026 15:19:21 +0800 Subject: [PATCH] fix(agent): defer subagent wait until turn exit --- nanobot/agent/loop.py | 76 ++++++++------ nanobot/agent/runner.py | 34 +++++- tests/agent/test_runner_core.py | 3 + tests/agent/test_runner_injections.py | 124 ++++++++++++++++++++++ tests/agent/tools/test_subagent_tools.py | 125 +++++++++++++++++++---- 5 files changed, 306 insertions(+), 56 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 140c3fbc1..d7efda7a2 100644 --- a/nanobot/agent/loop.py +++ b/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( diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 6a802e75e..a777bb6c6 100644 --- a/nanobot/agent/runner.py +++ b/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 diff --git a/tests/agent/test_runner_core.py b/tests/agent/test_runner_core.py index b16af1a42..7ca319a4b 100644 --- a/tests/agent/test_runner_core.py +++ b/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" diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index cafbc89e3..ea076aaa4 100644 --- a/tests/agent/test_runner_injections.py +++ b/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.""" diff --git a/tests/agent/tools/test_subagent_tools.py b/tests/agent/tools/test_subagent_tools.py index 0a90b1ff6..35b29e06f 100644 --- a/tests/agent/tools/test_subagent_tools.py +++ b/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