undo more unnecessary changes

Author: liramon1

packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts

@@ -6,7 +6,7 @@
 import assert from 'assert'
 import * as sinon from 'sinon'
 import { assertTelemetryCurried } from 'aws-core-vscode/test'
-import { CodeWhispererTracker } from 'aws-core-vscode/codewhisperer'
+import { AuthUtil, CodeWhispererTracker } from 'aws-core-vscode/codewhisperer'
 import { resetCodeWhispererGlobalVariables, createAcceptedSuggestionEntry } from 'aws-core-vscode/test'
 import { globals } from 'aws-core-vscode/shared'
 
@@ -93,8 +93,7 @@ describe('codewhispererTracker', function () {
                 codewhispererModificationPercentage: 1,
                 codewhispererCompletionType: 'Line',
                 codewhispererLanguage: 'java',
-                // TODO: fix this
-                // credentialStartUrl: AuthUtil.instance.connection?.startUrl,
+                credentialStartUrl: AuthUtil.instance.connection?.startUrl,
                 codewhispererCharactersAccepted: suggestion.originalString.length,
                 codewhispererCharactersModified: 0,
             })

packages/core/src/amazonqFeatureDev/client/featureDev.ts

@@ -51,7 +51,7 @@ const writeAPIRetryOptions = {
 
 // Create a client for featureDev proxy client based off of aws sdk v2
 export async function createFeatureDevProxyClient(options?: Partial<ServiceOptions>): Promise<FeatureDevProxyClient> {
-    const bearerToken = await AuthUtil.instance.getBearerToken()
+    const bearerToken = await AuthUtil.instance.getToken()
     const cwsprConfig = getCodewhispererConfig()
     return (await globals.sdkClientBuilder.createAwsService(
         Service,

packages/core/src/codewhisperer/client/codewhisperer.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AWSError, HttpRequest, Service } from 'aws-sdk'
+import { AWSError, Credentials, Service } from 'aws-sdk'
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
 import * as CodeWhispererUserClient from './codewhispereruserclient'
@@ -13,8 +13,8 @@ import { hasVendedIamCredentials } from '../../auth/auth'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { PromiseResult } from 'aws-sdk/lib/request'
 import { AuthUtil } from '../util/authUtil'
-import userApiConfig = require('./user-service-2.json')
 import apiConfig = require('./service-2.json')
+import userApiConfig = require('./user-service-2.json')
 import { session } from '../util/codeWhispererSession'
 import { getLogger } from '../../shared/logger/logger'
 import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
@@ -84,13 +84,14 @@ export type Imports = CodeWhispererUserClient.Imports
 export class DefaultCodeWhispererClient {
     private async createSdkClient(): Promise<CodeWhispererClient> {
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
+        const credential = await AuthUtil.instance.getIamCredential()
         const cwsprConfig = getCodewhispererConfig()
         return (await globals.sdkClientBuilder.createAwsService(
             Service,
             {
                 apiConfig: apiConfig,
                 region: cwsprConfig.region,
-                credentials: undefined,
+                credentials: new Credentials({ accessKeyId: credential.accessKeyId, secretAccessKey: credential.secretAccessKey }),
                 endpoint: cwsprConfig.endpoint,
                 onRequestSetup: [
                     (req) => {
@@ -108,9 +109,9 @@ export class DefaultCodeWhispererClient {
                                     resp.error?.code === 'AccessDeniedException' &&
                                     resp.error.message.match(/expired/i)
                                 ) {
-                                    // AuthUtil.instance.reauthenticate().catch((e) => {
-                                    //     getLogger().error('reauthenticate failed: %s', (e as Error).message)
-                                    // })
+                                    AuthUtil.instance.reauthenticate()?.catch((e) => {
+                                        getLogger().error('reauthenticate failed: %s', (e as Error).message)
+                                    })
                                     resp.error.retryable = true
                                 }
                             })
@@ -125,7 +126,7 @@ export class DefaultCodeWhispererClient {
     async createUserSdkClient(maxRetries?: number): Promise<CodeWhispererUserClient> {
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
         session.setFetchCredentialStart()
-        const credential = await AuthUtil.instance.getBearerToken()
+        const bearerToken = await AuthUtil.instance.getToken()
 
         session.setSdkApiCallStart()
         const cwsprConfig = getCodewhispererConfig()
@@ -137,9 +138,9 @@ export class DefaultCodeWhispererClient {
                 endpoint: cwsprConfig.endpoint,
                 maxRetries: maxRetries,
                 onRequestSetup: [
-                    (req: any) => {
-                        req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => {
-                            httpRequest.headers['Authorization'] = `Bearer ${credential}`
+                    (req) => {
+                        req.on('build', ({ httpRequest }) => {
+                            httpRequest.headers['Authorization'] = `Bearer ${bearerToken}`
                         })
                         if (req.operation === 'generateCompletions') {
                             req.on('build', () => {

packages/core/src/codewhisperer/region/regionProfileManager.ts

@@ -428,12 +428,12 @@ export class RegionProfileManager {
     async _createQClient(region: string, endpoint: string): Promise<CodeWhispererUserClient> {
         let authConfig: ServiceOptions = {}
         if (this.authProvider.isSsoSession()) {
-            const credential = await this.authProvider.getBearerToken()
+            const token = await this.authProvider.getToken()
             authConfig = {
                 onRequestSetup: [
                     (req: any) => {
                         req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => {
-                            httpRequest.headers['Authorization'] = `Bearer ${credential}`
+                            httpRequest.headers['Authorization'] = `Bearer ${token}`
                         })
                     },
                 ],

packages/core/src/codewhisperer/util/authUtil.ts

@@ -61,7 +61,7 @@ export interface IAuthProvider {
     isIdcConnection(): boolean
     isSsoSession(): boolean
     isIamSession(): boolean
-    getBearerToken(): Promise<string>
+    getToken(): Promise<string>
     getIamCredential(): Promise<IamCredentials>
     readonly profileName: string
     readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string }
@@ -202,7 +202,7 @@ export class AuthUtil implements IAuthProvider {
         return response
     }
 
-    async getBearerToken() {
+    async getToken() {
         if (this.isSsoSession()) {
             return (await (this.session as SsoLogin).getCredential()).credential
         } else {

packages/core/src/shared/clients/codewhispererChatClient.ts

@@ -10,7 +10,7 @@ import { getUserAgent } from '../telemetry/util'
 
 // Create a client for featureDev streaming based off of aws sdk v3
 export async function createCodeWhispererChatStreamingClient(): Promise<CodeWhispererStreaming> {
-    const bearerToken = await AuthUtil.instance.getBearerToken()
+    const bearerToken = await AuthUtil.instance.getToken()
     const cwsprConfig = getCodewhispererConfig()
     const streamingClient = new CodeWhispererStreaming({
         region: cwsprConfig.region,