fix(webui): merge length recovery stream segments

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,

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

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:

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

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),
     ]
 
 

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([
         {