refactor(ai-chat): enhance AIChainOfThought and ToolInvocationComponent for improved reasoning handling

- Updated the `AIChainOfThought` component to support a new `ChainReasoningPart` type, allowing for better integration of tool and reasoning parts.
- Refactored the logic for determining when the current chain reasoning is finished, ensuring accurate state management.
- Enhanced the `AIMessageParts` component to utilize the new `ChainReasoningPart` type, improving the handling of message parts.
- Modified the `ToolInvocationComponent` to streamline the rendering of tool invocation details, improving UI consistency and clarity.

These changes aim to enhance the functionality and maintainability of the AI chat components, providing a more robust user experience.

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

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

@@ -2,13 +2,17 @@ import type { CollapseCssRef } from "@follow/components/ui/collapse/CollapseCss.
 import { CollapseCss, CollapseCssGroup } from "@follow/components/ui/collapse/CollapseCss.js"
 import { ShinyText } from "@follow/components/ui/shiny-text/ShinyText.js"
 import { cn } from "@follow/utils"
-import type { ReasoningUIPart } from "ai"
+import type { BizUITools } from "@folo-services/ai-tools"
+import type { ReasoningUIPart, ToolUIPart } from "ai"
+import { isToolUIPart } from "ai"
 import * as React from "react"
 
+import { ToolInvocationComponent } from "../message/ToolInvocationComponent"
 import { AIReasoningPart } from "./AIReasoningPart"
 
+export type ChainReasoningPart = ReasoningUIPart | ToolUIPart<BizUITools>
 interface AIChainOfThoughtProps {
-  groups: ReadonlyArray<ReasoningUIPart>
+  groups: ReadonlyArray<ChainReasoningPart>
   isStreaming?: boolean
   className?: string
 }
