Chat: preserve unsent chat messages when switching tabs

vscode/webviews/Chat.story.tsx

@@ -1,4 +1,8 @@
-import { ExtensionAPIProviderForTestsOnly, MOCK_API } from '@sourcegraph/prompt-editor'
+import {
+    ExtensionAPIProviderForTestsOnly,
+    MOCK_API,
+    type SerializedPromptEditorValue,
+} from '@sourcegraph/prompt-editor'
 import type { Meta, StoryObj } from '@storybook/react'
 import { Observable } from 'observable-fns'
 import { Chat } from './Chat'
@@ -27,6 +31,7 @@ const meta: Meta<typeof Chat> = {
             onMessage: () => () => {},
         },
         setView: () => {},
+        updateEditorStateOnChange: () => {},
     } satisfies React.ComponentProps<typeof Chat>,
 
     decorators: [VSCodeWebview],
@@ -76,3 +81,12 @@ export const EmptyWithNoPrompts: StoryObj<typeof meta> = {
 }
 
 export const Disabled: StoryObj<typeof meta> = { args: { chatEnabled: false } }
+
+export const UnsentMessagePreservation: StoryObj<typeof meta> = {
+    args: {
+        transcript: [],
+        updateEditorStateOnChange: (index: number, state: SerializedPromptEditorValue) => {
+            console.log(`Editor state updated for message ${index}:`, state)
+        },
+    },
+}

vscode/webviews/Chat.tsx

@@ -7,6 +7,7 @@ import type {
     CodyIDE,
     Guardrails,
     PromptString,
+    SerializedPromptEditorValue,
 } from '@sourcegraph/cody-shared'
 import { Transcript, focusLastHumanMessageEditor } from './chat/Transcript'
 import type { VSCodeWrapper } from './utils/VSCodeApi'
