feat(amazonq): add button to view diff in IDE (aws#5338)

Add `Accept Diff` & `View Diff` button to Amazon Q for auto generated
code.



https://github.com/user-attachments/assets/2bd6cf19-0649-4355-8ac5-7eb87bfe27fb



## License

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

packages/amazonq/.changes/next-release/Feature-3ebcda93-ca02-4682-b475-dead26e94900.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Add buttons to code blocks to view and accept diffs."
+}

packages/amazonq/src/extension.ts

@@ -20,11 +20,14 @@ import {
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
 import {
+    amazonQDiffScheme,
     DefaultAWSClientBuilder,
     DefaultAwsContext,
     ExtContext,
     RegionProvider,
     Settings,
+    VirtualFileSystem,
+    VirtualMemoryFile,
     activateLogger,
     activateTelemetry,
     env,
@@ -136,6 +139,14 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // Handle Amazon Q Extension un-installation.
     setupUninstallHandler(VSCODE_EXTENSION_ID.amazonq, context.extension.packageJSON.version, context)
 
+    const vfs = new VirtualFileSystem()
+
+    // Register an empty file that's used when a to open a diff
+    vfs.registerProvider(
+        vscode.Uri.from({ scheme: amazonQDiffScheme, path: 'empty' }),
+        new VirtualMemoryFile(new Uint8Array())
+    )
+
     // Hide the Amazon Q tree in toolkit explorer
     await setContext('aws.toolkit.amazonq.dismissed', true)
 

packages/amazonq/test/e2e/amazonq/featureDev.test.ts

@@ -110,7 +110,7 @@ describe('Amazon Q Feature Dev', function () {
 
     beforeEach(() => {
         registerAuthHook('amazonq-test-account')
-        framework = new qTestingFramework('featuredev', true)
+        framework = new qTestingFramework('featuredev', true, [])
         tab = framework.createTab()
     })
 
@@ -135,7 +135,7 @@ describe('Amazon Q Feature Dev', function () {
         it('Does NOT show /dev when feature dev is NOT enabled', () => {
             // The beforeEach registers a framework which accepts requests. If we don't dispose before building a new one we have duplicate messages
             framework.dispose()
-            framework = new qTestingFramework('featuredev', false)
+            framework = new qTestingFramework('featuredev', false, [])
             const tab = framework.createTab()
             const command = tab.findCommand('/dev')
             if (command.length > 0) {

packages/amazonq/test/e2e/amazonq/framework/framework.ts

@@ -12,6 +12,7 @@ import * as vscode from 'vscode'
 import { MynahUI, MynahUIProps } from '@aws/mynah-ui'
 import { DefaultAmazonQAppInitContext, TabType, createMynahUI } from 'aws-core-vscode/amazonq'
 import { Messenger, MessengerOptions } from './messenger'
+import { FeatureContext } from 'aws-core-vscode/shared'
 
 /**
  * Abstraction over Amazon Q to make e2e testing easier
@@ -23,7 +24,7 @@ export class qTestingFramework {
 
     lastEventId: string = ''
 
-    constructor(featureName: TabType, amazonQEnabled: boolean) {
+    constructor(featureName: TabType, amazonQEnabled: boolean, featureConfigsSerialized: [string, FeatureContext][]) {
         /**
          * Instantiate the UI and override the postMessage to publish using the app message
          * publishers directly.
@@ -42,7 +43,8 @@ export class qTestingFramework {
                     appMessagePublisher.publish(message)
                 },
             },
-            amazonQEnabled
+            amazonQEnabled,
+            featureConfigsSerialized
         )
         this.mynahUI = ui.mynahUI
         this.mynahUIProps = (this.mynahUI as any).props

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

@@ -4,8 +4,28 @@
  */
 
 import * as vscode from 'vscode'
+import path from 'path'
 import { Position, TextEditor, window } from 'vscode'
 import { getLogger } from '../../../shared/logger'
+import { amazonQDiffScheme, amazonQTabSuffix } from '../../../shared/constants'
+import { disposeOnEditorClose } from '../../../shared/utilities/editorUtilities'
+import {
+    applyChanges,
+    createTempFileForDiff,
+    getSelectionFromRange,
+} from '../../../shared/utilities/textDocumentUtilities'
+import { extractFileAndCodeSelectionFromMessage, fs, getErrorMsg, getIndentedCode, ToolkitError } from '../../../shared'
+
+class ContentProvider implements vscode.TextDocumentContentProvider {
+    constructor(private uri: vscode.Uri) {}
+
+    provideTextDocumentContent(_uri: vscode.Uri) {
+        return fs.readFileText(this.uri.fsPath)
+    }
+}
+
+const chatDiffCode = 'ChatDiff'
+const ChatDiffError = ToolkitError.named(chatDiffCode)
 
 export class EditorContentController {
     /* *
@@ -52,4 +72,106 @@ export class EditorContentController {
                 )
         }
     }
+
+    /**
+     * Accepts and applies a diff to a file, then closes the associated diff view tab.
+     *
+     * @param {any} message - The message containing diff information.
+     * @returns {Promise<void>} A promise that resolves when the diff is applied and the tab is closed.
+     *
+     * @description
+     * This method performs the following steps:
+     * 1. Extracts file path and selection from the message.
+     * 2. If valid file path, non-empty code, and selection are present:
+     *    a. Opens the document.
+     *    b. Gets the indented code to update.
+     *    c. Applies the changes to the document.
+     *    d. Attempts to close the diff view tab for the file.
+     *
+     * @throws {Error} If there's an issue opening the document or applying changes.
+     */
+    public async acceptDiff(message: any) {
+        const errorNotification = 'Unable to Apply code changes.'
+        const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
+
+        if (filePath && message?.code?.trim().length > 0 && selection) {
+            try {
+                const doc = await vscode.workspace.openTextDocument(filePath)
+
+                const code = getIndentedCode(message, doc, selection)
+                const range = getSelectionFromRange(doc, selection)
+                await applyChanges(doc, range, code)
+
+                // Sets the editor selection from the start of the given range, extending it by the number of lines in the code till the end of the last line
+                const editor = await vscode.window.showTextDocument(doc)
+                editor.selection = new vscode.Selection(
+                    range.start,
+                    new Position(range.start.line + code.split('\n').length, Number.MAX_SAFE_INTEGER)
+                )
+
+                // 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}`) {
+                            const tabClosed = vscode.window.tabGroups.close(tab)
+                            if (!tabClosed) {
+                                getLogger().error(
+                                    '%s: Unable to close the diff view tab for %s',
+                                    chatDiffCode,
+                                    tab.label
+                                )
+                            }
+                        }
+                    })
+                )
+            } catch (error) {
+                void vscode.window.showInformationMessage(errorNotification)
+                const wrappedError = ChatDiffError.chain(error, `Failed to Accept Diff`, { code: chatDiffCode })
+                getLogger().error('%s: Failed to open diff view %s', chatDiffCode, getErrorMsg(wrappedError, true))
+                throw wrappedError
+            }
+        }
+    }
+
+    /**
+     * Displays a diff view comparing proposed changes with the existing file.
+     *
+     * How is diff generated:
+     * 1. Creates a temporary file as a clone of the original file.
+     * 2. Applies the proposed changes to the temporary file within the selected range.
+     * 3. Opens a diff view comparing original file to the temporary file.
+     *
+     * This approach ensures that the diff view only shows the changes proposed by Amazon Q,
+     * isolating them from any other modifications in the original file.
+     *
+     * @param message the message from Amazon Q chat
+     */
+    public async viewDiff(message: any, scheme: string = amazonQDiffScheme) {
+        const errorNotification = 'Unable to Open Diff.'
+        const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
+
+        try {
+            if (filePath && message?.code?.trim().length > 0 && selection) {
+                const originalFileUri = vscode.Uri.file(filePath)
+                const uri = await createTempFileForDiff(originalFileUri, message, selection, scheme)
+
+                // Register content provider and show diff
+                const contentProvider = new ContentProvider(uri)
+                const disposable = vscode.workspace.registerTextDocumentContentProvider(scheme, contentProvider)
+                await vscode.commands.executeCommand(
+                    'vscode.diff',
+                    originalFileUri,
+                    uri,
+                    `${path.basename(filePath)} ${amazonQTabSuffix}`
+                )
+
+                disposeOnEditorClose(uri, disposable)
+            }
+        } catch (error) {
+            void vscode.window.showInformationMessage(errorNotification)
+            const wrappedError = ChatDiffError.chain(error, `Failed to Open Diff View`, { code: chatDiffCode })
+            getLogger().error('%s: Failed to open diff view %s', chatDiffCode, getErrorMsg(wrappedError, true))
+            throw wrappedError
+        }
+    }
 }

packages/core/src/amazonq/index.ts

@@ -28,17 +28,22 @@ export { listCodeWhispererCommandsWalkthrough } from '../codewhisperer/ui/status
 export { focusAmazonQPanel, focusAmazonQPanelKeybinding } from '../codewhispererChat/commands/registerCommands'
 export { TryChatCodeLensProvider, tryChatCodeLensCommand } from '../codewhispererChat/editor/codelens'
 export { createAmazonQUri, openDiff, openDeletedDiff, getOriginalFileUri, getFileDiffUris } from './commons/diff'
+import { FeatureContext } from '../shared'
 
 /**
  * main from createMynahUI is a purely browser dependency. Due to this
  * we need to create a wrapper function that will dynamically execute it
  * while only running on browser instances (like the e2e tests). If we
  * just export it regularly we will get "ReferenceError: self is not defined"
  */
-export function createMynahUI(ideApi: any, amazonQEnabled: boolean) {
+export function createMynahUI(
+    ideApi: any,
+    amazonQEnabled: boolean,
+    featureConfigsSerialized: [string, FeatureContext][]
+) {
     if (typeof window !== 'undefined') {
         const mynahUI = require('./webview/ui/main')
-        return mynahUI.createMynahUI(ideApi, amazonQEnabled)
+        return mynahUI.createMynahUI(ideApi, amazonQEnabled, featureConfigsSerialized)
     }
     throw new Error('Not implemented for node')
 }

packages/core/src/amazonq/util/functionUtils.ts

@@ -0,0 +1,24 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Converts an array of key-value pairs into a Map object.
+ *
+ * @param {[unknown, unknown][]} arr - An array of tuples, where each tuple represents a key-value pair.
+ * @returns {Map<unknown, unknown>} A new Map object created from the input array.
+ *                                  If the conversion fails, an empty Map is returned.
+ *
+ * @example
+ * const array = [['key1', 'value1'], ['key2', 'value2']];
+ * const map = tryNewMap(array);
+ * // map is now a Map object with entries: { 'key1' => 'value1', 'key2' => 'value2' }
+ */
+export function tryNewMap(arr: [unknown, unknown][]) {
+    try {
+        return new Map(arr)
+    } catch (error) {
+        return new Map()
+    }
+}

packages/core/src/amazonq/webview/generators/webViewContent.ts

@@ -6,7 +6,7 @@
 import path from 'path'
 import { Uri, Webview } from 'vscode'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
-import { globals } from '../../../shared'
+import { FeatureConfigProvider, FeatureContext, globals } from '../../../shared'
 
 export class WebViewContentGenerator {
     public async generate(extensionURI: Uri, webView: Webview): Promise<string> {
@@ -45,14 +45,23 @@ export class WebViewContentGenerator {
             Uri.joinPath(globals.context.extensionUri, 'resources', 'css', 'amazonq-webview.css')
         )
 
+        let featureConfigs = new Map<string, FeatureContext>()
+        try {
+            await FeatureConfigProvider.instance.fetchFeatureConfigs()
+            featureConfigs = FeatureConfigProvider.getFeatureConfigs()
+        } catch (error) {
+            // eslint-disable-next-line aws-toolkits/no-console-log
+            console.error('Error fetching feature configs:', error)
+        }
+
         return `
         <script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
         <link rel="stylesheet" href="${cssEntrypoint.toString()}">
         <script type="text/javascript">
             const init = () => {
                 createMynahUI(acquireVsCodeApi(), ${
                     (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
-                });
+                },${JSON.stringify(Array.from(featureConfigs.entries()))});
             }
     </script>
         `

packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts

@@ -18,7 +18,7 @@ interface ChatPayload {
 export interface ConnectorProps {
     sendMessageToExtension: (message: ExtensionMessage) => void
     onMessageReceived?: (tabID: string, messageData: any, needToShowAPIDocsTab: boolean) => void
-    onChatAnswerReceived?: (tabID: string, message: CWCChatItem) => void
+    onChatAnswerReceived?: (tabID: string, message: CWCChatItem, messageData: any) => void
     onCWCContextCommandMessage: (message: CWCChatItem, command?: string) => string | undefined
     onError: (tabID: string, message: string, title: string) => void
     onWarning: (tabID: string, message: string, title: string) => void
@@ -308,7 +308,7 @@ export class Connector {
                     content: messageData.relatedSuggestions,
                 }
             }
-            this.onChatAnswerReceived(messageData.tabID, answer)
+            this.onChatAnswerReceived(messageData.tabID, answer, messageData)
 
             // Exit the function if we received an answer from AI
             if (
@@ -336,7 +336,7 @@ export class Connector {
                           }
                         : undefined,
             }
-            this.onChatAnswerReceived(messageData.tabID, answer)
+            this.onChatAnswerReceived(messageData.tabID, answer, messageData)
 
             return
         }
@@ -347,13 +347,17 @@ export class Connector {
             return
         }
 
-        this.onChatAnswerReceived(messageData.tabID, {
-            type: ChatItemType.ANSWER,
-            messageId: messageData.triggerID,
-            body: messageData.message,
-            followUp: this.followUpGenerator.generateAuthFollowUps('cwc', messageData.authType),
-            canBeVoted: false,
-        })
+        this.onChatAnswerReceived(
+            messageData.tabID,
+            {
+                type: ChatItemType.ANSWER,
+                messageId: messageData.triggerID,
+                body: messageData.message,
+                followUp: this.followUpGenerator.generateAuthFollowUps('cwc', messageData.authType),
+                canBeVoted: false,
+            },
+            messageData
+        )
 
         return
     }