From 38fb154ffe2287d8bf87d63f5736a55391225fef Mon Sep 17 00:00:00 2001
From: Bohdan Triapitsyn <artmore@protonmail.com>
Date: Wed, 17 Dec 2025 02:38:33 +0200
Subject: [PATCH] feat: improve message body rendering with enhanced expansion
 handling and animation triggers

---
 .../components/chat/message/MessageBody.tsx   |  5 ++--
 .../chat/message/parts/ProgressiveGroup.tsx   | 30 ++++++++++++-------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/packages/ui/src/components/chat/message/MessageBody.tsx b/packages/ui/src/components/chat/message/MessageBody.tsx
index 3395fe46..5eb67ad3 100644
--- a/packages/ui/src/components/chat/message/MessageBody.tsx
+++ b/packages/ui/src/components/chat/message/MessageBody.tsx
@@ -44,7 +44,6 @@ const useMigrationTimer = (
     React.useEffect(() => {
         if (!turnGroupingContext) return;
         if (!turnGroupingContext.isWorking) return;
-        if (turnGroupingContext.isGroupExpanded) return;
         if (!hasPreviewableParts) return;
         if (timerStartedRef.current) return;
 
@@ -71,7 +70,7 @@ const useMigrationTimer = (
 
     React.useEffect(() => {
         if (!turnGroupingContext) return;
-        if (!turnGroupingContext.isWorking || turnGroupingContext.isGroupExpanded || !hasPreviewableParts) {
+        if (!turnGroupingContext.isWorking || !hasPreviewableParts) {
             if (timerRef.current) {
                 window.clearTimeout(timerRef.current);
                 timerRef.current = null;
@@ -611,7 +610,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
     const previewableActivityPartsForMessage = React.useMemo(() => {
         if (!turnGroupingContext) return [];
         if (!shouldShowActivityGroup) return [];
-        if (!turnGroupingContext.isWorking || turnGroupingContext.isGroupExpanded) {
+        if (!turnGroupingContext.isWorking) {
             return [];
         }
 
diff --git a/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx b/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx
index bc55ccac..6410d3bf 100644
--- a/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx
+++ b/packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx
@@ -99,30 +99,38 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
 }) => {
     const previousExpandedRef = React.useRef<boolean | undefined>(isExpanded);
 
+    // Track expansion count to force re-mount of items when group expands from collapsed
+    const [expansionKey, setExpansionKey] = React.useState(0);
+
     React.useEffect(() => {
         if (previousExpandedRef.current === isExpanded) return;
+        const wasCollapsed = previousExpandedRef.current === false;
         previousExpandedRef.current = isExpanded;
         onContentChange?.('structural');
+
+        // Increment key when expanding to trigger fresh animations
+        if (isExpanded && wasCollapsed) {
+            setExpansionKey((k) => k + 1);
+        }
     }, [isExpanded, onContentChange]);
 
+
     const displayParts = React.useMemo(() => {
         if (!isWorking) {
-
-            return sortPartsByTime(parts);
-        }
-
-        if (isExpanded) {
-
             return sortPartsByTime(parts);
         }
 
+        // While turn is working, only show parts that have been "previewed".
+        // Collapsed mode previews them in-chat first, then migrates into Activity.
+        // Summary/Detailed modes skip in-chat preview, but still use the same migration gate.
         return sortPartsByTime(
             parts.filter((activity) => {
                 const partId = activity.part.id;
                 return partId && previewedPartIds.has(activity.id);
             })
         );
-    }, [parts, isWorking, isExpanded, previewedPartIds]);
+    }, [parts, isWorking, previewedPartIds]);
+
 
     const summary = getGroupSummary(displayParts);
     const toolConnections = getToolConnections(displayParts);
@@ -206,10 +214,12 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
                         const partId = activity.part.id || `group-part-${index}`;
                         const connection = toolConnections[partId];
 
+                        const animationKey = `${partId}-exp${expansionKey}`;
+
                         switch (activity.kind) {
                             case 'tool':
                                 return (
-                                    <FadeInOnReveal key={partId}>
+                                    <FadeInOnReveal key={animationKey}>
                                         <ToolPart
                                             part={activity.part as ToolPartType}
                                             isExpanded={expandedTools.has(partId)}
@@ -225,7 +235,7 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
 
                             case 'reasoning':
                                 return (
-                                    <FadeInOnReveal key={partId}>
+                                    <FadeInOnReveal key={animationKey}>
                                         <ReasoningPart
                                             part={activity.part}
                                             messageId={activity.messageId}
@@ -236,7 +246,7 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
 
                             case 'justification':
                                 return (
-                                    <FadeInOnReveal key={partId}>
+                                    <FadeInOnReveal key={animationKey}>
                                         <JustificationBlock
                                             part={activity.part}
                                             messageId={activity.messageId}