@@ -156,44 +156,79 @@ function toToolData(tc: NonNullable<ContentBlock['toolCall']>): ToolCallData {
156156 */
157157function parseBlocks ( blocks : ContentBlock [ ] ) : MessageSegment [ ] {
158158 const segments : MessageSegment [ ] = [ ]
159- let group : AgentGroupSegment | null = null
160- const pushGroup = ( nextGroup : AgentGroupSegment , isOpen = false ) => {
161- segments . push ( { ...nextGroup , isOpen } )
159+ const groupsByKey = new Map < string , AgentGroupSegment > ( )
160+ let activeGroupKey : string | null = null
161+
162+ const groupKey = ( name : string , parentToolCallId : string | undefined ) =>
163+ parentToolCallId ? `${ name } :${ parentToolCallId } ` : `${ name } :legacy`
164+
165+ const resolveGroupKey = ( name : string , parentToolCallId : string | undefined ) => {
166+ if ( parentToolCallId ) return groupKey ( name , parentToolCallId )
167+ if ( activeGroupKey && groupsByKey . get ( activeGroupKey ) ?. agentName === name ) {
168+ return activeGroupKey
169+ }
170+ for ( const [ key , g ] of groupsByKey ) {
171+ if ( g . agentName === name && g . isOpen ) return key
172+ }
173+ return groupKey ( name , undefined )
162174 }
163175
164- const lastSubagentBlockIndex = new Map < string , number > ( )
165- for ( let i = 0 ; i < blocks . length ; i ++ ) {
166- const block = blocks [ i ]
167- if ( block . type === 'subagent' && block . content ) {
168- lastSubagentBlockIndex . set ( block . content , i )
176+ const ensureGroup = ( name : string , parentToolCallId : string | undefined ) : AgentGroupSegment => {
177+ const key = resolveGroupKey ( name , parentToolCallId )
178+ let g = groupsByKey . get ( key )
179+ if ( ! g ) {
180+ g = {
181+ type : 'agent_group' ,
182+ id : `agent-${ key } ` ,
183+ agentName : name ,
184+ agentLabel : resolveAgentLabel ( name ) ,
185+ items : [ ] ,
186+ isDelegating : false ,
187+ isOpen : false ,
188+ }
189+ segments . push ( g )
190+ groupsByKey . set ( key , g )
191+ }
192+ return g
193+ }
194+
195+ const findGroupForSubagentChunk = (
196+ parentToolCallId : string | undefined
197+ ) : AgentGroupSegment | undefined => {
198+ if ( parentToolCallId ) {
199+ for ( const [ key , g ] of groupsByKey ) {
200+ if ( key . endsWith ( `:${ parentToolCallId } ` ) ) return g
201+ }
169202 }
203+ if ( activeGroupKey ) return groupsByKey . get ( activeGroupKey )
204+ return undefined
170205 }
171- const hasSubagentBlockAfter = ( name : string , index : number ) : boolean => {
172- const last = lastSubagentBlockIndex . get ( name )
173- return last !== undefined && last > index
206+
207+ const flushLanes = ( ) => {
208+ groupsByKey . clear ( )
209+ activeGroupKey = null
174210 }
175211
176212 for ( let i = 0 ; i < blocks . length ; i ++ ) {
177213 const block = blocks [ i ]
178214
179215 if ( block . type === 'subagent_text' || block . type === 'subagent_thinking' ) {
180- if ( ! block . content || ! group ) continue
181- group . isDelegating = false
182- const lastItem = group . items [ group . items . length - 1 ]
216+ if ( ! block . content ) continue
217+ const g = findGroupForSubagentChunk ( block . parentToolCallId )
218+ if ( ! g ) continue
219+ g . isDelegating = false
220+ const lastItem = g . items [ g . items . length - 1 ]
183221 if ( lastItem ?. type === 'text' ) {
184222 lastItem . content += block . content
185223 } else {
186- group . items . push ( { type : 'text' , content : block . content } )
224+ g . items . push ( { type : 'text' , content : block . content } )
187225 }
188226 continue
189227 }
190228
191229 if ( block . type === 'thinking' ) {
192230 if ( ! block . content ?. trim ( ) ) continue
193- if ( group ) {
194- pushGroup ( group )
195- group = null
196- }
231+ flushLanes ( )
197232 const last = segments [ segments . length - 1 ]
198233 if ( last ?. type === 'thinking' && last . endedAt === undefined ) {
199234 last . content += block . content
@@ -213,21 +248,19 @@ function parseBlocks(blocks: ContentBlock[]): MessageSegment[] {
213248 if ( block . type === 'text' ) {
214249 if ( ! block . content ) continue
215250 if ( block . subagent ) {
216- if ( group && group . agentName === block . subagent ) {
217- group . isDelegating = false
218- const lastItem = group . items [ group . items . length - 1 ]
251+ const g = groupsByKey . get ( resolveGroupKey ( block . subagent , block . parentToolCallId ) )
252+ if ( g ) {
253+ g . isDelegating = false
254+ const lastItem = g . items [ g . items . length - 1 ]
219255 if ( lastItem ?. type === 'text' ) {
220256 lastItem . content += block . content
221257 } else {
222- group . items . push ( { type : 'text' , content : block . content } )
258+ g . items . push ( { type : 'text' , content : block . content } )
223259 }
224260 continue
225261 }
226262 }
227- if ( group ) {
228- pushGroup ( group )
229- group = null
230- }
263+ flushLanes ( )
231264 const last = segments [ segments . length - 1 ]
232265 if ( last ?. type === 'text' ) {
233266 last . content += block . content
@@ -240,34 +273,22 @@ function parseBlocks(blocks: ContentBlock[]): MessageSegment[] {
240273 if ( block . type === 'subagent' ) {
241274 if ( ! block . content ) continue
242275 const key = block . content
243- if ( group && group . agentName === key ) continue
244-
245- const dispatchToolName = SUBAGENT_DISPATCH_TOOLS [ key ]
246276 let inheritedDelegation = false
247- if ( group && dispatchToolName ) {
248- const last : AgentGroupItem | undefined = group . items [ group . items . length - 1 ]
249- if ( last ?. type === 'tool' && last . data . toolName === dispatchToolName ) {
250- inheritedDelegation = ! isToolDone ( last . data . status ) && Boolean ( last . data . streamingArgs )
251- group . items . pop ( )
252- }
253- if ( group . items . length > 0 ) {
254- pushGroup ( group )
277+ const dispatchToolName = SUBAGENT_DISPATCH_TOOLS [ key ]
278+ if ( dispatchToolName ) {
279+ const mship = groupsByKey . get ( groupKey ( 'mothership' , undefined ) )
280+ if ( mship ) {
281+ const last = mship . items [ mship . items . length - 1 ]
282+ if ( last ?. type === 'tool' && last . data . toolName === dispatchToolName ) {
283+ inheritedDelegation = ! isToolDone ( last . data . status ) && Boolean ( last . data . streamingArgs )
284+ mship . items . pop ( )
285+ }
255286 }
256- group = null
257- } else if ( group ) {
258- pushGroup ( group )
259- group = null
260- }
261-
262- group = {
263- type : 'agent_group' ,
264- id : `agent-${ key } -${ i } ` ,
265- agentName : key ,
266- agentLabel : resolveAgentLabel ( key ) ,
267- items : [ ] ,
268- isDelegating : inheritedDelegation ,
269- isOpen : false ,
270287 }
288+ const g = ensureGroup ( key , block . parentToolCallId )
289+ if ( inheritedDelegation ) g . isDelegating = true
290+ g . isOpen = block . endedAt === undefined
291+ activeGroupKey = resolveGroupKey ( key , block . parentToolCallId )
271292 continue
272293 }
273294
@@ -279,98 +300,72 @@ function parseBlocks(blocks: ContentBlock[]): MessageSegment[] {
279300 const isDispatch = SUBAGENT_KEYS . has ( tc . name ) && ! tc . calledBy
280301
281302 if ( isDispatch ) {
282- if ( hasSubagentBlockAfter ( tc . name , i ) ) {
283- continue
284- }
285- if ( ! group || group . agentName !== tc . name ) {
286- if ( group ) {
287- pushGroup ( group )
288- group = null
289- }
290- group = {
291- type : 'agent_group' ,
292- id : `agent-${ tc . name } -${ i } ` ,
293- agentName : tc . name ,
294- agentLabel : resolveAgentLabel ( tc . name ) ,
295- items : [ ] ,
296- isDelegating : false ,
297- isOpen : false ,
298- }
299- }
300- group . isDelegating = isDelegatingTool ( tc )
303+ const g = ensureGroup ( tc . name , tc . id )
304+ g . isDelegating = isDelegatingTool ( tc )
305+ g . isOpen = g . isDelegating
301306 continue
302307 }
303308
304309 const tool = toToolData ( tc )
305310
306- if ( tc . calledBy && group && group . agentName === tc . calledBy ) {
307- group . isDelegating = false
308- group . items . push ( { type : 'tool' , data : tool } )
309- } else if ( tc . calledBy ) {
310- if ( group ) {
311- pushGroup ( group )
312- group = null
313- }
314- group = {
315- type : 'agent_group' ,
316- id : `agent-${ tc . calledBy } -${ i } ` ,
317- agentName : tc . calledBy ,
318- agentLabel : resolveAgentLabel ( tc . calledBy ) ,
319- items : [ { type : 'tool' , data : tool } ] ,
320- isDelegating : false ,
321- isOpen : false ,
322- }
311+ if ( tc . calledBy ) {
312+ const g = ensureGroup ( tc . calledBy , block . parentToolCallId )
313+ g . isDelegating = false
314+ if ( block . parentToolCallId ) g . isOpen = true
315+ g . items . push ( { type : 'tool' , data : tool } )
316+ activeGroupKey = resolveGroupKey ( tc . calledBy , block . parentToolCallId )
323317 } else {
324- if ( group && group . agentName === 'mothership' ) {
325- group . items . push ( { type : 'tool' , data : tool } )
326- } else {
327- if ( group ) {
328- pushGroup ( group )
329- group = null
330- }
331- group = {
332- type : 'agent_group' ,
333- id : `agent-mothership-${ i } ` ,
334- agentName : 'mothership' ,
335- agentLabel : 'Mothership' ,
336- items : [ { type : 'tool' , data : tool } ] ,
337- isDelegating : false ,
338- isOpen : false ,
339- }
340- }
318+ const g = ensureGroup ( 'mothership' , undefined )
319+ g . items . push ( { type : 'tool' , data : tool } )
341320 }
342321 continue
343322 }
344323
345324 if ( block . type === 'options' ) {
346325 if ( ! block . options ?. length ) continue
347- if ( group ) {
348- pushGroup ( group )
349- group = null
350- }
326+ flushLanes ( )
351327 segments . push ( { type : 'options' , items : block . options } )
352328 continue
353329 }
354330
355331 if ( block . type === 'subagent_end' ) {
356- if ( group ) {
357- pushGroup ( group )
358- group = null
332+ if ( block . parentToolCallId ) {
333+ for ( const [ key , g ] of groupsByKey ) {
334+ if ( key . endsWith ( `:${ block . parentToolCallId } ` ) ) {
335+ g . isOpen = false
336+ g . isDelegating = false
337+ }
338+ }
339+ }
340+ if ( block . parentToolCallId && activeGroupKey ?. endsWith ( `:${ block . parentToolCallId } ` ) ) {
341+ activeGroupKey = null
342+ } else if ( ! block . parentToolCallId ) {
343+ for ( const g of groupsByKey . values ( ) ) {
344+ if ( g . agentName !== 'mothership' ) {
345+ g . isOpen = false
346+ g . isDelegating = false
347+ }
348+ }
349+ activeGroupKey = null
359350 }
360351 continue
361352 }
362353
363354 if ( block . type === 'stopped' ) {
364- if ( group ) {
365- pushGroup ( group )
366- group = null
367- }
355+ flushLanes ( )
368356 segments . push ( { type : 'stopped' } )
369357 }
370358 }
371359
372- if ( group ) pushGroup ( group , true )
373- return segments
360+ const visibleSegments = segments . filter (
361+ ( segment ) =>
362+ segment . type !== 'agent_group' ||
363+ segment . items . length > 0 ||
364+ segment . isDelegating ||
365+ segment . isOpen
366+ )
367+
368+ return visibleSegments
374369}
375370
376371/**
@@ -443,12 +438,6 @@ export function MessageContent({
443438 isStreaming &&
444439 ! hasTrailingContent &&
445440 ( lastSegment . type === 'thinking' || hasSubagentEnded || allLastGroupToolsDone )
446- const lastOpenSubagentGroupId = [ ...segments ]
447- . reverse ( )
448- . find (
449- ( segment ) : segment is AgentGroupSegment =>
450- segment . type === 'agent_group' && segment . agentName !== 'mothership' && segment . isOpen
451- ) ?. id
452441
453442 return (
454443 < div className = 'space-y-[10px]' >
@@ -503,8 +492,8 @@ export function MessageContent({
503492 items = { segment . items }
504493 isDelegating = { segment . isDelegating }
505494 isStreaming = { isStreaming }
506- autoCollapse = { allToolsDone && hasFollowingText }
507- defaultExpanded = { segment . id === lastOpenSubagentGroupId }
495+ autoCollapse = { ! segment . isOpen && allToolsDone && hasFollowingText }
496+ defaultExpanded = { segment . isOpen }
508497 />
509498 </ div >
510499 )
0 commit comments