packages/core/src/codewhisperer/service/recommendationHandler.ts (622 lines of code) (raw):

/*! * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. * SPDX-License-Identifier: Apache-2.0 */ import * as vscode from 'vscode' import { extensionVersion } from '../../shared/vscode/env' import { RecommendationsList, DefaultCodeWhispererClient, CognitoCredentialsError } from '../client/codewhisperer' import * as EditorContext from '../util/editorContext' import * as CodeWhispererConstants from '../models/constants' import { ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model' import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { AWSError } from 'aws-sdk' import { isAwsError } from '../../shared/errors' import { TelemetryHelper } from '../util/telemetryHelper' import { getLogger } from '../../shared/logger/logger' import { hasVendedIamCredentials } from '../../auth/auth' import { asyncCallWithTimeout, isInlineCompletionEnabled, isVscHavingRegressionInlineCompletionApi, } from '../util/commonUtil' import { showTimedMessage } from '../../shared/utilities/messages' import { CodewhispererAutomatedTriggerType, CodewhispererCompletionType, CodewhispererGettingStartedTask, CodewhispererTriggerType, telemetry, } from '../../shared/telemetry/telemetry' import { CodeWhispererCodeCoverageTracker } from '../tracker/codewhispererCodeCoverageTracker' import { invalidCustomizationMessage } from '../models/constants' import { getSelectedCustomization, switchToBaseCustomizationAndNotify } from '../util/customizationUtil' import { session } from '../util/codeWhispererSession' import { Commands } from '../../shared/vscode/commands2' import globals from '../../shared/extensionGlobals' import { noSuggestions, updateInlineLockKey } from '../models/constants' import AsyncLock from 'async-lock' import { AuthUtil } from '../util/authUtil' import { CWInlineCompletionItemProvider } from './inlineCompletionItemProvider' import { application } from '../util/codeWhispererApplication' import { openUrl } from '../../shared/utilities/vsCodeUtils' import { indent } from '../../shared/utilities/textUtilities' import path from 'path' import { isIamConnection } from '../../auth/connection' import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker' /** * This class is for getRecommendation/listRecommendation API calls and its states * It does not contain UI/UX related logic */ /** * Commands as a level of indirection so that declare doesn't intercept any registrations for the * language server implementation. * * Otherwise you'll get: * "Unable to launch amazonq language server: Command "aws.amazonq.rejectCodeSuggestion" has already been declared by the Toolkit" */ function createCommands() { // below commands override VS Code inline completion commands const prevCommand = Commands.declare('editor.action.inlineSuggest.showPrevious', () => async () => { await RecommendationHandler.instance.showRecommendation(-1) }) const nextCommand = Commands.declare('editor.action.inlineSuggest.showNext', () => async () => { await RecommendationHandler.instance.showRecommendation(1) }) const rejectCommand = Commands.declare('aws.amazonq.rejectCodeSuggestion', () => async () => { telemetry.record({ traceId: TelemetryHelper.instance.traceId, }) await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') RecommendationHandler.instance.reportUserDecisions(-1) await Commands.tryExecute('aws.amazonq.refreshAnnotation') }) return { prevCommand, nextCommand, rejectCommand, } } const lock = new AsyncLock({ maxPending: 1 }) export class RecommendationHandler { public lastInvocationTime: number // TODO: remove this requestId public requestId: string private nextToken: string private cancellationToken: vscode.CancellationTokenSource private _onDidReceiveRecommendation: vscode.EventEmitter<void> = new vscode.EventEmitter<void>() public readonly onDidReceiveRecommendation: vscode.Event<void> = this._onDidReceiveRecommendation.event private inlineCompletionProvider?: CWInlineCompletionItemProvider private inlineCompletionProviderDisposable?: vscode.Disposable private reject: vscode.Disposable private next: vscode.Disposable private prev: vscode.Disposable private _timer?: NodeJS.Timer documentUri: vscode.Uri | undefined = undefined constructor() { this.requestId = '' this.nextToken = '' this.lastInvocationTime = performance.now() - CodeWhispererConstants.invocationTimeIntervalThreshold * 1000 this.cancellationToken = new vscode.CancellationTokenSource() this.prev = new vscode.Disposable(() => {}) this.next = new vscode.Disposable(() => {}) this.reject = new vscode.Disposable(() => {}) } static #instance: RecommendationHandler public static get instance() { return (this.#instance ??= new this()) } isValidResponse(): boolean { return session.recommendations.some((r) => r.content.trim() !== '') } async getServerResponse( triggerType: CodewhispererTriggerType, isManualTriggerOn: boolean, promise: Promise<any> ): Promise<any> { const timeoutMessage = hasVendedIamCredentials() ? 'Generate recommendation timeout.' : 'List recommendation timeout' if (isManualTriggerOn && triggerType === 'OnDemand' && hasVendedIamCredentials()) { return vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, title: CodeWhispererConstants.pendingResponse, cancellable: false, }, async () => { return await asyncCallWithTimeout( promise, timeoutMessage, CodeWhispererConstants.promiseTimeoutLimit * 1000 ) } ) } return await asyncCallWithTimeout(promise, timeoutMessage, CodeWhispererConstants.promiseTimeoutLimit * 1000) } async getTaskTypeFromEditorFileName(filePath: string): Promise<CodewhispererGettingStartedTask | undefined> { if (filePath.includes('CodeWhisperer_generate_suggestion')) { return 'autoTrigger' } else if (filePath.includes('CodeWhisperer_manual_invoke')) { return 'manualTrigger' } else if (filePath.includes('CodeWhisperer_use_comments')) { return 'commentAsPrompt' } else if (filePath.includes('CodeWhisperer_navigate_suggestions')) { return 'navigation' } else if (filePath.includes('Generate_unit_tests')) { return 'unitTest' } else { return undefined } } async getRecommendations( client: DefaultCodeWhispererClient, editor: vscode.TextEditor, triggerType: CodewhispererTriggerType, config: ConfigurationEntry, autoTriggerType?: CodewhispererAutomatedTriggerType, pagination: boolean = true, page: number = 0, generate: boolean = isIamConnection(AuthUtil.instance.conn) ): Promise<GetRecommendationsResponse> { let invocationResult: 'Succeeded' | 'Failed' = 'Failed' let errorMessage: string | undefined = undefined let errorCode: string | undefined = undefined if (!editor) { return Promise.resolve<GetRecommendationsResponse>({ result: invocationResult, errorMessage: errorMessage, recommendationCount: 0, }) } let recommendations: RecommendationsList = [] let requestId = '' let sessionId = '' let reason = '' let startTime = 0 let latency = 0 let nextToken = '' let shouldRecordServiceInvocation = true session.language = runtimeLanguageContext.getLanguageContext( editor.document.languageId, path.extname(editor.document.fileName) ).language session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName) if (pagination && !generate) { if (page === 0) { session.requestContext = await EditorContext.buildListRecommendationRequest( editor as vscode.TextEditor, this.nextToken, config.isSuggestionsWithCodeReferencesEnabled ) } else { session.requestContext = { request: { ...session.requestContext.request, // Putting nextToken assignment in the end so it overwrites the existing nextToken nextToken: this.nextToken, }, supplementalMetadata: session.requestContext.supplementalMetadata, } } } else { session.requestContext = await EditorContext.buildGenerateRecommendationRequest(editor as vscode.TextEditor) } const request = session.requestContext.request // record preprocessing end time TelemetryHelper.instance.setPreprocessEndTime() // set start pos for non pagination call or first pagination call if (!pagination || (pagination && page === 0)) { session.startPos = editor.selection.active session.startCursorOffset = editor.document.offsetAt(session.startPos) session.leftContextOfCurrentLine = EditorContext.getLeftContext(editor, session.startPos.line) session.triggerType = triggerType session.autoTriggerType = autoTriggerType /** * Validate request */ if (!EditorContext.validateRequest(request)) { getLogger().verbose('Invalid Request: %O', request) const languageName = request.fileContext.programmingLanguage.languageName if (!runtimeLanguageContext.isLanguageSupported(languageName)) { errorMessage = `${languageName} is currently not supported by Amazon Q inline suggestions` } return Promise.resolve<GetRecommendationsResponse>({ result: invocationResult, errorMessage: errorMessage, recommendationCount: 0, }) } } try { startTime = performance.now() this.lastInvocationTime = startTime const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request) const codewhispererPromise = pagination && !generate ? client.listRecommendations(mappedReq) : client.generateRecommendations(mappedReq) const resp = await this.getServerResponse(triggerType, config.isManualTriggerEnabled, codewhispererPromise) TelemetryHelper.instance.setSdkApiCallEndTime() latency = startTime !== 0 ? performance.now() - startTime : 0 if ('recommendations' in resp) { recommendations = (resp && resp.recommendations) || [] } else { recommendations = (resp && resp.completions) || [] } invocationResult = 'Succeeded' requestId = resp?.$response && resp?.$response?.requestId nextToken = resp?.nextToken ? resp?.nextToken : '' sessionId = resp?.$response?.httpResponse?.headers['x-amzn-sessionid'] TelemetryHelper.instance.setFirstResponseRequestId(requestId) if (page === 0) { session.setTimeToFirstRecommendation(performance.now()) } if (nextToken === '') { TelemetryHelper.instance.setAllPaginationEndTime() } } catch (error) { if (error instanceof CognitoCredentialsError) { shouldRecordServiceInvocation = false } if (latency === 0) { latency = startTime !== 0 ? performance.now() - startTime : 0 } getLogger().error('amazonq inline-suggest: Invocation Exception : %s', (error as Error).message) if (isAwsError(error)) { errorMessage = error.message requestId = error.requestId || '' errorCode = error.code reason = `CodeWhisperer Invocation Exception: ${error?.code ?? error?.name ?? 'unknown'}` await this.onThrottlingException(error, triggerType) if (error?.code === 'AccessDeniedException' && errorMessage?.includes('no identity-based policy')) { getLogger().error('amazonq inline-suggest: AccessDeniedException : %s', (error as Error).message) void vscode.window .showErrorMessage(`CodeWhisperer: ${error?.message}`, CodeWhispererConstants.settingsLearnMore) .then(async (resp) => { if (resp === CodeWhispererConstants.settingsLearnMore) { void openUrl(vscode.Uri.parse(CodeWhispererConstants.learnMoreUri)) } }) await vscode.commands.executeCommand('aws.amazonq.enableCodeSuggestions', false) } } else { errorMessage = error instanceof Error ? error.message : String(error) reason = error ? String(error) : 'unknown' } } finally { const timezone = Intl.DateTimeFormat().resolvedOptions().timeZone let msg = indent( `codewhisperer: request-id: ${requestId}, timestamp(epoch): ${Date.now()}, timezone: ${timezone}, datetime: ${new Date().toLocaleString([], { timeZone: timezone })}, vscode version: '${vscode.version}', extension version: '${extensionVersion}', filename: '${EditorContext.getFileName(editor)}', left context of line: '${session.leftContextOfCurrentLine}', line number: ${session.startPos.line}, character location: ${session.startPos.character}, latency: ${latency} ms. Recommendations:`, 4, true ).trimStart() for (const [index, item] of recommendations.entries()) { msg += `\n ${index.toString().padStart(2, '0')}: ${indent(item.content, 8, true).trim()}` session.requestIdList.push(requestId) } getLogger().debug(msg) if (invocationResult === 'Succeeded') { CodeWhispererCodeCoverageTracker.getTracker(session.language)?.incrementServiceInvocationCount() UserWrittenCodeTracker.instance.onQFeatureInvoked() } else { if ( (errorMessage?.includes(invalidCustomizationMessage) && errorCode === 'AccessDeniedException') || errorCode === 'ResourceNotFoundException' ) { getLogger() .debug(`The selected customization is no longer available. Retrying with the default model. Failed request id: ${requestId}`) await switchToBaseCustomizationAndNotify() await this.getRecommendations( client, editor, triggerType, config, autoTriggerType, pagination, page, true ) } } if (shouldRecordServiceInvocation) { TelemetryHelper.instance.recordServiceInvocationTelemetry( requestId, sessionId, session.recommendations.length + recommendations.length - 1, invocationResult, latency, session.language, session.taskType, reason, session.requestContext.supplementalMetadata ) } } if (this.isCancellationRequested()) { return Promise.resolve<GetRecommendationsResponse>({ result: invocationResult, errorMessage: errorMessage, recommendationCount: session.recommendations.length, }) } const typedPrefix = editor.document .getText(new vscode.Range(session.startPos, editor.selection.active)) .replace('\r\n', '\n') if (recommendations.length > 0) { TelemetryHelper.instance.setTypeAheadLength(typedPrefix.length) // mark suggestions that does not match typeahead when arrival as Discard // these suggestions can be marked as Showed if typeahead can be removed with new inline API for (const [i, r] of recommendations.entries()) { const recommendationIndex = i + session.recommendations.length if ( !r.content.startsWith(typedPrefix) && session.getSuggestionState(recommendationIndex) === undefined ) { session.setSuggestionState(recommendationIndex, 'Discard') } session.setCompletionType(recommendationIndex, r) } session.recommendations = pagination ? session.recommendations.concat(recommendations) : recommendations if (isInlineCompletionEnabled() && this.hasAtLeastOneValidSuggestion(typedPrefix)) { this._onDidReceiveRecommendation.fire() } } this.requestId = requestId session.sessionId = sessionId this.nextToken = nextToken // send Empty userDecision event if user receives no recommendations in this session at all. if (invocationResult === 'Succeeded' && nextToken === '') { // case 1: empty list of suggestion [] if (session.recommendations.length === 0) { session.requestIdList.push(requestId) // Received an empty list of recommendations TelemetryHelper.instance.recordUserDecisionTelemetryForEmptyList( session.requestIdList, sessionId, page, runtimeLanguageContext.getLanguageContext( editor.document.languageId, path.extname(editor.document.fileName) ).language, session.requestContext.supplementalMetadata ) } // case 2: non empty list of suggestion but with (a) empty content or (b) non-matching typeahead else if (!this.hasAtLeastOneValidSuggestion(typedPrefix)) { this.reportUserDecisions(-1) } } return Promise.resolve<GetRecommendationsResponse>({ result: invocationResult, errorMessage: errorMessage, recommendationCount: session.recommendations.length, }) } hasAtLeastOneValidSuggestion(typedPrefix: string): boolean { return session.recommendations.some((r) => r.content.trim() !== '' && r.content.startsWith(typedPrefix)) } cancelPaginatedRequest() { this.nextToken = '' this.cancellationToken.cancel() } isCancellationRequested() { return this.cancellationToken.token.isCancellationRequested } checkAndResetCancellationTokens() { if (this.isCancellationRequested()) { this.cancellationToken.dispose() this.cancellationToken = new vscode.CancellationTokenSource() this.nextToken = '' return true } return false } /** * Clear recommendation state */ clearRecommendations() { session.requestIdList = [] session.recommendations = [] session.suggestionStates = new Map<number, string>() session.completionTypes = new Map<number, CodewhispererCompletionType>() this.requestId = '' session.sessionId = '' this.nextToken = '' session.requestContext.supplementalMetadata = undefined } async clearInlineCompletionStates() { try { vsCodeState.isCodeWhispererEditing = false application()._clearCodeWhispererUIListener.fire() this.cancelPaginatedRequest() this.clearRecommendations() this.disposeInlineCompletion() await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar') // fix a regression that requires user to hit Esc twice to clear inline ghost text // because disposing a provider does not clear the UX if (isVscHavingRegressionInlineCompletionApi()) { await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') } } finally { this.clearRejectionTimer() } } reportDiscardedUserDecisions() { for (const [i, _] of session.recommendations.entries()) { session.setSuggestionState(i, 'Discard') } this.reportUserDecisions(-1) } /** * Emits telemetry reflecting user decision for current recommendation. */ reportUserDecisions(acceptIndex: number) { if (session.sessionId === '' || this.requestId === '') { return } TelemetryHelper.instance.recordUserDecisionTelemetry( session.requestIdList, session.sessionId, session.recommendations, acceptIndex, session.recommendations.length, session.completionTypes, session.suggestionStates, session.requestContext.supplementalMetadata ) if (isInlineCompletionEnabled()) { this.clearInlineCompletionStates().catch((e) => { getLogger().error('clearInlineCompletionStates failed: %s', (e as Error).message) }) } } hasNextToken(): boolean { return this.nextToken !== '' } canShowRecommendationInIntelliSense( editor: vscode.TextEditor, showPrompt: boolean = false, response: GetRecommendationsResponse ): boolean { const reject = () => { this.reportUserDecisions(-1) } if (!this.isValidResponse()) { if (showPrompt) { void showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 3000) } reject() return false } // do not show recommendation if cursor is before invocation position // also mark as Discard if (editor.selection.active.isBefore(session.startPos)) { for (const [i, _] of session.recommendations.entries()) { session.setSuggestionState(i, 'Discard') } reject() return false } // do not show recommendation if typeahead does not match // also mark as Discard const typedPrefix = editor.document.getText( new vscode.Range( session.startPos.line, session.startPos.character, editor.selection.active.line, editor.selection.active.character ) ) if (!session.recommendations[0].content.startsWith(typedPrefix.trimStart())) { for (const [i, _] of session.recommendations.entries()) { session.setSuggestionState(i, 'Discard') } reject() return false } return true } async onThrottlingException(awsError: AWSError, triggerType: CodewhispererTriggerType) { if ( awsError.code === 'ThrottlingException' && awsError.message.includes(CodeWhispererConstants.throttlingMessage) ) { if (triggerType === 'OnDemand') { void vscode.window.showErrorMessage(CodeWhispererConstants.freeTierLimitReached) } vsCodeState.isFreeTierLimitReached = true } } public disposeInlineCompletion() { this.inlineCompletionProviderDisposable?.dispose() this.inlineCompletionProvider = undefined } private disposeCommandOverrides() { this.prev.dispose() this.reject.dispose() this.next.dispose() } // These commands override the vs code inline completion commands // They are subscribed when suggestion starts and disposed when suggestion is accepted/rejected // to avoid impacting other plugins or user who uses this API private registerCommandOverrides() { const { prevCommand, nextCommand, rejectCommand } = createCommands() this.prev = prevCommand.register() this.next = nextCommand.register() this.reject = rejectCommand.register() } subscribeSuggestionCommands() { this.disposeCommandOverrides() this.registerCommandOverrides() globals.context.subscriptions.push(this.prev) globals.context.subscriptions.push(this.next) globals.context.subscriptions.push(this.reject) } async showRecommendation(indexShift: number, noSuggestionVisible: boolean = false) { await lock.acquire(updateInlineLockKey, async () => { if (!vscode.window.state.focused) { this.reportDiscardedUserDecisions() return } const inlineCompletionProvider = new CWInlineCompletionItemProvider( this.inlineCompletionProvider?.getActiveItemIndex, indexShift, session.recommendations, this.requestId, session.startPos, this.nextToken ) this.inlineCompletionProviderDisposable?.dispose() // when suggestion is active, registering a new provider will let VS Code invoke inline API automatically this.inlineCompletionProviderDisposable = vscode.languages.registerInlineCompletionItemProvider( Object.assign([], CodeWhispererConstants.platformLanguageIds), inlineCompletionProvider ) this.inlineCompletionProvider = inlineCompletionProvider if (isVscHavingRegressionInlineCompletionApi() && !noSuggestionVisible) { // fix a regression in new VS Code when disposing and re-registering // a new provider does not auto refresh the inline suggestion widget // by manually refresh it await vscode.commands.executeCommand('editor.action.inlineSuggest.hide') await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger') } if (noSuggestionVisible) { await vscode.commands.executeCommand(`editor.action.inlineSuggest.trigger`) this.sendPerceivedLatencyTelemetry() } }) } async onEditorChange() { this.reportUserDecisions(-1) } async onFocusChange() { this.reportUserDecisions(-1) } async onCursorChange(e: vscode.TextEditorSelectionChangeEvent) { // we do not want to reset the states for keyboard events because they can be typeahead if ( e.kind !== vscode.TextEditorSelectionChangeKind.Keyboard && vscode.window.activeTextEditor === e.textEditor ) { application()._clearCodeWhispererUIListener.fire() // when cursor change due to mouse movement we need to reset the active item index for inline if (e.kind === vscode.TextEditorSelectionChangeKind.Mouse) { this.inlineCompletionProvider?.clearActiveItemIndex() } } } isSuggestionVisible(): boolean { return this.inlineCompletionProvider?.getActiveItemIndex !== undefined } async tryShowRecommendation() { const editor = vscode.window.activeTextEditor if (editor === undefined) { return } if (this.isSuggestionVisible()) { // do not force refresh the tooltip to avoid suggestion "flashing" return } if ( editor.selection.active.isBefore(session.startPos) || editor.document.uri.fsPath !== this.documentUri?.fsPath ) { for (const [i, _] of session.recommendations.entries()) { session.setSuggestionState(i, 'Discard') } this.reportUserDecisions(-1) } else if (session.recommendations.length > 0) { await this.showRecommendation(0, true) } } private clearRejectionTimer() { if (this._timer !== undefined) { clearInterval(this._timer) this._timer = undefined } } private sendPerceivedLatencyTelemetry() { if (vscode.window.activeTextEditor) { const languageContext = runtimeLanguageContext.getLanguageContext( vscode.window.activeTextEditor.document.languageId, vscode.window.activeTextEditor.document.fileName.substring( vscode.window.activeTextEditor.document.fileName.lastIndexOf('.') + 1 ) ) telemetry.codewhisperer_perceivedLatency.emit({ codewhispererRequestId: this.requestId, codewhispererSessionId: session.sessionId, codewhispererTriggerType: session.triggerType, codewhispererCompletionType: session.getCompletionType(0), codewhispererCustomizationArn: getSelectedCustomization().arn, codewhispererLanguage: languageContext.language, duration: performance.now() - this.lastInvocationTime, passive: true, credentialStartUrl: AuthUtil.instance.startUrl, result: 'Succeeded', }) } } }