From 27a00c7a4f1bf952b2b224b49546f8d23dde7904 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 18:35:41 +0800 Subject: [PATCH] fix(webui): merge length recovery stream segments --- nanobot/agent/hook.py | 1 + nanobot/agent/loop.py | 31 ++++++-- nanobot/agent/progress_hook.py | 8 +- nanobot/agent/runner.py | 1 + nanobot/agent/turn_delivery.py | 11 ++- nanobot/bus/outbound_events.py | 2 + nanobot/channels/base.py | 4 + nanobot/channels/discord/runtime.py | 1 + nanobot/channels/feishu/runtime.py | 1 + nanobot/channels/manager.py | 24 +++++- nanobot/channels/matrix/runtime.py | 1 + nanobot/channels/mattermost/runtime.py | 1 + nanobot/channels/telegram/runtime.py | 1 + nanobot/channels/websocket/runtime.py | 9 ++- .../websocket/tests/test_websocket_channel.py | 33 +++++++++ nanobot/channels/weixin/runtime.py | 1 + nanobot/webui/transcript.py | 8 +- tests/agent/test_loop_progress.py | 46 ++++++++++++ tests/agent/test_runner_hooks.py | 3 + tests/bus/test_outbound_events.py | 11 ++- .../test_channel_manager_delta_coalescing.py | 19 ++++- tests/channels/test_channel_plugins.py | 16 ++-- tests/utils/test_webui_transcript.py | 20 +++++ webui/src/hooks/useNanobotStream.ts | 26 +++++-- webui/src/lib/types.ts | 2 + webui/src/tests/useNanobotStream.test.tsx | 73 +++++++++++++++++++ 26 files changed, 325 insertions(+), 29 deletions(-) diff --git a/nanobot/agent/hook.py b/nanobot/agent/hook.py index 12a4e717..8e4dd5ff 100644 --- a/nanobot/agent/hook.py +++ b/nanobot/agent/hook.py @@ -25,6 +25,7 @@ class AgentHookContext: tool_events: list[dict[str, str]] = field(default_factory=list) streamed_content: bool = False streamed_reasoning: bool = False + stream_continues_current_message: bool = False final_content: str | None = None stop_reason: str | None = None error: str | None = None diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index f079867b..e972b6be 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import dataclasses +import inspect import os import time from collections.abc import Mapping @@ -860,9 +861,9 @@ class AgentLoop: """Run the agent iteration loop. *on_stream*: called with each content delta during streaming. - *on_stream_end(resuming)*: called when a streaming session finishes. - ``resuming=True`` means tool calls follow (spinner should restart); - ``resuming=False`` means this is the final response. + *on_stream_end(resuming, merge_next)*: called when a streaming session finishes. + ``resuming=True`` means the active turn continues. ``merge_next=True`` means + the next text segment belongs to the same user-visible assistant message. Returns (final_content, tools_used, messages, stop_reason, had_injections). """ @@ -1385,6 +1386,19 @@ class AgentLoop: if ctx.on_stream is not None: stream_callback = ctx.on_stream stream_end_callback = ctx.on_stream_end + stream_end_accepts_merge_next = False + if stream_end_callback is not None: + try: + stream_end_signature = inspect.signature(stream_end_callback) + stream_end_accepts_merge_next = ( + "merge_next" in stream_end_signature.parameters + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in stream_end_signature.parameters.values() + ) + ) + except (TypeError, ValueError): + pass segment_streamed_content = False async def _tracked_stream(delta: str) -> None: @@ -1393,12 +1407,19 @@ class AgentLoop: segment_streamed_content = True await stream_callback(delta) - async def _tracked_stream_end(*, resuming: bool = False) -> None: + async def _tracked_stream_end( + *, + resuming: bool = False, + merge_next: 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) + if merge_next and stream_end_accepts_merge_next: + await stream_end_callback(resuming=resuming, merge_next=True) + else: + await stream_end_callback(resuming=resuming) ctx.on_stream = _tracked_stream ctx.on_stream_end = _tracked_stream_end diff --git a/nanobot/agent/progress_hook.py b/nanobot/agent/progress_hook.py index 205c2019..826093d9 100644 --- a/nanobot/agent/progress_hook.py +++ b/nanobot/agent/progress_hook.py @@ -85,7 +85,13 @@ class AgentProgressHook(AgentHook): async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: await self.emit_reasoning_end() if self._on_stream_end: - await self._on_stream_end(resuming=resuming) + kwargs: dict[str, bool] = {"resuming": resuming} + if ( + context.stream_continues_current_message + and self._on_progress_accepts(self._on_stream_end, "merge_next") + ): + kwargs["merge_next"] = True + await self._on_stream_end(**kwargs) self._stream_buf = "" self._think_extractor.reset() diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 18e93aad..3d217822 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -599,6 +599,7 @@ class AgentRunner: _MAX_LENGTH_RECOVERIES, ) if hook.wants_streaming(): + context.stream_continues_current_message = True await hook.on_stream_end(context, resuming=True) messages.append(build_assistant_message( clean, diff --git a/nanobot/agent/turn_delivery.py b/nanobot/agent/turn_delivery.py index 2eeb5aab..a415551d 100644 --- a/nanobot/agent/turn_delivery.py +++ b/nanobot/agent/turn_delivery.py @@ -285,7 +285,12 @@ class TurnDelivery: ) ) - async def _publish_stream_end(self, *, resuming: bool = False) -> None: + async def _publish_stream_end( + self, + *, + resuming: bool = False, + merge_next: bool = False, + ) -> None: await self.bus.publish_outbound( outbound_message_for_event( channel=self.delivery_message.channel, @@ -293,8 +298,10 @@ class TurnDelivery: event=StreamEndEvent( stream_id=self._stream_id(), resuming=resuming, + merge_next=merge_next, ), metadata=self.delivery_message.metadata, ) ) - self._stream_segment += 1 + if not merge_next: + self._stream_segment += 1 diff --git a/nanobot/bus/outbound_events.py b/nanobot/bus/outbound_events.py index 3b03068e..1a5b8d55 100644 --- a/nanobot/bus/outbound_events.py +++ b/nanobot/bus/outbound_events.py @@ -46,6 +46,7 @@ class StreamEndEvent(OutboundEvent): content: str = "" stream_id: str | None = None resuming: bool = False + merge_next: bool = False @dataclass(frozen=True) @@ -176,6 +177,7 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None: content=msg.content, stream_id=_metadata_str(meta, "_stream_id"), resuming=bool(meta.get("_resuming")), + merge_next=bool(meta.get("_merge_next")), ) if meta.get("_stream_delta"): return StreamDeltaEvent( diff --git a/nanobot/channels/base.py b/nanobot/channels/base.py index c4f51b7a..01a794a4 100644 --- a/nanobot/channels/base.py +++ b/nanobot/channels/base.py @@ -110,6 +110,7 @@ class BaseChannel(ABC): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Deliver a streaming text chunk. @@ -118,6 +119,9 @@ class BaseChannel(ABC): Stateful implementations should key buffers by ``stream_id`` rather than only by ``chat_id`` when it is provided. + + ``merge_next`` marks a resumable provider boundary whose next text + segment belongs to the same user-visible message. """ pass diff --git a/nanobot/channels/discord/runtime.py b/nanobot/channels/discord/runtime.py index c6d858d6..b20baca1 100644 --- a/nanobot/channels/discord/runtime.py +++ b/nanobot/channels/discord/runtime.py @@ -489,6 +489,7 @@ class DiscordChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive Discord delivery: send once, then edit until the stream ends.""" client = self._client diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index 51f6d8b5..1541fd97 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -2216,6 +2216,7 @@ class FeishuChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive streaming via CardKit: create card on first delta, stream-update on subsequent. diff --git a/nanobot/channels/manager.py b/nanobot/channels/manager.py index 03fe89be..3fbe32ed 100644 --- a/nanobot/channels/manager.py +++ b/nanobot/channels/manager.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import hashlib +import inspect from collections.abc import Callable, Iterable from contextlib import suppress from pathlib import Path @@ -763,13 +764,29 @@ class ChannelManager: msg: OutboundMessage, event: StreamDeltaEvent | StreamEndEvent, ) -> None: + kwargs: dict[str, Any] = { + "stream_id": event.stream_id, + "stream_end": isinstance(event, StreamEndEvent), + "resuming": event.resuming if isinstance(event, StreamEndEvent) else False, + } + if isinstance(event, StreamEndEvent) and event.merge_next: + try: + signature = inspect.signature(channel.send_delta) + if ( + "merge_next" in signature.parameters + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in signature.parameters.values() + ) + ): + kwargs["merge_next"] = True + except (TypeError, ValueError): + pass await channel.send_delta( msg.chat_id, msg.content, msg.metadata, - stream_id=event.stream_id, - stream_end=isinstance(event, StreamEndEvent), - resuming=event.resuming if isinstance(event, StreamEndEvent) else False, + **kwargs, ) @staticmethod @@ -850,6 +867,7 @@ class ChannelManager: final_event = StreamEndEvent( stream_id=next_stream_id, resuming=next_event.resuming, + merge_next=next_event.merge_next, ) # Stream ended - stop coalescing this stream break diff --git a/nanobot/channels/matrix/runtime.py b/nanobot/channels/matrix/runtime.py index 2544d64e..f0cccfb4 100644 --- a/nanobot/channels/matrix/runtime.py +++ b/nanobot/channels/matrix/runtime.py @@ -598,6 +598,7 @@ class MatrixChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: relates_to = self._build_thread_relates_to(metadata) diff --git a/nanobot/channels/mattermost/runtime.py b/nanobot/channels/mattermost/runtime.py index 756daed7..d15bcd3f 100644 --- a/nanobot/channels/mattermost/runtime.py +++ b/nanobot/channels/mattermost/runtime.py @@ -515,6 +515,7 @@ class MattermostChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: if not self._http_client: return diff --git a/nanobot/channels/telegram/runtime.py b/nanobot/channels/telegram/runtime.py index 3a802074..4183a82c 100644 --- a/nanobot/channels/telegram/runtime.py +++ b/nanobot/channels/telegram/runtime.py @@ -923,6 +923,7 @@ class TelegramChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Progressive message editing: send on first delta, edit on subsequent ones.""" if not self._app: diff --git a/nanobot/channels/websocket/runtime.py b/nanobot/channels/websocket/runtime.py index 6741ad43..cdd8e4f8 100644 --- a/nanobot/channels/websocket/runtime.py +++ b/nanobot/channels/websocket/runtime.py @@ -995,13 +995,18 @@ class WebSocketChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: conns = list(self._subs.get(chat_id, ())) meta = metadata or {} stream_key = (chat_id, str(stream_id or "")) if stream_end: body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id} - buffered = self._stream_text_buffers.pop(stream_key, []) + buffered = ( + self._stream_text_buffers.setdefault(stream_key, []) + if merge_next + else self._stream_text_buffers.pop(stream_key, []) + ) if delta: buffered.append(delta) full_text = "".join(buffered) @@ -1019,6 +1024,8 @@ class WebSocketChannel(BaseChannel): body["stream_id"] = stream_id if stream_end and resuming: body["resuming"] = True + if stream_end and merge_next: + body["merge_next"] = True self._transcripts.prepare_and_append( chat_id, body, diff --git a/nanobot/channels/websocket/tests/test_websocket_channel.py b/nanobot/channels/websocket/tests/test_websocket_channel.py index fcefdb23..92eb84bb 100644 --- a/nanobot/channels/websocket/tests/test_websocket_channel.py +++ b/nanobot/channels/websocket/tests/test_websocket_channel.py @@ -1350,6 +1350,39 @@ async def test_send_delta_marks_resuming_stream_end() -> None: assert payload["resuming"] is True +@pytest.mark.asyncio +async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None: + bus = MagicMock() + channel = WebSocketChannel( + {"enabled": True, "allowFrom": ["*"], "streaming": True}, + bus, + gateway=_basic_handler(bus), + ) + mock_ws = AsyncMock() + channel._attach(mock_ws, "chat-1") + + await channel.send_delta("chat-1", "first ", stream_id="sid") + await channel.send_delta( + "chat-1", + "", + stream_id="sid", + stream_end=True, + resuming=True, + merge_next=True, + ) + await channel.send_delta("chat-1", "second", stream_id="sid") + await channel.send_delta("chat-1", "", stream_id="sid", stream_end=True) + + payloads = [json.loads(call.args[0]) for call in mock_ws.send.await_args_list] + assert payloads[1]["merge_next"] is True + assert payloads[1]["resuming"] is True + assert [payload["text"] for payload in payloads if payload["event"] == "delta"] == [ + "first ", + "second", + ] + assert ("chat-1", "sid") not in channel._stream_text_buffers + + @pytest.mark.asyncio async def test_send_delta_stream_end_includes_inline_final_text() -> None: bus = MagicMock() diff --git a/nanobot/channels/weixin/runtime.py b/nanobot/channels/weixin/runtime.py index eab09c84..976013c0 100644 --- a/nanobot/channels/weixin/runtime.py +++ b/nanobot/channels/weixin/runtime.py @@ -1243,6 +1243,7 @@ class WeixinChannel(BaseChannel): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: """Deliver a streamed reply to WeChat. diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index 63a9aac9..2c8a7329 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -1770,6 +1770,7 @@ def replay_transcript_to_ui_messages( buffer_message_id = None buffer_parts = [] continue + merge_next = rec.get("resuming") is True and rec.get("merge_next") is True final_text = rec.get("text") if isinstance(final_text, str): if buffer_message_id is None: @@ -1794,8 +1795,11 @@ def replay_transcript_to_ui_messages( **_turn_fields(rec, "answer"), } break - buffer_message_id = None - buffer_parts = [] + if merge_next: + buffer_parts = [final_text] + if not merge_next: + buffer_message_id = None + buffer_parts = [] continue if ev == "reasoning_delta": diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 3b490645..aca35927 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -534,6 +534,52 @@ class TestToolEventProgress: assert turn_end_msgs[0].content == "" provider.chat_with_retry.assert_not_awaited() + @pytest.mark.asyncio + async def test_length_recovery_keeps_one_user_visible_stream( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "test-model" + responses = iter([ + LLMResponse(content="first-", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + ]) + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(responses) + await on_content_delta(response.content or "") + return response + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + _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="give a long answer", + metadata={"_wants_stream": True}, + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)] + endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)] + + assert [event.content for event in deltas] == ["first-", "second"] + assert [event.resuming for event in endings] == [True, False] + assert [event.merge_next for event in endings] == [True, False] + assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id} + @pytest.mark.asyncio async def test_non_streamed_finalization_is_delivered_as_regular_message( self, diff --git a/tests/agent/test_runner_hooks.py b/tests/agent/test_runner_hooks.py index 0bc8a6fe..508f663e 100644 --- a/tests/agent/test_runner_hooks.py +++ b/tests/agent/test_runner_hooks.py @@ -151,6 +151,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont provider = MagicMock(spec=LLMProvider) streamed: list[str] = [] endings: list[bool] = [] + merge_next: list[bool] = [] responses = iter([ LLMResponse(content="first ", finish_reason="length"), LLMResponse(content="second", finish_reason="stop"), @@ -175,6 +176,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: endings.append(resuming) + merge_next.append(context.stream_continues_current_message) runner = AgentRunner() result = await runner.run(make_run_spec(provider, @@ -189,6 +191,7 @@ async def test_runner_length_recovery_streams_segments_once_and_returns_all_cont assert result.final_content == "first second" assert streamed == ["first ", "second"] assert endings == [True, False] + assert merge_next == [True, False] provider.chat_with_retry.assert_not_awaited() diff --git a/tests/bus/test_outbound_events.py b/tests/bus/test_outbound_events.py index 0beb7e60..99eb259c 100644 --- a/tests/bus/test_outbound_events.py +++ b/tests/bus/test_outbound_events.py @@ -94,7 +94,12 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None: channel="websocket", chat_id="chat-1", content="", - metadata={"_stream_end": True, "_stream_id": "s1", "_resuming": True}, + metadata={ + "_stream_end": True, + "_stream_id": "s1", + "_resuming": True, + "_merge_next": True, + }, ) delta_event = outbound_event_from_message(delta) @@ -106,6 +111,7 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None: assert isinstance(end_event, StreamEndEvent) assert end_event.stream_id == "s1" assert end_event.resuming is True + assert end_event.merge_next is True def test_legacy_webui_runtime_metadata_flags_create_runtime_events() -> None: @@ -221,7 +227,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None: updated = replace_outbound_event( msg, - StreamEndEvent(stream_id="s1", resuming=True), + StreamEndEvent(stream_id="s1", resuming=True, merge_next=True), content="hello world", ) @@ -230,6 +236,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None: assert isinstance(updated.event, StreamEndEvent) assert updated.event.stream_id == "s1" assert updated.event.resuming is True + assert updated.event.merge_next is True def test_streamed_response_event_keeps_final_content_outside_event_payload() -> None: diff --git a/tests/channels/test_channel_manager_delta_coalescing.py b/tests/channels/test_channel_manager_delta_coalescing.py index de60124e..8a009d84 100644 --- a/tests/channels/test_channel_manager_delta_coalescing.py +++ b/tests/channels/test_channel_manager_delta_coalescing.py @@ -49,6 +49,7 @@ class MockChannel(BaseChannel): stream_id=None, stream_end=False, resuming=False, + merge_next=False, ): return await self._send_delta_mock( chat_id, @@ -57,6 +58,7 @@ class MockChannel(BaseChannel): stream_id=stream_id, stream_end=stream_end, resuming=resuming, + merge_next=merge_next, ) @@ -92,11 +94,17 @@ def _end( chat_id: str = "chat1", stream_id: str | None = None, resuming: bool = False, + merge_next: bool = False, ): return outbound_message_for_event( channel="mock", chat_id=chat_id, - event=StreamEndEvent(content=content, stream_id=stream_id, resuming=resuming), + event=StreamEndEvent( + content=content, + stream_id=stream_id, + resuming=resuming, + merge_next=merge_next, + ), ) @@ -137,6 +145,7 @@ class TestDeltaCoalescing: stream_id=None, stream_end=False, resuming=False, + merge_next=False, ) @pytest.mark.asyncio @@ -184,13 +193,19 @@ class TestDeltaCoalescing: @pytest.mark.asyncio async def test_stream_end_terminates_coalescing(self, manager, bus): await bus.publish_outbound(_delta("Hello")) - await bus.publish_outbound(_end(" world")) + await bus.publish_outbound(_end( + " world", + resuming=True, + merge_next=True, + )) first_msg = await bus.consume_outbound() merged, pending = manager._coalesce_stream_deltas(first_msg) assert merged.content == "Hello world" assert isinstance(merged.event, StreamEndEvent) + assert merged.event.resuming is True + assert merged.event.merge_next is True assert len(pending) == 0 @pytest.mark.asyncio diff --git a/tests/channels/test_channel_plugins.py b/tests/channels/test_channel_plugins.py index 50ff6905..93a8b31d 100644 --- a/tests/channels/test_channel_plugins.py +++ b/tests/channels/test_channel_plugins.py @@ -2818,7 +2818,7 @@ async def test_send_with_retry_no_retry_when_max_is_zero(): @pytest.mark.asyncio async def test_send_with_retry_calls_send_delta(): """_send_with_retry should call send_delta for stream delta events.""" - calls: list[tuple[str, str, str | None, bool, bool]] = [] + calls: list[tuple[str, str, str | None, bool, bool, bool]] = [] class _StreamingChannel(BaseChannel): name = "streaming" @@ -2842,8 +2842,9 @@ async def test_send_with_retry_calls_send_delta(): stream_id: str | None = None, stream_end: bool = False, resuming: bool = False, + merge_next: bool = False, ) -> None: - calls.append((chat_id, delta, stream_id, stream_end, resuming)) + calls.append((chat_id, delta, stream_id, stream_end, resuming, merge_next)) fake_config = SimpleNamespace( channels=ChannelsConfig(send_max_retries=3), @@ -2865,13 +2866,18 @@ async def test_send_with_retry_calls_send_delta(): end = outbound_message_for_event( channel="streaming", chat_id="123", - event=StreamEndEvent(content="", stream_id="s1", resuming=True), + event=StreamEndEvent( + content="", + stream_id="s1", + resuming=True, + merge_next=True, + ), ) await mgr._send_with_retry(mgr.channels["streaming"], end) assert calls == [ - ("123", "test delta", "s1", False, False), - ("123", "", "s1", True, True), + ("123", "test delta", "s1", False, False, False), + ("123", "", "s1", True, True, True), ] diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index 067a45fa..921982f6 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -1121,6 +1121,26 @@ def test_replay_keeps_interrupted_pre_tool_text_in_activity() -> None: assert msgs[2]["content"] == "Done. Open index.html to play." +def test_replay_merges_length_recovery_segments_into_one_assistant_message() -> None: + msgs = replay_transcript_to_ui_messages([ + {"event": "delta", "chat_id": "t-stream", "text": "first "}, + { + "event": "stream_end", + "chat_id": "t-stream", + "text": "first ", + "resuming": True, + "merge_next": True, + }, + {"event": "delta", "chat_id": "t-stream", "text": "second"}, + {"event": "stream_end", "chat_id": "t-stream"}, + {"event": "turn_end", "chat_id": "t-stream"}, + ]) + + assert len(msgs) == 1 + assert msgs[0]["role"] == "assistant" + assert msgs[0]["content"] == "first second" + + def test_replay_tool_events_dedupes_finish_after_start() -> None: msgs = replay_transcript_to_ui_messages([ { diff --git a/webui/src/hooks/useNanobotStream.ts b/webui/src/hooks/useNanobotStream.ts index 4d2ee120..1f18e200 100644 --- a/webui/src/hooks/useNanobotStream.ts +++ b/webui/src/hooks/useNanobotStream.ts @@ -26,7 +26,7 @@ import type { } from "@/lib/types"; interface StreamBuffer { - /** ID of the assistant message currently receiving deltas (cleared on ``stream_end``). */ + /** ID of the assistant message currently receiving deltas (cleared when its segment closes). */ messageId: string; } @@ -780,15 +780,20 @@ export function useNanobotStream( ?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current, turn); if (targetIndex !== null) { const target = next[targetIndex]; - next = replaceMessageAt(next, targetIndex, { + const merged = { ...target, content: finalAnswerText, isStreaming: true, ...turn, - }); + }; + next = replaceMessageAt(next, targetIndex, merged); + if (!options?.closeAnswerSegment) { + closedAssistantStreamIdsRef.current.delete(merged.id); + activeAssistantRef.current = { id: merged.id, index: targetIndex }; + buffer.current = { messageId: merged.id }; + } } else { const id = crypto.randomUUID(); - closedAssistantStreamIdsRef.current.add(id); next = [ ...next, { @@ -800,6 +805,12 @@ export function useNanobotStream( createdAt: Date.now(), }, ]; + if (options?.closeAnswerSegment) { + closedAssistantStreamIdsRef.current.add(id); + } else { + activeAssistantRef.current = { id, index: next.length - 1 }; + buffer.current = { messageId: id }; + } } } if (options?.closeAnswerSegment) closeActiveAssistantStream(); @@ -911,8 +922,9 @@ export function useNanobotStream( if (ev.event === "stream_end") { const turn = turnFieldsFromEvent(ev, "answer"); + const mergeNext = ev.resuming === true && ev.merge_next === true; flushPendingStreamEvents({ - closeAnswerSegment: true, + closeAnswerSegment: !mergeNext, ...(typeof ev.text === "string" ? { finalAnswerText: ev.text } : {}), turn, }); @@ -920,7 +932,9 @@ export function useNanobotStream( if (ev.resuming) { cancelStreamEndTimer(); setIsStreaming(true); - setMessages((prev) => finalizeStreamedTurn(prev, turn)); + if (!mergeNext) { + setMessages((prev) => finalizeStreamedTurn(prev, turn)); + } return; } scheduleStreamEndTimer(turn); diff --git a/webui/src/lib/types.ts b/webui/src/lib/types.ts index 5a9ce09e..1db55490 100644 --- a/webui/src/lib/types.ts +++ b/webui/src/lib/types.ts @@ -1118,6 +1118,8 @@ export type InboundEvent = text?: string; /** This answer segment ended, but the active agent turn will continue. */ resuming?: boolean; + /** The next answer segment continues this same assistant message. */ + merge_next?: boolean; } & InboundTurnMetadata) | ({ event: "reasoning_delta"; diff --git a/webui/src/tests/useNanobotStream.test.tsx b/webui/src/tests/useNanobotStream.test.tsx index c5cd49a3..ebd2db0d 100644 --- a/webui/src/tests/useNanobotStream.test.tsx +++ b/webui/src/tests/useNanobotStream.test.tsx @@ -2009,6 +2009,79 @@ describe("useNanobotStream", () => { expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true); }); + it("keeps length-recovery segments in one assistant message", async () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-length", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + + act(() => { + result.current.send("give a long answer"); + }); + const activeTurnId = fake.client.sendMessage.mock.calls.at(-1)![3]?.turnId; + + act(() => { + fake.emit("chat-length", { + event: "delta", + chat_id: "chat-length", + text: "first ", + turn_id: activeTurnId, + }); + }); + await flushStreamFrame(); + const assistantId = result.current.messages[1].id; + + act(() => { + fake.emit("chat-length", { + event: "stream_end", + chat_id: "chat-length", + text: "first ", + resuming: true, + merge_next: true, + turn_id: activeTurnId, + }); + }); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first ", + isStreaming: true, + }); + + act(() => { + fake.emit("chat-length", { + event: "delta", + chat_id: "chat-length", + text: "second", + turn_id: activeTurnId, + }); + }); + await flushStreamFrame(); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first second", + isStreaming: true, + }); + + act(() => { + fake.emit("chat-length", { + event: "turn_end", + chat_id: "chat-length", + turn_id: activeTurnId, + }); + }); + + expect(result.current.messages).toHaveLength(2); + expect(result.current.messages[1]).toMatchObject({ + id: assistantId, + content: "first second", + isStreaming: false, + }); + }); + it("keeps streaming alive across stream_end when tool activity follows", async () => { const fake = fakeClient(); const onTurnEnd = vi.fn();