diff --git a/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.test.tsx b/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.test.tsx index 12843e93bb..5033502768 100644 --- a/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.test.tsx +++ b/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( + , + ); + + // 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 diff --git a/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.tsx b/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.tsx index 75f61a641d..09de20f07a 100644 --- a/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.tsx +++ b/ui/app/(prowler)/lighthouse/_components/chat/message-markdown.tsx @@ -44,6 +44,9 @@ export function MessageMarkdown({ return (
{ 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(); diff --git a/ui/app/(prowler)/lighthouse/_lib/chat-store.ts b/ui/app/(prowler)/lighthouse/_lib/chat-store.ts index 8b8cecf3ed..987f216100 100644 --- a/ui/app/(prowler)/lighthouse/_lib/chat-store.ts +++ b/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 = () => ({}), ): Promise => { 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(); }; diff --git a/ui/changelog.d/lighthouse-stream-end-flicker.fixed.md b/ui/changelog.d/lighthouse-stream-end-flicker.fixed.md new file mode 100644 index 0000000000..008c589084 --- /dev/null +++ b/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