@@ -17,14 +21,34 @@ export const AIChainOfThought: React.FC<AIChainOfThoughtProps> = React.memo(
     const collapseId = React.useMemo(() => `chain-${Math.random().toString(36).slice(2)}`, [])
 
     const collapseRef = React.useRef<CollapseCssRef>(null)
-    const lastPartText = groups.at?.(-1)?.text
+
     const currentChainReasoningIsFinished = React.useMemo(() => {
-      return groups.every((part) => part.state === "done")
+      let allDone = true
+      for (const part of groups) {
+        if (isToolUIPart(part)) {
+          continue
+        }
+        if (part.state !== "done") {
+          allDone = false
+          break
+        }
+      }
+      return allDone
     }, [groups])
     const currentReasoningTitle = React.useMemo(() => {
       if (!isStreaming) return null
+
+      const lastPart = groups.at?.(-1)
+
+      if (!lastPart) return null
+
+      if (isToolUIPart(lastPart)) {
+        return `Calling ${lastPart.type.replace("tool-", "")}`
+      }
+
+      const lastPartText = lastPart.text
       return extractHeading(lastPartText)
-    }, [isStreaming, lastPartText])
+    }, [groups, isStreaming])
 
     React.useEffect(() => {
       collapseRef.current?.setIsOpened(!currentChainReasoningIsFinished)
@@ -65,20 +89,22 @@ export const AIChainOfThought: React.FC<AIChainOfThoughtProps> = React.memo(
             <div className="relative">
               <div aria-hidden className="border-fill absolute inset-y-0 left-2 border-l" />
               {groups.map((part, index) => {
+                const innerCollapseId = `${collapseId}-${index}`
+                if (isToolUIPart(part)) {
+                  return <ToolInvocationComponent key={innerCollapseId} part={part} />
+                }
                 const mergedText = part.text
 
                 const title = extractHeading(part.text)
                 const groupStreaming = part.state === "streaming"
 
-                const innerCollapseId = `${collapseId}-${index}`
-
                 return (
                   <div key={innerCollapseId} className="relative pb-3 pl-8 last:pb-0">
                     <div
                       aria-hidden
                       className={cn(
                         "absolute left-2 top-2 size-2 -translate-x-1/2 rounded-full border",
-                        groupStreaming ? "border-accent bg-accent" : "border-fill bg-fill-vibrant",
+                        "border-fill bg-fill-vibrant",
                       )}
                     >
                       <i className="i-mgc-brain-cute-re absolute top-1/2 -translate-x-1/4 -translate-y-1/2" />

apps/desktop/layer/renderer/src/modules/ai-chat/components/message/AIMessageParts.tsx

@@ -21,6 +21,7 @@ import {
   AIDisplayFeedPart,
   AIDisplaySubscriptionsPart,
 } from "../displays"
+import type { ChainReasoningPart } from "../displays/AIChainOfThought"
 import { AIMarkdownStreamingMessage } from "./AIMarkdownMessage"
 import { ToolInvocationComponent } from "./ToolInvocationComponent"
 
@@ -51,29 +52,47 @@ export const AIMessageParts: React.FC<AIMessagePartsProps> = React.memo(
     }, [chatStatus, isLastMessage, shouldStreamingAnimation])
 
     const displayParts = React.useMemo(() => {
-      const parts = [] as (ReasoningUIPart[] | TextUIPart | ToolUIPart<BizUITools>)[]
+      const parts = [] as (ChainReasoningPart[] | TextUIPart | ToolUIPart<BizUITools>)[]
 
-      const chainReasoningParts: ReasoningUIPart[] = []
+      let chainReasoningParts: ChainReasoningPart[] | null = null
       for (const part of message.parts) {
-        if (part.type !== "reasoning") {
-          parts.push(chainReasoningParts.concat())
-          chainReasoningParts.length = 0
+        const isReasoning = part.type === "reasoning" && !!(part as ReasoningUIPart).text
+        const isTool = part.type.startsWith("tool-")
+
+        if (isReasoning) {
+          if (!chainReasoningParts) {
+            chainReasoningParts = []
+            // insert by reference once; keep appending to the same array thereafter
+            parts.push(chainReasoningParts)
+          }
+          chainReasoningParts.push(part as ReasoningUIPart)
+          continue
         }
+
+        if (isTool) {
+          if (chainReasoningParts && chainReasoningParts.length > 0) {
+            chainReasoningParts.push(part as ToolUIPart<BizUITools>)
+          } else {
+            parts.push(part as ToolUIPart<BizUITools>)
+          }
+          continue
+        }
+
+        // Only add text to top-level; do not break an active chain
         if (part.type === "text") {
           parts.push(part)
-        } else if (part.type.startsWith("tool-")) {
-          parts.push(part as ToolUIPart<BizUITools>)
-        } else if (part.type === "reasoning" && part.text) {
-          chainReasoningParts.push(part as ReasoningUIPart)
+          continue
         }
+
+        // Unknown/meta parts (e.g., step-start, source) are skipped here without breaking an active chain
       }
 
-      if (chainReasoningParts.length > 0) {
-        parts.push(chainReasoningParts.concat())
-      }
+      // No final flush needed; chain array already referenced in parts
       return parts
     }, [message.parts])
 
+    // console.info("displayParts", displayParts)
+
     const lowPriorityParts = React.useDeferredValue(displayParts)
 
     return (
@@ -82,7 +101,7 @@ export const AIMessageParts: React.FC<AIMessagePartsProps> = React.memo(
           const partKey = `${message.id}-${index}`
 
           if (Array.isArray(partOrParts)) {
-            const reasoningParts = partOrParts as ReasoningUIPart[]
+            const reasoningParts = partOrParts as ChainReasoningPart[]
             return (
               <AIChainOfThought
                 key={partKey}

apps/desktop/layer/renderer/src/modules/ai-chat/components/message/ToolInvocationComponent.tsx

@@ -15,77 +15,83 @@ export const ToolInvocationComponent: React.FC<ToolInvocationComponentProps> = R
     const hasResult = "output" in part && part.output
     const hasArgs = "input" in part && part.input
 
+    const isCalling = part.state === "input-streaming"
+
     // Generate a unique value for this accordion item
     const accordionValue = `tool-${"toolCallId" in part ? part.toolCallId : Math.random()}`
 
     return (
-      <div
-        className={`min-w-0 max-w-full text-left ${hasError ? "border-red/30" : "border-border"}`}
-      >
-        <div className="w-[calc(var(--ai-chat-message-container-width,65ch))] max-w-full" />
-
-        <div className="w-full">
-          <CollapseCssGroup>
-            <CollapseCss
-              collapseId={accordionValue}
-              hideArrow
-              className="group border-none"
-              title={
-                <div className="group flex h-6 min-w-0 flex-1 items-center justify-between py-0">
-                  <div className="text-text-secondary flex items-center gap-2 text-xs">
-                    <i
-                      className={hasError ? "i-mgc-close-cute-re text-red" : "i-mgc-tool-cute-re"}
-                    />
-                    <span>{hasError ? "Tool Failed:" : "Tool Called:"}</span>
-                    <h4 className={`truncate font-medium ${hasError ? "text-red" : "text-text"}`}>
-                      {toolName}
-                    </h4>
-                  </div>
-
-                  {/* Custom arrow that only shows on hover */}
-                  <div className="ml-auto flex items-center justify-center opacity-0 transition-opacity duration-200 group-hover:opacity-100">
-                    <i className="i-mgc-right-cute-re size-3 shrink-0 transition-transform duration-200 group-data-[state=open]:rotate-90" />
-                  </div>
-                </div>
-              }
-              contentClassName="pb-0 pt-2"
-            >
-              <div className="space-y-2 text-xs">
-                {/* Show tool arguments if available */}
-                {hasArgs ? (
-                  <div>
-                    <div className="text-text-secondary mb-1 font-medium">Arguments:</div>
-                    <JsonHighlighter
-                      className="text-text-tertiary bg-fill-secondary overflow-x-auto rounded p-2 text-[11px]"
-                      json={JSON.stringify(part.input, null, 2)}
-                    />
-                  </div>
-                ) : null}
-
-                {/* Show tool result if available */}
-                {hasResult ? (
-                  <div>
-                    <div className="text-text-secondary mb-1 font-medium">Result:</div>
-                    <JsonHighlighter
-                      className="text-text-tertiary bg-fill-secondary overflow-x-auto rounded p-2 text-[11px]"
-                      json={JSON.stringify(part.output, null, 2)}
-                    />
-                  </div>
-                ) : null}
-
-                {/* Show error if available */}
-                {hasError && "errorText" in part ? (
-                  <div>
-                    <div className="text-red mb-1 font-medium">Error:</div>
-                    <pre className="text-red bg-red/10 overflow-x-auto rounded p-2 text-[11px]">
-                      {String(part.errorText)}
-                    </pre>
-                  </div>
-                ) : null}
-              </div>
-            </CollapseCss>
-          </CollapseCssGroup>
+      <div className="relative pb-3 pl-8 last:pb-0">
+        <div
+          aria-hidden
+          className={`absolute left-2 top-2 size-2 -translate-x-1/2 rounded-full border ${
+            hasError ? "border-red bg-red" : "border-fill bg-fill-vibrant"
+          }`}
+        >
+          <i
+            className={`absolute top-1/2 -translate-x-1/4 -translate-y-1/2 ${
+              hasError ? "i-mgc-close-cute-re" : "i-mgc-tool-cute-re"
+            }`}
+          />
         </div>
+
+        <CollapseCssGroup>
+          <CollapseCss
+            collapseId={accordionValue}
+            hideArrow
+            className="group/collapse border-none"
+            title={
+              <div className="group/tool flex h-6 min-w-0 flex-1 items-center py-0">
+                <div className="text-text-secondary flex items-center gap-2 text-xs">
+                  <span>
+                    {hasError ? "Tool Failed:" : isCalling ? "Tool Calling:" : "Tool Called:"}
+                  </span>
+                  <span className={`truncate font-medium ${hasError ? "text-red" : "text-text"}`}>
+                    {toolName}
+                  </span>
+                </div>
+                <div className="ml-2 flex items-center justify-center opacity-0 transition-opacity duration-200 group-hover/tool:opacity-100">
+                  <i className="i-mgc-right-cute-re size-3 shrink-0 transition-transform duration-200 group-data-[state=open]/collapse:rotate-90" />
+                </div>
+              </div>
+            }
+            contentClassName="pb-0 pt-2"
+          >
+            <div className="space-y-2 text-xs">
+              {/* Show tool arguments if available */}
+              {hasArgs ? (
+                <div>
+                  <div className="text-text-secondary mb-1 font-medium">Arguments:</div>
+                  <JsonHighlighter
+                    className="text-text-tertiary bg-fill-secondary overflow-x-auto rounded p-2 text-[11px]"
+                    json={JSON.stringify(part.input, null, 2)}
+                  />
+                </div>
+              ) : null}
+
+              {/* Show tool result if available */}
+              {hasResult ? (
+                <div>
+                  <div className="text-text-secondary mb-1 font-medium">Result:</div>
+                  <JsonHighlighter
+                    className="text-text-tertiary bg-fill-secondary overflow-x-auto rounded p-2 text-[11px]"
+                    json={JSON.stringify(part.output, null, 2)}
+                  />
+                </div>
+              ) : null}
+
+              {/* Show error if available */}
+              {hasError && "errorText" in part ? (
+                <div>
+                  <div className="text-red mb-1 font-medium">Error:</div>
+                  <pre className="text-red bg-red/10 overflow-x-auto rounded p-2 text-[11px]">
+                    {String(part.errorText)}
+                  </pre>
+                </div>
+              ) : null}
+            </div>
+          </CollapseCss>
+        </CollapseCssGroup>
       </div>
     )
   },

apps/mobile/ios/Folo/Info.plist

@@ -53,7 +53,7 @@
 		</dict>
 	</array>
 	<key>CFBundleVersion</key>
-	<string>135</string>
+	<string>136</string>
 	<key>ITSAppUsesNonExemptEncryption</key>
 	<false/>
 	<key>LSApplicationCategoryType</key>