fix(chat): overlay live parts on every streaming tail message

The streaming tail only overlaid the actively streaming message with live
parts from the sync store. When a turn moved to its next step message, the
previous message fell back to its lagging base record, briefly dropping its
completed tool parts — remounting them and replaying their reveal
animation once the record caught up.

Overlay live parts for every message of the streaming tail via a new
useSessionPartsForMessages hook, guarded so an empty live array never
erases parts the record does have. Also key generate-effect glyphs by
index so appended text does not replay earlier characters, and keep tool
row reveal wrappers mounted unconditionally.

packages/ui/src/components/chat/MessageList.tsx

@@ -19,7 +19,7 @@ import { hasPendingUserSendAnimation, consumePendingUserSendAnimation } from '@/
 import { streamPerfCount, streamPerfMark, streamPerfMeasure } from '@/stores/utils/streamDebug';
 import type { StreamPhase } from './message/types';
 import { useGlobalSessionsStore } from '@/stores/useGlobalSessionsStore';
-import { useSessionParts } from '@/sync/sync-context';
+import { useSessionPartsForMessages } from '@/sync/sync-context';
 import type { ReviewTransferDirection } from '@/lib/reviewFlow';
 import { resolveChatListAnchoredEndSpace, resolveTimelineIsAtEnd } from './lib/scroll/timelineScrollAnchoring';
 import {
@@ -479,6 +479,7 @@ const TurnBlock = React.memo(({
     activeStreamingPhase,
     reviewTransferDirection,
 }: TurnBlockProps) => {
+
     const planModeEnabled = useFeatureFlagsStore((state) => state.planModeEnabled);
     const userMessageHidden = React.useMemo(
         () => isHiddenUserMessage(turn.userMessage, { planModeEnabled }),
@@ -1121,15 +1122,22 @@ const StreamingTailContent: React.FC<{
     activeStreamingPhase,
     reviewTransferDirection,
 }) => {
-    const liveParts = useSessionParts(activeStreamingMessageId ?? '', directory);
+    // Overlay live parts on every message of the tail, not only the one
+    // currently streaming: a finished step message's base record can lag the
+    // part store, and rendering it from that stale snapshot briefly unmounts
+    // its completed tool parts when the stream hands off to the next message.
+    const tailMessageIds = React.useMemo(() => {
+        if (entry.kind === 'turn') return entry.turn.assistantMessageIds;
+        return [entry.message.info.id];
+    }, [entry]);
+    const livePartsByMessageId = useSessionPartsForMessages(tailMessageIds, directory);
     const planModeEnabled = useFeatureFlagsStore((state) => state.planModeEnabled);
     const liveEntry = React.useMemo(() => buildLiveStreamingEntry(entry, {
-        activeStreamingMessageId,
+        livePartsByMessageId,
-        liveParts,
         showTextJustificationActivity: chatRenderMode === 'sorted',
         showTurnChangedFiles,
         mergeHiddenUserTurns: { planModeEnabled },
-    }), [activeStreamingMessageId, chatRenderMode, entry, liveParts, showTurnChangedFiles, planModeEnabled]);
+    }), [chatRenderMode, entry, livePartsByMessageId, showTurnChangedFiles, planModeEnabled]);
 
     return (
         <MessageListEntry

packages/ui/src/components/chat/lib/turns/streamingTailEntry.test.ts

@@ -64,8 +64,7 @@ describe('buildLiveStreamingEntry', () => {
         const entry = turnEntry(assistant);
 
         const next = buildLiveStreamingEntry(entry, {
-            activeStreamingMessageId: 'assistant_other',
+            livePartsByMessageId: { assistant_other: [textPart('part_live', 'live')] },
-            liveParts: [textPart('part_live', 'live')],
             showTextJustificationActivity: true,
             showTurnChangedFiles: false,
         });
@@ -79,8 +78,7 @@ describe('buildLiveStreamingEntry', () => {
         const liveParts = [reasoningPart('part_1_live', 'thinking')];
 
         const next = buildLiveStreamingEntry(entry, {
-            activeStreamingMessageId: 'assistant_1',
+            livePartsByMessageId: { assistant_1: liveParts },
-            liveParts,
             showTextJustificationActivity: true,
             showTurnChangedFiles: false,
         });
@@ -102,8 +100,7 @@ describe('buildLiveStreamingEntry', () => {
         const liveParts = [textPart('part_1_live', 'live')];
 
         const next = buildLiveStreamingEntry(entry, {
-            activeStreamingMessageId: 'assistant_1',
+            livePartsByMessageId: { assistant_1: liveParts },
-            liveParts,
             showTextJustificationActivity: false,
             showTurnChangedFiles: false,
         });
@@ -121,8 +118,7 @@ describe('buildLiveStreamingEntry', () => {
         const synthetic = syntheticTextPart('part_synthetic', 'hidden while streaming');
 
         const next = buildLiveStreamingEntry(entry, {
-            activeStreamingMessageId: 'assistant_1',
+            livePartsByMessageId: { assistant_1: [synthetic, visible] },
-            liveParts: [synthetic, visible],
             showTextJustificationActivity: true,
             showTurnChangedFiles: false,
         });
@@ -131,4 +127,39 @@ describe('buildLiveStreamingEntry', () => {
         if (next.kind !== 'turn') return;
         expect(next.turn.assistantMessages[0]?.parts).toEqual([visible]);
     });
+
+    test('keeps a finished step message on its live parts after the stream moves on', () => {
+        const finished = message('assistant_1', 'assistant', 'user_1', []);
+        const streaming = message('assistant_2', 'assistant', 'user_1', []);
+        const entry = turnEntry(finished);
+        if (entry.kind !== 'turn') return;
+        entry.turn.assistantMessageIds = ['assistant_1', 'assistant_2'];
+        entry.turn.assistantMessages = [finished, streaming];
+        const finishedLive = [textPart('part_tool_done', 'tool output')];
+        const streamingLive = [textPart('part_streaming', 'streaming')];
+
+        const next = buildLiveStreamingEntry(entry, {
+            livePartsByMessageId: { assistant_1: finishedLive, assistant_2: streamingLive },
+            showTextJustificationActivity: true,
+            showTurnChangedFiles: false,
+        });
+
+        expect(next.kind).toBe('turn');
+        if (next.kind !== 'turn') return;
+        expect(next.turn.assistantMessages[0]?.parts).toEqual(finishedLive);
+        expect(next.turn.assistantMessages[1]?.parts).toEqual(streamingLive);
+    });
+
+    test('never erases record parts with an empty live array', () => {
+        const assistant = message('assistant_1', 'assistant', 'user_1', [textPart('part_1', 'kept')]);
+        const entry = turnEntry(assistant);
+
+        const next = buildLiveStreamingEntry(entry, {
+            livePartsByMessageId: { assistant_1: [] },
+            showTextJustificationActivity: true,
+            showTurnChangedFiles: false,
+        });
+
+        expect(next).toBe(entry);
+    });
 });

packages/ui/src/components/chat/lib/turns/streamingTailEntry.ts

@@ -15,8 +15,13 @@ export type StreamingTailEntry =
     | { kind: 'turn'; key: string; turn: TurnRecord; isLastTurn: boolean };
 
 type BuildLiveStreamingEntryOptions = {
-    activeStreamingMessageId: string | null | undefined;
+    // Live parts for EVERY message of the streaming tail, not only the one
-    liveParts: Part[];
+    // currently streaming: when the stream moves to the next step message, the
+    // previous message's base record can still lag behind the part store, and
+    // rendering it from that stale snapshot briefly drops its completed tool
+    // parts — remounting them (and replaying their reveal animation) once the
+    // record catches up.
+    livePartsByMessageId: Readonly<Record<string, Part[]>>;
     showTextJustificationActivity: boolean;
     showTurnChangedFiles: boolean;
     mergeHiddenUserTurns?: { planModeEnabled: boolean };
@@ -24,10 +29,12 @@ type BuildLiveStreamingEntryOptions = {
 
 const withLiveParts = (
     message: ChatMessageEntry,
-    activeStreamingMessageId: string,
+    livePartsByMessageId: Readonly<Record<string, Part[]>>,
-    liveParts: Part[],
 ): ChatMessageEntry => {
-    if (message.info.id !== activeStreamingMessageId || message.parts === liveParts) {
+    const liveParts = livePartsByMessageId[message.info.id];
+    // An empty live array is ambiguous — the store may simply not have loaded
+    // this message's parts — and must never erase parts the record does have.
+    if (!liveParts || liveParts.length === 0 || message.parts === liveParts) {
         return message;
     }
 
@@ -41,13 +48,10 @@ export const buildLiveStreamingEntry = <TEntry extends StreamingTailEntry>(
     entry: TEntry,
     options: BuildLiveStreamingEntryOptions,
 ): TEntry => {
-    const activeStreamingMessageId = options.activeStreamingMessageId;
+    const livePartsByMessageId = options.livePartsByMessageId;
-    if (!activeStreamingMessageId) {
-        return entry;
-    }
 
     if (entry.kind === 'ungrouped') {
-        const message = withLiveParts(entry.message, activeStreamingMessageId, options.liveParts);
+        const message = withLiveParts(entry.message, livePartsByMessageId);
         if (message === entry.message) {
             return entry;
         }
@@ -59,7 +63,7 @@ export const buildLiveStreamingEntry = <TEntry extends StreamingTailEntry>(
 
     let changed = false;
     const assistantMessages = entry.turn.assistantMessages.map((message) => {
-        const next = withLiveParts(message, activeStreamingMessageId, options.liveParts);
+        const next = withLiveParts(message, livePartsByMessageId);
         if (next !== message) {
             changed = true;
         }

packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx

@@ -378,7 +378,6 @@ interface ExpandableToolRowProps {
     onShowPopup: (content: ToolPopupContent) => void;
     onContentChange?: (reason?: ContentChangeReason) => void;
     animateTailText: boolean;
-    animateRows: boolean;
 }
 
 const ExpandableToolRow: React.FC<ExpandableToolRowProps> = ({
@@ -389,7 +388,6 @@ const ExpandableToolRow: React.FC<ExpandableToolRowProps> = ({
     onShowPopup,
     onContentChange,
     animateTailText,
-    animateRows,
 }) => {
     const handleToggle = React.useCallback(() => {
         onToggleTool(activity.id);
@@ -407,17 +405,17 @@ const ExpandableToolRow: React.FC<ExpandableToolRowProps> = ({
         />
     );
 
-    const maybeWrapped = animateTailText ? (
+    // Wrappers are unconditional: a conditional wrapper changes the element
-        <ToolRevealOnMount animate={true} wipe>
+    // type at this position when animateTailText/animateRows flip (message
-            {content}
+    // completion), remounting the tool subtree and replaying the reveal wipe.
-        </ToolRevealOnMount>
+    // Both wrappers are inert with animation off.
-    ) : content;
+    return (
-
+        <FadeInOnReveal>
-    if (!animateRows) {
+            <ToolRevealOnMount animate={animateTailText} wipe>
-        return maybeWrapped;
+                {content}
-    }
+            </ToolRevealOnMount>
-
+        </FadeInOnReveal>
-    return <FadeInOnReveal>{maybeWrapped}</FadeInOnReveal>;
+    );
 };
 
 const MemoExpandableToolRow = React.memo(ExpandableToolRow, (prev, next) => {
@@ -427,7 +425,6 @@ const MemoExpandableToolRow = React.memo(ExpandableToolRow, (prev, next) => {
         && prev.onShowPopup === next.onShowPopup
         && prev.onContentChange === next.onContentChange
         && prev.animateTailText === next.animateTailText
-        && prev.animateRows === next.animateRows
         && prev.activity.id === next.activity.id
         && prev.activity.kind === next.activity.kind
         && prev.activity.endedAt === next.activity.endedAt
@@ -438,14 +435,12 @@ interface StaticGroupedToolRowProps {
     toolName: string;
     activities: TurnActivityPart[];
     animateTailText: boolean;
-    animateRows: boolean;
 }
 
 const StaticGroupedToolRow: React.FC<StaticGroupedToolRowProps> = ({
     toolName,
     activities,
     animateTailText,
-    animateRows,
 }) => {
     const content = (
         <StaticToolRow
@@ -455,23 +450,22 @@ const StaticGroupedToolRow: React.FC<StaticGroupedToolRowProps> = ({
         />
     );
 
-    const maybeWrapped = animateTailText ? (
+    // Wrappers are unconditional: a conditional wrapper changes the element
-        <ToolRevealOnMount animate={true} wipe>
+    // type at this position when animateTailText/animateRows flip (message
-            {content}
+    // completion), remounting the tool subtree and replaying the reveal wipe.
-        </ToolRevealOnMount>
+    // Both wrappers are inert with animation off.
-    ) : content;
+    return (
-
+        <FadeInOnReveal>
-    if (!animateRows) {
+            <ToolRevealOnMount animate={animateTailText} wipe>
-        return maybeWrapped;
+                {content}
-    }
+            </ToolRevealOnMount>
-
+        </FadeInOnReveal>
-    return <FadeInOnReveal>{maybeWrapped}</FadeInOnReveal>;
+    );
 };
 
 const MemoStaticGroupedToolRow = React.memo(StaticGroupedToolRow, (prev, next) => {
     return prev.toolName === next.toolName
         && prev.animateTailText === next.animateTailText
-        && prev.animateRows === next.animateRows
         && areActivityListsEqual(prev.activities, next.activities);
 });
 
@@ -926,7 +920,6 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
                         onShowPopup={onShowPopup}
                         onContentChange={onContentChange}
                         animateTailText={Boolean(animatedToolIds?.has(row.activity.id))}
-                        animateRows={animateRows}
                     />
                 );
 
@@ -937,7 +930,6 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
                         toolName={row.toolName}
                         activities={row.activities}
                         animateTailText={row.activities.some((activity) => animatedToolIds?.has(activity.id))}
-                        animateRows={animateRows}
                     />
                 );
 
@@ -952,7 +944,6 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
                         onShowPopup={onShowPopup}
                         onContentChange={onContentChange}
                         animateTailText={Boolean(animatedToolIds?.has(row.activity.id))}
-                        animateRows={animateRows}
                     />
                 );
 

packages/ui/src/components/ui/text.tsx

@@ -43,7 +43,11 @@ const variants = [
           {textContent.split("").map((char, index) => (
             <motion.span
               {...props}
-              key={char + String(index)}
+              // Index-only: keying by character remounted every span when
+              // the text mutated (a tool title resolving on completion) and
+              // replayed the whole fade. Same-index spans update in place;
+              // appended characters still mount with the reveal.
+              key={index}
               className={cn(
                 "inline-block whitespace-pre align-baseline"
               )}

packages/ui/src/sync/sync-context.tsx

@@ -2647,6 +2647,48 @@ export function useSessionParts(messageID: string, directory?: string) {
   )
 }
 
+const EMPTY_PARTS_BY_MESSAGE: Record<string, Part[]> = {}
+
+/**
+ * Get parts for several messages at once, keyed by message id. The snapshot
+ * keeps its identity until one of the requested part arrays changes, so a
+ * streaming turn can overlay every one of its step messages — not only the
+ * currently streaming one — without tearing between them when the stream
+ * moves to the next message.
+ */
+export function useSessionPartsForMessages(messageIDs: readonly string[], directory?: string): Record<string, Part[]> {
+  const store = useDirectoryStore(directory)
+  const cacheRef = React.useRef<{ ids: readonly string[]; parts: Record<string, Part[]> } | null>(null)
+  const getSnapshot = useCallback(() => {
+    if (messageIDs.length === 0) return EMPTY_PARTS_BY_MESSAGE
+    const state = store.getState()
+    const cached = cacheRef.current
+    if (
+      cached
+      && cached.ids === messageIDs
+      && messageIDs.every((id) => (state.part[id] ?? EMPTY_PARTS) === (cached.parts[id] ?? EMPTY_PARTS))
+    ) {
+      return cached.parts
+    }
+    const parts: Record<string, Part[]> = {}
+    for (const id of messageIDs) parts[id] = state.part[id] ?? EMPTY_PARTS
+    cacheRef.current = { ids: messageIDs, parts }
+    return parts
+  }, [messageIDs, store])
+  const subscribe = useCallback((notify: () => void) => {
+    if (messageIDs.length === 0) return () => undefined
+    return store.subscribe((state, previous) => {
+      for (const id of messageIDs) {
+        if (state.part[id] !== previous.part[id]) {
+          notify()
+          return
+        }
+      }
+    })
+  }, [messageIDs, store])
+  return React.useSyncExternalStore(subscribe, getSnapshot, getSnapshot)
+}
+
 /** Get status for a specific session */
 export function useSessionStatus(sessionID: string, directory?: string) {
   const store = useDirectoryStore(directory)