feat: add activity group anchor message handling and improve standalone tool rendering logic

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

@@ -43,6 +43,10 @@ export interface TurnGroupingContext {
     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,10 +67,17 @@ interface TurnActivityInfo {
     hasReasoning: boolean;
     summaryBody?: string;
     diffStats?: TurnDiffStats;
+    activityGroupAnchorMessageId?: string;
 }
 
 const ENABLE_TEXT_JUSTIFICATION_ACTIVITY = false;
 
+const ACTIVITY_STANDALONE_TOOL_NAMES = new Set<string>(['task']);
+
+const isActivityStandaloneTool = (toolName: unknown): boolean => {
+    return typeof toolName === 'string' && ACTIVITY_STANDALONE_TOOL_NAMES.has(toolName.toLowerCase());
+};
+
 export const detectTurns = (messages: ChatMessageEntry[]): Turn[] => {
     const result: Turn[] = [];
     let currentTurn: Turn | null = null;
@@ -255,12 +266,36 @@ 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>();
+    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);
+    });
+
+    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;
+        }
+    }
+
     return {
         activityParts,
         hasTools,
         hasReasoning,
         summaryBody,
         diffStats,
+        activityGroupAnchorMessageId,
     };
 };
 
@@ -383,6 +418,7 @@ export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingRe
                 hasTools,
                 hasReasoning,
                 diffStats,
+                activityGroupAnchorMessageId: activityInfo?.activityGroupAnchorMessageId,
                 isWorking: isTurnWorking,
                 isGroupExpanded: uiState.isExpanded,
                 previewedPartIds: uiState.previewedPartIds,

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

@@ -94,6 +94,12 @@ const useMigrationTimer = (
     return { isAnimating };
 };
 
+const ACTIVITY_STANDALONE_TOOL_NAMES = new Set<string>(['task']);
+
+const isActivityStandaloneTool = (toolName: unknown): boolean => {
+    return typeof toolName === 'string' && ACTIVITY_STANDALONE_TOOL_NAMES.has(toolName.toLowerCase());
+};
+
 interface MessageBodyProps {
     messageId: string;
     parts: Part[];
@@ -541,7 +547,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
     const toolConnections = React.useMemo(() => {
         const connections: Record<string, { hasPrev: boolean; hasNext: boolean }> = {};
         const displayableTools = toolParts.filter((toolPart) => {
-            if (toolPart.tool === 'task') {
+            if (isActivityStandaloneTool(toolPart.tool)) {
                 return false;
             }
             if (shouldHoldTools) {
@@ -584,13 +590,13 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
             ? activityPartsForTurn.filter((activity) => activity.kind === 'tool')
             : activityPartsForTurn;
 
-        // Task tool gets its own progressive card (not part of Activity group).
+        // Tools rendered standalone are excluded from Activity group.
         return base.filter((activity) => {
             if (activity.kind !== 'tool') {
                 return true;
             }
             const toolName = (activity.part as ToolPartType).tool;
-            return !(typeof toolName === 'string' && toolName.toLowerCase() === 'task');
+            return !isActivityStandaloneTool(toolName);
         });
     }, [activityPartsForTurn, showReasoningTraces, turnGroupingContext]);
 
@@ -629,7 +635,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
 
             if (activity.kind === 'tool') {
                 const toolPart = part as ToolPartType;
-                if (toolPart.tool === 'task') {
+                if (isActivityStandaloneTool(toolPart.tool)) {
                     return;
                 }
                 if (shouldHoldTools) return;
@@ -671,15 +677,59 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
 
     const { isAnimating: isMessageAnimating } = useMigrationTimer(turnGroupingContext, previewableActivityPartIds);
 
+    const shouldRenderActivityGroup = Boolean(
+        turnGroupingContext &&
+            turnGroupingContext.activityGroupAnchorMessageId === messageId &&
+            shouldShowActivityGroup &&
+            visibleActivityPartsForTurn.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[] = [];
 
-        if (
+        const pushActivityGroup = () => {
-            turnGroupingContext &&
+            if (!turnGroupingContext || !shouldRenderActivityGroup) {
-            turnGroupingContext.isFirstAssistantInTurn &&
+                return;
-            shouldShowActivityGroup &&
+            }
-            visibleActivityPartsForTurn.length > 0
-        ) {
             rendered.push(
                 <ProgressiveGroup
                     key="progressive-group"
@@ -697,16 +747,19 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
                     diffStats={turnGroupingContext.diffStats}
                 />
             );
+        };
+
+        if (!shouldPlaceActivityAfterStandaloneTools) {
+            pushActivityGroup();
         }
 
-        // Task tool: show immediately and update progressively from metadata.
+        // Standalone tools: rendered outside Activity group
-        const taskTools = toolParts.filter((toolPart) => toolPart.tool === 'task');
+        standaloneToolParts.forEach((standaloneToolPart) => {
-        taskTools.forEach((taskPart) => {
             rendered.push(
-                <FadeInOnReveal key={`task-${taskPart.id}`}>
+                <FadeInOnReveal key={`standalone-tool-${standaloneToolPart.id}`}>
                     <ToolPart
-                        part={taskPart}
+                        part={standaloneToolPart}
-                        isExpanded={expandedTools.has(taskPart.id)}
+                        isExpanded={expandedTools.has(standaloneToolPart.id)}
                         onToggle={onToggleTool}
                         syntaxTheme={syntaxTheme}
                         isMobile={isMobile}
@@ -718,6 +771,10 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
             );
         });
 
+        if (shouldPlaceActivityAfterStandaloneTools) {
+            pushActivityGroup();
+        }
+
         const partsWithTime: Array<{
             part: Part;
             index: number;
@@ -742,7 +799,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
                 if (activity.kind === 'tool') {
                     const toolPart = part as ToolPartType;
 
-                    if (toolPart.tool === 'task') {
+                    if (isActivityStandaloneTool(toolPart.tool)) {
                         return;
                     }
 
@@ -805,7 +862,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
                 case 'tool': {
                     const toolPart = part as ToolPartType;
 
-                    if (toolPart.tool === 'task') {
+                    if (isActivityStandaloneTool(toolPart.tool)) {
                         break;
                     }
 
@@ -973,6 +1030,9 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         visibleActivityPartsForTurn,
         visibleParts,
         toolParts,
+        standaloneToolParts,
+        shouldRenderActivityGroup,
+        shouldPlaceActivityAfterStandaloneTools,
     ]);
 
     const userMessageId = turnGroupingContext?.turnId;