start modifying auth2 consumers

liramon1 · liramon1 · commit ff181f49c734 · 2025-06-16T15:31:07.000-04:00
diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts
@@ -51,18 +51,22 @@ 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.getToken()
+    const credential = await AuthUtil.instance.getCredential()
+    // TODO: handle IAM credentials when IAM version of API JSON file is generated
+    if (credential !== 'string') {
+        throw new Error('Feature dev does not support IAM credentials')
+    }
     const cwsprConfig = getCodewhispererConfig()
     return (await globals.sdkClientBuilder.createAwsService(
         Service,
         {
             apiConfig: apiConfig,
             region: cwsprConfig.region,
             endpoint: cwsprConfig.endpoint,
-            token: new Token({ token: bearerToken }),
             httpOptions: {
                 connectTimeout: 10000, // 10 seconds, 3 times P99 API latency
             },
+            token: new Token({ token: credential }),
             ...options,
         } as ServiceOptions,
         undefined
diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts
@@ -61,6 +61,7 @@ import { ToolkitError } from '../shared/errors'
 import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
 import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache'
 import { VSCODE_EXTENSION_ID } from '../shared/extensions'
+import { IamCredentials } from '@aws/language-server-runtimes-types'
 
 export const notificationTypes = {
     updateIamCredential: new RequestType<UpdateCredentialsParams, ResponseMessage, Error>(
@@ -217,15 +218,15 @@ export class LanguageClientAuth {
         return { profile, ssoSession }
     }
 
-    updateBearerToken(request: UpdateCredentialsParams) {
+    updateBearerToken(request: UpdateCredentialsParams | undefined) {
         return this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request)
     }
 
     deleteBearerToken() {
         return this.client.sendNotification(bearerCredentialsDeleteNotificationType.method)
     }
 
-    updateIamCredential(request: UpdateCredentialsParams) {
+    updateIamCredential(request: UpdateCredentialsParams | undefined) {
         return this.client.sendRequest(iamCredentialsUpdateRequestType.method, request)
     }
 
@@ -283,6 +284,10 @@ export abstract class BaseLogin {
     abstract reauthenticate(): Promise<GetSsoTokenResult | GetIamCredentialResult | undefined>
     abstract logout(): void
     abstract restore(): void
+    abstract getCredential(): Promise<{
+        credential: string | IamCredentials
+        updateCredentialsParams: UpdateCredentialsParams
+    }>
 
     get data() {
         return this._data
@@ -402,11 +407,11 @@ export class SsoLogin extends BaseLogin {
      * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API
      * with encrypted token
      */
-    async getToken() {
+    async getCredential() {
         const response = await this._getSsoToken(false)
         const accessToken = await this.decrypt(response.ssoToken.accessToken)
         return {
-            token: accessToken,
+            credential: accessToken,
             updateCredentialsParams: response.updateCredentialsParams,
         }
     }
@@ -548,18 +553,17 @@ export class IamLogin extends BaseLogin {
      * Returns both the decrypted IAM credential and the payload to send to the `updateCredentials` LSP API
      * with encrypted credential
      */
-    async getCredentials() {
+    async getCredential() {
         const response = await this._getIamCredential(false)
-        const accessKey = await this.decrypt(response.credentials.accessKeyId)
-        const secretKey = await this.decrypt(response.credentials.secretAccessKey)
-        let sessionToken: string | undefined
-        if (response.credentials.sessionToken) {
-            sessionToken = await this.decrypt(response.credentials.sessionToken)
+        const credentials: IamCredentials = {
+            accessKeyId: await this.decrypt(response.credentials.accessKeyId),
+            secretAccessKey: await this.decrypt(response.credentials.secretAccessKey),
+            sessionToken: response.credentials.sessionToken
+                ? await this.decrypt(response.credentials.sessionToken)
+                : undefined,
         }
         return {
-            accessKey: accessKey,
-            secretKey: secretKey,
-            sessionToken: sessionToken,
+            credential: credentials,
             updateCredentialsParams: response.updateCredentialsParams,
         }
     }
diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { AWSError, Credentials, Service } from 'aws-sdk'
+import { AWSError, HttpRequest, 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 apiConfig = require('./service-2.json')
 import userApiConfig = require('./user-service-2.json')
+import apiConfig = require('./service-2.json')
 import { session } from '../util/codeWhispererSession'
 import { getLogger } from '../../shared/logger/logger'
 import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
@@ -83,8 +83,6 @@ export type Imports = CodeWhispererUserClient.Imports
 
 export class DefaultCodeWhispererClient {
     private async createSdkClient(): Promise<CodeWhispererClient> {
-        throw new Error('Do not call this function until IAM is supported by LSP identity server')
-
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
         const cwsprConfig = getCodewhispererConfig()
         return (await globals.sdkClientBuilder.createAwsService(
@@ -127,7 +125,11 @@ export class DefaultCodeWhispererClient {
     async createUserSdkClient(maxRetries?: number): Promise<CodeWhispererUserClient> {
         const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
         session.setFetchCredentialStart()
-        const bearerToken = await AuthUtil.instance.getToken()
+        const credential = await AuthUtil.instance.getCredential()
+        if (typeof credential !== 'string') {
+            throw new TypeError('Cannot create user SDK client from IAM credentials')
+        }
+
         session.setSdkApiCallStart()
         const cwsprConfig = getCodewhispererConfig()
         return (await globals.sdkClientBuilder.createAwsService(
@@ -137,11 +139,10 @@ export class DefaultCodeWhispererClient {
                 region: cwsprConfig.region,
                 endpoint: cwsprConfig.endpoint,
                 maxRetries: maxRetries,
-                credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }),
                 onRequestSetup: [
-                    (req) => {
-                        req.on('build', ({ httpRequest }) => {
-                            httpRequest.headers['Authorization'] = `Bearer ${bearerToken}`
+                    (req: any) => {
+                        req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => {
+                            httpRequest.headers['Authorization'] = `Bearer ${credential}`
                         })
                         if (req.operation === 'generateCompletions') {
                             req.on('build', () => {
@@ -160,6 +161,88 @@ export class DefaultCodeWhispererClient {
         return AuthUtil.instance.isConnected() // TODO: Handle IAM credentials
     }
 
+    // private async createServiceSdkClient(credential: IamCredentials): Promise<CodeWhispererClient> {
+    //     const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
+    //     const cwsprConfig = getCodewhispererConfig()
+    //     return (await globals.sdkClientBuilder.createAwsService(
+    //         Service,
+    //         {
+    //             apiConfig: apiConfig,
+    //             region: cwsprConfig.region,
+    //             credentials: undefined,
+    //             endpoint: cwsprConfig.endpoint,
+    //             onRequestSetup: [
+    //                 (req) => {
+    //                     if (req.operation === 'listRecommendations') {
+    //                         req.on('build', () => {
+    //                             req.httpRequest.headers['x-amzn-codewhisperer-optout'] = `${isOptedOut}`
+    //                         })
+    //                     }
+    //                     // This logic is for backward compatability with legacy SDK v2 behavior for refreshing
+    //                     // credentials. Once the Toolkit adds a file watcher for credentials it won't be needed.
+
+    //                     if (hasVendedIamCredentials()) {
+    //                         req.on('retry', (resp) => {
+    //                             if (
+    //                                 resp.error?.code === 'AccessDeniedException' &&
+    //                                 resp.error.message.match(/expired/i)
+    //                             ) {
+    //                                 // AuthUtil.instance.reauthenticate().catch((e) => {
+    //                                 //     getLogger().error('reauthenticate failed: %s', (e as Error).message)
+    //                                 // })
+    //                                 resp.error.retryable = true
+    //                             }
+    //                         })
+    //                     }
+    //                 },
+    //             ],
+    //         } as ServiceOptions,
+    //         undefined
+    //     )) as CodeWhispererClient
+    // }
+
+    // private async createUserServiceSdkClient(
+    //     credential: string,
+    //     maxRetries?: number
+    // ): Promise<CodeWhispererUserClient> {
+    //     const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled()
+    //     session.setFetchCredentialStart()
+    //     session.setSdkApiCallStart()
+    //     const cwsprConfig = getCodewhispererConfig()
+    //     return (await globals.sdkClientBuilder.createAwsService(
+    //         Service,
+    //         {
+    //             apiConfig: userApiConfig,
+    //             region: cwsprConfig.region,
+    //             endpoint: cwsprConfig.endpoint,
+    //             maxRetries: maxRetries,
+    //             onRequestSetup: [
+    //                 (req: any) => {
+    //                     req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => {
+    //                         httpRequest.headers['Authorization'] = `Bearer ${credential}`
+    //                     })
+    //                     if (req.operation === 'generateCompletions') {
+    //                         req.on('build', () => {
+    //                             req.httpRequest.headers['x-amzn-codewhisperer-optout'] = `${isOptedOut}`
+    //                             req.httpRequest.headers['Connection'] = keepAliveHeader
+    //                         })
+    //                     }
+    //                 },
+    //             ],
+    //         } as ServiceOptions,
+    //         undefined
+    //     )) as CodeWhispererUserClient
+    // }
+
+    // async createSdkClient(maxRetries?: number): Promise<CodeWhispererUserClient | CodeWhispererClient> {
+    //     const credential = await AuthUtil.instance.getCredential()
+    //     if (typeof credential === 'string') {
+    //         return this.createUserServiceSdkClient(credential, maxRetries)
+    //     } else {
+    //         return this.createServiceSdkClient(credential)
+    //     }
+    // }
+
     public async generateRecommendations(
         request: GenerateRecommendationsRequest
     ): Promise<GenerateRecommendationsResponse> {
diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts
@@ -11,9 +11,10 @@ import { showConfirmationMessage } from '../../shared/utilities/messages'
 import globals from '../../shared/extensionGlobals'
 import { once } from '../../shared/utilities/functionUtils'
 import CodeWhispererUserClient from '../client/codewhispereruserclient'
-import { Credentials, Service } from 'aws-sdk'
+import { Credentials, HttpRequest, Service } from 'aws-sdk'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
-import userApiConfig = require('../client/user-service-2.json')
+import tokenApiConfig = require('../client/user-service-2.json')
+import iamApiConfig = require('../client/service-2.json')
 import { createConstantMap } from '../../shared/utilities/tsUtils'
 import { getLogger } from '../../shared/logger/logger'
 import { pageableToCollection } from '../../shared/utilities/collectionUtils'
@@ -425,19 +426,31 @@ export class RegionProfileManager {
 
     // Visible for testing only, do not use this directly, please use createQClient(profile)
     async _createQClient(region: string, endpoint: string): Promise<CodeWhispererUserClient> {
-        const token = await this.authProvider.getToken()
+        const credential = await this.authProvider.getCredential()
+        const authConfig: ServiceOptions =
+            typeof credential === 'string'
+                ? {
+                      onRequestSetup: [
+                          (req: any) => {
+                              req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => {
+                                  httpRequest.headers['Authorization'] = `Bearer ${credential}`
+                              })
+                          },
+                      ],
+                  }
+                : {
+                      credentials: new Credentials({
+                          accessKeyId: credential.accessKeyId,
+                          secretAccessKey: credential.secretAccessKey,
+                          sessionToken: credential.sessionToken,
+                      }),
+                  }
+        const apiConfig = typeof credential === 'string' ? tokenApiConfig : iamApiConfig
         const serviceOption: ServiceOptions = {
-            apiConfig: userApiConfig,
+            apiConfig: apiConfig,
             region: region,
             endpoint: endpoint,
-            credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }),
-            onRequestSetup: [
-                (req) => {
-                    req.on('build', ({ httpRequest }) => {
-                        httpRequest.headers['Authorization'] = `Bearer ${token}`
-                    })
-                },
-            ],
+            ...authConfig,
         } as ServiceOptions
 
         const c = (await globals.sdkClientBuilder.createAwsService(
diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -44,6 +44,7 @@ import {
     GetSsoTokenResult,
     GetIamCredentialResult,
     SsoTokenSourceKind,
+    IamCredentials,
 } from '@aws/language-server-runtimes/server-interface'
 
 const localize = nls.loadMessageBundle()
@@ -59,7 +60,8 @@ export interface IAuthProvider {
     isBuilderIdConnection(): boolean
     isIdcConnection(): boolean
     isSsoSession(): boolean
-    getToken(): Promise<string>
+    isIamSession(): boolean
+    getCredential(): Promise<string | IamCredentials>
     readonly profileName: string
     readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string }
 }
@@ -199,11 +201,11 @@ export class AuthUtil implements IAuthProvider {
         return response
     }
 
-    async getToken() {
-        if (this.isSsoSession()) {
-            return (await (this.session as SsoLogin).getToken()).token
+    async getCredential() {
+        if (this.session) {
+            return (await this.session.getCredential()).credential
         } else {
-            throw new ToolkitError('Cannot get token for non-SSO session.')
+            throw new ToolkitError('Cannot get credential without logging in.')
         }
     }
 
@@ -336,11 +338,12 @@ export class AuthUtil implements IAuthProvider {
 
     private async stateChangeHandler(e: AuthStateEvent) {
         if (e.state === 'refreshed') {
-            const params = this.isSsoSession()
-                ? (await (this.session as SsoLogin).getToken()).updateCredentialsParams
-                : undefined
-            await this.lspAuth.updateBearerToken(params!)
-            return
+            const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined
+            if (this.isSsoSession()) {
+                await this.lspAuth.updateBearerToken(params)
+            } else if (this.isIamSession()) {
+                await this.lspAuth.updateIamCredential(params)
+            }
         } else {
             this.logger.info(`codewhisperer: connection changed to ${e.state}`)
             await this.refreshState(e.state)
@@ -356,8 +359,12 @@ export class AuthUtil implements IAuthProvider {
             }
         }
         if (state === 'connected') {
-            const bearerTokenParams = (await (this.session as SsoLogin).getToken()).updateCredentialsParams
-            await this.lspAuth.updateBearerToken(bearerTokenParams)
+            const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined
+            if (this.isSsoSession()) {
+                await this.lspAuth.updateBearerToken(params)
+            } else if (this.isIamSession()) {
+                await this.lspAuth.updateIamCredential(params)
+            }
 
             if (this.isIdcConnection()) {
                 await this.regionProfileManager.restoreProfileSelection()
@@ -400,7 +407,7 @@ export class AuthUtil implements IAuthProvider {
                 credentialStartUrl: AuthUtil.instance.connection?.startUrl,
                 awsRegion: AuthUtil.instance.connection?.region,
             }
-        } else if (!AuthUtil.instance.isSsoSession) {
+        } else if (this.isIamSession()) {
             return {
                 credentialSourceId: 'sharedCredentials',
             }
diff --git a/packages/core/src/shared/clients/codewhispererChatClient.ts b/packages/core/src/shared/clients/codewhispererChatClient.ts
