fix(amazonq): fix iam credential update logic to use custom comparator and added buffer time in cred validation

Author: parameja1

packages/amazonq/src/lsp/auth.ts

@@ -17,7 +17,7 @@ import * as crypto from 'crypto'
 import { LanguageClient } from 'vscode-languageclient'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
-import { onceChanged } from 'aws-core-vscode/utils'
+import { onceChanged, onceChangedWithComparator } from 'aws-core-vscode/utils'
 import { getLogger, oneMinute, isSageMaker } from 'aws-core-vscode/shared'
 import { isSsoConnection, isIamConnection } from 'aws-core-vscode/auth'
 
@@ -108,7 +108,21 @@ export class AmazonQLspAuth {
         this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
     }
 
-    public updateIamCredentials = onceChanged(this._updateIamCredentials.bind(this))
+    private areCredentialsEqual(creds1: any, creds2: any): boolean {
+        if (!creds1 && !creds2) return true
+        if (!creds1 || !creds2) return false
+
+        return (
+            creds1.accessKeyId === creds2.accessKeyId &&
+            creds1.secretAccessKey === creds2.secretAccessKey &&
+            creds1.sessionToken === creds2.sessionToken
+        )
+    }
+
+    public updateIamCredentials = onceChangedWithComparator(
+        this._updateIamCredentials.bind(this),
+        ([prevCreds], [currentCreds]) => this.areCredentialsEqual(prevCreds, currentCreds)
+    )
     private async _updateIamCredentials(credentials: any) {
         getLogger().info(
             `[SageMaker Debug] Updating IAM credentials - credentials received: ${credentials ? 'YES' : 'NO'}`

packages/core/src/auth/auth.ts

@@ -862,6 +862,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     private async createCachedCredentials(provider: CredentialsProvider) {
         const providerId = provider.getCredentialsId()
+        getLogger().debug(`credentials: create cache credentials for ${provider.getProviderType()}`)
         globals.loginManager.store.invalidateCredentials(providerId)
         const { credentials, endpointUrl } = await globals.loginManager.store.upsertCredentials(providerId, provider)
         await globals.loginManager.validateCredentials(credentials, endpointUrl, provider.getDefaultRegion())
@@ -874,6 +875,7 @@ export class Auth implements AuthService, ConnectionManager {
         if (creds !== undefined && creds.credentialsHashCode === provider.getHashCode()) {
             return creds.credentials
         }
+        return undefined
     }
 
     private readonly getToken = keyedDebounce(this._getToken.bind(this))

packages/core/src/auth/credentials/store.ts

@@ -8,6 +8,7 @@ import globals from '../../shared/extensionGlobals'
 import { getLogger } from '../../shared/logger/logger'
 import { asString, CredentialsProvider, CredentialsId } from '../providers/credentials'
 import { CredentialsProviderManager } from '../providers/credentialsProviderManager'
+// import { get } from 'lodash'
 
 export interface CachedCredentials {
     credentials: AWS.Credentials
@@ -31,11 +32,17 @@ export class CredentialsStore {
      * If the expiration property does not exist, it is assumed to never expire.
      */
     public isValid(key: string): boolean {
+        // Apply 60-second buffer similar to SSO token expiry logic
+        const expirationBufferMs = 60000
+
         if (this.credentialsCache[key]) {
             const expiration = this.credentialsCache[key].credentials.expiration
-            return expiration !== undefined ? expiration >= new globals.clock.Date() : true
+            const now = new globals.clock.Date()
+            const bufferedNow = new globals.clock.Date(now.getTime() + expirationBufferMs)
+            const isValid = expiration !== undefined ? expiration >= bufferedNow : true
+            return isValid
         }
-
+        getLogger().debug(`credentials: no credentials found for ${key}`)
         return false
     }
 

packages/core/src/shared/utilities/functionUtils.ts

@@ -63,6 +63,32 @@ export function onceChanged<T, U extends any[]>(fn: (...args: U) => T): (...args
             : ((val = fn(...args)), (ran = true), (prevArgs = args.map(String).join(':')), val)
 }
 
+/**
+ * Creates a function that runs only if the args changed versus the previous invocation,
+ * using a custom comparator function for argument comparison.
+ *
+ * @param fn The function to wrap
+ * @param comparator Function that returns true if arguments are equal
+ */
+export function onceChangedWithComparator<T, U extends any[]>(
+    fn: (...args: U) => T,
+    comparator: (prev: U, current: U) => boolean
+): (...args: U) => T {
+    let val: T
+    let ran = false
+    let prevArgs: U
+
+    return (...args) => {
+        if (ran && comparator(prevArgs, args)) {
+            return val
+        }
+        val = fn(...args)
+        ran = true
+        prevArgs = args
+        return val
+    }
+}
+
 /**
  * Creates a new function that stores the result of a call.
  *

packages/core/src/test/shared/utilities/functionUtils.test.ts

@@ -4,7 +4,13 @@
  */
 
 import assert from 'assert'
-import { once, onceChanged, debounce, oncePerUniqueArg } from '../../../shared/utilities/functionUtils'
+import {
+    once,
+    onceChanged,
+    debounce,
+    oncePerUniqueArg,
+    onceChangedWithComparator,
+} from '../../../shared/utilities/functionUtils'
 import { installFakeClock } from '../../testUtil'
 
 describe('functionUtils', function () {
@@ -49,6 +55,32 @@ describe('functionUtils', function () {
         assert.strictEqual(counter, 3)
     })
 
+    it('onceChangedWithComparator()', function () {
+        let counter = 0
+        const credentialsEqual = ([prev]: [any], [current]: [any]) => {
+            if (!prev && !current) return true
+            if (!prev || !current) return false
+            return prev.accessKeyId === current.accessKeyId && prev.secretAccessKey === current.secretAccessKey
+        }
+        const fn = onceChangedWithComparator((creds: any) => void counter++, credentialsEqual)
+
+        const creds1 = { accessKeyId: 'key1', secretAccessKey: 'secret1' }
+        const creds2 = { accessKeyId: 'key1', secretAccessKey: 'secret1' }
+        const creds3 = { accessKeyId: 'key2', secretAccessKey: 'secret2' }
+
+        fn(creds1)
+        assert.strictEqual(counter, 1)
+
+        fn(creds2) // Same values, should not execute
+        assert.strictEqual(counter, 1)
+
+        fn(creds3) // Different values, should execute
+        assert.strictEqual(counter, 2)
+
+        fn(creds3) // Same as previous, should not execute
+        assert.strictEqual(counter, 2)
+    })
+
     it('oncePerUniqueArg()', function () {
         let counter = 0
         const fn = oncePerUniqueArg((s: string) => {