refactor(cli): consolidate block utils and add tree traversal primitives (Commit 2.2)

Creates centralized block-tree-utils.ts with reusable tree traversal primitives:
- traverseBlocks: visit all blocks with early exit support
- findBlockByPredicate: find first matching block
- mapBlocks: transform blocks with reference preservation
- updateBlockById: update by any ID type (agentId, toolCallId, thinkingId, id)
- updateAgentBlockById: update specifically agent blocks
- toggleBlockCollapse: toggle collapsed state with userOpened tracking

Comprehensive tests (33 tests) including:
- Deep nesting (3+ levels) coverage
- Proper type narrowing (isAgentBlock guard instead of "as any")
- Multiple toggle cycle verification
- Parent + children transformation tests

Author: brandonkachen

cli/src/hooks/use-chat-messages.ts

@@ -5,6 +5,7 @@
 
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 
+import { toggleBlockCollapse, updateBlockById } from '../utils/block-tree-utils'
 import { buildMessageTree } from '../utils/message-tree-utils'
 
 import type { ChatMessage, ContentBlock } from '../types/chat'
@@ -105,78 +106,19 @@ export function useChatMessages({
           // Handle blocks within messages
           if (!message.blocks) return message
 
-          const updateBlocksRecursively = (
-            blocks: ContentBlock[],
-          ): ContentBlock[] => {
-            let foundTarget = false
-            const result = blocks.map((block) => {
-              // Handle thinking blocks - just match by thinkingId
-              if (block.type === 'text' && block.thinkingId === id) {
-                foundTarget = true
-                const wasCollapsed = block.isCollapsed ?? false
-                return {
-                  ...block,
-                  isCollapsed: !wasCollapsed,
-                  userOpened: wasCollapsed, // Mark as user-opened if expanding
-                }
-              }
-
-              // Handle agent blocks
-              if (block.type === 'agent' && block.agentId === id) {
-                foundTarget = true
-                const wasCollapsed = block.isCollapsed ?? false
-                return {
-                  ...block,
-                  isCollapsed: !wasCollapsed,
-                  userOpened: wasCollapsed, // Mark as user-opened if expanding
-                }
-              }
-
-              // Handle tool blocks
-              if (block.type === 'tool' && block.toolCallId === id) {
-                foundTarget = true
-                const wasCollapsed = block.isCollapsed ?? false
-                return {
-                  ...block,
-                  isCollapsed: !wasCollapsed,
-                  userOpened: wasCollapsed, // Mark as user-opened if expanding
-                }
-              }
-
-              // Handle agent-list blocks
-              if (block.type === 'agent-list' && block.id === id) {
-                foundTarget = true
-                const wasCollapsed = block.isCollapsed ?? false
-                return {
-                  ...block,
-                  isCollapsed: !wasCollapsed,
-                  userOpened: wasCollapsed, // Mark as user-opened if expanding
-                }
-              }
-
-              // Recursively update nested blocks inside agent blocks
-              if (block.type === 'agent' && block.blocks) {
-                const updatedBlocks = updateBlocksRecursively(block.blocks)
-                // Only create new block if nested blocks actually changed
-                if (updatedBlocks !== block.blocks) {
-                  foundTarget = true
-                  return {
-                    ...block,
-                    blocks: updatedBlocks,
-                  }
-                }
-              }
-
-              return block
-            })
-
-            // Return original array reference if nothing changed
-            return foundTarget ? result : blocks
-          }
+          // Use unified block tree utility + shared collapse helper
+          const updatedBlocks = updateBlockById(
+            message.blocks,
+            id,
+            toggleBlockCollapse,
+          )
+
+          // Only create new message if blocks actually changed
+          if (updatedBlocks === message.blocks) return message
 
           return {
             ...message,
-            blocks: updateBlocksRecursively(message.blocks),
+            blocks: updatedBlocks,
           }
         })
       })

cli/src/types/chat.ts

@@ -141,6 +141,17 @@ export type ContentBlock =
   | ToolContentBlock
   | PlanContentBlock
 
+/**
+ * Block types that support collapse state (isCollapsed/userOpened).
+ * Used for type-safe collapse toggling operations.
+ */
+export type CollapsibleBlock =
+  | AgentContentBlock
+  | AgentListContentBlock
+  | ImageContentBlock
+  | TextContentBlock
+  | ToolContentBlock
+
 export type AgentMessage = {
   agentName: string
   agentType: string
@@ -225,3 +236,16 @@ export function isAskUserBlock(
 export function isImageBlock(block: ContentBlock): block is ImageContentBlock {
   return block.type === 'image'
 }
+
+/**
+ * Type guard for blocks that support collapse state (isCollapsed/userOpened).
+ */
+export function isCollapsibleBlock(block: ContentBlock): block is CollapsibleBlock {
+  return (
+    block.type === 'agent' ||
+    block.type === 'agent-list' ||
+    block.type === 'image' ||
+    block.type === 'text' ||
+    block.type === 'tool'
+  )
+}