From cca4e380da484cb4e73fdba6a149b8ad81b063b4 Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Wed, 17 Sep 2025 16:23:40 -0700 Subject: [PATCH 1/8] fix(amazonq): fix iam credential update logic to use custom comparator and added buffer time in cred validation --- packages/amazonq/src/lsp/auth.ts | 18 ++++++++-- packages/core/src/auth/auth.ts | 2 ++ packages/core/src/auth/credentials/store.ts | 11 ++++-- .../src/shared/utilities/functionUtils.ts | 26 ++++++++++++++ .../shared/utilities/functionUtils.test.ts | 34 ++++++++++++++++++- 5 files changed, 86 insertions(+), 5 deletions(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index 161ba4d9762..b2e1df7786f 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/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'}` diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index 9a050c9f185..501577ee09e 100644 --- a/packages/core/src/auth/auth.ts +++ b/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)) diff --git a/packages/core/src/auth/credentials/store.ts b/packages/core/src/auth/credentials/store.ts index d99595a3877..7392daa5203 100644 --- a/packages/core/src/auth/credentials/store.ts +++ b/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 } diff --git a/packages/core/src/shared/utilities/functionUtils.ts b/packages/core/src/shared/utilities/functionUtils.ts index 214721b1cdb..fa0e61847bb 100644 --- a/packages/core/src/shared/utilities/functionUtils.ts +++ b/packages/core/src/shared/utilities/functionUtils.ts @@ -63,6 +63,32 @@ export function onceChanged(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( + 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. * diff --git a/packages/core/src/test/shared/utilities/functionUtils.test.ts b/packages/core/src/test/shared/utilities/functionUtils.test.ts index b675fe74feb..06acb2bef00 100644 --- a/packages/core/src/test/shared/utilities/functionUtils.test.ts +++ b/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) => { From b64cf2b5c00053e9b94c2353eba3ca28fbd429e6 Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Wed, 17 Sep 2025 16:32:57 -0700 Subject: [PATCH 2/8] fix(amazonq): remove unused import --- packages/core/src/auth/credentials/store.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/core/src/auth/credentials/store.ts b/packages/core/src/auth/credentials/store.ts index 7392daa5203..f9d44408679 100644 --- a/packages/core/src/auth/credentials/store.ts +++ b/packages/core/src/auth/credentials/store.ts @@ -8,7 +8,6 @@ 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 From 1cd2f619462310e3da6dcc805f303e5127eb222f Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Wed, 17 Sep 2025 16:44:55 -0700 Subject: [PATCH 3/8] fix(amazonq): add change log --- .../Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json diff --git a/packages/amazonq/.changes/next-release/Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json b/packages/amazonq/.changes/next-release/Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json new file mode 100644 index 00000000000..3bb0428ebac --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-da0e805d-faab-4274-b37a-943c7263e42b.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Amazon Q automatically refreshes expired IAM Credentials in Sagemaker instances" +} From c634328e2eb4c3d0bb386acd596adef6a445b42a Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Wed, 17 Sep 2025 17:01:37 -0700 Subject: [PATCH 4/8] fix linting --- packages/amazonq/src/lsp/auth.ts | 8 ++++++-- .../core/src/test/shared/utilities/functionUtils.test.ts | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index b2e1df7786f..779625b56af 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -109,8 +109,12 @@ export class AmazonQLspAuth { } private areCredentialsEqual(creds1: any, creds2: any): boolean { - if (!creds1 && !creds2) return true - if (!creds1 || !creds2) return false + if (!creds1 && !creds2) { + return true + } + if (!creds1 || !creds2) { + return false + } return ( creds1.accessKeyId === creds2.accessKeyId && diff --git a/packages/core/src/test/shared/utilities/functionUtils.test.ts b/packages/core/src/test/shared/utilities/functionUtils.test.ts index 06acb2bef00..3ba11518414 100644 --- a/packages/core/src/test/shared/utilities/functionUtils.test.ts +++ b/packages/core/src/test/shared/utilities/functionUtils.test.ts @@ -58,8 +58,12 @@ describe('functionUtils', function () { it('onceChangedWithComparator()', function () { let counter = 0 const credentialsEqual = ([prev]: [any], [current]: [any]) => { - if (!prev && !current) return true - if (!prev || !current) return false + 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) From 7b4b8e4341e75a8d7720d4bd15267072e163f264 Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Wed, 17 Sep 2025 20:16:23 -0700 Subject: [PATCH 5/8] refactor redundant checks --- packages/amazonq/src/lsp/auth.ts | 5 +---- packages/core/src/auth/credentials/store.ts | 3 +-- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index 779625b56af..af23844d37c 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -109,11 +109,8 @@ export class AmazonQLspAuth { } private areCredentialsEqual(creds1: any, creds2: any): boolean { - if (!creds1 && !creds2) { - return true - } if (!creds1 || !creds2) { - return false + return creds1 === creds2 } return ( diff --git a/packages/core/src/auth/credentials/store.ts b/packages/core/src/auth/credentials/store.ts index f9d44408679..9fd73c9130c 100644 --- a/packages/core/src/auth/credentials/store.ts +++ b/packages/core/src/auth/credentials/store.ts @@ -38,8 +38,7 @@ export class CredentialsStore { const expiration = this.credentialsCache[key].credentials.expiration const now = new globals.clock.Date() const bufferedNow = new globals.clock.Date(now.getTime() + expirationBufferMs) - const isValid = expiration !== undefined ? expiration >= bufferedNow : true - return isValid + return expiration !== undefined ? expiration >= bufferedNow : true } getLogger().debug(`credentials: no credentials found for ${key}`) return false From 02d484830d4da3ead5e96a61e0c467ec0b7e3a01 Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Thu, 18 Sep 2025 08:52:54 -0700 Subject: [PATCH 6/8] fix(amazonq): move method to shared util --- packages/amazonq/src/lsp/auth.ts | 16 ++-------------- packages/core/src/auth/auth.ts | 1 - packages/core/src/auth/connection.ts | 12 ++++++++++++ 3 files changed, 14 insertions(+), 15 deletions(-) diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts index af23844d37c..5dfc269f5ef 100644 --- a/packages/amazonq/src/lsp/auth.ts +++ b/packages/amazonq/src/lsp/auth.ts @@ -19,7 +19,7 @@ import { AuthUtil } from 'aws-core-vscode/codewhisperer' import { Writable } from 'stream' 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' +import { isSsoConnection, isIamConnection, areCredentialsEqual } from 'aws-core-vscode/auth' export const encryptionKey = crypto.randomBytes(32) @@ -108,21 +108,9 @@ export class AmazonQLspAuth { this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`) } - private areCredentialsEqual(creds1: any, creds2: any): boolean { - if (!creds1 || !creds2) { - return creds1 === creds2 - } - - 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) + ([prevCreds], [currentCreds]) => areCredentialsEqual(prevCreds, currentCreds) ) private async _updateIamCredentials(credentials: any) { getLogger().info( diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index 501577ee09e..053df50321e 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -875,7 +875,6 @@ 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)) diff --git a/packages/core/src/auth/connection.ts b/packages/core/src/auth/connection.ts index fca46da53e4..fea929fc8af 100644 --- a/packages/core/src/auth/connection.ts +++ b/packages/core/src/auth/connection.ts @@ -71,6 +71,18 @@ export const isBuilderIdConnection = (conn?: Connection): conn is SsoConnection export const isValidCodeCatalystConnection = (conn?: Connection): conn is SsoConnection => isSsoConnection(conn) && hasScopes(conn, scopesCodeCatalyst) +export const areCredentialsEqual = (creds1: any, creds2: any): boolean => { + if (!creds1 || !creds2) { + return creds1 === creds2 + } + + return ( + creds1.accessKeyId === creds2.accessKeyId && + creds1.secretAccessKey === creds2.secretAccessKey && + creds1.sessionToken === creds2.sessionToken + ) +} + export function hasScopes(target: SsoConnection | SsoProfile | string[], scopes: string[]): boolean { return scopes?.every((s) => (Array.isArray(target) ? target : target.scopes)?.includes(s)) } From e49ec58aa36409db0dab94e38a5c21028aace8a7 Mon Sep 17 00:00:00 2001 From: Jayakrishna P Date: Thu, 18 Sep 2025 09:13:55 -0700 Subject: [PATCH 7/8] add missing export for util --- packages/core/src/auth/index.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/core/src/auth/index.ts b/packages/core/src/auth/index.ts index c180d603c67..a5a3ca0edd9 100644 --- a/packages/core/src/auth/index.ts +++ b/packages/core/src/auth/index.ts @@ -19,6 +19,7 @@ export { getTelemetryMetadataForConn, isIamConnection, isSsoConnection, + areCredentialsEqual, } from './connection' export { Auth } from './auth' export { CredentialsStore } from './credentials/store' From 9614843c33b8b1667dea6b24b50c2e778d49a583 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Thu, 18 Sep 2025 12:38:04 -0700 Subject: [PATCH 8/8] ci: empty commit