Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
79 changes: 18 additions & 61 deletions packages/amazonq/src/app/inline/completion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ import { RecommendationService } from './recommendationService'
import {
CodeWhispererConstants,
ReferenceHoverProvider,
ReferenceInlineProvider,
ReferenceLogViewProvider,
ImportAdderProvider,
CodeSuggestionsState,
Expand Down Expand Up @@ -158,39 +157,6 @@ export class InlineCompletionManager implements Disposable {
this.languageClient.sendNotification(this.logSessionResultMessageName, params)
}
commands.registerCommand('aws.amazonq.rejectCodeSuggestion', onInlineRejection)

/*
We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item
To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
*/

const swapProviderAndShow = async () => {
Copy link
Contributor

Choose a reason for hiding this comment

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

cc @zixlin7 was this stuff added to support inline completions that come after the initial request? or is it no longer needed? As far as I can see from the demo we still have pagination

Copy link
Contributor Author

@Hweinstock Hweinstock May 16, 2025

Choose a reason for hiding this comment

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

My understanding of the old implementation is that we fetch all the suggestions, add them to the session, and then request them from the session one at a time (by triggering the provideInlineCompletion via overwriting the scroll commands). This is where the state tracking comes in to manage what the "active" suggestion is. We also add 'placeholder' items so that VSC knows we can scroll.

Since we're fetching all the suggestions anyway, we can return those directly and avoid 'placeholder' items, and additional triggers of provideInlineCompletion.

However, I don't have all the context on the original reason for this decision.

When testing this, I also noticed a slight bug on current prod where it always shows were on suggestion 1, even when we scroll.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Checked with Zoe and sounds like this is a hack to handle the codelenses for references and imports.

  • references are confirmed working in fix(amazonq): code reference codelense works #7331.
  • import codelense is unknown. Note: There weren't any reported issue in first three bug bashes, so either this is not a common case, or already works. However, still worth following up to confirm this.

await commands.executeCommand('editor.action.inlineSuggest.hide')
this.disposable.dispose()
this.disposable = languages.registerInlineCompletionItemProvider(
CodeWhispererConstants.platformLanguageIds,
new AmazonQInlineCompletionItemProvider(
this.languageClient,
this.recommendationService,
this.sessionManager,
this.inlineTutorialAnnotation,
false
)
)
await commands.executeCommand('editor.action.inlineSuggest.trigger')
}

const prevCommandHandler = async () => {
this.sessionManager.decrementActiveIndex()
await swapProviderAndShow()
}
commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler)

const nextCommandHandler = async () => {
this.sessionManager.incrementActiveIndex()
await swapProviderAndShow()
}
commands.registerCommand('editor.action.inlineSuggest.showNext', nextCommandHandler)
}
}

Expand All @@ -199,8 +165,7 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
private readonly languageClient: LanguageClient,
private readonly recommendationService: RecommendationService,
private readonly sessionManager: SessionManager,
private readonly inlineTutorialAnnotation: InlineTutorialAnnotation,
private readonly isNewSession: boolean = true
private readonly inlineTutorialAnnotation: InlineTutorialAnnotation
) {}

