Skip to content

Commit aff0d3b

Browse files
brandonkachenbrandonkachen
committed
feat: implement nested agent display with parent-child tracking
This adds hierarchical agent nesting support so spawned agents appear inside their parent agent blocks in the UI, making agent relationships clearer. 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent a802691 commit aff0d3b

File tree

7 files changed

+89
-20
lines changed

7 files changed

+89
-20
lines changed

‎backend/src/tools/handlers/tool/spawn-agent-utils.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,7 @@ export async function executeSubagent(
338338
agentType: agentTemplate.id,
339339
displayName: agentTemplate.displayName,
340340
onlyChild: isOnlyChild,
341+
parentAgentId: parentAgentState.agentId,
341342
})
342343

343344
// Import loopAgentSteps dynamically to avoid circular dependency
@@ -354,6 +355,7 @@ export async function executeSubagent(
354355
agentType: agentTemplate.id,
355356
displayName: agentTemplate.displayName,
356357
onlyChild: isOnlyChild,
358+
parentAgentId: parentAgentState.agentId,
357359
})
358360

359361
if (result.agentState.runId) {

‎backend/src/tools/handlers/tool/spawn-agents.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,18 @@ export const handleSpawnAgents = ((
158158
return
159159
}
160160

161+
// Don't overwrite agentId for events that already have the correct agent ID
162+
// (subagent_start/finish, tool_call, tool_result from nested agents)
163+
if (
164+
chunk.type === 'subagent_start' ||
165+
chunk.type === 'subagent_finish' ||
166+
chunk.type === 'tool_call' ||
167+
chunk.type === 'tool_result'
168+
) {
169+
writeToClient(chunk)
170+
return
171+
}
172+
161173
const eventWithAgent = {
162174
...chunk,
163175
agentId: subAgentState.agentId,

‎backend/src/tools/stream-parser.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,7 @@ export async function processStreamWithTools(
160160
toolCallId,
161161
toolName,
162162
input,
163+
agentId: agentState.agentId,
163164
})
164165
} else {
165166
// First non-str_replace tool marks end of str_replace phase

‎backend/src/tools/tool-executor.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ export function executeToolCall<T extends ToolName>(
193193
toolCallId: toolCall.toolCallId,
194194
toolName,
195195
input: toolCall.input,
196+
agentId: state.agentState?.agentId,
196197
})
197198

198199
toolCalls.push(toolCall)
@@ -429,6 +430,7 @@ export async function executeCustomToolCall(
429430
toolCallId: toolCall.toolCallId,
430431
toolName,
431432
input: toolCall.input,
433+
agentId: state.agentState?.agentId,
432434
})
433435

434436
toolCalls.push(toolCall)

‎cli/src/hooks/use-send-message.ts‎

Lines changed: 68 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,26 @@ const hiddenToolNames = new Set<ToolName | 'spawn_agent_inline'>([
2424
'spawn_agents',
2525
])
2626

