From b850a04e21c9f053df2943e892d744b1ec2611be Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Wed, 24 Dec 2025 03:27:51 +0200 Subject: [PATCH] fix: turn grouping and activity rendering refactor --- .../ui/src/components/chat/ChatMessage.tsx | 46 +++++- packages/ui/src/components/chat/StatusRow.tsx | 16 ++- .../components/chat/hooks/useTurnGrouping.ts | 134 ++++++++++++++---- .../components/chat/message/MessageBody.tsx | 114 +++++++-------- 4 files changed, 214 insertions(+), 96 deletions(-) diff --git a/packages/ui/src/components/chat/ChatMessage.tsx b/packages/ui/src/components/chat/ChatMessage.tsx index 77fe5811..d0af45e1 100644 --- a/packages/ui/src/components/chat/ChatMessage.tsx +++ b/packages/ui/src/components/chat/ChatMessage.tsx @@ -434,6 +434,9 @@ const ChatMessage: React.FC = ({ 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 = ({ 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 = ({ 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); diff --git a/packages/ui/src/components/chat/StatusRow.tsx b/packages/ui/src/components/chat/StatusRow.tsx index a49c4874..83a8ae08 100644 --- a/packages/ui/src/components/chat/StatusRow.tsx +++ b/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 = ({ isWorking, statusText, + isGenericStatus, isWaitingForPermission, wasAborted, abortActive, @@ -122,7 +124,17 @@ export const StatusRow: React.FC = ({ 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(null); @@ -201,10 +213,12 @@ export const StatusRow: React.FC = ({ ) : null} diff --git a/packages/ui/src/components/chat/hooks/useTurnGrouping.ts b/packages/ui/src/components/chat/hooks/useTurnGrouping.ts index aa8a9e1a..2fcbbb78 100644 --- a/packages/ui/src/components/chat/hooks/useTurnGrouping.ts +++ b/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. - // Excludes standalone tools (rendered outside Activity group). - const activityCountByMessage = new Map(); + const activityGroupSegments: TurnActivityGroup[] = []; + + const activityByPart = new WeakMap(); 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(); + const taskOrder: string[] = []; + const partsByAfterTool = new Map(); + + 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(); + 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 = [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, diff --git a/packages/ui/src/components/chat/message/MessageBody.tsx b/packages/ui/src/components/chat/message/MessageBody.tsx index 3c86808b..347e4fa7 100644 --- a/packages/ui/src/components/chat/message/MessageBody.tsx +++ b/packages/ui/src/components/chat/message/MessageBody.tsx @@ -597,6 +597,11 @@ const AssistantMessageBody: React.FC> = ({ 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(); activityPartsForMessage.forEach((activity) => { @@ -605,6 +610,7 @@ const AssistantMessageBody: React.FC> = ({ return map; }, [activityPartsForMessage]); + const visibleActivityPartsForTurn = React.useMemo(() => { if (!turnGroupingContext) return []; @@ -628,7 +634,12 @@ const AssistantMessageBody: React.FC> = ({ 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> = ({ 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(null); - const activityGroupFirstVisibleAtRef = React.useRef(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( - - ); + + 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( + + ); + }); }; - 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( @@ -791,11 +779,9 @@ const AssistantMessageBody: React.FC> = ({ /> ); - }); - if (shouldPlaceActivityAfterStandaloneTools) { - pushActivityGroup(); - } + renderActivitySegments(standaloneToolPart.id); + }); const partsWithTime: Array<{ part: Part; @@ -1026,6 +1012,7 @@ const AssistantMessageBody: React.FC> = ({ return rendered; }, [ activityPartsByPart, + activityGroupSegmentsForMessage, copiedCode, copiedMessage, expandedTools, @@ -1054,7 +1041,6 @@ const AssistantMessageBody: React.FC> = ({ toolParts, standaloneToolParts, shouldRenderActivityGroup, - shouldPlaceActivityAfterStandaloneTools, ]); const userMessageId = turnGroupingContext?.turnId;