Fix thinking block collapse/expansion

Author: jahooma

cli/src/chat.tsx

@@ -286,7 +286,7 @@ export const Chat = ({
           ): ContentBlock[] => {
             let foundTarget = false
             const result = blocks.map((block) => {
-              // Handle thinking blocks (grouped text blocks)
+              // Handle thinking blocks - just match by thinkingId
               if (block.type === 'text' && block.thinkingId === id) {
                 foundTarget = true
                 const wasCollapsed = block.isCollapsed ?? false
@@ -330,7 +330,7 @@ export const Chat = ({
                 }
               }
 
-              // Recursively update nested blocks
+              // 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

cli/src/components/blocks/thinking-block.tsx

@@ -10,8 +10,6 @@ const NESTED_WIDTH_OFFSET = 10
 
 interface ThinkingBlockProps {
   blocks: Extract<ContentBlock, { type: 'text' }>[]
-  keyPrefix: string
-  startIndex: number
   onToggleCollapsed: (id: string) => void
   availableWidth: number
   isNested: boolean
@@ -20,28 +18,28 @@ interface ThinkingBlockProps {
 export const ThinkingBlock = memo(
   ({
     blocks,
-    keyPrefix,
-    startIndex,
     onToggleCollapsed,
     availableWidth,
     isNested,
   }: ThinkingBlockProps) => {
-    const thinkingId = `${keyPrefix}-thinking-${startIndex}`
+    const firstBlock = blocks[0]
+    const thinkingId = firstBlock?.thinkingId
     const combinedContent = blocks
       .map((b) => b.content)
       .join('')
       .trim()
 
-    const firstBlock = blocks[0]
     const isCollapsed = firstBlock?.isCollapsed ?? true
     const offset = isNested ? NESTED_WIDTH_OFFSET : WIDTH_OFFSET
     const availWidth = Math.max(10, availableWidth - offset)
 
     const handleToggle = useCallback(() => {
-      onToggleCollapsed(thinkingId)
+      if (thinkingId) {
+        onToggleCollapsed(thinkingId)
+      }
     }, [onToggleCollapsed, thinkingId])
 
-    if (!combinedContent) {
+    if (!combinedContent || !thinkingId) {
       return null
     }
 

cli/src/components/message-block.tsx

@@ -421,10 +421,8 @@ const AgentBody = memo(
 
         nodes.push(
           <ThinkingBlock
-            key={`${keyPrefix}-thinking-${start}`}
+            key={reasoningBlocks[0]?.thinkingId ?? `${keyPrefix}-thinking-${start}`}
             blocks={reasoningBlocks}
-            keyPrefix={keyPrefix}
-            startIndex={start}
             onToggleCollapsed={onToggleCollapsed}
             availableWidth={availableWidth}
             isNested={true}
@@ -1234,10 +1232,8 @@ const BlocksRenderer = memo(
 
         nodes.push(
           <ThinkingBlock
-            key={`${messageId}-thinking-${start}`}
+            key={reasoningBlocks[0]?.thinkingId ?? `${messageId}-thinking-${start}`}
             blocks={reasoningBlocks}
-            keyPrefix={messageId}
-            startIndex={start}
             onToggleCollapsed={onToggleCollapsed}
             availableWidth={availableWidth}
             isNested={false}

cli/src/components/thinking.tsx

@@ -44,70 +44,35 @@ export const Thinking = memo(
         }}
         onClick={onToggle}
       >
+        <text style={{ fg: theme.foreground }}>
+          <span>• </span>
+          <span attributes={TextAttributes.BOLD}>Thinking</span>
+        </text>
         {isCollapsed ? (
           lines.length > 0 && (
-            <box
-              style={{
-                flexDirection: 'row',
-                gap: 0,
-                alignItems: 'stretch',
-                marginTop: 0,
-              }}
-            >
-              <box
-                style={{
-                  width: 1,
-                  backgroundColor: theme.muted,
-                  marginTop: 0,
-                  marginBottom: 0,
-                }}
-              />
-              <box
-                style={{
-                  paddingLeft: 1,
-                  flexGrow: 1,
-                }}
-              >
-                <text
-                  style={{
-                    wrapMode: 'word',
-                    fg: theme.muted,
-                  }}
-                  attributes={TextAttributes.ITALIC}
-                >
-                  {hasMore ? '...' + lines.join(' ') : lines.join(' ')}
-                </text>
-              </box>
-            </box>
-          )
-        ) : (
-          <box
-            style={{
-              flexDirection: 'row',
-              gap: 0,
-              alignItems: 'stretch',
-              marginTop: 0,
-            }}
-          >
-            <box
-              style={{
-                width: 1,
-                backgroundColor: theme.muted,
-                marginTop: 0,
-                marginBottom: 0,
-              }}
-            />
-            <box style={{ paddingLeft: 1, flexGrow: 1 }}>
+            <box style={{ paddingLeft: 2 }}>
               <text
                 style={{
                   wrapMode: 'word',
                   fg: theme.muted,
                 }}
                 attributes={TextAttributes.ITALIC}
               >
-                {content}
+                {hasMore ? '...' + lines.join(' ') : lines.join(' ')}
               </text>
             </box>
+          )
+        ) : (
+          <box style={{ paddingLeft: 2 }}>
+            <text
+              style={{
+                wrapMode: 'word',
+                fg: theme.muted,
+              }}
+              attributes={TextAttributes.ITALIC}
+            >
+              {content}
+            </text>
           </box>
         )}
       </Button>

cli/src/utils/block-operations.ts

@@ -13,6 +13,12 @@ import type {
 } from '../types/chat'
 import { logger } from './logger'
 
+let thinkingIdCounter = 0
+const generateThinkingId = (): string => {
+  thinkingIdCounter++
+  return `thinking-${thinkingIdCounter}`
+}
+
 type AgentTextUpdate =
   | { type: 'text'; mode: 'append'; content: string }
   | { type: 'text'; mode: 'replace'; content: string }
@@ -91,12 +97,14 @@ const isOpenThinkingBlock = (block: ContentBlock | undefined): boolean => {
 const createReasoningBlock = (
   content: string,
   thinkingOpen: boolean,
+  thinkingId: string,
 ): TextContentBlock => ({
   type: 'text',
   content,
   textType: 'reasoning',
   isCollapsed: true,
   thinkingOpen,
+  thinkingId,
 })
 
 /**
@@ -242,7 +250,7 @@ const appendTextWithThinkParsingToBlocks = (
       const thinkingOpen =
         isLastSegment && !textToParse.endsWith(THINK_CLOSE_TAG)
       if (thinkingOpen) {
-        nextBlocks.push(createReasoningBlock(segment.content, thinkingOpen))
+        nextBlocks.push(createReasoningBlock(segment.content, thinkingOpen, generateThinkingId()))
       }
     } else {
       const prevBlock = nextBlocks[nextBlocks.length - 1]
@@ -295,6 +303,7 @@ export const appendTextToRootStream = (
       content: delta.text,
       textType: 'reasoning',
       isCollapsed: true,
+      thinkingId: generateThinkingId(),
     }
 
     return [...nextBlocks, newBlock]