feat: passing partialResultToken for the next trigger if user accepts EDITS suggestion

Author: atonaamz

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

@@ -326,7 +326,8 @@ export async function displaySvgDecoration(
                         selectedCompletionInfo: undefined,
                     },
                     new vscode.CancellationTokenSource().token,
-                    { emitTelemetry: false, showUi: false }
+                    { emitTelemetry: false, showUi: false },
+                    session.editsStreakPartialResultToken
                 )
             }
         },

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

@@ -221,7 +221,8 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
         position: Position,
         context: InlineCompletionContext,
         token: CancellationToken,
-        getAllRecommendationsOptions?: GetAllRecommendationsOptions
+        getAllRecommendationsOptions?: GetAllRecommendationsOptions,
+        editsStreakPartialResultToken?: number | string
     ): Promise<InlineCompletionItem[]> {
         getLogger().info('_provideInlineCompletionItems called with: %O', {
             documentUri: document.uri.toString(),
@@ -323,7 +324,8 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
                 position,
                 context,
                 token,
-                getAllRecommendationsOptions
+                getAllRecommendationsOptions,
+                editsStreakPartialResultToken
             )
             // get active item from session for displaying
             const items = this.sessionManager.getActiveRecommendation()

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

@@ -42,18 +42,22 @@ export class RecommendationService {
         position: Position,
         context: InlineCompletionContext,
         token: CancellationToken,
-        options: GetAllRecommendationsOptions = { emitTelemetry: true, showUi: true }
+        options: GetAllRecommendationsOptions = { emitTelemetry: true, showUi: true },
+        editsStreakPartialResultToken?: number | string
     ) {
         // Record that a regular request is being made
         this.cursorUpdateRecorder?.recordCompletionRequest()
 
-        const request: InlineCompletionWithReferencesParams = {
+        let request: InlineCompletionWithReferencesParams = {
             textDocument: {
                 uri: document.uri.toString(),
             },
             position,
             context,
         }
+        if (editsStreakPartialResultToken) {
+            request = { ...request, partialResultToken: editsStreakPartialResultToken }
+        }
         const requestStartTime = globals.clock.Date.now()
         const statusBar = CodeWhispererStatusBarManager.instance
 
@@ -113,18 +117,26 @@ export class RecommendationService {
             )
 
             // If there are more results to fetch, handle them in the background
-            try {
-                while (result.partialResultToken) {
-                    const paginatedRequest = { ...request, partialResultToken: result.partialResultToken }
-                    result = await languageClient.sendRequest(
-                        inlineCompletionWithReferencesRequestType.method,
-                        paginatedRequest,
-                        token
-                    )
-                    this.sessionManager.updateSessionSuggestions(result.items)
+            const isInlineEdit = result.items.find((item) => item.isInlineEdit) ? true : false
+            // Skip fetching for more items if the suggesion is EDITS. If it is EDITS suggestion, only fetching for more
+            // suggestions when the user start to accept a suggesion.
+            if (!isInlineEdit) {
+                try {
+                    while (result.partialResultToken) {
+                        const paginatedRequest = { ...request, partialResultToken: result.partialResultToken }
+                        result = await languageClient.sendRequest(
+                            inlineCompletionWithReferencesRequestType.method,
+                            paginatedRequest,
+                            token
+                        )
+                        this.sessionManager.updateSessionSuggestions(result.items)
+                    }
+                } catch (error) {
+                    languageClient.warn(`Error when getting suggestions: ${error}`)
                 }
-            } catch (error) {
-                languageClient.warn(`Error when getting suggestions: ${error}`)
+            } else {
+                // save editsStreakPartialResultToken for the next EDITS suggestion trigger if user accepts
+                this.sessionManager.updateActiveEditsStreakPartialResultToken(result.partialResultToken)
             }
 
             // Close session and finalize telemetry regardless of pagination path

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

@@ -14,6 +14,8 @@ export interface CodeWhispererSession {
     requestStartTime: number
     firstCompletionDisplayLatency?: number
     startPosition: vscode.Position
+    // partialResultToken for the next trigger if user accepts an EDITS suggestion
+    editsStreakPartialResultToken?: number | string
 }
 
 export class SessionManager {
@@ -69,6 +71,13 @@ export class SessionManager {
         this._acceptedSuggestionCount += 1
     }
 
+    public updateActiveEditsStreakPartialResultToken(partialResultToken?: number | string) {
+        if (!this.activeSession || !partialResultToken) {
+            return
+        }
+        this.activeSession.editsStreakPartialResultToken = partialResultToken
+    }
+
     public clear() {
         this.activeSession = undefined
     }