Fix Cody Chat UI rendering reconciliation (sourcegraph#5850)

Fixes part of
https://linear.app/sourcegraph/issue/CODY-3943/optimize-performance-for-large-chat-windows-in-cody

It looks like since sourcegraph#5036 we've
lost memoization in a few places. As a result, our chat UI doesn't
perform well with more than 4-5 messages (depending on how long these
chat messages are)

This PR tries to improve it by fixing memoization on expensive
components (Transcript, HumanMessage, ...etc)

| Before | After |
| ------- | ------- |
| <video
src="https://github.com/user-attachments/assets/79e7b93c-ac82-4a8c-9b2e-6c712af3fb2c"
/> | <video
src="https://github.com/user-attachments/assets/223b3513-0f30-47cc-b24c-c54f4861329c"
/> |

This PR also fixes a problem with one-box intention detection;
previously, it produced a state where we didn't have a follow-up message
UI during intent resolution.

## Test plan
- Check with the React profile that we don't render components during
message response in places where we don't need to re-render anything
- Check that intent detection in one-box doesn't produce UI flashes

vscode/src/chat/chat-view/ChatController.ts

@@ -751,7 +751,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                           .then(async response => {
                               signal.throwIfAborted()
                               this.chatBuilder.setLastMessageIntent(response?.intent)
-                              this.postViewTranscript()
+                              this.postEmptyMessageInProgress(model)
                               return response
                           })
                           .catch(() => undefined)

vscode/webviews/chat/Transcript.tsx

@@ -11,7 +11,7 @@ import { clsx } from 'clsx'
 import debounce from 'lodash/debounce'
 import isEqual from 'lodash/isEqual'
 import { Search } from 'lucide-react'
-import { type FC, memo, useCallback, useMemo, useRef, useState } from 'react'
+import { type FC, type MutableRefObject, memo, useCallback, useMemo, useRef } from 'react'
 import type { UserAccountInfo } from '../Chat'
 import type { ApiPostMessage } from '../Chat'
 import { Button } from '../components/shadcn/ui/button'
@@ -180,7 +180,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         smartApplyEnabled,
     } = props
 
-    const [intentResults, setIntentResults] = useState<
+    const [intentResults, setIntentResults] = useMutatedValue<
         | { intent: ChatMessage['intent']; allScores?: { intent: string; score: number }[] }
         | undefined
         | null
@@ -193,28 +193,27 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
             editHumanMessage({
                 messageIndexInTranscript: humanMessage.index,
                 editorValue,
-                intent: intentFromSubmit || intentResults?.intent,
-                intentScores: intentFromSubmit ? undefined : intentResults?.allScores,
+                intent: intentFromSubmit || intentResults.current?.intent,
+                intentScores: intentFromSubmit ? undefined : intentResults.current?.allScores,
                 manuallySelectedIntent: !!intentFromSubmit,
             })
         },
-        [humanMessage, intentResults]
+        [humanMessage.index, intentResults]
     )
 
     const onFollowupSubmit = useCallback(
         (editorValue: SerializedPromptEditorValue, intentFromSubmit?: ChatMessage['intent']): void => {
             submitHumanMessage({
                 editorValue,
-                intent: intentFromSubmit || intentResults?.intent,
-                intentScores: intentFromSubmit ? undefined : intentResults?.allScores,
+                intent: intentFromSubmit || intentResults.current?.intent,
+                intentScores: intentFromSubmit ? undefined : intentResults.current?.allScores,
                 manuallySelectedIntent: !!intentFromSubmit,
             })
         },
         [intentResults]
     )
 
     const extensionAPI = useExtensionAPI()
-
     const experimentalOneBoxEnabled = useExperimentalOneBox()
 
     const onChange = useMemo(() => {
@@ -240,7 +239,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                     })
             }
         }, 300)
