fix(chat): avoid final stream remount

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;
       }
 

src/components/chat/AgentWorkspace.test.tsx

@@ -7,6 +7,8 @@ import { AgentWorkspace } from "./AgentWorkspace";
 import type { Message } from "./GlobalChatbox.types";
 
 const renderCounts = new Map<string, number>();
+const mountCounts = new Map<string, number>();
+const unmountCounts = new Map<string, number>();
 const streamingFlags = new Map<string, boolean>();
 
 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 <div data-testid={`turn-${message.id}`}>{message.content}</div>;
@@ -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(
+      <AgentWorkspace
+        {...defaultProps}
+        isStreaming
+        messages={[userMessage, assistantMessage]}
+      />,
+    );
+
+    expect(streamingFlags.get("assistant-1")).toBe(true);
+
+    rerender(
+      <AgentWorkspace
+        {...defaultProps}
+        isStreaming={false}
+        messages={[userMessage, assistantMessage]}
+      />,
+    );
+
+    expect(mountCounts.get("assistant-1")).toBe(1);
+    expect(unmountCounts.get("assistant-1") ?? 0).toBe(0);
+    expect(streamingFlags.get("assistant-1")).toBe(false);
+  });
 });

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) => (
-        <AgentTurn
+        <TurnItem
           key={message.id}
           message={message}
-          isStreaming={isStreaming}
+          isStreaming={isStreaming && message.id === streamingMessageId}
           messageSpeechState={speakingMessageId === message.id ? speechState : "idle"}
           onSpeak={onSpeak}
           onPause={onPauseSpeech}
@@ -103,6 +107,7 @@ const TurnList = React.memo(
   (prevProps, nextProps) =>
     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<HTMLDivElement>) => {
       if (!onScrollStateChange) return;
@@ -368,26 +371,9 @@ export const AgentWorkspace = ({
           {messages.length > 0 ? (
             <Box sx={{ display: "flex", flexDirection: "column", gap: 2 }}>
               <TurnList
-                messages={historyMessages}
-                isStreaming={false}
-                speakingMessageId={speakingMessageId}
-                speechState={speechState}
-                onSpeak={onSpeak}
-                onPauseSpeech={onPauseSpeech}
-                onResumeSpeech={onResumeSpeech}
-                onStopSpeech={onStopSpeech}
-                isTtsSupported={isTtsSupported}
-                onCreateBranch={onCreateBranch}
-                onReplyPermission={onReplyPermission}
-                onReplyQuestion={onReplyQuestion}
-                onRejectQuestion={onRejectQuestion}
-              />
-
-              {streamingMessage ? (
-                <Box sx={{ width: "100%" }}>
-                  <TurnList
-                    messages={[streamingMessage]}
+                messages={messages}
                 isStreaming={isStreaming}
+                streamingMessageId={streamingMessage?.id ?? null}
                 speakingMessageId={speakingMessageId}
                 speechState={speechState}
                 onSpeak={onSpeak}
@@ -402,8 +388,6 @@ export const AgentWorkspace = ({
               />
             </Box>
           ) : null}
-            </Box>
-          ) : null}
         </>
       )}