fix: Make thinking/reasoning blocks display consistently and make the show justification setting work. (#332)

packages/ui/src/components/chat/contexts/TurnGroupingContext.tsx

@@ -236,8 +236,6 @@ const isActivityStandaloneTool = (toolName: unknown): boolean => {
     return typeof toolName === 'string' && ACTIVITY_STANDALONE_TOOL_NAMES.has(toolName.toLowerCase());
 };
 
-const ENABLE_TEXT_JUSTIFICATION_ACTIVITY = false;
-
 const extractFinalAssistantText = (turn: Turn): string | undefined => {
     for (const assistantMsg of turn.assistantMessages) {
         const infoFinish = (assistantMsg.info as { finish?: string | null | undefined }).finish;
@@ -255,7 +253,7 @@ const extractFinalAssistantText = (turn: Turn): string | undefined => {
     return undefined;
 };
 
-const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
+const getTurnActivityInfo = (turn: Turn, showTextJustificationActivity: boolean): TurnActivityInfo => {
     interface SummaryDiff {
         additions?: number | null | undefined;
         deletions?: number | null | undefined;
@@ -303,13 +301,33 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
         });
     });
 
+    // Find the LAST assistant message that has text content - this is the summary
+    // All other text messages are justification (yapping during work)
+    let lastTextMessageId: string | undefined;
+    for (let i = turn.assistantMessages.length - 1; i >= 0; i--) {
+        const msg = turn.assistantMessages[i];
+        if (!msg) continue;
+        const hasText = msg.parts.some((p) => {
+            if (p.type !== 'text') return false;
+            const text = (p as { text?: string; content?: string }).text ?? 
+                        (p as { text?: string; content?: string }).content;
+            return typeof text === 'string' && text.trim().length > 0;
+        });
+        if (hasText) {
+            lastTextMessageId = msg.info.id;
+            break;
+        }
+    }
+
     const activityParts: TurnActivityPart[] = [];
     let syntheticIdCounter = 0;
 
     turn.assistantMessages.forEach((msg) => {
         const messageId = msg.info.id;
-        const infoFinish = (msg.info as { finish?: string | null | undefined }).finish;
+        
-        const hasStopFinishInMessage = ENABLE_TEXT_JUSTIFICATION_ACTIVITY ? infoFinish === 'stop' : false;
+        // Only the LAST message with text is the summary (not justification)
+        // All earlier text messages are justification
+        const isFinalSummaryMessage = messageId === lastTextMessageId;
 
         msg.parts.forEach((part) => {
             const baseId = (typeof part.id === 'string' && part.id.trim().length > 0)
@@ -351,10 +369,10 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
             }
 
             if (
-                ENABLE_TEXT_JUSTIFICATION_ACTIVITY &&
+                showTextJustificationActivity &&
                 part.type === 'text' &&
                 (hasTools || hasReasoning) &&
-                !hasStopFinishInMessage
+                !isFinalSummaryMessage
             ) {
                 const text = (part as { text?: string | null | undefined; content?: string | null | undefined }).text ??
                     (part as { text?: string | null | undefined; content?: string | null | undefined }).content;
@@ -483,9 +501,10 @@ const buildNeighborMap = (messages: ChatMessageEntry[]): Map<string, NeighborInf
 export const TurnGroupingProvider: React.FC<TurnGroupingProviderProps> = ({ messages, children }) => {
     const { isWorking: sessionIsWorking } = useCurrentSessionActivity();
     const toolCallExpansion = useUIStore((state) => state.toolCallExpansion);
+    const showTextJustificationActivity = useUIStore((state) => state.showTextJustificationActivity);
     const defaultActivityExpanded = toolCallExpansion === 'activity' || toolCallExpansion === 'detailed';
 
-    // Static data - only changes when messages change
+    // Static data - only changes when messages change or justification setting changes
     const staticValue = React.useMemo<TurnGroupingStaticData>(() => {
         const turns = detectTurns(messages);
         const lastTurnId = turns.length > 0 ? turns[turns.length - 1]!.turnId : null;
@@ -500,7 +519,7 @@ export const TurnGroupingProvider: React.FC<TurnGroupingProviderProps> = ({ mess
 
         const turnActivityInfo = new Map<string, TurnActivityInfo>();
         turns.forEach((turn) => {
-            turnActivityInfo.set(turn.turnId, getTurnActivityInfo(turn));
+            turnActivityInfo.set(turn.turnId, getTurnActivityInfo(turn, showTextJustificationActivity));
         });
 
         const messageNeighbors = buildNeighborMap(messages);
@@ -524,7 +543,7 @@ export const TurnGroupingProvider: React.FC<TurnGroupingProviderProps> = ({ mess
             defaultActivityExpanded,
             messageNeighbors,
         };
-    }, [messages, defaultActivityExpanded]);
+    }, [messages, defaultActivityExpanded, showTextJustificationActivity]);
 
     // UI state for expansion toggles
     const [turnUiStates, setTurnUiStates] = React.useState<Map<string, { isExpanded: boolean }>>(

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

@@ -120,8 +120,7 @@ const extractFinalAssistantText = (turn: Turn): string | undefined => {
     return undefined;
 };
 
-const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
+const getTurnActivityInfo = (turn: Turn, showTextJustificationActivity: boolean): TurnActivityInfo => {
-    const showTextJustificationActivity = useUIStore.getState().showTextJustificationActivity;
     interface SummaryDiff {
         additions?: number | null | undefined;
         deletions?: number | null | undefined;
@@ -179,15 +178,33 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
         });
     });
 
+    // Find the LAST assistant message that has text content - this is the summary
+    // All other text messages are justification (yapping during work)
+    let lastTextMessageId: string | undefined;
+    for (let i = turn.assistantMessages.length - 1; i >= 0; i--) {
+        const msg = turn.assistantMessages[i];
+        if (!msg) continue;
+        const hasText = msg.parts.some((p) => {
+            if (p.type !== 'text') return false;
+            const text = (p as { text?: string; content?: string }).text ?? 
+                        (p as { text?: string; content?: string }).content;
+            return typeof text === 'string' && text.trim().length > 0;
+        });
+        if (hasText) {
+            lastTextMessageId = msg.info.id;
+            break;
+        }
+    }
+
     const activityParts: TurnActivityPart[] = [];
     let syntheticIdCounter = 0;
 
     turn.assistantMessages.forEach((msg) => {
         const messageId = msg.info.id;
-        const infoFinish = (msg.info as { finish?: string | null | undefined }).finish;
+        
-        const hasStopFinishInMessage = showTextJustificationActivity
+        // Only the LAST message with text is the summary (not justification)
-            ? infoFinish === 'stop'
+        // All earlier text messages are justification
-            : false;
+        const isFinalSummaryMessage = messageId === lastTextMessageId;
 
         msg.parts.forEach((part) => {
             const baseId =
@@ -235,7 +252,7 @@ const getTurnActivityInfo = (turn: Turn): TurnActivityInfo => {
                 showTextJustificationActivity &&
                 part.type === 'text' &&
                 (hasTools || hasReasoning) &&
-                !hasStopFinishInMessage
+                !isFinalSummaryMessage
             ) {
                 const text =
                     (part as { text?: string | null | undefined; content?: string | null | undefined }).text ??
@@ -375,6 +392,7 @@ interface UseTurnGroupingResult {
 
 export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingResult => {
     const { isWorking: sessionIsWorking } = useCurrentSessionActivity();
+    const showTextJustificationActivity = useUIStore((state) => state.showTextJustificationActivity);
 
     const turns = React.useMemo(() => detectTurns(messages), [messages]);
 
@@ -397,10 +415,10 @@ export const useTurnGrouping = (messages: ChatMessageEntry[]): UseTurnGroupingRe
     const turnActivityInfo = React.useMemo(() => {
         const map = new Map<string, TurnActivityInfo>();
         turns.forEach((turn) => {
-            map.set(turn.turnId, getTurnActivityInfo(turn));
+            map.set(turn.turnId, getTurnActivityInfo(turn, showTextJustificationActivity));
         });
         return map;
-    }, [turns]);
+    }, [turns, showTextJustificationActivity]);
 
     const [turnUiStates, setTurnUiStates] = React.useState<Map<string, TurnUiState>>(
         () => new Map()

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

@@ -4,6 +4,7 @@ import type { Part } from '@opencode-ai/sdk/v2';
 import UserTextPart from './parts/UserTextPart';
 import ToolPart from './parts/ToolPart';
 import ProgressiveGroup from './parts/ProgressiveGroup';
+import ReasoningPart from './parts/ReasoningPart';
 import { MessageFilesDisplay } from '../FileAttachment';
 import type { ToolPart as ToolPartType } from '@opencode-ai/sdk/v2';
 import type { StreamPhase, ToolPopupContent, AgentMentionInfo } from './types';
@@ -556,8 +557,11 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
     const visibleActivityPartsForTurn = React.useMemo(() => {
         if (!turnGroupingContext) return [];
 
+        // Filter out reasoning if showReasoningTraces is off.
+        // Justification parts are already filtered at the source (useTurnGrouping)
+        // based on showTextJustificationActivity, so we keep them here.
         const base = !showReasoningTraces
-            ? activityPartsForTurn.filter((activity) => activity.kind === 'tool')
+            ? activityPartsForTurn.filter((activity) => activity.kind !== 'reasoning')
             : activityPartsForTurn;
 
         // Tools rendered standalone are excluded from Activity group.
@@ -581,10 +585,18 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
             (segment) => segment.afterToolPartId !== null
         );
 
-        if (visibleActivityPartsForTurn.length > 1 || (hasTaskSplitSegments && visibleActivityPartsForTurn.length > 0)) {
+        const hasReasoningActivity = visibleActivityPartsForTurn.some(
+            (activity) => activity.kind === 'reasoning'
+        );
+
+        if (
+            visibleActivityPartsForTurn.length > 1 ||
+            (hasTaskSplitSegments && visibleActivityPartsForTurn.length > 0) ||
+            hasReasoningActivity
+        ) {
             setHasEverHadMultipleVisibleActivities(true);
         }
-    }, [turnGroupingContext, visibleActivityPartsForTurn.length]);
+    }, [turnGroupingContext, visibleActivityPartsForTurn]);
 
     const shouldShowActivityGroup = Boolean(turnGroupingContext && hasEverHadMultipleVisibleActivities);
 
@@ -611,8 +623,11 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
             activityGroupSegmentsForMessage
                 .filter((segment) => (segment.afterToolPartId ?? null) === afterToolPartId)
                 .forEach((segment) => {
+                    // Filter out reasoning if showReasoningTraces is off.
+                    // Justification parts are already filtered at the source (useTurnGrouping)
+                    // based on showTextJustificationActivity, so we keep them here.
                     const visibleSegmentParts = !showReasoningTraces
-                        ? segment.parts.filter((activity) => activity.kind === 'tool')
+                        ? segment.parts.filter((activity) => activity.kind !== 'reasoning')
                         : segment.parts;
 
                     if (visibleSegmentParts.length === 0) {
@@ -638,6 +653,8 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         };
 
         // Activity groups and standalone tasks are interleaved in message order.
+        // Note: Reasoning parts are rendered in the visibleParts.forEach loop below
+        // when Activity group isn't showing, to maintain proper ordering with tools.
         renderActivitySegments(null);
 
         standaloneToolParts.forEach((standaloneToolPart) => {
@@ -715,6 +732,23 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
 
                     element = toolElement;
                     endTime = isFinalized && typeof time?.end === 'number' ? time.end : null;
+                } else if (activity.kind === 'reasoning' && showReasoningTraces) {
+                    // Fallback rendering for reasoning when Activity group isn't shown
+                    const time = (part as { time?: { end?: number | null | undefined } | null | undefined }).time;
+                    const partEndTime = typeof time?.end === 'number' ? time.end : null;
+
+                    const reasoningElement = (
+                        <FadeInOnReveal key={`reasoning-${activity.id}`}>
+                            <ReasoningPart
+                                part={part}
+                                messageId={messageId}
+                                onContentChange={onContentChange}
+                            />
+                        </FadeInOnReveal>
+                    );
+
+                    element = reasoningElement;
+                    endTime = partEndTime;
                 }
 
                 if (element) {
@@ -752,6 +786,7 @@ const AssistantMessageBody: React.FC<Omit<MessageBodyProps, 'isUser'>> = ({
         expandedTools,
         isMobile,
         isToolFinalized,
+        messageId,
         onContentChange,
         onShowPopup,
         onToggleTool,

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

@@ -33,8 +33,8 @@ const JustificationBlock: React.FC<JustificationBlockProps> = ({
     const rawText = partWithText.text || partWithText.content || '';
     const textContent = React.useMemo(() => cleanJustificationText(rawText), [rawText]);
 
-    const timeInfo = 'time' in part ? (part.time as { start: number; end?: number }) : null;
+    // Don't render if there's no text content
-    if (!timeInfo?.end) {
+    if (!textContent || textContent.trim().length === 0) {
         return null;
     }
 

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

@@ -160,8 +160,9 @@ const ReasoningPart: React.FC<ReasoningPartProps> = ({
     const rawText = partWithText.text || partWithText.content || '';
     const textContent = React.useMemo(() => cleanReasoningText(rawText), [rawText]);
 
-    const timeInfo = 'time' in part ? (part.time as { start: number; end?: number }) : null;
+    // Show reasoning even if time.end isn't set yet (during streaming)
-    if (!timeInfo?.end) {
+    // Only hide if there's no text content
+    if (!textContent || textContent.trim().length === 0) {
         return null;
     }
 

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

@@ -10,6 +10,7 @@ import { toolDisplayStyles } from '@/lib/typography';
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
 import { useDirectoryStore } from '@/stores/useDirectoryStore';
 import { useSessionStore } from '@/stores/useSessionStore';
+import { useUIStore } from '@/stores/useUIStore';
 import { ScrollableOverlay } from '@/components/ui/ScrollableOverlay';
 import type { ContentChangeReason } from '@/hooks/useChatScrollManager';
 
@@ -1068,6 +1069,22 @@ const ToolPart: React.FC<ToolPartProps> = ({ part, isExpanded, onToggle, syntaxT
     const description = getToolDescription(part, state, isMobile, currentDirectory);
     const displayName = getToolMetadata(part.tool).displayName;
     
+    // Get justification text (tool title/description) when setting is enabled
+    const showTextJustificationActivity = useUIStore((state) => state.showTextJustificationActivity);
+    const justificationText = React.useMemo(() => {
+        if (!showTextJustificationActivity) return null;
+        // Get title or description from state - this is the "yapping" text like "Shows system information"
+        const title = (stateWithData as { title?: string }).title;
+        if (typeof title === 'string' && title.trim().length > 0) {
+            return title;
+        }
+        const inputDesc = input?.description;
+        if (typeof inputDesc === 'string' && inputDesc.trim().length > 0) {
+            return inputDesc;
+        }
+        return null;
+    }, [showTextJustificationActivity, stateWithData, input]);
+
     const runtime = React.useContext(RuntimeAPIContext);
 
     const handleMainClick = (e: React.MouseEvent) => {
@@ -1147,7 +1164,12 @@ const ToolPart: React.FC<ToolPartProps> = ({ part, isExpanded, onToggle, syntaxT
                 </div>
 
                 <div className="flex items-center gap-1 flex-1 min-w-0 typography-meta" style={{ color: 'var(--tools-description)' }}>
-                    {description && (
+                    {justificationText && (
+                        <span className={cn("truncate italic", isMobile && "max-w-[120px]")} style={{ color: 'var(--tools-description)', opacity: 0.8 }}>
+                            {justificationText}
+                        </span>
+                    )}
+                    {!justificationText && description && (
                         <span className={cn("truncate", isMobile && "max-w-[120px]")}>
                             {description}
                         </span>