provideInlineCompletionItems = debounce(
Expand All @@ -217,27 +182,24 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
): Promise<InlineCompletionItem[]> {
try {
vsCodeState.isRecommendationsActive = true
if (this.isNewSession) {
const isAutoTrigger = context.triggerKind === InlineCompletionTriggerKind.Automatic
if (isAutoTrigger && !CodeSuggestionsState.instance.isSuggestionsEnabled()) {
// return early when suggestions are disabled with auto trigger
return []
}

// tell the tutorial that completions has been triggered
await this.inlineTutorialAnnotation.triggered(context.triggerKind)
TelemetryHelper.instance.setInvokeSuggestionStartTime()
TelemetryHelper.instance.setTriggerType(context.triggerKind)

// make service requests if it's a new session
await this.recommendationService.getAllRecommendations(
this.languageClient,
document,
position,
context,
token
)
const isAutoTrigger = context.triggerKind === InlineCompletionTriggerKind.Automatic
if (isAutoTrigger && !CodeSuggestionsState.instance.isSuggestionsEnabled()) {
// return early when suggestions are disabled with auto trigger
return []
}

// tell the tutorial that completions has been triggered
await this.inlineTutorialAnnotation.triggered(context.triggerKind)
TelemetryHelper.instance.setInvokeSuggestionStartTime()
TelemetryHelper.instance.setTriggerType(context.triggerKind)

await this.recommendationService.getAllRecommendations(
this.languageClient,
document,
position,
context,
token
)
// get active item from session for displaying
const items = this.sessionManager.getActiveRecommendation()
const session = this.sessionManager.getActiveSession()
Expand Down Expand Up @@ -271,11 +233,6 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
}
item.range = new Range(cursorPosition, cursorPosition)
item.insertText = typeof item.insertText === 'string' ? item.insertText : item.insertText.value
ReferenceInlineProvider.instance.setInlineReference(
cursorPosition.line,
item.insertText,
item.references
)
ImportAdderProvider.instance.onShowRecommendation(document, cursorPosition.line, item)
}
return items as InlineCompletionItem[]
Expand Down
42 changes: 1 addition & 41 deletions packages/amazonq/src/app/inline/sessionManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ interface CodeWhispererSession {

export class SessionManager {
private activeSession?: CodeWhispererSession
private activeIndex: number = 0
private _acceptedSuggestionCount: number = 0

constructor() {}
Expand All @@ -35,7 +34,6 @@ export class SessionManager {
requestStartTime,
firstCompletionDisplayLatency,
}
this.activeIndex = 0
}

public closeSession() {
Expand All @@ -56,45 +54,8 @@ export class SessionManager {
this.activeSession.suggestions = [...this.activeSession.suggestions, ...suggestions]
}

public incrementActiveIndex() {
const suggestionCount = this.activeSession?.suggestions?.length
if (!suggestionCount) {
return
}
this.activeIndex === suggestionCount - 1 ? suggestionCount - 1 : this.activeIndex++
}

public decrementActiveIndex() {
this.activeIndex === 0 ? 0 : this.activeIndex--
}

/*
We have to maintain the active suggestion index ourselves because VS Code doesn't expose which suggestion it's currently showing
In order to keep track of the right suggestion state, and for features such as reference tracker, this hack is still needed
*/

public getActiveRecommendation(): InlineCompletionItemWithReferences[] {
let suggestionCount = this.activeSession?.suggestions.length
if (!suggestionCount) {
return []
}
if (suggestionCount === 1 && this.activeSession?.isRequestInProgress) {
suggestionCount += 1
}

const activeSuggestion = this.activeSession?.suggestions[this.activeIndex]
if (!activeSuggestion) {
return []
}
const items = [activeSuggestion]
// to make the total number of suggestions match the actual number
for (let i = 1; i < suggestionCount; i++) {
items.push({
...activeSuggestion,
insertText: `${i}`,
})
}
return items
return this.activeSession?.suggestions ?? []
}

public get acceptedSuggestionCount(): number {
Expand All @@ -107,6 +68,5 @@ export class SessionManager {

public clear() {
this.activeSession = undefined
this.activeIndex = 0
}
}
89 changes: 6 additions & 83 deletions packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,7 @@ import { AmazonQInlineCompletionItemProvider, InlineCompletionManager } from '..
import { RecommendationService } from '../../../../../src/app/inline/recommendationService'
import { SessionManager } from '../../../../../src/app/inline/sessionManager'
import { createMockDocument, createMockTextEditor, getTestWindow, installFakeClock } from 'aws-core-vscode/test'
import {
noInlineSuggestionsMsg,
ReferenceHoverProvider,
ReferenceInlineProvider,
ReferenceLogViewProvider,
} from 'aws-core-vscode/codewhisperer'
import { noInlineSuggestionsMsg, ReferenceHoverProvider, ReferenceLogViewProvider } from 'aws-core-vscode/codewhisperer'
import { InlineGeneratingMessage } from '../../../../../src/app/inline/inlineGeneratingMessage'
import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker'
import { InlineTutorialAnnotation } from '../../../../../src/app/inline/tutorials/inlineTutorialAnnotation'
Expand Down Expand Up @@ -230,46 +225,6 @@ describe('InlineCompletionManager', () => {
assert(registerProviderStub.calledTwice) // Once in constructor, once after rejection
})
})

describe('previous command', () => {
it('should register and handle previous command correctly', async () => {
const prevCommandCall = registerCommandStub
.getCalls()
.find((call) => call.args[0] === 'editor.action.inlineSuggest.showPrevious')

assert(prevCommandCall, 'Previous command should be registered')

if (prevCommandCall) {
const handler = prevCommandCall.args[1]
await handler()

assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
assert(disposableStub.calledOnce)
assert(registerProviderStub.calledTwice)
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger'))
}
})
})

describe('next command', () => {
it('should register and handle next command correctly', async () => {
const nextCommandCall = registerCommandStub
.getCalls()
.find((call) => call.args[0] === 'editor.action.inlineSuggest.showNext')

assert(nextCommandCall, 'Next command should be registered')

if (nextCommandCall) {
const handler = nextCommandCall.args[1]
await handler()

assert(executeCommandStub.calledWith('editor.action.inlineSuggest.hide'))
assert(disposableStub.calledOnce)
assert(registerProviderStub.calledTwice)
assert(executeCommandStub.calledWith('editor.action.inlineSuggest.trigger'))
}
})
})
})

