refactor: simplify collapse handling with callback wrapper pattern

Replace prop drilling of isUserCollapsingRef with a callback wrapper pattern. Created setCollapsedAgentsWithFlag wrapper that automatically sets the flag before updating state, eliminating the need to pass the ref through multiple component layers.

Changes:
- Add setCollapsedAgentsWithFlag and handleCollapseToggle in chat.tsx
- Remove isUserCollapsingRef prop from MessageRenderer, MessageBlock, and AgentBranchItem
- Update use-message-renderer to use timerStartTime consistently
- Fix status indicator tests to use proper ElapsedTimeTracker mock

Author: brandonkachen

cli/src/chat.tsx

@@ -228,6 +228,42 @@ export const Chat = ({
     isUserCollapsingRef.current = false
   }, [collapsedAgents])
 
+  // Wrapper for setCollapsedAgents that sets the flag to prevent auto-scroll
+  const setCollapsedAgentsWithFlag = useCallback(
+    (action: React.SetStateAction<Set<string>>) => {
+      isUserCollapsingRef.current = true
+      setCollapsedAgents(action)
+    },
+    [setCollapsedAgents],
+  )
+
+  const handleCollapseToggle = useCallback(
+    (id: string) => {
+      const wasCollapsed = collapsedAgents.has(id)
+
+      setCollapsedAgentsWithFlag((prev) => {
+        const next = new Set(prev)
+        if (next.has(id)) {
+          next.delete(id)
+        } else {
+          next.add(id)
+        }
+        return next
+      })
+
+      setUserOpenedAgents((prev) => {
+        const next = new Set(prev)
+        if (wasCollapsed) {
+          next.add(id)
+        } else {
+          next.delete(id)
+        }
+        return next
+      })
+    },
+    [collapsedAgents, setCollapsedAgentsWithFlag, setUserOpenedAgents],
+  )
+
   const { scrollToLatest, scrollboxProps, isAtBottom } = useChatScrollbox(
     scrollRef,
     messages,
@@ -592,11 +628,11 @@ export const Chat = ({
             streamingAgents={streamingAgents}
             isWaitingForResponse={isWaitingForResponse}
             timerStartTime={timerStartTime}
-            setCollapsedAgents={setCollapsedAgents}
+            onCollapseToggle={handleCollapseToggle}
+            setCollapsedAgents={setCollapsedAgentsWithFlag}
             setFocusedAgentId={setFocusedAgentId}
             userOpenedAgents={userOpenedAgents}
             setUserOpenedAgents={setUserOpenedAgents}
-            isUserCollapsingRef={isUserCollapsingRef}
             onBuildFast={handleBuildFast}
             onBuildMax={handleBuildMax}
           />

cli/src/components/__tests__/status-indicator.timer.test.tsx

@@ -15,12 +15,17 @@ import '../../state/theme-store' // Initialize theme store
 import { renderToStaticMarkup } from 'react-dom/server'
 
 import * as codebuffClient from '../../utils/codebuff-client'
+import type { ElapsedTimeTracker } from '../../hooks/use-elapsed-time'
 
-const createTimerStartTime = (
+const createMockTimer = (
   elapsedSeconds: number,
   started: boolean,
-): number | null =>
-  started ? Date.now() - elapsedSeconds * 1000 : null
+): ElapsedTimeTracker => ({
+  startTime: started ? Date.now() - elapsedSeconds * 1000 : null,
+  elapsedSeconds,
+  start: () => {},
+  stop: () => {},
+})
 
 describe('StatusIndicator timer rendering', () => {
   let getClientSpy: ReturnType<typeof spyOn>
@@ -40,7 +45,7 @@ describe('StatusIndicator timer rendering', () => {
       <StatusIndicator
         clipboardMessage={null}
         isActive={true}
-        timerStartTime={createTimerStartTime(5, true)}
+        timer={createMockTimer(5, true)}
         nextCtrlCWillExit={false}
       />,
     )
@@ -51,7 +56,7 @@ describe('StatusIndicator timer rendering', () => {
       <StatusIndicator
         clipboardMessage={null}
         isActive={false}
-        timerStartTime={createTimerStartTime(0, false)}
+        timer={createMockTimer(0, false)}
         nextCtrlCWillExit={false}
       />,
     )
@@ -64,7 +69,7 @@ describe('StatusIndicator timer rendering', () => {
       <StatusIndicator
         clipboardMessage="Copied!"
         isActive={true}
-        timerStartTime={createTimerStartTime(12, true)}
+        timer={createMockTimer(12, true)}
         nextCtrlCWillExit={false}
       />,
     )

cli/src/components/agent-branch-item.tsx

@@ -18,7 +18,6 @@ interface AgentBranchItemProps {
   statusIndicator?: string
   onToggle?: () => void
   titleSuffix?: string
-  isUserCollapsingRef?: React.MutableRefObject<boolean>
 }
 
 export const AgentBranchItem = ({
@@ -35,7 +34,6 @@ export const AgentBranchItem = ({
   statusIndicator = '●',
   onToggle,
   titleSuffix,
-  isUserCollapsingRef,
 }: AgentBranchItemProps) => {
   const theme = useTheme()
 
@@ -289,13 +287,7 @@ export const AgentBranchItem = ({
                   alignSelf: 'flex-end',
                   marginTop: 1,
                 }}
-                onMouseDown={() => {
-                  // Set flag to prevent auto-scroll during user-initiated collapse
-                  if (isUserCollapsingRef) {
-                    isUserCollapsingRef.current = true
-                  }
-                  onToggle()
-                }}
+                onMouseDown={onToggle}
               >
                 <text
                   fg={theme.secondary}

cli/src/components/message-renderer.tsx

@@ -25,11 +25,11 @@ interface MessageRendererProps {
   streamingAgents: Set<string>
   isWaitingForResponse: boolean
   timerStartTime: number | null
+  onCollapseToggle: (id: string) => void
   setCollapsedAgents: React.Dispatch<React.SetStateAction<Set<string>>>
   setFocusedAgentId: React.Dispatch<React.SetStateAction<string | null>>
   userOpenedAgents: Set<string>
   setUserOpenedAgents: React.Dispatch<React.SetStateAction<Set<string>>>
-  isUserCollapsingRef: React.MutableRefObject<boolean>
   onBuildFast: () => void
   onBuildMax: () => void
 }
@@ -46,43 +46,14 @@ export const MessageRenderer = (props: MessageRendererProps): ReactNode => {
     streamingAgents,
     isWaitingForResponse,
     timerStartTime,
+    onCollapseToggle,
     setCollapsedAgents,
     setFocusedAgentId,
     setUserOpenedAgents,
-    isUserCollapsingRef,
     onBuildFast,
     onBuildMax,
   } = props
 
-  const onToggleCollapsed = useCallback(
-    (id: string) => {
-      const wasCollapsed = collapsedAgents.has(id)
-
-      // Set flag to prevent auto-scroll during user-initiated collapse
-      isUserCollapsingRef.current = true
-
-      setCollapsedAgents((prev) => {
-        const next = new Set(prev)
-        if (next.has(id)) {
-          next.delete(id)
-        } else {
-          next.add(id)
-        }
-        return next
-      })
-      setUserOpenedAgents((prev) => {
-        const next = new Set(prev)
-        if (wasCollapsed) {
-          next.add(id)
-        } else {
-          next.delete(id)
-        }
-        return next
-      })
-    },
-    [collapsedAgents, setCollapsedAgents, setUserOpenedAgents, isUserCollapsingRef],
-  )
-
   return (
     <>
       {topLevelMessages.map((message, idx) => {
@@ -105,7 +76,7 @@ export const MessageRenderer = (props: MessageRendererProps): ReactNode => {
             setFocusedAgentId={setFocusedAgentId}
             isWaitingForResponse={isWaitingForResponse}
             timerStartTime={timerStartTime}
-            onToggleCollapsed={onToggleCollapsed}
+            onToggleCollapsed={onCollapseToggle}
             onBuildFast={onBuildFast}
             onBuildMax={onBuildMax}
           />

cli/src/hooks/use-message-renderer.tsx

@@ -30,7 +30,6 @@ interface UseMessageRendererProps {
   setFocusedAgentId: React.Dispatch<React.SetStateAction<string | null>>
   userOpenedAgents: Set<string>
   setUserOpenedAgents: React.Dispatch<React.SetStateAction<Set<string>>>
-  isUserCollapsingRef: React.MutableRefObject<boolean>
 }
 
 export const useMessageRenderer = (
@@ -51,7 +50,6 @@ export const useMessageRenderer = (
     setFocusedAgentId,
     userOpenedAgents,
     setUserOpenedAgents,
-    isUserCollapsingRef,
   } = props
 
   return useMemo(() => {
@@ -103,9 +101,6 @@ export const useMessageRenderer = (
 
         const wasCollapsed = collapsedAgents.has(message.id)
 
-        // Set flag to prevent auto-scroll during user-initiated collapse
-        isUserCollapsingRef.current = true
-
         setCollapsedAgents((prev) => {
           const next = new Set(prev)
 
@@ -381,7 +376,7 @@ export const useMessageRenderer = (
                     isComplete={message.isComplete}
                     completionTime={message.completionTime}
                     credits={message.credits}
-                    timer={timer}
+                    timerStartTime={timer.startTime}
                     textColor={textColor}
                     timestampColor={timestampColor}
                     markdownOptions={markdownOptions}
@@ -391,10 +386,7 @@ export const useMessageRenderer = (
                     streamingAgents={streamingAgents}
                     onToggleCollapsed={(id: string) => {
                       const wasCollapsed = collapsedAgents.has(id)
-                      
-                      // Set flag to prevent auto-scroll during user-initiated collapse
-                      isUserCollapsingRef.current = true
-                      
+
                       setCollapsedAgents((prev) => {
                         const next = new Set(prev)
                         if (next.has(id)) {
@@ -447,7 +439,7 @@ export const useMessageRenderer = (
                   isComplete={message.isComplete}
                   completionTime={message.completionTime}
                   credits={message.credits}
-                  timer={timer}
+                  timerStartTime={timer.startTime}
                   textColor={textColor}
                   timestampColor={timestampColor}
                   markdownOptions={markdownOptions}
@@ -457,10 +449,7 @@ export const useMessageRenderer = (
                   streamingAgents={streamingAgents}
                   onToggleCollapsed={(id: string) => {
                     const wasCollapsed = collapsedAgents.has(id)
-                    
-                    // Set flag to prevent auto-scroll during user-initiated collapse
-                    isUserCollapsingRef.current = true
-                    
+
                     setCollapsedAgents((prev) => {
                       const next = new Set(prev)
                       if (next.has(id)) {
@@ -484,7 +473,6 @@ export const useMessageRenderer = (
                       return next
                     })
                   }}
-                  isUserCollapsingRef={isUserCollapsingRef}
                 />
               </box>
             )}
@@ -524,6 +512,5 @@ export const useMessageRenderer = (
     setUserOpenedAgents,
     setFocusedAgentId,
     userOpenedAgents,
-    isUserCollapsingRef,
   ])
 }