feat: add acceptDiff button to chat

Author: ivikash

packages/amazonq/src/extension.ts

@@ -20,7 +20,7 @@ import {
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
 import {
-    amazonQDiffScheme,
+    AMAZON_Q_DIFF_SCHEME,
     DefaultAWSClientBuilder,
     DefaultAwsContext,
     ExtContext,
@@ -140,7 +140,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
 
     // Register an empty file that's used when a to open a diff
     vfs.registerProvider(
-        vscode.Uri.from({ scheme: amazonQDiffScheme, path: 'empty' }),
+        vscode.Uri.from({ scheme: AMAZON_Q_DIFF_SCHEME, path: 'empty' }),
         new VirtualMemoryFile(new Uint8Array())
     )
 

packages/core/src/amazonq/commons/controllers/contentController.ts

@@ -6,10 +6,24 @@
 import * as vscode from 'vscode'
 import fs from 'fs-extra'
 import path from 'path'
-import { commands, Position, TextEditor, Uri, window, workspace } from 'vscode'
+import { Position, TextEditor, window } from 'vscode'
 import { getLogger } from '../../../shared/logger'
-import { amazonQDiffScheme, getNonexistentFilename, tempDirPath } from '../../../shared'
+import { AMAZON_Q_DIFF_SCHEME, amazonQTabSuffix } from '../../../shared/constants'
+import { disposeOnEditorClose } from '../../../shared/utilities/editorUtilities'
+import {
+    applyChanges,
+    createTempFileForDiff,
+    getSelectionFromRange,
+} from '../../../shared/utilities/textDocumentUtilities'
+import { extractFileAndCodeSelectionFromMessage, getIndentedCode } from '../diff'
 
+class ContentProvider implements vscode.TextDocumentContentProvider {
+    constructor(private tempFileUri: vscode.Uri) {}
+
+    provideTextDocumentContent(_uri: vscode.Uri) {
+        return fs.readFileSync(this.tempFileUri.fsPath, 'utf-8')
+    }
+}
 export class EditorContentController {
     /* *
      *  Insert the Amazon Q chat written code to the cursor position
@@ -56,6 +70,31 @@ export class EditorContentController {
         }
     }
 
+    /**
+     * Accept code changes received by Amazon Q
+     *
+     * @param message the message from Amazon Q chat
+     */
+    public async acceptDiff(message: any) {
+        const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
+
+        if (filePath && message?.code?.trim().length > 0 && selection) {
+            const doc = await vscode.workspace.openTextDocument(filePath)
+            const codeToUpdate = getIndentedCode(message, doc, selection)
+            const range = getSelectionFromRange(doc, selection)
+            await applyChanges(doc, range, codeToUpdate)
+
+            // If vscode.diff is open for the filePath then close it.
+            vscode.window.tabGroups.all.flatMap(({ tabs }) =>
+                tabs.map((tab) => {
+                    if (tab.label === `${path.basename(filePath)} ${amazonQTabSuffix}`) {
+                        void vscode.window.tabGroups.close(tab)
+                    }
+                })
+            )
+        }
+    }
+
     /**
      * Displays a diff view comparing proposed changes with the existing file.
      *
@@ -70,63 +109,26 @@ export class EditorContentController {
      * @param message the message from Amazon Q chat
      */
     public async viewDiff(message: any) {
-        const { filePath } = message?.context?.activeFileContext || {}
-        const selection = message?.context?.focusAreaContext?.selectionInsideExtendedCodeBlock as vscode.Selection
+        const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
 
         if (filePath && message?.code?.trim().length > 0 && selection) {
-            const id = Date.now()
-
-            const originalFileUri = Uri.file(filePath)
-            const fileNameWithExtension = path.basename(originalFileUri.path)
-
-            const fileName = path.parse(fileNameWithExtension).name
-            const fileExtension = path.extname(fileNameWithExtension)
-
-            // Create a new file in the temp directory
-            const tempFile = await getNonexistentFilename(tempDirPath, `${fileName}_proposed-${id}`, fileExtension, 99)
-            const tempFilePath = path.join(tempDirPath, tempFile)
-
-            // Create a new URI for the temp file
-            const diffScheme = amazonQDiffScheme
-            const tempFileUri = Uri.parse(`${diffScheme}:${tempFilePath}`)
-
-            // Write the initial code to the temp file
-            fs.writeFileSync(tempFilePath, fs.readFileSync(originalFileUri.fsPath, 'utf-8'))
-            const doc = await workspace.openTextDocument(tempFileUri.path)
-
-            // Apply the edit to the temp file
-            const edit = new vscode.WorkspaceEdit()
-            edit.replace(doc.uri, selection, message.code)
+            const originalFileUri = vscode.Uri.file(filePath)
+            const tempFileUri = await createTempFileForDiff(originalFileUri, message, selection)
 
-            const successfulEdit = await workspace.applyEdit(edit)
-            if (successfulEdit) {
-                await doc.save()
-                fs.writeFileSync(tempFilePath, doc.getText())
-            }
-            class ContentProvider {
-                provideTextDocumentContent(_uri: Uri) {
-                    return fs.readFileSync(tempFilePath, 'utf-8')
-                }
-            }
-            const contentProvider = new ContentProvider()
-            const disposable = workspace.registerTextDocumentContentProvider(diffScheme, contentProvider)
-
-            await commands.executeCommand(
+            // Register content provider and show diff
+            const contentProvider = new ContentProvider(tempFileUri)
+            const disposable = vscode.workspace.registerTextDocumentContentProvider(
+                AMAZON_Q_DIFF_SCHEME,
+                contentProvider
+            )
+            await vscode.commands.executeCommand(
                 'vscode.diff',
                 originalFileUri,
                 tempFileUri,
-                `${fileName}${fileExtension} (Generated by Amazon Q)`
+                `${path.basename(filePath)} ${amazonQTabSuffix}`
             )
 
-            vscode.window.onDidChangeVisibleTextEditors(() => {
-                if (
-                    !vscode.window.visibleTextEditors.some(
-                        (editor) => editor.document.uri.toString() === tempFileUri.toString()
-                    )
-                ) {
-                    disposable.dispose()
-                }
-            })
+            disposeOnEditorClose(tempFileUri, disposable)
         }
     }
 }

packages/core/src/amazonq/commons/diff.ts

@@ -3,9 +3,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { existsSync } from 'fs'
 import * as vscode from 'vscode'
+import fs from 'fs'
 import { featureDevScheme } from '../../amazonqFeatureDev/constants'
+import { indent } from '../../shared'
 
 export async function openDiff(leftPath: string, rightPath: string, tabId: string) {
     const { left, right } = getFileDiffUris(leftPath, rightPath, tabId)
@@ -18,7 +19,7 @@ export async function openDeletedDiff(filePath: string, name: string, tabId: str
 }
 
 export function getOriginalFileUri(fullPath: string, tabId: string) {
-    return existsSync(fullPath) ? vscode.Uri.file(fullPath) : createAmazonQUri('empty', tabId)
+    return fs.existsSync(fullPath) ? vscode.Uri.file(fullPath) : createAmazonQUri('empty', tabId)
 }
 
 export function getFileDiffUris(leftPath: string, rightPath: string, tabId: string) {
@@ -32,3 +33,34 @@ export function createAmazonQUri(path: string, tabId: string) {
     // TODO change the featureDevScheme to a more general amazon q scheme
     return vscode.Uri.from({ scheme: featureDevScheme, path, query: `tabID=${tabId}` })
 }
+
+/**
+ * Extracts the file path and selection context from the message.
+ *
+ * @param {any} message - The message object containing the file and selection context.
+ * @returns {Object} - An object with `filePath` and `selection` properties.
+ */
+export function extractFileAndCodeSelectionFromMessage(message: any) {
+    const filePath = message?.context?.activeFileContext?.filePath
+    const selection = message?.context?.focusAreaContext?.selectionInsideExtendedCodeBlock as vscode.Selection
+    return { filePath, selection }
+}
+
+/**
+ * Indents the given code based on the current document's indentation at the selection start.
+ *
+ * @param {any} message - The message object containing the code.
+ * @param {vscode.TextDocument} doc - The VSCode document where the code is applied.
+ * @param {vscode.Selection} selection - The selection range in the document.
+ * @returns {string} - The processed code to be applied to the document.
+ */
+export function getIndentedCode(message: any, doc: vscode.TextDocument, selection: vscode.Selection) {
+    const indentRange = new vscode.Range(new vscode.Position(selection.start.line, 0), selection.active)
+    let indentation = doc.getText(indentRange)
+
+    if (indentation.trim().length !== 0) {
+        indentation = ' '.repeat(indentation.length - indentation.trimStart().length)
+    }
+
+    return indent(message.code, indentation.length)
+}

packages/core/src/amazonq/webview/ui/commands.ts

@@ -16,6 +16,7 @@ type MessageCommand =
     | 'open-diff'
     | 'code_was_copied_to_clipboard'
     | 'insert_code_at_cursor_position'
+    | 'accept_diff'
     | 'view_diff'
     | 'stop-response'
     | 'trigger-tabID-received'

packages/core/src/amazonq/webview/ui/connector.ts

@@ -250,6 +250,35 @@ export class Connector {
         }
     }
 
+    onAcceptDiff = (
+        tabId: string,
+        messageId: string,
+        actionId: string,
+        data?: string,
+        code?: string,
+        type?: CodeSelectionType,
+        referenceTrackerInformation?: ReferenceTrackerInformation[],
+        eventId?: string,
+        codeBlockIndex?: number,
+        totalCodeBlocks?: number
+    ) => {
+        const tabType = this.tabsStorage.getTab(tabId)?.type
+        this.sendMessageToExtension({
+            tabType,
+            tabID: tabId,
+            command: 'accept_diff',
+            messageId,
+            actionId,
+            data,
+            code,
+            type,
+            referenceTrackerInformation,
+            eventId,
+            codeBlockIndex,
+            totalCodeBlocks,
+        })
+    }
+
     onViewDiff = (
         tabId: string,
         messageId: string,

packages/core/src/amazonq/webview/ui/main.ts

@@ -80,7 +80,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
     // eslint-disable-next-line prefer-const
     let messageController: MessageController
 
-    function shouldDisplayViewDiff(messageData: any) {
+    function shouldDisplayDiff(messageData: any) {
         const tab = tabsStorage.getTab(messageData?.tabID || '')
         const allowedCommands = [
             'aws.amazonq.refactorCode',
@@ -242,12 +242,19 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
                     ...item,
                     messageId: item.messageId,
                     codeBlockActions: {
-                        ...(shouldDisplayViewDiff(messageData) // update the condition for fix, refactor etc.
+                        ...(shouldDisplayDiff(messageData)
                             ? {
+                                  'insert-to-cursor': undefined,
+                                  'accept-diff': {
+                                      id: 'accept-diff',
+                                      label: 'Apply Diff',
+                                      icon: MynahIcons.OK_CIRCLED,
+                                      data: messageData,
+                                  },
                                   'view-diff': {
                                       id: 'view-diff',
                                       label: 'View Diff',
-                                      icon: MynahIcons.OK_CIRCLED,
+                                      icon: MynahIcons.EYE,
                                       data: messageData,
                                   },
                               }
@@ -483,19 +490,37 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
             codeBlockIndex?: number,
             totalCodeBlocks?: number
         ) => {
-            if (actionId === 'view-diff') {
-                connector.onViewDiff(
-                    tabId,
-                    messageId,
-                    actionId,
-                    data,
-                    code,
-                    type,
-                    referenceTrackerInformation,
-                    eventId,
-                    codeBlockIndex,
-                    totalCodeBlocks
-                )
+            switch (actionId) {
+                case 'accept-diff':
+                    connector.onAcceptDiff(
+                        tabId,
+                        messageId,
+                        actionId,
+                        data,
+                        code,
+                        type,
+                        referenceTrackerInformation,
+                        eventId,
+                        codeBlockIndex,
+                        totalCodeBlocks
+                    )
+                    break
+                case 'view-diff':
+                    connector.onViewDiff(
+                        tabId,
+                        messageId,
+                        actionId,
+                        data,
+                        code,
+                        type,
+                        referenceTrackerInformation,
+                        eventId,
+                        codeBlockIndex,
+                        totalCodeBlocks
+                    )
+                    break
+                default:
+                    break
             }
         },
         onCodeInsertToCursorPosition: connector.onCodeInsertToCursorPosition,

packages/core/src/codewhispererChat/app.ts

@@ -25,6 +25,7 @@ import {
     TabCreatedMessage,
     TriggerTabIDReceived,
     UIFocusMessage,
+    AcceptDiff,
 } from './controllers/chat/model'
 import { EditorContextCommand, registerCommands } from './commands/registerCommands'
 
@@ -35,6 +36,7 @@ export function init(appContext: AmazonQAppInitContext) {
         processTabClosedMessage: new EventEmitter<TabClosedMessage>(),
         processTabChangedMessage: new EventEmitter<TabChangedMessage>(),
         processInsertCodeAtCursorPosition: new EventEmitter<InsertCodeAtCursorPosition>(),
+        processAcceptDiff: new EventEmitter<AcceptDiff>(),
         processViewDiff: new EventEmitter<ViewDiff>(),
         processCopyCodeToClipboard: new EventEmitter<CopyCodeToClipboard>(),
         processContextMenuCommand: new EventEmitter<EditorContextCommand>(),
@@ -64,6 +66,7 @@ export function init(appContext: AmazonQAppInitContext) {
         processInsertCodeAtCursorPosition: new MessageListener<InsertCodeAtCursorPosition>(
             cwChatControllerEventEmitters.processInsertCodeAtCursorPosition
         ),
+        processAcceptDiff: new MessageListener<AcceptDiff>(cwChatControllerEventEmitters.processAcceptDiff),
         processViewDiff: new MessageListener<ViewDiff>(cwChatControllerEventEmitters.processViewDiff),
         processCopyCodeToClipboard: new MessageListener<CopyCodeToClipboard>(
             cwChatControllerEventEmitters.processCopyCodeToClipboard
@@ -111,6 +114,7 @@ export function init(appContext: AmazonQAppInitContext) {
         processInsertCodeAtCursorPosition: new MessagePublisher<InsertCodeAtCursorPosition>(
             cwChatControllerEventEmitters.processInsertCodeAtCursorPosition
         ),
+        processAcceptDiff: new MessagePublisher<AcceptDiff>(cwChatControllerEventEmitters.processAcceptDiff),
         processViewDiff: new MessagePublisher<ViewDiff>(cwChatControllerEventEmitters.processViewDiff),
         processCopyCodeToClipboard: new MessagePublisher<CopyCodeToClipboard>(
             cwChatControllerEventEmitters.processCopyCodeToClipboard