fix(codewhisperer): latency caused by credentials (#2946)

Fixes performance issue from cognito credentials

Author: zixlin7

.changes/next-release/Bug Fix-225a8639-7052-4759-9d2f-1823e527728c.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "latency caused by unnecessary refreshing of credentials"
+}

src/codewhisperer/client/codewhisperer.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AWSError, Service } from 'aws-sdk'
+import { AWSError, CognitoIdentityCredentials, Service } from 'aws-sdk'
 import apiConfig = require('./service-2.json')
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
@@ -39,8 +39,23 @@ export type ListCodeScanFindingsRequest = Readonly<CodeWhispererClient.ListCodeS
 export type ArtifactType = Readonly<CodeWhispererClient.ArtifactType>
 export type ArtifactMap = Readonly<CodeWhispererClient.ArtifactMap>
 export class DefaultCodeWhispererClient {
+    private credentials?: CognitoIdentityCredentials
+
     private async createSdkClient(): Promise<CodeWhispererClient> {
-        const credentials = !isCloud9() ? await getCognitoCredentials() : undefined
+        try {
+            if (!this.credentials && !isCloud9()) {
+                this.credentials = await getCognitoCredentials()
+            }
+            if (this.credentials && this.credentials.needsRefresh()) {
+                await new Promise<void>((resolve, reject) =>
+                    this.credentials?.refresh(e => (e ? reject(e) : resolve()))
+                )
+            }
+        } catch (e) {
+            getLogger().debug('Error when getting or refreshing Cognito credentials', e)
+            throw new CognitoCredentialsError('Error when getting or refreshing Cognito credentials')
+        }
+
         const accessToken = globals.context.globalState.get<string | undefined>(CodeWhispererConstants.accessToken)
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
 
@@ -49,7 +64,7 @@ export class DefaultCodeWhispererClient {
             {
                 apiConfig: apiConfig,
                 region: CodeWhispererConstants.region,
-                credentials: credentials,
+                credentials: this.credentials,
                 endpoint: CodeWhispererConstants.endpoint,
                 onRequestSetup: [
                     req => {
@@ -132,3 +147,5 @@ export class DefaultCodeWhispererClient {
         return (await this.createSdkClient()).listCodeScanFindings(request).promise()
     }
 }
+
+export class CognitoCredentialsError extends Error {}

src/codewhisperer/service/recommendationHandler.ts

@@ -5,7 +5,12 @@
 
 import * as vscode from 'vscode'
 import { extensionVersion } from '../../shared/vscode/env'
-import { RecommendationsList, DefaultCodeWhispererClient, Recommendation } from '../client/codewhisperer'
+import {
+    RecommendationsList,
+    DefaultCodeWhispererClient,
+    Recommendation,
+    CognitoCredentialsError,
+} from '../client/codewhisperer'
 import * as EditorContext from '../util/editorContext'
 import * as CodeWhispererConstants from '../models/constants'
 import { ConfigurationEntry } from '../models/model'
@@ -192,6 +197,9 @@ export class RecommendationHandler {
                 }
             }
         } catch (error) {
+            if (error instanceof CognitoCredentialsError) {
+                shouldRecordServiceInvocation = false
+            }
             if (latency === 0) {
                 latency = startTime !== 0 ? performance.now() - startTime : 0
             }

src/codewhisperer/util/cognitoIdentity.ts

@@ -29,7 +29,6 @@ export const getCognitoCredentials = async (): Promise<CognitoIdentityCredential
         }
 
         const credentials = new CognitoIdentityCredentials({ IdentityId: identityId }, { region })
-        await credentials.getPromise()
         return credentials
     } catch (err) {
         getLogger().error(`Failed to initialize Cognito identity for CodeWhisperer: ${err} in region: ${region}`)