fix: turn grouping and activity rendering refactor

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

@@ -434,6 +434,9 @@ const ChatMessage: React.FC<ChatMessageProps> = ({
         return freshnessDetector.shouldAnimateMessage(message.info, currentSessionId || message.info.sessionID);
     }, [message.info, currentSessionId, isUser]);
 
+    // Track if this message should show header to prevent flickering
+    const shouldShowHeaderRef = React.useRef(false);
+
     const previousRole = React.useMemo(() => {
         if (!previousMessage) return null;
         return deriveMessageRole(previousMessage.info);
@@ -444,12 +447,6 @@ const ChatMessage: React.FC<ChatMessageProps> = ({
         return deriveMessageRole(nextMessage.info);
     }, [nextMessage]);
 
-    const shouldShowHeader = React.useMemo(() => {
-        if (isUser) return true;
-        if (!previousRole) return true;
-        return previousRole.isUser;
-    }, [isUser, previousRole]);
-
     const isFollowedByAssistant = React.useMemo(() => {
         if (isUser) return false;
         if (!nextRole) return false;
@@ -466,6 +463,43 @@ const ChatMessage: React.FC<ChatMessageProps> = ({
         return isStreamingMessage ? 'streaming' : 'completed';
     }, [isMessageCompleted, lifecyclePhase, isStreamingMessage]);
 
+    const shouldShowHeader = React.useMemo(() => {
+        if (isUser) return true;
+        
+        // Use turn grouping context if available for more precise control
+        const headerMessageId = turnGroupingContext?.headerMessageId;
+        if (headerMessageId) {
+            // For turn grouping: only show header for the first assistant message in the turn
+            const isFirstAssistantInTurn = message.info.id === headerMessageId;
+            
+            if (isFirstAssistantInTurn) {
+                // For completed messages, always show header (historical messages)
+                if (streamPhase === 'completed') {
+                    return true;
+                }
+                
+                // For streaming messages: show header when streaming starts and keep it visible
+                const isCurrentlyStreaming = streamPhase === 'streaming' || streamPhase === 'cooldown';
+                const hasStartedStreaming = shouldShowHeaderRef.current;
+                
+                // Update the ref when streaming starts
+                if (isCurrentlyStreaming && !hasStartedStreaming) {
+                    shouldShowHeaderRef.current = true;
+                }
+                
+                // Show header if streaming has started or is currently active
+                return hasStartedStreaming || isCurrentlyStreaming;
+            }
+            
+            // For non-first assistant messages, don't show header
+            return false;
+        }
+        
+        // Fallback to original logic when turn grouping is not available
+        if (!previousRole) return true;
+        return previousRole.isUser;
+    }, [isUser, previousRole, turnGroupingContext, streamPhase, message.info]);
+
     const handleCopyCode = React.useCallback((code: string) => {
         navigator.clipboard.writeText(code);
         setCopiedCode(code);

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

@@ -49,6 +49,7 @@ interface StatusRowProps {
   // Working state
   isWorking: boolean;
   statusText: string | null;
+  isGenericStatus?: boolean;
   isWaitingForPermission?: boolean;
   wasAborted?: boolean;
   abortActive?: boolean;
@@ -64,6 +65,7 @@ interface StatusRowProps {
 export const StatusRow: React.FC<StatusRowProps> = ({
   isWorking,
   statusText,
+  isGenericStatus,
   isWaitingForPermission,
   wasAborted,
   abortActive,
@@ -122,7 +124,17 @@ export const StatusRow: React.FC<StatusRowProps> = ({
   const hasTodos = visibleTodos.length > 0;
   // Original logic from ChatInput
   const shouldRenderPlaceholder = !showAbortStatus && (wasAborted || !abortActive);
-  const hasContent = isWorking || hasTodos || showAbortStatus;
+  
+  // Track if placeholder is showing result (done/aborted) to keep StatusRow mounted
+  const [placeholderShowingResult, setPlaceholderShowingResult] = React.useState(false);
+  
+  // Keep StatusRow rendered while:
+  // - isWorking (active session)
+  // - isComplete (showing "Done" result)
+  // - wasAborted (showing "Aborted" result)
+  // - placeholderShowingResult (placeholder still displaying result)
+  // - hasTodos or showAbortStatus
+  const hasContent = isWorking || isComplete || wasAborted || placeholderShowingResult || hasTodos || showAbortStatus;
 
   // Close popover when clicking outside
   const popoverRef = React.useRef<HTMLDivElement>(null);
@@ -201,10 +213,12 @@ export const StatusRow: React.FC<StatusRowProps> = ({
             <WorkingPlaceholder
               key={currentSessionId ?? "no-session"}
               statusText={statusText}
+              isGenericStatus={isGenericStatus}
               isWaitingForPermission={isWaitingForPermission}
               wasAborted={wasAborted}
               completionId={completionId ?? null}
               isComplete={isComplete}
+              onResultVisibilityChange={setPlaceholderShowingResult}
             />
           ) : null}
         </div>

packages/ui/src/components/chat/hooks/useTurnGrouping.ts

@@ -31,6 +31,13 @@ interface TurnDiffStats {
     files: number;
 }
 
+export interface TurnActivityGroup {
+    id: string;
+    anchorMessageId: string;
+    afterToolPartId: string | null;
+    parts: TurnActivityPart[];
+}
+
 export interface TurnGroupingContext {
     turnId: string;
     isFirstAssistantInTurn: boolean;
@@ -39,14 +46,12 @@ export interface TurnGroupingContext {
     summaryBody?: string;
 
     activityParts: TurnActivityPart[];
+    activityGroupSegments: TurnActivityGroup[];
+    headerMessageId?: string;
     hasTools: boolean;
     hasReasoning: boolean;
     diffStats?: TurnDiffStats;
 
-    // Message that should render the Activity group for this turn.
-    // Chosen as the first assistant message where the turn reaches 2+ activities.
-    activityGroupAnchorMessageId?: string;
-
     isWorking: boolean;
     isGroupExpanded: boolean;
 
@@ -63,11 +68,11 @@ interface TurnUiState {
 
 interface TurnActivityInfo {
     activityParts: TurnActivityPart[];
+    activityGroupSegments: TurnActivityGroup[];
     hasTools: boolean;
     hasReasoning: boolean;
     summaryBody?: string;
     diffStats?: TurnDiffStats;
-    activityGroupAnchorMessageId?: string;
 }
 
 const ENABLE_TEXT_JUSTIFICATION_ACTIVITY = false;
@@ -266,36 +271,112 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
         });
     });
 
-    // Pick the first assistant message where the turn reaches 2+ activities.
+    const activityGroupSegments: TurnActivityGroup[] = [];
-    // Excludes standalone tools (rendered outside Activity group).
+
-    const activityCountByMessage = new Map<string, number>();
+    const activityByPart = new WeakMap<Part, TurnActivityPart>();
     activityParts.forEach((activity) => {
-        if (activity.kind === 'tool') {
+        activityByPart.set(activity.part, activity);
-            const toolName = (activity.part as { tool?: unknown }).tool;
-            if (isActivityStandaloneTool(toolName)) {
-                return;
-            }
-        }
-        activityCountByMessage.set(activity.messageId, (activityCountByMessage.get(activity.messageId) ?? 0) + 1);
     });
 
-    let activityGroupAnchorMessageId: string | undefined;
+    const taskMessageById = new Map<string, string>();
-    let cumulative = 0;
+    const taskOrder: string[] = [];
-    for (const msg of turn.assistantMessages) {
+    const partsByAfterTool = new Map<string | null, TurnActivityPart[]>();
-        cumulative += activityCountByMessage.get(msg.info.id) ?? 0;
+
-        if (cumulative >= 2) {
+    let currentAfterToolPartId: string | null = null;
-            activityGroupAnchorMessageId = msg.info.id;
+
-            break;
+    turn.assistantMessages.forEach((msg) => {
+        const messageId = msg.info.id;
+
+        msg.parts.forEach((part) => {
+            if (part.type === 'tool') {
+                const toolName = (part as { tool?: unknown }).tool;
+                if (isActivityStandaloneTool(toolName)) {
+                    const toolPartId = typeof part.id === 'string' && part.id.trim().length > 0
+                        ? part.id
+                        : `${messageId}-task-${taskOrder.length + 1}`;
+
+                    if (!taskMessageById.has(toolPartId)) {
+                        taskMessageById.set(toolPartId, messageId);
+                        taskOrder.push(toolPartId);
+                    }
+
+                    currentAfterToolPartId = toolPartId;
+                    return;
+                }
+            }
+
+            const activity = activityByPart.get(part);
+            if (!activity) {
+                return;
+            }
+
+            if (activity.kind === 'tool') {
+                const toolName = (activity.part as { tool?: unknown }).tool;
+                if (isActivityStandaloneTool(toolName)) {
+                    return;
+                }
+            }
+
+            const list = partsByAfterTool.get(currentAfterToolPartId) ?? [];
+            list.push(activity);
+            partsByAfterTool.set(currentAfterToolPartId, list);
+        });
+    });
+
+    const pickAnchorForStartSegment = (segmentParts: TurnActivityPart[]): string | undefined => {
+        if (segmentParts.length === 0) return undefined;
+
+        const countByMessage = new Map<string, number>();
+        segmentParts.forEach((activity) => {
+            countByMessage.set(activity.messageId, (countByMessage.get(activity.messageId) ?? 0) + 1);
+        });
+
+        let firstWithAny: string | undefined;
+        let cumulative = 0;
+        for (const msg of turn.assistantMessages) {
+            const count = countByMessage.get(msg.info.id) ?? 0;
+            if (count > 0 && !firstWithAny) {
+                firstWithAny = msg.info.id;
+            }
+            cumulative += count;
+            if (cumulative >= 2) {
+                return msg.info.id;
+            }
         }
-    }
+        return firstWithAny;
+    };
+
+    const orderedKeys: Array<string | null> = [null, ...taskOrder];
+
+    orderedKeys.forEach((afterToolPartId) => {
+        const segmentParts = partsByAfterTool.get(afterToolPartId) ?? [];
+        if (segmentParts.length === 0) {
+            return;
+        }
+
+        const anchorMessageId = afterToolPartId === null
+            ? pickAnchorForStartSegment(segmentParts)
+            : taskMessageById.get(afterToolPartId);
+
+        if (!anchorMessageId) {
+            return;
+        }
+
+        activityGroupSegments.push({
+            id: `${turn.turnId}:${anchorMessageId}:${afterToolPartId ?? 'start'}`,
+            anchorMessageId,
+            afterToolPartId,
+            parts: segmentParts,
+        });
+    });
 
     return {
         activityParts,
+        activityGroupSegments,
         hasTools,
         hasReasoning,
         summaryBody,
         diffStats,
-        activityGroupAnchorMessageId,
     };
 };
 
@@ -396,6 +477,7 @@ export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingRe
 
             const activityInfo = turnActivityInfo.get(turn.turnId);
             const activityParts = activityInfo?.activityParts ?? [];
+            const activityGroupSegments = activityInfo?.activityGroupSegments ?? [];
             const hasTools = Boolean(activityInfo?.hasTools);
             const hasReasoning = Boolean(activityInfo?.hasReasoning);
             const summaryBody = activityInfo?.summaryBody;
@@ -405,6 +487,7 @@ export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingRe
             const isFirstAssistantInTurn = messageId === firstAssistantId;
             const lastAssistantId = turn.assistantMessages[turn.assistantMessages.length - 1]?.info.id;
             const isLastAssistantInTurn = messageId === lastAssistantId;
+            const headerMessageId = firstAssistantId;
 
             const uiState = getOrCreateTurnState(turn.turnId);
             const isTurnWorking = sessionIsWorking && lastTurnId === turn.turnId;
@@ -415,10 +498,11 @@ export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingRe
                 isLastAssistantInTurn,
                 summaryBody,
                 activityParts,
+                activityGroupSegments,
+                headerMessageId,
                 hasTools,
                 hasReasoning,
                 diffStats,
-                activityGroupAnchorMessageId: activityInfo?.activityGroupAnchorMessageId,
                 isWorking: isTurnWorking,
                 isGroupExpanded: uiState.isExpanded,
                 previewedPartIds: uiState.previewedPartIds,

packages/ui/src/components/chat/message/MessageBody.tsx

@@ -597,6 +597,11 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         return activityPartsForTurn.filter((activity) => activity.messageId === messageId);
     }, [activityPartsForTurn, messageId, turnGroupingContext]);
 
+    const activityGroupSegmentsForMessage = React.useMemo(() => {
+        if (!turnGroupingContext) return [];
+        return turnGroupingContext.activityGroupSegments.filter((segment) => segment.anchorMessageId === messageId);
+    }, [messageId, turnGroupingContext]);
+
     const activityPartsByPart = React.useMemo(() => {
         const map = new Map<Part, (typeof activityPartsForMessage)[number]>();
         activityPartsForMessage.forEach((activity) => {
@@ -605,6 +610,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         return map;
     }, [activityPartsForMessage]);
 
+
     const visibleActivityPartsForTurn = React.useMemo(() => {
         if (!turnGroupingContext) return [];
 
@@ -628,7 +634,12 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         if (!turnGroupingContext) {
             return;
         }
-        if (visibleActivityPartsForTurn.length > 1) {
+
+        const hasTaskSplitSegments = turnGroupingContext.activityGroupSegments.some(
+            (segment) => segment.afterToolPartId !== null
+        );
+
+        if (visibleActivityPartsForTurn.length > 1 || (hasTaskSplitSegments && visibleActivityPartsForTurn.length > 0)) {
             setHasEverHadMultipleVisibleActivities(true);
         }
     }, [turnGroupingContext, visibleActivityPartsForTurn.length]);
@@ -701,81 +712,58 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
 
     const shouldRenderActivityGroup = Boolean(
         turnGroupingContext &&
-            turnGroupingContext.activityGroupAnchorMessageId === messageId &&
             shouldShowActivityGroup &&
-            visibleActivityPartsForTurn.length > 0
+            visibleActivityPartsForTurn.length > 0 &&
+            activityGroupSegmentsForMessage.length > 0
     );
 
     const standaloneToolParts = React.useMemo(() => {
         return toolParts.filter((toolPart) => isActivityStandaloneTool(toolPart.tool));
     }, [toolParts]);
 
-    const isActivityGroupVisibleNow = React.useMemo(() => {
-        if (!turnGroupingContext || !shouldRenderActivityGroup) {
-            return false;
-        }
-        if (!turnGroupingContext.isWorking) {
-            return true;
-        }
-        const previewed = turnGroupingContext.previewedPartIds;
-        return visibleActivityPartsForTurn.some((activity) => previewed.has(activity.id));
-    }, [shouldRenderActivityGroup, turnGroupingContext, visibleActivityPartsForTurn]);
-
-    const standaloneToolsFirstVisibleAtRef = React.useRef<number | null>(null);
-    const activityGroupFirstVisibleAtRef = React.useRef<number | null>(null);
-
-    React.useEffect(() => {
-        standaloneToolsFirstVisibleAtRef.current = null;
-        activityGroupFirstVisibleAtRef.current = null;
-    }, [messageId]);
-
-    const now = Date.now();
-    if (standaloneToolParts.length > 0 && standaloneToolsFirstVisibleAtRef.current === null) {
-        standaloneToolsFirstVisibleAtRef.current = now;
-    }
-    if (isActivityGroupVisibleNow && activityGroupFirstVisibleAtRef.current === null) {
-        activityGroupFirstVisibleAtRef.current = now;
-    }
-
-    const shouldPlaceActivityAfterStandaloneTools = Boolean(
-        standaloneToolParts.length > 0 &&
-            isActivityGroupVisibleNow &&
-            typeof standaloneToolsFirstVisibleAtRef.current === 'number' &&
-            typeof activityGroupFirstVisibleAtRef.current === 'number' &&
-            activityGroupFirstVisibleAtRef.current > standaloneToolsFirstVisibleAtRef.current
-    );
 
     const renderedParts = React.useMemo(() => {
         const rendered: React.ReactNode[] = [];
 
-        const pushActivityGroup = () => {
+        const renderActivitySegments = (afterToolPartId: string | null) => {
             if (!turnGroupingContext || !shouldRenderActivityGroup) {
                 return;
             }
-            rendered.push(
+
-                <ProgressiveGroup
+            activityGroupSegmentsForMessage
-                    key="progressive-group"
+                .filter((segment) => (segment.afterToolPartId ?? null) === afterToolPartId)
-                    parts={visibleActivityPartsForTurn}
+                .forEach((segment) => {
-                    isExpanded={turnGroupingContext.isGroupExpanded}
+                    const visibleSegmentParts = !showReasoningTraces
-                    onToggle={turnGroupingContext.toggleGroup}
+                        ? segment.parts.filter((activity) => activity.kind === 'tool')
-                    syntaxTheme={syntaxTheme}
+                        : segment.parts;
-                    isMobile={isMobile}
+
-                    expandedTools={expandedTools}
+                    if (visibleSegmentParts.length === 0) {
-                    onToggleTool={onToggleTool}
+                        return;
-                    onShowPopup={onShowPopup}
+                    }
-                    onContentChange={onContentChange}
+
-                    isWorking={turnGroupingContext.isWorking}
+                    rendered.push(
-                    previewedPartIds={turnGroupingContext.previewedPartIds}
+                        <ProgressiveGroup
-                    diffStats={turnGroupingContext.diffStats}
+                            key={`progressive-group-${segment.id}`}
-                />
+                            parts={visibleSegmentParts}
-            );
+                            isExpanded={turnGroupingContext.isGroupExpanded}
+                            onToggle={turnGroupingContext.toggleGroup}
+                            syntaxTheme={syntaxTheme}
+                            isMobile={isMobile}
+                            expandedTools={expandedTools}
+                            onToggleTool={onToggleTool}
+                            onShowPopup={onShowPopup}
+                            onContentChange={onContentChange}
+                            isWorking={turnGroupingContext.isWorking}
+                            previewedPartIds={turnGroupingContext.previewedPartIds}
+                            diffStats={turnGroupingContext.diffStats}
+                        />
+                    );
+                });
         };
 
-        if (!shouldPlaceActivityAfterStandaloneTools) {
+        // Activity groups and standalone tasks are interleaved in message order.
-            pushActivityGroup();
+        renderActivitySegments(null);
-        }
 
-        // Standalone tools: rendered outside Activity group
         standaloneToolParts.forEach((standaloneToolPart) => {
             rendered.push(
                 <FadeInOnReveal key={`standalone-tool-${standaloneToolPart.id}`}>
@@ -791,11 +779,9 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
                     />
                 </FadeInOnReveal>
             );
-        });
 
-        if (shouldPlaceActivityAfterStandaloneTools) {
+            renderActivitySegments(standaloneToolPart.id);
-            pushActivityGroup();
+        });
-        }
 
         const partsWithTime: Array<{
             part: Part;
@@ -1026,6 +1012,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         return rendered;
     }, [
         activityPartsByPart,
+        activityGroupSegmentsForMessage,
         copiedCode,
         copiedMessage,
         expandedTools,
@@ -1054,7 +1041,6 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         toolParts,
         standaloneToolParts,
         shouldRenderActivityGroup,
-        shouldPlaceActivityAfterStandaloneTools,
     ]);
 
     const userMessageId = turnGroupingContext?.turnId;