SearchOrAskAi Error Handling Improvements

src/Elastic.Documentation.Site/Assets/eui-icons-cache.ts

@@ -26,6 +26,8 @@ import { icon as EuiIconThumbUp } from '@elastic/eui/es/components/icon/assets/t
 import { icon as EuiIconTrash } from '@elastic/eui/es/components/icon/assets/trash'
 import { icon as EuiIconUser } from '@elastic/eui/es/components/icon/assets/user'
 import { icon as EuiIconWrench } from '@elastic/eui/es/components/icon/assets/wrench'
+import { icon as EuiIconArrowStart } from '@elastic/eui/es/components/icon/assets/arrowStart'
+import { icon as EuiIconArrowEnd } from '@elastic/eui/es/components/icon/assets/arrowEnd'
 import { appendIconComponentCache } from '@elastic/eui/es/components/icon/icon'
 
 appendIconComponentCache({
@@ -57,4 +59,6 @@ appendIconComponentCache({
     copy: EuiIconCopy,
     play: EuiIconPlay,
     sortUp: EuiIconSortUp,
+    arrowStart: EuiIconArrowStart,
+    arrowEnd: EuiIconArrowEnd,
 })

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

@@ -37,7 +37,7 @@ const ALL_SUGGESTIONS: AskAiSuggestion[] = [
     },
 ]
 
-export const AskAiSuggestions = () => {
+export const AskAiSuggestions = ({ disabled }: { disabled?: boolean }) => {
     const { submitQuestion } = useChatActions()
     const { setModalMode } = useModalActions()
     const { euiTheme } = useEuiTheme()
@@ -64,9 +64,12 @@ export const AskAiSuggestions = () => {
                         fullWidth
                         size="s"
                         onClick={() => {
-                            submitQuestion(suggestion.question)
-                            setModalMode('askAi')
+                            if (!disabled) {
+                                submitQuestion(suggestion.question)
+                                setModalMode('askAi')
+                            }
                         }}
+                        disabled={disabled}
                     >
                         {suggestion.question}
                     </EuiButton>

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

@@ -1,4 +1,6 @@
 /** @jsxImportSource @emotion/react */
+import { SearchOrAskAiErrorCallout } from '../SearchOrAskAiErrorCallout'
+import { useIsCooldownActive } from '../hooks/useIsCooldownActive'
 import { AiProviderSelector } from './AiProviderSelector'
 import { AskAiSuggestions } from './AskAiSuggestions'
 import { ChatMessageList } from './ChatMessageList'
@@ -15,7 +17,6 @@ import {
     EuiTitle,
 } from '@elastic/eui'
 import { css } from '@emotion/react'
-import * as React from 'react'
 import { useCallback, useEffect, useRef, useState } from 'react'
 
 const containerStyles = css`
@@ -48,11 +49,22 @@ const scrollToBottom = (container: HTMLDivElement | null) => {
 }
 
 // Header shown when a conversation exists
-const NewConversationHeader = ({ onClick }: { onClick: () => void }) => (
+const NewConversationHeader = ({
+    onClick,
+    disabled,
+}: {
+    onClick: () => void
+    disabled?: boolean
+}) => (
     <EuiFlexItem grow={false}>
         <EuiFlexGroup justifyContent="flexEnd" alignItems="center">
             <EuiFlexItem grow={false}>
-                <EuiButtonEmpty size="xs" onClick={onClick} iconType="refresh">
+                <EuiButtonEmpty
+                    size="xs"
+                    onClick={onClick}
+                    iconType="refresh"
+                    disabled={disabled}
+                >
                     New conversation
                 </EuiButtonEmpty>
             </EuiFlexItem>
@@ -63,21 +75,56 @@ const NewConversationHeader = ({ onClick }: { onClick: () => void }) => (
 
 export const Chat = () => {
     const messages = useChatMessages()
-    const { submitQuestion, clearChat } = useChatActions()
+    const { submitQuestion, clearChat, clearNon429Errors } = useChatActions()
+    const isCooldownActive = useIsCooldownActive()
     const inputRef = useRef<HTMLInputElement>(null)
     const scrollRef = useRef<HTMLDivElement>(null)
     const lastMessageStatusRef = useRef<string | null>(null)
+    const abortFunctionRef = useRef<(() => void) | null>(null)
     const [inputValue, setInputValue] = useState('')
+    const [hasClearedError, setHasClearedError] = useState(false)
 
     const dynamicScrollableStyles = css`
         ${scrollableStyles}
         ${useEuiOverflowScroll('y', true)}
     `
 
+    // Check if there's an active streaming query
+    const isStreaming =
+        messages.length > 0 &&
+        messages[messages.length - 1].type === 'ai' &&
+        messages[messages.length - 1].status === 'streaming'
+
+    // Handle abort function from StreamingAiMessage
+    const handleAbortReady = useCallback((abort: () => void) => {
+        abortFunctionRef.current = abort
+    }, [])
+
+    // Clear abort function when streaming ends
+    useEffect(() => {
+        if (!isStreaming) {
+            abortFunctionRef.current = null
+        }
+    }, [isStreaming])
+
+    useEffect(() => {
+        if (inputValue && !hasClearedError) {
+            clearNon429Errors()
+            setHasClearedError(true)
+        } else if (!inputValue) {
+            setHasClearedError(false)
+        }
+    }, [inputValue, clearNon429Errors, hasClearedError])
+
     const handleSubmit = useCallback(
         (question: string) => {
             if (!question.trim()) return
 
+            // Prevent submission during countdown
+            if (isCooldownActive) {
+                return
+            }
+
             submitQuestion(question.trim())
 
             if (inputRef.current) {
@@ -88,9 +135,19 @@ export const Chat = () => {
             // Scroll to bottom after new message
             setTimeout(() => scrollToBottom(scrollRef.current), 100)
         },
-        [submitQuestion]
+        [submitQuestion, isCooldownActive]
     )
 
+    const handleButtonClick = useCallback(() => {
+        if (isStreaming && abortFunctionRef.current) {
+            // Interrupt current query
+            abortFunctionRef.current()
+            abortFunctionRef.current = null
+        } else if (inputRef.current) {
+            handleSubmit(inputRef.current.value)
+        }
+    }, [isStreaming, handleSubmit])
+
     // Refocus input when AI answer transitions to complete
     useEffect(() => {
         if (messages.length > 0) {
@@ -126,43 +183,59 @@ export const Chat = () => {
             <EuiSpacer size="m" />
 
             {messages.length > 0 && (
-                <NewConversationHeader onClick={clearChat} />
+                <NewConversationHeader
+                    onClick={clearChat}
+                    disabled={isCooldownActive}
+                />
             )}
 
             <EuiFlexItem grow={true} css={scrollContainerStyles}>
                 <div ref={scrollRef} css={dynamicScrollableStyles}>
                     {messages.length === 0 ? (
-                        <EuiEmptyPrompt
-                            iconType="logoElastic"
-                            title={
-                                <h2>Hi! I'm the Elastic Docs AI Assistant</h2>
-                            }
-                            body={
-                                <>
-                                    <p>
-                                        I can help answer your questions about
-                                        Elastic documentation. <br />
-                                        Ask me anything about Elasticsearch,
-                                        Kibana, Observability, Security, and
-                                        more.
-                                    </p>
-                                    <EuiSpacer size="m" />
-                                    <AiProviderSelector />
-                                </>
-                            }
-                            footer={
-                                <>
-                                    <EuiTitle size="xxs">
-                                        <h3>Try asking me:</h3>
-                                    </EuiTitle>
-                                    <EuiSpacer size="s" />
-                                    <AskAiSuggestions />
-                                </>
-                            }
-                        />
+                        <>
+                            <EuiEmptyPrompt
+                                iconType="logoElastic"
+                                title={
+                                    <h2>
+                                        Hi! I'm the Elastic Docs AI Assistant
+                                    </h2>
+                                }
+                                body={
+                                    <>
+                                        <p>
+                                            I can help answer your questions
+                                            about Elastic documentation. <br />
+                                            Ask me anything about Elasticsearch,
+                                            Kibana, Observability, Security, and
+                                            more.
+                                        </p>
+                                        <EuiSpacer size="m" />
+                                        <AiProviderSelector />
+                                    </>
+                                }
+                                footer={
+                                    <>
+                                        <EuiTitle size="xxs">
+                                            <h3>Try asking me:</h3>
+                                        </EuiTitle>
+                                        <EuiSpacer size="s" />
+                                        <AskAiSuggestions
+                                            disabled={isCooldownActive}
+                                        />
+                                    </>
+                                }
+                            />
+                            {/* Show error callout when there's a cooldown, even on initial page */}
+                            <div css={messagesStyles}>
+                                <SearchOrAskAiErrorCallout error={null} />
+                            </div>
+                        </>
                     ) : (
                         <div css={messagesStyles}>
-                            <ChatMessageList messages={messages} />
+                            <ChatMessageList
+                                messages={messages}
+                                onAbortReady={handleAbortReady}
+                            />
                         </div>
                     )}
                 </div>
@@ -187,9 +260,12 @@ export const Chat = () => {
                                 handleSubmit(e.currentTarget.value)
                             }
                         }}
+                        disabled={isCooldownActive}
                     />
                     <EuiButtonIcon
-                        aria-label="Send message"
+                        aria-label={
+                            isStreaming ? 'Interrupt query' : 'Send message'
+                        }
                         css={css`
                             position: absolute;
                             right: 8px;
@@ -198,13 +274,12 @@ export const Chat = () => {
                             border-radius: 9999px;
                         `}
                         color="primary"
-                        iconType="sortUp"
-                        display={inputValue.trim() ? 'fill' : 'base'}
-                        onClick={() => {
-                            if (inputRef.current) {
-                                handleSubmit(inputRef.current.value)
-                            }
-                        }}
+                        iconType={isStreaming ? 'cross' : 'sortUp'}
+                        display={
+                            inputValue.trim() || isStreaming ? 'fill' : 'base'
+                        }
+                        onClick={handleButtonClick}
+                        disabled={isCooldownActive}
                     ></EuiButtonIcon>
                 </div>
                 <EuiSpacer size="m" />