Improve assistant message action placement for split responses (#1032)

* Improve assistant message action placement for split responses

* Add toggle for split assistant message actions

* Address split message actions review feedback

* Fix inline assistant actions

---------

Co-authored-by: vhqtvn <8930337+vhqtvn@users.noreply.github.com>
Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>

packages/ui/src/components/chat/message/parts/JustificationBlock.tsx

@@ -23,12 +23,14 @@ interface JustificationBlockProps {
     part: Part;
     messageId: string;
     onContentChange?: (reason?: ContentChangeReason) => void;
+    actions?: React.ReactNode;
 }
 
 const JustificationBlock: React.FC<JustificationBlockProps> = ({
     part,
     messageId,
     onContentChange,
+    actions,
 }) => {
     const chatRenderMode = useUIStore((state) => state.chatRenderMode);
     const partWithText = part as PartWithText;
@@ -49,6 +51,7 @@ const JustificationBlock: React.FC<JustificationBlockProps> = ({
             blockId={part.id || `${messageId}-justification`}
             time={time}
             showDuration={chatRenderMode !== 'sorted'}
+            actions={actions}
         />
     );
 };

packages/ui/src/components/chat/message/parts/ProgressiveGroup.tsx

@@ -37,6 +37,7 @@ interface ProgressiveGroupProps {
     showHeader: boolean;
     animateRows?: boolean;
     animatedToolIds?: Set<string>;
+    renderJustificationActions?: (activity: TurnActivityPart) => React.ReactNode;
 }
 
 const isActivityRunning = (activity: TurnActivityPart): boolean => {
@@ -754,15 +755,17 @@ const InlineReasoningBlock = React.memo(({ activity, onContentChange }: {
 /**
  * Inline justification text block — rendered as normal assistant text between tools.
  */
-const InlineJustificationBlock = React.memo(({ activity, onContentChange }: {
+const InlineJustificationBlock = React.memo(({ activity, onContentChange, actions }: {
     activity: TurnActivityPart;
     onContentChange?: (reason?: ContentChangeReason) => void;
+    actions?: React.ReactNode;
 }) => {
     return (
         <JustificationBlock
             part={activity.part}
             messageId={activity.messageId}
             onContentChange={onContentChange}
+            actions={actions}
         />
     );
 });
@@ -782,6 +785,7 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
     showHeader,
     animateRows = true,
     animatedToolIds,
+    renderJustificationActions,
 }) => {
     void _streamPhase;
     const previewCount = showHeader && !isExpanded
@@ -849,6 +853,7 @@ const ProgressiveGroup: React.FC<ProgressiveGroupProps> = ({
                         <InlineJustificationBlock
                             activity={row.activity}
                             onContentChange={onContentChange}
+                            actions={renderJustificationActions?.(row.activity)}
                         />
                     </>
                 );

packages/ui/src/components/chat/message/parts/ReasoningPart.tsx

@@ -81,6 +81,7 @@ type ReasoningTimelineBlockProps = {
     time?: { start?: number; end?: number };
     showDuration?: boolean;
     isStreaming?: boolean;
+    actions?: React.ReactNode;
 };
 
 export const ReasoningTimelineBlock: React.FC<ReasoningTimelineBlockProps> = ({
@@ -91,6 +92,7 @@ export const ReasoningTimelineBlock: React.FC<ReasoningTimelineBlockProps> = ({
     time,
     showDuration = true,
     isStreaming = false,
+    actions,
 }) => {
     const [isExpanded, setIsExpanded] = React.useState(false);
 
@@ -111,7 +113,7 @@ export const ReasoningTimelineBlock: React.FC<ReasoningTimelineBlockProps> = ({
     }
 
     return (
-        <div className="my-1" data-reasoning-block-id={blockId}>
+        <div className="my-1" data-reasoning-block-id={blockId} data-message-text-export-root="true">
             <div
                 className={cn(
                     'group/tool flex items-center gap-2 pr-2 pl-px py-1.5 rounded-xl cursor-pointer'
@@ -171,13 +173,22 @@ export const ReasoningTimelineBlock: React.FC<ReasoningTimelineBlockProps> = ({
                         outerClassName="max-h-80"
                         className="p-0"
                     >
-                        <MarkdownRenderer
-                            content={text}
-                            messageId={blockId}
-                            isAnimated={false}
-                            isStreaming={isStreaming}
-                            variant="reasoning"
-                        />
+                        <div data-message-text-export-source="true">
+                            <MarkdownRenderer
+                                content={text}
+                                messageId={blockId}
+                                isAnimated={false}
+                                isStreaming={isStreaming}
+                                variant="reasoning"
+                            />
+                        </div>
+                        {actions ? (
+                            <div className="mt-2 mb-1 flex items-center justify-start gap-1.5" data-message-actions="true">
+                                <div className="flex items-center gap-1.5" data-message-action-group="true">
+                                    {actions}
+                                </div>
+                            </div>
+                        ) : null}
                     </ScrollableOverlay>
                 </div>
             )}