Refactorings and re-render optimizations

Author: reakaleek

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

@@ -4,6 +4,13 @@ import { EuiButton, useEuiTheme } from '@elastic/eui'
 import { css } from '@emotion/react'
 import * as React from 'react'
 
+const buttonStyles = css`
+    border: none;
+    & > span {
+        justify-content: flex-start;
+    }
+`
+
 export interface AskAiSuggestion {
     question: string
 }
@@ -16,15 +23,14 @@ export const AskAiSuggestions = (props: Props) => {
     const { submitQuestion } = useChatActions()
     const { setModalMode } = useModalActions()
     const { euiTheme } = useEuiTheme()
-    const buttonCss = css`
-        border: none;
-        & > span {
-            justify-content: flex-start;
-        }
+    
+    const dynamicButtonStyles = css`
+        ${buttonStyles}
         svg {
             color: ${euiTheme.colors.textSubdued};
         }
     `
+    
     return (
         <ul>
             {Array.from(props.suggestions).map((suggestion) => (
@@ -34,7 +40,7 @@ export const AskAiSuggestions = (props: Props) => {
                         color="text"
                         fullWidth
                         size="s"
-                        css={buttonCss}
+                        css={dynamicButtonStyles}
                         onClick={() => {
                             submitQuestion(suggestion.question)
                             setModalMode('askAi')

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

@@ -17,6 +17,29 @@ import { css } from '@emotion/react'
 import * as React from 'react'
 import { useCallback, useEffect, useRef } from 'react'
 
+const containerStyles = css`
+    height: 100%;
+    max-height: 70vh;
+    overflow: hidden;
+`
+
+const scrollContainerStyles = css`
+    position: relative;
+    overflow: hidden;
+`
+
+const scrollableStyles = css`
+    height: 100%;
+    overflow-y: auto;
+    scrollbar-gutter: stable;
+    padding: 1rem;
+`
+
+const messagesStyles = css`
+    max-width: 800px;
+    margin: 0 auto;
+`
+
 // Small helper for scroll behavior
 const scrollToBottom = (container: HTMLDivElement | null) => {
     if (!container) return
@@ -44,28 +67,9 @@ export const Chat = () => {
     const scrollRef = useRef<HTMLDivElement>(null)
     const lastMessageStatusRef = useRef<string | null>(null)
 
-    const containerStyles = css`
-        height: 100%;
-        max-height: 70vh;
-        overflow: hidden;
-    `
-
-    const scrollContainerStyles = css`
-        position: relative;
-        overflow: hidden;
-    `
-
-    const scrollableStyles = css`
-        height: 100%;
-        overflow-y: auto;
-        scrollbar-gutter: stable;
+    const dynamicScrollableStyles = css`
+        ${scrollableStyles}
         ${useEuiOverflowScroll('y', true)}
-        padding: 1rem;
-    `
-
-    const messagesStyles = css`
-        max-width: 800px;
-        margin: 0 auto;
     `
 
     const handleSubmit = useCallback(
@@ -123,7 +127,7 @@ export const Chat = () => {
 
             {/* Messages */}
             <EuiFlexItem grow={true} css={scrollContainerStyles}>
-                <div ref={scrollRef} css={scrollableStyles}>
+                <div ref={scrollRef} css={dynamicScrollableStyles}>
                     {messages.length === 0 ? (
                         <EuiEmptyPrompt
                             iconType="logoElastic"

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

@@ -39,24 +39,23 @@ const createMarkedInstance = () => {
     return new Marked({ renderer })
 }
 
-const markedInstance = createMarkedInstance() // Created once globally
+const markedInstance = createMarkedInstance()
 
 interface ChatMessageProps {
     message: ChatMessageType
     llmMessages?: LlmGatewayMessage[]
+    streamingContent?: string
     error?: Error | null
     onRetry?: () => void
 }
 
-// Helper function to accumulate AI message content
 const getAccumulatedContent = (messages: LlmGatewayMessage[]) => {
     return messages
-        .filter((m) => m.type === 'ai_message_chunk') // Only accumulate chunks, not the final message
+        .filter((m) => m.type === 'ai_message_chunk')
         .map((m) => m.data.content)
         .join('')
 }
 
-// Derived state helper for readability
 const getMessageState = (message: ChatMessageType) => ({
     isUser: message.type === 'user',
     isLoading: message.status === 'streaming',
@@ -127,6 +126,7 @@ const ActionBar = ({
 export const ChatMessage = ({
     message,
     llmMessages = [],
+    streamingContent,
     error,
     onRetry,
 }: ChatMessageProps) => {
@@ -166,27 +166,22 @@ export const ChatMessage = ({
         )
     }
 
-    // AI message
-    const content =
-        llmMessages.length > 0
-            ? getAccumulatedContent(llmMessages)
-            : message.content
+    const content = streamingContent 
+        || (llmMessages.length > 0 ? getAccumulatedContent(llmMessages) : message.content)
 
     const hasError = message.status === 'error' || !!error
 
-    // Memoize the parsed HTML to avoid re-parsing on every render
     const parsed = useMemo(() => {
         const html = markedInstance.parse(content) as string
         return DOMPurify.sanitize(html)
     }, [content])
+    
     const ref = React.useRef<HTMLDivElement>(null)
 
-    // Initialize copy buttons after DOM is updated
     useEffect(() => {
         if (isComplete && ref.current) {
             const timer = setTimeout(() => {
                 try {
-                    // Use the modified initCopyButton with scoped element
                     initCopyButton(
                         '.highlight pre',
                         ref.current!,
@@ -198,7 +193,7 @@ export const ChatMessage = ({
             }, 100)
             return () => clearTimeout(timer)
         }
-    }, [isComplete]) // Only depend on isComplete, not parsed
+    }, [isComplete])
 
     return (
         <EuiFlexGroup

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

@@ -27,7 +27,6 @@ export const StreamingAiMessage = ({
         onMessage: (llmMessage) => {
             if (llmMessage.type === 'ai_message_chunk') {
                 contentRef.current += llmMessage.data.content
-                updateAiMessage(message.id, contentRef.current, 'streaming')
             } else if (llmMessage.type === 'agent_end') {
                 updateAiMessage(message.id, contentRef.current, 'complete')
             }
@@ -41,8 +40,6 @@ export const StreamingAiMessage = ({
         },
     })
 
-    // Send question when this is the last message and status is streaming
-    // Use store-level tracking so it persists across remounts
     useEffect(() => {
         if (
             isLast &&
@@ -68,6 +65,7 @@ export const StreamingAiMessage = ({
         <ChatMessage
             message={message}
             llmMessages={isLast ? llmMessages : []}
+            streamingContent={isLast && message.status === 'streaming' ? contentRef.current : undefined}
         />
     )
 }

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

@@ -7,6 +7,10 @@ import { EuiFieldSearch, EuiSpacer, EuiButton, useEuiTheme } from '@elastic/eui'
 import { css } from '@emotion/react'
 import * as React from 'react'
 
+const askAiButtonStyles = css`
+    font-weight: bold;
+`
+
 export const Search = () => {
     const searchTerm = useSearchTerm()
     const { setSearchTerm } = useSearchActions()
@@ -53,15 +57,10 @@ export const Search = () => {
 }
 
 const AskAiButton = ({ term, onAsk }: { term: string; onAsk: () => void }) => {
-    const { euiTheme } = useEuiTheme()
     return (
         <EuiButton iconType="newChat" fullWidth onClick={onAsk}>
             Ask AI about{' '}
-            <span
-                css={css`
-                    font-weight: ${euiTheme.font.weight.bold};
-                `}
-            >
+            <span css={askAiButtonStyles}>
                 "{term}"
             </span>
         </EuiButton>

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

@@ -14,7 +14,7 @@ import { css } from '@emotion/react'
 import * as React from 'react'
 import { useMemo } from 'react'
 
-export const SearchOrAskAiModal = () => {
+export const SearchOrAskAiModal = React.memo(() => {
     const modalMode = useModalMode()
     const { setModalMode } = useModalActions()
 
@@ -24,18 +24,13 @@ export const SearchOrAskAiModal = () => {
                 id: 'search',
                 name: 'Search',
                 prepend: <EuiIcon type="search" />,
-                content: <Search />,
+                content: <></>,
             },
             {
                 id: 'askAi',
                 name: 'Ask AI',
                 prepend: <EuiIcon type="sparkles" />,
-                content: (
-                    <>
-                        <EuiSpacer size="m" />
-                        <Chat />
-                    </>
-                ),
+                content: <></>,
             },
         ],
         []
@@ -50,10 +45,20 @@ export const SearchOrAskAiModal = () => {
                 selectedTab={selectedTab}
                 onTabClick={(tab) => setModalMode(tab.id as 'search' | 'askAi')}
             />
+            
+            {modalMode === 'search' ? (
+                <Search />
+            ) : (
+                <>
+                    <EuiSpacer size="m" />
+                    <Chat />
+                </>
+            )}
+            
             <ModalFooter />
         </>
     )
-}
+})
 
 const ModalFooter = () => {
     return (