cli: collapse code reviewers within code-reviewer-multi-prompt

Author: jahooma

cli/src/utils/constants.ts

@@ -32,6 +32,40 @@ export const shouldCollapseByDefault = (agentType: string): boolean => {
   )
 }
 
+/**
+ * Rules for collapsing child agents when spawned by specific parent agents.
+ * Key: parent agent type pattern, Value: array of child agent type patterns to collapse
+ */
+export const PARENT_CHILD_COLLAPSE_RULES: Record<string, string[]> = {
+  'code-reviewer-multi-prompt': ['code-reviewer'],
+}
+
+/**
+ * Check if a child agent should be collapsed when spawned by a specific parent
+ */
+export const shouldCollapseForParent = (
+  childAgentType: string,
+  parentAgentType: string | undefined,
+): boolean => {
+  if (!parentAgentType) {
+    return false
+  }
+
+  for (const [parentPattern, childPatterns] of Object.entries(
+    PARENT_CHILD_COLLAPSE_RULES,
+  )) {
+    if (parentAgentType.includes(parentPattern)) {
+      for (const childPattern of childPatterns) {
+        if (childAgentType.includes(childPattern)) {
+          return true
+        }
+      }
+    }
+  }
+
+  return false
+}
+
 // Agent IDs that should render as simple text instead of full agent boxes
 export const SIMPLE_TEXT_AGENT_IDS = [
   'best-of-n-selector',

cli/src/utils/message-block-helpers.ts

@@ -1,7 +1,7 @@
 import { isEqual } from 'lodash'
 
 import { formatToolOutput } from './codebuff-client'
-import { shouldCollapseByDefault } from './constants'
+import { shouldCollapseByDefault, shouldCollapseForParent } from './constants'
 
 import type {
   ContentBlock,
@@ -250,6 +250,30 @@ export const appendInterruptionNotice = (
   return [...blocks, interruptionNotice]
 }
 
+/**
+ * Recursively finds an agent block by ID and returns its agent type.
+ * Returns undefined if not found.
+ */
+export const findAgentTypeById = (
+  blocks: ContentBlock[],
+  agentId: string,
+): string | undefined => {
+  for (const block of blocks) {
+    if (block.type === 'agent') {
+      if (block.agentId === agentId) {
+        return block.agentType
+      }
+      if (block.blocks) {
+        const found = findAgentTypeById(block.blocks, agentId)
+        if (found) {
+          return found
+        }
+      }
+    }
+  }
+  return undefined
+}
+
 /**
  * Options for creating an agent content block.
  */
@@ -262,6 +286,8 @@ export interface CreateAgentBlockOptions {
   spawnToolCallId?: string
   /** The index within the spawn_agents call */
   spawnIndex?: number
+  /** The agent type of the parent agent that spawned this one */
+  parentAgentType?: string
 }
 
 /**
@@ -270,7 +296,10 @@ export interface CreateAgentBlockOptions {
 export const createAgentBlock = (
   options: CreateAgentBlockOptions,
 ): AgentContentBlock => {
-  const { agentId, agentType, prompt, params, spawnToolCallId, spawnIndex } = options
+  const { agentId, agentType, prompt, params, spawnToolCallId, spawnIndex, parentAgentType } = options
+  const shouldCollapse =
+    shouldCollapseByDefault(agentType || '') ||
+    shouldCollapseForParent(agentType || '', parentAgentType)
   return {
     type: 'agent',
     agentId,
@@ -283,7 +312,7 @@ export const createAgentBlock = (
     ...(params && { params }),
     ...(spawnToolCallId && { spawnToolCallId }),
     ...(spawnIndex !== undefined && { spawnIndex }),
-    ...(shouldCollapseByDefault(agentType || '') && { isCollapsed: true }),
+    ...(shouldCollapse && { isCollapsed: true }),
   }
 }
 

cli/src/utils/sdk-event-handlers.ts

@@ -10,6 +10,7 @@ import {
   createAgentBlock,
   extractPlanFromBuffer,
   extractSpawnAgentResultContent,
+  findAgentTypeById,
   insertPlanBlock,
   nestBlockUnderParent,
   transformAskUserBlocks,
@@ -216,14 +217,20 @@ const handleSubagentStart = (
     'Creating new agent block (no spawn_agents match)',
   )
 
-  const newAgentBlock = createAgentBlock({
-    agentId: event.agentId,
-    agentType: event.agentType || '',
-    prompt: event.prompt,
-    params: event.params,
-  })
-
   state.message.updater.updateAiMessageBlocks((blocks) => {
+    // Look up the parent agent's type if there's a parent agent ID
+    const parentAgentType = event.parentAgentId
+      ? findAgentTypeById(blocks, event.parentAgentId)
+      : undefined
+
+    const newAgentBlock = createAgentBlock({
+      agentId: event.agentId,
+      agentType: event.agentType || '',
+      prompt: event.prompt,
+      params: event.params,
+      parentAgentType,
+    })
+
     if (event.parentAgentId) {
       const { blocks: nestedBlocks, parentFound } = nestBlockUnderParent(
         blocks,
@@ -273,6 +280,11 @@ const handleSpawnAgentsToolCall = (
   })
 
   state.message.updater.updateAiMessageBlocks((blocks) => {
+    // Look up the parent agent's type if there's a parent agent ID
+    const parentAgentType = event.agentId
+      ? findAgentTypeById(blocks, event.agentId)
+      : undefined
+
     const newAgentBlocks: ContentBlock[] = agents
       .map((agent: any, originalIndex: number) => ({ agent, originalIndex }))
       .filter(({ agent }) => !shouldHideAgent(agent.agent_type || ''))
@@ -283,6 +295,7 @@ const handleSpawnAgentsToolCall = (
           prompt: agent.prompt,
           spawnToolCallId: event.toolCallId,
           spawnIndex: originalIndex,
+          parentAgentType,
         }),
       )