Fix subagent streaming and parent tracking for nested agents

🤖 Generated with Codebuff Co-Authored-By: Codebuff
<noreply@codebuff.com>

Author: brandonkachen

backend/src/client-wrapper.ts

@@ -196,14 +196,23 @@ export function sendSubagentChunkWs(
     ws: WebSocket
   } & ParamsOf<SendSubagentChunkFn>,
 ): ReturnType<SendSubagentChunkFn> {
-  const { ws, userInputId, agentId, agentType, chunk, prompt } = params
+  const {
+    ws,
+    userInputId,
+    agentId,
+    agentType,
+    chunk,
+    prompt,
+    forwardToPrompt = true,
+  } = params
   return sendAction(ws, {
     type: 'subagent-response-chunk',
     userInputId,
     agentId,
     agentType,
     chunk,
     prompt,
+    forwardToPrompt,
   })
 }
 

backend/src/run-programmatic-step.ts

@@ -273,11 +273,13 @@ export async function runProgrammaticStep(
           role: 'assistant' as const,
           content: toolCallString,
         })
-        state.sendSubagentChunk({
+        // Optional call handles both top-level and nested agents
+        state.sendSubagentChunk?.({
           userInputId,
           agentId: state.agentState.agentId,
           agentType: state.agentState.agentType!,
           chunk: toolCallString,
+          forwardToPrompt: !state.agentState.parentId,
         })
       }
 
@@ -331,10 +333,7 @@ export async function runProgrammaticStep(
             }
 
             // Add parentAgentId to tool calls and results from nested agents
-            if (
-              chunk.type === 'tool_call' ||
-              chunk.type === 'tool_result'
-            ) {
+            if (chunk.type === 'tool_call' || chunk.type === 'tool_result') {
               // Only add parentAgentId if it's not already set
               if (!(chunk as any).parentAgentId) {
                 logger.debug(

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

@@ -26,6 +26,7 @@ export type SendSubagentChunk = (data: {
   agentType: string
   chunk: string
   prompt?: string
+  forwardToPrompt?: boolean
 }) => void
 
 type ToolName = 'spawn_agents'
@@ -158,42 +159,49 @@ export const handleSpawnAgents = ((
                 return
               }
 
-              // For nested agent events, add parentAgentId to enable proper nesting in UI
-              if (
-                chunk.type === 'subagent_start' ||
-                chunk.type === 'subagent_finish'
-              ) {
-                logger.debug(
-                  {
-                    eventType: chunk.type,
-                    agentId: chunk.agentId,
-                    parentId: subAgentState.agentId,
-                    parentAgentId: subAgentState.agentId,
-                  },
-                  `spawn-agents: Adding parentAgentId to ${chunk.type} event`,
-                )
-                writeToClient({
-                  ...chunk,
-                  parentAgentId: subAgentState.agentId,
-                })
+              if (chunk.type === 'text') {
+                if (chunk.text) {
+                  sendSubagentChunk({
+                    userInputId,
+                    agentId: subAgentState.agentId,
+                    agentType,
+                    chunk: chunk.text,
+                    prompt,
+                  })
+                }
                 return
               }
 
-              // For tool calls and results from nested agents, preserve the agentId but add parentAgentId
-              if (chunk.type === 'tool_call' || chunk.type === 'tool_result') {
-                logger.debug(
-                  {
-                    eventType: chunk.type,
-                    agentId: (chunk as any).agentId,
-                    parentId: subAgentState.agentId,
-                    parentAgentId: subAgentState.agentId,
-                  },
-                  `spawn-agents: Adding parentAgentId to ${chunk.type} event`,
-                )
-                writeToClient({
-                  ...chunk,
-                  parentAgentId: subAgentState.agentId,
-                })
+              // Add parentAgentId for proper nesting in UI
+              const ensureParentAgentId = () => {
+                if (
+                  chunk.type === 'subagent_start' ||
+                  chunk.type === 'subagent_finish'
+                ) {
+                  return (
+                    chunk.parentAgentId ??
+                    subAgentState.parentId ??
+                    parentAgentState?.agentId
+                  )
+                }
+                if (
+                  chunk.type === 'tool_call' ||
+                  chunk.type === 'tool_result'
+                ) {
+                  return (chunk as any).parentAgentId ?? subAgentState.agentId
+                }
+                return undefined
+              }
+
+              const parentAgentId = ensureParentAgentId()
+              if (
+                parentAgentId !== undefined &&
+                (chunk.type === 'subagent_start' ||
+                  chunk.type === 'subagent_finish' ||
+                  chunk.type === 'tool_call' ||
+                  chunk.type === 'tool_result')
+              ) {
+                writeToClient({ ...chunk, parentAgentId })
                 return
               }
 

common/src/actions.ts

@@ -143,6 +143,7 @@ export const SERVER_ACTION_SCHEMA = z.discriminatedUnion('type', [
     agentType: z.string(),
     chunk: z.string(),
     prompt: z.string().optional(),
+    forwardToPrompt: z.boolean().optional(),
   }),
   z.object({
     type: z.literal('handlesteps-log-chunk'),

common/src/types/contracts/client.ts

@@ -35,6 +35,7 @@ export type SendSubagentChunkFn = (params: {
   agentType: string
   chunk: string
   prompt?: string | undefined
+  forwardToPrompt?: boolean
 }) => void
 
 export type HandleStepsLogChunkFn = (params: {

npm-app/src/client.ts

@@ -209,6 +209,10 @@ export class Client {
   private responseComplete: boolean = false
   private userInputId: string | undefined
   private currentOnChunk: ((chunk: string | PrintModeEvent) => void) | undefined
+  private onlyChildAgents: Set<string> = new Set()
+  private emitPromptChunkToParser:
+    | ((textChunk: string) => boolean)
+    | undefined
 
   public usageData: UsageData = {
     usage: 0,
@@ -952,11 +956,25 @@ export class Client {
     })
     // Handle subagent streaming messages
     this.webSocket.subscribe('subagent-response-chunk', (action) => {
-      const { agentId, agentType, chunk, prompt } = action
+      const {
+        agentId,
+        agentType,
+        chunk,
+        prompt,
+        forwardToPrompt,
+      } = action
 
       // Store the chunk locally
       storeSubagentChunk({ agentId, agentType, chunk, prompt })
 
+      if (
+        forwardToPrompt !== false &&
+        this.onlyChildAgents.has(agentId) &&
+        this.emitPromptChunkToParser
+      ) {
+        this.emitPromptChunkToParser(chunk)
+      }
+
       // Refresh display if we're currently viewing this agent
       refreshSubagentDisplay(agentId)
     })
@@ -1340,6 +1358,44 @@ export class Client {
 
     this.userInputId = userInputId
     this.currentOnChunk = onChunk
+    this.onlyChildAgents.clear()
+
+    const emitChunkToParser = (textChunk: string) => {
+      rawChunkBuffer.push(textChunk)
+
+      const trimmed = textChunk.trim()
+      for (const tag of ONE_TIME_TAGS) {
+        if (trimmed.startsWith(`<${tag}>`) && trimmed.endsWith(closeXml(tag))) {
+          if (this.oneTimeFlags[tag]) {
+            return true
+          }
+          Spinner.get().stop()
+          const warningMessage = trimmed
+            .replace(`<${tag}>`, '')
+            .replace(closeXml(tag), '')
+          process.stdout.write(yellow(`\n\n${warningMessage}\n\n`))
+          this.oneTimeFlags[tag as (typeof ONE_TIME_LABELS)[number]] = true
+          return true
+        }
+      }
+
+      try {
+        xmlStreamParser.write(textChunk, 'utf8')
+      } catch (e) {
+        logger.error(
+          {
+            errorMessage: e instanceof Error ? e.message : String(e),
+            errorStack: e instanceof Error ? e.stack : undefined,
+            chunk: textChunk,
+          },
+          'Error writing chunk to XML stream parser',
+        )
+      }
+
+      return false
+    }
+
+    this.emitPromptChunkToParser = emitChunkToParser
 
     const stopResponse = () => {
       responseStopped = true
@@ -1349,6 +1405,8 @@ export class Client {
       this.currentOnChunk = undefined
 
       xmlStreamParser.destroy()
+      this.emitPromptChunkToParser = undefined
+      this.onlyChildAgents.clear()
 
       const additionalMessages = prompt
         ? [
@@ -1394,45 +1452,50 @@ export class Client {
 
     unsubscribeChunks = this.webSocket.subscribe('response-chunk', (a) => {
       if (a.userInputId !== userInputId) return
-      if (typeof a.chunk === 'string') {
-        const { chunk } = a
-
-        rawChunkBuffer.push(chunk)
+      const incomingChunk = a.chunk
 
-        const trimmed = chunk.trim()
-        for (const tag of ONE_TIME_TAGS) {
+      const updateOnlyChildAgents = (event: PrintModeEvent) => {
+        if ('agentId' in event && event.agentId) {
           if (
-            trimmed.startsWith(`<${tag}>`) &&
-            trimmed.endsWith(closeXml(tag))
+            event.type === 'subagent_start' &&
+            'onlyChild' in event &&
+            event.onlyChild
           ) {
-            if (this.oneTimeFlags[tag]) {
-              return
-            }
-            Spinner.get().stop()
-            const warningMessage = trimmed
-              .replace(`<${tag}>`, '')
-              .replace(closeXml(tag), '')
-            process.stdout.write(yellow(`\n\n${warningMessage}\n\n`))
-            this.oneTimeFlags[tag as (typeof ONE_TIME_LABELS)[number]] = true
-            return
+            this.onlyChildAgents.add(event.agentId)
+          } else if (event.type === 'subagent_finish') {
+            this.onlyChildAgents.delete(event.agentId)
           }
         }
+      }
 
-        try {
-          xmlStreamParser.write(chunk, 'utf8')
-        } catch (e) {
-          logger.error(
-            {
-              errorMessage: e instanceof Error ? e.message : String(e),
-              errorStack: e instanceof Error ? e.stack : undefined,
-              chunk,
-            },
-            'Error writing chunk to XML stream parser',
-          )
+      if (typeof incomingChunk === 'string') {
+        emitChunkToParser(incomingChunk)
+        return
+      }
+
+      if (incomingChunk.type === 'text') {
+        printModeLog(incomingChunk)
+        // Skip nested subagent text from streaming output
+        if (incomingChunk.agentId) return
+        emitChunkToParser(incomingChunk.text)
+        return
+      }
+
+      if (incomingChunk.type === 'error') {
+        const errorText = `${yellow(incomingChunk.message)}\n`
+        const handler = this.currentOnChunk
+        if (handler) {
+          handler(errorText)
+        } else {
+          Spinner.get().stop()
+          DiffManager.receivedResponse()
+          process.stdout.write(errorText)
         }
-      } else {
-        onChunk(a.chunk)
       }
+
+      updateOnlyChildAgents(incomingChunk)
+
+      onChunk(incomingChunk)
     })
 
     let stepsCount = 0
@@ -1445,6 +1508,8 @@ export class Client {
 
         if (action.promptId !== userInputId) return
         this.responseComplete = true
+        this.emitPromptChunkToParser = undefined
+        this.onlyChildAgents.clear()
 
         Spinner.get().stop()
 
@@ -1552,6 +1617,7 @@ Go to https://www.codebuff.com/config for more information.`) +
 
         unsubscribeChunks()
         unsubscribeComplete()
+        this.onlyChildAgents.clear()
 
         // Clear the onChunk callback when response is complete
         this.currentOnChunk = undefined