diff --git a/web-app/src/containers/ThinkingBlock.tsx b/web-app/src/containers/ThinkingBlock.tsx index d59b41415..3107bff20 100644 --- a/web-app/src/containers/ThinkingBlock.tsx +++ b/web-app/src/containers/ThinkingBlock.tsx @@ -156,7 +156,7 @@ const ThinkingBlock = ({
diff --git a/web-app/src/containers/ThreadContent.tsx b/web-app/src/containers/ThreadContent.tsx index 047c86006..6811c8ab5 100644 --- a/web-app/src/containers/ThreadContent.tsx +++ b/web-app/src/containers/ThreadContent.tsx @@ -140,44 +140,98 @@ export const ThreadContent = memo( return { files: [], cleanPrompt: text } }, [text, item.role]) - const { reasoningSegment, textSegment } = useMemo(() => { - let reasoningSegment = undefined - let textSegment = text - - // Check for completed think tag format - console.log(textSegment) + const { + finalOutputText, + streamedReasoningText, + isReasoningActiveLoading, + hasReasoningSteps, + } = useMemo(() => { const thinkStartTag = '' const thinkEndTag = '' + let currentFinalText = '' + let currentReasoning = '' + let hasSteps = false - const firstThinkIndex = text.indexOf(thinkStartTag) - const lastThinkEndIndex = text.lastIndexOf(thinkEndTag) + const firstThinkStart = text.indexOf(thinkStartTag) + const lastThinkStart = text.lastIndexOf(thinkStartTag) + const lastThinkEnd = text.lastIndexOf(thinkEndTag) - if (firstThinkIndex !== -1 && lastThinkEndIndex > firstThinkIndex) { - // If multiple ... blocks exist sequentially, we capture the entire span - // from the start of the first tag to the end of the last tag. - const splitIndex = lastThinkEndIndex + thinkEndTag.length + // Check if there's an unclosed tag + const hasOpenThink = lastThinkStart > lastThinkEnd - reasoningSegment = text.slice(firstThinkIndex, splitIndex) - textSegment = text.slice(splitIndex).trim() + if (firstThinkStart === -1) { + // No tags at all - everything is final output + currentFinalText = text + } else if (hasOpenThink && isStreamingThisThread) { + // CASE 1: There's an open tag during streaming + // Everything from FIRST onward is reasoning + hasSteps = true - return { reasoningSegment, textSegment } - } - // If streaming, and we see the opening tag, the entire message is reasoningSegment - const hasThinkTagStart = - text.includes(thinkStartTag) && !text.includes(thinkEndTag) + // Text before first is final output + currentFinalText = text.substring(0, firstThinkStart) - if (hasThinkTagStart) { - reasoningSegment = text - textSegment = '' - return { reasoningSegment, textSegment } + // Everything from first onward is reasoning + const reasoningText = text.substring(firstThinkStart) + + // Extract content from all blocks (both closed and open) + const reasoningRegex = /([\s\S]*?)(?:<\/think>|$)/g + const matches = [...reasoningText.matchAll(reasoningRegex)] + const reasoningParts = matches.map((match) => cleanReasoning(match[1])) + currentReasoning = reasoningParts.join('\n\n') + } else { + // CASE 2: All tags are closed + // Extract reasoning from inside tags, everything else is final output + hasSteps = true + + const reasoningRegex = /[\s\S]*?<\/think>/g + const matches = [...text.matchAll(reasoningRegex)] + + let lastIndex = 0 + + // Build final output from text between/outside blocks + for (const match of matches) { + currentFinalText += text.substring(lastIndex, match.index) + lastIndex = match.index + match[0].length + } + + // Add remaining text after last + currentFinalText += text.substring(lastIndex) + + // Extract reasoning content + const reasoningParts = matches.map((match) => { + const content = match[0].replace(/|<\/think>/g, '') + return cleanReasoning(content) + }) + currentReasoning = reasoningParts.join('\n\n') } - // Default: No reasoning found, or it's a message composed entirely of final text. - return { reasoningSegment: undefined, textSegment: text } - }, [text]) + // Check for tool calls + const isToolCallsPresent = !!( + item.metadata && + 'tool_calls' in item.metadata && + Array.isArray(item.metadata.tool_calls) && + item.metadata.tool_calls.length > 0 + ) - // Check if reasoning segment is actually present (i.e., non-empty string) - const hasReasoning = !!reasoningSegment + hasSteps = hasSteps || isToolCallsPresent + + // Loading if streaming and no final output yet + const loading = + isStreamingThisThread && currentFinalText.trim().length === 0 + + return { + finalOutputText: currentFinalText.trim(), + streamedReasoningText: currentReasoning, + isReasoningActiveLoading: loading, + hasReasoningSteps: hasSteps, + } + }, [item.content, isStreamingThisThread, item.metadata, text]) + + 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) @@ -249,12 +303,6 @@ 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 @@ -273,6 +321,8 @@ export const ThreadContent = memo( 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 = '' @@ -366,10 +416,10 @@ export const ThreadContent = memo( }) }) } - } else { - console.debug('Fallback mode!!!!') - // FALLBACK PATH: No streamEvents - use old paragraph-splitting logic - const rawReasoningContent = cleanReasoning(reasoningSegment || '') + } else if (isMessageFinalized) { + // FALLBACK PATH: No streamEvents - use split text for content construction + + const rawReasoningContent = streamedReasoningText || '' const reasoningParagraphs = rawReasoningContent ? rawReasoningContent .split(/\n\s*\n/) @@ -442,9 +492,9 @@ export const ThreadContent = memo( const totalTime = item.metadata?.totalThinkingTime as number | undefined const lastStepType = steps[steps.length - 1]?.type - if (!isStreamingThisThread && (hasReasoning || isToolCalls)) { + if (!isStreamingThisThread && hasReasoningSteps) { const endsInToolOutputWithoutFinalText = - lastStepType === 'tool_output' && textSegment.length === 0 + lastStepType === 'tool_output' && finalOutputText.length === 0 if (!endsInToolOutputWithoutFinalText) { steps.push({ @@ -458,22 +508,34 @@ export const ThreadContent = memo( return steps }, [ item, - reasoningSegment, isStreamingThisThread, - hasReasoning, - isToolCalls, - textSegment, + hasReasoningSteps, + finalOutputText, + streamedReasoningText, ]) // END: Constructing allSteps - // Determine if reasoning phase is actively loading - // Loading is true only if streaming is happening AND we haven't started outputting final text yet. - const isReasoningActiveLoading = - isStreamingThisThread && textSegment.length === 0 + // ==================================================================== + // FIX: Determine which text prop to pass to ThinkingBlock + // If we have streamEvents, rely on 'steps' and pass an empty text buffer. + const streamingTextBuffer = useMemo(() => { + const streamEvents = item.metadata?.streamEvents - // Determine if we should show the thinking block (has reasoning OR tool calls OR currently loading reasoning) + // 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 '' + } + + // Otherwise, rely on the raw text buffer for rendering (used during initial stream fallback) + return streamedReasoningText + }, [item.metadata?.streamEvents, streamedReasoningText]) // Use the object reference for dependency array + // ==================================================================== + + // Determine if we should show the thinking block const shouldShowThinkingBlock = - hasReasoning || isToolCalls || isReasoningActiveLoading + hasReasoningSteps || isToolCalls || isReasoningActiveLoading return ( @@ -614,19 +676,25 @@ export const ThreadContent = memo( )} - + {!isReasoningActiveLoading && finalOutputText.length > 0 && ( + + )} {!isToolCalls && (
diff --git a/web-app/src/lib/completion.ts b/web-app/src/lib/completion.ts index 3b71333c0..fabde868d 100644 --- a/web-app/src/lib/completion.ts +++ b/web-app/src/lib/completion.ts @@ -407,9 +407,6 @@ export const extractToolCall = ( return calls } -// Keep track of total tool steps to prevent infinite loops -let toolStepCounter = 0 - /** * Helper function to check if a tool call is a browser MCP tool * @param toolName - The name of the tool @@ -533,6 +530,12 @@ const filterOldProactiveScreenshots = (builder: CompletionMessagesBuilder) => { * @param approvedTools * @param showModal * @param allowAllMCPPermissions + * @param thread + * @param provider + * @param tools + * @param updateStreamingUI + * @param maxToolSteps + * @param currentStepCount - Internal counter for recursive calls (do not set manually) * @param isProactiveMode */ export const postMessageProcessing = async ( @@ -552,16 +555,20 @@ export const postMessageProcessing = async ( tools: MCPTool[] = [], updateStreamingUI?: (content: ThreadMessage) => void, maxToolSteps: number = 20, + currentStepCount: number = 0, isProactiveMode: boolean = false ): Promise => { - // Reset counter at the start of a new message processing chain - if (toolStepCounter === 0) { - toolStepCounter = 0 - } - // Handle completed tool calls if (calls.length > 0) { - toolStepCounter++ + // 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() @@ -687,6 +694,7 @@ export const postMessageProcessing = async ( toolCallEntry.response = result toolCallEntry.state = 'ready' if (updateStreamingUI) updateStreamingUI({ ...message }) // Show result + const streamEvents = (message.metadata?.streamEvents || []) as any[] streamEvents.push({ timestamp: Date.now(), @@ -701,13 +709,16 @@ export const postMessageProcessing = async ( // 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) { @@ -722,12 +733,8 @@ export const postMessageProcessing = async ( // update message metadata } - if ( - thread && - provider && - !abortController.signal.aborted && - toolStepCounter < maxToolSteps - ) { + // Process follow-up completion if conditions are met + if (thread && provider && !abortController.signal.aborted) { try { const messagesWithToolResults = builder.getMessages() @@ -750,6 +757,7 @@ export const postMessageProcessing = async ( ) if (isCompletionResponse(followUpCompletion)) { + // Handle non-streaming response const choice = followUpCompletion.choices[0] const content = choice?.message?.content if (content) followUpText = content as string @@ -759,6 +767,7 @@ export const postMessageProcessing = async ( if (textContent?.text) textContent.text.value += followUpText if (updateStreamingUI) updateStreamingUI({ ...message }) } else { + // Handle streaming response const reasoningProcessor = new ReasoningProcessor() for await (const chunk of followUpCompletion) { if (abortController.signal.aborted) break @@ -772,9 +781,9 @@ export const postMessageProcessing = async ( if (deltaContent) { textContent.text.value += deltaContent followUpText += deltaContent - console.log(`delta content from followup:\n${deltaContent}`) } } + if (deltaReasoning) { streamEvents.push({ timestamp: Date.now(), @@ -782,6 +791,7 @@ export const postMessageProcessing = async ( data: { content: deltaReasoning }, }) } + const initialToolCallCount = newToolCalls.length if (chunk.choices[0]?.delta?.tool_calls) { @@ -795,6 +805,7 @@ export const postMessageProcessing = async ( }) } } + // Ensure the metadata is updated before calling updateStreamingUI message.metadata = { ...(message.metadata ?? {}), @@ -802,26 +813,27 @@ export const postMessageProcessing = async ( } if (updateStreamingUI) { - // FIX: Create a new object reference for the content array + // 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 })), // Shallow copy array and its parts + content: message.content.map((c) => ({ ...c })), } updateStreamingUI(uiMessage) } } + if (textContent?.text && updateStreamingUI) { - // FIX: Create a new object reference for the content array - // This forces the memoized component to detect the change in the mutated text + // Final UI update after streaming completes const uiMessage: ThreadMessage = { ...message, - content: message.content.map((c) => ({ ...c })), // Shallow copy array and its parts + content: message.content.map((c) => ({ ...c })), } updateStreamingUI(uiMessage) } } + // Recursively process new tool calls if any if (newToolCalls.length > 0) { builder.addAssistantMessage(followUpText, undefined, newToolCalls) await postMessageProcessing( @@ -837,6 +849,7 @@ export const postMessageProcessing = async ( tools, updateStreamingUI, maxToolSteps, + nextStepCount, // Pass the incremented step count isProactiveMode ) } @@ -846,11 +859,23 @@ export const postMessageProcessing = async ( '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, + } } } } - // Reset counter when the chain is fully resolved - toolStepCounter = 0 return message }