feat(amazonq): re-add activeState/line trackers

jpinkney-aws · jpinkney-aws · commit 41c6e5625b81 · 2025-05-07T21:05:44.000-04:00
diff --git a/packages/amazonq/src/app/inline/completion.ts b/packages/amazonq/src/app/inline/completion.ts
@@ -33,19 +33,25 @@ import {
     ImportAdderProvider,
     CodeSuggestionsState,
 } from 'aws-core-vscode/codewhisperer'
+import { ActiveStateTracker } from './stateTracker/activeStateTracker'
+import { LineTracker } from './stateTracker/lineTracker'
 
 export class InlineCompletionManager implements Disposable {
     private disposable: Disposable
     private inlineCompletionProvider: AmazonQInlineCompletionItemProvider
     private languageClient: LanguageClient
     private sessionManager: SessionManager
     private recommendationService: RecommendationService
+    private lineTracker: LineTracker
+    private activeStateTracker: ActiveStateTracker
     private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
 
     constructor(languageClient: LanguageClient) {
         this.languageClient = languageClient
         this.sessionManager = new SessionManager()
-        this.recommendationService = new RecommendationService(this.sessionManager)
+        this.lineTracker = new LineTracker()
+        this.activeStateTracker = new ActiveStateTracker(this.lineTracker)
+        this.recommendationService = new RecommendationService(this.sessionManager, this.activeStateTracker)
         this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(
             languageClient,
             this.recommendationService,
@@ -55,11 +61,15 @@ export class InlineCompletionManager implements Disposable {
             CodeWhispererConstants.platformLanguageIds,
             this.inlineCompletionProvider
         )
+
+        this.lineTracker.ready()
     }
 
     public dispose(): void {
         if (this.disposable) {
             this.disposable.dispose()
+            this.activeStateTracker.dispose()
+            this.lineTracker.dispose()
         }
     }
 
diff --git a/packages/amazonq/src/app/inline/recommendationService.ts b/packages/amazonq/src/app/inline/recommendationService.ts
@@ -11,9 +11,13 @@ import {
 import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
 import { SessionManager } from './sessionManager'
+import { ActiveStateTracker } from './stateTracker/activeStateTracker'
 
 export class RecommendationService {
-    constructor(private readonly sessionManager: SessionManager) {}
+    constructor(
+        private readonly sessionManager: SessionManager,
+        private readonly activeStateTracker: ActiveStateTracker
+    ) {}
 
     async getAllRecommendations(
         languageClient: LanguageClient,
@@ -31,6 +35,8 @@ export class RecommendationService {
         }
         const requestStartTime = Date.now()
 
+        await this.activeStateTracker.showGenerating(context.triggerKind)
+
         // Handle first request
         const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest(
             inlineCompletionWithReferencesRequestType as any,
@@ -54,6 +60,8 @@ export class RecommendationService {
         } else {
             this.sessionManager.closeSession()
         }
+
+        this.activeStateTracker.hideGenerating()
     }
 
     private async processRemainingRequests(
diff --git a/packages/amazonq/src/app/inline/stateTracker/activeStateTracker.ts b/packages/amazonq/src/app/inline/stateTracker/activeStateTracker.ts
@@ -0,0 +1,95 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { editorUtilities } from 'aws-core-vscode/shared'
+import * as vscode from 'vscode'
+import { LineSelection, LineTracker } from './lineTracker'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { cancellableDebounce } from 'aws-core-vscode/utils'
+
+export class ActiveStateTracker implements vscode.Disposable {
+    private readonly _disposable: vscode.Disposable
+
+    private readonly cwLineHintDecoration: vscode.TextEditorDecorationType =
+        vscode.window.createTextEditorDecorationType({
+            after: {
+                margin: '0 0 0 3em',
+                contentText: 'Amazon Q is generating...',
+                textDecoration: 'none',
+                fontWeight: 'normal',
+                fontStyle: 'normal',
+                color: 'var(--vscode-editorCodeLens-foreground)',
+            },
+            rangeBehavior: vscode.DecorationRangeBehavior.OpenOpen,
+            isWholeLine: true,
+        })
+
+    constructor(private readonly lineTracker: LineTracker) {
+        this._disposable = vscode.Disposable.from(
+            AuthUtil.instance.auth.onDidChangeConnectionState(async (e) => {
+                if (e.state !== 'authenticating') {
+                    this.hideGenerating()
+                }
+            }),
+            AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(async () => {
+                this.hideGenerating()
+            })
+        )
+    }
+
+    dispose() {
+        this._disposable.dispose()
+    }
+
+    readonly refreshDebounced = cancellableDebounce(async () => {
+        await this._refresh(true)
+    }, 1000)
+
+    async showGenerating(triggerType: vscode.InlineCompletionTriggerKind) {
+        if (triggerType === vscode.InlineCompletionTriggerKind.Invoke) {
+            // if user triggers on demand, immediately update the UI and cancel the previous debounced update if there is one
+            this.refreshDebounced.cancel()
+            await this._refresh(true)
+        } else {
+            await this.refreshDebounced.promise()
+        }
+    }
+
+    async _refresh(shouldDisplay: boolean) {
+        const editor = vscode.window.activeTextEditor
+        if (!editor) {
+            return
+        }
+
+        const selections = this.lineTracker.selections
+        if (!editor || !selections || !editorUtilities.isTextEditor(editor)) {
+            this.hideGenerating()
+            return
+        }
+
+        if (!AuthUtil.instance.isConnectionValid()) {
+            this.hideGenerating()
+            return
+        }
+
+        await this.updateDecorations(editor, selections, shouldDisplay)
+    }
+
+    hideGenerating() {
+        vscode.window.activeTextEditor?.setDecorations(this.cwLineHintDecoration, [])
+    }
+
+    async updateDecorations(editor: vscode.TextEditor, lines: LineSelection[], shouldDisplay: boolean) {
+        const range = editor.document.validateRange(
+            new vscode.Range(lines[0].active, lines[0].active, lines[0].active, lines[0].active)
+        )
+
+        if (shouldDisplay) {
+            editor.setDecorations(this.cwLineHintDecoration, [range])
+        } else {
+            editor.setDecorations(this.cwLineHintDecoration, [])
+        }
+    }
+}
diff --git a/packages/amazonq/src/app/inline/stateTracker/lineTracker.ts b/packages/amazonq/src/app/inline/stateTracker/lineTracker.ts
@@ -0,0 +1,178 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { editorUtilities, setContext } from 'aws-core-vscode/shared'
+
+export interface LineSelection {
+    anchor: number
+    active: number
+}
+
+export interface LinesChangeEvent {
+    readonly editor: vscode.TextEditor | undefined
+    readonly selections: LineSelection[] | undefined
+
+    readonly reason: 'editor' | 'selection' | 'content'
+}
+
+/**
+ * This class providees a single interface to manage and access users' "line" selections
+ * Callers could use it by subscribing onDidChangeActiveLines to do UI updates or logic needed to be executed when line selections get changed
+ */
+export class LineTracker implements vscode.Disposable {
+    private _onDidChangeActiveLines = new vscode.EventEmitter<LinesChangeEvent>()
+    get onDidChangeActiveLines(): vscode.Event<LinesChangeEvent> {
+        return this._onDidChangeActiveLines.event
+    }
+
+    private _editor: vscode.TextEditor | undefined
+    private _disposable: vscode.Disposable | undefined
+
+    private _selections: LineSelection[] | undefined
+    get selections(): LineSelection[] | undefined {
+        return this._selections
+    }
+
+    private _onReady: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
+    get onReady(): vscode.Event<void> {
+        return this._onReady.event
+    }
+
+    private _ready: boolean = false
+    get isReady() {
+        return this._ready
+    }
+
+    constructor() {
+        this._disposable = vscode.Disposable.from(
+            vscode.window.onDidChangeActiveTextEditor(async (e) => {
+                await this.onActiveTextEditorChanged(e)
+            }),
+            vscode.window.onDidChangeTextEditorSelection(async (e) => {
+                await this.onTextEditorSelectionChanged(e)
+            }),
+            vscode.workspace.onDidChangeTextDocument((e) => {
+                this.onContentChanged(e)
+            })
+        )
+
+        queueMicrotask(async () => await this.onActiveTextEditorChanged(vscode.window.activeTextEditor))
+    }
+
+    dispose() {
+        this._disposable?.dispose()
+    }
+
+    ready() {
+        if (this._ready) {
+            throw new Error('Linetracker is already activated')
+        }
+
+        this._ready = true
+        queueMicrotask(() => this._onReady.fire())
+    }
+
+    // @VisibleForTesting
+    async onActiveTextEditorChanged(editor: vscode.TextEditor | undefined) {
+        if (editor === this._editor) {
+            return
+        }
+
+        this._editor = editor
+        this._selections = toLineSelections(editor?.selections)
+        if (this._selections && this._selections[0]) {
+            const s = this._selections.map((item) => item.active + 1)
+            await setContext('codewhisperer.activeLine', s)
+        }
+
+        this.notifyLinesChanged('editor')
+    }
+
+    // @VisibleForTesting
+    async onTextEditorSelectionChanged(e: vscode.TextEditorSelectionChangeEvent) {
+        // If this isn't for our cached editor and its not a real editor -- kick out
+        if (this._editor !== e.textEditor && !editorUtilities.isTextEditor(e.textEditor)) {
+            return
+        }
+
+        const selections = toLineSelections(e.selections)
+        if (this._editor === e.textEditor && this.includes(selections)) {
+            return
+        }
+
+        this._editor = e.textEditor
+        this._selections = selections
+        if (this._selections && this._selections[0]) {
+            const s = this._selections.map((item) => item.active + 1)
+            await setContext('codewhisperer.activeLine', s)
+        }
+
+        this.notifyLinesChanged('selection')
+    }
+
+    // @VisibleForTesting
+    onContentChanged(e: vscode.TextDocumentChangeEvent) {
+        const editor = vscode.window.activeTextEditor
+        if (e.document === editor?.document && e.contentChanges.length > 0 && editorUtilities.isTextEditor(editor)) {
+            this._editor = editor
+            this._selections = toLineSelections(this._editor?.selections)
+
+            this.notifyLinesChanged('content')
+        }
+    }
+
+    notifyLinesChanged(reason: 'editor' | 'selection' | 'content') {
+        const e: LinesChangeEvent = { editor: this._editor, selections: this.selections, reason: reason }
+        this._onDidChangeActiveLines.fire(e)
+    }
+
+    includes(selections: LineSelection[]): boolean
+    includes(line: number, options?: { activeOnly: boolean }): boolean
+    includes(lineOrSelections: number | LineSelection[], options?: { activeOnly: boolean }): boolean {
+        if (typeof lineOrSelections !== 'number') {
+            return isIncluded(lineOrSelections, this._selections)
+        }
+
+        if (this._selections === undefined || this._selections.length === 0) {
+            return false
+        }
+
+        const line = lineOrSelections
+        const activeOnly = options?.activeOnly ?? true
+
+        for (const selection of this._selections) {
+            if (
+                line === selection.active ||
+                (!activeOnly &&
+                    ((selection.anchor >= line && line >= selection.active) ||
+                        (selection.active >= line && line >= selection.anchor)))
+            ) {
+                return true
+            }
+        }
+        return false
+    }
+}
+
+function isIncluded(selections: LineSelection[] | undefined, within: LineSelection[] | undefined): boolean {
+    if (selections === undefined && within === undefined) {
+        return true
+    }
+    if (selections === undefined || within === undefined || selections.length !== within.length) {
+        return false
+    }
+
+    return selections.every((s, i) => {
+        const match = within[i]
+        return s.active === match.active && s.anchor === match.anchor
+    })
+}
+
+function toLineSelections(selections: readonly vscode.Selection[]): LineSelection[]
+function toLineSelections(selections: readonly vscode.Selection[] | undefined): LineSelection[] | undefined
+function toLineSelections(selections: readonly vscode.Selection[] | undefined) {
+    return selections?.map((s) => ({ active: s.active.line, anchor: s.anchor.line }))
+}
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/completion.test.ts
@@ -15,6 +15,8 @@ import {
     ReferenceInlineProvider,
     ReferenceLogViewProvider,
 } from 'aws-core-vscode/codewhisperer'
+import { ActiveStateTracker } from '../../../../../src/app/inline/stateTracker/activeStateTracker'
+import { LineTracker } from '../../../../../src/app/inline/stateTracker/lineTracker'
 
 describe('InlineCompletionManager', () => {
     let manager: InlineCompletionManager
@@ -264,7 +266,9 @@ describe('InlineCompletionManager', () => {
             let setInlineReferenceStub: sinon.SinonStub
 
             beforeEach(() => {
-                recommendationService = new RecommendationService(mockSessionManager)
+                const lineTracker = new LineTracker()
+                const activeStateController = new ActiveStateTracker(lineTracker)
+                recommendationService = new RecommendationService(mockSessionManager, activeStateController)
                 setInlineReferenceStub = sandbox.stub(ReferenceInlineProvider.instance, 'setInlineReference')
 
                 mockSessionManager = {
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/inlineTracker.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/inlineTracker.test.ts
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts
