From b55b76d75574d74c1ff5356bace3b76ad12ef399 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 27 Jul 2026 01:33:04 +0800 Subject: [PATCH] fix(streaming): preserve recovered segments across channels --- nanobot/agent/runner.py | 17 +++++++ nanobot/channels/discord/runtime.py | 4 ++ .../discord/tests/test_discord_channel.py | 30 ++++++++++++ nanobot/channels/feishu/runtime.py | 4 ++ .../feishu/tests/test_feishu_streaming.py | 21 +++++++++ nanobot/channels/matrix/runtime.py | 4 ++ .../matrix/tests/test_matrix_channel.py | 23 +++++++++ nanobot/channels/telegram/runtime.py | 4 ++ .../telegram/tests/test_telegram_channel.py | 27 +++++++++++ nanobot/channels/weixin/runtime.py | 4 ++ .../weixin/tests/test_weixin_channel.py | 23 +++++++++ tests/agent/test_loop_progress.py | 47 +++++++++++++++++++ 12 files changed, 208 insertions(+) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index acc54abb..671208ce 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -613,6 +613,23 @@ class AgentRunner: await hook.after_iteration(context) continue + # Some streaming providers recover with a complete response but no + # content deltas. When an earlier length segment is already visible, + # emit this terminal segment into the same stream; otherwise the + # regular full response would duplicate the visible prefix. + if ( + length_recovery_parts + and hook.wants_streaming() + and not context.streamed_content + and response.finish_reason != "error" + and not is_blank_text(clean) + ): + await hook.on_stream( + context, + _restore_outer_whitespace(clean, original_content), + ) + context.streamed_content = True + assistant_message: dict[str, Any] | None = None if response.finish_reason != "error" and not is_blank_text(clean): assistant_message = build_assistant_message( diff --git a/nanobot/channels/discord/runtime.py b/nanobot/channels/discord/runtime.py index b20baca1..bab06fe2 100644 --- a/nanobot/channels/discord/runtime.py +++ b/nanobot/channels/discord/runtime.py @@ -497,6 +497,10 @@ class DiscordChannel(BaseChannel): self.logger.warning("client not ready; dropping stream delta") return + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: buf = self._stream_bufs.get(chat_id) if not buf or buf.message is None or not buf.text: diff --git a/nanobot/channels/discord/tests/test_discord_channel.py b/nanobot/channels/discord/tests/test_discord_channel.py index d86b56b7..a749363b 100644 --- a/nanobot/channels/discord/tests/test_discord_channel.py +++ b/nanobot/channels/discord/tests/test_discord_channel.py @@ -754,6 +754,36 @@ async def test_send_delta_streams_by_editing_message(monkeypatch) -> None: assert owner._stream_bufs == {} +@pytest.mark.asyncio +async def test_send_delta_merge_next_keeps_one_message(monkeypatch) -> None: + owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus()) + client = _FakeDiscordClient(owner, intents=None) + owner._client = client + owner._running = True + target = _FakeChannel(channel_id=123) + client.channels[123] = target + + times = iter([1.0, 3.0, 5.0]) + monkeypatch.setattr("nanobot.channels.discord.runtime.time.monotonic", lambda: next(times, 5.0)) + + await owner.send_delta( + "123", + "first-", + stream_id="s1", + stream_end=True, + merge_next=True, + ) + await owner.send_delta("123", "second", stream_id="s1") + await owner.send_delta("123", "", stream_id="s1", stream_end=True) + + assert target.sent_payloads == [{"content": "first-"}] + assert target.sent_messages[0].edits == [ + {"content": "first-second"}, + {"content": "first-second"}, + ] + assert owner._stream_bufs == {} + + @pytest.mark.asyncio async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None: owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus()) diff --git a/nanobot/channels/feishu/runtime.py b/nanobot/channels/feishu/runtime.py index 1541fd97..a9753e6a 100644 --- a/nanobot/channels/feishu/runtime.py +++ b/nanobot/channels/feishu/runtime.py @@ -2232,6 +2232,10 @@ class FeishuChannel(BaseChannel): rid_type = "chat_id" if chat_id.startswith("oc_") else "open_id" # --- stream end: final update or fallback --- + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: message_id = meta.get("message_id") # Only finalize the OnIt -> DONE reaction transition on the truly diff --git a/nanobot/channels/feishu/tests/test_feishu_streaming.py b/nanobot/channels/feishu/tests/test_feishu_streaming.py index 7540d5a3..2b67f5a5 100644 --- a/nanobot/channels/feishu/tests/test_feishu_streaming.py +++ b/nanobot/channels/feishu/tests/test_feishu_streaming.py @@ -285,6 +285,27 @@ class TestSendDelta: settings_call = ch._client.cardkit.v1.card.settings.call_args[0][0] assert settings_call.body.sequence == 5 # after final content seq 4 + @pytest.mark.asyncio + async def test_stream_end_merge_next_preserves_buffer(self): + ch = _make_channel() + ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf( + text="first-", + card_id="card_1", + sequence=3, + last_edit=time.monotonic(), + ) + + await ch.send_delta( + "oc_chat1", + "boundary", + stream_end=True, + merge_next=True, + ) + + assert ch._stream_bufs["oc_chat1"].text == "first-boundary" + ch._client.cardkit.v1.card_element.content.assert_not_called() + ch._client.cardkit.v1.card.settings.assert_not_called() + @pytest.mark.asyncio async def test_stream_end_fallback_when_no_card_id(self): """If card creation failed, stream_end falls back to a plain card message.""" diff --git a/nanobot/channels/matrix/runtime.py b/nanobot/channels/matrix/runtime.py index f0cccfb4..992aba7a 100644 --- a/nanobot/channels/matrix/runtime.py +++ b/nanobot/channels/matrix/runtime.py @@ -602,6 +602,10 @@ class MatrixChannel(BaseChannel): ) -> None: relates_to = self._build_thread_relates_to(metadata) + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: stream_key = _matrix_stream_key(chat_id, stream_id) buf = self._stream_bufs.pop(stream_key, None) diff --git a/nanobot/channels/matrix/tests/test_matrix_channel.py b/nanobot/channels/matrix/tests/test_matrix_channel.py index 58ed088b..bcc6f3af 100644 --- a/nanobot/channels/matrix/tests/test_matrix_channel.py +++ b/nanobot/channels/matrix/tests/test_matrix_channel.py @@ -1937,6 +1937,29 @@ async def test_send_delta_stream_end_replaces_existing_message() -> None: } +@pytest.mark.asyncio +async def test_send_delta_merge_next_preserves_buffer() -> None: + channel = MatrixChannel(_make_config(), MessageBus()) + client = _FakeAsyncClient("", "", "", None) + channel.client = client + channel._stream_bufs["!room:matrix.org"] = matrix_module._StreamBuf( + text="first-", + event_id="event-1", + last_edit=100.0, + ) + channel.monotonic_time = lambda: 100.1 + + await channel.send_delta( + "!room:matrix.org", + "boundary", + stream_end=True, + merge_next=True, + ) + + assert channel._stream_bufs["!room:matrix.org"].text == "first-boundary" + assert client.room_send_calls == [] + + @pytest.mark.asyncio async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None: channel = MatrixChannel(_make_config(), MessageBus()) diff --git a/nanobot/channels/telegram/runtime.py b/nanobot/channels/telegram/runtime.py index 4183a82c..9e42b2df 100644 --- a/nanobot/channels/telegram/runtime.py +++ b/nanobot/channels/telegram/runtime.py @@ -931,6 +931,10 @@ class TelegramChannel(BaseChannel): meta = metadata or {} int_chat_id = int(chat_id) + if stream_end and merge_next: + if not delta: + return + stream_end = False if stream_end: buf = self._stream_bufs.get(chat_id) if not buf or not buf.message_id or not buf.text: diff --git a/nanobot/channels/telegram/tests/test_telegram_channel.py b/nanobot/channels/telegram/tests/test_telegram_channel.py index ff62713a..498aa892 100644 --- a/nanobot/channels/telegram/tests/test_telegram_channel.py +++ b/nanobot/channels/telegram/tests/test_telegram_channel.py @@ -675,6 +675,33 @@ async def test_send_delta_stream_end_raises_and_keeps_buffer_on_failure() -> Non assert "123" in channel._stream_bufs +@pytest.mark.asyncio +async def test_send_delta_merge_next_preserves_buffer() -> None: + channel = TelegramChannel( + TelegramConfig(enabled=True, token="123:abc", allow_from=["*"]), + MessageBus(), + ) + channel._app = _FakeApp(lambda: None) + channel._app.bot.edit_message_text = AsyncMock() + channel._stream_bufs["123"] = _StreamBuf( + text="first-", + message_id=7, + last_edit=float("inf"), + stream_id="s:0", + ) + + await channel.send_delta( + "123", + "boundary", + stream_id="s:0", + stream_end=True, + merge_next=True, + ) + + assert channel._stream_bufs["123"].text == "first-boundary" + channel._app.bot.edit_message_text.assert_not_awaited() + + @pytest.mark.asyncio async def test_send_delta_stream_end_treats_not_modified_as_success() -> None: from telegram.error import BadRequest diff --git a/nanobot/channels/weixin/runtime.py b/nanobot/channels/weixin/runtime.py index 976013c0..ea5d88b5 100644 --- a/nanobot/channels/weixin/runtime.py +++ b/nanobot/channels/weixin/runtime.py @@ -1257,6 +1257,10 @@ class WeixinChannel(BaseChannel): return is_end = stream_end or bool(meta.get("_stream_end")) buffer_key = stream_id or chat_id + if is_end and merge_next: + if delta: + self._stream_buffers.setdefault(buffer_key, []).append(delta) + return # Accumulate intermediate deltas. The stream_end message's own content # (present when the manager coalesces deltas into the end message) is # folded into `full` below instead of appended here, so a send retry diff --git a/nanobot/channels/weixin/tests/test_weixin_channel.py b/nanobot/channels/weixin/tests/test_weixin_channel.py index d0c1c6e4..7058767e 100644 --- a/nanobot/channels/weixin/tests/test_weixin_channel.py +++ b/nanobot/channels/weixin/tests/test_weixin_channel.py @@ -1824,6 +1824,29 @@ async def test_stream_end_flushes_buffered_answer() -> None: assert "wx-user" not in channel._stream_buffers +@pytest.mark.asyncio +async def test_stream_end_merge_next_preserves_buffer_until_final_end() -> None: + channel, _bus = _make_channel() + channel._client = object() + channel._token = "token" + channel._context_tokens["wx-user"] = "ctx-1" + channel._context_token_at["wx-user"] = time.time() + channel._send_text = AsyncMock() + + await channel.send_delta( + "wx-user", + "first-", + stream_id="s1", + stream_end=True, + merge_next=True, + ) + await channel.send_delta("wx-user", "second", stream_id="s1") + await channel.send_delta("wx-user", "", stream_id="s1", stream_end=True) + + channel._send_text.assert_awaited_once_with("wx-user", "first-second", "ctx-1") + assert "s1" not in channel._stream_buffers + + @pytest.mark.asyncio async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None: channel, _bus = _make_channel() diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 39ecb727..220ae646 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -580,6 +580,53 @@ class TestToolEventProgress: 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_length_recovery_streams_non_delta_terminal_segment( + self, + tmp_path: Path, + ) -> None: + bus = MessageBus() + provider = MagicMock() + provider.supports_progress_deltas = True + provider.get_default_model.return_value = "test-model" + call_count = 0 + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + await on_content_delta("first-") + return LLMResponse(content="first-", finish_reason="length") + return LLMResponse(content="second", finish_reason="stop") + + 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)] + final = [m for m in outbound if m.content == "first-second"] + + assert [event.content for event in deltas] == ["first-", "second"] + assert [event.merge_next for event in endings] == [True, False] + assert len(final) == 1 + assert isinstance(final[0].event, StreamedResponseEvent) + @pytest.mark.asyncio async def test_length_recovery_at_max_iterations_streams_only_missing_tail( self,