Merge branch 'aws:master' into code-review-tool

Author: singhAws

package-lock.json

@@ -41,7 +41,7 @@
         "skippedTestReport": "ts-node ./scripts/skippedTestReport.ts ./packages/amazonq/test/e2e/"
     },
     "devDependencies": {
-        "@aws-toolkits/telemetry": "^1.0.326",
+        "@aws-toolkits/telemetry": "^1.0.328",
         "@playwright/browser-chromium": "^1.43.1",
         "@stylistic/eslint-plugin": "^2.11.0",
         "@types/he": "^1.2.3",

package.json

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

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

@@ -924,12 +924,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"
             },
             {

packages/amazonq/package.json

@@ -2,7 +2,6 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-
 import {
     InlineCompletionListWithReferences,
     InlineCompletionWithReferencesParams,
@@ -80,7 +79,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 +119,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 +131,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 +153,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/recommendationService.ts

@@ -24,6 +24,7 @@ export interface CodeWhispererSession {
 export class SessionManager {
     private activeSession?: CodeWhispererSession
     private _acceptedSuggestionCount: number = 0
+    private _refreshedSessions = new Set<string>()
 
     constructor() {}
 
@@ -86,4 +87,20 @@ 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')
+            if (this._refreshedSessions.size > 1000) {
+                this._refreshedSessions.clear()
+            }
+            this._refreshedSessions.add(this.activeSession.sessionId)
+        }
+    }
 }

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

@@ -266,6 +266,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')
         }),