Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jpinkney-aws merged 2 commits into from
Mar 25, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
014c8b4
feat(amazonq): handle reference for inline with language server
zixlin7 Mar 24, 2025
4b7ff3f
log reference on accept
zixlin7 Mar 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 32 additions & 5 deletions packages/amazonq/src/app/inline/completion.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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)

Expand Down Expand Up @@ -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(
Copy link
Contributor

@jpinkney-aws jpinkney-aws Mar 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not entirely familiar with the flow but is it something like this:

  • user gets inline suggestion
    • it may have a reference attached
  • if it does have a reference then show it as a code lens
  • if they click the codelens to add a code reference it adds it in the "Code reference log" panel and opens it up

Is that mostly correct?

Copy link
Contributor Author

@zixlin7 zixlin7 Mar 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes that is right, and I realized that I missed the logging reference on acceptance piece, just added it. Thanks for catching!

position.line,
item.insertText as string,
item.references
)
}
return items as InlineCompletionItem[]
}
Expand Down
163 changes: 158 additions & 5 deletions packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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(() => {
Expand All @@ -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(
Expand All @@ -78,7 +120,7 @@ describe('InlineCompletionManager', () => {
sinon.match({
sessionId,
completionSessionResult: {
[itemId]: {
[mockSuggestions[0].itemId]: {
seen: true,
accepted: true,
discarded: false,
Expand All @@ -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', () => {
Expand Down Expand Up @@ -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
)
)
})
})
})
})
Loading