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.
This commit is contained in:
Bohdan Triapitsyn
2026-08-25 14:07:33 +03:00
parent 6a944db746
commit 3703ba730e
6 changed files with 136 additions and 56 deletions
@@ -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 (
<MessageListEntry
@@ -64,8 +64,7 @@ describe('buildLiveStreamingEntry', () => {
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);
});
});
@@ -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<Record<string, Part[]>>;
showTextJustificationActivity: boolean;
showTurnChangedFiles: boolean;
mergeHiddenUserTurns?: { planModeEnabled: boolean };
@@ -24,10 +29,12 @@ type BuildLiveStreamingEntryOptions = {
const withLiveParts = (
message: ChatMessageEntry,
activeStreamingMessageId: string,
liveParts: Part[],
livePartsByMessageId: Readonly<Record<string, 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;
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 = <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;
}
@@ -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 ? (
<ToolRevealOnMount animate={true} wipe>
{content}
</ToolRevealOnMount>
) : content;
if (!animateRows) {
return maybeWrapped;
}
return <FadeInOnReveal>{maybeWrapped}</FadeInOnReveal>;
// 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 (
<FadeInOnReveal>
<ToolRevealOnMount animate={animateTailText} wipe>
{content}
</ToolRevealOnMount>
</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 ? (
<ToolRevealOnMount animate={true} wipe>
{content}
</ToolRevealOnMount>
) : content;
if (!animateRows) {
return maybeWrapped;
}
return <FadeInOnReveal>{maybeWrapped}</FadeInOnReveal>;
// 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 (
<FadeInOnReveal>
<ToolRevealOnMount animate={animateTailText} wipe>
{content}
</ToolRevealOnMount>
</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}
/>
);
+5 -1
View File
@@ -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"
)}
+42
View File
@@ -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)