Simplify streaming message

Author: reakaleek

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/ChatMessage.tsx

@@ -63,13 +63,6 @@ interface ChatMessageProps {
     onRetry?: () => void
 }
 
-const getAccumulatedContent = (messages: AskAiEvent[]) => {
-    return messages
-        .filter((m) => m.type === EventTypes.MESSAGE_CHUNK)
-        .map((m) => m.content)
-        .join('')
-}
-
 const splitContentAndReferences = (
     content: string
 ): { mainContent: string; referencesJson: string | null } => {
@@ -279,9 +272,9 @@ export const ChatMessage = ({
         )
     }
 
-    const content =
-        streamingContent ||
-        (events.length > 0 ? getAccumulatedContent(events) : message.content)
+    // Use streamingContent during streaming, otherwise use message.content from store
+    // message.content is updated atomically with status when CONVERSATION_END arrives
+    const content = streamingContent || message.content
 
     const hasError = message.status === 'error' || !!error
 

src/Elastic.Documentation.Site/Assets/web-components/SearchOrAskAi/AskAi/StreamingAiMessage.tsx

@@ -6,7 +6,6 @@ import {
     useConversationId,
 } from './chat.store'
 import { useAskAi } from './useAskAi'
-import * as React from 'react'
 import { useEffect, useRef } from 'react'
 
 interface StreamingAiMessageProps {
@@ -45,7 +44,7 @@ export const StreamingAiMessage = ({
                     'error'
                 )
             } else if (event.type === EventTypes.CONVERSATION_END) {
-                updateAiMessage(message.id, contentRef.current, 'complete')
+                updateAiMessage(message.id, message.content || contentRef.current, 'complete')
             }
         },
         onError: (error) => {
@@ -85,15 +84,16 @@ export const StreamingAiMessage = ({
         markMessageAsSent,
     ])
 
+    // Always use contentRef.current if it has content (regardless of status)
+    // This way we don't need to save to message.content and can just use streamingContent
+    const streamingContentToPass =
+        isLast && contentRef.current ? contentRef.current : undefined
+
     return (
         <ChatMessage
             message={message}
             events={isLast ? events : []}
-            streamingContent={
-                isLast && message.status === 'streaming'
-                    ? contentRef.current
-                    : undefined
-            }
+            streamingContent={streamingContentToPass}
         />
     )
 }