1- import React , { memo , Suspense } from 'react' ;
2- import { LazyMarkDown } from '@/components/ui/LazyMarkDown' ;
3- import { ThinkingBlock } from './ThinkingBlock' ;
4- import { PromptSuggestions } from './PromptSuggestions' ;
5- import { getToolComponent } from '@/components/chat/tools/registry' ;
1+ import React , { memo } from 'react' ;
2+ import { SegmentView } from './SegmentView' ;
3+ import { WorkedRollup } from './WorkedRollup' ;
64import { buildSegments } from './segmentBuilder' ;
75import { AgentToolsContext } from '@/contexts/AgentToolsContext' ;
86import type { AgentKind , AssistantStreamEvent } from '@/types/chat.types' ;
97import type { ToolAggregate } from '@/types/tools.types' ;
10- import { Spinner } from '@/components/ui/primitives/Spinner' ;
118
129interface MessageRendererProps {
1310 events : AssistantStreamEvent [ ] ;
1411 className ?: string ;
1512 isStreaming ?: boolean ;
1613 chatId ?: string ;
1714 isLastBotMessage ?: boolean ;
15+ durationMs ?: number | null ;
1816 onSuggestionSelect ?: ( suggestion : string ) => void ;
1917 agentKind ?: AgentKind ;
2018}
@@ -25,6 +23,7 @@ const MessageRendererInner: React.FC<MessageRendererProps> = ({
2523 isStreaming = false ,
2624 chatId,
2725 isLastBotMessage = false ,
26+ durationMs = null ,
2827 onSuggestionSelect,
2928 agentKind,
3029} ) => {
@@ -65,68 +64,54 @@ const MessageRendererInner: React.FC<MessageRendererProps> = ({
6564 [ segments ] ,
6665 ) ;
6766
67+ // Split a completed turn at its last tool/thinking segment: everything up to
68+ // it is the collapsible work trace, everything after is the final answer.
69+ // While streaming, the trace stays empty so the live work renders flat.
70+ const { traceSegments, tailSegments } = React . useMemo ( ( ) => {
71+ if ( isStreaming ) return { traceSegments : [ ] , tailSegments : segments } ;
72+ let traceEnd = - 1 ;
73+ for ( let i = segments . length - 1 ; i >= 0 ; i -- ) {
74+ if ( segments [ i ] . kind === 'tool' || segments [ i ] . kind === 'thinking' ) {
75+ traceEnd = i ;
76+ break ;
77+ }
78+ }
79+ if ( traceEnd < 0 ) return { traceSegments : [ ] , tailSegments : segments } ;
80+ return {
81+ traceSegments : segments . slice ( 0 , traceEnd + 1 ) ,
82+ tailSegments : segments . slice ( traceEnd + 1 ) ,
83+ } ;
84+ } , [ segments , isStreaming ] ) ;
85+
6886 return (
6987 < AgentToolsContext value = { agentTools } >
7088 < div className = { className } >
71- { segments . map ( ( segment ) => {
72- switch ( segment . kind ) {
73- case 'text' :
74- return (
75- < div
76- key = { segment . id }
77- className = "prose prose-sm dark:prose-invert max-w-none break-words"
78- >
79- < LazyMarkDown content = { segment . text } />
80- </ div >
81- ) ;
82- case 'thinking' : {
83- return (
84- < div key = { segment . id } className = "mb-2 mt-0.5" >
85- < ThinkingBlock
86- content = { segment . text }
87- isActiveThinking = { segment . eventIndex === activeThinkingIndex }
88- />
89- </ div >
90- ) ;
91- }
92- case 'tool' : {
93- const Component = getToolComponent ( segment . tool . name , agentKind ) ;
94- return (
95- < div key = { segment . id } className = "mb-2 mt-1" >
96- < Suspense
97- fallback = {
98- < div className = "flex items-center gap-2 rounded-lg border border-border/50 px-3 py-2 dark:border-border-dark/50" >
99- < Spinner
100- size = "sm"
101- className = "text-text-quaternary dark:text-text-dark-quaternary"
102- />
103- < span className = "text-xs text-text-tertiary dark:text-text-dark-tertiary" >
104- Loading tool output...
105- </ span >
106- </ div >
107- }
108- >
109- < Component tool = { segment . tool } chatId = { chatId } />
110- </ Suspense >
111- </ div >
112- ) ;
113- }
114- case 'suggestions' : {
115- if ( ! isLastBotMessage || ! onSuggestionSelect ) {
116- return null ;
117- }
118- return (
119- < PromptSuggestions
120- key = { segment . id }
121- suggestions = { segment . suggestions }
122- onSelect = { onSuggestionSelect }
123- />
124- ) ;
125- }
126- default :
127- return null ;
128- }
129- } ) }
89+ { traceSegments . length > 0 && (
90+ < WorkedRollup durationMs = { durationMs } >
91+ { traceSegments . map ( ( segment ) => (
92+ < SegmentView
93+ key = { segment . id }
94+ segment = { segment }
95+ chatId = { chatId }
96+ agentKind = { agentKind }
97+ activeThinkingIndex = { activeThinkingIndex }
98+ isLastBotMessage = { isLastBotMessage }
99+ onSuggestionSelect = { onSuggestionSelect }
100+ />
101+ ) ) }
102+ </ WorkedRollup >
103+ ) }
104+ { tailSegments . map ( ( segment ) => (
105+ < SegmentView
106+ key = { segment . id }
107+ segment = { segment }
108+ chatId = { chatId }
109+ agentKind = { agentKind }
110+ activeThinkingIndex = { activeThinkingIndex }
111+ isLastBotMessage = { isLastBotMessage }
112+ onSuggestionSelect = { onSuggestionSelect }
113+ />
114+ ) ) }
130115 </ div >
131116 </ AgentToolsContext >
132117 ) ;
0 commit comments