refactor(amazonq): separate edits from completion code path

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

@@ -274,15 +274,6 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             return []
         }
 
-        // yield event loop to let the document listen catch updates
-        await sleep(1)
-        // prevent user deletion invoking auto trigger
-        // this is a best effort estimate of deletion
-        if (this.documentEventListener.isLastEventDeletion(document.uri.fsPath)) {
-            getLogger().debug('Skip auto trigger when deleting code')
-            return []
-        }
-
         let logstr = `GenerateCompletion metadata:\\n`
         try {
             const t0 = performance.now()
@@ -370,8 +361,8 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
                 },
                 token,
                 isAutoTrigger,
-                getAllRecommendationsOptions,
-                this.documentEventListener.getLastDocumentChangeEvent(document.uri.fsPath)?.event
+                this.documentEventListener,
+                getAllRecommendationsOptions
             )
             // get active item from session for displaying
             const items = this.sessionManager.getActiveRecommendation()
@@ -404,7 +395,8 @@ ${itemLog}
 
             const cursorPosition = document.validatePosition(position)
 
-            if (position.isAfter(editor.selection.active)) {
+            // Edit suggestion works differently than completion suggestion, so even when it's a deletion and cause cursor to move back, we still allow the request to go through
+            if (position.isAfter(editor.selection.active) && items.length > 0 && !items[0].isInlineEdit) {
                 const params: LogInlineCompletionSessionResultsParams = {
                     sessionId: session.sessionId,
                     completionSessionResult: {
@@ -417,7 +409,7 @@ ${itemLog}
                 }
                 this.languageClient.sendNotification(this.logSessionResultMessageName, params)
                 this.sessionManager.clear()
-                logstr += `- cursor moved behind trigger position. Discarding suggestion...`
+                logstr += `- cursor moved behind trigger position. Discarding completion suggestion...`
                 return []
             }
 

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

@@ -2,12 +2,12 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import * as vscode from 'vscode'
 import {
     InlineCompletionListWithReferences,
     InlineCompletionWithReferencesParams,
     inlineCompletionWithReferencesRequestType,
     TextDocumentContentChangeEvent,
+    editCompletionRequestType,
 } from '@aws/language-server-runtimes/protocol'
 import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
@@ -20,7 +20,8 @@ import {
 } from 'aws-core-vscode/codewhisperer'
 import { TelemetryHelper } from './telemetryHelper'
 import { ICursorUpdateRecorder } from './cursorUpdateManager'
-import { getLogger } from 'aws-core-vscode/shared'
+import { getLogger, sleep } from 'aws-core-vscode/shared'
+import { DocumentEventListener } from './documentEventListener'
 import { getOpenFilesInWindow } from 'aws-core-vscode/utils'
 import { asyncCallWithTimeout } from '../../util/timeoutUtil'
 
@@ -66,9 +67,11 @@ export class RecommendationService {
         context: InlineCompletionContext,
         token: CancellationToken,
         isAutoTrigger: boolean,
-        options: GetAllRecommendationsOptions = { emitTelemetry: true, showUi: true },
-        documentChangeEvent?: vscode.TextDocumentChangeEvent
+        documentEventListener: DocumentEventListener,
+        options: GetAllRecommendationsOptions = { emitTelemetry: true, showUi: true }
     ) {
+        const documentChangeEvent = documentEventListener?.getLastDocumentChangeEvent(document.uri.fsPath)?.event
+
         // Record that a regular request is being made
         this.cursorUpdateRecorder?.recordCompletionRequest()
         const documentChangeParams = documentChangeEvent
@@ -119,7 +122,41 @@ export class RecommendationService {
             })
             const t0 = performance.now()
 
-            const result = await this.getRecommendationsWithTimeout(languageClient, request, token)
+            // Yield event loop to let the document listen catch updates
+            await sleep(1)
+            // Best effort estimate of deletion
+            const isTriggerByDeletion = documentEventListener.isLastEventDeletion(document.uri.fsPath)
+
+            const ps: Promise<InlineCompletionListWithReferences>[] = []
+            /**
+             * IsTriggerByDeletion is to prevent user deletion invoking Completions.
+             * PartialResultToken is not a hack for now since only Edits suggestion use partialResultToken across different calls of [getAllRecommendations],
+             * 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) {
+                const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
+                    inlineCompletionWithReferencesRequestType.method,
+                    request,
+                    token
+                )
+                ps.push(completionPromise)
+            }
+
+            /**
+             * Though Edit request is sent on keystrokes everytime, the language server will execute the request in a debounced manner so that it won't be immediately executed.
+             */
+            const editPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
+                editCompletionRequestType.method,
+                request,
+                token
+            )
+            ps.push(editPromise)
+
+            let result = await Promise.race(ps)
+            if (ps.length > 1 && result.items.length === 0) {
+                result = await editPromise
+            }
 
             getLogger().info('Received inline completion response from LSP: %O', {
                 sessionId: result.sessionId,

packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts

@@ -14,6 +14,10 @@ import { createMockDocument } from 'aws-core-vscode/test'
 import { CursorUpdateManager } from '../../../../../src/app/inline/cursorUpdateManager'
 import { CodeWhispererStatusBarManager } from 'aws-core-vscode/codewhisperer'
 import { globals } from 'aws-core-vscode/shared'
+import { DocumentEventListener } from '../../../../../src/app/inline/documentEventListener'
+
+const completionApi = 'aws/textDocument/inlineCompletionWithReferences'
+const editApi = 'aws/textDocument/editCompletion'
 
 describe('RecommendationService', () => {
     let languageClient: LanguageClient
@@ -28,6 +32,10 @@ describe('RecommendationService', () => {
     const mockPosition = { line: 0, character: 0 } as Position
     const mockContext = { triggerKind: InlineCompletionTriggerKind.Automatic, selectedCompletionInfo: undefined }
     const mockToken = { isCancellationRequested: false } as CancellationToken
+    const mockDocumentEventListener = {
+        isLastEventDeletion: (filepath: string) => false,
+        getLastDocumentChangeEvent: (filepath: string) => undefined,
+    } as DocumentEventListener
     const mockInlineCompletionItemOne = {
         insertText: 'ItemOne',
     } as InlineCompletionItem
@@ -134,12 +142,19 @@ describe('RecommendationService', () => {
                 mockPosition,
                 mockContext,
                 mockToken,
-                true
+                true,
+                mockDocumentEventListener
             )
 
             // Verify sendRequest was called with correct parameters
-            assert(sendRequestStub.calledOnce)
-            const requestArgs = sendRequestStub.firstCall.args[1]
+            const cs = sendRequestStub.getCalls()
+            const completionCalls = cs.filter((c) => c.firstArg === completionApi)
+            const editCalls = cs.filter((c) => c.firstArg === editApi)
+            assert.strictEqual(cs.length, 2)
+            assert.strictEqual(completionCalls.length, 1)
+            assert.strictEqual(editCalls.length, 1)
+
+            const requestArgs = completionCalls[0].args[1]
             assert.deepStrictEqual(requestArgs, {
                 textDocument: {
                     uri: 'file:///test.py',
@@ -177,12 +192,19 @@ describe('RecommendationService', () => {
                 mockPosition,
                 mockContext,
                 mockToken,
-                true
+                true,
+                mockDocumentEventListener
             )
 
             // Verify sendRequest was called with correct parameters
-            assert(sendRequestStub.calledTwice)
-            const firstRequestArgs = sendRequestStub.firstCall.args[1]
+            const cs = sendRequestStub.getCalls()
+            const completionCalls = cs.filter((c) => c.firstArg === completionApi)
+            const editCalls = cs.filter((c) => c.firstArg === editApi)
+            assert.strictEqual(cs.length, 3)
+            assert.strictEqual(completionCalls.length, 2)
+            assert.strictEqual(editCalls.length, 1)
+
+            const firstRequestArgs = completionCalls[0].args[1]
             const expectedRequestArgs = {
                 textDocument: {
                     uri: 'file:///test.py',
@@ -192,7 +214,7 @@ describe('RecommendationService', () => {
                 documentChangeParams: undefined,
                 openTabFilepaths: [],
             }
-            const secondRequestArgs = sendRequestStub.secondCall.args[1]
+            const secondRequestArgs = completionCalls[1].args[1]
             assert.deepStrictEqual(firstRequestArgs, expectedRequestArgs)
             assert.deepStrictEqual(secondRequestArgs, {
                 ...expectedRequestArgs,
@@ -218,7 +240,8 @@ describe('RecommendationService', () => {
                 mockPosition,
                 mockContext,
                 mockToken,
-                true
+                true,
+                mockDocumentEventListener
             )
 
             // Verify recordCompletionRequest was called
@@ -235,6 +258,7 @@ describe('RecommendationService', () => {
                 mockContext,
                 mockToken,
                 true,
+                mockDocumentEventListener,
                 {
                     showUi: false,
                     emitTelemetry: true,
@@ -254,7 +278,8 @@ describe('RecommendationService', () => {
                 mockPosition,
                 mockContext,
                 mockToken,
-                true
+                true,
+                mockDocumentEventListener
             )
 
             // Verify UI methods were called
@@ -286,6 +311,7 @@ describe('RecommendationService', () => {
                     mockContext,
                     mockToken,
                     true,
+                    mockDocumentEventListener,
                     options
                 )