27+
// Helper function to recursively update blocks
28+
const updateBlocksRecursively = (
29+
blocks: ContentBlock[],
30+
targetAgentId: string,
31+
updateFn: (block: ContentBlock) => ContentBlock,
32+
): ContentBlock[] => {
33+
return blocks.map((block) => {
34+
if (block.type === 'agent' && block.agentId === targetAgentId) {
35+
return updateFn(block)
36+
}
37+
if (block.type === 'agent' && block.blocks) {
38+
return {
39+
...block,
40+
blocks: updateBlocksRecursively(block.blocks, targetAgentId, updateFn),
41+
}
42+
}
43+
return block
44+
})
45+
}
46+
2747
interface UseSendMessageOptions {
2848
setMessages: React.Dispatch<React.SetStateAction<ChatMessage[]>>
2949
setFocusedAgentId: (id: string | null) => void
@@ -171,8 +191,11 @@ export const useSendMessage = ({
171191
setMessages((prev) =>
172192
prev.map((msg) => {
173193
if (msg.id === aiMessageId && msg.blocks) {
174-
const newBlocks = msg.blocks.map((block) => {
175-
if (block.type === 'agent' && block.agentId === agentId) {
194+
// Use recursive update to handle nested agents
195+
const newBlocks = updateBlocksRecursively(
196+
msg.blocks,
197+
agentId,
198+
(block) => {
176199
const agentBlocks: ContentBlock[] = block.blocks
177200
? [...block.blocks]
178201
: []
@@ -223,9 +246,9 @@ export const useSendMessage = ({
223246
})
224247
return { ...block, blocks: [...agentBlocks, update] }
225248
}
226-
}
227-
return block
228-
})
249+
return block
250+
},
251+
)
229252
return { ...msg, blocks: newBlocks }
230253
}
231254
return msg
@@ -467,6 +490,8 @@ export const useSendMessage = ({
467490
logger.info('subagent_start event', {
468491
agentId: event.agentId,
469492
agentType: event.agentType,
493+
parentAgentId: event.parentAgentId || 'ROOT',
494+
hasParentAgentId: !!event.parentAgentId,
470495
})
471496
activeSubagentsRef.current.add(event.agentId)
472497

@@ -529,6 +554,7 @@ export const useSendMessage = ({
529554
{
530555
agentId: event.agentId,
531556
agentType: event.agentType,
557+
parentAgentId: event.parentAgentId || 'ROOT',
532558
},
533559
)
534560
setMessages((prev) =>
@@ -551,6 +577,27 @@ export const useSendMessage = ({
551577
initialPrompt: '',
552578
}
553579

580+
// If parentAgentId exists, nest inside parent agent
581+
if (event.parentAgentId) {
582+
logger.info('Nesting agent inside parent', {
583+
childId: event.agentId,
584+
parentId: event.parentAgentId,
585+
})
586+
const updatedBlocks = updateBlocksRecursively(
587+
blocks,
588+
event.parentAgentId,
589+
(parentBlock) => ({
590+
...parentBlock,
591+
blocks: [
592+
...(parentBlock.blocks || []),
593+
newAgentBlock,
594+
],
595+
}),
596+
)
597+
return { ...msg, blocks: updatedBlocks }
598+
}
599+
600+
// No parent - add to top level
554601
return {
555602
...msg,
556603
blocks: [...blocks, newAgentBlock],
@@ -572,15 +619,12 @@ export const useSendMessage = ({
572619
setMessages((prev) =>
573620
prev.map((msg) => {
574621
if (msg.id === aiMessageId && msg.blocks) {
575-
const blocks = msg.blocks.map((block) => {
576-
if (
577-
block.type === 'agent' &&
578-
block.agentId === event.agentId
579-
) {
580-
return { ...block, status: 'complete' as const }
581-
}
582-
return block
583-
})
622+
// Use recursive update to handle nested agents
623+
const blocks = updateBlocksRecursively(
624+
msg.blocks,
625+
event.agentId,
626+
(block) => ({ ...block, status: 'complete' as const }),
627+
)
584628
return { ...msg, blocks }
585629
}
586630
return msg
@@ -597,6 +641,12 @@ export const useSendMessage = ({
597641

598642
if (event.type === 'tool_call' && event.toolCallId) {
599643
const { toolCallId, toolName, input, agentId } = event
644+
logger.info('tool_call event received', {
645+
toolCallId,
646+
toolName,
647+
agentId: agentId || 'ROOT',
648+
hasAgentId: !!agentId,
649+
})
600650

601651
if (toolName === 'spawn_agents' && input?.agents) {
602652
const agents = Array.isArray(input.agents) ? input.agents : []
@@ -678,12 +728,11 @@ export const useSendMessage = ({
678728
return msg
679729
}
680730

681-
const updatedBlocks: ContentBlock[] = msg.blocks.map(
731+
// Use recursive update to handle nested agents
732+
const updatedBlocks = updateBlocksRecursively(
733+
msg.blocks,
734+
agentId,
682735
(block) => {
683-
if (block.type !== 'agent' || block.agentId !== agentId) {
684-
return block
685-
}
686-
687736
const agentBlocks: ContentBlock[] = block.blocks
688737
? [...block.blocks]
689738
: []

‎common/src/types/print-mode.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ export const printModeToolCallSchema = z.object({
2929
toolCallId: z.string(),
3030
toolName: z.string(),
3131
input: z.record(z.string(), z.any()),
32+
agentId: z.string().optional(),
3233
})
3334
export type PrintModeToolCall = z.infer<typeof printModeToolCallSchema>
3435

@@ -59,6 +60,7 @@ export const printModeSubagentStartSchema = z.object({
5960
agentType: z.string(),
6061
displayName: z.string(),
6162
onlyChild: z.boolean(),
63+
parentAgentId: z.string().optional(),
6264
})
6365
export type PrintModeSubagentStart = z.infer<
6466
typeof printModeSubagentStartSchema
@@ -70,6 +72,7 @@ export const printModeSubagentFinishSchema = z.object({
7072
agentType: z.string(),
7173
displayName: z.string(),
7274
onlyChild: z.boolean(),
75+
parentAgentId: z.string().optional(),
7376
})
7477
export type PrintModeSubagentFinish = z.infer<
7578
typeof printModeSubagentFinishSchema

‎scripts/dev.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,4 @@ echo "Server is ready! Starting client..."
2727
rm -f "$READY_FILE"
2828

2929
# Start the client
30-
bun start-bin -- --cwd .. "$@"
30+
bun --cwd cli dev "$@"

0 commit comments

Comments
 (0)