fix(max): auto-scrolling experience (#32665)

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

Author: skoob13

frontend/__snapshots__/scenes-app-experiments--experiment-v-3-with-asymmetric-intervals--dark.png

@@ -13,6 +13,7 @@ import {
     formChunk,
     generationFailureChunk,
     humanMessage,
+    longResponseChunk,
 } from './__mocks__/chatResponse.mocks'
 import conversationList from './__mocks__/conversationList.json'
 import { MaxInstance, MaxInstanceProps } from './Max'
@@ -367,3 +368,43 @@ ThreadWithOpenedSuggestions.parameters = {
         waitForLoadersToDisappear: false,
     },
 }
+
+export const ThreadScrollsToBottomOnNewMessages: StoryFn = () => {
+    useStorybookMocks({
+        get: {
+            '/api/environments/:team_id/conversations/': () => [200, conversationList],
+        },
+        post: {
+            '/api/environments/:team_id/conversations/': (_, res, ctx) =>
+                res(ctx.delay(100), ctx.text(longResponseChunk)),
+        },
+    })
+
+    const { conversation } = useValues(maxLogic)
+    const { setConversationId } = useActions(maxLogic)
+    const logic = maxThreadLogic({ conversationId: 'poem', conversation })
+    const { threadRaw } = useValues(logic)
+    const { askMax } = useActions(logic)
+
+    useEffect(() => {
+        setConversationId('poem')
+    }, [setConversationId])
+
+    const messagesSet = threadRaw.length > 0
+    useEffect(() => {
+        if (messagesSet) {
+            askMax('This message must be on the top of the container')
+        }
+    }, [messagesSet, askMax])
+
+    return (
+        <div className="h-[800px] overflow-y-auto SidePanel3000__content">
+            <Template />
+        </div>
+    )
+}
+ThreadScrollsToBottomOnNewMessages.parameters = {
+    testOptions: {
+        waitForLoadersToDisappear: false,
+    },
+}

frontend/__snapshots__/scenes-app-max-ai--empty-thread-loading--dark.png

@@ -27,6 +27,7 @@ import { sidePanelLogic } from '~/layout/navigation-3000/sidepanel/sidePanelLogi
 import { SidePanelTab } from '~/types'
 
 import { AnimatedBackButton } from './components/AnimatedBackButton'
+import { ThreadAutoScroller } from './components/ThreadAutoScroller'
 import { ConversationHistory } from './ConversationHistory'
 import { HistoryPreview } from './HistoryPreview'
 import { Intro } from './Intro'
@@ -212,10 +213,11 @@ export const MaxInstance = React.memo(function MaxInstance({ sidePanel }: MaxIns
                         <HistoryPreview sidePanel={sidePanel} />
                     </div>
                 ) : (
-                    <>
+                    /** Must be the last child and be a direct descendant of the scrollable element */
+                    <ThreadAutoScroller>
                         <Thread />
                         <QuestionInput isFloating />
-                    </>
+                    </ThreadAutoScroller>
                 )}
             </BindLogic>
         </>

frontend/__snapshots__/scenes-app-max-ai--empty-thread-loading--light.png

