Skip to content

Commit bbf7b0e

Browse files
authored
feat(amazonq): re-add activeState/line trackers (#7257)
## Problem - "Amazon Q is generating ..." does not show with the lsp mode ## Solution - re-add the line tracker with tests - re-implement the activeState tracker ## Notes In master the active state tracker decides whether or not to show "Amazon Q is generating ..." by the following: - When a change is made, the auto trigger decides whether or not to start a recommendation request. When a recommendation requests eventually starts, an event is sent to the active state tracker to tell it to start showing the "Amazon Q is generating ..." message. When the first recommendation starts loading and the results are shown to the user another event is sent telling it to hide the message. It de-bounces this message showing every 1000ms so that the message is not constantly toggling on/off In this implementation its slightly different: - VSCode decides when to trigger the inline completion through their inline completion provider. From here we show the "Amazon Q is generating ... " message until the first recommendation is received from the language server and shown to the user. It still de-bounces this message every 1000ms so that users aren't shown it too often --- - 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.
1 parent 1a728f1 commit bbf7b0e

File tree

8 files changed

+605
-4
lines changed

8 files changed

+605
-4
lines changed

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,19 +33,25 @@ import {
3333
ImportAdderProvider,
3434
CodeSuggestionsState,
3535
} from 'aws-core-vscode/codewhisperer'
36+
import { ActiveStateTracker } from './stateTracker/activeStateTracker'
37+
import { LineTracker } from './stateTracker/lineTracker'
3638

3739
export class InlineCompletionManager implements Disposable {
3840
private disposable: Disposable
3941
private inlineCompletionProvider: AmazonQInlineCompletionItemProvider
4042
private languageClient: LanguageClient
4143
private sessionManager: SessionManager
4244
private recommendationService: RecommendationService
45+
private lineTracker: LineTracker
46+
private activeStateTracker: ActiveStateTracker
4347
private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
4448

4549
constructor(languageClient: LanguageClient) {
4650
this.languageClient = languageClient
4751
this.sessionManager = new SessionManager()
48-
this.recommendationService = new RecommendationService(this.sessionManager)
52+
this.lineTracker = new LineTracker()
53+
this.activeStateTracker = new ActiveStateTracker(this.lineTracker)
54+
this.recommendationService = new RecommendationService(this.sessionManager, this.activeStateTracker)
4955
this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(
5056
languageClient,
5157
this.recommendationService,
@@ -55,11 +61,15 @@ export class InlineCompletionManager implements Disposable {
5561
CodeWhispererConstants.platformLanguageIds,
5662
this.inlineCompletionProvider
5763
)
64+
65+
this.lineTracker.ready()
5866
}
5967

6068
public dispose(): void {
6169
if (this.disposable) {
6270
this.disposable.dispose()
71+
this.activeStateTracker.dispose()
72+
this.lineTracker.dispose()
6373
}
6474
}
6575

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,13 @@ import {
1111
import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
1212
import { LanguageClient } from 'vscode-languageclient'
1313
import { SessionManager } from './sessionManager'
14+
import { ActiveStateTracker } from './stateTracker/activeStateTracker'
1415

1516
export class RecommendationService {
16-
constructor(private readonly sessionManager: SessionManager) {}
17+
constructor(
18+
private readonly sessionManager: SessionManager,
19+
private readonly activeStateTracker: ActiveStateTracker
20+
) {}
1721

1822
async getAllRecommendations(
1923
languageClient: LanguageClient,
@@ -31,6 +35,8 @@ export class RecommendationService {
3135
}
3236
const requestStartTime = Date.now()
3337

38+
await this.activeStateTracker.showGenerating(context.triggerKind)
39+
3440
// Handle first request
3541
const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest(
3642
inlineCompletionWithReferencesRequestType as any,
@@ -54,6 +60,8 @@ export class RecommendationService {
5460
} else {
5561
this.sessionManager.closeSession()
5662
}
63+
64+
this.activeStateTracker.hideGenerating()
5765
}
5866

5967
private async processRemainingRequests(
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import { editorUtilities } from 'aws-core-vscode/shared'
7+
import * as vscode from 'vscode'
8+
import { LineSelection, LineTracker } from './lineTracker'
9+
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
10+
import { cancellableDebounce } from 'aws-core-vscode/utils'
11+
12+
export class ActiveStateTracker implements vscode.Disposable {
13+
private readonly _disposable: vscode.Disposable
14+
15+
private readonly cwLineHintDecoration: vscode.TextEditorDecorationType =
16+
vscode.window.createTextEditorDecorationType({
17+
after: {
18+
margin: '0 0 0 3em',
19+
contentText: 'Amazon Q is generating...',
20+
textDecoration: 'none',
21+
fontWeight: 'normal',
22+
fontStyle: 'normal',
23+
color: 'var(--vscode-editorCodeLens-foreground)',
24+
},
25+
rangeBehavior: vscode.DecorationRangeBehavior.OpenOpen,
26+
isWholeLine: true,
27+
})
28+
29+
constructor(private readonly lineTracker: LineTracker) {
30+
this._disposable = vscode.Disposable.from(
31+
AuthUtil.instance.auth.onDidChangeConnectionState(async (e) => {
32+
if (e.state !== 'authenticating') {
33+
this.hideGenerating()
34+
}
35+
}),
36+
AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(async () => {
37+
this.hideGenerating()
38+
})
39+
)
40+
}
41+
42+
dispose() {
43+
this._disposable.dispose()
44+
}
45+
46+
readonly refreshDebounced = cancellableDebounce(async () => {
47+
await this._refresh(true)
48+
}, 1000)
49+
50+
async showGenerating(triggerType: vscode.InlineCompletionTriggerKind) {
51+
if (triggerType === vscode.InlineCompletionTriggerKind.Invoke) {
52+
// if user triggers on demand, immediately update the UI and cancel the previous debounced update if there is one
53+
this.refreshDebounced.cancel()
54+
await this._refresh(true)
55+
} else {
56+
await this.refreshDebounced.promise()
57+
}
58+
}
59+
60+
async _refresh(shouldDisplay: boolean) {
61+
const editor = vscode.window.activeTextEditor
62+
if (!editor) {
63+
return
64+
}
65+
66+
const selections = this.lineTracker.selections
67+
if (!editor || !selections || !editorUtilities.isTextEditor(editor)) {
68+
this.hideGenerating()
69+
return
70+
}
71+
72+
if (!AuthUtil.instance.isConnectionValid()) {
73+
this.hideGenerating()
74+
return
75+
}
76+
77+
await this.updateDecorations(editor, selections, shouldDisplay)
78+
}
79+
80+
hideGenerating() {
81+
vscode.window.activeTextEditor?.setDecorations(this.cwLineHintDecoration, [])
82+
}
83+
84+
async updateDecorations(editor: vscode.TextEditor, lines: LineSelection[], shouldDisplay: boolean) {
85+
const range = editor.document.validateRange(
86+
new vscode.Range(lines[0].active, lines[0].active, lines[0].active, lines[0].active)
87+
)
88+
89+
if (shouldDisplay) {
90+
editor.setDecorations(this.cwLineHintDecoration, [range])
91+
} else {
92+
editor.setDecorations(this.cwLineHintDecoration, [])
93+
}
94+
}
95+
}
Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import * as vscode from 'vscode'
7+
import { editorUtilities, setContext } from 'aws-core-vscode/shared'
8+
9+
export interface LineSelection {
10+
anchor: number
11+
active: number
12+
}
13+
14+
export interface LinesChangeEvent {
15+
readonly editor: vscode.TextEditor | undefined
16+
readonly selections: LineSelection[] | undefined
17+
18+
readonly reason: 'editor' | 'selection' | 'content'
19+
}
20+
21+
/**
22+
* This class providees a single interface to manage and access users' "line" selections
23+
* Callers could use it by subscribing onDidChangeActiveLines to do UI updates or logic needed to be executed when line selections get changed
24+
*/
25+
export class LineTracker implements vscode.Disposable {
26+
private _onDidChangeActiveLines = new vscode.EventEmitter<LinesChangeEvent>()
27+
get onDidChangeActiveLines(): vscode.Event<LinesChangeEvent> {
28+
return this._onDidChangeActiveLines.event
29+
}
30+
31+
private _editor: vscode.TextEditor | undefined
32+
private _disposable: vscode.Disposable | undefined
33+
34+
private _selections: LineSelection[] | undefined
35+
get selections(): LineSelection[] | undefined {
36+
return this._selections
37+
}
38+
39+
private _onReady: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
40+
get onReady(): vscode.Event<void> {
41+
return this._onReady.event
42+
}
43+
44+
private _ready: boolean = false
45+
get isReady() {
46+
return this._ready
47+
}
48+
49+
constructor() {
50+
this._disposable = vscode.Disposable.from(
51+
vscode.window.onDidChangeActiveTextEditor(async (e) => {
52+
await this.onActiveTextEditorChanged(e)
53+
}),
54+
vscode.window.onDidChangeTextEditorSelection(async (e) => {
55+
await this.onTextEditorSelectionChanged(e)
56+
}),
57+
vscode.workspace.onDidChangeTextDocument((e) => {
58+
this.onContentChanged(e)
59+
})
60+
)
61+
62+
queueMicrotask(async () => await this.onActiveTextEditorChanged(vscode.window.activeTextEditor))
63+
}
64+
65+
dispose() {
66+
this._disposable?.dispose()
67+
}
68+
69+
ready() {
70+
if (this._ready) {
71+
throw new Error('Linetracker is already activated')
72+
}
73+
74+
this._ready = true
75+
queueMicrotask(() => this._onReady.fire())
76+
}
77+
78+
// @VisibleForTesting
79+
async onActiveTextEditorChanged(editor: vscode.TextEditor | undefined) {
80+
if (editor === this._editor) {
81+
return
82+
}
83+
84+
this._editor = editor
85+
this._selections = toLineSelections(editor?.selections)
86+
if (this._selections && this._selections[0]) {
87+
const s = this._selections.map((item) => item.active + 1)
88+
await setContext('codewhisperer.activeLine', s)
89+
}
90+
91+
this.notifyLinesChanged('editor')
92+
}
93+
94+
// @VisibleForTesting
95+
async onTextEditorSelectionChanged(e: vscode.TextEditorSelectionChangeEvent) {
96+
// If this isn't for our cached editor and its not a real editor -- kick out
97+
if (this._editor !== e.textEditor && !editorUtilities.isTextEditor(e.textEditor)) {
98+
return
99+
}
100+
101+
const selections = toLineSelections(e.selections)
102+
if (this._editor === e.textEditor && this.includes(selections)) {
103+
return
104+
}
105+
106+
this._editor = e.textEditor
107+
this._selections = selections
108+
if (this._selections && this._selections[0]) {
109+
const s = this._selections.map((item) => item.active + 1)
110+
await setContext('codewhisperer.activeLine', s)
111+
}
112+
113+
this.notifyLinesChanged('selection')
114+
}
115+
116+
// @VisibleForTesting
117+
onContentChanged(e: vscode.TextDocumentChangeEvent) {
118+
const editor = vscode.window.activeTextEditor
119+
if (e.document === editor?.document && e.contentChanges.length > 0 && editorUtilities.isTextEditor(editor)) {
120+
this._editor = editor
121+
this._selections = toLineSelections(this._editor?.selections)
122+
123+
this.notifyLinesChanged('content')
124+
}
125+
}
126+
127+
notifyLinesChanged(reason: 'editor' | 'selection' | 'content') {
128+
const e: LinesChangeEvent = { editor: this._editor, selections: this.selections, reason: reason }
129+
this._onDidChangeActiveLines.fire(e)
130+
}
131+
132+
includes(selections: LineSelection[]): boolean
133+
includes(line: number, options?: { activeOnly: boolean }): boolean
134+
includes(lineOrSelections: number | LineSelection[], options?: { activeOnly: boolean }): boolean {
135+
if (typeof lineOrSelections !== 'number') {
136+
return isIncluded(lineOrSelections, this._selections)
137+
}
138+
139+
if (this._selections === undefined || this._selections.length === 0) {
140+
return false
141+
}
142+
143+
const line = lineOrSelections
144+
const activeOnly = options?.activeOnly ?? true
145+
146+
for (const selection of this._selections) {
147+
if (
148+
line === selection.active ||
149+
(!activeOnly &&
150+
((selection.anchor >= line && line >= selection.active) ||
151+
(selection.active >= line && line >= selection.anchor)))
152+
) {
153+
return true
154+
}
155+
}
156+
return false
157+
}
158+
}
159+
160+
function isIncluded(selections: LineSelection[] | undefined, within: LineSelection[] | undefined): boolean {
161+
if (selections === undefined && within === undefined) {
162+
return true
163+
}
164+
if (selections === undefined || within === undefined || selections.length !== within.length) {
165+
return false
166+
}
167+
168+
return selections.every((s, i) => {
169+
const match = within[i]
170+
return s.active === match.active && s.anchor === match.anchor
171+
})
172+
}
173+
174+
function toLineSelections(selections: readonly vscode.Selection[]): LineSelection[]
175+
function toLineSelections(selections: readonly vscode.Selection[] | undefined): LineSelection[] | undefined
176+
function toLineSelections(selections: readonly vscode.Selection[] | undefined) {
177+
return selections?.map((s) => ({ active: s.active.line, anchor: s.anchor.line }))
178+
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import {
1515
ReferenceInlineProvider,
1616
ReferenceLogViewProvider,
1717
} from 'aws-core-vscode/codewhisperer'
18+
import { ActiveStateTracker } from '../../../../../src/app/inline/stateTracker/activeStateTracker'
19+
import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker'
1820

1921
describe('InlineCompletionManager', () => {
2022
let manager: InlineCompletionManager
@@ -264,7 +266,9 @@ describe('InlineCompletionManager', () => {
264266
let setInlineReferenceStub: sinon.SinonStub
265267

266268
beforeEach(() => {
267-
recommendationService = new RecommendationService(mockSessionManager)
269+
const lineTracker = new LineTracker()
270+
const activeStateController = new ActiveStateTracker(lineTracker)
271+
recommendationService = new RecommendationService(mockSessionManager, activeStateController)
268272
setInlineReferenceStub = sandbox.stub(ReferenceInlineProvider.instance, 'setInlineReference')
269273

270274
mockSessionManager = {

0 commit comments

Comments
 (0)