diff --git a/packages/amazonq/src/app/inline/completion.ts b/packages/amazonq/src/app/inline/completion.ts index 75017a4e337..fd5ef749d0d 100644 --- a/packages/amazonq/src/app/inline/completion.ts +++ b/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,7 +63,8 @@ export class InlineCompletionManager implements Disposable { public registerInlineCompletion() { const onInlineAcceptance = async ( sessionId: string, - itemId: string, + item: InlineCompletionItemWithReferences, + editor: TextEditor, requestStartTime: number, firstCompletionDisplayLatency?: number ) => { @@ -61,7 +72,7 @@ export class InlineCompletionManager implements Disposable { 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[] } diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts index d0af4acff71..d2182329e45 100644 --- a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts +++ b/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 + ) + ) + }) + }) + }) })