fix: refactor AIChainOfThought for improved collapse functionality

- Introduced a new `AIInnerReasoningPart` component to encapsulate the rendering logic for individual reasoning parts, enhancing modularity.
- Updated `AIChainOfThought` to utilize the new component and manage collapse state more effectively using a ref.
- Simplified the extraction of headings from reasoning text, improving readability and maintainability of the code.

These changes aim to enhance the user experience by providing clearer feedback during AI reasoning processes and improving the overall structure of the components.

Signed-off-by: Innei <tukon479@gmail.com>

apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx

@@ -1,3 +1,4 @@
+import type { CollapseCssRef } from "@follow/components/ui/collapse/CollapseCss.js"
 import { CollapseCss, CollapseCssGroup } from "@follow/components/ui/collapse/CollapseCss.js"
 import { cn } from "@follow/utils"
 import type { ReasoningUIPart } from "ai"
@@ -10,62 +11,41 @@ interface AIChainOfThoughtProps {
   isStreaming?: boolean
   className?: string
 }
-
 export const AIChainOfThought: React.FC<AIChainOfThoughtProps> = React.memo(
-  ({ groups, isStreaming = false, className }) => {
+  ({ groups, isStreaming, className }) => {
     const collapseId = React.useMemo(() => `chain-${Math.random().toString(36).slice(2)}`, [])
-    // Re-mount CollapseCssGroup when streaming state changes or when we need to force-open while streaming
-    const [remountTick, setRemountTick] = React.useState(0)
-    const groupKey = `${isStreaming ? "streaming" : "idle"}:${remountTick}`
 
+    const collapseRef = React.useRef<CollapseCssRef>(null)
     const lastPartText = groups.at?.(-1)?.text
+    const currentChainReasoningIsFinished = React.useMemo(() => {
+      return groups.every((part) => part.state === "done")
+    }, [groups])
     const currentReasoningTitle = React.useMemo(() => {
       if (!isStreaming) return null
       return extractHeading(lastPartText)
     }, [isStreaming, lastPartText])
 
+    React.useEffect(() => {
+      collapseRef.current?.setIsOpened(!currentChainReasoningIsFinished)
+    }, [collapseRef, currentChainReasoningIsFinished])
+
     if (!groups || groups.length === 0) return null
 
-    const extractHeading = (text?: string): string | undefined => {
-      if (!text) return
-      const lines = text.split(/\r?\n/)
-      for (const raw of lines) {
-        const line = raw.trim()
-        if (!line) continue
-        if (line.startsWith("#")) {
-          let idx = 0
-          while (idx < line.length && line.charAt(idx) === "#") idx++
-          let content = line.slice(idx).trim()
-          while (content.endsWith("#")) content = content.slice(0, -1).trim()
-          return content || undefined
-        }
-        if (line.startsWith("**") && line.endsWith("**") && line.length > 4) {
-          return line.slice(2, -2).trim() || undefined
-        }
-        break
-      }
-      return
-    }
     return (
       <div className={cn("border-border min-w-0 max-w-full text-left", className)}>
         <div className="w-[calc(var(--ai-chat-layout-width,65ch))] max-w-full" />
 
-        <CollapseCssGroup key={groupKey}>
+        <CollapseCssGroup>
           <CollapseCss
+            ref={collapseRef}
             hideArrow
             collapseId={collapseId}
-            defaultOpen={isStreaming}
-            onOpenChange={(opened) => {
-              // While streaming, keep it open and block manual collapse
-              if (isStreaming && !opened) {
-                setRemountTick((x) => x + 1)
-              }
-            }}
+            defaultOpen={!currentChainReasoningIsFinished}
             title={
               <div className="group flex h-6 min-w-0 flex-1 items-center py-0">
                 <div className="flex items-center gap-2 text-xs">
                   <span className="text-text-secondary">
-                    {isStreaming ? (
+                    {!currentChainReasoningIsFinished ? (
                       <span>
                         Thinking: <span className="font-medium">{currentReasoningTitle}</span>
                       </span>
@@ -98,42 +78,17 @@ export const AIChainOfThought: React.FC<AIChainOfThoughtProps> = React.memo(
                       aria-hidden
                       className={cn(
                         "absolute left-2 top-2 size-2 -translate-x-1/2 rounded-full border",
-                        groupStreaming ? "border-blue bg-blue" : "border-fill bg-fill-vibrant",
+                        groupStreaming ? "border-accent bg-accent" : "border-fill bg-fill-vibrant",
                       )}
                     >
                       <i className="i-mgc-brain-cute-re absolute top-1/2 -translate-x-1/4 -translate-y-1/2" />
                     </div>
 
-                    <CollapseCss
-                      hideArrow
-                      collapseId={innerCollapseId}
-                      defaultOpen={groupStreaming}
-                      onOpenChange={(opened) => {
-                        if (groupStreaming && !opened) {
-                          setRemountTick((x) => x + 1)
-                        }
-                      }}
-                      title={
-                        <div className="group/inner flex h-6 min-w-0 flex-1 items-center py-0">
-                          <div className="text-text-secondary flex items-center gap-2 text-xs">
-                            {title ? (
-                              <span className="truncate">
-                                {"Reason: "}
-                                <span className="text-text font-medium">{title}</span>
-                              </span>
-                            ) : (
-                              <span>{groupStreaming ? "Reasoning..." : "Reasoning"}</span>
-                            )}
-                          </div>
-                          <div className="ml-2 flex items-center justify-center opacity-0 transition-opacity duration-200 group-hover/inner:opacity-100">
-                            <i className="i-mgc-right-cute-re size-3 shrink-0 transition-transform duration-200 group-data-[state=open]/inner:rotate-90" />
-                          </div>
-                        </div>
-                      }
-                      className="group/inner w-full border-none"
-                    >
-                      <AIReasoningPart text={mergedText} isStreaming={groupStreaming} />
-                    </CollapseCss>
+                    <AIInnerReasoningPart
+                      title={title}
+                      text={mergedText}
+                      groupStreaming={groupStreaming}
+                    />
                   </div>
                 )
               })}
@@ -145,4 +100,67 @@ export const AIChainOfThought: React.FC<AIChainOfThoughtProps> = React.memo(
   },
 )
 
+const AIInnerReasoningPart: React.FC<{
+  title: string | undefined
+  text: string
+  groupStreaming: boolean
+}> = React.memo(({ title, text, groupStreaming }) => {
+  const id = React.useId()
+  const collapseRef = React.useRef<CollapseCssRef>(null)
+
+  React.useEffect(() => {
+    collapseRef.current?.setIsOpened(groupStreaming)
+  }, [groupStreaming, collapseRef])
+
+  return (
+    <CollapseCss
+      ref={collapseRef}
+      hideArrow
+      collapseId={id}
+      defaultOpen
+      title={
+        <div className="group/inner flex h-6 min-w-0 flex-1 items-center py-0">
+          <div className="text-text-secondary flex items-center gap-2 text-xs">
+            {title ? (
+              <span className="truncate">
+                {"Reason: "}
+                <span className="text-text font-medium">{title}</span>
+              </span>
+            ) : (
+              <span>{groupStreaming ? "Reasoning..." : "Reasoning"}</span>
+            )}
+          </div>
+          <div className="ml-2 flex items-center justify-center opacity-0 transition-opacity duration-200 group-hover/inner:opacity-100">
+            <i className="i-mgc-right-cute-re size-3 shrink-0 transition-transform duration-200 group-data-[state=open]/inner:rotate-90" />
+          </div>
+        </div>
+      }
+      className="group/inner w-full border-none"
+    >
+      <AIReasoningPart text={text} isStreaming={groupStreaming} />
+    </CollapseCss>
+  )
+})
+
 AIChainOfThought.displayName = "AIChainOfThought"
+
+const extractHeading = (text?: string): string | undefined => {
+  if (!text) return
+  const lines = text.split(/\r?\n/)
+  for (const raw of lines) {
+    const line = raw.trim()
+    if (!line) continue
+    if (line.startsWith("#")) {
+      let idx = 0
+      while (idx < line.length && line.charAt(idx) === "#") idx++
+      let content = line.slice(idx).trim()
+      while (content.endsWith("#")) content = content.slice(0, -1).trim()
+      return content || undefined
+    }
+    if (line.startsWith("**") && line.endsWith("**") && line.length > 4) {
+      return line.slice(2, -2).trim() || undefined
+    }
+    break
+  }
+  return
+}

packages/internal/components/src/ui/collapse/CollapseCss.tsx

@@ -70,8 +70,12 @@ interface CollapseProps {
   className?: string
   children: React.ReactNode
   innerClassName?: string
+  ref?: React.Ref<CollapseCssRef>
 }
 
+export interface CollapseCssRef {
+  setIsOpened: (isOpened: boolean) => void
+}
 export const CollapseCss: FC<CollapseProps> = ({
   title,
   hideArrow,
@@ -83,6 +87,7 @@ export const CollapseCss: FC<CollapseProps> = ({
   className,
   innerClassName,
   children,
+  ref,
 }) => {
   const reactId = React.useId()
   const id = collapseId ?? reactId
@@ -100,6 +105,11 @@ export const CollapseCss: FC<CollapseProps> = ({
     onOpenChange?.(newOpened)
   }, [id, isOpened, controlledIsOpened, setOpenState, onOpenChange])
 
+  React.useImperativeHandle(ref, () => ({
+    setIsOpened: (isOpened: boolean) => {
+      setOpenState(id, isOpened)
+    },
+  }))
   return (
     <div className={cn("flex flex-col", className)} data-state={isOpened ? "open" : "hidden"}>
       <div