Add additional memoization to different UI levels

lib/shared/src/lexicalEditor/editorState.ts

@@ -147,9 +147,9 @@ export function serializedPromptEditorStateFromText(text: string): SerializedPro
     }
 }
 
-export function serializedPromptEditorStateFromChatMessage(
+export const serializedPromptEditorStateFromChatMessage = (
     chatMessage: ChatMessage
-): SerializedPromptEditorState {
+): SerializedPromptEditorState => {
     function isCompatibleVersionEditorState(value: unknown): value is SerializedPromptEditorState {
         if (!value) {
             return false

vscode/webviews/App.tsx

@@ -53,6 +53,7 @@ export const App: React.FunctionComponent<{ vscodeAPI: VSCodeWrapper }> = ({ vsc
     const dispatchClientAction = useClientActionDispatcher()
 
     const clientConfigAttribution = clientConfig?.attribution ?? 'none'
+
     const guardrails = useMemo(() => {
         return createGuardrailsImpl(clientConfigAttribution, (snippet: string) => {
             vscodeAPI.postMessage({

vscode/webviews/Chat.story.tsx

@@ -29,7 +29,6 @@ const meta: Meta<typeof Chat> = {
             postMessage: () => {},
             onMessage: () => () => {},
         },
-        setView: () => {},
         models: FIXTURE_MODELS,
         guardrails: new MockNoGuardrails(),
     } satisfies React.ComponentProps<typeof Chat>,

vscode/webviews/Chat.tsx

@@ -15,7 +15,6 @@ import type {
 import styles from './Chat.module.css'
 import { Transcript } from './chat/Transcript'
 
-import type { View } from './tabs'
 import type { VSCodeWrapper } from './utils/VSCodeApi'
 import { SpanManager } from './utils/spanManager'
 import { getTraceparentFromSpanContext } from './utils/telemetry'
@@ -36,10 +35,7 @@ interface ChatboxProps {
     models: Model[]
     vscodeAPI: Pick<VSCodeWrapper, 'postMessage' | 'onMessage'>
     guardrails: Guardrails
-    showWelcomeMessage?: boolean
     showIDESnippetActions?: boolean
-    setView: (view: View) => void
-    isWorkspacesUpgradeCtaEnabled?: boolean
 }
 
 export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>> = ({
@@ -50,10 +46,7 @@ export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>
     vscodeAPI,
     chatEnabled = true,
     guardrails,
-    showWelcomeMessage = true,
     showIDESnippetActions = true,
-    setView,
-    isWorkspacesUpgradeCtaEnabled,
 }) => {
     const transcriptRef = useRef(transcript)
     transcriptRef.current = transcript

vscode/webviews/CodyPanel.tsx

@@ -5,6 +5,7 @@ import {
     type CodyNotice,
     FeatureFlag,
     type Guardrails,
+    type Model,
     type UserProductSubscription,
     type WebviewToExtensionAPI,
     firstValueFrom,
@@ -28,6 +29,8 @@ import { useUserAccountInfo } from './utils/useConfig'
 import { useFeatureFlag } from './utils/useFeatureFlags'
 import { TabViewContext } from './utils/useTabView'
 
+const DEFAULT_CHAT_MODELS: Model[] = []
+
 interface CodyPanelProps {
     view: View
     setView: (view: View) => void
@@ -93,11 +96,13 @@ export const CodyPanel: FunctionComponent<CodyPanelProps> = ({
             [api.mcpSettings]
         )
     )
-    // workspace upgrade eligibility should be that the flag is set, is on dotcom and only has one account. This prevents enterprise customers that are logged into multiple endpoints from seeing the CTA
+    // Workspace upgrade eligibility should be that the flag is set, is on dotcom and only has one account.
+    // This prevents enterprise customers that are logged into multiple endpoints from seeing the CTA
     const isWorkspacesUpgradeCtaEnabled =
         useFeatureFlag(FeatureFlag.SourcegraphTeamsUpgradeCTA) &&
         isDotComUser &&
         config.endpointHistory?.length === 1
+
     useEffect(() => {
         onExternalApiReady?.(externalAPI)
     }, [onExternalApiReady, externalAPI])
@@ -162,13 +167,10 @@ export const CodyPanel: FunctionComponent<CodyPanelProps> = ({
                             messageInProgress={messageInProgress}
                             transcript={transcript}
                             tokenUsage={tokenUsage}
-                            models={chatModels || []}
+                            models={chatModels || DEFAULT_CHAT_MODELS}
                             vscodeAPI={vscodeAPI}
                             guardrails={guardrails}
                             showIDESnippetActions={showIDESnippetActions}
-                            showWelcomeMessage={showWelcomeMessage}
-                            setView={setView}
-                            isWorkspacesUpgradeCtaEnabled={isWorkspacesUpgradeCtaEnabled}
                         />
                     )}
                     {view === View.History && (

vscode/webviews/chat/Transcript.tsx

@@ -201,7 +201,7 @@ export const Transcript: FC<TranscriptProps> = props => {
         []
     )
 
-    const scrollTotheBottom = useCallback(() => {
+    const scrollToBottom = useCallback(() => {
         scrollableContainer?.scroll({
             top: scrollableContainer?.scrollHeight,
             behavior: 'smooth',
@@ -285,7 +285,7 @@ export const Transcript: FC<TranscriptProps> = props => {
                     {scrollableContainer && <ScrollbarMarkers scrollContainer={scrollableContainer} />}
 
                     {!isAtBottomDebounced && interactions.length > 1 && (
-                        <ScrollDown onClick={scrollTotheBottom} />
+                        <ScrollDown onClick={scrollToBottom} />
                     )}
 
                     <div className="tw-bg-[var(--vscode-input-background)]">
@@ -409,11 +409,20 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
     } = props
 
     const { activeChatContext, setActiveChatContext } = props
+
     const humanEditorRef = useRef<PromptEditorRefAPI | null>(null)
     const lastEditorRef = useContext(LastEditorContext)
-    useImperativeHandle(parentEditorRef, () => humanEditorRef.current)
-
     const [selectedIntent, setSelectedIntent] = useState<ChatMessage['intent']>(humanMessage?.intent)
+    // We track, ephemerally, the code blocks that are being regenerated so
+    // we can show an accurate loading indicator or error message on those
+    // blocks.
+    const [regeneratingCodeBlocks, setRegeneratingCodeBlocks] = useState<RegeneratingCodeBlockState[]>(
+        []
+    )
+
+    const vscodeAPI = useMemo(() => getVSCodeAPI(), [])
+
+    useImperativeHandle(parentEditorRef, () => humanEditorRef.current)
 
     // Reset intent to 'chat' when there are no interactions (new chat)
     useEffect(() => {
@@ -482,7 +491,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         [humanMessage, setActiveChatContext, isLastSentInteraction, lastEditorRef]
     )
 
-    const vscodeAPI = getVSCodeAPI()
     const onStop = useCallback(() => {
         vscodeAPI.postMessage({
             command: 'abort',
@@ -500,17 +508,11 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
     const timeToFirstTokenSpan = useRef<Span>()
     const hasRecordedFirstToken = useRef(false)
 
-    const [isLoading, setIsLoading] = useState(assistantMessage?.isLoading)
-
     const [isThoughtProcessOpened, setThoughtProcessOpened] = useLocalStorage(
         'cody.thinking-space.open',
         true
     )
 
-    useEffect(() => {
-        setIsLoading(assistantMessage?.isLoading)
-    }, [assistantMessage])
-
     const humanMessageText = humanMessage.text
     const smartApplyWithInstruction = useMemo(() => {
         if (!smartApply) return undefined
@@ -617,10 +619,10 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
             context.with(activeChatContext, startRenderSpan)
         }
         // Case 2: End rendering if loading is complete and a render span exists
-        else if (!isLoading && renderSpan.current) {
+        else if (!assistantMessage?.isLoading && renderSpan.current) {
             endRenderSpan()
         }
-    }, [assistantMessage, activeChatContext, setActiveChatContext, spanManager, isLoading])
+    }, [assistantMessage, activeChatContext, setActiveChatContext, spanManager])
 
     const humanMessageInfo = useMemo(() => {
         // See SRCH-942: it's critical to memoize this value to avoid repeated
@@ -651,12 +653,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         [humanMessage, onUserAction, selectedIntent]
     )
 
-    // We track, ephemerally, the code blocks that are being regenerated so
-    // we can show an accurate loading indicator or error message on those
-    // blocks.
-    const [regeneratingCodeBlocks, setRegeneratingCodeBlocks] = useState<RegeneratingCodeBlockState[]>(
-        []
-    )
     useClientActionListener(
         { isActive: true, selector: event => Boolean(event.regenerateStatus) },
         useCallback(event => {
@@ -687,7 +683,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
 
     const onRegenerate = useCallback(
         (code: string, language?: string) => {
-            if (assistantMessage) {
+            if (assistantMessage?.index !== undefined) {
                 const id = uuid.v4()
                 regenerateCodeBlock({ id, code, language, index: assistantMessage.index })
                 setRegeneratingCodeBlocks(blocks => [
@@ -698,7 +694,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                 console.warn('tried to regenerate a code block, but there is no assistant message')
             }
         },
-        [assistantMessage]
+        [assistantMessage?.index]
     )
 
     const isAgenticMode = useMemo(
@@ -770,7 +766,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                         key={assistantMessage.index}
                         userInfo={userInfo}
                         models={models}
-                        chatEnabled={chatEnabled}
                         message={assistantMessage}
                         copyButtonOnSubmit={copyButtonOnSubmit}
                         insertButtonOnSubmit={insertButtonOnSubmit}
@@ -781,7 +776,6 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                         humanMessage={humanMessageInfo}
                         isLoading={isLastSentInteraction && assistantMessage.isLoading}
                         smartApply={isAgenticMode ? undefined : smartApplyWithInstruction}
-                        isLastSentInteraction={isLastSentInteraction}
                         setThoughtProcessOpened={setThoughtProcessOpened}
                         isThoughtProcessOpened={isThoughtProcessOpened}
                     />

vscode/webviews/chat/cells/messageCell/assistant/AssistantMessageCell.tsx

@@ -39,7 +39,6 @@ export const AssistantMessageCell: FunctionComponent<{
     humanMessage: PriorHumanMessageInfo | null
 
     userInfo: UserAccountInfo
-    chatEnabled: boolean
     isLoading: boolean
 
     copyButtonOnSubmit?: CodeBlockActionsProps['copyButtonOnSubmit']
@@ -54,7 +53,6 @@ export const AssistantMessageCell: FunctionComponent<{
 
     postMessage?: ApiPostMessage
     guardrails: Guardrails
-    isLastSentInteraction: boolean
 }> = memo(
     ({
         message,
@@ -69,7 +67,6 @@ export const AssistantMessageCell: FunctionComponent<{
         postMessage,
         guardrails,
         smartApply,
-        isLastSentInteraction: isLastInteraction,
         isThoughtProcessOpened,
         setThoughtProcessOpened,
     }) => {
@@ -79,9 +76,7 @@ export const AssistantMessageCell: FunctionComponent<{
         )
         const chatModel = useChatModelByID(message.model, models)
         const isAborted = isAbortErrorOrSocketHangUp(message.error)
-
         const hasLongerResponseTime = chatModel?.tags?.includes(ModelTag.StreamDisabled)
-
         const messageIntent = humanMessage?.intent ?? message.intent ?? 'chat'
         const isSearchIntent = messageIntent === 'search'
 

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

@@ -48,7 +48,7 @@ interface HumanMessageCellProps {
 /**
  * A component that displays a chat message from the human.
  */
-export const HumanMessageCell: FC<HumanMessageCellProps> = ({ message, ...otherProps }) => {
+export const HumanMessageCell: FC<HumanMessageCellProps> = memo(({ message, ...otherProps }) => {
     // Don't render the editor if the message text is explicitly undefined or empty,
     // and it's been sent already and it's not the last interaction (i.e. there is a tool result response).
     if (
@@ -61,14 +61,13 @@ export const HumanMessageCell: FC<HumanMessageCellProps> = ({ message, ...otherP
     }
 
     const messageJSON = JSON.stringify(message)
-
     const initialEditorState = useMemo(
         () => serializedPromptEditorStateFromChatMessage(JSON.parse(messageJSON)),
         [messageJSON]
     )
 
     return <HumanMessageCellContent {...otherProps} initialEditorState={initialEditorState} />
-}
+})
 
 type HumanMessageCellContent = {
     initialEditorState: SerializedPromptEditorState

vscode/webviews/components/MarkdownFromCody.tsx

@@ -126,6 +126,7 @@ export const MarkdownFromCody: FunctionComponent<{
 }
 
 let _markdownPluginProps: ReturnType<typeof markdownPluginProps> | undefined
+
 function markdownPluginProps(
     prefixRemarkPlugins: Pluggable[] = []
 ): Pick<ComponentProps<typeof Markdown>, 'rehypePlugins' | 'remarkPlugins'> {
@@ -181,5 +182,6 @@ function markdownPluginProps(
         ],
         remarkPlugins: [...prefixRemarkPlugins, remarkGFM, remarkAttachFilePathToCodeBlocks],
     }
+
     return _markdownPluginProps
 }