33 * SPDX-License-Identifier: Apache-2.0
44 */
55
6+ import { getLogger , setContext } from 'aws-core-vscode/shared'
67import { getLogger , setContext } from 'aws-core-vscode/shared'
78import * as vscode from 'vscode'
89import { diffLines } from 'diff'
@@ -18,6 +19,7 @@ export class EditDecorationManager {
1819 private currentRemovedCodeDecorations : vscode . DecorationOptions [ ] = [ ]
1920 private acceptHandler : ( ( ) => void ) | undefined
2021 private rejectHandler : ( ( ) => void ) | undefined
22+ private disposables : vscode . Disposable [ ] = [ ]
2123
2224 constructor ( ) {
2325 this . imageDecorationType = vscode . window . createTextEditorDecorationType ( {
@@ -27,8 +29,6 @@ export class EditDecorationManager {
2729 this . removedCodeDecorationType = vscode . window . createTextEditorDecorationType ( {
2830 backgroundColor : 'rgba(255, 0, 0, 0.2)' ,
2931 } )
30-
31- this . registerCommandHandlers ( )
3232 }
3333
3434 /**
@@ -104,6 +104,7 @@ export class EditDecorationManager {
104104 removedHighlights ?: vscode . DecorationOptions [ ]
105105 ) : void {
106106 // Clear any existing decorations
107+ this . registerCommandHandlers ( )
107108 this . clearDecorations ( editor )
108109
109110 // Set context to enable the Tab key handler
@@ -152,24 +153,30 @@ export class EditDecorationManager {
152153 */
153154 public registerCommandHandlers ( ) : void {
154155 // Register Tab key handler for accepting suggestion
155- vscode . commands . registerCommand ( 'aws.amazonq.inline.acceptEdit' , ( ) => {
156+ const acceptDisposable = vscode . commands . registerCommand ( 'aws.amazonq.inline.acceptEdit' , ( ) => {
156157 if ( this . acceptHandler ) {
157158 this . acceptHandler ( )
158159 }
159160 } )
161+ this . disposables . push ( acceptDisposable )
160162
161163 // Register Esc key handler for rejecting suggestion
162- vscode . commands . registerCommand ( 'aws.amazonq.inline.rejectEdit' , ( ) => {
164+ const rejectDisposable = vscode . commands . registerCommand ( 'aws.amazonq.inline.rejectEdit' , ( ) => {
163165 if ( this . rejectHandler ) {
164166 this . rejectHandler ( )
165167 }
166168 } )
169+ this . disposables . push ( rejectDisposable )
167170 }
168171
169172 /**
170173 * Disposes resources
171174 */
172175 public dispose ( ) : void {
176+ for ( const disposable of this . disposables ) {
177+ disposable . dispose ( )
178+ }
179+ this . disposables = [ ]
173180 this . imageDecorationType . dispose ( )
174181 this . removedCodeDecorationType . dispose ( )
175182 }
@@ -267,9 +274,18 @@ export async function displaySvgDecoration(
267274 // Calculate cursor position before replacing content
268275 const endPosition = getEndOfEditPosition ( originalCode , newCode )
269276
277+ // Replace content
278+
279+ // Calculate cursor position before replacing content
280+ const endPosition = getEndOfEditPosition ( originalCode , newCode )
281+
270282 // Replace content
271283 replaceEditorContent ( editor , newCode )
272284
285+ // Move cursor to end of the actual changed content
286+ editor . selection = new vscode . Selection ( endPosition , endPosition )
287+
288+
273289 // Move cursor to end of the actual changed content
274290 editor . selection = new vscode . Selection ( endPosition , endPosition )
275291
@@ -289,6 +305,7 @@ export async function displaySvgDecoration(
289305 deletedCharacterCount : deletedCharacterCount ,
290306 }
291307 languageClient . sendNotification ( 'aws/logInlineCompletionSessionResults' , params )
308+ decorationManager . dispose ( )
292309 } ,
293310 ( ) => {
294311 // Handle reject
@@ -307,6 +324,7 @@ export async function displaySvgDecoration(
307324 deletedCharacterCount : deletedCharacterCount ,
308325 }
309326 languageClient . sendNotification ( 'aws/logInlineCompletionSessionResults' , params )
327+ decorationManager . dispose ( )
310328 } ,
311329 originalCode ,
312330 newCode
0 commit comments