refactor: extract and clean <think> tags in ThreadContent

Add a helper `extractContentAndClean` that pulls out the content between `<think>` tags and removes all auxiliary tags from the final output.
Update the message rendering logic to use this helper for finalized messages that lack explicit stream events, ensuring that reasoning and final output are correctly separated and displayed.

Adjust the reasoning detection to consider extracted reasoning as well as stream events, clean the copy button to use the actual final output, and eliminate duplicate `StreamEvent` type definitions.

These changes improve message parsing accuracy and simplify the component’s handling of legacy messages that embed both reasoning and results in the same string.

web-app/src/containers/ScrollToBottom.tsx

@@ -8,6 +8,7 @@ import { cn } from '@/lib/utils'
 import { ArrowDown } from 'lucide-react'
 import { useTranslation } from '@/i18n/react-i18next-compat'
 import { useAppState } from '@/hooks/useAppState'
+import { MessageStatus } from '@janhq/core'
 
 const ScrollToBottom = ({
   threadId,
@@ -18,8 +19,10 @@ const ScrollToBottom = ({
 }) => {
   const { t } = useTranslation()
   const appMainViewBgColor = useAppearance((state) => state.appMainViewBgColor)
-  const { showScrollToBottomBtn, scrollToBottom } =
-    useThreadScrolling(threadId, scrollContainerRef)
+  const { showScrollToBottomBtn, scrollToBottom } = useThreadScrolling(
+    threadId,
+    scrollContainerRef
+  )
   const { messages } = useMessages(
     useShallow((state) => ({
       messages: state.messages[threadId],
@@ -28,12 +31,9 @@ const ScrollToBottom = ({
 
   const streamingContent = useAppState((state) => state.streamingContent)
 
+  const lastMsg = messages[messages.length - 1]
   const showGenerateAIResponseBtn =
-    (messages[messages.length - 1]?.role === 'user' ||
-      (messages[messages.length - 1]?.metadata &&
-        'tool_calls' in (messages[messages.length - 1].metadata ?? {}))) &&
-    !streamingContent
-
+    !!lastMsg && lastMsg.status !== MessageStatus.Ready && !streamingContent
   return (
     <div
       className={cn(

web-app/src/containers/ThinkingBlock.tsx

@@ -1,13 +1,43 @@
-import { ChevronDown, ChevronUp, Loader } from 'lucide-react'
+/* eslint-disable react-hooks/rules-of-hooks */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { ChevronDown, ChevronUp, Loader, Check } from 'lucide-react'
 import { create } from 'zustand'
 import { RenderMarkdown } from './RenderMarkdown'
-import { useAppState } from '@/hooks/useAppState'
 import { useTranslation } from '@/i18n/react-i18next-compat'
-import { extractThinkingContent } from '@/lib/utils'
+import { useMemo, useState } from 'react'
+import { cn } from '@/lib/utils'
+import ImageModal from '@/containers/dialogs/ImageModal'
+
+// Define ReActStep type (Reasoning-Action Step)
+type ReActStep = {
+  type: 'reasoning' | 'tool_call' | 'tool_output' | 'done'
+  content: string
+  metadata?: any
+  time?: number
+}
 
 interface Props {
   text: string
   id: string
+  steps?: ReActStep[] // Updated type
+  loading?: boolean
+  duration?: number
+  linkComponents?: object
+}
+
+// Utility function to safely parse JSON
+const safeParseJSON = (text: string) => {
+  try {
+    return JSON.parse(text)
+  } catch {
+    return null
+  }
+}
+
+// Utility to create data URL for images
+const createDataUrl = (base64Data: string, mimeType: string): string => {
+  if (base64Data.startsWith('data:')) return base64Data
+  return `data:${mimeType};base64,${base64Data}`
 }
 
 // Zustand store for thinking block state
@@ -27,54 +57,320 @@ const useThinkingStore = create<ThinkingBlockState>((set) => ({
     })),
 }))
 
-const ThinkingBlock = ({ id, text }: Props) => {
+// Helper to format duration in seconds
+const formatDuration = (ms: number) => {
+  if (ms > 0) {
+    return Math.round(ms / 1000)
+  }
+  return 0
+}
+
+const ThinkingBlock = ({
+  id,
+  steps = [],
+  loading: propLoading,
+  duration,
+  linkComponents,
+}: Props) => {
   const thinkingState = useThinkingStore((state) => state.thinkingState)
   const setThinkingState = useThinkingStore((state) => state.setThinkingState)
-  const isStreaming = useAppState((state) => !!state.streamingContent)
   const { t } = useTranslation()
-  // Check for thinking formats
-  const hasThinkTag = text.includes('<think>') && !text.includes('</think>')
-  const hasAnalysisChannel =
-    text.includes('<|channel|>analysis<|message|>') &&
-    !text.includes('<|start|>assistant<|channel|>final<|message|>')
-  const loading = (hasThinkTag || hasAnalysisChannel) && isStreaming
+
+  // Move useState for modal management to the top level of the component
+  const [modalImage, setModalImage] = useState<{
+    url: string
+    alt: string
+  } | null>(null)
+  const closeModal = () => setModalImage(null)
+  const handleImageClick = (url: string, alt: string) =>
+    setModalImage({ url, alt })
+
+  // Actual loading state comes from prop, determined by whether final text started streaming
+  const loading = propLoading
+
+  // Set default expansion state: collapsed if done (not loading).
+  // If loading transitions to false (textSegment starts), this defaults to collapsed.
   const isExpanded = thinkingState[id] ?? (loading ? true : false)
+
+  // Filter out the 'done' step for streaming display
+  const stepsWithoutDone = useMemo(
+    () => steps.filter((step) => step.type !== 'done'),
+    [steps]
+  )
+  const N = stepsWithoutDone.length
+
+  // Determine the step to display in the condensed streaming view
+  const activeStep = useMemo(() => {
+    if (!loading || N === 0) return null
+    return stepsWithoutDone[N - 1]
+  }, [loading, N, stepsWithoutDone])
+
+  // If not loading, and there are no steps, hide the block entirely.
+  const hasContent = steps.length > 0
+  if (!loading && !hasContent) return null
+
   const handleClick = () => {
-    const newExpandedState = !isExpanded
-    setThinkingState(id, newExpandedState)
+    // Only allow toggling expansion if not currently loading
+    // Also only allow if there is content (to prevent collapsing the simple 'Thinking')
+    if (!loading && hasContent) {
+      setThinkingState(id, !isExpanded)
+    }
   }
 
-  const thinkingContent = extractThinkingContent(text)
-  if (!thinkingContent) return null
+  // --- Rendering Functions for Expanded View ---
+  const renderStepContent = (
+    step: ReActStep,
+    index: number,
+    handleImageClick: (url: string, alt: string) => void,
+    t: (key: string) => string
+  ) => {
+    // Updated type
+    if (step.type === 'done') {
+      const timeInSeconds = formatDuration(step.time ?? 0)
+      const timeDisplay =
+        timeInSeconds > 0
+          ? `(${t('chat:for')} ${timeInSeconds} ${t('chat:seconds')})`
+          : ''
+
+      return (
+        <div
+          key={index}
+          className="flex items-center gap-1 text-accent transition-all"
+        >
+          <Check className="size-4" />
+          <span className="font-medium">{t('done')}</span>
+          {timeDisplay && (
+            <span className="text-main-view-fg/60 text-xs">{timeDisplay}</span>
+          )}
+        </div>
+      )
+    }
+
+    const parsed = safeParseJSON(step.content)
+    const mcpContent = parsed?.content ?? []
+    const hasImages =
+      Array.isArray(mcpContent) &&
+      mcpContent.some((c) => c.type === 'image' && c.data && c.mimeType)
+
+    let contentDisplay: React.ReactNode
+
+    if (step.type === 'tool_call') {
+      const args = step.metadata ? step.metadata : ''
+      contentDisplay = (
+        <>
+          <p className="font-medium text-main-view-fg/90">
+            Tool Input: <span className="text-accent">{step.content}</span>
+          </p>
+          {args && (
+            <div className="mt-1">
+              <RenderMarkdown
+                isWrapping={true}
+                content={'```json\n' + args + '\n```'}
+              />
+            </div>
+          )}
+        </>
+      )
+    } else if (step.type === 'tool_output') {
+      if (hasImages) {
+        // Display each image
+        contentDisplay = (
+          <>
+            <p className="font-medium text-main-view-fg/90">
+              Tool Output (Images):
+            </p>
+            <div className="mt-2 space-y-2">
+              {mcpContent.map((item: any, index: number) =>
+                item.type === 'image' && item.data && item.mimeType ? (
+                  <div key={index} className="my-2">
+                    <img
+                      src={createDataUrl(item.data, item.mimeType)}
+                      alt={`MCP Image ${index + 1}`}
+                      className="max-w-full max-h-64 object-contain rounded-md border border-main-view-fg/10 cursor-pointer hover:opacity-80 transition-opacity"
+                      onError={(e) => (e.currentTarget.style.display = 'none')}
+                      onClick={() =>
+                        handleImageClick(
+                          createDataUrl(item.data, item.mimeType),
+                          `MCP Image ${index + 1}`
+                        )
+                      }
+                    />
+                  </div>
+                ) : null
+              )}
+            </div>
+          </>
+        )
+      } else {
+        // Default behavior: wrap text in code block if no backticks
+        let content = step.content.substring(0, 1000)
+        if (!content.includes('```')) {
+          content = '```json\n' + content + '\n```'
+        }
+
+        contentDisplay = (
+          <>
+            <p className="font-medium text-main-view-fg/90">Tool Output:</p>
+            <div className="mt-1">
+              <RenderMarkdown
+                isWrapping={true}
+                content={content}
+                components={linkComponents}
+              />
+            </div>
+          </>
+        )
+      }
+    } else {
+      contentDisplay = (
+        <RenderMarkdown
+          isWrapping={true}
+          content={step.content}
+          components={linkComponents}
+        />
+      )
+    }
+
+    return (
+      <div key={index} className="text-main-view-fg/80">
+        {contentDisplay}
+      </div>
+    )
+  }
+
+  const headerTitle: string = useMemo(() => {
+    // Check if any step was a tool call
+    const hasToolCalls = steps.some((step) => step.type === 'tool_call')
+    const hasReasoning = steps.some((step) => step.type === 'reasoning')
+
+    const timeInSeconds = formatDuration(duration ?? 0)
+
+    if (loading) {
+      // Logic for streaming (loading) state:
+      if (activeStep) {
+        if (
+          activeStep.type === 'tool_call' ||
+          activeStep.type === 'tool_output'
+        ) {
+          return `${t('chat:calling_tool')}` // Use a specific translation key for tool
+        } else if (activeStep.type === 'reasoning') {
+          return `${t('chat:thinking')}` // Use the generic thinking key
+        }
+      }
+
+      // Fallback for isStreamingEmpty state (N=0)
+      return `${t('chat:thinking')}`
+    }
+
+    // Logic for finalized (not loading) state:
+    // Build label based on what steps occurred
+    let label = ''
+    if (hasReasoning && hasToolCalls) {
+      // Use a more descriptive label when both were involved
+      label = t('chat:thought_and_tool_call')
+    } else if (hasToolCalls) {
+      label = t('chat:tool_called')
+    } else {
+      label = t('chat:thought')
+    }
+
+    if (timeInSeconds > 0) {
+      return `${label} ${t('chat:for')} ${timeInSeconds} ${t('chat:seconds')}`
+    }
+    return label
+  }, [loading, duration, t, activeStep, steps])
 
   return (
     <div
-      className="mx-auto w-full cursor-pointer break-words"
-      onClick={handleClick}
+      className="mx-auto w-full break-words"
+      onClick={loading || !hasContent ? undefined : handleClick}
     >
-      <div className="mb-4 rounded-lg bg-main-view-fg/4 border border-dashed border-main-view-fg/10 p-2">
+      <div className="mb-4 rounded-lg bg-main-view-fg/4 p-2 transition-all duration-200">
         <div className="flex items-center gap-3">
           {loading && (
             <Loader className="size-4 animate-spin text-main-view-fg/60" />
           )}
-          <button className="flex items-center gap-2 focus:outline-none">
-            {isExpanded ? (
-              <ChevronUp className="size-4 text-main-view-fg/60" />
-            ) : (
-              <ChevronDown className="size-4 text-main-view-fg/60" />
-            )}
-            <span className="font-medium">
-              {loading ? t('common:thinking') : t('common:thought')}
+          <button
+            className="flex items-center gap-2 focus:outline-none"
+            disabled={loading || !hasContent}
+          >
+            {/* Display chevron only if not loading AND steps exist to expand */}
+            {!loading &&
+              hasContent &&
+              (isExpanded ? (
+                <ChevronUp className="size-4 text-main-view-fg/60 transition-transform duration-200" />
+              ) : (
+                <ChevronDown className="size-4 text-main-view-fg/60 transition-transform duration-200" />
+              ))}
+            <span className="font-medium transition-all duration-200">
+              {headerTitle}
             </span>
           </button>
         </div>
 
-        {isExpanded && (
-          <div className="mt-2 pl-6 pr-4 text-main-view-fg/60">
-            <RenderMarkdown content={thinkingContent} />
+        {/* Streaming/Condensed View - shows active step (N-1) */}
+        {/* This block handles both the N>0 case and the N=0 fallback, ensuring stability. */}
+        {loading && (activeStep || N === 0) && (
+          <div
+            key={`streaming-${N - 1}`}
+            className={cn(
+              'mt-4 pl-2 pr-4 text-main-view-fg/60',
+              // Only animate fade-in if it's not the very first step (N > 1)
+              N > 1 && 'animate-in fade-in slide-in-from-top-2 duration-300'
+            )}
+          >
+            <div className="relative border-main-view-fg/20">
+              {/* If N=0, just show the fallback text in the header and this area remains minimal. */}
+              {activeStep && (
+                <div className="relative pl-5">
+                  {/* Bullet point/Icon position relative to line */}
+                  <div
+                    className={cn(
+                      'absolute left-[-2px] top-1.5 size-2 rounded-full bg-main-view-fg/60',
+                      activeStep.type !== 'done' && 'animate-pulse' // Pulse if active/streaming
+                    )}
+                  />
+                  {/* Active step content */}
+                  {renderStepContent(activeStep, N - 1, handleImageClick, t)}
+                </div>
+              )}
+            </div>
+          </div>
+        )}
+
+        {/* Expanded View - shows all steps */}
+        {isExpanded && !loading && hasContent && (
+          <div className="mt-4 pl-2 pr-4 text-main-view-fg/60 animate-in fade-in slide-in-from-top-2 duration-300">
+            <div className="relative border-main-view-fg/20">
+              {steps.map((step, index) => (
+                <div
+                  key={index}
+                  className={cn(
+                    'relative pl-5 pb-2',
+                    'fade-in slide-in-from-left-1 duration-200',
+                    step.type !== 'done' &&
+                      'after:content-[] after:border-l after:border-dashed after:border-main-view-fg/20 after:absolute after:left-0.5 after:bottom-0 after:w-1 after:h-[calc(100%-8px)]'
+                  )}
+                  style={{ animationDelay: `${index * 50}ms` }}
+                >
+                  {/* Bullet point/Icon position relative to line */}
+                  <div
+                    className={cn(
+                      'absolute left-[-2px] top-1.5 size-2 rounded-full transition-colors duration-200',
+                      step.type === 'done' ? 'bg-accent' : 'bg-main-view-fg/60'
+                    )}
+                  />
+
+                  {/* Step Content */}
+                  {renderStepContent(step, index, handleImageClick, t)}
+                </div>
+              ))}
+            </div>
           </div>
         )}
       </div>
+      {/* Render ImageModal once at the top level */}
+      <ImageModal image={modalImage} onClose={closeModal} />
     </div>
   )
 }

web-app/src/containers/ThreadContent.tsx

@@ -7,7 +7,7 @@ import { useAppState } from '@/hooks/useAppState'
 import { cn } from '@/lib/utils'
 import { useMessages } from '@/hooks/useMessages'
 import ThinkingBlock from '@/containers/ThinkingBlock'
-import ToolCallBlock from '@/containers/ToolCallBlock'
+// import ToolCallBlock from '@/containers/ToolCallBlock'
 import { useChat } from '@/hooks/useChat'
 import {
   EditMessageDialog,
@@ -29,6 +29,62 @@ import { useModelProvider } from '@/hooks/useModelProvider'
 import { extractFilesFromPrompt } from '@/lib/fileMetadata'
 import { createImageAttachment } from '@/types/attachment'
 
+// Define ToolCall interface for type safety when accessing metadata
+interface ToolCall {
+  tool?: {
+    id: number
+    function?: {
+      name: string
+      arguments?: object | string
+    }
+  }
+  response?: any
+  state?: 'pending' | 'completed'
+}
+
+// Define ReActStep type (Reasoning-Action Step)
+type ReActStep = {
+  type: 'reasoning' | 'tool_call' | 'tool_output' | 'done'
+  content: string
+  metadata?: any
+  time?: number
+}
+
+const cleanReasoning = (content: string) => {
+  return content
+    .replace(/^<think>/, '') // Remove opening tag at start
+    .replace(/<\/think>$/, '') // Remove closing tag at end
+    .trim()
+}
+
+// Helper function to extract content within <think> tags and strip all auxiliary tags from the final output
+const extractContentAndClean = (
+  rawText: string
+): { reasoningText: string; finalOutput: string } => {
+  // Regex to match content within <think>...</think> tags
+  const thinkTagRegex = /<think>([\s\S]*?)<\/think>/g
+
+  let reasoningText = ''
+  let finalOutput = rawText
+
+  // Extract content within <think> tags for streamedReasoningText
+  const thinkMatches = [...rawText.matchAll(thinkTagRegex)]
+  if (thinkMatches.length > 0) {
+    // Join all reasoning parts separated by newlines
+    reasoningText = thinkMatches
+      .map((match) => match[1])
+      .join('\n\n')
+      .trim()
+  }
+
+  // 2. Strip ALL auxiliary tags from finalOutput
+  finalOutput = finalOutput
+    .replace(thinkTagRegex, '') // Remove <think> tags and content
+    .trim()
+
+  return { reasoningText, finalOutput }
+}
+
 const CopyButton = ({ text }: { text: string }) => {
   const [copied, setCopied] = useState(false)
   const { t } = useTranslation()
@@ -112,40 +168,71 @@ export const ThreadContent = memo(
       return { files: [], cleanPrompt: text }
     }, [text, item.role])
 
-    const { reasoningSegment, textSegment } = useMemo(() => {
-      // Check for thinking formats
-      const hasThinkTag = text.includes('<think>') && !text.includes('</think>')
-      const hasAnalysisChannel =
-        text.includes('<|channel|>analysis<|message|>') &&
-        !text.includes('<|start|>assistant<|channel|>final<|message|>')
+    type StreamEvent = {
+      timestamp: number
+      type: 'reasoning_chunk' | 'tool_call' | 'tool_output'
+      data: any
+    }
 
-      if (hasThinkTag || hasAnalysisChannel)
-        return { reasoningSegment: text, textSegment: '' }
+    const {
+      finalOutputText,
+      streamedReasoningText,
+      isReasoningActiveLoading,
+      hasReasoningSteps,
+    } = useMemo(() => {
+      let currentFinalText = text.trim()
+      let currentReasoning = '' // Reasoning is now only derived from streamEvents/allSteps
 
-      // Check for completed think tag format
-      const thinkMatch = text.match(/<think>([\s\S]*?)<\/think>/)
-      if (thinkMatch?.index !== undefined) {
-        const splitIndex = thinkMatch.index + thinkMatch[0].length
-        return {
-          reasoningSegment: text.slice(0, splitIndex),
-          textSegment: text.slice(splitIndex),
-        }
+      // Extract raw streamEvents and check for finalized state
+      const streamEvents = (item.metadata?.streamEvents as StreamEvent[]) || []
+      const isMessageFinalized = !isStreamingThisThread
+
+      // If the message is finalized AND there are no streamEvents,
+      // we assume the 'text' contains the full dump (reasoning + output + tool tags)
+      if (isMessageFinalized && streamEvents.length === 0) {
+        // Use the new helper to separate reasoning (from <think>) and clean the final output
+        const { reasoningText, finalOutput } = extractContentAndClean(text)
+        currentFinalText = finalOutput
+        currentReasoning = reasoningText
+      } else {
+        // Otherwise, trust the streamEvents path (if present) or the current text is the final output
+        // We clean the current text just in case, but it should be clean in streaming mode
+        const { finalOutput } = extractContentAndClean(text)
+        currentFinalText = finalOutput
       }
 
-      // Check for completed analysis channel format
-      const analysisMatch = text.match(
-        /<\|channel\|>analysis<\|message\|>([\s\S]*?)<\|start\|>assistant<\|channel\|>final<\|message\|>/
+      // Check for tool calls or reasoning events in metadata to determine steps/loading
+      const isToolCallsPresent = !!(
+        item.metadata &&
+        'tool_calls' in item.metadata &&
+        Array.isArray(item.metadata.tool_calls) &&
+        item.metadata.tool_calls.length > 0
       )
-      if (analysisMatch?.index !== undefined) {
-        const splitIndex = analysisMatch.index + analysisMatch[0].length
-        return {
-          reasoningSegment: text.slice(0, splitIndex),
-          textSegment: text.slice(splitIndex),
-        }
-      }
 
-      return { reasoningSegment: undefined, textSegment: text }
-    }, [text])
+      // Check for any reasoning chunks in the streamEvents OR if we extracted reasoning from text
+      const hasReasoningEvents =
+        streamEvents.some((e: StreamEvent) => e.type === 'reasoning_chunk') ||
+        currentReasoning.length > 0 // Added check for extracted reasoning
+
+      const hasSteps = isToolCallsPresent || hasReasoningEvents
+
+      // Loading if streaming, no final output yet, but we expect steps (reasoning or tool calls)
+      const loading =
+        isStreamingThisThread && currentFinalText.length === 0 && hasSteps
+
+      return {
+        finalOutputText: currentFinalText,
+        streamedReasoningText: currentReasoning,
+        isReasoningActiveLoading: loading,
+        hasReasoningSteps: hasSteps,
+      }
+    }, [item.metadata, text, isStreamingThisThread])
+
+    const isToolCalls =
+      item.metadata &&
+      'tool_calls' in item.metadata &&
+      Array.isArray(item.metadata.tool_calls) &&
+      item.metadata.tool_calls.length
 
     const getMessages = useMessages((state) => state.getMessages)
     const deleteMessage = useMessages((state) => state.deleteMessage)
@@ -164,7 +251,8 @@ export const ThreadContent = memo(
         deleteMessage(toSendMessage.thread_id, toSendMessage.id ?? '')
         // Extract text content and any attachments
         const rawText =
-          toSendMessage.content?.find((c) => c.type === 'text')?.text?.value || ''
+          toSendMessage.content?.find((c) => c.type === 'text')?.text?.value ||
+          ''
         const { cleanPrompt: textContent } = extractFilesFromPrompt(rawText)
         const attachments = toSendMessage.content
           ?.filter((c) => (c.type === 'image_url' && c.image_url?.url) || false)
@@ -203,7 +291,7 @@ export const ThreadContent = memo(
         while (toSendMessage && toSendMessage?.role !== 'user') {
           deleteMessage(toSendMessage.thread_id, toSendMessage.id ?? '')
           toSendMessage = threadMessages.pop()
-          // Stop deletion when encountering an assistant message that isn’t a tool call
+          // Stop deletion when encountering an assistant message that isn't a tool call
           if (
             toSendMessage &&
             toSendMessage.role === 'assistant' &&
@@ -216,16 +304,235 @@ export const ThreadContent = memo(
       }
     }, [deleteMessage, getMessages, item])
 
-    const isToolCalls =
-      item.metadata &&
-      'tool_calls' in item.metadata &&
-      Array.isArray(item.metadata.tool_calls) &&
-      item.metadata.tool_calls.length
-
     const assistant = item.metadata?.assistant as
       | { avatar?: React.ReactNode; name?: React.ReactNode }
       | undefined
 
+    // Constructing allSteps for ThinkingBlock - CHRONOLOGICAL approach
+    const allSteps: ReActStep[] = useMemo(() => {
+      const steps: ReActStep[] = []
+
+      // Get streamEvents from metadata (if available)
+      const streamEvents = (item.metadata?.streamEvents as StreamEvent[]) || []
+      const toolCalls = (item.metadata?.tool_calls || []) as ToolCall[]
+
+      const isMessageFinalized = !isStreamingThisThread
+
+      if (streamEvents.length > 0) {
+        // CHRONOLOGICAL PATH: Use streamEvents for true temporal order
+        let reasoningBuffer = ''
+
+        streamEvents.forEach((event) => {
+          switch (event.type) {
+            case 'reasoning_chunk':
+              // Accumulate reasoning chunks
+              reasoningBuffer += event.data.content
+              break
+
+            case 'tool_call':
+            case 'tool_output':
+              // Flush accumulated reasoning before tool event
+              if (reasoningBuffer.trim()) {
+                const cleanedBuffer = cleanReasoning(reasoningBuffer) // <--- Strip tags here
+
+                // Split accumulated reasoning by paragraphs for display
+                const paragraphs = cleanedBuffer
+                  .split(/\n\s*\n/)
+                  .filter((p) => p.trim().length > 0)
+
+                paragraphs.forEach((para) => {
+                  steps.push({
+                    type: 'reasoning',
+                    content: para.trim(),
+                  })
+                })
+
+                reasoningBuffer = ''
+              }
+
+              if (event.type === 'tool_call') {
+                // Add tool call
+                const toolCall = event.data.toolCall
+                steps.push({
+                  type: 'tool_call',
+                  content: toolCall?.function?.name || 'Tool Call',
+                  metadata:
+                    typeof toolCall?.function?.arguments === 'string'
+                      ? toolCall.function.arguments
+                      : JSON.stringify(
+                          toolCall?.function?.arguments || {},
+                          null,
+                          2
+                        ),
+                })
+              } else if (event.type === 'tool_output') {
+                // Add tool output
+                const result = event.data.result
+                let outputContent = JSON.stringify(result, null, 2) // Default fallback
+
+                const firstContentPart = result?.content?.[0]
+
+                if (firstContentPart?.type === 'text') {
+                  const textContent = firstContentPart.text
+                  // Robustly check for { value: string } structure or direct string
+                  if (
+                    typeof textContent === 'object' &&
+                    textContent !== null &&
+                    'value' in textContent
+                  ) {
+                    outputContent = textContent.value as string
+                  } else if (typeof textContent === 'string') {
+                    outputContent = textContent
+                  }
+                } else if (typeof result === 'string') {
+                  outputContent = result
+                }
+
+                steps.push({
+                  type: 'tool_output',
+                  content: outputContent,
+                })
+              }
+              break
+          }
+        })
+
+        // Flush any remaining reasoning at the end
+        if (reasoningBuffer.trim()) {
+          const cleanedBuffer = cleanReasoning(reasoningBuffer) // <--- Strip tags here
+          const paragraphs = cleanedBuffer
+            .split(/\n\s*\n/)
+            .filter((p) => p.trim().length > 0)
+
+          paragraphs.forEach((para) => {
+            steps.push({
+              type: 'reasoning',
+              content: para.trim(),
+            })
+          })
+        }
+      } else if (isMessageFinalized) {
+        // FALLBACK PATH: No streamEvents - use split text for content construction
+
+        const rawReasoningContent = streamedReasoningText || ''
+        const reasoningParagraphs = rawReasoningContent
+          ? rawReasoningContent // streamedReasoningText is now populated from <think> tags if present
+              .split(/\n\s*\n/)
+              .filter((s) => s.trim().length > 0)
+              .map((content) => content.trim())
+          : []
+
+        let reasoningIndex = 0
+
+        toolCalls.forEach((call) => {
+          // Add reasoning before this tool call
+          if (reasoningIndex < reasoningParagraphs.length) {
+            steps.push({
+              type: 'reasoning',
+              content: reasoningParagraphs[reasoningIndex],
+            })
+            reasoningIndex++
+          }
+
+          // Add tool call
+          steps.push({
+            type: 'tool_call',
+            content: call.tool?.function?.name || 'Tool Call',
+            metadata:
+              typeof call.tool?.function?.arguments === 'string'
+                ? call.tool.function.arguments
+                : JSON.stringify(call.tool?.function?.arguments || {}, null, 2),
+          })
+
+          // Add tool output
+          if (call.response) {
+            const result = call.response
+            let outputContent = JSON.stringify(result, null, 2)
+
+            const firstContentPart = result?.content?.[0]
+
+            if (firstContentPart?.type === 'text') {
+              const textContent = firstContentPart.text
+              if (
+                typeof textContent === 'object' &&
+                textContent !== null &&
+                'value' in textContent
+              ) {
+                outputContent = textContent.value as string
+              } else if (typeof textContent === 'string') {
+                outputContent = textContent
+              }
+            } else if (typeof result === 'string') {
+              outputContent = result
+            }
+
+            steps.push({
+              type: 'tool_output',
+              content: outputContent,
+            })
+          }
+        })
+
+        // Add remaining reasoning
+        while (reasoningIndex < reasoningParagraphs.length) {
+          steps.push({
+            type: 'reasoning',
+            content: reasoningParagraphs[reasoningIndex],
+          })
+          reasoningIndex++
+        }
+      }
+
+      // Add Done step
+      const totalTime = item.metadata?.totalThinkingTime as number | undefined
+      const lastStepType = steps[steps.length - 1]?.type
+
+      if (!isStreamingThisThread && hasReasoningSteps) {
+        const endsInToolOutputWithoutFinalText =
+          lastStepType === 'tool_output' && finalOutputText.length === 0
+
+        if (!endsInToolOutputWithoutFinalText) {
+          steps.push({
+            type: 'done',
+            content: 'Done',
+            time: totalTime,
+          })
+        }
+      }
+
+      return steps
+    }, [
+      item,
+      isStreamingThisThread,
+      hasReasoningSteps,
+      finalOutputText,
+      streamedReasoningText,
+    ])
+    // END: Constructing allSteps
+
+    // ====================================================================
+    // If we have streamEvents, rely on 'steps' and pass an empty text buffer.
+    const streamingTextBuffer = useMemo(() => {
+      const streamEvents = item.metadata?.streamEvents
+
+      // Check if streamEvents exists AND is an array AND has a length greater than 0
+      if (Array.isArray(streamEvents) && streamEvents.length > 0) {
+        // We are using the chronological path (allSteps) for rendering
+        // Return empty string to disable the ThinkingBlock's raw text buffer
+        return ''
+      }
+
+      // Since we no longer concatenate reasoning to the main text,
+      // the only time we'd rely on text buffer is if streamEvents fails to load.
+      // For robustness, we can simply return an empty string to force use of 'steps'.
+      return ''
+    }, [item.metadata?.streamEvents]) // Use the object reference for dependency array
+    // ====================================================================
+
+    // Determine if we should show the thinking block
+    const shouldShowThinkingBlock =
+      hasReasoningSteps || isToolCalls || isReasoningActiveLoading
+
     return (
       <Fragment>
         {item.role === 'user' && (
@@ -360,46 +667,33 @@ export const ThreadContent = memo(
               </div>
             )}
 
-            {reasoningSegment && (
+            {/* Single unified ThinkingBlock for both reasoning and tool calls */}
+            {shouldShowThinkingBlock && (
               <ThinkingBlock
                 id={
                   item.isLastMessage
-                    ? `${item.thread_id}-last-${reasoningSegment.slice(0, 50).replace(/\s/g, '').slice(-10)}`
+                    ? `${item.thread_id}-last-${(streamingTextBuffer || text).slice(0, 50).replace(/\s/g, '').slice(-10)}`
                     : `${item.thread_id}-${item.index ?? item.id}`
                 }
-                text={reasoningSegment}
+                // Pass the safe buffer
+                text={streamingTextBuffer}
+                steps={allSteps}
+                loading={isReasoningActiveLoading}
+                duration={
+                  item.metadata?.totalThinkingTime as number | undefined
+                }
+                linkComponents={linkComponents}
               />
             )}
 
-            <RenderMarkdown
-              content={textSegment.replace('</think>', '')}
-              components={linkComponents}
-            />
+            {!isReasoningActiveLoading && finalOutputText.length > 0 && (
+              <RenderMarkdown
+                content={finalOutputText}
+                components={linkComponents}
+              />
+            )}
 
-            {isToolCalls && item.metadata?.tool_calls ? (
-              <>
-                {(item.metadata.tool_calls as ToolCall[]).map((toolCall) => (
-                  <ToolCallBlock
-                    id={toolCall.tool?.id ?? 0}
-                    key={toolCall.tool?.id}
-                    name={
-                      (item.streamTools?.tool_calls?.function?.name ||
-                        toolCall.tool?.function?.name) ??
-                      ''
-                    }
-                    args={
-                      item.streamTools?.tool_calls?.function?.arguments ||
-                      toolCall.tool?.function?.arguments ||
-                      undefined
-                    }
-                    result={JSON.stringify(toolCall.response)}
-                    loading={toolCall.state === 'pending'}
-                  />
-                ))}
-              </>
-            ) : null}
-
-            {!isToolCalls && (
+            {
               <div className="flex items-center gap-2 text-main-view-fg/60 text-xs">
                 <div className={cn('flex items-center gap-2')}>
                   <div
@@ -414,10 +708,10 @@ export const ThreadContent = memo(
                         item.updateMessage && item.updateMessage(item, message)
                       }
                     />
-                    <CopyButton text={item.content?.[0]?.text.value || ''} />
+                    <CopyButton text={finalOutputText || ''} />{' '}
+                    {/* Use finalOutputText for copy */}
                     <DeleteMessageDialog onDelete={removeMessage} />
                     <MessageMetadataDialog metadata={item.metadata} />
-
                     {item.isLastMessage && selectedModel && (
                       <Tooltip>
                         <TooltipTrigger asChild>
@@ -443,7 +737,7 @@ export const ThreadContent = memo(
                   />
                 </div>
               </div>
-            )}
+            }
           </>
         )}
 

web-app/src/containers/dialogs/MessageMetadataDialog.tsx

@@ -1,4 +1,4 @@
-import { useState } from 'react'
+import { useState, useMemo } from 'react'
 import { useTranslation } from '@/i18n/react-i18next-compat'
 import {
   Dialog,
@@ -7,21 +7,96 @@ import {
   DialogTitle,
   DialogHeader,
 } from '@/components/ui/dialog'
-import { IconInfoCircle } from '@tabler/icons-react'
+import {
+  IconInfoCircle,
+  IconRobot,
+  IconGauge,
+  IconId,
+  IconCalendar,
+  IconTemperature,
+  IconHierarchy,
+  IconTool,
+  IconBoxMultiple,
+  IconRuler,
+  IconMessageCircle,
+} from '@tabler/icons-react'
 import {
   Tooltip,
   TooltipContent,
   TooltipTrigger,
 } from '@/components/ui/tooltip'
-import CodeEditor from '@uiw/react-textarea-code-editor'
-import '@uiw/react-textarea-code-editor/dist.css'
+// Removed CodeEditor and its styles
+
+// Type definitions for the provided metadata structure
+interface Parameters {
+  temperature: number
+  top_k: number
+  top_p: number
+}
+
+interface AssistantMetadata {
+  avatar: string
+  created_at: number
+  description: string
+  id: string
+  instructions: string
+  name: string
+  parameters: Parameters
+  tool_steps: number
+}
+
+interface TokenSpeedMetadata {
+  lastTimestamp: number
+  message: string
+  tokenCount: number
+  tokenSpeed: number
+}
+
+interface MessageMetadata {
+  assistant?: AssistantMetadata
+  tokenSpeed?: TokenSpeedMetadata
+}
 
 interface MessageMetadataDialogProps {
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  metadata: any
+  metadata: MessageMetadata | null | undefined // Use the specific interface
   triggerElement?: React.ReactNode
 }
 
+// --- Helper Components & Utilities ---
+
+// A utility component to display a single detail row
+const DetailItem: React.FC<{
+  icon: React.ReactNode
+  label: string
+  value: React.ReactNode
+}> = ({ icon, label, value }) => (
+  <div className="flex items-start text-sm p-2 bg-main-view-bg/5 rounded-md">
+    <div className="text-accent mr-3 flex-shrink-0">{icon}</div>
+    <div className="flex flex-col">
+      <span className="font-semibold text-main-view-fg/80">{label}:</span>
+      <span className="text-main-view-fg/90 whitespace-pre-wrap break-words">
+        {value}
+      </span>
+    </div>
+  </div>
+)
+
+// Helper for formatting timestamps
+const formatDate = (timestamp: number) => {
+  if (!timestamp) return 'N/A'
+  return new Intl.DateTimeFormat('en-US', {
+    year: 'numeric',
+    month: 'short',
+    day: 'numeric',
+    hour: 'numeric',
+    minute: 'numeric',
+    second: 'numeric',
+    timeZoneName: 'short',
+  }).format(new Date(timestamp))
+}
+
+// --- Main Component ---
+
 export function MessageMetadataDialog({
   metadata,
   triggerElement,
@@ -29,10 +104,12 @@ export function MessageMetadataDialog({
   const { t } = useTranslation()
   const [isOpen, setIsOpen] = useState(false)
 
+  const { assistant, tokenSpeed } = (metadata || {}) as MessageMetadata
+
   const defaultTrigger = (
     <Tooltip>
       <TooltipTrigger asChild>
-        <div 
+        <div
           className="outline-0 focus:outline-0 flex items-center gap-1 hover:text-accent transition-colors cursor-pointer group relative"
           role="button"
           tabIndex={0}
@@ -52,27 +129,127 @@ export function MessageMetadataDialog({
     </Tooltip>
   )
 
+  const formattedTokenSpeed = useMemo(() => {
+    if (tokenSpeed?.tokenSpeed === undefined) return 'N/A'
+    return (
+      new Intl.NumberFormat('en-US', {
+        style: 'decimal',
+        minimumFractionDigits: 2,
+        maximumFractionDigits: 2,
+      }).format(tokenSpeed.tokenSpeed) + ' tokens/s'
+    )
+  }, [tokenSpeed])
+
   return (
     <Dialog open={isOpen} onOpenChange={setIsOpen}>
       <DialogTrigger>{triggerElement || defaultTrigger}</DialogTrigger>
-      <DialogContent>
+      <DialogContent className="max-w-xl">
         <DialogHeader>
           <DialogTitle>{t('common:dialogs.messageMetadata.title')}</DialogTitle>
-          <div className="space-y-2 mt-4">
-            <div className="border border-main-view-fg/10 rounded-md">
-              <CodeEditor
-                value={JSON.stringify(metadata || {}, null, 2)}
-                language="json"
-                readOnly
-                data-color-mode="dark"
-                style={{
-                  fontSize: 12,
-                  backgroundColor: 'transparent',
-                  fontFamily: 'monospace',
-                }}
-                className="w-full h-full !text-sm "
-              />
-            </div>
+          <div className="space-y-6 mt-4 max-h-[70vh] overflow-y-auto pr-2">
+            {/* --- Assistant/Model Section --- */}
+            {assistant && (
+              <section>
+                <h3 className="flex items-center text-lg font-bold border-b border-main-view-fg/10 pb-2 mb-3">
+                  <IconRobot className="mr-2" size={20} />
+                  {t('common:dialogs.messageMetadata.model')}
+                </h3>
+                <div className="grid grid-cols-1 md:grid-cols-2 gap-3">
+                  <DetailItem
+                    icon={<IconRobot size={18} />}
+                    label={t('common:dialogs.messageMetadata.name')}
+                    value={`${assistant.avatar} ${assistant.name}`}
+                  />
+                  <DetailItem
+                    icon={<IconId size={18} />}
+                    label={t('common:dialogs.messageMetadata.id')}
+                    value={assistant.id}
+                  />
+                  <DetailItem
+                    icon={<IconCalendar size={18} />}
+                    label={t('common:dialogs.messageMetadata.createdAt')}
+                    value={formatDate(assistant.created_at)}
+                  />
+                  <DetailItem
+                    icon={<IconTool size={18} />}
+                    label={t('common:dialogs.messageMetadata.toolSteps')}
+                    value={assistant.tool_steps}
+                  />
+
+                  {/* Parameters */}
+                  <div className="col-span-1 md:col-span-2 grid grid-cols-3 gap-3">
+                    <DetailItem
+                      icon={<IconTemperature size={18} />}
+                      label={t('common:dialogs.messageMetadata.temperature')}
+                      value={assistant.parameters.temperature}
+                    />
+                    <DetailItem
+                      icon={<IconHierarchy size={18} />}
+                      label={t('common:dialogs.messageMetadata.topK')}
+                      value={assistant.parameters.top_k}
+                    />
+                    <DetailItem
+                      icon={<IconBoxMultiple size={18} />}
+                      label={t('common:dialogs.messageMetadata.topP')}
+                      value={assistant.parameters.top_p}
+                    />
+                  </div>
+
+                  {/* Description/Instructions */}
+                  {(assistant.description || assistant.instructions) && (
+                    <div className="col-span-1 md:col-span-2 space-y-3">
+                      {assistant.description && (
+                        <DetailItem
+                          icon={<IconMessageCircle size={18} />}
+                          label={t('common:dialogs.messageMetadata.description')}
+                          value={assistant.description}
+                        />
+                      )}
+                      {assistant.instructions && (
+                        <DetailItem
+                          icon={<IconMessageCircle size={18} />}
+                          label={t('common:dialogs.messageMetadata.instructions')}
+                          value={assistant.instructions}
+                        />
+                      )}
+                    </div>
+                  )}
+                </div>
+              </section>
+            )}
+
+            {/* --- Token Speed Section --- */}
+            {tokenSpeed && (
+              <section>
+                <h3 className="flex items-center text-lg font-bold border-b border-main-view-fg/10 pb-2 mb-3">
+                  <IconGauge className="mr-2" size={20} />
+                  {t('Performance')}
+                </h3>
+                <div className="grid grid-cols-1 md:grid-cols-3 gap-3">
+                  <DetailItem
+                    icon={<IconGauge size={18} />}
+                    label={t('common:dialogs.messageMetadata.tokenSpeed')}
+                    value={formattedTokenSpeed}
+                  />
+                  <DetailItem
+                    icon={<IconRuler size={18} />}
+                    label={t('common:dialogs.messageMetadata.tokenCount')}
+                    value={tokenSpeed.tokenCount}
+                  />
+                  <DetailItem
+                    icon={<IconCalendar size={18} />}
+                    label={t('common:dialogs.messageMetadata.lastUpdate')}
+                    value={formatDate(tokenSpeed.lastTimestamp)}
+                  />
+                </div>
+              </section>
+            )}
+
+            {!assistant && !tokenSpeed && (
+              <p className="text-center text-main-view-fg/70 py-4">
+                {t('common:dialogs.messageMetadata.noMetadataAvailable.')}
+              </p>
+            )}
           </div>
         </DialogHeader>
       </DialogContent>

web-app/src/hooks/useChat.ts

@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
 import { useCallback, useMemo } from 'react'
 import { flushSync } from 'react-dom'
 import { usePrompt } from './usePrompt'
@@ -41,6 +42,12 @@ import { TEMPORARY_CHAT_QUERY_ID, TEMPORARY_CHAT_ID } from '@/constants/chat'
 import { toast } from 'sonner'
 import { Attachment } from '@/types/attachment'
 
+type StreamEvent = {
+  timestamp: number
+  type: 'reasoning_chunk' | 'tool_call' | 'tool_output'
+  data: any
+}
+
 export const useChat = () => {
   const [
     updateTokenSpeed,
@@ -93,66 +100,74 @@ export const useChat = () => {
   const setModelLoadError = useModelLoad((state) => state.setModelLoadError)
   const router = useRouter()
 
-  const getCurrentThread = useCallback(async (projectId?: string) => {
-    let currentThread = retrieveThread()
+  const getCurrentThread = useCallback(
+    async (projectId?: string) => {
+      let currentThread = retrieveThread()
 
-    // Check if we're in temporary chat mode
-    const isTemporaryMode = window.location.search.includes(`${TEMPORARY_CHAT_QUERY_ID}=true`)
-
-    // Clear messages for existing temporary thread on reload to ensure fresh start
-    if (isTemporaryMode && currentThread?.id === TEMPORARY_CHAT_ID) {
-      setMessages(TEMPORARY_CHAT_ID, [])
-    }
-
-    if (!currentThread) {
-      // Get prompt directly from store when needed
-      const currentPrompt = usePrompt.getState().prompt
-      const currentAssistant = useAssistant.getState().currentAssistant
-      const assistants = useAssistant.getState().assistants
-      const selectedModel = useModelProvider.getState().selectedModel
-      const selectedProvider = useModelProvider.getState().selectedProvider
-
-      // Get project metadata if projectId is provided
-      let projectMetadata: { id: string; name: string; updated_at: number } | undefined
-      if (projectId) {
-        const project = await serviceHub.projects().getProjectById(projectId)
-        if (project) {
-          projectMetadata = {
-            id: project.id,
-            name: project.name,
-            updated_at: project.updated_at,
-          }
-        }
-      }
-
-      currentThread = await createThread(
-        {
-          id: selectedModel?.id ?? defaultModel(selectedProvider),
-          provider: selectedProvider,
-        },
-        isTemporaryMode ? 'Temporary Chat' : currentPrompt,
-        assistants.find((a) => a.id === currentAssistant?.id) || assistants[0],
-        projectMetadata,
-        isTemporaryMode // pass temporary flag
+      // Check if we're in temporary chat mode
+      const isTemporaryMode = window.location.search.includes(
+        `${TEMPORARY_CHAT_QUERY_ID}=true`
       )
 
-      // Clear messages for temporary chat to ensure fresh start on reload
+      // Clear messages for existing temporary thread on reload to ensure fresh start
       if (isTemporaryMode && currentThread?.id === TEMPORARY_CHAT_ID) {
         setMessages(TEMPORARY_CHAT_ID, [])
       }
 
-      // Set flag for temporary chat navigation
-      if (currentThread.id === TEMPORARY_CHAT_ID) {
-        sessionStorage.setItem('temp-chat-nav', 'true')
-      }
+      if (!currentThread) {
+        // Get prompt directly from store when needed
+        const currentPrompt = usePrompt.getState().prompt
+        const currentAssistant = useAssistant.getState().currentAssistant
+        const assistants = useAssistant.getState().assistants
+        const selectedModel = useModelProvider.getState().selectedModel
+        const selectedProvider = useModelProvider.getState().selectedProvider
 
-      router.navigate({
-        to: route.threadsDetail,
-        params: { threadId: currentThread.id },
-      })
-    }
-    return currentThread
-  }, [createThread, retrieveThread, router, setMessages, serviceHub])
+        // Get project metadata if projectId is provided
+        let projectMetadata:
+          | { id: string; name: string; updated_at: number }
+          | undefined
+        if (projectId) {
+          const project = await serviceHub.projects().getProjectById(projectId)
+          if (project) {
+            projectMetadata = {
+              id: project.id,
+              name: project.name,
+              updated_at: project.updated_at,
+            }
+          }
+        }
+
+        currentThread = await createThread(
+          {
+            id: selectedModel?.id ?? defaultModel(selectedProvider),
+            provider: selectedProvider,
+          },
+          isTemporaryMode ? 'Temporary Chat' : currentPrompt,
+          assistants.find((a) => a.id === currentAssistant?.id) ||
+            assistants[0],
+          projectMetadata,
+          isTemporaryMode // pass temporary flag
+        )
+
+        // Clear messages for temporary chat to ensure fresh start on reload
+        if (isTemporaryMode && currentThread?.id === TEMPORARY_CHAT_ID) {
+          setMessages(TEMPORARY_CHAT_ID, [])
+        }
+
+        // Set flag for temporary chat navigation
+        if (currentThread.id === TEMPORARY_CHAT_ID) {
+          sessionStorage.setItem('temp-chat-nav', 'true')
+        }
+
+        router.navigate({
+          to: route.threadsDetail,
+          params: { threadId: currentThread.id },
+        })
+      }
+      return currentThread
+    },
+    [createThread, retrieveThread, router, setMessages, serviceHub]
+  )
 
   const restartModel = useCallback(
     async (provider: ProviderObject, modelId: string) => {
@@ -271,6 +286,8 @@ export const useChat = () => {
       const selectedProvider = useModelProvider.getState().selectedProvider
       let activeProvider = getProviderByName(selectedProvider)
 
+      const streamEvents: StreamEvent[] = []
+
       resetTokenSpeed()
       if (!activeThread || !activeProvider) return
 
@@ -297,7 +314,9 @@ export const useChat = () => {
               updateAttachmentProcessing(img.name, 'processing')
             }
             // Upload image, get id/URL
-            const res = await serviceHub.uploads().ingestImage(activeThread.id, img)
+            const res = await serviceHub
+              .uploads()
+              .ingestImage(activeThread.id, img)
             processedAttachments.push({
               ...img,
               id: res.id,
@@ -313,7 +332,9 @@ export const useChat = () => {
               updateAttachmentProcessing(img.name, 'error')
             }
             const desc = err instanceof Error ? err.message : String(err)
-            toast.error('Failed to ingest image attachment', { description: desc })
+            toast.error('Failed to ingest image attachment', {
+              description: desc,
+            })
             return
           }
         }
@@ -394,6 +415,9 @@ export const useChat = () => {
       updateThreadTimestamp(activeThread.id)
       usePrompt.getState().setPrompt('')
       const selectedModel = useModelProvider.getState().selectedModel
+
+      const startTime = Date.now() // Start timer here
+
       try {
         if (selectedModel?.id) {
           updateLoadingModel(true)
@@ -410,10 +434,8 @@ export const useChat = () => {
         // Using addUserMessage to respect legacy code. Should be using the userContent above.
         if (troubleshooting) builder.addUserMessage(userContent)
 
-        let isCompleted = false
-
         // Filter tools based on model capabilities and available tools for this thread
-        let availableTools = selectedModel?.capabilities?.includes('tools')
+        const availableTools = selectedModel?.capabilities?.includes('tools')
           ? useAppState.getState().tools.filter((tool) => {
               const disabledTools = getDisabledToolsForThread(activeThread.id)
               return !disabledTools.includes(tool.name)
@@ -421,13 +443,21 @@ export const useChat = () => {
           : []
 
         // Check if proactive mode is enabled
-        const isProactiveMode = selectedModel?.capabilities?.includes('proactive') ?? false
+        const isProactiveMode =
+          selectedModel?.capabilities?.includes('proactive') ?? false
 
         // Proactive mode: Capture initial screenshot/snapshot before first LLM call
-        if (isProactiveMode && availableTools.length > 0 && !abortController.signal.aborted) {
-          console.log('Proactive mode: Capturing initial screenshots before LLM call')
+        if (
+          isProactiveMode &&
+          availableTools.length > 0 &&
+          !abortController.signal.aborted
+        ) {
+          console.log(
+            'Proactive mode: Capturing initial screenshots before LLM call'
+          )
           try {
-            const initialScreenshots = await captureProactiveScreenshots(abortController)
+            const initialScreenshots =
+              await captureProactiveScreenshots(abortController)
 
             // Add initial screenshots to builder
             for (const screenshot of initialScreenshots) {
@@ -441,131 +471,91 @@ export const useChat = () => {
           }
         }
 
-        let assistantLoopSteps = 0
+        // The agent logic is now self-contained within postMessageProcessing.
+        // We no longer need a `while` loop here.
 
-        while (
-          !isCompleted &&
-          !abortController.signal.aborted &&
-          activeProvider
-        ) {
-          const modelConfig = activeProvider.models.find(
-            (m) => m.id === selectedModel?.id
-          )
-          assistantLoopSteps += 1
+        if (abortController.signal.aborted || !activeProvider) return
 
-          const modelSettings = modelConfig?.settings
-            ? Object.fromEntries(
-                Object.entries(modelConfig.settings)
-                  .filter(
-                    ([key, value]) =>
-                      key !== 'ctx_len' &&
-                      key !== 'ngl' &&
-                      value.controller_props?.value !== undefined &&
-                      value.controller_props?.value !== null &&
-                      value.controller_props?.value !== ''
-                  )
-                  .map(([key, value]) => [key, value.controller_props?.value])
-              )
-            : undefined
+        const modelConfig = activeProvider.models.find(
+          (m) => m.id === selectedModel?.id
+        )
 
-          const completion = await sendCompletion(
-            activeThread,
-            activeProvider,
-            builder.getMessages(),
-            abortController,
-            availableTools,
-            currentAssistant?.parameters?.stream === false ? false : true,
-            {
-              ...modelSettings,
-              ...(currentAssistant?.parameters || {}),
-            } as unknown as Record<string, object>
-          )
+        const modelSettings = modelConfig?.settings
+          ? Object.fromEntries(
+              Object.entries(modelConfig.settings)
+                .filter(
+                  ([key, value]) =>
+                    key !== 'ctx_len' &&
+                    key !== 'ngl' &&
+                    value.controller_props?.value !== undefined &&
+                    value.controller_props?.value !== null &&
+                    value.controller_props?.value !== ''
+                )
+                .map(([key, value]) => [key, value.controller_props?.value])
+            )
+          : undefined
 
-          if (!completion) throw new Error('No completion received')
-          let accumulatedText = ''
-          const currentCall: ChatCompletionMessageToolCall | null = null
-          const toolCalls: ChatCompletionMessageToolCall[] = []
-          const timeToFirstToken = Date.now()
-          let tokenUsage: CompletionUsage | undefined = undefined
-          try {
-            if (isCompletionResponse(completion)) {
-              const message = completion.choices[0]?.message
-              accumulatedText = (message?.content as string) || ''
+        const completion = await sendCompletion(
+          activeThread,
+          activeProvider,
+          builder.getMessages(),
+          abortController,
+          availableTools,
+          currentAssistant?.parameters?.stream === false ? false : true,
+          {
+            ...modelSettings,
+            ...(currentAssistant?.parameters || {}),
+          } as unknown as Record<string, object>
+        )
 
-              // Handle reasoning field if there is one
-              const reasoning = extractReasoningFromMessage(message)
-              if (reasoning) {
-                accumulatedText =
-                  `<think>${reasoning}</think>` + accumulatedText
-              }
+        if (!completion) throw new Error('No completion received')
+        let accumulatedText = ''
+        const currentCall: ChatCompletionMessageToolCall | null = null
+        const toolCalls: ChatCompletionMessageToolCall[] = []
+        const timeToFirstToken = Date.now()
+        let tokenUsage: CompletionUsage | undefined = undefined
+        try {
+          if (isCompletionResponse(completion)) {
+            const message = completion.choices[0]?.message
+            accumulatedText = (message?.content as string) || ''
 
-              if (message?.tool_calls) {
-                toolCalls.push(...message.tool_calls)
-              }
-              if ('usage' in completion) {
-                tokenUsage = completion.usage
-              }
-            } else {
-              // High-throughput scheduler: batch UI updates on rAF (requestAnimationFrame)
-              let rafScheduled = false
-              let rafHandle: number | undefined
-              let pendingDeltaCount = 0
-              const reasoningProcessor = new ReasoningProcessor()
-              const scheduleFlush = () => {
-                if (rafScheduled || abortController.signal.aborted) return
-                rafScheduled = true
-                const doSchedule = (cb: () => void) => {
-                  if (typeof requestAnimationFrame !== 'undefined') {
-                    rafHandle = requestAnimationFrame(() => cb())
-                  } else {
-                    // Fallback for non-browser test environments
-                    const t = setTimeout(() => cb(), 0) as unknown as number
-                    rafHandle = t
-                  }
+            // Handle reasoning field if there is one
+            const reasoning = extractReasoningFromMessage(message)
+            if (reasoning) {
+              accumulatedText = `<think>${reasoning}</think>` + accumulatedText
+            }
+
+            if (message?.tool_calls) {
+              toolCalls.push(...message.tool_calls)
+            }
+            if ('usage' in completion) {
+              tokenUsage = completion.usage
+            }
+          } else {
+            // High-throughput scheduler: batch UI updates on rAF (requestAnimationFrame)
+            let rafScheduled = false
+            let rafHandle: number | undefined
+            let pendingDeltaCount = 0
+            const reasoningProcessor = new ReasoningProcessor()
+            const scheduleFlush = () => {
+              if (rafScheduled || abortController.signal.aborted) return
+              rafScheduled = true
+              const doSchedule = (cb: () => void) => {
+                if (typeof requestAnimationFrame !== 'undefined') {
+                  rafHandle = requestAnimationFrame(() => cb())
+                } else {
+                  // Fallback for non-browser test environments
+                  const t = setTimeout(() => cb(), 0) as unknown as number
+                  rafHandle = t
                 }
-                doSchedule(() => {
-                  // Check abort status before executing the scheduled callback
-                  if (abortController.signal.aborted) {
-                    rafScheduled = false
-                    return
-                  }
-
-                  const currentContent = newAssistantThreadContent(
-                    activeThread.id,
-                    accumulatedText,
-                    {
-                      tool_calls: toolCalls.map((e) => ({
-                        ...e,
-                        state: 'pending',
-                      })),
-                    }
-                  )
-                  updateStreamingContent(currentContent)
-                  if (tokenUsage) {
-                    setTokenSpeed(
-                      currentContent,
-                      tokenUsage.completion_tokens /
-                        Math.max((Date.now() - timeToFirstToken) / 1000, 1),
-                      tokenUsage.completion_tokens
-                    )
-                  } else if (pendingDeltaCount > 0) {
-                    updateTokenSpeed(currentContent, pendingDeltaCount)
-                  }
-                  pendingDeltaCount = 0
+              }
+              doSchedule(() => {
+                // Check abort status before executing the scheduled callback
+                if (abortController.signal.aborted) {
                   rafScheduled = false
-                })
-              }
-              const flushIfPending = () => {
-                if (!rafScheduled) return
-                if (
-                  typeof cancelAnimationFrame !== 'undefined' &&
-                  rafHandle !== undefined
-                ) {
-                  cancelAnimationFrame(rafHandle)
-                } else if (rafHandle !== undefined) {
-                  clearTimeout(rafHandle)
+                  return
                 }
-                // Do an immediate flush
+
                 const currentContent = newAssistantThreadContent(
                   activeThread.id,
                   accumulatedText,
@@ -574,6 +564,7 @@ export const useChat = () => {
                       ...e,
                       state: 'pending',
                     })),
+                    streamEvents: streamEvents,
                   }
                 )
                 updateStreamingContent(currentContent)
@@ -589,160 +580,232 @@ export const useChat = () => {
                 }
                 pendingDeltaCount = 0
                 rafScheduled = false
+              })
+            }
+            const flushIfPending = () => {
+              if (!rafScheduled) return
+              if (
+                typeof cancelAnimationFrame !== 'undefined' &&
+                rafHandle !== undefined
+              ) {
+                cancelAnimationFrame(rafHandle)
+              } else if (rafHandle !== undefined) {
+                clearTimeout(rafHandle)
               }
-              try {
-                for await (const part of completion) {
-                  // Check if aborted before processing each part
-                  if (abortController.signal.aborted) {
-                    break
-                  }
+              // Do an immediate flush
+              const currentContent = newAssistantThreadContent(
+                activeThread.id,
+                accumulatedText,
+                {
+                  tool_calls: toolCalls.map((e) => ({
+                    ...e,
+                    state: 'pending',
+                  })),
+                  streamEvents: streamEvents,
+                }
+              )
+              updateStreamingContent(currentContent)
+              if (tokenUsage) {
+                setTokenSpeed(
+                  currentContent,
+                  tokenUsage.completion_tokens /
+                    Math.max((Date.now() - timeToFirstToken) / 1000, 1),
+                  tokenUsage.completion_tokens
+                )
+              } else if (pendingDeltaCount > 0) {
+                updateTokenSpeed(currentContent, pendingDeltaCount)
+              }
+              pendingDeltaCount = 0
+              rafScheduled = false
+            }
+            try {
+              for await (const part of completion) {
+                // Check if aborted before processing each part
+                if (abortController.signal.aborted) {
+                  break
+                }
 
-                  // Handle prompt progress if available
-                  if ('prompt_progress' in part && part.prompt_progress) {
-                    // Force immediate state update to ensure we see intermediate values
-                    flushSync(() => {
-                      updatePromptProgress(part.prompt_progress)
+                // Handle prompt progress if available
+                if ('prompt_progress' in part && part.prompt_progress) {
+                  // Force immediate state update to ensure we see intermediate values
+                  flushSync(() => {
+                    updatePromptProgress(part.prompt_progress)
+                  })
+                  // Add a small delay to make progress visible
+                  await new Promise((resolve) => setTimeout(resolve, 100))
+                }
+
+                // Error message
+                if (!part.choices) {
+                  throw new Error(
+                    'message' in part
+                      ? (part.message as string)
+                      : (JSON.stringify(part) ?? '')
+                  )
+                }
+
+                if ('usage' in part && part.usage) {
+                  tokenUsage = part.usage
+                }
+                const deltaToolCalls = part.choices[0]?.delta?.tool_calls
+                if (deltaToolCalls) {
+                  const index = deltaToolCalls[0]?.index
+                  // Check if this chunk starts a brand new tool call
+                  const isNewToolCallStart =
+                    index !== undefined && toolCalls[index] === undefined
+
+                  extractToolCall(part, currentCall, toolCalls)
+
+                  if (isNewToolCallStart) {
+                    // Track tool call event only when it begins
+                    // toolCalls[index] is the newly created object due to extractToolCall
+                    streamEvents.push({
+                      timestamp: Date.now(),
+                      type: 'tool_call',
+                      data: { toolCall: toolCalls[index] },
                     })
-                    // Add a small delay to make progress visible
-                    await new Promise((resolve) => setTimeout(resolve, 100))
-                  }
-
-                  // Error message
-                  if (!part.choices) {
-                    throw new Error(
-                      'message' in part
-                        ? (part.message as string)
-                        : (JSON.stringify(part) ?? '')
-                    )
-                  }
-
-                  if ('usage' in part && part.usage) {
-                    tokenUsage = part.usage
-                  }
-
-                  if (part.choices[0]?.delta?.tool_calls) {
-                    extractToolCall(part, currentCall, toolCalls)
                     // Schedule a flush to reflect tool update
                     scheduleFlush()
                   }
-                  const deltaReasoning =
-                    reasoningProcessor.processReasoningChunk(part)
-                  if (deltaReasoning) {
-                    accumulatedText += deltaReasoning
-                    pendingDeltaCount += 1
-                    // Schedule flush for reasoning updates
-                    scheduleFlush()
-                  }
-                  const deltaContent = part.choices[0]?.delta?.content || ''
-                  if (deltaContent) {
-                    accumulatedText += deltaContent
-                    pendingDeltaCount += 1
-                    // Batch UI update on next animation frame
-                    scheduleFlush()
-                  }
                 }
-              } finally {
-                // Always clean up scheduled RAF when stream ends (either normally or via abort)
-                if (rafHandle !== undefined) {
-                  if (typeof cancelAnimationFrame !== 'undefined') {
-                    cancelAnimationFrame(rafHandle)
-                  } else {
-                    clearTimeout(rafHandle)
-                  }
-                  rafHandle = undefined
-                  rafScheduled = false
+                const deltaReasoning =
+                  reasoningProcessor.processReasoningChunk(part)
+                if (deltaReasoning) {
+                  // accumulatedText += deltaReasoning
+                  // Track reasoning event
+                  streamEvents.push({
+                    timestamp: Date.now(),
+                    type: 'reasoning_chunk',
+                    data: { content: deltaReasoning },
+                  })
+                  pendingDeltaCount += 1
+                  // Schedule flush for reasoning updates
+                  scheduleFlush()
                 }
-
-                // Only finalize and flush if not aborted
-                if (!abortController.signal.aborted) {
-                  // Finalize reasoning (close any open think tags)
-                  accumulatedText += reasoningProcessor.finalize()
-                  // Ensure any pending buffered content is rendered at the end
-                  flushIfPending()
+                const deltaContent = part.choices[0]?.delta?.content || ''
+                if (deltaContent) {
+                  accumulatedText += deltaContent
+                  pendingDeltaCount += 1
+                  // Batch UI update on next animation frame
+                  scheduleFlush()
                 }
               }
-            }
-          } catch (error) {
-            const errorMessage =
-              error && typeof error === 'object' && 'message' in error
-                ? error.message
-                : error
-            if (
-              typeof errorMessage === 'string' &&
-              errorMessage.includes(OUT_OF_CONTEXT_SIZE) &&
-              selectedModel
-            ) {
-              const method = await showIncreaseContextSizeModal()
-              if (method === 'ctx_len') {
-                /// Increase context size
-                activeProvider = await increaseModelContextSize(
-                  selectedModel.id,
-                  activeProvider
-                )
-                continue
-              } else if (method === 'context_shift' && selectedModel?.id) {
-                /// Enable context_shift
-                activeProvider = await toggleOnContextShifting(
-                  selectedModel?.id,
-                  activeProvider
-                )
-                continue
-              } else throw error
-            } else {
-              throw error
+            } finally {
+              // Always clean up scheduled RAF when stream ends (either normally or via abort)
+              if (rafHandle !== undefined) {
+                if (typeof cancelAnimationFrame !== 'undefined') {
+                  cancelAnimationFrame(rafHandle)
+                } else {
+                  clearTimeout(rafHandle)
+                }
+                rafHandle = undefined
+                rafScheduled = false
+              }
+
+              // Only finalize and flush if not aborted
+              if (!abortController.signal.aborted) {
+                // Finalize reasoning (close any open think tags)
+                // accumulatedText += reasoningProcessor.finalize()
+                // Ensure any pending buffered content is rendered at the end
+                flushIfPending()
+              }
             }
           }
-          // TODO: Remove this check when integrating new llama.cpp extension
+        } catch (error) {
+          const errorMessage =
+            error && typeof error === 'object' && 'message' in error
+              ? error.message
+              : error
           if (
-            accumulatedText.length === 0 &&
-            toolCalls.length === 0 &&
-            activeThread.model?.id &&
-            activeProvider?.provider === 'llamacpp'
+            typeof errorMessage === 'string' &&
+            errorMessage.includes(OUT_OF_CONTEXT_SIZE) &&
+            selectedModel
           ) {
-            await serviceHub
-              .models()
-              .stopModel(activeThread.model.id, 'llamacpp')
-            throw new Error('No response received from the model')
-          }
-
-          // Create a final content object for adding to the thread
-          const finalContent = newAssistantThreadContent(
-            activeThread.id,
-            accumulatedText,
-            {
-              tokenSpeed: useAppState.getState().tokenSpeed,
-              assistant: currentAssistant,
-            }
-          )
-
-          builder.addAssistantMessage(accumulatedText, undefined, toolCalls)
-
-          // Check if proactive mode is enabled for this model
-          const isProactiveMode = selectedModel?.capabilities?.includes('proactive') ?? false
-
-          const updatedMessage = await postMessageProcessing(
-            toolCalls,
-            builder,
-            finalContent,
-            abortController,
-            useToolApproval.getState().approvedTools,
-            allowAllMCPPermissions ? undefined : showApprovalModal,
-            allowAllMCPPermissions,
-            isProactiveMode
-          )
-          addMessage(updatedMessage ?? finalContent)
-          updateStreamingContent(emptyThreadContent)
-          updatePromptProgress(undefined)
-          updateThreadTimestamp(activeThread.id)
-
-          isCompleted = !toolCalls.length
-          // Do not create agent loop if there is no need for it
-          // Check if assistant loop steps are within limits
-          if (assistantLoopSteps >= (currentAssistant?.tool_steps ?? 20)) {
-            // Stop the assistant tool call if it exceeds the maximum steps
-            availableTools = []
+            const method = await showIncreaseContextSizeModal()
+            if (method === 'ctx_len') {
+              /// Increase context size
+              activeProvider = await increaseModelContextSize(
+                selectedModel.id,
+                activeProvider
+              )
+              // NOTE: This will exit and not retry. A more robust solution might re-call sendMessage.
+              // For this change, we keep the existing behavior.
+              return
+            } else if (method === 'context_shift' && selectedModel?.id) {
+              /// Enable context_shift
+              activeProvider = await toggleOnContextShifting(
+                selectedModel?.id,
+                activeProvider
+              )
+              // NOTE: See above comment about retry.
+              return
+            } else throw error
+          } else {
+            throw error
           }
         }
+        // TODO: Remove this check when integrating new llama.cpp extension
+        if (
+          accumulatedText.length === 0 &&
+          toolCalls.length === 0 &&
+          activeThread.model?.id &&
+          activeProvider?.provider === 'llamacpp'
+        ) {
+          await serviceHub.models().stopModel(activeThread.model.id, 'llamacpp')
+          throw new Error('No response received from the model')
+        }
+
+        const completionFinishTime = Date.now()
+        // Calculate the time taken for the initial completion (streaming or non-streaming)
+        const initialCompletionTime = completionFinishTime - startTime
+
+        const messageMetadata: Record<string, any> = {
+          tokenSpeed: useAppState.getState().tokenSpeed,
+          assistant: currentAssistant,
+          streamEvents, // Add chronological events
+        }
+
+        if (accumulatedText.includes('<think>') || toolCalls.length > 0) {
+          messageMetadata.totalThinkingTime = initialCompletionTime
+        }
+
+        // This is the message object that will be built upon by postMessageProcessing
+        const finalContent = newAssistantThreadContent(
+          activeThread.id,
+          accumulatedText,
+          messageMetadata
+        )
+
+        builder.addAssistantMessage(accumulatedText, undefined, toolCalls)
+        // All subsequent tool calls and follow-up completions will modify `finalContent`.
+        const updatedMessage = await postMessageProcessing(
+          toolCalls,
+          builder,
+          finalContent,
+          abortController,
+          useToolApproval.getState().approvedTools,
+          allowAllMCPPermissions ? undefined : showApprovalModal,
+          allowAllMCPPermissions,
+          activeThread,
+          activeProvider,
+          availableTools,
+          updateStreamingContent, // Pass the callback to update UI
+          currentAssistant?.tool_steps,
+          isProactiveMode
+        )
+
+        if (updatedMessage && updatedMessage.metadata) {
+          if (finalContent.metadata?.totalThinkingTime !== undefined) {
+            updatedMessage.metadata.totalThinkingTime =
+              finalContent.metadata.totalThinkingTime
+          }
+        }
+
+        // Add the single, final, composite message to the store.
+        addMessage(updatedMessage ?? finalContent)
+        updateStreamingContent(emptyThreadContent)
+        updatePromptProgress(undefined)
+        updateThreadTimestamp(activeThread.id)
       } catch (error) {
         if (!abortController.signal.aborted) {
           if (error && typeof error === 'object' && 'message' in error) {

web-app/src/lib/__tests__/completion.test.ts

@@ -9,7 +9,7 @@ import {
   normalizeTools,
   extractToolCall,
   postMessageProcessing,
-  captureProactiveScreenshots
+  captureProactiveScreenshots,
 } from '../completion'
 
 // Mock dependencies
@@ -87,10 +87,12 @@ vi.mock('@/hooks/useServiceHub', () => ({
     })),
     rag: vi.fn(() => ({
       getToolNames: vi.fn(() => Promise.resolve([])),
-      callTool: vi.fn(() => Promise.resolve({
-        content: [{ type: 'text', text: 'mock rag result' }],
-        error: '',
-      })),
+      callTool: vi.fn(() =>
+        Promise.resolve({
+          content: [{ type: 'text', text: 'mock rag result' }],
+          error: '',
+        })
+      ),
     })),
   })),
 }))
@@ -133,13 +135,15 @@ describe('completion.ts', () => {
       expect(result.type).toBe('text')
       expect(result.role).toBe('user')
       expect(result.thread_id).toBe('thread-123')
-      expect(result.content).toEqual([{
-        type: 'text',
-        text: {
-          value: 'Hello world',
-          annotations: [],
+      expect(result.content).toEqual([
+        {
+          type: 'text',
+          text: {
+            value: 'Hello world',
+            annotations: [],
+          },
         },
-      }])
+      ])
     })
 
     it('should handle empty text', () => {
@@ -147,13 +151,15 @@ describe('completion.ts', () => {
 
       expect(result.type).toBe('text')
       expect(result.role).toBe('user')
-      expect(result.content).toEqual([{
-        type: 'text',
-        text: {
-          value: '',
-          annotations: [],
+      expect(result.content).toEqual([
+        {
+          type: 'text',
+          text: {
+            value: '',
+            annotations: [],
+          },
         },
-      }])
+      ])
     })
   })
 
@@ -164,13 +170,15 @@ describe('completion.ts', () => {
       expect(result.type).toBe('text')
       expect(result.role).toBe('assistant')
       expect(result.thread_id).toBe('thread-123')
-      expect(result.content).toEqual([{
-        type: 'text',
-        text: {
-          value: 'AI response',
-          annotations: [],
+      expect(result.content).toEqual([
+        {
+          type: 'text',
+          text: {
+            value: 'AI response',
+            annotations: [],
+          },
         },
-      }])
+      ])
     })
   })
 
@@ -207,16 +215,20 @@ describe('completion.ts', () => {
   describe('extractToolCall', () => {
     it('should extract tool calls from message', () => {
       const message = {
-        choices: [{
-          delta: {
-            tool_calls: [{ 
-              id: 'call_1', 
-              type: 'function', 
-              index: 0,
-              function: { name: 'test', arguments: '{}' } 
-            }]
-          }
-        }]
+        choices: [
+          {
+            delta: {
+              tool_calls: [
+                {
+                  id: 'call_1',
+                  type: 'function',
+                  index: 0,
+                  function: { name: 'test', arguments: '{}' },
+                },
+              ],
+            },
+          },
+        ],
       }
       const calls = []
       const result = extractToolCall(message, null, calls)
@@ -226,9 +238,11 @@ describe('completion.ts', () => {
 
     it('should handle message without tool calls', () => {
       const message = {
-        choices: [{
-          delta: {}
-        }]
+        choices: [
+          {
+            delta: {},
+          },
+        ],
       }
       const calls = []
       const result = extractToolCall(message, null, calls)
@@ -245,23 +259,31 @@ describe('completion.ts', () => {
       const mockMcp = {
         getTools: mockGetTools,
         callToolWithCancellation: vi.fn(() => ({
-          promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
+          promise: Promise.resolve({
+            content: [{ type: 'text', text: 'result' }],
+            error: '',
+          }),
           cancel: vi.fn(),
-        }))
+        })),
       }
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => mockMcp,
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'browserbase_navigate', arguments: '{"url": "test.com"}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: {
+            name: 'browserbase_navigate',
+            arguments: '{"url": "test.com"}',
+          },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
@@ -284,30 +306,44 @@ describe('completion.ts', () => {
     it('should detect browserbase tools', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
       const mockCallTool = vi.fn(() => ({
-        promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
+        promise: Promise.resolve({
+          content: [{ type: 'text', text: 'result' }],
+          error: '',
+        }),
         cancel: vi.fn(),
       }))
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: () => Promise.resolve([]),
-          callToolWithCancellation: mockCallTool
+          callToolWithCancellation: mockCallTool,
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'browserbase_screenshot', arguments: '{}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: { name: 'browserbase_screenshot', arguments: '{}' },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
 
-      await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
+      await postMessageProcessing(
+        calls,
+        builder,
+        message,
+        abortController,
+        {},
+        undefined,
+        false,
+        true
+      )
 
       expect(mockCallTool).toHaveBeenCalled()
     })
@@ -315,30 +351,47 @@ describe('completion.ts', () => {
     it('should detect multi_browserbase tools', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
       const mockCallTool = vi.fn(() => ({
-        promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
+        promise: Promise.resolve({
+          content: [{ type: 'text', text: 'result' }],
+          error: '',
+        }),
         cancel: vi.fn(),
       }))
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: () => Promise.resolve([]),
-          callToolWithCancellation: mockCallTool
+          callToolWithCancellation: mockCallTool,
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'multi_browserbase_stagehand_navigate', arguments: '{}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: {
+            name: 'multi_browserbase_stagehand_navigate',
+            arguments: '{}',
+          },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
 
-      await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
+      await postMessageProcessing(
+        calls,
+        builder,
+        message,
+        abortController,
+        {},
+        undefined,
+        false,
+        true
+      )
 
       expect(mockCallTool).toHaveBeenCalled()
     })
@@ -350,26 +403,40 @@ describe('completion.ts', () => {
         mcp: () => ({
           getTools: mockGetTools,
           callToolWithCancellation: vi.fn(() => ({
-            promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
+            promise: Promise.resolve({
+              content: [{ type: 'text', text: 'result' }],
+              error: '',
+            }),
             cancel: vi.fn(),
-          }))
+          })),
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'fetch_url', arguments: '{"url": "test.com"}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: { name: 'fetch_url', arguments: '{"url": "test.com"}' },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
 
-      await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
+      await postMessageProcessing(
+        calls,
+        builder,
+        message,
+        abortController,
+        {},
+        undefined,
+        false,
+        true
+      )
 
       // Proactive screenshots should not be called for non-browser tools
       expect(mockGetTools).not.toHaveBeenCalled()
@@ -380,7 +447,9 @@ describe('completion.ts', () => {
     it('should capture screenshot and snapshot when available', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
       const mockScreenshotResult = {
-        content: [{ type: 'image', data: 'base64screenshot', mimeType: 'image/png' }],
+        content: [
+          { type: 'image', data: 'base64screenshot', mimeType: 'image/png' },
+        ],
         error: '',
       }
       const mockSnapshotResult = {
@@ -388,11 +457,14 @@ describe('completion.ts', () => {
         error: '',
       }
 
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'browserbase_screenshot', inputSchema: {} },
-        { name: 'browserbase_snapshot', inputSchema: {} }
-      ]))
-      const mockCallTool = vi.fn()
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([
+          { name: 'browserbase_screenshot', inputSchema: {} },
+          { name: 'browserbase_snapshot', inputSchema: {} },
+        ])
+      )
+      const mockCallTool = vi
+        .fn()
         .mockReturnValueOnce({
           promise: Promise.resolve(mockScreenshotResult),
           cancel: vi.fn(),
@@ -405,8 +477,8 @@ describe('completion.ts', () => {
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
-        })
+          callToolWithCancellation: mockCallTool,
+        }),
       } as any)
 
       const abortController = new AbortController()
@@ -420,15 +492,15 @@ describe('completion.ts', () => {
 
     it('should handle missing screenshot tool gracefully', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'some_other_tool', inputSchema: {} }
-      ]))
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([{ name: 'some_other_tool', inputSchema: {} }])
+      )
 
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: vi.fn()
-        })
+          callToolWithCancellation: vi.fn(),
+        }),
       } as any)
 
       const abortController = new AbortController()
@@ -439,9 +511,9 @@ describe('completion.ts', () => {
 
     it('should handle screenshot capture errors gracefully', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'browserbase_screenshot', inputSchema: {} }
-      ]))
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([{ name: 'browserbase_screenshot', inputSchema: {} }])
+      )
       const mockCallTool = vi.fn(() => ({
         promise: Promise.reject(new Error('Screenshot failed')),
         cancel: vi.fn(),
@@ -450,8 +522,8 @@ describe('completion.ts', () => {
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
-        })
+          callToolWithCancellation: mockCallTool,
+        }),
       } as any)
 
       const abortController = new AbortController()
@@ -463,22 +535,30 @@ describe('completion.ts', () => {
 
     it('should respect abort controller', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'browserbase_screenshot', inputSchema: {} }
-      ]))
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([{ name: 'browserbase_screenshot', inputSchema: {} }])
+      )
       const mockCallTool = vi.fn(() => ({
-        promise: new Promise((resolve) => setTimeout(() => resolve({
-          content: [{ type: 'image', data: 'base64', mimeType: 'image/png' }],
-          error: '',
-        }), 100)),
+        promise: new Promise((resolve) =>
+          setTimeout(
+            () =>
+              resolve({
+                content: [
+                  { type: 'image', data: 'base64', mimeType: 'image/png' },
+                ],
+                error: '',
+              }),
+            100
+          )
+        ),
         cancel: vi.fn(),
       }))
 
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
-        })
+          callToolWithCancellation: mockCallTool,
+        }),
       } as any)
 
       const abortController = new AbortController()
@@ -500,12 +580,15 @@ describe('completion.ts', () => {
             role: 'tool',
             content: [
               { type: 'text', text: 'Tool result' },
-              { type: 'image_url', image_url: { url: 'data:image/png;base64,old' } }
+              {
+                type: 'image_url',
+                image_url: { url: 'data:image/png;base64,old' },
+              },
             ],
-            tool_call_id: 'old_call'
+            tool_call_id: 'old_call',
           },
           { role: 'assistant', content: 'Response' },
-        ]
+        ],
       }
 
       expect(builder.messages).toHaveLength(3)
@@ -517,13 +600,19 @@ describe('completion.ts', () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
 
       const mockScreenshotResult = {
-        content: [{ type: 'image', data: 'proactive_screenshot', mimeType: 'image/png' }],
+        content: [
+          {
+            type: 'image',
+            data: 'proactive_screenshot',
+            mimeType: 'image/png',
+          },
+        ],
         error: '',
       }
 
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'browserbase_screenshot', inputSchema: {} }
-      ]))
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([{ name: 'browserbase_screenshot', inputSchema: {} }])
+      )
 
       let callCount = 0
       const mockCallTool = vi.fn(() => {
@@ -549,19 +638,24 @@ describe('completion.ts', () => {
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
+          callToolWithCancellation: mockCallTool,
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'browserbase_navigate', arguments: '{"url": "test.com"}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: {
+            name: 'browserbase_navigate',
+            arguments: '{"url": "test.com"}',
+          },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
@@ -574,7 +668,12 @@ describe('completion.ts', () => {
         {},
         undefined,
         false,
-        true
+        undefined, // thread
+        undefined, // provider
+        [], // tools
+        undefined, // updateStreamingUI
+        undefined, // maxToolSteps
+        true // isProactiveMode - Correctly set to true
       )
 
       // Should have called: 1) browser tool, 2) getTools, 3) proactive screenshot
@@ -586,9 +685,9 @@ describe('completion.ts', () => {
     it('should not trigger proactive screenshots when mode is disabled', async () => {
       const { getServiceHub } = await import('@/hooks/useServiceHub')
 
-      const mockGetTools = vi.fn(() => Promise.resolve([
-        { name: 'browserbase_screenshot', inputSchema: {} }
-      ]))
+      const mockGetTools = vi.fn(() =>
+        Promise.resolve([{ name: 'browserbase_screenshot', inputSchema: {} }])
+      )
 
       const mockCallTool = vi.fn(() => ({
         promise: Promise.resolve({
@@ -601,19 +700,21 @@ describe('completion.ts', () => {
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
+          callToolWithCancellation: mockCallTool,
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'browserbase_navigate', arguments: '{}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: { name: 'browserbase_navigate', arguments: '{}' },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
@@ -626,7 +727,12 @@ describe('completion.ts', () => {
         {},
         undefined,
         false,
-        false
+        undefined, // thread
+        undefined, // provider
+        [], // tools
+        undefined, // updateStreamingUI
+        undefined, // maxToolSteps
+        false // isProactiveMode - Correctly set to false
       )
 
       expect(mockCallTool).toHaveBeenCalledTimes(1)
@@ -648,19 +754,21 @@ describe('completion.ts', () => {
       vi.mocked(getServiceHub).mockReturnValue({
         mcp: () => ({
           getTools: mockGetTools,
-          callToolWithCancellation: mockCallTool
+          callToolWithCancellation: mockCallTool,
         }),
-        rag: () => ({ getToolNames: () => Promise.resolve([]) })
+        rag: () => ({ getToolNames: () => Promise.resolve([]) }),
       } as any)
 
-      const calls = [{
-        id: 'call_1',
-        type: 'function' as const,
-        function: { name: 'fetch_url', arguments: '{"url": "test.com"}' }
-      }]
+      const calls = [
+        {
+          id: 'call_1',
+          type: 'function' as const,
+          function: { name: 'fetch_url', arguments: '{"url": "test.com"}' },
+        },
+      ]
       const builder = {
         addToolMessage: vi.fn(),
-        getMessages: vi.fn(() => [])
+        getMessages: vi.fn(() => []),
       } as any
       const message = { thread_id: 'test-thread', metadata: {} } as any
       const abortController = new AbortController()
@@ -673,7 +781,12 @@ describe('completion.ts', () => {
         {},
         undefined,
         false,
-        true
+        undefined, // thread
+        undefined, // provider
+        [], // tools
+        undefined, // updateStreamingUI
+        undefined, // maxToolSteps
+        true // isProactiveMode - Still set to true, but the non-browser tool should skip the proactive step
       )
 
       expect(mockCallTool).toHaveBeenCalledTimes(1)

web-app/src/lib/completion.ts

@@ -41,6 +41,7 @@ import { useAppState } from '@/hooks/useAppState'
 import { injectFilesIntoPrompt } from './fileMetadata'
 import { Attachment } from '@/types/attachment'
 import { ModelCapabilities } from '@/types/models'
+import { ReasoningProcessor } from '@/utils/reasoning'
 
 export type ChatCompletionResponse =
   | chatCompletion
@@ -48,6 +49,12 @@ export type ChatCompletionResponse =
   | StreamCompletionResponse
   | CompletionResponse
 
+type ToolCallEntry = {
+  tool: object
+  response: any
+  state: 'pending' | 'ready'
+}
+
 /**
  * @fileoverview Helper functions for creating thread content.
  * These functions are used to create thread content objects
@@ -73,11 +80,14 @@ export const newUserThreadContent = (
       name: doc.name,
       type: doc.fileType,
       size: typeof doc.size === 'number' ? doc.size : undefined,
-      chunkCount: typeof doc.chunkCount === 'number' ? doc.chunkCount : undefined,
+      chunkCount:
+        typeof doc.chunkCount === 'number' ? doc.chunkCount : undefined,
     }))
 
   const textWithFiles =
-    docMetadata.length > 0 ? injectFilesIntoPrompt(content, docMetadata) : content
+    docMetadata.length > 0
+      ? injectFilesIntoPrompt(content, docMetadata)
+      : content
 
   const contentParts = [
     {
@@ -238,10 +248,8 @@ export const sendCompletion = async (
   const providerModelConfig = provider.models?.find(
     (model) => model.id === thread.model?.id || model.model === thread.model?.id
   )
-  const effectiveCapabilities = Array.isArray(
-    providerModelConfig?.capabilities
-  )
-    ? providerModelConfig?.capabilities ?? []
+  const effectiveCapabilities = Array.isArray(providerModelConfig?.capabilities)
+    ? (providerModelConfig?.capabilities ?? [])
     : getModelCapabilities(provider.provider, thread.model.id)
   const modelSupportsTools = effectiveCapabilities.includes(
     ModelCapabilities.TOOLS
@@ -254,7 +262,10 @@ export const sendCompletion = async (
       PlatformFeatures[PlatformFeature.ATTACHMENTS] &&
       modelSupportsTools
     ) {
-      const ragTools = await getServiceHub().rag().getTools().catch(() => [])
+      const ragTools = await getServiceHub()
+        .rag()
+        .getTools()
+        .catch(() => [])
       if (Array.isArray(ragTools) && ragTools.length) {
         usableTools = [...tools, ...ragTools]
       }
@@ -395,7 +406,6 @@ export const extractToolCall = (
   }
   return calls
 }
-
 /**
  * Helper function to check if a tool call is a browser MCP tool
  * @param toolName - The name of the tool
@@ -519,7 +529,13 @@ const filterOldProactiveScreenshots = (builder: CompletionMessagesBuilder) => {
  * @param approvedTools
  * @param showModal
  * @param allowAllMCPPermissions
+ * @param thread
+ * @param provider
+ * @param tools
+ * @param updateStreamingUI
+ * @param maxToolSteps
  * @param isProactiveMode
+ * @param currentStepCount - Internal counter for recursive calls (do not set manually)
  */
 export const postMessageProcessing = async (
   calls: ChatCompletionMessageToolCall[],
@@ -533,10 +549,30 @@ export const postMessageProcessing = async (
     toolParameters?: object
   ) => Promise<boolean>,
   allowAllMCPPermissions: boolean = false,
-  isProactiveMode: boolean = false
-) => {
+  thread?: Thread,
+  provider?: ModelProvider,
+  tools: MCPTool[] = [],
+  updateStreamingUI?: (content: ThreadMessage) => void,
+  maxToolSteps: number = 20,
+  isProactiveMode: boolean = false,
+  currentStepCount: number = 0
+): Promise<ThreadMessage> => {
+  // Initialize/get the current total thinking time from metadata
+  // This value is passed from sendMessage (initial completion time) or previous recursive call
+  let currentTotalTime = (message.metadata?.totalThinkingTime as number) ?? 0
+
   // Handle completed tool calls
-  if (calls.length) {
+  if (calls.length > 0) {
+    // Check limit BEFORE processing
+    if (currentStepCount >= maxToolSteps) {
+      console.warn(
+        `Reached maximum tool steps (${maxToolSteps}), stopping chain to prevent infinite loop`
+      )
+      return message
+    }
+
+    const nextStepCount = currentStepCount + 1
+
     // Fetch RAG tool names from RAG service
     let ragToolNames = new Set<string>()
     try {
@@ -546,43 +582,42 @@ export const postMessageProcessing = async (
       console.error('Failed to load RAG tool names:', e)
     }
     const ragFeatureAvailable =
-      useAttachments.getState().enabled && PlatformFeatures[PlatformFeature.ATTACHMENTS]
+      useAttachments.getState().enabled &&
+      PlatformFeatures[PlatformFeature.ATTACHMENTS]
+
+    const currentToolCalls =
+      message.metadata?.tool_calls && Array.isArray(message.metadata.tool_calls)
+        ? [...message.metadata.tool_calls]
+        : []
+
     for (const toolCall of calls) {
       if (abortController.signal.aborted) break
       const toolId = ulid()
-      const toolCallsMetadata =
-        message.metadata?.tool_calls &&
-        Array.isArray(message.metadata?.tool_calls)
-          ? message.metadata?.tool_calls
-          : []
+
+      const toolCallEntry: ToolCallEntry = {
+        tool: {
+          ...(toolCall as object),
+          id: toolId,
+        },
+        response: undefined,
+        state: 'pending' as 'pending' | 'ready',
+      }
+      currentToolCalls.push(toolCallEntry)
+
       message.metadata = {
         ...(message.metadata ?? {}),
-        tool_calls: [
-          ...toolCallsMetadata,
-          {
-            tool: {
-              ...(toolCall as object),
-              id: toolId,
-            },
-            response: undefined,
-            state: 'pending',
-          },
-        ],
+        tool_calls: currentToolCalls,
+        totalThinkingTime: currentTotalTime,
       }
+      if (updateStreamingUI) updateStreamingUI({ ...message }) // Show pending call
 
       // Check if tool is approved or show modal for approval
       let toolParameters = {}
       if (toolCall.function.arguments.length) {
         try {
-          console.log('Raw tool arguments:', toolCall.function.arguments)
           toolParameters = JSON.parse(toolCall.function.arguments)
-          console.log('Parsed tool parameters:', toolParameters)
         } catch (error) {
           console.error('Failed to parse tool arguments:', error)
-          console.error(
-            'Raw arguments that failed:',
-            toolCall.function.arguments
-          )
         }
       }
 
@@ -591,7 +626,6 @@ export const postMessageProcessing = async (
       const isRagTool = ragToolNames.has(toolName)
       const isBrowserTool = isBrowserMCPTool(toolName)
 
-      // Auto-approve RAG tools (local/safe operations), require permission for MCP tools
       const approved = isRagTool
         ? true
         : allowAllMCPPermissions ||
@@ -604,10 +638,16 @@ export const postMessageProcessing = async (
               )
             : true)
 
+      const toolExecutionStartTime = Date.now()
+
       const { promise, cancel } = isRagTool
         ? ragFeatureAvailable
           ? {
-              promise: getServiceHub().rag().callTool({ toolName, arguments: toolArgs, threadId: message.thread_id }),
+              promise: getServiceHub().rag().callTool({
+                toolName,
+                arguments: toolArgs,
+                threadId: message.thread_id,
+              }),
               cancel: async () => {},
             }
           : {
@@ -630,18 +670,15 @@ export const postMessageProcessing = async (
       useAppState.getState().setCancelToolCall(cancel)
 
       let result = approved
-        ? await promise.catch((e) => {
-            console.error('Tool call failed:', e)
-            return {
-              content: [
-                {
-                  type: 'text',
-                  text: `Error calling tool ${toolCall.function.name}: ${e.message ?? e}`,
-                },
-              ],
-              error: String(e?.message ?? e ?? 'Tool call failed'),
-            }
-          })
+        ? await promise.catch((e) => ({
+            content: [
+              {
+                type: 'text',
+                text: `Error calling tool ${toolCall.function.name}: ${e.message ?? e}`,
+              },
+            ],
+            error: String(e?.message ?? e ?? 'Tool call failed'),
+          }))
         : {
             content: [
               {
@@ -652,43 +689,50 @@ export const postMessageProcessing = async (
             error: 'disallowed',
           }
 
+      const toolExecutionTime = Date.now() - toolExecutionStartTime
+
       if (typeof result === 'string') {
         result = {
-          content: [
-            {
-              type: 'text',
-              text: result,
-            },
-          ],
+          content: [{ type: 'text', text: result }],
           error: '',
         }
       }
 
+      currentTotalTime += toolExecutionTime
+
+      // Update the entry in the metadata array
+      toolCallEntry.response = result
+      toolCallEntry.state = 'ready'
       message.metadata = {
         ...(message.metadata ?? {}),
-        tool_calls: [
-          ...toolCallsMetadata,
-          {
-            tool: {
-              ...toolCall,
-              id: toolId,
-            },
-            response: result,
-            state: 'ready',
-          },
-        ],
+        totalThinkingTime: currentTotalTime,
+      }
+      if (updateStreamingUI) updateStreamingUI({ ...message }) // Show result
+
+      const streamEvents = (message.metadata?.streamEvents || []) as any[]
+      streamEvents.push({
+        timestamp: Date.now(),
+        type: 'tool_output',
+        data: { result: result },
+      })
+      message.metadata = {
+        ...(message.metadata ?? {}),
+        streamEvents: streamEvents,
       }
       builder.addToolMessage(result as ToolResult, toolCall.id)
 
       // Proactive mode: Capture screenshot/snapshot after browser tool execution
       if (isProactiveMode && isBrowserTool && !abortController.signal.aborted) {
-        console.log('Proactive mode: Capturing screenshots after browser tool call')
+        console.log(
+          'Proactive mode: Capturing screenshots after browser tool call'
+        )
 
         // Filter out old screenshots before adding new ones
         filterOldProactiveScreenshots(builder)
 
         // Capture new screenshots
-        const proactiveScreenshots = await captureProactiveScreenshots(abortController)
+        const proactiveScreenshots =
+          await captureProactiveScreenshots(abortController)
 
         // Add proactive screenshots to builder
         for (const screenshot of proactiveScreenshots) {
@@ -702,6 +746,164 @@ export const postMessageProcessing = async (
 
       // update message metadata
     }
-    return message
+
+    // Process follow-up completion if conditions are met
+    if (thread && provider && !abortController.signal.aborted) {
+      try {
+        const messagesWithToolResults = builder.getMessages()
+
+        const followUpStartTime = Date.now()
+
+        const followUpCompletion = await sendCompletion(
+          thread,
+          provider,
+          messagesWithToolResults,
+          abortController,
+          tools,
+          true,
+          {}
+        )
+
+        let streamFinishTime = Date.now()
+
+        if (followUpCompletion) {
+          let followUpText = ''
+          const newToolCalls: ChatCompletionMessageToolCall[] = []
+          const streamEvents = (message.metadata?.streamEvents || []) as any[]
+          const textContent = message.content.find(
+            (c) => c.type === ContentType.Text
+          )
+
+          if (isCompletionResponse(followUpCompletion)) {
+            // Handle non-streaming response
+            const choice = followUpCompletion.choices[0]
+            const content = choice?.message?.content
+            if (content) followUpText = content as string
+            if (choice?.message?.tool_calls) {
+              newToolCalls.push(...choice.message.tool_calls)
+            }
+            if (textContent?.text) textContent.text.value += followUpText
+            if (updateStreamingUI) updateStreamingUI({ ...message })
+            streamFinishTime = Date.now()
+          } else {
+            // Handle streaming response
+            const reasoningProcessor = new ReasoningProcessor()
+            for await (const chunk of followUpCompletion) {
+              if (abortController.signal.aborted) break
+
+              const deltaReasoning =
+                reasoningProcessor.processReasoningChunk(chunk)
+              const deltaContent = chunk.choices[0]?.delta?.content || ''
+
+              if (textContent?.text) {
+                if (deltaContent) {
+                  textContent.text.value += deltaContent
+                  followUpText += deltaContent
+                }
+              }
+
+              if (deltaReasoning) {
+                streamEvents.push({
+                  timestamp: Date.now(),
+                  type: 'reasoning_chunk',
+                  data: { content: deltaReasoning },
+                })
+              }
+
+              const initialToolCallCount = newToolCalls.length
+
+              if (chunk.choices[0]?.delta?.tool_calls) {
+                extractToolCall(chunk, null, newToolCalls)
+                if (newToolCalls.length > initialToolCallCount) {
+                  // The new tool call is the last element added
+                  streamEvents.push({
+                    timestamp: Date.now(),
+                    type: 'tool_call',
+                    data: { toolCall: newToolCalls[newToolCalls.length - 1] },
+                  })
+                }
+              }
+
+              // Ensure the metadata is updated before calling updateStreamingUI
+              message.metadata = {
+                ...(message.metadata ?? {}),
+                streamEvents: streamEvents,
+                totalThinkingTime:
+                  currentTotalTime + (Date.now() - followUpStartTime), // Optimistic update
+              }
+
+              if (updateStreamingUI) {
+                // Create a new object reference for the content array
+                // This forces the memoized component to detect the change in the mutated text
+                const uiMessage: ThreadMessage = {
+                  ...message,
+                  content: message.content.map((c) => ({ ...c })),
+                }
+                updateStreamingUI(uiMessage)
+              }
+            }
+            streamFinishTime = Date.now()
+            if (textContent?.text && updateStreamingUI) {
+              // Final UI update after streaming completes
+              const uiMessage: ThreadMessage = {
+                ...message,
+                content: message.content.map((c) => ({ ...c })),
+              }
+              updateStreamingUI(uiMessage)
+            }
+          }
+
+          const followUpTotalTime = streamFinishTime - followUpStartTime
+          currentTotalTime += followUpTotalTime //
+          message.metadata = {
+            ...(message.metadata ?? {}),
+            totalThinkingTime: currentTotalTime,
+          }
+
+          // Recursively process new tool calls if any
+          if (newToolCalls.length > 0) {
+            builder.addAssistantMessage(followUpText, undefined, newToolCalls)
+            // Recursive call continues accumulation on the same message object
+            await postMessageProcessing(
+              newToolCalls,
+              builder,
+              message,
+              abortController,
+              approvedTools,
+              showModal,
+              allowAllMCPPermissions,
+              thread,
+              provider,
+              tools,
+              updateStreamingUI,
+              maxToolSteps,
+              isProactiveMode,
+              nextStepCount, // Pass the incremented step count
+            )
+          }
+        }
+      } catch (error) {
+        console.error(
+          'Failed to get follow-up completion after tool execution:',
+          String(error)
+        )
+        // Optionally add error to message metadata for UI display
+        const streamEvents = (message.metadata?.streamEvents || []) as any[]
+        streamEvents.push({
+          timestamp: Date.now(),
+          type: 'error',
+          data: {
+            message: 'Follow-up completion failed',
+            error: String(error),
+          },
+        })
+        message.metadata = {
+          ...(message.metadata ?? {}),
+          streamEvents: streamEvents,
+        }
+      }
+    }
   }
+
+  return message
 }

web-app/src/locales/en/chat.json

@@ -8,5 +8,13 @@
   },
   "sendMessage": "Send Message",
   "newConversation": "New Conversation",
-  "clearHistory": "Clear History"
-}
+  "clearHistory": "Clear History",
+  "thought_and_tool_call": "Thought and called tools",
+  "tool_called": "Called tools",
+  "calling_tool": "Calling a tool",
+  "thinking": "Thinking",
+  "thought": "Thought",
+  "for": "for",
+  "seconds": "seconds"
+}
+

web-app/src/locales/en/common.json

@@ -235,7 +235,21 @@
       "title": "Edit Message"
     },
     "messageMetadata": {
-      "title": "Message Metadata"
+      "title": "Message Metadata",
+      "model": "Model",
+      "name": "Name",
+      "id": "ID",
+      "createdAt": "Created At",
+      "toolSteps": "Tool Steps",
+      "temperature": "Temperature",
+      "topK": "Top K",
+      "topP": "Top P",
+      "description": "Description",
+      "instructions": "Instructions",
+      "tokenSpeed": "Token Speed",
+      "tokenCount": "Token Count",
+      "lastUpdate": "Last Update",
+      "noMessageMetadataAvailable": "No Message Metadata Available"
     }
   },
   "projects": {