fix: passing nextToken to Flare for Edits

Author: atonaamz

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

@@ -318,18 +318,20 @@ export async function displaySvgDecoration(
                 isInlineEdit: true,
             }
             languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
-            if (inlineCompletionProvider) {
-                await inlineCompletionProvider.provideInlineCompletionItems(
-                    editor.document,
-                    endPosition,
-                    {
-                        triggerKind: vscode.InlineCompletionTriggerKind.Automatic,
-                        selectedCompletionInfo: undefined,
-                    },
-                    new vscode.CancellationTokenSource().token,
-                    { emitTelemetry: false, showUi: false, editsStreakToken: session.editsStreakPartialResultToken }
-                )
-            }
+            session.triggerOnAcceptance = true
+            // VS Code triggers suggestion on every keystroke, temporarily disable trigger on acceptance
+            // if (inlineCompletionProvider) {
+            //     await inlineCompletionProvider.provideInlineCompletionItems(
+            //         editor.document,
+            //         endPosition,
+            //         {
+            //             triggerKind: vscode.InlineCompletionTriggerKind.Automatic,
+            //             selectedCompletionInfo: undefined,
+            //         },
+            //         new vscode.CancellationTokenSource().token,
+            //         { emitTelemetry: false, showUi: false, editsStreakToken: session.editsStreakPartialResultToken }
+            //     )
+            // }
         },
         async () => {
             // Handle reject

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

@@ -253,6 +253,12 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             const prevSessionId = prevSession?.sessionId
             const prevItemId = this.sessionManager.getActiveRecommendation()?.[0]?.itemId
             const prevStartPosition = prevSession?.startPosition
+            if (prevSession?.triggerOnAcceptance) {
+                getAllRecommendationsOptions = {
+                    ...getAllRecommendationsOptions,
+                    editsStreakToken: prevSession?.editsStreakPartialResultToken,
+                }
+            }
             const editor = window.activeTextEditor
             if (prevSession && prevSessionId && prevItemId && prevStartPosition) {
                 const prefix = document.getText(new Range(prevStartPosition, position))

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

@@ -77,6 +77,7 @@ export class RecommendationService {
                     textDocument: request.textDocument,
                     position: request.position,
                     context: request.context,
+                    partialResultToken: request.partialResultToken,
                 },
             })
             let result: InlineCompletionListWithReferences = await languageClient.sendRequest(

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

@@ -16,6 +16,7 @@ export interface CodeWhispererSession {
     startPosition: vscode.Position
     // partialResultToken for the next trigger if user accepts an EDITS suggestion
     editsStreakPartialResultToken?: number | string
+    triggerOnAcceptance?: boolean
 }
 
 export class SessionManager {