fix pagination

Author: leigaol

packages/amazonq/.changes/next-release/Bug Fix-91380b87-5955-4c15-b762-31e7f1c71575.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Render first response before receiving all paginated inline completion results"
+}

packages/amazonq/package.json

@@ -917,12 +917,12 @@
             },
             {
                 "key": "right",
-                "command": "editor.action.inlineSuggest.showNext",
+                "command": "aws.amazonq.showNext",
                 "when": "inlineSuggestionVisible && !editorReadonly && aws.codewhisperer.connected"
             },
             {
                 "key": "left",
-                "command": "editor.action.inlineSuggest.showPrevious",
+                "command": "aws.amazonq.showPrev",
                 "when": "inlineSuggestionVisible && !editorReadonly && aws.codewhisperer.connected"
             },
             {
@@ -1325,26 +1325,40 @@
                     "fontCharacter": "\\f1de"
                 }
             },
-            "aws-schemas-registry": {
+            "aws-sagemaker-code-editor": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1df"
                 }
             },
-            "aws-schemas-schema": {
+            "aws-sagemaker-jupyter-lab": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e0"
                 }
             },
-            "aws-stepfunctions-preview": {
+            "aws-schemas-registry": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e1"
                 }
+            },
+            "aws-schemas-schema": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e2"
+                }
+            },
+            "aws-stepfunctions-preview": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e3"
+                }
             }
         },
         "walkthroughs": [

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

@@ -2,7 +2,7 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-
+//import vscode from 'vscode'
 import {
     InlineCompletionListWithReferences,
     InlineCompletionWithReferencesParams,
@@ -80,7 +80,7 @@ export class RecommendationService {
                     nextToken: request.partialResultToken,
                 },
             })
-            let result: InlineCompletionListWithReferences = await languageClient.sendRequest(
+            const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
                 inlineCompletionWithReferencesRequestType.method,
                 request,
                 token
@@ -120,18 +120,10 @@ export class RecommendationService {
                 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}`)
+                if (result.partialResultToken) {
+                    this.processRemainingRequests(languageClient, request, result, token).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
@@ -140,11 +132,6 @@ export class RecommendationService {
                 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
-            this.sessionManager.closeSession()
-            TelemetryHelper.instance.setAllPaginationEndTime()
-            options.emitTelemetry && TelemetryHelper.instance.tryRecordClientComponentLatency()
         } catch (error: any) {
             getLogger().error('Error getting recommendations: %O', error)
             // bearer token expired
@@ -167,4 +154,31 @@ export class RecommendationService {
             }
         }
     }
+
+    private async processRemainingRequests(
+        languageClient: LanguageClient,
+        initialRequest: InlineCompletionWithReferencesParams,
+        firstResult: InlineCompletionListWithReferences,
+        token: CancellationToken
+    ): Promise<void> {
+        let nextToken = firstResult.partialResultToken
+        while (nextToken) {
+            const request = { ...initialRequest, partialResultToken: nextToken }
+
+            const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
+                inlineCompletionWithReferencesRequestType.method,
+                request,
+                token
+            )
+            this.sessionManager.updateSessionSuggestions(result.items)
+            nextToken = result.partialResultToken
+        }
+
+        this.sessionManager.closeSession()
+
+        // refresh inline completion items to render paginated responses
+        // All pagination requests completed
+        TelemetryHelper.instance.setAllPaginationEndTime()
+        TelemetryHelper.instance.tryRecordClientComponentLatency()
+    }
 }

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

@@ -24,6 +24,7 @@ export interface CodeWhispererSession {
 export class SessionManager {
     private activeSession?: CodeWhispererSession
     private _acceptedSuggestionCount: number = 0
+    private _refreshedSessions = new Set<string>()
 
     constructor() {}
 
@@ -51,6 +52,7 @@ export class SessionManager {
             return
         }
         this.activeSession.isRequestInProgress = false
+        console.log(`session closed!`)
     }
 
     public getActiveSession() {
@@ -86,4 +88,17 @@ export class SessionManager {
     public clear() {
         this.activeSession = undefined
     }
+
+    // re-render the session ghost text to display paginated responses once per completed session
+    public async maybeRefreshSessionUx() {
+        if (
+            this.activeSession &&
+            !this.activeSession.isRequestInProgress &&
+            !this._refreshedSessions.has(this.activeSession.sessionId)
+        ) {
+            await vscode.commands.executeCommand('editor.action.inlineSuggest.hide')
+            await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
+            this._refreshedSessions.add(this.activeSession.sessionId)
+        }
+    }
 }

packages/amazonq/src/lsp/client.ts

@@ -265,6 +265,14 @@ async function onLanguageServerReady(
 
     toDispose.push(
         inlineManager,
+        Commands.register('aws.amazonq.showPrev', async () => {
+            await sessionManager.maybeRefreshSessionUx()
+            await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
+        }),
+        Commands.register('aws.amazonq.showNext', async () => {
+            await sessionManager.maybeRefreshSessionUx()
+            await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
+        }),
         Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
             await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
         }),