fix: turn grouping and activity rendering refactor

2025-12-24 03:27:51 +02:00
parent 5a0a13b71c
commit b850a04e21
4 changed files with 214 additions and 96 deletions
@@ -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);
@@ -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>
@@ -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.
-    // Excludes standalone tools (rendered outside Activity group).
-    const activityCountByMessage = new Map<string, number>();
+    const activityGroupSegments: TurnActivityGroup[] = [];
+
+    const activityByPart = new WeakMap<Part, TurnActivityPart>();
    activityParts.forEach((activity) => {
-        if (activity.kind === 'tool') {
-            const toolName = (activity.part as { tool?: unknown }).tool;
-            if (isActivityStandaloneTool(toolName)) {
-                return;
-            }
-        }
-        activityCountByMessage.set(activity.messageId, (activityCountByMessage.get(activity.messageId) ?? 0) + 1);
+        activityByPart.set(activity.part, activity);
    });

-    let activityGroupAnchorMessageId: string | undefined;
-    let cumulative = 0;
-    for (const msg of turn.assistantMessages) {
-        cumulative += activityCountByMessage.get(msg.info.id) ?? 0;
-        if (cumulative >= 2) {
-            activityGroupAnchorMessageId = msg.info.id;
-            break;
+    const taskMessageById = new Map<string, string>();
+    const taskOrder: string[] = [];
+    const partsByAfterTool = new Map<string | null, TurnActivityPart[]>();
+
+    let currentAfterToolPartId: string | null = null;
+
+    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,
@@ -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
-                    key="progressive-group"
-                    parts={visibleActivityPartsForTurn}
-                    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}
-                />
-            );
+
+            activityGroupSegmentsForMessage
+                .filter((segment) => (segment.afterToolPartId ?? null) === afterToolPartId)
+                .forEach((segment) => {
+                    const visibleSegmentParts = !showReasoningTraces
+                        ? segment.parts.filter((activity) => activity.kind === 'tool')
+                        : segment.parts;
+
+                    if (visibleSegmentParts.length === 0) {
+                        return;
+                    }
+
+                    rendered.push(
+                        <ProgressiveGroup
+                            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) {
-            pushActivityGroup();
-        }
+        // Activity groups and standalone tasks are interleaved in message order.
+        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) {
-            pushActivityGroup();
-        }
+            renderActivitySegments(standaloneToolPart.id);
+        });

        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;