feat: support at-mentions context on edits

lib/ui/src/Chat.tsx

@@ -303,7 +303,11 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
                 : messageAtIndex?.text
             if (inputText) {
                 setFormInput(inputText)
+                if (messageAtIndex.contextFiles) {
+                    useOldChatMessageContext(messageAtIndex.contextFiles)
+                }
             }
+            // move focus back to chatbox
             setInputFocus(true)
         },
         [messageBeingEdited, setFormInput, setMessageBeingEdited, transcript]
@@ -345,6 +349,15 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
         return `@${displayPath(contextFile.uri)}${range}${symbolName}`
     }
 
+    // Add old context files from the transcript to the map
+    const useOldChatMessageContext = (oldContextFiles: ContextFile[]) => {
+        const contextFilesMap = new Map<string, ContextFile>()
+        for (const file of oldContextFiles) {
+            contextFilesMap.set(getContextFileDisplayText(file), file)
+        }
+        setChatContextFiles(contextFilesMap)
+    }
+
     /**
      * Callback function called when a chat context file is selected from the context selector.
      * This updates the chat input with the selected file context.
@@ -693,11 +706,7 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
                     } else {
                         setFormInput(newHistoryInput.inputText)
                         // chatContextFiles uses a map but history only stores a simple array.
-                        const contextFilesMap = new Map<string, ContextFile>()
-                        for (const file of newHistoryInput.inputContextFiles) {
-                            contextFilesMap.set(getContextFileDisplayText(file), file)
-                        }
-                        setChatContextFiles(contextFilesMap)
+                        useOldChatMessageContext(newHistoryInput.inputContextFiles)
                     }
                 }
             }
@@ -718,6 +727,7 @@ export const Chat: React.FunctionComponent<ChatProps> = ({
             onChatContextSelected,
             enableNewChatMode,
             resetContextSelection,
+            useOldChatMessageContext,
         ]
     )
 

vscode/test/e2e/chat-edits.test.ts

@@ -1,7 +1,7 @@
 import { expect } from '@playwright/test'
 
 import { sidebarSignin } from './common'
-import { test } from './helpers'
+import { test, withPlatformSlashes } from './helpers'
 
 const isPlatform = (platform: string) => process.platform === platform
 const isMac = isPlatform('darwin')
@@ -104,34 +104,26 @@ test('editing follow-up messages in chat view', async ({ page, sidebar }) => {
     await expect(cancelEditButton).not.toBeVisible()
     await expect(chatInput).toBeEmpty()
     await expect(chatFrame.getByText('Explain @Main.java')).toBeVisible()
-    await chatInput.press('Escape')
+
+    // Add a new at-file to an old messages
+    await chatInput.press(`${osKey}+k`)
+    await chatInput.type('and @vgo', { delay: 50 })
+    await chatInput.press('Tab')
+    await expect(chatInput).toHaveValue(
+        withPlatformSlashes('Explain @Main.java and @lib/batches/env/var.go ')
+    )
+    await chatInput.press('Enter')
+    // both main.java and var.go should be used
+    await expect(chatFrame.getByText(/Context: 2 files/)).toBeVisible()
+    await chatFrame.getByText(/Context: 2 files/).click()
+    await expect(chatFrame.getByRole('button', { name: 'Main.java' })).toBeVisible()
+    await expect(
+        chatFrame.getByRole('button', { name: withPlatformSlashes('lib/batches/env/var.go') })
+    ).toBeVisible()
 
     // Meta+/ also creates a new chat session
     await chatInput.press(`${osKey}+/`)
     await expect(chatFrame.getByText('The End')).not.toBeVisible()
 
-    // TODO (bee) - update after switching to the new keybinding for toggling "New Chat Mode"
     await expect(startNewChatButton).not.toBeVisible()
-    // // "MetaKey(MacOS)/Control" + "Shift" to toggle "New Chat Mode" on and off
-    // // When it's on, the submit button will be replaced with "Start New Chat" button
-    // await expect(submitMessageButton).toBeVisible()
-    // await chatInput.press(`${osKey}+Shift`)
-    // await expect(submitMessageButton).not.toBeVisible()
-    // await expect(startNewChatButton).toBeVisible()
-    // await chatInput.press(`${osKey}+Shift`)
-    // await expect(startNewChatButton).not.toBeVisible()
-    // await chatInput.press(`${osKey}+Shift`)
-    // await expect(startNewChatButton).toBeVisible()
-
-    // // With "New Chat Mode" enabled, submit a new message to start a new chat
-    // // The new message should be "The End"
-    // // And the last message should not be visible anymore
-    // await chatInput.fill('The End')
-    // await chatInput.press('Enter')
-    // await expect(chatFrame.getByText('The End')).toBeVisible()
-    // await expect(chatFrame.getByText('Explain @Main.java')).not.toBeVisible()
-
-    // // Meta+/ also creates a new chat session
-    // await chatInput.press(`${osKey}+/`)
-    // await expect(chatFrame.getByText('The End')).not.toBeVisible()
 })