describe('AmazonQInlineCompletionItemProvider', () => {
Expand All @@ -278,15 +233,13 @@ describe('InlineCompletionManager', () => {
let provider: AmazonQInlineCompletionItemProvider
let getAllRecommendationsStub: sinon.SinonStub
let recommendationService: RecommendationService
let setInlineReferenceStub: sinon.SinonStub
let inlineTutorialAnnotation: InlineTutorialAnnotation

beforeEach(() => {
const lineTracker = new LineTracker()
const activeStateController = new InlineGeneratingMessage(lineTracker)
inlineTutorialAnnotation = new InlineTutorialAnnotation(lineTracker, mockSessionManager)
recommendationService = new RecommendationService(mockSessionManager, activeStateController)
setInlineReferenceStub = sandbox.stub(ReferenceInlineProvider.instance, 'setInlineReference')

mockSessionManager = {
getActiveSession: getActiveSessionStub,
Expand Down Expand Up @@ -320,48 +273,21 @@ describe('InlineCompletionManager', () => {
assert(getAllRecommendationsStub.calledOnce)
assert.deepStrictEqual(items, mockSuggestions)
}),
it('should not call recommendation service for existing sessions', async () => {
provider = new AmazonQInlineCompletionItemProvider(
languageClient,
recommendationService,
mockSessionManager,
inlineTutorialAnnotation,
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,
inlineTutorialAnnotation,
false
inlineTutorialAnnotation
)
await provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken)
assert(setInlineReferenceStub.calledOnce)
assert(
setInlineReferenceStub.calledWithExactly(
mockPosition.line,
mockSuggestions[0].insertText,
fakeReferences
)
)
}),
it('should add a range to the completion item when missing', async function () {
provider = new AmazonQInlineCompletionItemProvider(
languageClient,
recommendationService,
mockSessionManager,
inlineTutorialAnnotation,
true
inlineTutorialAnnotation
)
getActiveRecommendationStub.returns([
{
Expand Down Expand Up @@ -391,8 +317,7 @@ describe('InlineCompletionManager', () => {
languageClient,
recommendationService,
mockSessionManager,
inlineTutorialAnnotation,
true
inlineTutorialAnnotation
)
const expectedText = 'this is my text'
getActiveRecommendationStub.returns([
Expand All @@ -415,8 +340,7 @@ describe('InlineCompletionManager', () => {
languageClient,
recommendationService,
mockSessionManager,
inlineTutorialAnnotation,
true
inlineTutorialAnnotation
)
getActiveRecommendationStub.returns([])
const messageShown = new Promise((resolve) =>
Expand Down Expand Up @@ -449,8 +373,7 @@ describe('InlineCompletionManager', () => {
languageClient,
recommendationService,
mockSessionManager,
inlineTutorialAnnotation,
false
inlineTutorialAnnotation
)
const p1 = provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken)
const p2 = provider.provideInlineCompletionItems(mockDocument, mockPosition, mockContext, mockToken)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,6 @@ describe('RecommendationService', () => {
...expectedRequestArgs,
partialResultToken: mockPartialResultToken,
})

// Verify session management
const items = sessionManager.getActiveRecommendation()
assert.deepStrictEqual(items, [mockInlineCompletionItemOne, { insertText: '1' } as InlineCompletionItem])
sessionManager.incrementActiveIndex()
const items2 = sessionManager.getActiveRecommendation()
assert.deepStrictEqual(items2, [mockInlineCompletionItemTwo, { insertText: '1' } as InlineCompletionItem])
})
})
})
Loading