Skip to content

Commit bf8f4dd

Browse files
committed
feat(amazonq): add button to view and accept diff
1 parent 401fc1e commit bf8f4dd

File tree

28 files changed

+717
-56
lines changed

28 files changed

+717
-56
lines changed
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "Add buttons to code blocks to view and accept diffs."
4+
}

packages/amazonq/src/extension.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,14 @@ import {
2020
import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
2121
import { CommonAuthWebview } from 'aws-core-vscode/login'
2222
import {
23+
amazonQDiffScheme,
2324
DefaultAWSClientBuilder,
2425
DefaultAwsContext,
2526
ExtContext,
2627
RegionProvider,
2728
Settings,
29+
VirtualFileSystem,
30+
VirtualMemoryFile,
2831
activateLogger,
2932
activateTelemetry,
3033
env,
@@ -133,6 +136,14 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
133136
// Handle Amazon Q Extension un-installation.
134137
setupUninstallHandler(VSCODE_EXTENSION_ID.amazonq, context.extension.packageJSON.version, context)
135138

139+
const vfs = new VirtualFileSystem()
140+
141+
// Register an empty file that's used when a to open a diff
142+
vfs.registerProvider(
143+
vscode.Uri.from({ scheme: amazonQDiffScheme, path: 'empty' }),
144+
new VirtualMemoryFile(new Uint8Array())
145+
)
146+
136147
// Hide the Amazon Q tree in toolkit explorer
137148
await setContext('aws.toolkit.amazonq.dismissed', true)
138149

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ describe('Amazon Q Feature Dev', function () {
110110

111111
beforeEach(() => {
112112
registerAuthHook('amazonq-test-account')
113-
framework = new qTestingFramework('featuredev', true)
113+
framework = new qTestingFramework('featuredev', true, [])
114114
tab = framework.createTab()
115115
})
116116

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

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import * as vscode from 'vscode'
1212
import { MynahUI, MynahUIProps } from '@aws/mynah-ui'
1313
import { DefaultAmazonQAppInitContext, TabType, createMynahUI } from 'aws-core-vscode/amazonq'
1414
import { Messenger, MessengerOptions } from './messenger'
15+
import { FeatureContext } from 'aws-core-vscode/shared'
1516

1617
/**
1718
* Abstraction over Amazon Q to make e2e testing easier
@@ -23,7 +24,7 @@ export class qTestingFramework {
2324

2425
lastEventId: string = ''
2526

26-
constructor(featureName: TabType, amazonQEnabled: boolean) {
27+
constructor(featureName: TabType, amazonQEnabled: boolean, featureConfigsSerialized: [string, FeatureContext][]) {
2728
/**
2829
* Instantiate the UI and override the postMessage to publish using the app message
2930
* publishers directly.
@@ -42,7 +43,8 @@ export class qTestingFramework {
4243
appMessagePublisher.publish(message)
4344
},
4445
},
45-
amazonQEnabled
46+
amazonQEnabled,
47+
featureConfigsSerialized
4648
)
4749
this.mynahUI = ui.mynahUI
4850
this.mynahUIProps = (this.mynahUI as any).props

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

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,29 @@
44
*/
55

66
import * as vscode from 'vscode'
7+
import * as fs from 'fs'
8+
import path from 'path'
79
import { Position, TextEditor, window } from 'vscode'
810
import { getLogger } from '../../../shared/logger'
11+
import { amazonQDiffScheme, amazonQTabSuffix } from '../../../shared/constants'
12+
import { disposeOnEditorClose } from '../../../shared/utilities/editorUtilities'
13+
import {
14+
applyChanges,
15+
createTempFileForDiff,
16+
getSelectionFromRange,
17+
} from '../../../shared/utilities/textDocumentUtilities'
18+
import { extractFileAndCodeSelectionFromMessage, getIndentedCode, ToolkitError } from '../../../shared'
19+
20+
class ContentProvider implements vscode.TextDocumentContentProvider {
21+
constructor(private uri: vscode.Uri) {}
22+
23+
provideTextDocumentContent(_uri: vscode.Uri) {
24+
return fs.readFileSync(this.uri.fsPath, 'utf-8')
25+
}
26+
}
27+
28+
const chatDiffCode = 'ChatDiff'
29+
const ChatDiffError = ToolkitError.named(chatDiffCode)
930

1031
export class EditorContentController {
1132
/* *
@@ -52,4 +73,98 @@ export class EditorContentController {
5273
)
5374
}
5475
}
76+
77+
/**
78+
* Accepts and applies a diff to a file, then closes the associated diff view tab.
79+
*
80+
* @param {any} message - The message containing diff information.
81+
* @returns {Promise<void>} A promise that resolves when the diff is applied and the tab is closed.
82+
*
83+
* @description
84+
* This method performs the following steps:
85+
* 1. Extracts file path and selection from the message.
86+
* 2. If valid file path, non-empty code, and selection are present:
87+
* a. Opens the document.
88+
* b. Gets the indented code to update.
89+
* c. Applies the changes to the document.
90+
* d. Attempts to close the diff view tab for the file.
91+
*
92+
* @throws {Error} If there's an issue opening the document or applying changes.
93+
*/
94+
public async acceptDiff(message: any) {
95+
const errorNotification = 'Unable to Apply code changes.'
96+
const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
97+
98+
if (filePath && message?.code?.trim().length > 0 && selection) {
99+
try {
100+
const doc = await vscode.workspace.openTextDocument(filePath) // can throw
101+
102+
const code = getIndentedCode(message, doc, selection)
103+
const range = getSelectionFromRange(doc, selection)
104+
await applyChanges(doc, range, code)
105+
106+
// 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
107+
const editor = await vscode.window.showTextDocument(doc)
108+
editor.selection = new vscode.Selection(
109+
range.start,
110+
new Position(range.start.line + code.split('\n').length, Number.MAX_SAFE_INTEGER)
111+
)
112+
113+
// If vscode.diff is open for the filePath then close it.
114+
vscode.window.tabGroups.all.flatMap(({ tabs }) =>
115+
tabs.map((tab) => {
116+
if (tab.label === `${path.basename(filePath)} ${amazonQTabSuffix}`) {
117+
const tabClosed = vscode.window.tabGroups.close(tab)
118+
if (!tabClosed) {
119+
getLogger().error('Diff: Unable to close the diff view tab for %s', tab.label)
120+
}
121+
}
122+
})
123+
)
124+
} catch (error) {
125+
void vscode.window.showInformationMessage(errorNotification)
126+
ChatDiffError.chain(error, `Failed to Accept Diff`, { code: chatDiffCode })
127+
}
128+
}
129+
}
130+
131+
/**
132+
* Displays a diff view comparing proposed changes with the existing file.
133+
*
134+
* How is diff generated:
135+
* 1. Creates a temporary file as a clone of the original file.
136+
* 2. Applies the proposed changes to the temporary file within the selected range.
137+
* 3. Opens a diff view comparing original file to the temporary file.
138+
*
139+
* This approach ensures that the diff view only shows the changes proposed by Amazon Q,
140+
* isolating them from any other modifications in the original file.
141+
*
142+
* @param message the message from Amazon Q chat
143+
*/
144+
public async viewDiff(message: any, scheme: string = amazonQDiffScheme) {
145+
const errorNotification = 'Unable to Open Diff.'
146+
const { filePath, selection } = extractFileAndCodeSelectionFromMessage(message)
147+
148+
try {
149+
if (filePath && message?.code?.trim().length > 0 && selection) {
150+
const originalFileUri = vscode.Uri.file(filePath)
151+
const uri = await createTempFileForDiff(originalFileUri, message, selection, scheme)
152+
153+
// Register content provider and show diff
154+
const contentProvider = new ContentProvider(uri)
155+
const disposable = vscode.workspace.registerTextDocumentContentProvider(scheme, contentProvider)
156+
await vscode.commands.executeCommand(
157+
'vscode.diff',
158+
originalFileUri,
159+
uri,
160+
`${path.basename(filePath)} ${amazonQTabSuffix}`
161+
)
162+
163+
disposeOnEditorClose(uri, disposable)
164+
}
165+
} catch (error) {
166+
void vscode.window.showInformationMessage(errorNotification)
167+
ChatDiffError.chain(error, `Failed to Open Diff View`, { code: chatDiffCode })
168+
}
169+
}
55170
}

packages/core/src/amazonq/index.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,22 @@ export { listCodeWhispererCommandsWalkthrough } from '../codewhisperer/ui/status
2828
export { focusAmazonQPanel, focusAmazonQPanelKeybinding } from '../codewhispererChat/commands/registerCommands'
2929
export { TryChatCodeLensProvider, tryChatCodeLensCommand } from '../codewhispererChat/editor/codelens'
3030
export { createAmazonQUri, openDiff, openDeletedDiff, getOriginalFileUri, getFileDiffUris } from './commons/diff'
31+
import { FeatureContext } from '../shared'
3132

3233
/**
3334
* main from createMynahUI is a purely browser dependency. Due to this
3435
* we need to create a wrapper function that will dynamically execute it
3536
* while only running on browser instances (like the e2e tests). If we
3637
* just export it regularly we will get "ReferenceError: self is not defined"
3738
*/
38-
export function createMynahUI(ideApi: any, amazonQEnabled: boolean) {
39+
export function createMynahUI(
40+
ideApi: any,
41+
amazonQEnabled: boolean,
42+
featureConfigsSerialized: [string, FeatureContext][]
43+
) {
3944
if (typeof window !== 'undefined') {
4045
const mynahUI = require('./webview/ui/main')
41-
return mynahUI.createMynahUI(ideApi, amazonQEnabled)
46+
return mynahUI.createMynahUI(ideApi, amazonQEnabled, featureConfigsSerialized)
4247
}
4348
throw new Error('Not implemented for node')
4449
}
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
/**
7+
* Converts an array of key-value pairs into a Map object.
8+
*
9+
* @param {[unknown, unknown][]} arr - An array of tuples, where each tuple represents a key-value pair.
10+
* @returns {Map<unknown, unknown>} A new Map object created from the input array.
11+
* If the conversion fails, an empty Map is returned.
12+
*
13+
* @example
14+
* const array = [['key1', 'value1'], ['key2', 'value2']];
15+
* const map = tryNewMap(array);
16+
* // map is now a Map object with entries: { 'key1' => 'value1', 'key2' => 'value2' }
17+
*/
18+
export function tryNewMap(arr: [unknown, unknown][]) {
19+
try {
20+
return new Map(arr)
21+
} catch (error) {
22+
return new Map()
23+
}
24+
}

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import path from 'path'
77
import { Uri, Webview } from 'vscode'
88
import { AuthUtil } from '../../../codewhisperer/util/authUtil'
9-
import { globals } from '../../../shared'
9+
import { FeatureConfigProvider, FeatureContext, globals } from '../../../shared'
1010

1111
export class WebViewContentGenerator {
1212
public async generate(extensionURI: Uri, webView: Webview): Promise<string> {
@@ -45,14 +45,23 @@ export class WebViewContentGenerator {
4545
Uri.joinPath(globals.context.extensionUri, 'resources', 'css', 'amazonq-webview.css')
4646
)
4747

48+
let featureConfigs = new Map<string, FeatureContext>()
49+
try {
50+
await FeatureConfigProvider.instance.fetchFeatureConfigs()
51+
featureConfigs = FeatureConfigProvider.getFeatureConfigs()
52+
} catch (error) {
53+
// eslint-disable-next-line aws-toolkits/no-console-log
54+
console.error('Error fetching feature configs:', error)
55+
}
56+
4857
return `
4958
<script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
5059
<link rel="stylesheet" href="${cssEntrypoint.toString()}">
5160
<script type="text/javascript">
5261
const init = () => {
5362
createMynahUI(acquireVsCodeApi(), ${
5463
(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
55-
});
64+
},${JSON.stringify(Array.from(featureConfigs.entries()))});
5665
}
5766
</script>
5867
`

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

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ interface ChatPayload {
1919
export interface ConnectorProps {
2020
sendMessageToExtension: (message: ExtensionMessage) => void
2121
onMessageReceived?: (tabID: string, messageData: any, needToShowAPIDocsTab: boolean) => void
22-
onChatAnswerReceived?: (tabID: string, message: ChatItem) => void
22+
onChatAnswerReceived?: (tabID: string, message: ChatItem, messageData: any) => void
2323
onCWCContextCommandMessage: (message: ChatItem, command?: string) => string | undefined
2424
onError: (tabID: string, message: string, title: string) => void
2525
onWarning: (tabID: string, message: string, title: string) => void
@@ -282,7 +282,7 @@ export class Connector {
282282
content: messageData.relatedSuggestions,
283283
}
284284
}
285-
this.onChatAnswerReceived(messageData.tabID, answer)
285+
this.onChatAnswerReceived(messageData.tabID, answer, messageData)
286286

287287
// Exit the function if we received an answer from AI
288288
if (
@@ -310,7 +310,7 @@ export class Connector {
310310
: undefined,
311311
traceId: messageData.traceId,
312312
}
313-
this.onChatAnswerReceived(messageData.tabID, answer)
313+
this.onChatAnswerReceived(messageData.tabID, answer, messageData)
314314

315315
return
316316
}
@@ -321,13 +321,17 @@ export class Connector {
321321
return
322322
}
323323

324-
this.onChatAnswerReceived(messageData.tabID, {
325-
type: ChatItemType.ANSWER,
326-
messageId: messageData.triggerID,
327-
body: messageData.message,
328-
followUp: this.followUpGenerator.generateAuthFollowUps('cwc', messageData.authType),
329-
canBeVoted: false,
330-
})
324+
this.onChatAnswerReceived(
325+
messageData.tabID,
326+
{
327+
type: ChatItemType.ANSWER,
328+
messageId: messageData.triggerID,
329+
body: messageData.message,
330+
followUp: this.followUpGenerator.generateAuthFollowUps('cwc', messageData.authType),
331+
canBeVoted: false,
332+
},
333+
messageData
334+
)
331335

332336
return
333337
}

0 commit comments

Comments
 (0)