feat(amazonq): support multiple uri in a single diff scheme (aws#7167)

ctlai95 · web-flow · commit 947c471967e0 · 2025-04-25T23:14:33.000-04:00
## Problem The current viewDiff implementation is restricted to only allowing comparing the current file contents. This means it's not possible to view a diff for a change that was applied at a previous point in time. ## Solution - Create a new `DiffContentProvider` that can support showing multiple URIs in the same scheme. The current `ContentProvider` has a limitation that it can only show a single URI. - Refactor the code to use a defined type `ViewDiffMessage` instead of `any` - Change the logic so that diffs are stored in memory instead of a temp file on disk --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts
@@ -59,7 +59,12 @@ import * as jose from 'jose'
 import { AmazonQChatViewProvider } from './webviewProvider'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { amazonQDiffScheme, AmazonQPromptSettings, messages, openUrl } from 'aws-core-vscode/shared'
-import { DefaultAmazonQAppInitContext, messageDispatcher, EditorContentController } from 'aws-core-vscode/amazonq'
+import {
+    DefaultAmazonQAppInitContext,
+    messageDispatcher,
+    EditorContentController,
+    ViewDiffMessage,
+} from 'aws-core-vscode/amazonq'
 import { telemetry, TelemetryBase } from 'aws-core-vscode/telemetry'
 import { isValidResponseError } from './error'
 
@@ -449,17 +454,24 @@ export function registerMessageListeners(
             new vscode.Position(0, 0),
             new vscode.Position(doc.lineCount - 1, doc.lineAt(doc.lineCount - 1).text.length)
         )
-        await ecc.viewDiff(
-            {
-                context: {
-                    activeFileContext: { filePath: params.originalFileUri },
-                    focusAreaContext: { selectionInsideExtendedCodeBlock: entireDocumentSelection },
+        const viewDiffMessage: ViewDiffMessage = {
+            context: {
+                activeFileContext: {
+                    filePath: params.originalFileUri,
+                    fileText: params.originalFileContent ?? '',
+                    fileLanguage: undefined,
+                    matchPolicy: undefined,
+                },
+                focusAreaContext: {
+                    selectionInsideExtendedCodeBlock: entireDocumentSelection,
+                    codeBlock: '',
+                    extendedCodeBlock: '',
+                    names: undefined,
                 },
-                code: params.fileContent ?? '',
             },
-            amazonQDiffScheme,
-            true
-        )
+            code: params.fileContent ?? '',
+        }
+        await ecc.viewDiff(viewDiffMessage, amazonQDiffScheme)
     })
 
     languageClient.onNotification(chatUpdateNotificationType.method, (params: ChatUpdateParams) => {
diff --git a/packages/core/package.json b/packages/core/package.json
@@ -444,7 +444,7 @@
         "@aws/chat-client": "^0.1.4",
         "@aws/chat-client-ui-types": "^0.1.24",
         "@aws/language-server-runtimes": "^0.2.70",
-        "@aws/language-server-runtimes-types": "^0.1.21",
+        "@aws/language-server-runtimes-types": "^0.1.26",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^10.0.2",
         "@types/adm-zip": "^0.4.34",
diff --git a/packages/core/src/amazonq/commons/controllers/contentController.ts b/packages/core/src/amazonq/commons/controllers/contentController.ts
@@ -11,14 +11,19 @@ import { amazonQDiffScheme, amazonQTabSuffix } from '../../../shared/constants'
 import { disposeOnEditorClose } from '../../../shared/utilities/editorUtilities'
 import {
     applyChanges,
-    createTempFileForDiff,
+    createTempUrisForDiff,
     getIndentedCode,
     getSelectionFromRange,
 } from '../../../shared/utilities/textDocumentUtilities'
 import { ToolkitError, getErrorMsg } from '../../../shared/errors'
 import fs from '../../../shared/fs/fs'
 import { extractFileAndCodeSelectionFromMessage } from '../../../shared/utilities/textUtilities'
 import { UserWrittenCodeTracker } from '../../../codewhisperer/tracker/userWrittenCodeTracker'
+import type { ViewDiff } from '../../../codewhispererChat/controllers/chat/model'
+import type { TriggerEvent } from '../../../codewhispererChat/storages/triggerEvents'
+import { DiffContentProvider } from './diffContentProvider'
+
+export type ViewDiffMessage = Pick<ViewDiff, 'code'> & Partial<Pick<TriggerEvent, 'context'>>
 
 export class ContentProvider implements vscode.TextDocumentContentProvider {
     constructor(private uri: vscode.Uri) {}
@@ -155,26 +160,35 @@ export class EditorContentController {
      * isolating them from any other modifications in the original file.
      *
      * @param message the message from Amazon Q chat
+     * @param scheme the URI scheme to use for the diff view
      */
-    public async viewDiff(message: any, scheme: string = amazonQDiffScheme, reverseOrder = false) {
+    public async viewDiff(message: ViewDiffMessage, scheme: string = amazonQDiffScheme) {
         const errorNotification = 'Unable to Open Diff.'
-        const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
+        const { filePath, fileText, selection } = extractFileAndCodeSelectionFromMessage(message)
 
         try {
             if (filePath && message?.code !== undefined && 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 contentProvider = new DiffContentProvider()
                 const disposable = vscode.workspace.registerTextDocumentContentProvider(scheme, contentProvider)
+
+                const [originalFileUri, modifiedFileUri] = await createTempUrisForDiff(
+                    filePath,
+                    fileText,
+                    message,
+                    selection,
+                    scheme,
+                    contentProvider
+                )
+
                 await vscode.commands.executeCommand(
                     'vscode.diff',
-                    ...(reverseOrder ? [uri, originalFileUri] : [originalFileUri, uri]),
+                    originalFileUri,
+                    modifiedFileUri,
                     `${path.basename(filePath)} ${amazonQTabSuffix}`
                 )
 
-                disposeOnEditorClose(uri, disposable)
+                disposeOnEditorClose(originalFileUri, disposable)
             }
         } catch (error) {
             void vscode.window.showInformationMessage(errorNotification)
diff --git a/packages/core/src/amazonq/commons/controllers/diffContentProvider.ts b/packages/core/src/amazonq/commons/controllers/diffContentProvider.ts
@@ -0,0 +1,51 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import vscode from 'vscode'
+import { getLogger } from '../../../shared/logger/logger'
+
+/**
+ * A TextDocumentContentProvider that can handle multiple URIs with the same scheme.
+ * This provider maintains a mapping of URIs to their content.
+ */
+export class DiffContentProvider implements vscode.TextDocumentContentProvider {
+    private contentMap = new Map<string, string>()
+    private _onDidChange = new vscode.EventEmitter<vscode.Uri>()
+
+    public readonly onDidChange = this._onDidChange.event
+
+    /**
+     * Register content for a specific URI
+     * @param uri The URI to register content for
+     * @param content The content to serve for this URI
+     */
+    public registerContent(uri: vscode.Uri, content: string): void {
+        this.contentMap.set(uri.toString(), content)
+        this._onDidChange.fire(uri)
+    }
+
+    /**
+     * Unregister a URI
+     * @param uri The URI to unregister
+     */
+    public unregisterUri(uri: vscode.Uri): void {
+        this.contentMap.delete(uri.toString())
+    }
+
+    /**
+     * Provides the content for a given URI
+     * @param uri The URI to provide content for
+     * @returns The content as a string
+     */
+    public provideTextDocumentContent(uri: vscode.Uri): string {
+        const content = this.contentMap.get(uri.toString())
+
+        if (content === undefined) {
+            getLogger().warn('No content registered for URI: %s', uri.toString())
+            return ''
+        }
+
+        return content
+    }
+}
diff --git a/packages/core/src/amazonq/index.ts b/packages/core/src/amazonq/index.ts
@@ -47,7 +47,7 @@ export * as authConnection from '../auth/connection'
 export * as featureConfig from './webview/generators/featureConfig'
 export * as messageDispatcher from './webview/messages/messageDispatcher'
 import { FeatureContext } from '../shared/featureConfig'
-export { EditorContentController } from './commons/controllers/contentController'
+export { EditorContentController, ViewDiffMessage } from './commons/controllers/contentController'
 
 /**
  * main from createMynahUI is a purely browser dependency. Due to this
diff --git a/packages/core/src/shared/utilities/textDocumentUtilities.ts b/packages/core/src/shared/utilities/textDocumentUtilities.ts
@@ -11,6 +11,8 @@ import { getLogger } from '../logger/logger'
 import fs from '../fs/fs'
 import { ToolkitError } from '../errors'
 import { indent } from './textUtilities'
+import { ViewDiffMessage } from '../../amazonq/commons/controllers/contentController'
+import { DiffContentProvider } from '../../amazonq/commons/controllers/diffContentProvider'
 
 /**
  * Finds occurences of text in a document. Currently only used for highlighting cloudwatchlogs data.
@@ -123,13 +125,13 @@ export async function applyChanges(doc: vscode.TextDocument, range: vscode.Range
  * and applying the proposed changes within the selected range.
  *
  * @param {vscode.Uri} originalFileUri - The URI of the original file.
- * @param {any} message - The message object containing the proposed code changes.
+ * @param {ViewDiffMessage} message - The message object containing the proposed code changes.
  * @param {vscode.Selection} selection - The selection range in the document where the changes are applied.
  * @returns {Promise<vscode.Uri>} - A promise that resolves to the URI of the temporary file.
  */
 export async function createTempFileForDiff(
     originalFileUri: vscode.Uri,
-    message: any,
+    message: ViewDiffMessage,
     selection: vscode.Selection,
     scheme: string
 ): Promise<vscode.Uri> {
@@ -168,6 +170,53 @@ export async function createTempFileForDiff(
     return tempFileUri
 }
 
+/**
+ * Creates temporary URIs for diff comparison and registers their content with a DiffContentProvider.
+ * This approach avoids writing to the file system by keeping content in memory.
+ *
+ * @param filePath The path of the original file (used for naming)
+ * @param fileText Optional content of the original file (if not provided, will be read from filePath)
+ * @param message The message object containing the proposed code changes
+ * @param selection The selection range where changes should be applied
+ * @param scheme The URI scheme to use
+ * @param diffProvider The content provider to register URIs with
+ * @returns A promise that resolves to a tuple of [originalUri, modifiedUri]
+ */
+export async function createTempUrisForDiff(
+    filePath: string,
+    fileText: string | undefined,
+    message: ViewDiffMessage,
+    selection: vscode.Selection,
+    scheme: string,
+    diffProvider: DiffContentProvider
+): Promise<[vscode.Uri, vscode.Uri]> {
+    const originalFile = _path.parse(filePath)
+    const id = Date.now()
+
+    // Create URIs with the custom scheme
+    const originalFileUri = vscode.Uri.parse(`${scheme}:/${originalFile.name}_original-${id}${originalFile.ext}`)
+    const modifiedFileUri = vscode.Uri.parse(`${scheme}:/${originalFile.name}_proposed-${id}${originalFile.ext}`)
+
+    // Get the original content
+    const contentToUse = fileText ?? (await fs.readFileText(filePath))
+
+    // Register the original content
+    diffProvider.registerContent(originalFileUri, contentToUse)
+
+    const indentedCode = getIndentedCodeFromOriginalContent(message, contentToUse, selection)
+    const lines = contentToUse.split('\n')
+
+    // Create the modified content
+    const beforeLines = lines.slice(0, selection.start.line)
+    const afterLines = lines.slice(selection.end.line + 1)
+    const modifiedContent = [...beforeLines, indentedCode, ...afterLines].join('\n')
+
+    // Register the modified content
+    diffProvider.registerContent(modifiedFileUri, modifiedContent)
+
+    return [originalFileUri, modifiedFileUri]
+}
+
 /**
  * Indents the given code based on the current document's indentation at the selection start.
  *
@@ -176,7 +225,7 @@ export async function createTempFileForDiff(
  * @param selection The selection range in the document.
  * @returns The processed code to be applied to the document.
  */
-export function getIndentedCode(message: any, doc: vscode.TextDocument, selection: vscode.Selection) {
+export function getIndentedCode(message: ViewDiffMessage, 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)
 
@@ -187,6 +236,26 @@ export function getIndentedCode(message: any, doc: vscode.TextDocument, selectio
     return indent(message.code, indentation.length)
 }
 
+/**
+ * Indents the given code based on the indentation of the original content at the selection start.
+ *
+ * @param message The message object containing the code.
+ * @param originalContent The original content of the document.
+ * @param selection The selection range in the document.
+ * @returns The processed code to be applied to the document.
+ */
+export function getIndentedCodeFromOriginalContent(
+    message: ViewDiffMessage,
+    originalContent: string,
+    selection: vscode.Selection
+) {
+    const lines = originalContent.split('\n')
+    const selectionStartLine = lines[selection.start.line] || ''
+    const indentMatch = selectionStartLine.match(/^(\s*)/)
+    const indentation = indentMatch ? indentMatch[1] : ''
+    return indent(message.code, indentation.length)
+}
+
 export async function showFile(uri: vscode.Uri) {
     const doc = await vscode.workspace.openTextDocument(uri)
     await vscode.window.showTextDocument(doc, { preview: false })
diff --git a/packages/core/src/shared/utilities/textUtilities.ts b/packages/core/src/shared/utilities/textUtilities.ts
@@ -8,6 +8,7 @@ import * as crypto from 'crypto'
 import * as fs from 'fs' // eslint-disable-line no-restricted-imports
 import { default as stripAnsi } from 'strip-ansi'
 import { getLogger } from '../logger/logger'
+import { ViewDiffMessage } from '../../amazonq/commons/controllers/contentController'
 
 /**
  * Truncates string `s` if it has or exceeds `n` chars.
@@ -268,10 +269,11 @@ export function decodeBase64(base64Str: string): string {
  * @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) {
+export function extractFileAndCodeSelectionFromMessage(message: ViewDiffMessage) {
     const filePath = message?.context?.activeFileContext?.filePath
+    const fileText = message?.context?.activeFileContext?.fileText
     const selection = message?.context?.focusAreaContext?.selectionInsideExtendedCodeBlock as vscode.Selection
-    return { filePath, selection }
+    return { filePath, fileText, selection }
 }
 
 export function matchesPattern(source: string, target: string | RegExp) {
diff --git a/packages/core/src/test/amazonq/commons/controllers/contentController.test.ts b/packages/core/src/test/amazonq/commons/controllers/contentController.test.ts
diff --git a/packages/core/src/test/shared/utilities/textDocumentUtilities.test.ts b/packages/core/src/test/shared/utilities/textDocumentUtilities.test.ts