@@ -80,9 +80,9 @@ export function Thread(): JSX.Element | null {
             ) : threadGrouped.length > 0 ? (
                 threadGrouped.map((group, index) => (
                     <MessageGroup
-                        key={index}
+                        // Reset the components when the thread changes
+                        key={`${conversationId}-${index}`}
                         messages={group}
-                        index={index}
                         isFinal={index === threadGrouped.length - 1}
                     />
                 ))
@@ -130,7 +130,6 @@ function MessageGroupContainer({
 interface MessageGroupProps {
     messages: ThreadMessage[]
     isFinal: boolean
-    index: number
 }
 
 function MessageGroup({ messages, isFinal: isFinalGroup }: MessageGroupProps): JSX.Element {
@@ -260,7 +259,7 @@ const MessageTemplate = React.forwardRef<HTMLDivElement, MessageTemplateProps>(f
     return (
         <div
             className={twMerge(
-                'flex flex-col gap-px w-full break-words',
+                'flex flex-col gap-px w-full break-words scroll-mt-12',
                 type === 'human' ? 'items-end' : 'items-start',
                 className
             )}
@@ -360,7 +359,7 @@ function AssistantMessageForm({ form }: AssistantMessageFormProps): JSX.Element
     )
 }
 
-function VisualizationAnswer({
+const VisualizationAnswer = React.memo(function VisualizationAnswer({
     message,
     status,
     isEditingInsight,
@@ -458,7 +457,7 @@ function VisualizationAnswer({
                   </MessageTemplate>
               </>
           )
-}
+})
 
 function RetriableFailureActions(): JSX.Element {
     const { retryLastMessage } = useActions(maxThreadLogic)

frontend/__snapshots__/scenes-app-max-ai--thread--dark.png

@@ -93,3 +93,16 @@ const formMessage: AssistantMessage = {
 }
 
 export const formChunk = generateChunk(['event: message', `data: ${JSON.stringify(formMessage)}`])
+
+export const longMessage: AssistantMessage = {
+    type: AssistantMessageType.Assistant,
+    content: 'This\n\nis\n\na\n\nlong\n\nmessage\n\nthat\n\nshould\n\nbe\n\nsplit\n\ninto\n\nmultiple\n\nlines',
+    id: 'assistant-2',
+}
+
+export const longResponseChunk = generateChunk([
+    'event: message',
+    `data: ${JSON.stringify(humanMessage)}`,
+    'event: message',
+    `data: ${JSON.stringify(longMessage)}`,
+])

frontend/__snapshots__/scenes-app-max-ai--thread--light.png

@@ -0,0 +1,118 @@
+import { useValues } from 'kea'
+import { useEffect, useRef } from 'react'
+
+import { getScrollableContainer } from '../maxLogic'
+import { maxThreadLogic } from '../maxThreadLogic'
+
+export function ThreadAutoScroller({ children }: { children: React.ReactNode }): JSX.Element {
+    const { streamingActive, threadGrouped } = useValues(maxThreadLogic)
+
+    const scrollOrigin = useRef({ user: false, programmatic: false, resizing: false })
+    const sentinelRef = useRef<HTMLDivElement | null>(null)
+
+    const scrollToBottom = useRef(() => {
+        if (!sentinelRef.current) {
+            return
+        }
+
+        // Get the root scrollable element
+        const scrollableContainer = getScrollableContainer(sentinelRef.current)
+
+        if (scrollableContainer) {
+            // Lock the scroll listener, so we don't detect the programmatic scroll as a user scroll
+            scrollOrigin.current.programmatic = true
+            scrollableContainer.scrollTo({ top: scrollableContainer.scrollHeight })
+            // Reset the scroll when we're done
+            requestAnimationFrame(() => {
+                scrollOrigin.current.programmatic = false
+            })
+        }
+    }).current // Keep the stable reference
+
+    useEffect(() => {
+        const scrollableContainer = getScrollableContainer(sentinelRef.current)
+        if (!sentinelRef.current || !streamingActive || !scrollableContainer) {
+            return
+        }
+
+        // Detect if the user has scrolled the content during generation,
+        // so we can stop auto-scrolling
+        function scrollListener(event: Event): void {
+            if (
+                scrollOrigin.current.programmatic ||
+                scrollOrigin.current.resizing ||
+                !sentinelRef.current ||
+                !event.target
+            ) {
+                return
+            }
+
+            const scrollableContainer = event.target as HTMLElement
+            // Can be tracked through the IntersectionObserver, but the intersection observer event is fired after the scroll event,
+            // so it adds an annoying delay.
+            const isAtBottom =
+                sentinelRef.current.getBoundingClientRect().top <= scrollableContainer.getBoundingClientRect().bottom
+
+            if (!isAtBottom) {
+                scrollOrigin.current.user = true
+            }
+        }
+        scrollableContainer.addEventListener('scroll', scrollListener)
+
+        // When the thread is resized during generation, we need to scroll to the bottom
+        let resizeTimeout: NodeJS.Timeout | null = null
+        // eslint-disable-next-line compat/compat
+        const resizeObserver = new ResizeObserver(() => {
+            if (scrollOrigin.current.user) {
+                return
+            }
+
+            if (resizeTimeout) {
+                clearTimeout(resizeTimeout)
+                resizeTimeout = null
+            }
+
+            scrollOrigin.current.resizing = true
+            scrollToBottom()
+            // Block the scroll listener from firing
+            resizeTimeout = setTimeout(() => {
+                scrollOrigin.current.resizing = false
+                resizeTimeout = null
+            }, 50)
+        })
+        resizeObserver.observe(sentinelRef.current)
+
+        // Reset the user interaction if we've reached the bottom of the screen
+        // eslint-disable-next-line compat/compat
+        const intersectionObserver = new IntersectionObserver(([entries]) => {
+            if (!scrollOrigin.current.programmatic && scrollOrigin.current.user && entries.isIntersecting) {
+                scrollOrigin.current.user = false
+            }
+        })
+        intersectionObserver.observe(sentinelRef.current)
+
+        return () => {
+            resizeObserver.disconnect()
+            intersectionObserver.disconnect()
+            scrollableContainer.removeEventListener('scroll', scrollListener)
+            scrollOrigin.current = { user: false, programmatic: false, resizing: false }
+            if (resizeTimeout) {
+                clearTimeout(resizeTimeout)
+            }
+        }
+    }, [streamingActive, scrollToBottom])
+
+    useEffect(() => {
+        if (!streamingActive || scrollOrigin.current.user) {
+            return
+        }
+        scrollToBottom()
+    }, [streamingActive, scrollToBottom, threadGrouped]) // Scroll when the thread updates
+
+    return (
+        <>
+            {children}
+            <div id="max-sentinel" className="pointer-events-none h-0" ref={sentinelRef} />
+        </>
+    )
+}

frontend/__snapshots__/scenes-app-max-ai--thread-scrolls-to-bottom-on-new-messages--dark.png

@@ -186,7 +186,7 @@ export const maxLogic = kea<maxLogicType>([
                     // eslint-disable-next-line @typescript-eslint/no-unused-vars -- Used for conversation restoration
                     _?: {
                         /** If true, the current thread will not be updated with the retrieved conversation. */
-                        doNotUpdateCurrentThread: boolean
+                        doNotUpdateCurrentThread?: boolean
                     }
                 ) => {
                     const response = await api.conversations.list()
@@ -249,7 +249,7 @@ export const maxLogic = kea<maxLogicType>([
                 return (
                     !conversationHistory.length &&
                     conversationHistoryLoading &&
-                    conversationId &&
+                    !!conversationId &&
                     !isTempId(conversationId) &&
                     !conversation
                 )
@@ -460,7 +460,7 @@ export const maxLogic = kea<maxLogicType>([
     permanentlyMount(), // Prevent state from being reset when Max is unmounted, especially key in the side panel
 ])
 
-function getScrollableContainer(element?: Element | null): HTMLElement | null {
+export function getScrollableContainer(element?: Element | null): HTMLElement | null {
     if (!element) {
         return null
     }

frontend/__snapshots__/scenes-app-max-ai--thread-scrolls-to-bottom-on-new-messages--light.png

@@ -35,13 +35,7 @@ import { Conversation, ConversationDetail, ConversationStatus } from '~/types'
 import { maxGlobalLogic } from './maxGlobalLogic'
 import { maxLogic } from './maxLogic'
 import type { maxThreadLogicType } from './maxThreadLogicType'
-import {
-    isAssistantMessage,
-    isAssistantToolCallMessage,
-    isHumanMessage,
-    isReasoningMessage,
-    isVisualizationMessage,
-} from './utils'
+import { isAssistantMessage, isAssistantToolCallMessage, isHumanMessage, isReasoningMessage } from './utils'
 
 export type MessageStatus = 'loading' | 'completed' | 'error'
 
@@ -106,7 +100,6 @@ export const maxThreadLogic = kea<maxThreadLogicType>([
                 'prependOrReplaceConversation as updateGlobalConversationCache',
                 'setActiveStreamingThreads',
                 'setConversationId',
-                'scrollThreadToBottom',
                 'setAutoRun',
             ],
         ],
@@ -363,18 +356,6 @@ export const maxThreadLogic = kea<maxThreadLogicType>([
             }
         },
 
-        addMessage: (payload) => {
-            if (isHumanMessage(payload.message) || isVisualizationMessage(payload.message)) {
-                actions.scrollThreadToBottom()
-            }
-        },
-
-        replaceMessage: (payload) => {
-            if (isVisualizationMessage(payload.message)) {
-                actions.scrollThreadToBottom()
-            }
-        },
-
         completeThreadGeneration: () => {
             // Update the conversation history to include the new conversation
             actions.loadConversationHistory({ doNotUpdateCurrentThread: true })