Various improvements to CodeWhisperer codebase (#3823)

1. Refactor some duplicated code
2. Refactor some condition checking logic to improve
readiability
3. Moved some of the states from RecommendationHandler
to CodeWhisperer session since they are session scope(continued
work), as well as removed some others that are not needed.
4. Fixed an issue where CodeWhisperer would try to get editor context
(left/right context) for every request within a trigger session.
These info only need to be retrieved once in the beginning of the session.
5. Fixed an issue where CodeWhisperer would try to use the language
type of the current editor for sending telemetry; There could be cases when
the current editor is no longer the original editor where invocation happens,
leading to the telemetry to have the incorrect language type.
6. Fixed an issue where CodeWhisperer would keep stale suggestions as global
states after the pagination has been requested to cancel. Add a check right
after the service response and return immediately if the current pagination
session is cancelled.
7. Right before showing the inline suggestions, do a check for the current
window focus state. If it's not focused, mark all suggestions as discarded
and exit this session.
8. Created a CodeWhispererApplication class to managed some CW specific
application-wide listeners/variables.

Author: andrewyuq

src/codewhisperer/activation.ts

@@ -75,6 +75,10 @@ export async function activate(context: ExtContext): Promise<void> {
      */
     const client = new codewhispererClient.DefaultCodeWhispererClient()
 
+    // Service initialization
+    ReferenceInlineProvider.instance
+    ImportAdderProvider.instance
+
     context.extensionContext.subscriptions.push(
         /**
          * Configuration change
@@ -176,7 +180,7 @@ export async function activate(context: ExtContext): Promise<void> {
             if (isInlineCompletionEnabled() && e.uri.fsPath !== InlineCompletionService.instance.filePath()) {
                 return
             }
-            RecommendationHandler.instance.reportUserDecisionOfRecommendation(vscode.window.activeTextEditor, -1)
+            RecommendationHandler.instance.reportUserDecisions(-1)
         }),
 
         vscode.languages.registerHoverProvider(
@@ -259,49 +263,49 @@ export async function activate(context: ExtContext): Promise<void> {
                 await RecommendationHandler.instance.onCursorChange(e)
             }),
             vscode.workspace.onDidChangeTextDocument(async e => {
+                const editor = vscode.window.activeTextEditor
+                if (!editor) {
+                    return
+                }
+                if (e.document !== editor.document) {
+                    return
+                }
+                if (!runtimeLanguageContext.isLanguageSupported(e.document.languageId)) {
+                    return
+                }
+
                 /**
                  * CodeWhisperer security panel dynamic handling
                  */
-                if (e.document === vscode.window.activeTextEditor?.document) {
-                    disposeSecurityDiagnostic(e)
-                }
+                disposeSecurityDiagnostic(e)
 
                 CodeWhispererCodeCoverageTracker.getTracker(e.document.languageId)?.countTotalTokens(e)
 
                 /**
                  * Handle this keystroke event only when
-                 * 1. It is in current non plaintext active editor
-                 * 2. It is not a backspace
-                 * 3. It is not caused by CodeWhisperer editing
-                 * 4. It is not from undo/redo.
+                 * 1. It is not a backspace
+                 * 2. It is not caused by CodeWhisperer editing
+                 * 3. It is not from undo/redo.
                  */
-                if (
-                    e.document === vscode.window.activeTextEditor?.document &&
-                    runtimeLanguageContext.isLanguageSupported(e.document.languageId) &&
-                    e.contentChanges.length != 0 &&
-                    !vsCodeState.isCodeWhispererEditing
-                ) {
-                    if (vsCodeState.lastUserModificationTime) {
-                        TelemetryHelper.instance.setTimeSinceLastModification(
-                            performance.now() - vsCodeState.lastUserModificationTime
-                        )
-                    }
-                    vsCodeState.lastUserModificationTime = performance.now()
-                    /**
-                     * Important:  Doing this sleep(10) is to make sure
-                     * 1. this event is processed by vs code first
-                     * 2. editor.selection.active has been successfully updated by VS Code
-                     * Then this event can be processed by our code.
-                     */
-                    await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
-                    if (!RecommendationHandler.instance.isSuggestionVisible()) {
-                        await KeyStrokeHandler.instance.processKeyStroke(
-                            e,
-                            vscode.window.activeTextEditor,
-                            client,
-                            await getConfigEntry()
-                        )
-                    }
+                if (e.contentChanges.length === 0 || vsCodeState.isCodeWhispererEditing) {
+                    return
+                }
+
+                if (vsCodeState.lastUserModificationTime) {
+                    TelemetryHelper.instance.setTimeSinceLastModification(
+                        performance.now() - vsCodeState.lastUserModificationTime
+                    )
+                }
+                vsCodeState.lastUserModificationTime = performance.now()
+                /**
+                 * Important:  Doing this sleep(10) is to make sure
+                 * 1. this event is processed by vs code first
+                 * 2. editor.selection.active has been successfully updated by VS Code
+                 * Then this event can be processed by our code.
+                 */
+                await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
+                if (!RecommendationHandler.instance.isSuggestionVisible()) {
+                    await KeyStrokeHandler.instance.processKeyStroke(e, editor, client, await getConfigEntry())
                 }
             })
         )
@@ -327,39 +331,33 @@ export async function activate(context: ExtContext): Promise<void> {
              * Automated trigger
              */
             vscode.workspace.onDidChangeTextDocument(async e => {
+                const editor = vscode.window.activeTextEditor
+                if (!editor) {
+                    return
+                }
+                if (e.document !== editor.document) {
+                    return
+                }
+                if (!runtimeLanguageContext.isLanguageSupported(e.document.languageId)) {
+                    return
+                }
                 /**
                  * CodeWhisperer security panel dynamic handling
                  */
-                if (e.document === vscode.window.activeTextEditor?.document) {
-                    if (isCloud9()) {
-                        securityPanelViewProvider.disposeSecurityPanelItem(e, vscode.window.activeTextEditor)
-                    } else {
-                        disposeSecurityDiagnostic(e)
-                    }
-                }
-
+                securityPanelViewProvider.disposeSecurityPanelItem(e, editor)
                 CodeWhispererCodeCoverageTracker.getTracker(e.document.languageId)?.countTotalTokens(e)
 
-                if (
-                    e.document === vscode.window.activeTextEditor?.document &&
-                    runtimeLanguageContext.isLanguageSupported(e.document.languageId) &&
-                    e.contentChanges.length != 0 &&
-                    !vsCodeState.isCodeWhispererEditing
-                ) {
-                    /**
-                     * Important:  Doing this sleep(10) is to make sure
-                     * 1. this event is processed by vs code first
-                     * 2. editor.selection.active has been successfully updated by VS Code
-                     * Then this event can be processed by our code.
-                     */
-                    await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
-                    await KeyStrokeHandler.instance.processKeyStroke(
-                        e,
-                        vscode.window.activeTextEditor,
-                        client,
-                        await getConfigEntry()
-                    )
+                if (e.contentChanges.length != 0 && !vsCodeState.isCodeWhispererEditing) {
+                    return
                 }
+                /**
+                 * Important:  Doing this sleep(10) is to make sure
+                 * 1. this event is processed by vs code first
+                 * 2. editor.selection.active has been successfully updated by VS Code
+                 * Then this event can be processed by our code.
+                 */
+                await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
+                await KeyStrokeHandler.instance.processKeyStroke(e, editor, client, await getConfigEntry())
             }),
 
             /**
@@ -389,7 +387,7 @@ export async function activate(context: ExtContext): Promise<void> {
 }
 
 export async function shutdown() {
-    RecommendationHandler.instance.reportUserDecisionOfRecommendation(vscode.window.activeTextEditor, -1)
+    RecommendationHandler.instance.reportUserDecisions(-1)
     CodeWhispererTracker.getTracker().shutdown()
 }
 

src/codewhisperer/commands/invokeRecommendation.ts

@@ -4,7 +4,7 @@
  */
 
 import * as vscode from 'vscode'
-import { vsCodeState, ConfigurationEntry } from '../models/model'
+import { vsCodeState, ConfigurationEntry, GetRecommendationsResponse } from '../models/model'
 import { resetIntelliSenseState } from '../util/globalStateUtil'
 import { DefaultCodeWhispererClient } from '../client/codewhisperer'
 import { isCloud9 } from '../../shared/extensionUtilities'
@@ -64,8 +64,9 @@ export async function invokeRecommendation(
         vsCodeState.isIntelliSenseActive = false
         RecommendationHandler.instance.isGenerateRecommendationInProgress = true
         try {
+            let response: GetRecommendationsResponse
             if (isCloud9('classic') || isIamConnection(AuthUtil.instance.conn)) {
-                await RecommendationHandler.instance.getRecommendations(
+                response = await RecommendationHandler.instance.getRecommendations(
                     client,
                     editor,
                     'OnDemand',
@@ -77,7 +78,7 @@ export async function invokeRecommendation(
                 if (AuthUtil.instance.isConnectionExpired()) {
                     await AuthUtil.instance.showReauthenticatePrompt()
                 }
-                await RecommendationHandler.instance.getRecommendations(
+                response = await RecommendationHandler.instance.getRecommendations(
                     client,
                     editor,
                     'OnDemand',
@@ -86,7 +87,7 @@ export async function invokeRecommendation(
                     true
                 )
             }
-            if (RecommendationHandler.instance.canShowRecommendationInIntelliSense(editor, true)) {
+            if (RecommendationHandler.instance.canShowRecommendationInIntelliSense(editor, true, response)) {
                 await vscode.commands.executeCommand('editor.action.triggerSuggest').then(() => {
                     vsCodeState.isIntelliSenseActive = true
                 })

src/codewhisperer/commands/onAcceptance.ts

@@ -83,8 +83,5 @@ export async function onAcceptance(acceptanceEntry: OnRecommendationAcceptanceEn
     }
 
     // at the end of recommendation acceptance, report user decisions and clear recommendations.
-    RecommendationHandler.instance.reportUserDecisionOfRecommendation(
-        acceptanceEntry.editor,
-        acceptanceEntry.acceptIndex
-    )
+    RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex)
 }

src/codewhisperer/commands/onInlineAcceptance.ts

@@ -129,9 +129,6 @@ export async function onInlineAcceptance(
             )
         }
 
-        RecommendationHandler.instance.reportUserDecisionOfRecommendation(
-            acceptanceEntry.editor,
-            acceptanceEntry.acceptIndex
-        )
+        RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex)
     }
 }

src/codewhisperer/models/constants.ts

@@ -217,7 +217,7 @@ export const artifactTypeSource = 'SourceCode'
 export const codeScanFindingsSchema = 'codescan/findings/1.0'
 
 // wait time for editor to update editor.selection.active (in milliseconds)
-export const vsCodeCursorUpdateDelay = 3
+export const vsCodeCursorUpdateDelay = 10
 
 export const reloadWindow = 'Reload Now'
 

src/codewhisperer/models/model.ts

@@ -36,6 +36,12 @@ export const vsCodeState: VsCodeState = {
     lastUserModificationTime: 0,
 }
 
+// This response struct can contain more info as needed
+export interface GetRecommendationsResponse {
+    readonly result: 'Succeeded' | 'Failed'
+    readonly errorMessage: string | undefined
+}
+
 export interface AcceptedSuggestionEntry {
     readonly time: Date
     readonly fileUrl: vscode.Uri

src/codewhisperer/service/importAdderProvider.ts

@@ -8,6 +8,7 @@ import { isCloud9 } from '../../shared/extensionUtilities'
 import { Recommendation } from '../client/codewhisperer'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { findLineToInsertImportStatement } from '../util/importAdderUtil'
+import { application } from '../util/codeWhispererApplication'
 
 /**
  * ImportAdderProvider
@@ -27,6 +28,12 @@ export class ImportAdderProvider implements vscode.CodeLensProvider {
 
     static #instance: ImportAdderProvider
 
+    constructor() {
+        application().clearCodeWhispererUIListener(_ => {
+            this.clear()
+        })
+    }
+
     public static get instance() {
         return (this.#instance ??= new this())
     }

src/codewhisperer/service/inlineCompletionItemProvider.ts

@@ -10,6 +10,7 @@ import { TelemetryHelper } from '../util/telemetryHelper'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { ReferenceInlineProvider } from './referenceInlineProvider'
 import { ImportAdderProvider } from './importAdderProvider'
+import { application } from '../util/codeWhispererApplication'
 
 export class CWInlineCompletionItemProvider implements vscode.InlineCompletionItemProvider {
     private activeItemIndex: number | undefined
@@ -133,8 +134,7 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
         token: vscode.CancellationToken
     ): vscode.ProviderResult<vscode.InlineCompletionItem[] | vscode.InlineCompletionList> {
         if (position.line < 0 || position.isBefore(this.startPos)) {
-            ReferenceInlineProvider.instance.removeInlineReference()
-            ImportAdderProvider.instance.clear()
+            application()._clearCodeWhispererUIListener.fire()
             this.activeItemIndex = undefined
             return
         }
@@ -164,7 +164,7 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
             ImportAdderProvider.instance.onShowRecommendation(document, this.startPos.line, r)
             this.nextMove = 0
             TelemetryHelper.instance.setFirstSuggestionShowTime()
-            TelemetryHelper.instance.tryRecordClientComponentLatency(document.languageId)
+            TelemetryHelper.instance.tryRecordClientComponentLatency()
             this._onDidShow.fire()
             if (matchedCount >= 2 || this.nextToken !== '') {
                 const result = [item]
@@ -175,8 +175,7 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
             }
             return [item]
         }
-        ReferenceInlineProvider.instance.removeInlineReference()
-        ImportAdderProvider.instance.clear()
+        application()._clearCodeWhispererUIListener.fire()
         this.activeItemIndex = undefined
         return []
     }

src/codewhisperer/service/inlineCompletionService.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import { ConfigurationEntry, vsCodeState } from '../models/model'
+import { ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model'
 import * as CodeWhispererConstants from '../models/constants'
 import { DefaultCodeWhispererClient } from '../client/codewhisperer'
 import { RecommendationHandler } from './recommendationHandler'
@@ -15,6 +15,7 @@ import { AuthUtil } from '../util/authUtil'
 import { shared } from '../../shared/utilities/functionUtils'
 import { ClassifierTrigger } from './classifierTrigger'
 import { session } from '../util/codeWhispererSession'
+import { noSuggestions } from '../models/constants'
 
 const performance = globalThis.performance ?? require('perf_hooks').performance
 
@@ -80,6 +81,10 @@ export class InlineCompletionService {
         ) {
             return
         }
+
+        // Call report user decisions once to report recommendations leftover from last invocation.
+        RecommendationHandler.instance.reportUserDecisions(-1)
+
         this.setCodeWhispererStatusBarLoading()
         if (ClassifierTrigger.instance.shouldInvokeClassifier(editor.document.languageId)) {
             ClassifierTrigger.instance.recordClassifierResultForAutoTrigger(editor, autoTriggerType, event)
@@ -96,10 +101,14 @@ export class InlineCompletionService {
         TelemetryHelper.instance.setInvocationStartTime(performance.now())
         RecommendationHandler.instance.checkAndResetCancellationTokens()
         RecommendationHandler.instance.documentUri = editor.document.uri
+        let response: GetRecommendationsResponse = {
+            result: 'Failed',
+            errorMessage: undefined,
+        }
         try {
             let page = 0
             while (page < this.maxPage) {
-                await RecommendationHandler.instance.getRecommendations(
+                response = await RecommendationHandler.instance.getRecommendations(
                     client,
                     editor,
                     triggerType,
@@ -109,8 +118,8 @@ export class InlineCompletionService {
                     page
                 )
                 if (RecommendationHandler.instance.checkAndResetCancellationTokens()) {
-                    // RecommendationHandler.instance.reportUserDecisionOfRecommendation(editor, -1)
-                    // vscode.commands.executeCommand('aws.codeWhisperer.refreshStatusBar')
+                    RecommendationHandler.instance.reportUserDecisions(-1)
+                    vscode.commands.executeCommand('aws.codeWhisperer.refreshStatusBar')
                     TelemetryHelper.instance.setIsRequestCancelled(true)
                     return
                 }
@@ -125,13 +134,9 @@ export class InlineCompletionService {
         }
         vscode.commands.executeCommand('aws.codeWhisperer.refreshStatusBar')
         if (triggerType === 'OnDemand' && session.recommendations.length === 0) {
-            if (RecommendationHandler.instance.errorMessagePrompt !== '') {
-                showTimedMessage(RecommendationHandler.instance.errorMessagePrompt, 2000)
-            } else {
-                showTimedMessage(CodeWhispererConstants.noSuggestions, 2000)
-            }
+            showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 2000)
         }
-        TelemetryHelper.instance.tryRecordClientComponentLatency(editor.document.languageId)
+        TelemetryHelper.instance.tryRecordClientComponentLatency()
     }
 
     setCodeWhispererStatusBarLoading() {