@@ -32,6 +33,7 @@ interface ChatboxProps {
     showIDESnippetActions?: boolean
     setView: (view: View) => void
     smartApplyEnabled?: boolean
+    updateEditorStateOnChange: (index: number, state: SerializedPromptEditorValue) => void
 }
 
 export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>> = ({
@@ -45,6 +47,7 @@ export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>
     showIDESnippetActions = true,
     setView,
     smartApplyEnabled,
+    updateEditorStateOnChange,
 }) => {
     const telemetryRecorder = useTelemetryRecorder()
 
@@ -224,6 +227,7 @@ export const Chat: React.FunctionComponent<React.PropsWithChildren<ChatboxProps>
                 postMessage={postMessage}
                 guardrails={guardrails}
                 smartApplyEnabled={smartApplyEnabled}
+                updateEditorStateOnChange={updateEditorStateOnChange}
             />
             {transcript.length === 0 && showWelcomeMessage && (
                 <>

vscode/webviews/CodyPanel.tsx

@@ -1,6 +1,20 @@
-import { type AuthStatus, type ClientCapabilities, CodyIDE } from '@sourcegraph/cody-shared'
+import {
+    type AuthStatus,
+    type ChatMessage,
+    type ClientCapabilities,
+    CodyIDE,
+    type SerializedPromptEditorValue,
+} from '@sourcegraph/cody-shared'
 import type React from 'react'
-import { type ComponentProps, type FunctionComponent, useMemo, useRef } from 'react'
+import {
+    type ComponentProps,
+    type FunctionComponent,
+    useCallback,
+    useEffect,
+    useMemo,
+    useRef,
+    useState,
+} from 'react'
 import type { ConfigurationSubsetForWebview, LocalEnv } from '../src/chat/protocol'
 import styles from './App.module.css'
 import { Chat } from './Chat'
@@ -55,6 +69,39 @@ export const CodyPanel: FunctionComponent<
 }) => {
     const tabContainerRef = useRef<HTMLDivElement>(null)
 
+    const [activeTranscript, setActiveTranscript] = useState<ChatMessage[] | undefined>(transcript)
+    const [storedTranscriptState, setStoredTranscriptState] = useState(transcript)
+
+    // Update the Transcript State for each input box value change.
+    const updateEditorStateOnChange = useCallback(
+        (index: number, newEditorValue: SerializedPromptEditorValue) => {
+            setStoredTranscriptState(prev => {
+                const updated = [...prev]
+                updated[index] = {
+                    ...updated[index],
+                    editorState: newEditorValue.editorState,
+                    speaker: 'human',
+                }
+                return updated
+            })
+        },
+        []
+    )
+
+    // Reset transcripts on new transcript change.
+    useEffect(() => {
+        setActiveTranscript(undefined)
+        setStoredTranscriptState(transcript)
+    }, [transcript])
+
+    // Set the current transcript to the transcript with the last stored editor states when switching to a different tab.
+    // This ensures the editor states are preserved when switching back to the chat tab.
+    useEffect(() => {
+        if (view !== View.Chat) {
+            setActiveTranscript(storedTranscriptState)
+        }
+    }, [view, storedTranscriptState])
+
     return (
         <TabViewContext.Provider value={useMemo(() => ({ view, setView }), [view, setView])}>
             <TabRoot
@@ -76,9 +123,10 @@ export const CodyPanel: FunctionComponent<
                 <TabContainer value={view} ref={tabContainerRef}>
                     {view === View.Chat && (
                         <Chat
+                            transcript={activeTranscript ?? transcript}
+                            updateEditorStateOnChange={updateEditorStateOnChange}
                             chatEnabled={chatEnabled}
                             messageInProgress={messageInProgress}
-                            transcript={transcript}
                             vscodeAPI={vscodeAPI}
                             guardrails={attributionEnabled ? guardrails : undefined}
                             showIDESnippetActions={showIDESnippetActions}

vscode/webviews/chat/Transcript.story.tsx

@@ -36,6 +36,7 @@ const meta: Meta<typeof Transcript> = {
         userInfo: FIXTURE_USER_ACCOUNT_INFO,
         postMessage: () => {},
         chatEnabled: true,
+        updateEditorStateOnChange: () => {},
     } satisfies ComponentProps<typeof Transcript>,
 
     decorators: [

vscode/webviews/chat/Transcript.test.tsx

@@ -15,6 +15,7 @@ const PROPS: Omit<ComponentProps<typeof Transcript>, 'transcript'> = {
     userInfo: FIXTURE_USER_ACCOUNT_INFO,
     chatEnabled: true,
     postMessage: () => {},
+    updateEditorStateOnChange: () => {},
 }
 
 vi.mock('@vscode/webview-ui-toolkit/react', () => ({
@@ -398,4 +399,43 @@ describe('transcriptToInteractionPairs', () => {
             },
         ])
     })
+
+    test('preserves editor state when switching tabs', async () => {
+        const waitForDebounce = () => new Promise(resolve => setTimeout(resolve, 350)) // 350ms to be safe as we debounce on change.
+
+        const humanMessage: ChatMessage = {
+            speaker: 'human',
+            text: ps`Foo`,
+            contextFiles: [],
+        }
+        const assistantMessage: ChatMessage = { speaker: 'assistant', text: ps`Bar` }
+        const updateEditorStateOnChange = vi.fn()
+        const { container, rerender } = render(
+            <Transcript
+                {...PROPS}
+                transcript={[humanMessage, assistantMessage]}
+                updateEditorStateOnChange={updateEditorStateOnChange}
+            />
+        )
+
+        const editor = container.querySelector<EditorHTMLElement>(
+            '[role="row"]:last-child [data-lexical-editor="true"]'
+        )! as EditorHTMLElement
+        editor.textContent = 'Unsent message'
+
+        await typeInEditor(editor, 'Unsent message updated')
+
+        // Wait for the debounce
+        await waitForDebounce()
+
+        rerender(
+            <Transcript
+                {...PROPS}
+                transcript={[humanMessage, assistantMessage]}
+                updateEditorStateOnChange={updateEditorStateOnChange}
+            />
+        )
+
+        expect(editor.textContent).toBe('Unsent message updated')
+    })
 })

vscode/webviews/chat/Transcript.tsx

@@ -42,6 +42,8 @@ interface TranscriptProps {
     insertButtonOnSubmit?: CodeBlockActionsProps['insertButtonOnSubmit']
     smartApply?: CodeBlockActionsProps['smartApply']
     smartApplyEnabled?: boolean
+
+    updateEditorStateOnChange: (index: number, state: SerializedPromptEditorValue) => void
 }
 
 export const Transcript: FC<TranscriptProps> = props => {
@@ -57,6 +59,7 @@ export const Transcript: FC<TranscriptProps> = props => {
         insertButtonOnSubmit,
         smartApply,
         smartApplyEnabled,
+        updateEditorStateOnChange,
     } = props
 
     const interactions = useMemo(
@@ -92,6 +95,7 @@ export const Transcript: FC<TranscriptProps> = props => {
                     )}
                     smartApply={smartApply}
                     smartApplyEnabled={smartApplyEnabled}
+                    updateEditorStateOnChange={updateEditorStateOnChange}
                 />
             ))}
         </div>
@@ -160,6 +164,8 @@ interface TranscriptInteractionProps
     isLastInteraction: boolean
     isLastSentInteraction: boolean
     priorAssistantMessageIsLoading: boolean
+
+    updateEditorStateOnChange: (index: number, state: SerializedPromptEditorValue) => void
 }
 
 const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
@@ -178,6 +184,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         copyButtonOnSubmit,
         smartApply,
         smartApplyEnabled,
+        updateEditorStateOnChange,
     } = props
 
     const [intentResults, setIntentResults] = useState<
@@ -221,6 +228,10 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
         return debounce(async (editorValue: SerializedPromptEditorValue) => {
             setIntentResults(undefined)
 
+            if (editorValue.text) {
+                updateEditorStateOnChange(humanMessage.index, editorValue)
+            }
+
             if (!experimentalOneBoxEnabled) {
                 return
             }
@@ -240,7 +251,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
                     })
             }
         }, 300)
-    }, [experimentalOneBoxEnabled, extensionAPI])
+    }, [humanMessage.index, experimentalOneBoxEnabled, extensionAPI, updateEditorStateOnChange])
 
     const onStop = useCallback(() => {
         getVSCodeAPI().postMessage({