feat(amazonq): discard edit suggestions if it's displayed less than 1 sec when a completion arrives

Author: andrewyuq

packages/amazonq/src/app/inline/EditRendering/displayImage.ts

@@ -24,7 +24,7 @@ export class EditDecorationManager {
     private currentImageDecoration: vscode.DecorationOptions | undefined
     private currentRemovedCodeDecorations: vscode.DecorationOptions[] = []
     private acceptHandler: (() => void) | undefined
-    private rejectHandler: (() => void) | undefined
+    private rejectHandler: ((isDiscard: boolean) => void) | undefined
 
     constructor() {
         this.registerCommandHandlers()
@@ -131,7 +131,7 @@ export class EditDecorationManager {
         svgImage: vscode.Uri,
         startLine: number,
         onAccept: () => Promise<void>,
-        onReject: () => Promise<void>,
+        onReject: (isDiscard: boolean) => Promise<void>,
         originalCode: string,
         newCode: string,
         originalCodeHighlightRanges: Array<{ line: number; start: number; end: number }>
@@ -185,9 +185,9 @@ export class EditDecorationManager {
         })
 
         // Register Esc key handler for rejecting suggestion
-        vscode.commands.registerCommand('aws.amazonq.inline.rejectEdit', () => {
+        vscode.commands.registerCommand('aws.amazonq.inline.rejectEdit', (isDiscard: boolean = false) => {
             if (this.rejectHandler) {
-                this.rejectHandler()
+                this.rejectHandler(isDiscard)
             }
         })
     }
@@ -416,20 +416,31 @@ export async function displaySvgDecoration(
             //     )
             // }
         },
-        async () => {
+        async (isDiscard: boolean) => {
             // Handle reject
-            getLogger().info('Edit suggestion rejected')
+            if (isDiscard) {
+                getLogger().info('Edit suggestion discarded')
+            } else {
+                getLogger().info('Edit suggestion rejected')
+            }
             await decorationManager.clearDecorations(editor)
             documentChangeListener.dispose()
             cursorChangeListener.dispose()
+            const suggestionState = isDiscard
+                ? {
+                      seen: false,
+                      accepted: false,
+                      discarded: true,
+                  }
+                : {
+                      seen: true,
+                      accepted: false,
+                      discarded: false,
+                  }
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
-                    [item.itemId]: {
-                        seen: true,
-                        accepted: false,
-                        discarded: false,
-                    },
+                    [item.itemId]: suggestionState,
                 },
                 totalSessionDisplayTime: Date.now() - session.requestStartTime,
                 firstCompletionDisplayLatency: session.firstCompletionDisplayLatency,

packages/amazonq/src/app/inline/editSuggestionState.ts

@@ -8,12 +8,18 @@
  */
 export class EditSuggestionState {
     private static isEditSuggestionCurrentlyActive = false
+    private static displayStartTime = performance.now()
 
     static setEditSuggestionActive(active: boolean): void {
         this.isEditSuggestionCurrentlyActive = active
+        this.displayStartTime = performance.now()
     }
 
     static isEditSuggestionActive(): boolean {
         return this.isEditSuggestionCurrentlyActive
     }
+
+    static isEditSuggestionDisplayingOverOneSecond(): boolean {
+        return this.isEditSuggestionActive() && performance.now() - this.displayStartTime > 1000
+    }
 }

packages/amazonq/src/app/inline/recommendationService.ts

@@ -8,8 +8,9 @@ import {
     inlineCompletionWithReferencesRequestType,
     TextDocumentContentChangeEvent,
     editCompletionRequestType,
+    LogInlineCompletionSessionResultsParams,
 } from '@aws/language-server-runtimes/protocol'
-import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
+import { CancellationToken, InlineCompletionContext, Position, TextDocument, commands } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
 import { SessionManager } from './sessionManager'
 import {
@@ -24,8 +25,8 @@ import { getLogger } from 'aws-core-vscode/shared'
 import { DocumentEventListener } from './documentEventListener'
 import { getOpenFilesInWindow } from 'aws-core-vscode/utils'
 import { asyncCallWithTimeout } from '../../util/timeoutUtil'
-import { EditSuggestionState } from './editSuggestionState'
 import { extractFileContextInNotebooks } from './notebookUtil'
+import { EditSuggestionState } from './editSuggestionState'
 
 export interface GetAllRecommendationsOptions {
     emitTelemetry?: boolean
@@ -137,7 +138,7 @@ export class RecommendationService {
              * Completions use PartialResultToken with single 1 call of [getAllRecommendations].
              * Edits leverage partialResultToken to achieve EditStreak such that clients can pull all continuous suggestions generated by the model within 1 EOS block.
              */
-            if (!isTriggerByDeletion && !request.partialResultToken && !EditSuggestionState.isEditSuggestionActive()) {
+            if (!isTriggerByDeletion && !request.partialResultToken) {
                 const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
                     inlineCompletionWithReferencesRequestType.method,
                     request,
@@ -187,6 +188,31 @@ export class RecommendationService {
                 })),
             })
 
+            // Completion will not be rendered if an edit suggestion has been active for longer than 1 second
+            if (EditSuggestionState.isEditSuggestionDisplayingOverOneSecond()) {
+                const session = this.sessionManager.getActiveSession()
+                if (!session) {
+                    return []
+                }
+                const itemId = this.sessionManager.getActiveRecommendation()?.[0]?.itemId
+                const params: LogInlineCompletionSessionResultsParams = {
+                    sessionId: session.sessionId,
+                    completionSessionResult: {
+                        [itemId]: {
+                            seen: false,
+                            accepted: false,
+                            discarded: true,
+                        },
+                    },
+                }
+                languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
+                this.sessionManager.clear()
+                return []
+            } else if (EditSuggestionState.isEditSuggestionActive()) {
+                // discard the current edit suggestion if its display time is less than 1 sec
+                await commands.executeCommand('aws.amazonq.inline.rejectEdit', true)
+            }
+
             TelemetryHelper.instance.setSdkApiCallEndTime()
             TelemetryHelper.instance.setSessionId(result.sessionId)
             if (result.items.length > 0 && result.items[0].itemId !== undefined) {