From 8be258212e0377f8876c079a51ff8473a2e51760 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sat, 23 May 2026 01:43:48 +0800 Subject: [PATCH] fix(webui): handle final stream image rewrites --- nanobot/channels/websocket.py | 5 +++- nanobot/webui/transcript.py | 18 ++++++++++++ tests/channels/test_websocket_channel.py | 34 +++++++++++++++++++++++ tests/utils/test_webui_transcript.py | 11 ++++++++ webui/src/hooks/useNanobotStream.ts | 29 +++++++++++++------ webui/src/tests/useNanobotStream.test.tsx | 22 +++++++++++++++ 6 files changed, 110 insertions(+), 9 deletions(-) diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index d0add2ccd..856274090 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -1744,7 +1744,10 @@ class WebSocketChannel(BaseChannel): stream_key = (chat_id, str(meta.get("_stream_id") or "")) if meta.get("_stream_end"): body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id} - full_text = "".join(self._stream_text_buffers.pop(stream_key, [])) + buffered = self._stream_text_buffers.pop(stream_key, []) + if delta: + buffered.append(delta) + full_text = "".join(buffered) rewritten = self._rewrite_local_markdown_images(full_text) if rewritten != full_text: body["text"] = rewritten diff --git a/nanobot/webui/transcript.py b/nanobot/webui/transcript.py index 7a8073bdf..90f2e3b09 100644 --- a/nanobot/webui/transcript.py +++ b/nanobot/webui/transcript.py @@ -502,6 +502,24 @@ def replay_transcript_to_ui_messages( buffer_message_id = None buffer_parts = [] continue + final_text = rec.get("text") + if isinstance(final_text, str): + if buffer_message_id is None: + buffer_message_id = _new_id("buf", idx) + messages.append( + { + "id": buffer_message_id, + "role": "assistant", + "content": final_text, + "isStreaming": True, + "createdAt": _ts_base + idx, + }, + ) + else: + for i, m in enumerate(messages): + if m.get("id") == buffer_message_id: + messages[i] = {**m, "content": final_text, "isStreaming": True} + break buffer_message_id = None buffer_parts = [] continue diff --git a/tests/channels/test_websocket_channel.py b/tests/channels/test_websocket_channel.py index 21f22d363..f40ec4872 100644 --- a/tests/channels/test_websocket_channel.py +++ b/tests/channels/test_websocket_channel.py @@ -512,6 +512,40 @@ async def test_send_delta_stream_end_rewrites_local_markdown_image(monkeypatch, assert final["text"].startswith("![Diagram](/api/media/") +@pytest.mark.asyncio +async def test_send_delta_stream_end_rewrites_inline_final_text(monkeypatch, tmp_path) -> None: + bus = MagicMock() + workspace = tmp_path / "workspace" + workspace.mkdir() + (workspace / "diagram.png").write_bytes(b"\x89PNG\r\n\x1a\nimage") + media = tmp_path / "media" + + def fake_media_dir(channel: str | None = None): + path = media / channel if channel else media + path.mkdir(parents=True, exist_ok=True) + return path + + monkeypatch.setattr("nanobot.channels.websocket.get_media_dir", fake_media_dir) + channel = WebSocketChannel( + {"enabled": True, "allowFrom": ["*"], "streaming": True}, + bus, + workspace_path=workspace, + ) + mock_ws = AsyncMock() + channel._attach(mock_ws, "chat-1") + + await channel.send_delta( + "chat-1", + "![Diagram](diagram.png)", + {"_stream_delta": True, "_stream_end": True, "_stream_id": "sid"}, + ) + + mock_ws.send.assert_awaited_once() + final = json.loads(mock_ws.send.await_args.args[0]) + assert final["event"] == "stream_end" + assert final["text"].startswith("![Diagram](/api/media/") + + @pytest.mark.asyncio async def test_send_reasoning_delta_emits_streaming_frame() -> None: bus = MagicMock() diff --git a/tests/utils/test_webui_transcript.py b/tests/utils/test_webui_transcript.py index b84035072..6400900a8 100644 --- a/tests/utils/test_webui_transcript.py +++ b/tests/utils/test_webui_transcript.py @@ -55,6 +55,17 @@ def test_replay_augments_assistant_text() -> None: assert msgs[1]["content"] == "![Diagram](/api/media/sig/payload)" +def test_replay_uses_stream_end_final_text() -> None: + msgs = replay_transcript_to_ui_messages( + [ + {"event": "user", "chat_id": "t-img", "text": "draw"}, + {"event": "stream_end", "chat_id": "t-img", "text": "![Diagram](/api/media/sig/payload)"}, + ], + ) + + assert msgs[1]["content"] == "![Diagram](/api/media/sig/payload)" + + def test_replay_file_edit_event_creates_file_activity(tmp_path, monkeypatch) -> None: monkeypatch.setattr("nanobot.config.paths.get_data_dir", lambda: tmp_path) key = "websocket:t-file" diff --git a/webui/src/hooks/useNanobotStream.ts b/webui/src/hooks/useNanobotStream.ts index 90be81733..6efb60a4e 100644 --- a/webui/src/hooks/useNanobotStream.ts +++ b/webui/src/hooks/useNanobotStream.ts @@ -520,15 +520,28 @@ export function useNanobotStream( resolveActiveAssistantIndex(next) ?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current) ?? findLatestAssistantAnswerIndex(next); - if (targetIndex !== null) { - const target = next[targetIndex]; - next = replaceMessageAt(next, targetIndex, { - ...target, - content: finalAnswerText, - isStreaming: true, - }); + if (targetIndex !== null) { + const target = next[targetIndex]; + next = replaceMessageAt(next, targetIndex, { + ...target, + content: finalAnswerText, + isStreaming: true, + }); + } else { + const id = crypto.randomUUID(); + closedAssistantStreamIdsRef.current.add(id); + next = [ + ...next, + { + id, + role: "assistant", + content: finalAnswerText, + isStreaming: true, + createdAt: Date.now(), + }, + ]; + } } - } if (options?.closeAnswerSegment) closeActiveAssistantStream(); return next; }); diff --git a/webui/src/tests/useNanobotStream.test.tsx b/webui/src/tests/useNanobotStream.test.tsx index 8dd10e82a..3467fda33 100644 --- a/webui/src/tests/useNanobotStream.test.tsx +++ b/webui/src/tests/useNanobotStream.test.tsx @@ -1298,6 +1298,28 @@ describe("useNanobotStream", () => { }); }); + it("creates an assistant bubble from final stream_end text without prior delta", () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-stream-end-only", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + + act(() => { + fake.emit("chat-stream-end-only", { + event: "stream_end", + chat_id: "chat-stream-end-only", + text: "![Diagram](/api/media/sig/payload)", + }); + }); + + expect(result.current.messages).toHaveLength(1); + expect(result.current.messages[0]).toMatchObject({ + role: "assistant", + content: "![Diagram](/api/media/sig/payload)", + isStreaming: true, + }); + }); + it("stamps latency on the last assistant bubble from turn_end", () => { const fake = fakeClient(); const { result } = renderHook(() => useNanobotStream("chat-lat", EMPTY_MESSAGES), {