fix(amazonq): handle IAM credentials expiration field to be aws sdk versions compatible and add refresh logic to codewhisperer IAM client (#2349)

## Problem
Chat and InlineSuggestions doesnt work in Sagemaker instances after initial credential expiration since refresh logic isnt using expiration field properly

## Solution
- Standardizes IAM credentials expiration field handling across the codebase to support both AWS SDK v2 (`expireTime`) and v3 (`expiration`) formats
- Add credential fetch and refresh lazy loading logic to Codewhisperer IAM client to ensure fresh credentials are used

Author: parameja1

core/aws-lsp-core/src/credentials/credentialsProvider.ts

@@ -4,6 +4,7 @@ export interface IamCredentials {
     accessKeyId: string
     secretAccessKey: string
     sessionToken?: string
+    expiration?: Date // v3 format
 }
 
 export interface BearerToken {

core/aws-lsp-core/src/credentials/ideCredentialsProvider.test.ts

@@ -0,0 +1,48 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as assert from 'assert'
+import { IdeCredentialsProvider } from './ideCredentialsProvider'
+import { IamCredentials } from './credentialsProvider'
+import { Connection } from 'vscode-languageserver'
+import * as sinon from 'sinon'
+
+describe('IdeCredentialsProvider', function () {
+    let provider: IdeCredentialsProvider
+    let mockConnection: sinon.SinonStubbedInstance<Connection>
+
+    beforeEach(function () {
+        mockConnection = {
+            console: {
+                info: sinon.stub(),
+                log: sinon.stub(),
+                warn: sinon.stub(),
+                error: sinon.stub(),
+            },
+        } as any
+        provider = new IdeCredentialsProvider(mockConnection as any)
+    })
+
+    describe('validateIamCredentialsFields', function () {
+        it('throws error when accessKeyId is missing', function () {
+            const credentials = {
+                secretAccessKey: 'secret',
+            } as IamCredentials
+
+            assert.throws(() => provider['validateIamCredentialsFields'](credentials), /Missing property: accessKeyId/)
+        })
+
+        it('throws error when secretAccessKey is missing', function () {
+            const credentials = {
+                accessKeyId: 'key',
+            } as IamCredentials
+
+            assert.throws(
+                () => provider['validateIamCredentialsFields'](credentials),
+                /Missing property: secretAccessKey/
+            )
+        })
+    })
+})

core/aws-lsp-core/src/credentials/ideCredentialsProvider.ts

@@ -58,7 +58,15 @@ export class IdeCredentialsProvider implements CredentialsProvider {
             credentialsProtocolMethodNames.iamCredentialsUpdate,
             async (request: UpdateCredentialsRequest) => {
                 try {
-                    const iamCredentials = await this.decodeCredentialsRequestToken<IamCredentials>(request)
+                    const rawCredentials = await this.decodeCredentialsRequestToken<
+                        IamCredentials & { expireTime?: Date }
+                    >(request)
+
+                    // Normalize legacy expireTime field to standard expiration field
+                    const iamCredentials: IamCredentials = {
+                        ...rawCredentials,
+                        expiration: rawCredentials.expiration || rawCredentials.expireTime,
+                    }
 
                     this.validateIamCredentialsFields(iamCredentials)
 

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts

@@ -78,8 +78,13 @@ export class AmazonQIAMServiceManager extends BaseAmazonQServiceManager<
         return this.cachedStreamingClient
     }
 
-    public handleOnCredentialsDeleted(_type: CredentialsType): void {
-        return
+    public handleOnCredentialsDeleted(type: CredentialsType): void {
+        if (type === 'iam') {
+            this.cachedCodewhispererService?.abortInflightRequests()
+            this.cachedCodewhispererService = undefined
+            this.cachedStreamingClient?.abortInflightRequests()
+            this.cachedStreamingClient = undefined
+        }
     }
 
     public override handleOnUpdateConfiguration(

server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts

@@ -234,7 +234,46 @@ export class CodeWhispererServiceIAM extends CodeWhispererServiceBase {
             region: this.codeWhispererRegion,
             endpoint: this.codeWhispererEndpoint,
             credentialProvider: new CredentialProviderChain([
-                () => credentialsProvider.getCredentials('iam') as Credentials,
+                () => {
+                    const credentials = new Credentials({
+                        accessKeyId: '',
+                        secretAccessKey: '',
+                        sessionToken: '',
+                    })
+
+                    credentials.get = callback => {
+                        logging.info('CodeWhispererServiceIAM: Attempting to get credentials')
+
+                        Promise.resolve(credentialsProvider.getCredentials('iam'))
+                            .then((creds: any) => {
+                                logging.info('CodeWhispererServiceIAM: Successfully got credentials')
+
+                                credentials.accessKeyId = creds.accessKeyId as string
+                                credentials.secretAccessKey = creds.secretAccessKey as string
+                                credentials.sessionToken = creds.sessionToken as string
+                                credentials.expireTime = creds.expiration as Date
+                                callback()
+                            })
+                            .catch(err => {
+                                logging.error(`CodeWhispererServiceIAM: Failed to get credentials: ${err.message}`)
+                                callback(err)
+                            })
+                    }
+
+                    credentials.needsRefresh = () => {
+                        return (
+                            !credentials.accessKeyId ||
+                            !credentials.secretAccessKey ||
+                            (credentials.expireTime && credentials.expireTime.getTime() - Date.now() < 60000)
+                        ) // 1 min buffer
+                    }
+
+                    credentials.refresh = callback => {
+                        credentials.get(callback)
+                    }
+
+                    return credentials
+                },
             ]),
         }
         this.client = createCodeWhispererSigv4Client(options, sdkInitializator, logging)

server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts

@@ -272,19 +272,19 @@ describe('StreamingClientServiceIAM', () => {
         expect(streamingClientServiceIAM['inflightRequests'].size).to.eq(0)
     })
 
-    it('uses expireTime from credentials when available', async () => {
+    it('uses expiration from credentials when available', async () => {
         // Get the credential provider function from the client config
         const credentialProvider = streamingClientServiceIAM.client.config.credentials
         expect(credentialProvider).to.not.be.undefined
 
         // Reset call count on the stub
         features.credentialsProvider.getCredentials.resetHistory()
 
-        // Set up credentials with expireTime
+        // Set up credentials with expiration
         const futureDate = new Date(Date.now() + 3600000) // 1 hour in the future
         const CREDENTIALS_WITH_EXPIRY = {
             ...MOCKED_IAM_CREDENTIALS,
-            expireTime: futureDate.toISOString(),
+            expiration: futureDate,
         }
         features.credentialsProvider.getCredentials.withArgs('iam').returns(CREDENTIALS_WITH_EXPIRY)
 
@@ -293,34 +293,29 @@ describe('StreamingClientServiceIAM', () => {
         await clock.tickAsync(TIME_TO_ADVANCE_MS)
         const credentials = await credentialsPromise
 
-        // Verify expiration is set to the expireTime from credentials
+        // Verify expiration is set to the expiration from credentials
         expect(credentials.expiration).to.be.instanceOf(Date)
         expect(credentials.expiration.getTime()).to.equal(futureDate.getTime())
     })
 
-    it('falls back to current date when expireTime is not available', async () => {
+    it('forces refresh when expiration is not available in credentials', async () => {
         // Get the credential provider function from the client config
         const credentialProvider = streamingClientServiceIAM.client.config.credentials
         expect(credentialProvider).to.not.be.undefined
 
         // Reset call count on the stub
         features.credentialsProvider.getCredentials.resetHistory()
 
-        // Set up credentials without expireTime
+        // Set up credentials without expiration
         features.credentialsProvider.getCredentials.withArgs('iam').returns(MOCKED_IAM_CREDENTIALS)
 
-        // Set a fixed time for testing
-        const fixedNow = new Date()
-        clock.tick(0) // Ensure clock is at the fixed time
-
         // Call the credential provider
         const credentialsPromise = (credentialProvider as any)()
         await clock.tickAsync(TIME_TO_ADVANCE_MS)
         const credentials = await credentialsPromise
 
-        // Verify expiration is set to current date when expireTime is missing
+        // Verify expiration is set to current date to force refresh when not provided in credentials
         expect(credentials.expiration).to.be.instanceOf(Date)
-        // The expiration should be very close to the current time
-        expect(credentials.expiration.getTime()).to.be.closeTo(fixedNow.getTime(), 100)
+        expect(credentials.expiration.getTime()).to.be.closeTo(Date.now(), 1000)
     })
 })

server/aws-lsp-codewhisperer/src/shared/streamingClientService.ts

@@ -12,7 +12,12 @@ import {
     SendMessageCommandInput as SendMessageCommandInputQDeveloperStreaming,
     SendMessageCommandOutput as SendMessageCommandOutputQDeveloperStreaming,
 } from '@amzn/amazon-q-developer-streaming-client'
-import { CredentialsProvider, SDKInitializator, Logging } from '@aws/language-server-runtimes/server-interface'
+import {
+    CredentialsProvider,
+    SDKInitializator,
+    Logging,
+    IamCredentials,
+} from '@aws/language-server-runtimes/server-interface'
 import { getBearerTokenFromProvider, isUsageLimitError } from './utils'
 import { ConfiguredRetryStrategy } from '@aws-sdk/util-retry'
 import { CredentialProviderChain, Credentials } from 'aws-sdk'
@@ -188,13 +193,17 @@ export class StreamingClientServiceIAM extends StreamingClientServiceBase {
 
         // Create a credential provider that fetches fresh credentials on each request
         const iamCredentialProvider: AwsCredentialIdentityProvider = async (): Promise<AwsCredentialIdentity> => {
-            const creds = (await credentialsProvider.getCredentials('iam')) as Credentials
+            const creds = (await credentialsProvider.getCredentials('iam')) as IamCredentials
             logging.log(`Fetching new IAM credentials`)
+            if (!creds) {
+                logging.log('Failed to fetch IAM credentials: No IAM credentials found')
+                throw new Error('No IAM credentials found')
+            }
             return {
                 accessKeyId: creds.accessKeyId,
                 secretAccessKey: creds.secretAccessKey,
                 sessionToken: creds.sessionToken,
-                expiration: creds.expireTime ? new Date(creds.expireTime) : new Date(), // Force refresh on each request if creds do not have expiration time
+                expiration: creds.expiration ? new Date(creds.expiration) : new Date(), // Force refresh if expiration field is not available
             }
         }