diff --git a/src/components/chat/AgentTurn.tsx b/src/components/chat/AgentTurn.tsx index 7c9fa7b..4f2a1f8 100644 --- a/src/components/chat/AgentTurn.tsx +++ b/src/components/chat/AgentTurn.tsx @@ -114,7 +114,7 @@ const StreamingMarkdownBlock = ({ React.useLayoutEffect(() => { setStreamTextState((current) => { - if (current.displayText === text && current.animatedTailLength === 0) { + if (current.displayText === text) { return current; } diff --git a/src/components/chat/AgentWorkspace.test.tsx b/src/components/chat/AgentWorkspace.test.tsx index a28aae5..134d6ac 100644 --- a/src/components/chat/AgentWorkspace.test.tsx +++ b/src/components/chat/AgentWorkspace.test.tsx @@ -7,6 +7,8 @@ import { AgentWorkspace } from "./AgentWorkspace"; import type { Message } from "./GlobalChatbox.types"; const renderCounts = new Map(); +const mountCounts = new Map(); +const unmountCounts = new Map(); const streamingFlags = new Map(); jest.mock("next/image", () => ({ @@ -38,6 +40,12 @@ jest.mock("./GlobalChatbox.parts", () => ({ jest.mock("./AgentTurn", () => ({ AgentTurn: ({ message, isStreaming }: { message: Message; isStreaming: boolean }) => { + React.useEffect(() => { + mountCounts.set(message.id, (mountCounts.get(message.id) ?? 0) + 1); + return () => { + unmountCounts.set(message.id, (unmountCounts.get(message.id) ?? 0) + 1); + }; + }, [message.id]); renderCounts.set(message.id, (renderCounts.get(message.id) ?? 0) + 1); streamingFlags.set(message.id, isStreaming); return
{message.content}
; @@ -62,6 +70,8 @@ describe("AgentWorkspace", () => { beforeEach(() => { renderCounts.clear(); + mountCounts.clear(); + unmountCounts.clear(); streamingFlags.clear(); }); @@ -123,4 +133,39 @@ describe("AgentWorkspace", () => { expect(streamingFlags.get("assistant-1")).toBe(false); expect(streamingFlags.get("assistant-2")).toBe(true); }); + + it("does not remount the streaming assistant turn when streaming finishes", () => { + const userMessage: Message = { + id: "user-1", + role: "user", + content: "question", + }; + const assistantMessage: Message = { + id: "assistant-1", + role: "assistant", + content: "final answer", + }; + + const { rerender } = render( + , + ); + + expect(streamingFlags.get("assistant-1")).toBe(true); + + rerender( + , + ); + + expect(mountCounts.get("assistant-1")).toBe(1); + expect(unmountCounts.get("assistant-1") ?? 0).toBe(0); + expect(streamingFlags.get("assistant-1")).toBe(false); + }); }); diff --git a/src/components/chat/AgentWorkspace.tsx b/src/components/chat/AgentWorkspace.tsx index 0b0dab9..00a63a0 100644 --- a/src/components/chat/AgentWorkspace.tsx +++ b/src/components/chat/AgentWorkspace.tsx @@ -40,6 +40,7 @@ type AgentWorkspaceProps = { type TurnListProps = { messages: Message[]; isStreaming: boolean; + streamingMessageId: string | null; speakingMessageId: string | null; speechState: SpeechState; onSpeak: (messageId: string, text: string) => void; @@ -60,9 +61,12 @@ const sameMessages = (left: Message[], right: Message[]) => left.length === right.length && left.every((message, index) => message === right[index]); +const TurnItem = React.memo(AgentTurn); + const TurnListInner = ({ messages, isStreaming, + streamingMessageId, speakingMessageId, speechState, onSpeak, @@ -78,10 +82,10 @@ const TurnListInner = ({ return ( <> {messages.map((message) => ( - sameMessages(prevProps.messages, nextProps.messages) && prevProps.isStreaming === nextProps.isStreaming && + prevProps.streamingMessageId === nextProps.streamingMessageId && prevProps.speakingMessageId === nextProps.speakingMessageId && prevProps.speechState === nextProps.speechState && prevProps.onSpeak === nextProps.onSpeak && @@ -326,8 +331,6 @@ export const AgentWorkspace = ({ isStreaming && messages.at(-1)?.role === "assistant" ? messages.at(-1) : undefined; - const historyMessages = - streamingMessage !== undefined ? messages.slice(0, -1) : messages; const handleScroll = React.useCallback( (event: React.UIEvent) => { if (!onScrollStateChange) return; @@ -368,8 +371,9 @@ export const AgentWorkspace = ({ {messages.length > 0 ? ( - - {streamingMessage ? ( - - - - ) : null} ) : null}