diff --git a/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx b/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx
index aeb232500..2b80bd327 100644
--- a/apps/desktop/layer/renderer/src/modules/ai-chat/components/displays/AIChainOfThought.tsx
+++ b/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
+}
diff --git a/packages/internal/components/src/ui/collapse/CollapseCss.tsx b/packages/internal/components/src/ui/collapse/CollapseCss.tsx
index f6765e773..d7f0807bc 100644
--- a/packages/internal/components/src/ui/collapse/CollapseCss.tsx
+++ b/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