Skip to content

Commit 4180a7c

Browse files
committed
fix rendering logic
1 parent 340c154 commit 4180a7c

10 files changed

Lines changed: 268 additions & 176 deletions

File tree

apps/sim/app/workspace/[workspaceId]/home/components/message-content/message-content.tsx

Lines changed: 115 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -156,44 +156,79 @@ function toToolData(tc: NonNullable<ContentBlock['toolCall']>): ToolCallData {
156156
*/
157157
function 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

Comments
 (0)