fix(amazonq): correctly update the on show inlay hints for code reference and imports

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

@@ -434,7 +434,6 @@ ${itemLog}
                     }
                     item.range = new Range(cursorPosition, cursorPosition)
                     itemsMatchingTypeahead.push(item)
-                    ImportAdderProvider.instance.onShowRecommendation(document, cursorPosition.line, item)
                 }
             }
 
@@ -458,6 +457,7 @@ ${itemLog}
                 return []
             }
 
+            this.sessionManager.updateCodeReferenceAndImports()
             // suggestions returned here will be displayed on screen
             return itemsMatchingTypeahead as InlineCompletionItem[]
         } catch (e) {

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

@@ -4,7 +4,12 @@
  */
 import * as vscode from 'vscode'
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes-types'
-import { FileDiagnostic, getDiagnosticsOfCurrentFile } from 'aws-core-vscode/codewhisperer'
+import {
+    FileDiagnostic,
+    getDiagnosticsOfCurrentFile,
+    ImportAdderProvider,
+    ReferenceInlineProvider,
+} from 'aws-core-vscode/codewhisperer'
 
 // TODO: add more needed data to the session interface
 export interface CodeWhispererSession {
@@ -25,7 +30,7 @@ export class SessionManager {
     private activeSession?: CodeWhispererSession
     private _acceptedSuggestionCount: number = 0
     private _refreshedSessions = new Set<string>()
-
+    private _currentSuggestionIndex = 0
     constructor() {}
 
     public startSession(
@@ -45,6 +50,7 @@ export class SessionManager {
             firstCompletionDisplayLatency,
             diagnosticsBeforeAccept,
         }
+        this._currentSuggestionIndex = 0
     }
 
     public closeSession() {
@@ -86,6 +92,8 @@ export class SessionManager {
 
     public clear() {
         this.activeSession = undefined
+        this._currentSuggestionIndex = 0
+        this.clearReferenceInlineHintsAndImportHints()
     }
 
     // re-render the session ghost text to display paginated responses once per completed session
@@ -103,4 +111,60 @@ export class SessionManager {
             this._refreshedSessions.add(this.activeSession.sessionId)
         }
     }
+
+    public onNextSuggestion() {
+        if (this.activeSession?.suggestions && this.activeSession?.suggestions.length > 0) {
+            this._currentSuggestionIndex = (this._currentSuggestionIndex + 1) % this.activeSession.suggestions.length
+            this.updateCodeReferenceAndImports()
+        }
+    }
+
+    public onPrevSuggestion() {
+        if (this.activeSession?.suggestions && this.activeSession.suggestions.length > 0) {
+            this._currentSuggestionIndex =
+                (this._currentSuggestionIndex - 1 + this.activeSession.suggestions.length) %
+                this.activeSession.suggestions.length
+            this.updateCodeReferenceAndImports()
+        }
+    }
+
+    private clearReferenceInlineHintsAndImportHints() {
+        ReferenceInlineProvider.instance.removeInlineReference()
+        ImportAdderProvider.instance.clear()
+    }
+
+    // Ideally use this API handleDidShowCompletionItem
+    // https://github.com/microsoft/vscode/blob/main/src/vscode-dts/vscode.proposed.inlineCompletionsAdditions.d.ts#L83
+    updateCodeReferenceAndImports() {
+        try {
+            this.clearReferenceInlineHintsAndImportHints()
+            if (
+                this.activeSession?.suggestions &&
+                this.activeSession.suggestions[this._currentSuggestionIndex] &&
+                this.activeSession.suggestions.length > 0
+            ) {
+                const reference = this.activeSession.suggestions[this._currentSuggestionIndex].references
+                const insertText = this.activeSession.suggestions[this._currentSuggestionIndex].insertText
+                if (reference && reference.length > 0) {
+                    const insertTextStr =
+                        typeof insertText === 'string' ? insertText : (insertText.value ?? String(insertText))
+
+                    ReferenceInlineProvider.instance.setInlineReference(
+                        this.activeSession.startPosition.line,
+                        insertTextStr,
+                        reference
+                    )
+                }
+                if (vscode.window.activeTextEditor) {
+                    ImportAdderProvider.instance.onShowRecommendation(
+                        vscode.window.activeTextEditor.document,
+                        this.activeSession.startPosition.line,
+                        this.activeSession.suggestions[this._currentSuggestionIndex]
+                    )
+                }
+            }
+        } catch {
+            // do nothing as this is not critical path
+        }
+    }
 }

packages/amazonq/src/lsp/client.ts

@@ -309,10 +309,12 @@ async function onLanguageServerReady(
         Commands.register('aws.amazonq.showPrev', async () => {
             await sessionManager.maybeRefreshSessionUx()
             await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
+            sessionManager.onPrevSuggestion()
         }),
         Commands.register('aws.amazonq.showNext', async () => {
             await sessionManager.maybeRefreshSessionUx()
             await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
+            sessionManager.onNextSuggestion()
         }),
         Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
             await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')

packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts

@@ -256,6 +256,7 @@ describe('InlineCompletionManager', () => {
                     getActiveSession: getActiveSessionStub,
                     getActiveRecommendation: getActiveRecommendationStub,
                     clear: () => {},
+                    updateCodeReferenceAndImports: () => {},
                 } as unknown as SessionManager
 
                 getActiveSessionStub.returns({