-    }, [experimentalOneBoxEnabled, extensionAPI])
+    }, [experimentalOneBoxEnabled, extensionAPI, setIntentResults])
 
     const onStop = useCallback(() => {
         getVSCodeAPI().postMessage({
@@ -290,7 +289,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         [reSubmitWithIntent]
     )
 
-    const resetIntent = useCallback(() => setIntentResults(undefined), [])
+    const resetIntent = useCallback(() => setIntentResults(undefined), [setIntentResults])
 
     return (
         <>
@@ -312,6 +311,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                 className={!isFirstInteraction && isLastInteraction ? 'tw-mt-auto' : ''}
                 onEditorFocusChange={resetIntent}
             />
+
             {experimentalOneBoxEnabled && humanMessage.intent && (
                 <InfoMessage>
                     {humanMessage.intent === 'search' ? (
@@ -389,6 +389,17 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
     )
 }, isEqual)
 
+function useMutatedValue<T>(value?: T): [MutableRefObject<T | undefined>, setValue: (value: T) => void] {
+    const valueRef = useRef<T | undefined>(value)
+
+    return [
+        valueRef,
+        useCallback(value => {
+            valueRef.current = value
+        }, []),
+    ]
+}
+
 // TODO(sqs): Do this the React-y way.
 export function focusLastHumanMessageEditor(): void {
     const elements = document.querySelectorAll<HTMLElement>('[data-lexical-editor]')

vscode/webviews/chat/cells/messageCell/human/HumanMessageCell.tsx

@@ -1,12 +1,13 @@
 import {
     type ChatMessage,
+    type SerializedPromptEditorState,
     type SerializedPromptEditorValue,
     serializedPromptEditorStateFromChatMessage,
 } from '@sourcegraph/cody-shared'
 import type { PromptEditorRefAPI } from '@sourcegraph/prompt-editor'
 import isEqual from 'lodash/isEqual'
 import { ColumnsIcon } from 'lucide-react'
-import { type FunctionComponent, memo, useMemo } from 'react'
+import { type FC, memo, useMemo } from 'react'
 import type { UserAccountInfo } from '../../../../Chat'
 import { UserAvatar } from '../../../../components/UserAvatar'
 import { BaseMessageCell, MESSAGE_CELL_AVATAR_SIZE } from '../BaseMessageCell'
@@ -16,10 +17,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from '../../../../components/
 import { getVSCodeAPI } from '../../../../utils/VSCodeApi'
 import { useConfig } from '../../../../utils/useConfig'
 
-/**
- * A component that displays a chat message from the human.
- */
-export const HumanMessageCell: FunctionComponent<{
+interface HumanMessageCellProps {
     message: ChatMessage
     userInfo: UserAccountInfo
     chatEnabled: boolean
@@ -47,9 +45,29 @@ export const HumanMessageCell: FunctionComponent<{
 
     /** For use in storybooks only. */
     __storybook__focus?: boolean
-}> = memo(
-    ({
-        message,
+}
+
+/**
+ * A component that displays a chat message from the human.
+ */
+export const HumanMessageCell: FC<HumanMessageCellProps> = ({ message, ...otherProps }) => {
+    const messageJSON = JSON.stringify(message)
+
+    const initialEditorState = useMemo(
+        () => serializedPromptEditorStateFromChatMessage(JSON.parse(messageJSON)),
+        [messageJSON]
+    )
+
+    return <HumanMessageCellContent {...otherProps} initialEditorState={initialEditorState} />
+}
+
+type HumanMessageCellContent = { initialEditorState: SerializedPromptEditorState } & Omit<
+    HumanMessageCellProps,
+    'message'
+>
+const HumanMessageCellContent = memo<HumanMessageCellContent>(props => {
+    const {
+        initialEditorState,
         userInfo,
         chatEnabled = true,
         isFirstMessage,
@@ -65,50 +83,43 @@ export const HumanMessageCell: FunctionComponent<{
         editorRef,
         __storybook__focus,
         onEditorFocusChange,
-    }) => {
-        const messageJSON = JSON.stringify(message)
-        const initialEditorState = useMemo(
-            () => serializedPromptEditorStateFromChatMessage(JSON.parse(messageJSON)),
-            [messageJSON]
-        )
-
-        return (
-            <BaseMessageCell
-                speakerIcon={
-                    <UserAvatar
-                        user={userInfo.user}
-                        size={MESSAGE_CELL_AVATAR_SIZE}
-                        sourcegraphGradientBorder={true}
-                    />
-                }
-                speakerTitle={userInfo.user.displayName ?? userInfo.user.username}
-                cellAction={isFirstMessage && <OpenInNewEditorAction />}
-                content={
-                    <HumanMessageEditor
-                        userInfo={userInfo}
-                        initialEditorState={initialEditorState}
-                        placeholder={isFirstMessage ? 'Ask...' : 'Ask a followup...'}
-                        isFirstMessage={isFirstMessage}
-                        isSent={isSent}
-                        isPendingPriorResponse={isPendingPriorResponse}
-                        onChange={onChange}
-                        onSubmit={onSubmit}
-                        onStop={onStop}
-                        disabled={!chatEnabled}
-                        isFirstInteraction={isFirstInteraction}
-                        isLastInteraction={isLastInteraction}
-                        isEditorInitiallyFocused={isEditorInitiallyFocused}
-                        editorRef={editorRef}
-                        __storybook__focus={__storybook__focus}
-                        onEditorFocusChange={onEditorFocusChange}
-                    />
-                }
-                className={className}
-            />
-        )
-    },
-    isEqual
-)
+    } = props
+
+    return (
+        <BaseMessageCell
+            speakerIcon={
+                <UserAvatar
+                    user={userInfo.user}
+                    size={MESSAGE_CELL_AVATAR_SIZE}
+                    sourcegraphGradientBorder={true}
+                />
+            }
+            speakerTitle={userInfo.user.displayName ?? userInfo.user.username}
+            cellAction={isFirstMessage && <OpenInNewEditorAction />}
+            content={
+                <HumanMessageEditor
+                    userInfo={userInfo}
+                    initialEditorState={initialEditorState}
+                    placeholder={isFirstMessage ? 'Ask...' : 'Ask a followup...'}
+                    isFirstMessage={isFirstMessage}
+                    isSent={isSent}
+                    isPendingPriorResponse={isPendingPriorResponse}
+                    onChange={onChange}
+                    onSubmit={onSubmit}
+                    onStop={onStop}
+                    disabled={!chatEnabled}
+                    isFirstInteraction={isFirstInteraction}
+                    isLastInteraction={isLastInteraction}
+                    isEditorInitiallyFocused={isEditorInitiallyFocused}
+                    editorRef={editorRef}
+                    __storybook__focus={__storybook__focus}
+                    onEditorFocusChange={onEditorFocusChange}
+                />
+            }
+            className={className}
+        />
+    )
+}, isEqual)
 
 const OpenInNewEditorAction = () => {
     const {

vscode/webviews/tabs/TabsBar.tsx

@@ -19,12 +19,13 @@ import { getVSCodeAPI } from '../utils/VSCodeApi'
 import { View } from './types'
 
 import { CodyIDE, FeatureFlag, isDefined } from '@sourcegraph/cody-shared'
-import { type FC, Fragment, forwardRef, useCallback, useMemo, useState } from 'react'
+import { type FC, Fragment, forwardRef, memo, useCallback, useMemo, useState } from 'react'
 import { Kbd } from '../components/Kbd'
 import { Tooltip, TooltipContent, TooltipTrigger } from '../components/shadcn/ui/tooltip'
 import { useConfig } from '../utils/useConfig'
 
 import { useExtensionAPI } from '@sourcegraph/prompt-editor'
+import { isEqual } from 'lodash'
 import { downloadChatHistory } from '../chat/downloadChatHistory'
 import { Button } from '../components/shadcn/ui/button'
 import { useFeatureFlag } from '../utils/useFeatureFlags'
@@ -68,7 +69,8 @@ interface TabConfig {
     subActions?: TabSubAction[]
 }
 
-export const TabsBar: React.FC<TabsBarProps> = ({ currentView, setView, IDE }) => {
+export const TabsBar = memo<TabsBarProps>(props => {
+    const { currentView, setView, IDE } = props
     const tabItems = useTabs({ IDE })
     const {
         config: { webviewType, multipleWebviewsEnabled },
@@ -183,7 +185,7 @@ export const TabsBar: React.FC<TabsBarProps> = ({ currentView, setView, IDE }) =
             </Tabs.List>
         </div>
     )
-}
+}, isEqual)
 
 interface ActionButtonWithConfirmationProps {
     title: string

vscode/webviews/utils/debugger.ts

@@ -0,0 +1,23 @@
+import { useEffect, useRef } from 'react'
+
+export function useTracePropsUpdate(props: any) {
+    const prev = useRef(props)
+
+    useEffect(() => {
+        const changedProps = Object.entries(props).reduce(
+            (prop, [k, v]) => {
+                if (prev.current[k] !== v) {
+                    prop[k] = [prev.current[k], v]
+                }
+                return prop
+            },
+            {} as Record<any, any>
+        )
+
+        if (Object.keys(changedProps).length > 0) {
+            console.log('Changed props:', changedProps)
+        }
+
+        prev.current = props
+    })
+}

vscode/webviews/utils/useConfig.tsx

@@ -5,6 +5,7 @@ import {
     type ReactNode,
     createContext,
     useContext,
+    useMemo,
 } from 'react'
 import type { ExtensionMessage } from '../../src/chat/protocol'
 import type { UserAccountInfo } from '../Chat'
@@ -43,18 +44,22 @@ export function useConfig(): Config {
 }
 
 export function useUserAccountInfo(): UserAccountInfo {
-    const value = useConfig()
-    if (!value.authStatus.authenticated) {
+    const { authStatus, isDotComUser, clientCapabilities, userProductSubscription } = useConfig()
+
+    if (!authStatus.authenticated) {
         throw new Error(
             'useUserAccountInfo must be used within a ConfigProvider with authenticated user'
         )
     }
-    return {
-        isCodyProUser: isCodyProUser(value.authStatus, value.userProductSubscription ?? null),
-        // Receive this value from the extension backend to make it work
-        // with E2E tests where change the DOTCOM_URL via the env variable TESTING_DOTCOM_URL.
-        isDotComUser: value.isDotComUser,
-        user: value.authStatus,
-        IDE: value.clientCapabilities.agentIDE,
-    }
+    return useMemo<UserAccountInfo>(
+        () => ({
+            isCodyProUser: isCodyProUser(authStatus, userProductSubscription ?? null),
+            // Receive this value from the extension backend to make it work
+            // with E2E tests where change the DOTCOM_URL via the env variable TESTING_DOTCOM_URL.
+            isDotComUser: isDotComUser,
+            user: authStatus,
+            IDE: clientCapabilities.agentIDE,
+        }),
+        [authStatus, isDotComUser, clientCapabilities, userProductSubscription]
+    )
 }