From 5f054c0e74fb708f815150f40e15d857b0bf08fa Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 12:17:13 +0800 Subject: [PATCH] fix(agent): deliver non-streamed finalization responses --- nanobot/agent/loop.py | 32 ++++++++++++-- tests/agent/test_loop_progress.py | 49 +++++++++++++++++++++ tests/agent/test_loop_runner_integration.py | 9 +++- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 29f946b2..50df4949 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -149,6 +149,7 @@ class TurnContext: all_messages: list[dict[str, Any]] = field(default_factory=list) stop_reason: str = "" had_injections: bool = False + streamed_content: bool = False input_persisted_early: bool = False save_skip: int = 0 @@ -1304,6 +1305,29 @@ class AgentLoop: hook_factories=list(hook_factories or []), tools=tools, ) + # A streaming callback may be present even when the final text comes from a + # non-streaming recovery. Only the last completed segment can suppress the + # regular outbound message. + if ctx.on_stream is not None: + stream_callback = ctx.on_stream + stream_end_callback = ctx.on_stream_end + segment_streamed_content = False + + async def _tracked_stream(delta: str) -> None: + nonlocal segment_streamed_content + if delta: + segment_streamed_content = True + await stream_callback(delta) + + async def _tracked_stream_end(*, resuming: bool = False) -> None: + nonlocal segment_streamed_content + ctx.streamed_content = segment_streamed_content + segment_streamed_content = False + if stream_end_callback is not None: + await stream_end_callback(resuming=resuming) + + ctx.on_stream = _tracked_stream + ctx.on_stream_end = _tracked_stream_end while ctx.state is not TurnState.DONE: handler_name = f"_state_{ctx.state.name.lower()}" @@ -1366,7 +1390,7 @@ class AgentLoop: all_msgs: list[dict[str, Any]], stop_reason: str, had_injections: bool, - on_stream: Callable[[str], Awaitable[None]] | None, + streamed_content: bool, *, turn_latency_ms: int | None = None, ) -> OutboundMessage | None: @@ -1381,7 +1405,7 @@ class AgentLoop: event = None meta = dict(msg.metadata or {}) - if on_stream is not None and stop_reason not in {"error", "tool_error"}: + if streamed_content and stop_reason not in {"error", "tool_error"}: event = StreamedResponseEvent() if turn_latency_ms is not None: meta["latency_ms"] = int(turn_latency_ms) @@ -1626,7 +1650,7 @@ class AgentLoop: ctx.outbound = ctx.delivery.background_response( ctx.final_content, stop_reason=ctx.stop_reason, - streamed=ctx.on_stream is not None, + streamed=ctx.streamed_content, latency_ms=ctx.turn_latency_ms, ) return "ok" @@ -1636,7 +1660,7 @@ class AgentLoop: ctx.all_messages, ctx.stop_reason, ctx.had_injections, - ctx.on_stream, + ctx.streamed_content, turn_latency_ms=ctx.turn_latency_ms, ) if ctx.ephemeral and ctx.outbound is not None: diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 3ea6d953..3b490645 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -534,6 +534,55 @@ class TestToolEventProgress: assert turn_end_msgs[0].content == "" provider.chat_with_retry.assert_not_awaited() + @pytest.mark.asyncio + async def test_non_streamed_finalization_is_delivered_as_regular_message( + self, + tmp_path: Path, + ) -> None: + """A no-tools finalization must not be dropped after empty stream retries.""" + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "openai-codex/gpt-5.5" + provider.chat_stream_with_retry = AsyncMock(side_effect=[ + LLMResponse(content=None, tool_calls=[]), + LLMResponse(content=None, tool_calls=[]), + ]) + provider.chat_with_retry = AsyncMock( + return_value=LLMResponse(content="final answer", tool_calls=[]), + ) + loop = AgentLoop( + bus=bus, + provider=provider, + workspace=tmp_path, + model="openai-codex/gpt-5.5", + ) + _attach_webui_runtime_events(loop, bus) + loop.tools.get_definitions = MagicMock(return_value=[]) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + + await loop._dispatch(InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="chat1", + content="say hello", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + assert not any(isinstance(message.event, StreamDeltaEvent) for message in outbound) + assert len([ + message for message in outbound if isinstance(message.event, StreamEndEvent) + ]) == 3 + final = [message for message in outbound if message.content == "final answer"] + assert len(final) == 1 + assert final[0].event is None + provider.chat_stream_with_retry.assert_awaited() + provider.chat_with_retry.assert_awaited_once() + @pytest.mark.asyncio async def test_independent_late_subagent_result_gets_complete_webui_turn( self, diff --git a/tests/agent/test_loop_runner_integration.py b/tests/agent/test_loop_runner_integration.py index 53e8971a..da8a0893 100644 --- a/tests/agent/test_loop_runner_integration.py +++ b/tests/agent/test_loop_runner_integration.py @@ -512,7 +512,7 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path): )], usage={}, ) - provider.chat_stream_with_retry = AsyncMock(side_effect=[ + responses = iter([ tool_call_resp, LLMResponse( content="I cannot access private URLs. Please share the local file.", @@ -521,6 +521,13 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path): ), ]) + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(responses) + await on_content_delta(response.content) + return response + + provider.chat_stream_with_retry = AsyncMock(side_effect=chat_stream_with_retry) + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.prepare_call = MagicMock(return_value=(None, {}, None))