From 3703ba730e36157f01da7b3ce7f894d6c4e25d34 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Tue, 25 Aug 2026 14:07:33 +0300 Subject: [PATCH] fix(chat): overlay live parts on every streaming tail message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../ui/src/components/chat/MessageList.tsx | 18 +++++-- .../chat/lib/turns/streamingTailEntry.test.ts | 47 +++++++++++++--- .../chat/lib/turns/streamingTailEntry.ts | 26 +++++---- .../chat/message/parts/ProgressiveGroup.tsx | 53 ++++++++----------- packages/ui/src/components/ui/text.tsx | 6 ++- packages/ui/src/sync/sync-context.tsx | 42 +++++++++++++++ 6 files changed, 136 insertions(+), 56 deletions(-) diff --git a/packages/ui/src/components/chat/MessageList.tsx b/packages/ui/src/components/chat/MessageList.tsx index 131319cd..7d3c73d5 100644 --- a/packages/ui/src/components/chat/MessageList.tsx +++ b/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, - liveParts, + livePartsByMessageId, showTextJustificationActivity: chatRenderMode === 'sorted', showTurnChangedFiles, mergeHiddenUserTurns: { planModeEnabled }, - }), [activeStreamingMessageId, chatRenderMode, entry, liveParts, showTurnChangedFiles, planModeEnabled]); + }), [chatRenderMode, entry, livePartsByMessageId, showTurnChangedFiles, planModeEnabled]); return ( { const entry = turnEntry(assistant); const next = buildLiveStreamingEntry(entry, { - activeStreamingMessageId: 'assistant_other', - liveParts: [textPart('part_live', 'live')], + livePartsByMessageId: { assistant_other: [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', - liveParts, + livePartsByMessageId: { assistant_1: 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', - liveParts, + livePartsByMessageId: { assistant_1: 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', - liveParts: [synthetic, visible], + livePartsByMessageId: { assistant_1: [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); + }); }); diff --git a/packages/ui/src/components/chat/lib/turns/streamingTailEntry.ts b/packages/ui/src/components/chat/lib/turns/streamingTailEntry.ts index 2bf7dd2a..755a06f3 100644 --- a/packages/ui/src/components/chat/lib/turns/streamingTailEntry.ts +++ b/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; - liveParts: Part[]; + // Live parts for EVERY message of the streaming tail, not only the one + // 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>; showTextJustificationActivity: boolean; showTurnChangedFiles: boolean; mergeHiddenUserTurns?: { planModeEnabled: boolean }; @@ -24,10 +29,12 @@ type BuildLiveStreamingEntryOptions = { const withLiveParts = ( message: ChatMessageEntry, - activeStreamingMessageId: string, - liveParts: Part[], + livePartsByMessageId: Readonly>, ): 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 = ( entry: TEntry, options: BuildLiveStreamingEntryOptions, ): TEntry => { - const activeStreamingMessageId = options.activeStreamingMessageId; - if (!activeStreamingMessageId) { - return entry; - } + const livePartsByMessageId = options.livePartsByMessageId; 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 = ( 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; } diff --git a/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx b/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx index 03faddb5..ea5e1966 100644 --- a/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx +++ b/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 = ({ @@ -389,7 +388,6 @@ const ExpandableToolRow: React.FC = ({ onShowPopup, onContentChange, animateTailText, - animateRows, }) => { const handleToggle = React.useCallback(() => { onToggleTool(activity.id); @@ -407,17 +405,17 @@ const ExpandableToolRow: React.FC = ({ /> ); - const maybeWrapped = animateTailText ? ( - - {content} - - ) : content; - - if (!animateRows) { - return maybeWrapped; - } - - return {maybeWrapped}; + // Wrappers are unconditional: a conditional wrapper changes the element + // type at this position when animateTailText/animateRows flip (message + // completion), remounting the tool subtree and replaying the reveal wipe. + // Both wrappers are inert with animation off. + return ( + + + {content} + + + ); }; 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 = ({ toolName, activities, animateTailText, - animateRows, }) => { const content = ( = ({ /> ); - const maybeWrapped = animateTailText ? ( - - {content} - - ) : content; - - if (!animateRows) { - return maybeWrapped; - } - - return {maybeWrapped}; + // Wrappers are unconditional: a conditional wrapper changes the element + // type at this position when animateTailText/animateRows flip (message + // completion), remounting the tool subtree and replaying the reveal wipe. + // Both wrappers are inert with animation off. + return ( + + + {content} + + + ); }; 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 = ({ onShowPopup={onShowPopup} onContentChange={onContentChange} animateTailText={Boolean(animatedToolIds?.has(row.activity.id))} - animateRows={animateRows} /> ); @@ -937,7 +930,6 @@ const ProgressiveGroup: React.FC = ({ 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 = ({ onShowPopup={onShowPopup} onContentChange={onContentChange} animateTailText={Boolean(animatedToolIds?.has(row.activity.id))} - animateRows={animateRows} /> ); diff --git a/packages/ui/src/components/ui/text.tsx b/packages/ui/src/components/ui/text.tsx index 48f70e5c..11ebafd2 100644 --- a/packages/ui/src/components/ui/text.tsx +++ b/packages/ui/src/components/ui/text.tsx @@ -43,7 +43,11 @@ const variants = [ {textContent.split("").map((char, index) => ( = {} + +/** + * 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 { + const store = useDirectoryStore(directory) + const cacheRef = React.useRef<{ ids: readonly string[]; parts: Record } | 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 = {} + 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)