diff --git a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
index 28612d825b5..dc45becfb13 100644
--- a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
+++ b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
@@ -122,19 +122,19 @@ export class EditDecorationManager {
     /**
      * Displays an edit suggestion as an SVG image in the editor and highlights removed code
      */
-    public displayEditSuggestion(
+    public async displayEditSuggestion(
         editor: vscode.TextEditor,
         svgImage: vscode.Uri,
         startLine: number,
-        onAccept: () => void,
-        onReject: () => void,
+        onAccept: () => Promise<void>,
+        onReject: () => Promise<void>,
         originalCode: string,
         newCode: string,
         originalCodeHighlightRanges: Array<{ line: number; start: number; end: number }>
-    ): void {
-        this.clearDecorations(editor)
+    ): Promise<void> {
+        await this.clearDecorations(editor)
 
-        void setContext('aws.amazonq.editSuggestionActive' as any, true)
+        await setContext('aws.amazonq.editSuggestionActive' as any, true)
 
         this.acceptHandler = onAccept
         this.rejectHandler = onReject
@@ -157,14 +157,14 @@ export class EditDecorationManager {
     /**
      * Clears all edit suggestion decorations
      */
-    public clearDecorations(editor: vscode.TextEditor): void {
+    public async clearDecorations(editor: vscode.TextEditor): Promise<void> {
         editor.setDecorations(this.imageDecorationType, [])
         editor.setDecorations(this.removedCodeDecorationType, [])
         this.currentImageDecoration = undefined
         this.currentRemovedCodeDecorations = []
         this.acceptHandler = undefined
         this.rejectHandler = undefined
-        void setContext('aws.amazonq.editSuggestionActive' as any, false)
+        await setContext('aws.amazonq.editSuggestionActive' as any, false)
     }
 
     /**
@@ -285,7 +285,7 @@ export async function displaySvgDecoration(
 ) {
     const originalCode = editor.document.getText()
 
-    decorationManager.displayEditSuggestion(
+    await decorationManager.displayEditSuggestion(
         editor,
         svgImage,
         startLine,
@@ -303,7 +303,7 @@ export async function displaySvgDecoration(
             // Move cursor to end of the actual changed content
             editor.selection = new vscode.Selection(endPosition, endPosition)
 
-            decorationManager.clearDecorations(editor)
+            await decorationManager.clearDecorations(editor)
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
@@ -330,10 +330,10 @@ export async function displaySvgDecoration(
                 )
             }
         },
-        () => {
+        async () => {
             // Handle reject
             getLogger().info('Edit suggestion rejected')
-            decorationManager.clearDecorations(editor)
+            await decorationManager.clearDecorations(editor)
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
diff --git a/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts b/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts
index df4fac09c28..b88e30487a5 100644
--- a/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts
+++ b/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts
@@ -57,7 +57,7 @@ describe('EditDecorationManager', function () {
         sandbox.restore()
     })
 
-    it('should display SVG decorations in the editor', function () {
+    it('should display SVG decorations in the editor', async function () {
         // Create a fake SVG image URI
         const svgUri = vscode.Uri.parse('file:///path/to/image.svg')
 
@@ -69,7 +69,7 @@ describe('EditDecorationManager', function () {
         editorStub.setDecorations.reset()
 
         // Call displayEditSuggestion
-        manager.displayEditSuggestion(
+        await manager.displayEditSuggestion(
             editorStub as unknown as vscode.TextEditor,
             svgUri,
             0,
@@ -94,7 +94,7 @@ describe('EditDecorationManager', function () {
     })
 
     // Helper function to setup edit suggestion test
-    function setupEditSuggestionTest() {
+    async function setupEditSuggestionTest() {
         // Create a fake SVG image URI
         const svgUri = vscode.Uri.parse('file:///path/to/image.svg')
 
@@ -103,7 +103,7 @@ describe('EditDecorationManager', function () {
         const rejectHandler = sandbox.stub()
 
         // Display the edit suggestion
-        manager.displayEditSuggestion(
+        await manager.displayEditSuggestion(
             editorStub as unknown as vscode.TextEditor,
             svgUri,
             0,
@@ -117,8 +117,8 @@ describe('EditDecorationManager', function () {
         return { acceptHandler, rejectHandler }
     }
 
-    it('should trigger accept handler when command is executed', function () {
-        const { acceptHandler, rejectHandler } = setupEditSuggestionTest()
+    it('should trigger accept handler when command is executed', async function () {
+        const { acceptHandler, rejectHandler } = await setupEditSuggestionTest()
 
         // Find the command handler that was registered for accept
         const acceptCommandArgs = commandsStub.registerCommand.args.find(
@@ -138,8 +138,8 @@ describe('EditDecorationManager', function () {
         }
     })
 
-    it('should trigger reject handler when command is executed', function () {
-        const { acceptHandler, rejectHandler } = setupEditSuggestionTest()
+    it('should trigger reject handler when command is executed', async function () {
+        const { acceptHandler, rejectHandler } = await setupEditSuggestionTest()
 
         // Find the command handler that was registered for reject
         const rejectCommandArgs = commandsStub.registerCommand.args.find(
@@ -159,12 +159,12 @@ describe('EditDecorationManager', function () {
         }
     })
 
-    it('should clear decorations when requested', function () {
+    it('should clear decorations when requested', async function () {
         // Reset the setDecorations stub to clear any previous calls
         editorStub.setDecorations.reset()
 
         // Call clearDecorations
-        manager.clearDecorations(editorStub as unknown as vscode.TextEditor)
+        await manager.clearDecorations(editorStub as unknown as vscode.TextEditor)
 
         // Verify decorations were cleared
         assert.strictEqual(editorStub.setDecorations.callCount, 2)