feat(amazonq): handle reference for inline with language server (aws#6837)

## Problem
missing log references for inline with language server

## Solution
add them

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: zixlin7

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

@@ -14,12 +14,22 @@ import {
     commands,
     languages,
     Disposable,
+    window,
+    TextEditor,
 } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
-import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
+import {
+    InlineCompletionItemWithReferences,
+    LogInlineCompletionSessionResultsParams,
+} from '@aws/language-server-runtimes/protocol'
 import { SessionManager } from './sessionManager'
 import { RecommendationService } from './recommendationService'
-import { CodeWhispererConstants } from 'aws-core-vscode/codewhisperer'
+import {
+    CodeWhispererConstants,
+    ReferenceHoverProvider,
+    ReferenceInlineProvider,
+    ReferenceLogViewProvider,
+} from 'aws-core-vscode/codewhisperer'
 
 export class InlineCompletionManager implements Disposable {
     private disposable: Disposable
@@ -53,15 +63,16 @@ export class InlineCompletionManager implements Disposable {
     public registerInlineCompletion() {
         const onInlineAcceptance = async (
             sessionId: string,
-            itemId: string,
+            item: InlineCompletionItemWithReferences,
+            editor: TextEditor,
             requestStartTime: number,
             firstCompletionDisplayLatency?: number
         ) => {
             // TODO: also log the seen state for other suggestions in session
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: sessionId,
                 completionSessionResult: {
-                    [itemId]: {
+                    [item.itemId]: {
                         seen: true,
                         accepted: true,
                         discarded: false,
@@ -76,6 +87,15 @@ export class InlineCompletionManager implements Disposable {
                 CodeWhispererConstants.platformLanguageIds,
                 this.inlineCompletionProvider
             )
+            if (item.references && item.references.length) {
+                const referenceLog = ReferenceLogViewProvider.getReferenceLog(
+                    item.insertText as string,
+                    item.references,
+                    editor
+                )
+                ReferenceLogViewProvider.instance.addReferenceLog(referenceLog)
+                ReferenceHoverProvider.instance.addCodeReferences(item.insertText as string, item.references)
+            }
         }
         commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
 
@@ -170,17 +190,24 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
         if (!session || !items.length) {
             return []
         }
+        const editor = window.activeTextEditor
         for (const item of items) {
             item.command = {
                 command: 'aws.amazonq.acceptInline',
                 title: 'On acceptance',
                 arguments: [
                     session.sessionId,
-                    item.itemId,
+                    item,
+                    editor,
                     session.requestStartTime,
                     session.firstCompletionDisplayLatency,
                 ],
             }
+            ReferenceInlineProvider.instance.setInlineReference(
+                position.line,
+                item.insertText as string,
+                item.references
+            )
         }
         return items as InlineCompletionItem[]
     }

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

@@ -3,10 +3,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import sinon from 'sinon'
-import { commands, languages } from 'vscode'
+import { CancellationToken, commands, languages, Position } from 'vscode'
 import assert from 'assert'
 import { LanguageClient } from 'vscode-languageclient'
-import { InlineCompletionManager } from '../../../../../src/app/inline/completion'
+import { AmazonQInlineCompletionItemProvider, InlineCompletionManager } from '../../../../../src/app/inline/completion'
+import { RecommendationService } from '../../../../../src/app/inline/recommendationService'
+import { SessionManager } from '../../../../../src/app/inline/sessionManager'
+import { createMockDocument, createMockTextEditor } from 'aws-core-vscode/test'
+import {
+    ReferenceHoverProvider,
+    ReferenceInlineProvider,
+    ReferenceLogViewProvider,
+} from 'aws-core-vscode/codewhisperer'
 
 describe('InlineCompletionManager', () => {
     let manager: InlineCompletionManager
@@ -19,6 +27,32 @@ describe('InlineCompletionManager', () => {
     let sandbox: sinon.SinonSandbox
     let getActiveSessionStub: sinon.SinonStub
     let getActiveRecommendationStub: sinon.SinonStub
+    let logReferenceStub: sinon.SinonStub
+    let getReferenceStub: sinon.SinonStub
+    let hoverReferenceStub: sinon.SinonStub
+    const mockDocument = createMockDocument()
+    const mockEditor = createMockTextEditor()
+    const mockPosition = { line: 0, character: 0 } as Position
+    const mockContext = { triggerKind: 1, selectedCompletionInfo: undefined }
+    const mockToken = { isCancellationRequested: false } as CancellationToken
+    const fakeReferences = [
+        {
+            message: '',
+            licenseName: 'TEST_LICENSE',
+            repository: 'TEST_REPO',
+            recommendationContentSpan: {
+                start: 0,
+                end: 10,
+            },
+        },
+    ]
+    const mockSuggestions = [
+        {
+            itemId: 'test-item',
+            insertText: 'test',
+            references: fakeReferences,
+        },
+    ]
 
     beforeEach(() => {
         sandbox = sinon.createSandbox()
@@ -41,6 +75,9 @@ describe('InlineCompletionManager', () => {
         manager = new InlineCompletionManager(languageClient)
         getActiveSessionStub = sandbox.stub(manager['sessionManager'], 'getActiveSession')
         getActiveRecommendationStub = sandbox.stub(manager['sessionManager'], 'getActiveRecommendation')
+        getReferenceStub = sandbox.stub(ReferenceLogViewProvider, 'getReferenceLog')
+        logReferenceStub = sandbox.stub(ReferenceLogViewProvider.instance, 'addReferenceLog')
+        hoverReferenceStub = sandbox.stub(ReferenceHoverProvider.instance, 'addCodeReferences')
     })
 
     afterEach(() => {
@@ -65,11 +102,16 @@ describe('InlineCompletionManager', () => {
                     ?.find((call) => call.args[0] === 'aws.amazonq.acceptInline')?.args[1]
 
                 const sessionId = 'test-session'
-                const itemId = 'test-item'
                 const requestStartTime = Date.now() - 1000
                 const firstCompletionDisplayLatency = 500
 
-                await acceptanceHandler(sessionId, itemId, requestStartTime, firstCompletionDisplayLatency)
+                await acceptanceHandler(
+                    sessionId,
+                    mockSuggestions[0],
+                    mockEditor,
+                    requestStartTime,
+                    firstCompletionDisplayLatency
+                )
 
                 assert(sendNotificationStub.calledOnce)
                 assert(
@@ -78,7 +120,7 @@ describe('InlineCompletionManager', () => {
                         sinon.match({
                             sessionId,
                             completionSessionResult: {
-                                [itemId]: {
+                                [mockSuggestions[0].itemId]: {
                                     seen: true,
                                     accepted: true,
                                     discarded: false,
@@ -91,6 +133,39 @@ describe('InlineCompletionManager', () => {
                 assert(disposableStub.calledOnce)
                 assert(registerProviderStub.calledTwice) // Once in constructor, once after acceptance
             })
+
+            it('should log reference if there is any', async () => {
+                const acceptanceHandler = registerCommandStub
+                    .getCalls()
+                    ?.find((call) => call.args[0] === 'aws.amazonq.acceptInline')?.args[1]
+
+                const sessionId = 'test-session'
+                const requestStartTime = Date.now() - 1000
+                const firstCompletionDisplayLatency = 500
+                const mockReferenceLog = 'test reference log'
+                getReferenceStub.returns(mockReferenceLog)
+
+                await acceptanceHandler(
+                    sessionId,
+                    mockSuggestions[0],
+                    mockEditor,
+                    requestStartTime,
+                    firstCompletionDisplayLatency
+                )
+
+                assert(getReferenceStub.calledOnce)
+                assert(
+                    getReferenceStub.calledWith(
+                        mockSuggestions[0].insertText,
+                        mockSuggestions[0].references,
+                        mockEditor
+                    )
+                )
+                assert(logReferenceStub.calledOnce)
+                assert(logReferenceStub.calledWith(mockReferenceLog))
+                assert(hoverReferenceStub.calledOnce)
+                assert(hoverReferenceStub.calledWith(mockSuggestions[0].insertText, mockSuggestions[0].references))
+            })
         })
 
         describe('onInlineRejection', () => {
@@ -179,4 +254,82 @@ describe('InlineCompletionManager', () => {
             })
         })
     })
+
+    describe('AmazonQInlineCompletionItemProvider', () => {
+        describe('provideInlineCompletionItems', () => {
+            let mockSessionManager: SessionManager
+            let provider: AmazonQInlineCompletionItemProvider
+            let getAllRecommendationsStub: sinon.SinonStub
+            let recommendationService: RecommendationService
+            let setInlineReferenceStub: sinon.SinonStub
+
+            beforeEach(() => {
+                recommendationService = new RecommendationService(mockSessionManager)
+                setInlineReferenceStub = sandbox.stub(ReferenceInlineProvider.instance, 'setInlineReference')
+
+                mockSessionManager = {
+                    getActiveSession: getActiveSessionStub,
+                    getActiveRecommendation: getActiveRecommendationStub,
+                } as unknown as SessionManager
+
+                getActiveSessionStub.returns({
+                    sessionId: 'test-session',
+                    suggestions: mockSuggestions,
+                    isRequestInProgress: false,
+                    requestStartTime: Date.now(),
+                })
+                getActiveRecommendationStub.returns(mockSuggestions)
+                getAllRecommendationsStub = sandbox.stub(recommendationService, 'getAllRecommendations')
+                getAllRecommendationsStub.resolves()
+            }),
+                it('should call recommendation service to get new suggestions for new sessions', async () => {
+                    provider = new AmazonQInlineCompletionItemProvider(
+                        languageClient,
+                        recommendationService,
+                        mockSessionManager
+                    )
+                    const items = await provider.provideInlineCompletionItems(
+                        mockDocument,
+                        mockPosition,
+                        mockContext,
+                        mockToken
+                    )
+                    assert(getAllRecommendationsStub.calledOnce)
+                    assert.deepStrictEqual(items, mockSuggestions)
+                }),
+                it('should not call recommendation service for existing sessions', async () => {
+                    provider = new AmazonQInlineCompletionItemProvider(
+                        languageClient,
+                        recommendationService,
+                        mockSessionManager,
+                        false
+                    )
+                    const items = await provider.provideInlineCompletionItems(
+                        mockDocument,
+                        mockPosition,
+                        mockContext,
+                        mockToken
+                    )
+                    assert(getAllRecommendationsStub.notCalled)
+                    assert.deepStrictEqual(items, mockSuggestions)
+                }),
+                it('should handle reference if there is any', async () => {
+                    provider = new AmazonQInlineCompletionItemProvider(
+                        languageClient,
+                        recommendationService,
+                        mockSessionManager,
+                        false
+                    )
+                    await provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken)
+                    assert(setInlineReferenceStub.calledOnce)
+                    assert(
+                        setInlineReferenceStub.calledWithExactly(
+                            mockPosition.line,
+                            mockSuggestions[0].insertText,
+                            fakeReferences
+                        )
+                    )
+                })
+        })
+    })
 })