Add additional memoization to different UI levels

lib/shared/src/lexicalEditor/editorState.ts

@@ -11,6 +11,7 @@ import { type ContextItem, ContextItemSource } from '../codebase-context/message
 import type { RangeData } from '../common/range'
 import { displayPath } from '../editor/displayPath'
 import type { PromptString } from '../prompt/prompt-string'
+import { memoize } from '../utils'
 import { AT_MENTION_SERIALIZED_PREFIX, deserializeParagraph } from './atMentionsSerializer'
 import {
     CONTEXT_ITEM_MENTION_NODE_TYPE,
@@ -147,36 +148,36 @@ export function serializedPromptEditorStateFromText(text: string): SerializedPro
     }
 }
 
-export function serializedPromptEditorStateFromChatMessage(
-    chatMessage: ChatMessage
-): SerializedPromptEditorState {
-    function isCompatibleVersionEditorState(value: unknown): value is SerializedPromptEditorState {
-        if (!value) {
-            return false
-        }
+export const serializedPromptEditorStateFromChatMessage = memoize(
+    (chatMessage: ChatMessage): SerializedPromptEditorState => {
+        function isCompatibleVersionEditorState(value: unknown): value is SerializedPromptEditorState {
+            if (!value) {
+                return false
+            }
 
-        const editorState = value as SerializedPromptEditorState
+            const editorState = value as SerializedPromptEditorState
 
-        // We can read this if the version of the serialized text is compatible
-        // or its minimum version is compatible.
-        return (
-            SUPPORTED_READER_VERSIONS.includes(editorState.v) ||
-            SUPPORTED_READER_VERSIONS.includes(editorState.minReaderV ?? DEFAULT_MIN_READER_V)
-        )
-    }
+            // We can read this if the version of the serialized text is compatible
+            // or its minimum version is compatible.
+            return (
+                SUPPORTED_READER_VERSIONS.includes(editorState.v) ||
+                SUPPORTED_READER_VERSIONS.includes(editorState.minReaderV ?? DEFAULT_MIN_READER_V)
+            )
+        }
 
-    if (isCompatibleVersionEditorState(chatMessage.editorState)) {
-        return chatMessage.editorState
-    }
+        if (isCompatibleVersionEditorState(chatMessage.editorState)) {
+            return chatMessage.editorState
+        }
 
-    // Fall back to using plain text for chat messages that don't have a serialized Lexical editor
-    // state that we recognize.
-    //
-    // It would be smoother to automatically import or convert textual @-mentions to the Lexical
-    // mention nodes, but that would add a lot of extra complexity for the relatively rare use case
-    // of editing old messages in your chat history.
-    return serializedPromptEditorStateFromText(chatMessage.text ? chatMessage.text.toString() : '')
-}
+        // Fall back to using plain text for chat messages that don't have a serialized Lexical editor
+        // state that we recognize.
+        //
+        // It would be smoother to automatically import or convert textual @-mentions to the Lexical
+        // mention nodes, but that would add a lot of extra complexity for the relatively rare use case
+        // of editing old messages in your chat history.
+        return serializedPromptEditorStateFromText(chatMessage.text ? chatMessage.text.toString() : '')
+    }
+)
 
 export function contextItemsFromPromptEditorValue(
     state: SerializedPromptEditorState

lib/shared/src/utils.ts

@@ -1,3 +1,4 @@
+import isEqual from 'lodash/isEqual'
 import { logError } from './logger'
 
 type PromiseResolverFn<T, E> = (value?: T, err?: E) => void
@@ -187,3 +188,21 @@ export type PartialDeep<T> = {
           ? PartialDeep<T[P]>
           : T[P]
 }
+
+export function memoize<T extends (...args: any[]) => any>(
+    func: T
+): (...args: Parameters<T>) => ReturnType<T> {
+    let lastArguments: any[] | null = null
+    let lastCalculatedValue: ReturnType<T> | null = null
+
+    return (...args: Parameters<T>): ReturnType<T> => {
+        if (isEqual(lastArguments, args)) {
+            return lastCalculatedValue as ReturnType<T>
+        }
+
+        lastArguments = args
+        lastCalculatedValue = func(args)
+
+        return lastCalculatedValue as ReturnType<T>
+    }
+}

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

@@ -61,7 +61,6 @@ export const HumanMessageCell: FC<HumanMessageCellProps> = ({ message, ...otherP
     }
 
     const messageJSON = JSON.stringify(message)
-
     const initialEditorState = useMemo(
         () => serializedPromptEditorStateFromChatMessage(JSON.parse(messageJSON)),
         [messageJSON]

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
 }