fix(webui): merge length recovery stream segments

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

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

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

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,

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

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(

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
 

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

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.
 

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

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)
 

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

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:

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,

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

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.
 

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":