66import * as vscode from 'vscode'
77import fs from 'fs-extra'
88import path from 'path'
9- import { commands , Position , TextEditor , Uri , window , workspace } from 'vscode'
9+ import { Position , TextEditor , window } from 'vscode'
1010import { getLogger } from '../../../shared/logger'
11- import { amazonQDiffScheme , getNonexistentFilename , tempDirPath } from '../../../shared'
11+ import { AMAZON_Q_DIFF_SCHEME , 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 } from '../../../shared'
1219
20+ class ContentProvider implements vscode . TextDocumentContentProvider {
21+ constructor ( private tempFileUri : vscode . Uri ) { }
22+
23+ provideTextDocumentContent ( _uri : vscode . Uri ) {
24+ return fs . readFileSync ( this . tempFileUri . fsPath , 'utf-8' )
25+ }
26+ }
1327export class EditorContentController {
1428 /* *
1529 * Insert the Amazon Q chat written code to the cursor position
@@ -56,6 +70,31 @@ export class EditorContentController {
5670 }
5771 }
5872
73+ /**
74+ * Accept code changes received by Amazon Q
75+ *
76+ * @param message the message from Amazon Q chat
77+ */
78+ public async acceptDiff ( message : any ) {
79+ const { filePath, selection } = extractFileAndCodeSelectionFromMessage ( message )
80+
81+ if ( filePath && message ?. code ?. trim ( ) . length > 0 && selection ) {
82+ const doc = await vscode . workspace . openTextDocument ( filePath )
83+ const codeToUpdate = getIndentedCode ( message , doc , selection )
84+ const range = getSelectionFromRange ( doc , selection )
85+ await applyChanges ( doc , range , codeToUpdate )
86+
87+ // If vscode.diff is open for the filePath then close it.
88+ vscode . window . tabGroups . all . flatMap ( ( { tabs } ) =>
89+ tabs . map ( ( tab ) => {
90+ if ( tab . label === `${ path . basename ( filePath ) } ${ amazonQTabSuffix } ) {
91+ void vscode . window . tabGroups . close ( tab )
92+ }
93+ } )
94+ )
95+ }
96+ }
97+
5998 /**
6099 * Displays a diff view comparing proposed changes with the existing file.
61100 *
@@ -70,63 +109,26 @@ export class EditorContentController {
70109 * @param message the message from Amazon Q chat
71110 */
72111 public async viewDiff ( message : any ) {
73- const { filePath } = message ?. context ?. activeFileContext || { }
74- const selection = message ?. context ?. focusAreaContext ?. selectionInsideExtendedCodeBlock as vscode . Selection
112+ const { filePath, selection } = extractFileAndCodeSelectionFromMessage ( message )
75113
76114 if ( filePath && message ?. code ?. trim ( ) . length > 0 && selection ) {
77- const id = Date . now ( )
78-
79- const originalFileUri = Uri . file ( filePath )
80- const fileNameWithExtension = path . basename ( originalFileUri . path )
81-
82- const fileName = path . parse ( fileNameWithExtension ) . name
83- const fileExtension = path . extname ( fileNameWithExtension )
84-
85- // Create a new file in the temp directory
86- const tempFile = await getNonexistentFilename ( tempDirPath , `${ fileName } ${ id } , fileExtension , 99 )
87- const tempFilePath = path . join ( tempDirPath , tempFile )
88-
89- // Create a new URI for the temp file
90- const diffScheme = amazonQDiffScheme
91- const tempFileUri = Uri . parse ( `${ diffScheme } ${ tempFilePath } )
92-
93- // Write the initial code to the temp file
94- fs . writeFileSync ( tempFilePath , fs . readFileSync ( originalFileUri . fsPath , 'utf-8' ) )
95- const doc = await workspace . openTextDocument ( tempFileUri . path )
96-
97- // Apply the edit to the temp file
98- const edit = new vscode . WorkspaceEdit ( )
99- edit . replace ( doc . uri , selection , message . code )
115+ const originalFileUri = vscode . Uri . file ( filePath )
116+ const tempFileUri = await createTempFileForDiff ( originalFileUri , message , selection )
100117
101- const successfulEdit = await workspace . applyEdit ( edit )
102- if ( successfulEdit ) {
103- await doc . save ( )
104- fs . writeFileSync ( tempFilePath , doc . getText ( ) )
105- }
106- class ContentProvider {
107- provideTextDocumentContent ( _uri : Uri ) {
108- return fs . readFileSync ( tempFilePath , 'utf-8' )
109- }
110- }
111- const contentProvider = new ContentProvider ( )
112- const disposable = workspace . registerTextDocumentContentProvider ( diffScheme , contentProvider )
113-
114- await commands . executeCommand (
118+ // Register content provider and show diff
119+ const contentProvider = new ContentProvider ( tempFileUri )
120+ const disposable = vscode . workspace . registerTextDocumentContentProvider (
121+ AMAZON_Q_DIFF_SCHEME ,
122+ contentProvider
123+ )
124+ await vscode . commands . executeCommand (
115125 'vscode.diff' ,
116126 originalFileUri ,
117127 tempFileUri ,
118- `${ fileName } ${ fileExtension } (Generated by Amazon Q) `
128+ `${ path . basename ( filePath ) } ${ amazonQTabSuffix }
119129 )
120130
121- vscode . window . onDidChangeVisibleTextEditors ( ( ) => {
122- if (
123- ! vscode . window . visibleTextEditors . some (
124- ( editor ) => editor . document . uri . toString ( ) === tempFileUri . toString ( )
125- )
126- ) {
127- disposable . dispose ( )
128- }
129- } )
131+ disposeOnEditorClose ( tempFileUri , disposable )
130132 }
131133 }
132134}
0 commit comments