Fix flaky test: Replace setTimeout with useLayoutEffect for scroll behavior

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

@@ -54,6 +54,17 @@ describe('Chat Component', () => {
     beforeEach(() => {
         jest.clearAllMocks()
         resetStores()
+
+        // Mock scrollTo to prevent flaky test failures in CI
+        // The handleSubmit function schedules a setTimeout that calls scrollTo
+        // In test environments, scrollTo may not be available on DOM elements
+        HTMLElement.prototype.scrollTo = jest.fn()
+    })
+
+    afterEach(() => {
+        // Clean up the mock
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        delete (HTMLElement.prototype as any).scrollTo
     })
 
     describe('Empty state', () => {
@@ -258,6 +269,134 @@ describe('Chat Component', () => {
                 expect(input.value).toBe('')
             })
         })
+
+        it('should scroll to user message when question is submitted', async () => {
+            // Arrange
+            jest.useFakeTimers()
+            try {
+                const user = userEvent.setup({ delay: null })
+                const question = 'What is Kibana?'
+                const scrollToSpy = jest.fn()
+
+                // Set up existing messages so scroll container exists
+                chatStore.setState({
+                    chatMessages: [
+                        {
+                            id: '1',
+                            type: 'user',
+                            content: 'What is Elasticsearch?',
+                            conversationId: 'thread-1',
+                            timestamp: Date.now(),
+                        },
+                        {
+                            id: '2',
+                            type: 'ai',
+                            content:
+                                'Elasticsearch is a distributed search engine...',
+                            conversationId: 'thread-1',
+                            timestamp: Date.now(),
+                            status: 'complete',
+                        },
+                    ],
+                    conversationId: 'thread-1',
+                })
+
+                // Act
+                const { container } = renderWithQueryClient(<Chat />)
+
+                // Wait for component to render and find scroll container
+                await waitFor(() => {
+                    expect(
+                        screen.getByText('What is Elasticsearch?')
+                    ).toBeInTheDocument()
+                })
+
+                // Find the scroll container (the div that will have messages)
+                const scrollContainer =
+                    (container
+                        .querySelector('[data-message-type]')
+                        ?.closest('div[style*="overflow"]') as HTMLElement) ||
+                    (Array.from(container.querySelectorAll('div')).find(
+                        (el) => {
+                            const style = window.getComputedStyle(el)
+                            return style.overflowY === 'auto'
+                        }
+                    ) as HTMLElement)
+
+                if (scrollContainer) {
+                    scrollContainer.scrollTo = scrollToSpy
+                    scrollContainer.scrollTop = 0
+                    // Mock getBoundingClientRect for scroll calculations
+                    scrollContainer.getBoundingClientRect = jest.fn(() => ({
+                        top: 0,
+                        left: 0,
+                        right: 100,
+                        bottom: 500,
+                        width: 100,
+                        height: 500,
+                        x: 0,
+                        y: 0,
+                        toJSON: jest.fn(),
+                    })) as jest.MockedFunction<() => DOMRect>
+                }
+
+                const input = screen.getByPlaceholderText(
+                    /Ask the Elastic Docs AI Assistant/i
+                )
+                await user.type(input, question)
+                await user.keyboard('{Enter}')
+
+                // Wait for message to be added
+                await waitFor(() => {
+                    const messages = chatStore.getState().chatMessages
+                    expect(messages.length).toBeGreaterThan(2)
+                    expect(messages[messages.length - 2].type).toBe('user')
+                    expect(messages[messages.length - 2].content).toBe(question)
+                })
+
+                // Wait for the new user message to be rendered in the DOM
+                await waitFor(() => {
+                    expect(screen.getByText(question)).toBeInTheDocument()
+                })
+
+                // Mock getBoundingClientRect for the new user message
+                const newUserMessage = container.querySelector(
+                    '[data-message-type="user"]:last-of-type'
+                ) as HTMLElement
+                if (newUserMessage) {
+                    newUserMessage.getBoundingClientRect = jest.fn(() => ({
+                        top: 100,
+                        left: 0,
+                        right: 100,
+                        bottom: 200,
+                        width: 100,
+                        height: 100,
+                        x: 0,
+                        y: 100,
+                        toJSON: jest.fn(),
+                    })) as jest.MockedFunction<() => DOMRect>
+                }
+
+                // Advance timers to trigger requestAnimationFrame
+                jest.advanceTimersByTime(100)
+                await Promise.resolve() // Let React process the state update
+
+                // Run all pending requestAnimationFrame callbacks
+                jest.runAllTimers()
+
+                // Verify scrollTo was called
+                if (scrollContainer) {
+                    expect(scrollToSpy).toHaveBeenCalledWith(
+                        expect.objectContaining({
+                            top: expect.any(Number),
+                            behavior: 'smooth',
+                        })
+                    )
+                }
+            } finally {
+                jest.useRealTimers()
+            }
+        })
     })
 
     describe('Close modal', () => {

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

@@ -29,7 +29,14 @@ import {
     useEuiTheme,
 } from '@elastic/eui'
 import { css } from '@emotion/react'
-import { RefObject, useCallback, useEffect, useRef, useState } from 'react'
+import {
+    RefObject,
+    useCallback,
+    useEffect,
+    useLayoutEffect,
+    useRef,
+    useState,
+} from 'react'
 
 export const Chat = () => {
     const { euiTheme } = useEuiTheme()
@@ -86,7 +93,6 @@ export const Chat = () => {
 
             <ChatInputArea
                 inputRef={inputRef}
-                scrollRef={scrollRef}
                 onMetaSemicolon={handleMetaSemicolon}
                 onStateChange={setScrollAreaProps}
             />
@@ -200,6 +206,25 @@ const ChatScrollArea = ({
     const messages = useChatMessages()
     const { euiTheme } = useEuiTheme()
     const spacerHeight = useSpacerHeight(scrollRef, isStreaming, messages)
+    const lastUserMessageCountRef = useRef(0)
+
+    // Scroll to user message when a new one is added
+    useLayoutEffect(() => {
+        const userMessageCount = messages.filter(
+            (m) => m.type === 'user'
+        ).length
+        if (userMessageCount > lastUserMessageCountRef.current) {
+            // New user message added - scroll after DOM updates but before paint
+            // useLayoutEffect runs synchronously after DOM mutations, perfect for scrolling
+            requestAnimationFrame(() => {
+                if (scrollRef.current) {
+                    const scrollMargin = parseInt(euiTheme.size.l, 10)
+                    scrollUserMessageToTop(scrollRef.current, scrollMargin)
+                }
+            })
+        }
+        lastUserMessageCountRef.current = userMessageCount
+    }, [messages, scrollRef, euiTheme.size.l])
 
     const scrollableStyles = css`
         height: 100%;
@@ -228,7 +253,6 @@ const ChatScrollArea = ({
 
 interface ChatInputAreaProps {
     inputRef: RefObject<HTMLTextAreaElement>
-    scrollRef: RefObject<HTMLDivElement>
     onMetaSemicolon?: () => void
     onStateChange?: (state: {
         onAbortReady: (abort: () => void) => void
@@ -238,7 +262,6 @@ interface ChatInputAreaProps {
 
 const ChatInputArea = ({
     inputRef,
-    scrollRef,
     onMetaSemicolon,
     onStateChange,
 }: ChatInputAreaProps) => {
@@ -251,7 +274,7 @@ const ChatInputArea = ({
         handleAbortReady,
         isStreaming,
         isCooldownActive,
-    } = useChatSubmit(scrollRef)
+    } = useChatSubmit()
 
     useEffect(() => {
         onStateChange?.({
@@ -283,13 +306,12 @@ const ChatInputArea = ({
     )
 }
 
-function useChatSubmit(scrollRef: RefObject<HTMLDivElement | null>) {
+function useChatSubmit() {
     const { submitQuestion, clearNon429Errors, cancelStreaming } =
         useChatActions()
     const isCooldownActive = useIsAskAiCooldownActive()
     const isStreaming = useIsStreaming()
-    const { euiTheme } = useEuiTheme()
-    const scrollMargin = parseInt(euiTheme.size.l, 10)
+    // Note: Scrolling is now handled in ChatScrollArea via useLayoutEffect
 
     const [inputValue, setInputValue] = useState('')
     const abortRef = useRef<(() => void) | null>(null)
@@ -308,15 +330,8 @@ function useChatSubmit(scrollRef: RefObject<HTMLDivElement | null>) {
             clearNon429Errors()
             submitQuestion(trimmed)
             setInputValue('')
-
-            // Scroll to position the user message at the top of the viewport
-            setTimeout(() => {
-                if (scrollRef.current) {
-                    scrollUserMessageToTop(scrollRef.current, scrollMargin)
-                }
-            }, 100)
         },
-        [submitQuestion, isCooldownActive, clearNon429Errors, scrollRef]
+        [submitQuestion, isCooldownActive, clearNon429Errors]
     )
 
     const handleAbort = useCallback(() => {