feat(amazonq): passing partialResultToken for the next trigger if user accepts an EDITS suggestion

Author: atonaamz

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

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

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

@@ -20,6 +20,7 @@ import { globals, getLogger } from 'aws-core-vscode/shared'
 export interface GetAllRecommendationsOptions {
     emitTelemetry?: boolean
     showUi?: boolean
+    editsStreakToken?: number | string
 }
 
 export class RecommendationService {
@@ -47,13 +48,16 @@ export class RecommendationService {
         // 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 (options.editsStreakToken) {
+            request = { ...request, partialResultToken: options.editsStreakToken }
+        }
         const requestStartTime = globals.clock.Date.now()
         const statusBar = CodeWhispererStatusBarManager.instance
 
@@ -112,19 +116,31 @@ export class RecommendationService {
                 firstCompletionDisplayLatency
             )
 
-            // 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.some((item) => item.isInlineEdit)
+            if (!isInlineEdit) {
+                // If the suggestion is COMPLETIONS and there are more results to fetch, handle them in the background
+                getLogger().info(
+                    'Suggestion type is COMPLETIONS. Start fetching for more items if partialResultToken exists.'
+                )
+                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 {
+                // 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.
+                // Save editsStreakPartialResultToken for the next EDITS suggestion trigger if user accepts.
+                getLogger().info('Suggestion type is EDITS. Skip fetching for more items.')
+                this.sessionManager.updateActiveEditsStreakToken(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 updateActiveEditsStreakToken(partialResultToken?: number | string) {
+        if (!this.activeSession || !partialResultToken) {
+            return
+        }
+        this.activeSession.editsStreakPartialResultToken = partialResultToken
+    }
+
     public clear() {
         this.activeSession = undefined
     }