fix(ui): stop Lighthouse replies from blinking when the stream ends (#13180)

ui/app/(prowler)/lighthouse/_components/chat/message-markdown.test.tsx

@@ -1,9 +1,27 @@
 import { render, waitFor } from "@testing-library/react";
+import { renderToStaticMarkup } from "react-dom/server";
 import { describe, expect, it, vi } from "vitest";
 
 import { MessageMarkdown } from "./message-markdown";
 
 describe("MessageMarkdown", () => {
+  describe("when a finished message is first rendered", () => {
+    it("should paint the text without waiting for an effect", () => {
+      // Given
+      const markdown = "The **AWS** resource is private.";
+
+      // When
+      // Server rendering runs no effects, so it shows the first paint.
+      const firstPaint = document.createElement("div");
+      firstPaint.innerHTML = renderToStaticMarkup(
+        <MessageMarkdown text={markdown} />,
+      );
+
+      // Then
+      expect(firstPaint).toHaveTextContent("The AWS resource is private.");
+    });
+  });
+
   describe("when streamed inline code is incomplete", () => {
     it("should not render a code block inside a paragraph", async () => {
       // Given

ui/app/(prowler)/lighthouse/_components/chat/message-markdown.tsx

@@ -44,6 +44,9 @@ export function MessageMarkdown({
   return (
     <div className="lighthouse-markdown max-w-full min-w-0 overflow-x-auto">
       <Streamdown
+        // Streaming mode mounts empty and fills in a transition; finished
+        // messages render static so swapping in the persisted reply never blinks.
+        mode={isStreaming ? "streaming" : "static"}
         parseIncompleteMarkdown
         components={STREAMDOWN_COMPONENTS}
         shikiTheme={["github-light", "github-dark"]}

ui/app/(prowler)/lighthouse/_lib/chat-store.test.ts

@@ -351,6 +351,43 @@ describe("createLighthouseChatStore", () => {
     expect(store.getState().streamState.activeTaskId).toBeNull();
   });
 
+  it("swaps the live reply for the persisted one in a single update", async () => {
+    // Given
+    const store = makeStore();
+    await store.getState().submitMessage("Summarize findings");
+    eventSources[0].emit("message.delta", { content: "Persisted answer" });
+    getMessagesMock.mockResolvedValue({
+      data: [message("message-1", "assistant", "Persisted answer")],
+    });
+    const visibleReplies: string[] = [];
+    store.subscribe((state) => {
+      const hasPersistedReply = state.messages.some(
+        (item) => item.role === "assistant",
+      );
+      const hasLiveReply = state.streamState.assistantText !== "";
+      visibleReplies.push(
+        hasPersistedReply && hasLiveReply
+          ? "both"
+          : hasPersistedReply
+            ? "persisted"
+            : hasLiveReply
+              ? "live"
+              : "none",
+      );
+    });
+
+    // When
+    eventSources[0].emit("message.end", { message_id: "message-1" });
+
+    // Then
+    await vi.waitFor(() => expect(visibleReplies.at(-1)).toBe("persisted"));
+    // Neither both replies nor an empty gap at any point of the swap.
+    const transitions = visibleReplies.filter(
+      (reply, index) => reply !== visibleReplies[index - 1],
+    );
+    expect(transitions).toEqual(["live", "persisted"]);
+  });
+
   it("keeps the persisted user message ratable after a terminal error", async () => {
     // Given
     const store = makeStore();

ui/app/(prowler)/lighthouse/_lib/chat-store.ts

@@ -185,6 +185,10 @@ export function createLighthouseChatStore(
     const refreshMessages = async (
       sessionId: string,
       shouldApply: () => boolean = () => true,
+      // Committed in the same update as the messages.
+      getExtraState: (
+        messages: LighthouseV2Message[],
+      ) => Partial<LighthouseChatState> = () => ({}),
     ): Promise<boolean> => {
       const result = await getLighthouseV2Messages(sessionId);
       // The fetch is async, so a reset (new chat, or archiving this session)
@@ -192,7 +196,7 @@ export function createLighthouseChatStore(
       // repopulating a chat that no longer points at this session.
       if (sessionId !== get().activeSessionId || !shouldApply()) return false;
       if ("data" in result) {
-        set({ messages: result.data });
+        set({ messages: result.data, ...getExtraState(result.data) });
         return true;
       }
       return false;
@@ -206,17 +210,15 @@ export function createLighthouseChatStore(
       // it would erase the new optimistic message and provisional task id.
       const noNewerSubmission = () =>
         !get().isSubmitting && !get().streamState.activeTaskId;
-      const refreshed = await refreshMessages(sessionId, noNewerSubmission);
-      if (refreshed) {
-        const failedOutcomeMessageId =
+      // Clearing the stream with the persisted reply keeps the UI from ever
+      // rendering both replies at once.
+      await refreshMessages(sessionId, noNewerSubmission, (messages) => ({
+        streamState: createInitialLighthouseV2StreamState(),
+        failedOutcomeMessageId:
           event.type === LIGHTHOUSE_V2_SSE_EVENT.ERROR
-            ? findLastPersistedUserMessageId(get().messages)
-            : null;
-        set({
-          streamState: createInitialLighthouseV2StreamState(),
-          failedOutcomeMessageId,
-        });
-      }
+            ? findLastPersistedUserMessageId(messages)
+            : null,
+      }));
       notifyLighthouseV2SessionsChanged();
     };
 

ui/changelog.d/lighthouse-stream-end-flicker.fixed.md

@@ -0,0 +1 @@
+Lighthouse AI reply blanking out for a frame when the response finishes streaming, which made the whole text blink and reappear