fix(webui): handle final stream image rewrites

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

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

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

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"

webui/src/hooks/useNanobotStream.ts

@@ -520,15 +520,28 @@ export function useNanobotStream(
           resolveActiveAssistantIndex(next)
           ?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current)
           ?? findLatestAssistantAnswerIndex(next);
-        if (targetIndex !== null) {
+          if (targetIndex !== null) {
-          const target = next[targetIndex];
+            const target = next[targetIndex];
-          next = replaceMessageAt(next, targetIndex, {
+            next = replaceMessageAt(next, targetIndex, {
-            ...target,
+              ...target,
-            content: finalAnswerText,
+              content: finalAnswerText,
-            isStreaming: true,
+              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;
     });

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