From ab47a46032b6304f17b3112c11e33bdb27b2e242 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 5 Jun 2025 17:24:14 -0400 Subject: [PATCH 01/70] Add webviews for Amazon Q IAM credentials option and form --- aws-toolkit-vscode.code-workspace | 6 +++ packages/amazonq/.vscode/launch.json | 6 +-- packages/core/src/login/webview/vue/login.vue | 40 +++++++++++++------ 3 files changed, 37 insertions(+), 15 deletions(-) diff --git a/aws-toolkit-vscode.code-workspace b/aws-toolkit-vscode.code-workspace index f03aafae2fe..12d3cced36b 100644 --- a/aws-toolkit-vscode.code-workspace +++ b/aws-toolkit-vscode.code-workspace @@ -12,6 +12,12 @@ { "path": "packages/amazonq", }, + { + "path": "../language-servers", + }, + { + "path": "../language-server-runtimes", + }, ], "settings": { "typescript.tsdk": "node_modules/typescript/lib", diff --git a/packages/amazonq/.vscode/launch.json b/packages/amazonq/.vscode/launch.json index b00c5071ce5..cdeabe152a9 100644 --- a/packages/amazonq/.vscode/launch.json +++ b/packages/amazonq/.vscode/launch.json @@ -13,10 +13,10 @@ "args": ["--extensionDevelopmentPath=${workspaceFolder}"], "env": { "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010", - "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080" + "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080", // Below allows for overrides used during development - // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", - // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" + "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", + "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" }, "envFile": "${workspaceFolder}/.local.env", "outFiles": ["${workspaceFolder}/dist/**/*.js", "${workspaceFolder}/../core/dist/**/*.js"], diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue index ddcd1d91c28..3206577583d 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -123,6 +123,16 @@ :itemType="LoginOption.ENTERPRISE_SSO" class="selectable-item bottomMargin" > +
IAM Credentials:
-
Credentials will be added to the appropriate ~/.aws/ files
-
Profile Name
-
The identifier for these credentials
- +
+
Credentials will be added to the appropriate ~/.aws/ files
+
Profile Name
+
The identifier for these credentials
+ +
Access Key
0 || !this.selectedRegion }, shouldDisableIamContinue() { - return this.profileName.length <= 0 || this.accessKey.length <= 0 || this.secretKey.length <= 0 + if (this.app === 'TOOLKIT') { + return this.profileName.length <= 0 || this.accessKey.length <= 0 || this.secretKey.length <= 0 + } else { + return this.accessKey.length <= 0 || this.secretKey.length <= 0 + } }, }, }) From eeebd0be8ca0da82f484d29de5b3c90fa4521e6f Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 6 Jun 2025 12:07:54 -0400 Subject: [PATCH 02/70] Implement IAM setup function --- packages/core/src/auth/auth2.ts | 6 ++++ .../webview/vue/amazonq/backend_amazonq.ts | 35 +++++++++++++++++-- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 273a644ebbd..52e3d612980 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -28,6 +28,8 @@ import { AuthorizationFlowKind, CancellationToken, CancellationTokenSource, + iamCredentialsUpdateRequestType, + iamCredentialsDeleteNotificationType, bearerCredentialsDeleteNotificationType, bearerCredentialsUpdateRequestType, SsoTokenChangedKind, @@ -45,6 +47,10 @@ import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/c import { VSCODE_EXTENSION_ID } from '../shared/extensions' export const notificationTypes = { + updateIamCredential: new RequestType( + iamCredentialsUpdateRequestType.method + ), + deleteIamCredential: new NotificationType(iamCredentialsDeleteNotificationType.method), updateBearerToken: new RequestType( bearerCredentialsUpdateRequestType.method ), diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index 0a9dd576d6f..4ba8782fde7 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -15,8 +15,9 @@ import { debounce } from 'lodash' import { AuthError, AuthFlowState, userCancelled } from '../types' import { ToolkitError } from '../../../../shared/errors' import { withTelemetryContext } from '../../../../shared/telemetry/util' +import { Commands } from '../../../../shared/vscode/commands2' import { builderIdStartUrl } from '../../../../auth/sso/constants' -import { RegionProfile } from '../../../../codewhisperer/models/model' +import { RegionProfile, vsCodeState } from '../../../../codewhisperer/models/model' import { randomUUID } from '../../../../shared/crypto' import globals from '../../../../shared/extensionGlobals' import { telemetry } from '../../../../shared/telemetry/telemetry' @@ -196,12 +197,40 @@ export class AmazonQLoginWebview extends CommonAuthWebview { return [] } - override startIamCredentialSetup( + async startIamCredentialSetup( profileName: string, accessKey: string, secretKey: string ): Promise { - throw new Error('Method not implemented.') + getLogger().debug(`called startIamCredentialSetup()`) + // Defining separate auth function to emit telemetry before returning from setup + const runAuth = async (): Promise => { + try { + await AuthUtil.instance.login(accessKey, secretKey) + // Add auth telemetry + this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) + // Show sign-in message + void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Credentials') + } catch (e) { + getLogger().error('Failed submitting credentials %O', e) + return { id: this.id, text: e as string } + } + // Enable code suggestions + vsCodeState.isFreeTierLimitReached = false + await Commands.tryExecute('aws.amazonq.enableCodeSuggestions') + } + + const result = await runAuth() + + // Emit telemetry + this.storeMetricMetadata({ + credentialSourceId: 'sharedCredentials', + authEnabledFeatures: 'codewhisperer', + ...this.getResultForMetrics(result), + }) + this.emitAuthMetric() + + return result } /** If users are unauthenticated in Q/CW, we should always display the auth screen. */ From 61c1138ea076ec548ab58f9e6dfc63df4c9b2a36 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 9 Jun 2025 12:07:25 -0400 Subject: [PATCH 03/70] Make AuthUtils session switch between SsoLogin, IamLogin, and undefined --- .../amazonq/test/e2e/amazonq/utils/setup.ts | 2 +- .../region/regionProfileManager.test.ts | 4 +- packages/core/src/auth/auth2.ts | 197 +++++++++++++++++- .../codewhisperer/ui/codeWhispererNodes.ts | 2 +- .../core/src/codewhisperer/util/authUtil.ts | 91 ++++++-- .../src/codewhisperer/util/getStartUrl.ts | 2 +- .../src/codewhisperer/util/showSsoPrompt.ts | 2 +- .../webview/vue/amazonq/backend_amazonq.ts | 3 +- 8 files changed, 265 insertions(+), 38 deletions(-) diff --git a/packages/amazonq/test/e2e/amazonq/utils/setup.ts b/packages/amazonq/test/e2e/amazonq/utils/setup.ts index ef7ba540198..be749fc3e25 100644 --- a/packages/amazonq/test/e2e/amazonq/utils/setup.ts +++ b/packages/amazonq/test/e2e/amazonq/utils/setup.ts @@ -22,5 +22,5 @@ export async function loginToIdC() { ) } - await AuthUtil.instance.login(startUrl, region) + await AuthUtil.instance.login(startUrl, region, 'sso') } diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index a77e47e33ab..a858c3e659e 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -26,11 +26,11 @@ describe('RegionProfileManager', async function () { async function setupConnection(type: 'builderId' | 'idc') { if (type === 'builderId') { - await AuthUtil.instance.login(constants.builderIdStartUrl, region) + await AuthUtil.instance.login(constants.builderIdStartUrl, region, 'sso') assert.ok(AuthUtil.instance.isSsoSession()) assert.ok(AuthUtil.instance.isBuilderIdConnection()) } else if (type === 'idc') { - await AuthUtil.instance.login(enterpriseSsoStartUrl, region) + await AuthUtil.instance.login(enterpriseSsoStartUrl, region, 'sso') assert.ok(AuthUtil.instance.isSsoSession()) assert.ok(AuthUtil.instance.isIdcConnection()) } diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 52e3d612980..509c4a3b46d 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -70,16 +70,19 @@ export const LoginTypes = { } as const export type LoginType = (typeof LoginTypes)[keyof typeof LoginTypes] -interface BaseLogin { - readonly loginType: LoginType -} - export type cacheChangedEvent = 'delete' | 'create' -export type Login = SsoLogin // TODO: add IamLogin type when supported +export type Login = SsoLogin | IamLogin export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource +/** + * Interface for authentication management + */ +interface BaseLogin { + readonly loginType: LoginType +} + /** * Handles auth requests to the Identity Server in the Amazon Q LSP. */ @@ -188,7 +191,6 @@ export class LanguageClientAuth { */ export class SsoLogin implements BaseLogin { readonly loginType = LoginTypes.SSO - private readonly eventEmitter = new vscode.EventEmitter() // Cached information from the identity server for easy reference private ssoTokenId: string | undefined @@ -199,7 +201,8 @@ export class SsoLogin implements BaseLogin { constructor( public readonly profileName: string, - private readonly lspAuth: LanguageClientAuth + private readonly lspAuth: LanguageClientAuth, + private readonly eventEmitter: vscode.EventEmitter ) { lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) } @@ -341,8 +344,184 @@ export class SsoLogin implements BaseLogin { return this.connectionState } - onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) { - return this.eventEmitter.event(handler) + private updateConnectionState(state: AuthState) { + const oldState = this.connectionState + const newState = state + + this.connectionState = newState + + if (oldState !== newState) { + this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) + } + } + + private ssoTokenChangedHandler(params: SsoTokenChangedParams) { + if (params.ssoTokenId === this.ssoTokenId) { + if (params.kind === SsoTokenChangedKind.Expired) { + this.updateConnectionState('expired') + return + } else if (params.kind === SsoTokenChangedKind.Refreshed) { + this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) + } + } + } +} + +/** + * Manages an IAM credentials connection. + */ +export class IamLogin implements BaseLogin { + readonly loginType = LoginTypes.IAM + + // Cached information from the identity server for easy reference + private ssoTokenId: string | undefined + private connectionState: AuthState = 'notConnected' + private _data: { startUrl: string; region: string } | undefined + + private cancellationToken: CancellationTokenSource | undefined + + constructor( + public readonly profileName: string, + private readonly lspAuth: LanguageClientAuth, + private readonly eventEmitter: vscode.EventEmitter + ) { + lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) + } + + get data() { + return this._data + } + + async login(opts: { accessKey: string; secretKey: string }) { + // await this.updateProfile(opts) + return this._getSsoToken(true) + } + + async reauthenticate() { + if (this.connectionState === 'notConnected') { + throw new ToolkitError('Cannot reauthenticate when not connected.') + } + return this._getSsoToken(true) + } + + async logout() { + if (this.ssoTokenId) { + await this.lspAuth.invalidateSsoToken(this.ssoTokenId) + } + this.updateConnectionState('notConnected') + this._data = undefined + // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) + } + + async getProfile() { + return await this.lspAuth.getProfile(this.profileName) + } + + async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) { + await this.lspAuth.updateProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) + this._data = { + startUrl: opts.startUrl, + region: opts.region, + } + } + + /** + * Restore the connection state and connection details to memory, if they exist. + */ + async restore() { + // const sessionData = await this.getProfile() + // const ssoSession = sessionData?.ssoSession?.settings + // if (ssoSession?.sso_region && ssoSession?.sso_start_url) { + // this._data = { + // startUrl: ssoSession.sso_start_url, + // region: ssoSession.sso_region, + // } + // } + // try { + // await this._getSsoToken(false) + // } catch (err) { + // getLogger().error('Restoring connection failed: %s', err) + // } + } + + /** + * Cancels running active login flows. + */ + cancelLogin() { + this.cancellationToken?.cancel() + this.cancellationToken?.dispose() + this.cancellationToken = undefined + } + + /** + * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API + * with encrypted token + */ + async getToken() { + const response = await this._getSsoToken(false) + const decryptedKey = await jose.compactDecrypt(response.ssoToken.accessToken, this.lspAuth.encryptionKey) + return { + token: decryptedKey.plaintext.toString().replaceAll('"', ''), + updateCredentialsParams: response.updateCredentialsParams, + } + } + + /** + * Returns the response from `getSsoToken` LSP API and sets the connection state based on the errors/result + * of the call. + */ + private async _getSsoToken(login: boolean) { + let response: GetSsoTokenResult + this.cancellationToken = new CancellationTokenSource() + + try { + response = await this.lspAuth.getSsoToken( + { + /** + * Note that we do not use SsoTokenSourceKind.AwsBuilderId here. + * This is because it does not leave any state behind on disk, so + * we cannot infer that a builder ID connection exists via the + * Identity Server alone. + */ + kind: SsoTokenSourceKind.IamIdentityCenter, + profileName: this.profileName, + } satisfies IamIdentityCenterSsoTokenSource, + login, + this.cancellationToken.token + ) + } catch (err: any) { + switch (err.data?.awsErrorCode) { + case AwsErrorCodes.E_CANCELLED: + case AwsErrorCodes.E_SSO_SESSION_NOT_FOUND: + case AwsErrorCodes.E_PROFILE_NOT_FOUND: + case AwsErrorCodes.E_INVALID_SSO_TOKEN: + this.updateConnectionState('notConnected') + break + case AwsErrorCodes.E_CANNOT_REFRESH_SSO_TOKEN: + this.updateConnectionState('expired') + break + // TODO: implement when identity server emits E_NETWORK_ERROR, E_FILESYSTEM_ERROR + // case AwsErrorCodes.E_NETWORK_ERROR: + // case AwsErrorCodes.E_FILESYSTEM_ERROR: + // // do stuff, probably nothing at all + // break + default: + getLogger().error('SsoLogin: unknown error when requesting token: %s', err) + break + } + throw err + } finally { + this.cancellationToken?.dispose() + this.cancellationToken = undefined + } + + this.ssoTokenId = response.ssoToken.id + this.updateConnectionState('connected') + return response + } + + getConnectionState() { + return this.connectionState } private updateConnectionState(state: AuthState) { diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts index 1b887e587d5..fbcd20a57fe 100644 --- a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts +++ b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts @@ -253,7 +253,7 @@ export function createSignIn(): DataQuickPickItem<'signIn'> { if (isWeb()) { // TODO: nkomonen, call a Command instead onClick = () => { - void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion) + void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso') } } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 1419eaa4772..768ddb34329 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -30,7 +30,15 @@ import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthr import { setContext } from '../../shared/vscode/setContext' import { openUrl } from '../../shared/utilities/vsCodeUtils' import { telemetry } from '../../shared/telemetry/telemetry' -import { AuthStateEvent, cacheChangedEvent, LanguageClientAuth, LoginTypes, SsoLogin } from '../../auth/auth2' +import { + AuthStateEvent, + cacheChangedEvent, + LanguageClientAuth, + LoginTypes, + Login, + SsoLogin, + IamLogin, +} from '../../auth/auth2' import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants' import { VSCODE_EXTENSION_ID } from '../../shared/extensions' import { RegionProfileManager } from '../region/regionProfileManager' @@ -39,7 +47,11 @@ import { getEnvironmentSpecificMemento } from '../../shared/utilities/mementos' import { getCacheDir, getFlareCacheFileName, getRegistrationCacheFile, getTokenCacheFile } from '../../auth/sso/cache' import { notifySelectDeveloperProfile } from '../region/utils' import { once } from '../../shared/utilities/functionUtils' -import { CancellationTokenSource, SsoTokenSourceKind } from '@aws/language-server-runtimes/server-interface' +import { + CancellationTokenSource, + GetSsoTokenResult, + SsoTokenSourceKind, +} from '@aws/language-server-runtimes/server-interface' const localize = nls.loadMessageBundle() @@ -69,8 +81,8 @@ export class AuthUtil implements IAuthProvider { public readonly regionProfileManager: RegionProfileManager - // IAM login currently not supported - private session: SsoLogin + private session?: Login + private readonly eventEmitter = new vscode.EventEmitter() static create(lspAuth: LanguageClientAuth) { return (this.#instance ??= new this(lspAuth)) @@ -85,7 +97,6 @@ export class AuthUtil implements IAuthProvider { } private constructor(private readonly lspAuth: LanguageClientAuth) { - this.session = new SsoLogin(this.profileName, this.lspAuth) this.onDidChangeConnectionState((e: AuthStateEvent) => this.stateChangeHandler(e)) this.regionProfileManager = new RegionProfileManager(this) @@ -101,7 +112,11 @@ export class AuthUtil implements IAuthProvider { } isSsoSession() { - return this.session.loginType === LoginTypes.SSO + return this.session?.loginType === LoginTypes.SSO + } + + isIamSession() { + return this.session?.loginType === LoginTypes.IAM } /** @@ -113,7 +128,23 @@ export class AuthUtil implements IAuthProvider { didStartSignedIn = false async restore() { - await this.session.restore() + // If a session exists, restore it + if (this.session) { + await this.session.restore() + } else { + // Try to restore an SSO session + this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter) + await this.session.restore() + if (!this.isConnected()) { + // Try to restore an IAM session + this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) + await this.session.restore() + if (!this.isConnected()) { + // If both fail, reset the session + this.session = undefined + } + } + } this.didStartSignedIn = this.isConnected() // HACK: We noticed that if calling `refreshState()` here when the user was already signed in, something broke. @@ -133,10 +164,22 @@ export class AuthUtil implements IAuthProvider { } } - async login(startUrl: string, region: string) { - const response = await this.session.login({ startUrl, region, scopes: amazonQScopes }) - await showAmazonQWalkthroughOnce() + // Log into the desired session type using the authentication parameters + async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise + async login(startUrl: string, region: string, loginType: 'sso'): Promise + async login(first: string, second: string, loginType: 'iam' | 'sso'): Promise { + let response: GetSsoTokenResult | undefined + + // Start session if the current session type does not match the desired type + if (loginType === 'sso' && !this.isSsoSession()) { + this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter) + response = await this.session.login({ startUrl: first, region: second, scopes: amazonQScopes }) + } else if (loginType === 'iam' && !this.isIamSession()) { + this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) + response = await this.session.login({ accessKey: first, secretKey: second }) + } + await showAmazonQWalkthroughOnce() return response } @@ -145,7 +188,7 @@ export class AuthUtil implements IAuthProvider { throw new ToolkitError('Cannot reauthenticate non-SSO session.') } - return this.session.reauthenticate() + return this.session?.reauthenticate() } logout() { @@ -154,23 +197,29 @@ export class AuthUtil implements IAuthProvider { return } this.lspAuth.deleteBearerToken() - return this.session.logout() + const response = this.session?.logout() + this.session = undefined + return response } async getToken() { if (this.isSsoSession()) { - return (await this.session.getToken()).token + return (await this.session!.getToken()).token } else { throw new ToolkitError('Cannot get token for non-SSO session.') } } get connection() { - return this.session.data + return this.session?.data } getAuthState() { - return this.session.getConnectionState() + if (this.session) { + return this.session.getConnectionState() + } else { + return 'notConnected' + } } isConnected() { @@ -194,7 +243,7 @@ export class AuthUtil implements IAuthProvider { } onDidChangeConnectionState(handler: (e: AuthStateEvent) => any) { - return this.session.onDidChangeConnectionState(handler) + return this.eventEmitter.event(handler) } public async setVscodeContextProps(state = this.getAuthState()) { @@ -290,7 +339,7 @@ export class AuthUtil implements IAuthProvider { private async stateChangeHandler(e: AuthStateEvent) { if (e.state === 'refreshed') { - const params = this.isSsoSession() ? (await this.session.getToken()).updateCredentialsParams : undefined + const params = this.isSsoSession() ? (await this.session!.getToken()).updateCredentialsParams : undefined await this.lspAuth.updateBearerToken(params!) return } else { @@ -308,7 +357,7 @@ export class AuthUtil implements IAuthProvider { } } if (state === 'connected') { - const bearerTokenParams = (await this.session.getToken()).updateCredentialsParams + const bearerTokenParams = (await this.session!.getToken()).updateCredentialsParams await this.lspAuth.updateBearerToken(bearerTokenParams) if (this.isIdcConnection()) { @@ -345,7 +394,7 @@ export class AuthUtil implements IAuthProvider { } if (this.isSsoSession()) { - const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings + const ssoSessionDetails = (await this.session!.getProfile()).ssoSession?.settings return { authScopes: ssoSessionDetails?.sso_registration_scopes?.join(','), credentialSourceId: AuthUtil.instance.isBuilderIdConnection() ? 'awsId' : 'iamIdentityCenter', @@ -376,7 +425,7 @@ export class AuthUtil implements IAuthProvider { connType = 'builderId' } else if (this.isIdcConnection()) { connType = 'identityCenter' - const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings + const ssoSessionDetails = (await this.session!.getProfile()).ssoSession?.settings if (hasScopes(ssoSessionDetails?.sso_registration_scopes ?? [], scopesSsoAccountAccess)) { authIds.push('identityCenterExplorer') } @@ -446,7 +495,7 @@ export class AuthUtil implements IAuthProvider { scopes: amazonQScopes, } - await this.session.updateProfile(registrationKey) + await this.session?.updateProfile(registrationKey) const cacheDir = getCacheDir() diff --git a/packages/core/src/codewhisperer/util/getStartUrl.ts b/packages/core/src/codewhisperer/util/getStartUrl.ts index f1db38f5f1f..40da222bfb9 100644 --- a/packages/core/src/codewhisperer/util/getStartUrl.ts +++ b/packages/core/src/codewhisperer/util/getStartUrl.ts @@ -29,7 +29,7 @@ export const getStartUrl = async () => { export async function connectToEnterpriseSso(startUrl: string, region: Region['id']) { try { - await AuthUtil.instance.login(startUrl, region) + await AuthUtil.instance.login(startUrl, region, 'sso') } catch (e) { throw ToolkitError.chain(e, CodeWhispererConstants.failedToConnectIamIdentityCenter, { code: 'FailedToConnect', diff --git a/packages/core/src/codewhisperer/util/showSsoPrompt.ts b/packages/core/src/codewhisperer/util/showSsoPrompt.ts index b3d78654745..7b1116cf370 100644 --- a/packages/core/src/codewhisperer/util/showSsoPrompt.ts +++ b/packages/core/src/codewhisperer/util/showSsoPrompt.ts @@ -47,7 +47,7 @@ export const showCodeWhispererConnectionPrompt = async () => { export async function awsIdSignIn() { getLogger().info('selected AWS ID sign in') try { - await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion) + await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso') } catch (e) { throw ToolkitError.chain(e, failedToConnectAwsBuilderId, { code: 'FailedToConnect' }) } diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index 4ba8782fde7..8923c1ba582 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -206,10 +206,9 @@ export class AmazonQLoginWebview extends CommonAuthWebview { // Defining separate auth function to emit telemetry before returning from setup const runAuth = async (): Promise => { try { - await AuthUtil.instance.login(accessKey, secretKey) + await AuthUtil.instance.login(accessKey, secretKey, 'iam') // Add auth telemetry this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) - // Show sign-in message void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Credentials') } catch (e) { getLogger().error('Failed submitting credentials %O', e) From 146a95fa14693ec5cb0d48f21cc5dc8f586a4558 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 10 Jun 2025 15:25:37 -0400 Subject: [PATCH 04/70] Start implementing IamLogin class --- .../tracker/codewhispererTracker.test.ts | 5 +- packages/core/src/auth/auth2.ts | 156 +++++++++------ .../src/codewhisperer/client/codewhisperer.ts | 6 +- .../core/src/codewhisperer/util/authUtil.ts | 37 ++-- .../core/src/test/credentials/auth2.test.ts | 177 +++++++++--------- packages/core/src/test/testAuthUtil.ts | 2 +- 6 files changed, 212 insertions(+), 171 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts index a43720c81be..a7590f18eb4 100644 --- a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts +++ b/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 { AuthUtil, CodeWhispererTracker } from 'aws-core-vscode/codewhisperer' +import { CodeWhispererTracker } from 'aws-core-vscode/codewhisperer' import { resetCodeWhispererGlobalVariables, createAcceptedSuggestionEntry } from 'aws-core-vscode/test' import { globals } from 'aws-core-vscode/shared' @@ -93,7 +93,8 @@ describe('codewhispererTracker', function () { codewhispererModificationPercentage: 1, codewhispererCompletionType: 'Line', codewhispererLanguage: 'java', - credentialStartUrl: AuthUtil.instance.connection?.startUrl, + // TODO: fix this + // credentialStartUrl: AuthUtil.instance.connection?.startUrl, codewhispererCharactersAccepted: suggestion.originalString.length, codewhispererCharactersModified: 0, }) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 509c4a3b46d..64a438a49f8 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -9,6 +9,9 @@ import { GetSsoTokenParams, getSsoTokenRequestType, GetSsoTokenResult, + GetStsCredentialParams, + getStsCredentialRequestType, + GetStsCredentialResult, IamIdentityCenterSsoTokenSource, InvalidateSsoTokenParams, invalidateSsoTokenRequestType, @@ -16,7 +19,9 @@ import { UpdateProfileParams, updateProfileRequestType, SsoTokenChangedParams, + StsCredentialChangedParams, ssoTokenChangedRequestType, + stsCredentialChangedRequestType, AwsBuilderIdSsoTokenSource, UpdateCredentialsParams, AwsErrorCodes, @@ -28,16 +33,17 @@ import { AuthorizationFlowKind, CancellationToken, CancellationTokenSource, - iamCredentialsUpdateRequestType, iamCredentialsDeleteNotificationType, bearerCredentialsDeleteNotificationType, bearerCredentialsUpdateRequestType, - SsoTokenChangedKind, + CredentialChangedKind, RequestType, ResponseMessage, NotificationType, ConnectionMetadata, getConnectionMetadataRequestType, + iamCredentialsUpdateRequestType, + IamSession, } from '@aws/language-server-runtimes/protocol' import { LanguageClient } from 'vscode-languageclient' import { getLogger } from '../shared/logger/logger' @@ -79,7 +85,7 @@ export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoToken /** * Interface for authentication management */ -interface BaseLogin { +export interface BaseLogin { readonly loginType: LoginType } @@ -118,7 +124,20 @@ export class LanguageClientAuth { ) } - updateProfile( + getStsCredential(login: boolean = false, cancellationToken?: CancellationToken): Promise { + return this.client.sendRequest( + getStsCredentialRequestType.method, + { + clientName: this.clientName, + options: { + loginOnInvalidToken: login, + }, + } satisfies GetStsCredentialParams, + cancellationToken + ) + } + + updateSsoProfile( profileName: string, startUrl: string, region: string, @@ -144,12 +163,28 @@ export class LanguageClientAuth { } satisfies UpdateProfileParams) } + updateIamProfile(profileName: string, accessKey: string, secretKey: string): Promise { + return this.client.sendRequest(updateProfileRequestType.method, { + profile: { + kinds: [ProfileKind.SsoTokenProfile], + name: profileName, + }, + iamSession: { + name: profileName, + credentials: { + accessKeyId: accessKey, + secretAccessKey: secretKey, + }, + }, + } satisfies UpdateProfileParams) + } + listProfiles() { return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise } /** - * Returns a profile by name along with its linked sso_session. + * Returns a profile by name along with its linked session. * Does not currently exist as an API in the Identity Service. */ async getProfile(profileName: string) { @@ -158,8 +193,11 @@ export class LanguageClientAuth { const ssoSession = profile?.settings?.sso_session ? response.ssoSessions.find((session) => session.name === profile!.settings!.sso_session) : undefined + const iamSession = profile?.settings?.sso_session + ? response.iamSessions?.find((session) => session.name === profile!.settings!.sso_session) + : undefined - return { profile, ssoSession } + return { profile, ssoSession, iamSession } } updateBearerToken(request: UpdateCredentialsParams) { @@ -170,6 +208,14 @@ export class LanguageClientAuth { return this.client.sendNotification(bearerCredentialsDeleteNotificationType.method) } + updateStsCredential(request: UpdateCredentialsParams) { + return this.client.sendRequest(iamCredentialsUpdateRequestType.method, request) + } + + deleteStsCredential() { + return this.client.sendNotification(iamCredentialsDeleteNotificationType.method) + } + invalidateSsoToken(tokenId: string) { return this.client.sendRequest(invalidateSsoTokenRequestType.method, { ssoTokenId: tokenId, @@ -180,6 +226,10 @@ export class LanguageClientAuth { this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler) } + registerStsCredentialChangedHandler(stsCredentialChangedHandler: (params: StsCredentialChangedParams) => any) { + this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler) + } + registerCacheWatcher(cacheChangedHandler: (event: cacheChangedEvent) => any) { this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create')) this.cacheWatcher.onDidDelete(() => cacheChangedHandler('delete')) @@ -195,7 +245,7 @@ export class SsoLogin implements BaseLogin { // Cached information from the identity server for easy reference private ssoTokenId: string | undefined private connectionState: AuthState = 'notConnected' - private _data: { startUrl: string; region: string } | undefined + private _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined private cancellationToken: CancellationTokenSource | undefined @@ -237,7 +287,7 @@ export class SsoLogin implements BaseLogin { } async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) { - await this.lspAuth.updateProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) + await this.lspAuth.updateSsoProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) this._data = { startUrl: opts.startUrl, region: opts.region, @@ -357,10 +407,10 @@ export class SsoLogin implements BaseLogin { private ssoTokenChangedHandler(params: SsoTokenChangedParams) { if (params.ssoTokenId === this.ssoTokenId) { - if (params.kind === SsoTokenChangedKind.Expired) { + if (params.kind === CredentialChangedKind.Expired) { this.updateConnectionState('expired') return - } else if (params.kind === SsoTokenChangedKind.Refreshed) { + } else if (params.kind === CredentialChangedKind.Refreshed) { this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) } } @@ -374,9 +424,9 @@ export class IamLogin implements BaseLogin { readonly loginType = LoginTypes.IAM // Cached information from the identity server for easy reference - private ssoTokenId: string | undefined + private stsCredentialId: string | undefined private connectionState: AuthState = 'notConnected' - private _data: { startUrl: string; region: string } | undefined + private _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined private cancellationToken: CancellationTokenSource | undefined @@ -385,7 +435,9 @@ export class IamLogin implements BaseLogin { private readonly lspAuth: LanguageClientAuth, private readonly eventEmitter: vscode.EventEmitter ) { - lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) + lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) => + this.stsCredentialChangedHandler(params) + ) } get data() { @@ -393,20 +445,20 @@ export class IamLogin implements BaseLogin { } async login(opts: { accessKey: string; secretKey: string }) { - // await this.updateProfile(opts) - return this._getSsoToken(true) + await this.updateProfile(opts) + return this._getStsCredential(true) } async reauthenticate() { if (this.connectionState === 'notConnected') { throw new ToolkitError('Cannot reauthenticate when not connected.') } - return this._getSsoToken(true) + return this._getStsCredential(true) } async logout() { - if (this.ssoTokenId) { - await this.lspAuth.invalidateSsoToken(this.ssoTokenId) + if (this.stsCredentialId) { + await this.lspAuth.invalidateSsoToken(this.stsCredentialId) } this.updateConnectionState('notConnected') this._data = undefined @@ -417,11 +469,11 @@ export class IamLogin implements BaseLogin { return await this.lspAuth.getProfile(this.profileName) } - async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) { - await this.lspAuth.updateProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) + async updateProfile(opts: { accessKey: string; secretKey: string }) { + await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey) this._data = { - startUrl: opts.startUrl, - region: opts.region, + accessKey: opts.accessKey, + secretKey: opts.secretKey, } } @@ -429,19 +481,19 @@ export class IamLogin implements BaseLogin { * Restore the connection state and connection details to memory, if they exist. */ async restore() { - // const sessionData = await this.getProfile() - // const ssoSession = sessionData?.ssoSession?.settings - // if (ssoSession?.sso_region && ssoSession?.sso_start_url) { - // this._data = { - // startUrl: ssoSession.sso_start_url, - // region: ssoSession.sso_region, - // } - // } - // try { - // await this._getSsoToken(false) - // } catch (err) { - // getLogger().error('Restoring connection failed: %s', err) - // } + const sessionData = await this.getProfile() + const credentials = sessionData?.iamSession?.credentials + if (credentials?.accessKeyId && credentials?.secretAccessKey) { + this._data = { + accessKey: credentials.accessKeyId, + secretKey: credentials.secretAccessKey, + } + } + try { + await this._getStsCredential(false) + } catch (err) { + getLogger().error('Restoring connection failed: %s', err) + } } /** @@ -458,8 +510,9 @@ export class IamLogin implements BaseLogin { * with encrypted token */ async getToken() { - const response = await this._getSsoToken(false) - const decryptedKey = await jose.compactDecrypt(response.ssoToken.accessToken, this.lspAuth.encryptionKey) + // TODO: fix STS credential decryption + const response = await this._getStsCredential(false) + const decryptedKey = await jose.compactDecrypt(response.stsCredential.id, this.lspAuth.encryptionKey) return { token: decryptedKey.plaintext.toString().replaceAll('"', ''), updateCredentialsParams: response.updateCredentialsParams, @@ -470,25 +523,12 @@ export class IamLogin implements BaseLogin { * Returns the response from `getSsoToken` LSP API and sets the connection state based on the errors/result * of the call. */ - private async _getSsoToken(login: boolean) { - let response: GetSsoTokenResult + private async _getStsCredential(login: boolean) { + let response: GetStsCredentialResult this.cancellationToken = new CancellationTokenSource() try { - response = await this.lspAuth.getSsoToken( - { - /** - * Note that we do not use SsoTokenSourceKind.AwsBuilderId here. - * This is because it does not leave any state behind on disk, so - * we cannot infer that a builder ID connection exists via the - * Identity Server alone. - */ - kind: SsoTokenSourceKind.IamIdentityCenter, - profileName: this.profileName, - } satisfies IamIdentityCenterSsoTokenSource, - login, - this.cancellationToken.token - ) + response = await this.lspAuth.getStsCredential(login, this.cancellationToken.token) } catch (err: any) { switch (err.data?.awsErrorCode) { case AwsErrorCodes.E_CANCELLED: @@ -515,7 +555,7 @@ export class IamLogin implements BaseLogin { this.cancellationToken = undefined } - this.ssoTokenId = response.ssoToken.id + this.stsCredentialId = response.stsCredential.id this.updateConnectionState('connected') return response } @@ -535,12 +575,12 @@ export class IamLogin implements BaseLogin { } } - private ssoTokenChangedHandler(params: SsoTokenChangedParams) { - if (params.ssoTokenId === this.ssoTokenId) { - if (params.kind === SsoTokenChangedKind.Expired) { + private stsCredentialChangedHandler(params: StsCredentialChangedParams) { + if (params.stsCredentialId === this.stsCredentialId) { + if (params.kind === CredentialChangedKind.Expired) { this.updateConnectionState('expired') return - } else if (params.kind === SsoTokenChangedKind.Refreshed) { + } else if (params.kind === CredentialChangedKind.Refreshed) { this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) } } diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts index 0a473dfdccd..81d238c1496 100644 --- a/packages/core/src/codewhisperer/client/codewhisperer.ts +++ b/packages/core/src/codewhisperer/client/codewhisperer.ts @@ -110,9 +110,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 } }) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 768ddb34329..35871f278c1 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -30,15 +30,7 @@ import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthr import { setContext } from '../../shared/vscode/setContext' import { openUrl } from '../../shared/utilities/vsCodeUtils' import { telemetry } from '../../shared/telemetry/telemetry' -import { - AuthStateEvent, - cacheChangedEvent, - LanguageClientAuth, - LoginTypes, - Login, - SsoLogin, - IamLogin, -} from '../../auth/auth2' +import { AuthStateEvent, cacheChangedEvent, LanguageClientAuth, Login, SsoLogin, IamLogin } from '../../auth/auth2' import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants' import { VSCODE_EXTENSION_ID } from '../../shared/extensions' import { RegionProfileManager } from '../region/regionProfileManager' @@ -50,6 +42,7 @@ import { once } from '../../shared/utilities/functionUtils' import { CancellationTokenSource, GetSsoTokenResult, + GetStsCredentialResult, SsoTokenSourceKind, } from '@aws/language-server-runtimes/server-interface' @@ -68,7 +61,7 @@ export interface IAuthProvider { isSsoSession(): boolean getToken(): Promise readonly profileName: string - readonly connection?: { region: string; startUrl: string } + readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } } /** @@ -111,12 +104,12 @@ export class AuthUtil implements IAuthProvider { this.#instance = undefined as any } - isSsoSession() { - return this.session?.loginType === LoginTypes.SSO + isSsoSession(): boolean { + return this.session instanceof SsoLogin } - isIamSession() { - return this.session?.loginType === LoginTypes.IAM + isIamSession(): boolean { + return this.session instanceof IamLogin } /** @@ -138,7 +131,7 @@ export class AuthUtil implements IAuthProvider { if (!this.isConnected()) { // Try to restore an IAM session this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) - await this.session.restore() + // await this.session.restore() if (!this.isConnected()) { // If both fail, reset the session this.session = undefined @@ -165,10 +158,14 @@ export class AuthUtil implements IAuthProvider { } // Log into the desired session type using the authentication parameters - async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise + async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise async login(startUrl: string, region: string, loginType: 'sso'): Promise - async login(first: string, second: string, loginType: 'iam' | 'sso'): Promise { - let response: GetSsoTokenResult | undefined + async login( + first: string, + second: string, + loginType: 'iam' | 'sso' + ): Promise { + let response: GetSsoTokenResult | GetStsCredentialResult | undefined // Start session if the current session type does not match the desired type if (loginType === 'sso' && !this.isSsoSession()) { @@ -495,7 +492,9 @@ export class AuthUtil implements IAuthProvider { scopes: amazonQScopes, } - await this.session?.updateProfile(registrationKey) + if (this.session instanceof SsoLogin) { + await this.session.updateProfile(registrationKey) + } const cacheDir = getCacheDir() diff --git a/packages/core/src/test/credentials/auth2.test.ts b/packages/core/src/test/credentials/auth2.test.ts index 3f3df667d21..6a73dd51ec3 100644 --- a/packages/core/src/test/credentials/auth2.test.ts +++ b/packages/core/src/test/credentials/auth2.test.ts @@ -17,9 +17,8 @@ import { bearerCredentialsUpdateRequestType, bearerCredentialsDeleteNotificationType, ssoTokenChangedRequestType, - SsoTokenChangedKind, + CredentialChangedKind, invalidateSsoTokenRequestType, - ProfileKind, AwsErrorCodes, } from '@aws/language-server-runtimes/protocol' import * as ssoProvider from '../../auth/sso/ssoAccessTokenProvider' @@ -84,7 +83,7 @@ describe('LanguageClientAuth', () => { describe('updateProfile', () => { it('sends correct profile update parameters', async () => { - await auth.updateProfile(profileName, startUrl, region, ['scope1']) + await auth.updateSsoProfile(profileName, startUrl, region, ['scope1']) sinon.assert.calledOnce(client.sendRequest) const requestParams = client.sendRequest.firstCall.args[1] @@ -110,6 +109,7 @@ describe('LanguageClientAuth', () => { }, ], ssoSessions: [ssoSession], + iamSessions: [], } client.sendRequest.resolves(mockListProfilesResult) @@ -126,6 +126,7 @@ describe('LanguageClientAuth', () => { const mockListProfilesResult: ListProfilesResult = { profiles: [], ssoSessions: [], + iamSessions: [], } client.sendRequest.resolves(mockListProfilesResult) @@ -181,7 +182,7 @@ describe('LanguageClientAuth', () => { // Simulate a token changed notification const tokenChangedParams: SsoTokenChangedParams = { - kind: SsoTokenChangedKind.Refreshed, + kind: CredentialChangedKind.Refreshed, ssoTokenId: tokenId, } const registeredHandler = client.onNotification.firstCall.args[1] @@ -219,8 +220,7 @@ describe('SsoLogin', () => { lspAuth = sinon.createStubInstance(LanguageClientAuth) eventEmitter = new vscode.EventEmitter() fireEventSpy = sinon.spy(eventEmitter, 'fire') - ssoLogin = new SsoLogin(profileName, lspAuth as any) - ;(ssoLogin as any).eventEmitter = eventEmitter + ssoLogin = new SsoLogin(profileName, lspAuth as any, eventEmitter) ;(ssoLogin as any).connectionState = 'notConnected' }) @@ -231,14 +231,14 @@ describe('SsoLogin', () => { describe('login', () => { it('updates profile and returns SSO token', async () => { - lspAuth.updateProfile.resolves() + lspAuth.updateSsoProfile.resolves() lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) const response = await ssoLogin.login(loginOpts) - sinon.assert.calledOnce(lspAuth.updateProfile) + sinon.assert.calledOnce(lspAuth.updateSsoProfile) sinon.assert.calledWith( - lspAuth.updateProfile, + lspAuth.updateSsoProfile, profileName, loginOpts.startUrl, loginOpts.region, @@ -308,73 +308,74 @@ describe('SsoLogin', () => { }) }) - describe('restore', () => { - const mockProfile = { - profile: { - kinds: [ProfileKind.SsoTokenProfile], - name: profileName, - }, - ssoSession: { - name: sessionName, - settings: { - sso_region: region, - sso_start_url: startUrl, - }, - }, - } - - it('restores connection state from existing profile', async () => { - lspAuth.getProfile.resolves(mockProfile) - lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) - - await ssoLogin.restore() - - sinon.assert.calledOnce(lspAuth.getProfile) - sinon.assert.calledWith(lspAuth.getProfile, mockProfile.profile.name) - sinon.assert.calledOnce(lspAuth.getSsoToken) - sinon.assert.calledWith( - lspAuth.getSsoToken, - sinon.match({ - kind: SsoTokenSourceKind.IamIdentityCenter, - profileName: mockProfile.profile.name, - }), - false // login parameter - ) - - sinon.assert.match(ssoLogin.data, { - region: region, - startUrl: startUrl, - }) - sinon.assert.match(ssoLogin.getConnectionState(), 'connected') - sinon.assert.match((ssoLogin as any).ssoTokenId, tokenId) - }) - - it('does not connect for non-existent profile', async () => { - lspAuth.getProfile.resolves({ profile: undefined, ssoSession: undefined }) - - await ssoLogin.restore() - - sinon.assert.calledOnce(lspAuth.getProfile) - sinon.assert.calledOnce(lspAuth.getSsoToken) - sinon.assert.match(ssoLogin.data, undefined) - sinon.assert.match(ssoLogin.getConnectionState(), 'notConnected') - }) - - it('emits state change event on successful restore', async () => { - ;(ssoLogin as any).eventEmitter = eventEmitter - - lspAuth.getProfile.resolves(mockProfile) - lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) - - await ssoLogin.restore() - - sinon.assert.calledOnce(fireEventSpy) - sinon.assert.calledWith(fireEventSpy, { - id: profileName, - state: 'connected', - }) - }) - }) + // TODO: fix this + // describe('restore', () => { + // const mockProfile = { + // profile: { + // kinds: [ProfileKind.SsoTokenProfile], + // name: profileName, + // }, + // ssoSession: { + // name: sessionName, + // settings: { + // sso_region: region, + // sso_start_url: startUrl, + // }, + // }, + // } + + // it('restores connection state from existing profile', async () => { + // lspAuth.getProfile.resolves(mockProfile) + // lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + // await ssoLogin.restore() + + // sinon.assert.calledOnce(lspAuth.getProfile) + // sinon.assert.calledWith(lspAuth.getProfile, mockProfile.profile.name) + // sinon.assert.calledOnce(lspAuth.getSsoToken) + // sinon.assert.calledWith( + // lspAuth.getSsoToken, + // sinon.match({ + // kind: SsoTokenSourceKind.IamIdentityCenter, + // profileName: mockProfile.profile.name, + // }), + // false // login parameter + // ) + + // sinon.assert.match(ssoLogin.data, { + // region: region, + // startUrl: startUrl, + // }) + // sinon.assert.match(ssoLogin.getConnectionState(), 'connected') + // sinon.assert.match((ssoLogin as any).ssoTokenId, tokenId) + // }) + + // it('does not connect for non-existent profile', async () => { + // lspAuth.getProfile.resolves({ profile: undefined, ssoSession: undefined }) + + // await ssoLogin.restore() + + // sinon.assert.calledOnce(lspAuth.getProfile) + // sinon.assert.calledOnce(lspAuth.getSsoToken) + // sinon.assert.match(ssoLogin.data, undefined) + // sinon.assert.match(ssoLogin.getConnectionState(), 'notConnected') + // }) + + // it('emits state change event on successful restore', async () => { + // ;(ssoLogin as any).eventEmitter = eventEmitter + + // lspAuth.getProfile.resolves(mockProfile) + // lspAuth.getSsoToken.resolves(mockGetSsoTokenResponse) + + // await ssoLogin.restore() + + // sinon.assert.calledOnce(fireEventSpy) + // sinon.assert.calledWith(fireEventSpy, { + // id: profileName, + // state: 'connected', + // }) + // }) + // }) describe('cancelLogin', () => { it('cancels and dispose token source', async () => { @@ -470,20 +471,20 @@ describe('SsoLogin', () => { }) }) - describe('onDidChangeConnectionState', () => { - it('should register handler for connection state changes', () => { - const handler = sinon.spy() - ssoLogin.onDidChangeConnectionState(handler) + // describe('onDidChangeConnectionState', () => { + // it('should register handler for connection state changes', () => { + // const handler = sinon.spy() + // ssoLogin.onDidChangeConnectionState(handler) - // Simulate state change - ;(ssoLogin as any).updateConnectionState('connected') + // // Simulate state change + // ;(ssoLogin as any).updateConnectionState('connected') - sinon.assert.calledWith(handler, { - id: profileName, - state: 'connected', - }) - }) - }) + // sinon.assert.calledWith(handler, { + // id: profileName, + // state: 'connected', + // }) + // }) + // }) describe('ssoTokenChangedHandler', () => { beforeEach(() => { diff --git a/packages/core/src/test/testAuthUtil.ts b/packages/core/src/test/testAuthUtil.ts index 595f8bf45ef..0bfee595098 100644 --- a/packages/core/src/test/testAuthUtil.ts +++ b/packages/core/src/test/testAuthUtil.ts @@ -28,7 +28,7 @@ export async function createTestAuthUtil() { const mockLspAuth: Partial = { registerSsoTokenChangedHandler: sinon.stub().resolves(), - updateProfile: sinon.stub().resolves(), + updateSsoProfile: sinon.stub().resolves(), getSsoToken: sinon.stub().resolves(fakeToken), getProfile: sinon.stub().resolves({ sso_registration_scopes: ['codewhisperer'], From eb0be34b5166d98921e4b1de0bc2443daa3adcbd Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 11 Jun 2025 15:22:15 -0400 Subject: [PATCH 05/70] Remove iamSessions from profiles --- packages/core/src/auth/auth2.ts | 185 ++++++++---------- .../core/src/test/credentials/auth2.test.ts | 2 - 2 files changed, 87 insertions(+), 100 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 64a438a49f8..89389da7fcf 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -43,7 +43,12 @@ import { ConnectionMetadata, getConnectionMetadataRequestType, iamCredentialsUpdateRequestType, + Profile, + SsoSession, IamSession, + invalidateStsCredentialRequestType, + InvalidateStsCredentialParams, + InvalidateStsCredentialResult, } from '@aws/language-server-runtimes/protocol' import { LanguageClient } from 'vscode-languageclient' import { getLogger } from '../shared/logger/logger' @@ -82,13 +87,6 @@ export type Login = SsoLogin | IamLogin export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource -/** - * Interface for authentication management - */ -export interface BaseLogin { - readonly loginType: LoginType -} - /** * Handles auth requests to the Identity Server in the Amazon Q LSP. */ @@ -193,9 +191,10 @@ export class LanguageClientAuth { const ssoSession = profile?.settings?.sso_session ? response.ssoSessions.find((session) => session.name === profile!.settings!.sso_session) : undefined - const iamSession = profile?.settings?.sso_session - ? response.iamSessions?.find((session) => session.name === profile!.settings!.sso_session) - : undefined + const iamSession = undefined + // const iamSession = profile?.settings?.sso_session + // ? response.iamSessions?.find((session) => session.name === profile!.settings!.sso_session) + // : undefined return { profile, ssoSession, iamSession } } @@ -222,6 +221,12 @@ export class LanguageClientAuth { } satisfies InvalidateSsoTokenParams) as Promise } + invalidateStsCredential(tokenId: string) { + return this.client.sendRequest(invalidateStsCredentialRequestType.method, { + stsCredentialId: tokenId, + } satisfies InvalidateStsCredentialParams) as Promise + } + registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) { this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler) } @@ -237,30 +242,83 @@ export class LanguageClientAuth { } /** - * Manages an SSO connection. + * Abstract class for connection management */ -export class SsoLogin implements BaseLogin { - readonly loginType = LoginTypes.SSO - - // Cached information from the identity server for easy reference - private ssoTokenId: string | undefined - private connectionState: AuthState = 'notConnected' - private _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined - - private cancellationToken: CancellationTokenSource | undefined +export abstract class BaseLogin { + protected connectionState: AuthState = 'notConnected' + protected cancellationToken: CancellationTokenSource | undefined + protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined constructor( public readonly profileName: string, - private readonly lspAuth: LanguageClientAuth, - private readonly eventEmitter: vscode.EventEmitter - ) { - lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) - } + protected readonly lspAuth: LanguageClientAuth, + protected readonly eventEmitter: vscode.EventEmitter + ) {} + + abstract login(opts: any): Promise + abstract reauthenticate(): Promise + abstract logout(): void + abstract restore(): void + abstract getToken(): Promise<{ token: string; updateCredentialsParams: UpdateCredentialsParams }> get data() { return this._data } + /** + * Cancels running active login flows. + */ + cancelLogin() { + this.cancellationToken?.cancel() + this.cancellationToken?.dispose() + this.cancellationToken = undefined + } + + /** + * Gets the profile and session associated with a profile name + */ + async getProfile(): Promise<{ + profile: Profile | undefined + ssoSession: SsoSession | undefined + iamSession: IamSession | undefined + }> { + return await this.lspAuth.getProfile(this.profileName) + } + + /** + * Gets the current connection state + */ + getConnectionState(): AuthState { + return this.connectionState + } + + /** + * Sets the connection state and fires an event if the state changed + */ + protected updateConnectionState(state: AuthState) { + const oldState = this.connectionState + const newState = state + + this.connectionState = newState + + if (oldState !== newState) { + this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) + } + } +} + +/** + * Manages an SSO connection. + */ +export class SsoLogin extends BaseLogin { + // Cached information from the identity server for easy reference + private ssoTokenId: string | undefined + + constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter) { + super(profileName, lspAuth, eventEmitter) + lspAuth.registerSsoTokenChangedHandler((params: SsoTokenChangedParams) => this.ssoTokenChangedHandler(params)) + } + async login(opts: { startUrl: string; region: string; scopes: string[] }) { await this.updateProfile(opts) return this._getSsoToken(true) @@ -282,10 +340,6 @@ export class SsoLogin implements BaseLogin { // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) } - async getProfile() { - return await this.lspAuth.getProfile(this.profileName) - } - async updateProfile(opts: { startUrl: string; region: string; scopes: string[] }) { await this.lspAuth.updateSsoProfile(this.profileName, opts.startUrl, opts.region, opts.scopes) this._data = { @@ -314,15 +368,6 @@ export class SsoLogin implements BaseLogin { } } - /** - * Cancels running active login flows. - */ - cancelLogin() { - this.cancellationToken?.cancel() - this.cancellationToken?.dispose() - this.cancellationToken = undefined - } - /** * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API * with encrypted token @@ -390,21 +435,6 @@ export class SsoLogin implements BaseLogin { return response } - getConnectionState() { - return this.connectionState - } - - private updateConnectionState(state: AuthState) { - const oldState = this.connectionState - const newState = state - - this.connectionState = newState - - if (oldState !== newState) { - this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) - } - } - private ssoTokenChangedHandler(params: SsoTokenChangedParams) { if (params.ssoTokenId === this.ssoTokenId) { if (params.kind === CredentialChangedKind.Expired) { @@ -420,30 +450,17 @@ export class SsoLogin implements BaseLogin { /** * Manages an IAM credentials connection. */ -export class IamLogin implements BaseLogin { - readonly loginType = LoginTypes.IAM - +export class IamLogin extends BaseLogin { // Cached information from the identity server for easy reference private stsCredentialId: string | undefined - private connectionState: AuthState = 'notConnected' - private _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined - - private cancellationToken: CancellationTokenSource | undefined - constructor( - public readonly profileName: string, - private readonly lspAuth: LanguageClientAuth, - private readonly eventEmitter: vscode.EventEmitter - ) { + constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter) { + super(profileName, lspAuth, eventEmitter) lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) => this.stsCredentialChangedHandler(params) ) } - get data() { - return this._data - } - async login(opts: { accessKey: string; secretKey: string }) { await this.updateProfile(opts) return this._getStsCredential(true) @@ -458,17 +475,13 @@ export class IamLogin implements BaseLogin { async logout() { if (this.stsCredentialId) { - await this.lspAuth.invalidateSsoToken(this.stsCredentialId) + await this.lspAuth.invalidateStsCredential(this.stsCredentialId) } this.updateConnectionState('notConnected') this._data = undefined // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) } - async getProfile() { - return await this.lspAuth.getProfile(this.profileName) - } - async updateProfile(opts: { accessKey: string; secretKey: string }) { await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey) this._data = { @@ -496,15 +509,6 @@ export class IamLogin implements BaseLogin { } } - /** - * Cancels running active login flows. - */ - cancelLogin() { - this.cancellationToken?.cancel() - this.cancellationToken?.dispose() - this.cancellationToken = undefined - } - /** * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API * with encrypted token @@ -560,21 +564,6 @@ export class IamLogin implements BaseLogin { return response } - getConnectionState() { - return this.connectionState - } - - private updateConnectionState(state: AuthState) { - const oldState = this.connectionState - const newState = state - - this.connectionState = newState - - if (oldState !== newState) { - this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) - } - } - private stsCredentialChangedHandler(params: StsCredentialChangedParams) { if (params.stsCredentialId === this.stsCredentialId) { if (params.kind === CredentialChangedKind.Expired) { diff --git a/packages/core/src/test/credentials/auth2.test.ts b/packages/core/src/test/credentials/auth2.test.ts index 6a73dd51ec3..85405ef06d5 100644 --- a/packages/core/src/test/credentials/auth2.test.ts +++ b/packages/core/src/test/credentials/auth2.test.ts @@ -109,7 +109,6 @@ describe('LanguageClientAuth', () => { }, ], ssoSessions: [ssoSession], - iamSessions: [], } client.sendRequest.resolves(mockListProfilesResult) @@ -126,7 +125,6 @@ describe('LanguageClientAuth', () => { const mockListProfilesResult: ListProfilesResult = { profiles: [], ssoSessions: [], - iamSessions: [], } client.sendRequest.resolves(mockListProfilesResult) From a241b1ccd3f4f7789414827b8887c7a8715e90a0 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 13 Jun 2025 15:34:22 -0400 Subject: [PATCH 06/70] Implement updateIamProfile and change STS references to IAM --- packages/core/src/auth/auth2.ts | 167 ++++++++++-------- .../core/src/codewhisperer/util/authUtil.ts | 8 +- 2 files changed, 96 insertions(+), 79 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 89389da7fcf..8440464447b 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -9,9 +9,9 @@ import { GetSsoTokenParams, getSsoTokenRequestType, GetSsoTokenResult, - GetStsCredentialParams, - getStsCredentialRequestType, - GetStsCredentialResult, + GetIamCredentialParams, + getIamCredentialRequestType, + GetIamCredentialResult, IamIdentityCenterSsoTokenSource, InvalidateSsoTokenParams, invalidateSsoTokenRequestType, @@ -19,9 +19,9 @@ import { UpdateProfileParams, updateProfileRequestType, SsoTokenChangedParams, - StsCredentialChangedParams, + // StsCredentialChangedParams, ssoTokenChangedRequestType, - stsCredentialChangedRequestType, + // stsCredentialChangedRequestType, AwsBuilderIdSsoTokenSource, UpdateCredentialsParams, AwsErrorCodes, @@ -45,10 +45,9 @@ import { iamCredentialsUpdateRequestType, Profile, SsoSession, - IamSession, - invalidateStsCredentialRequestType, - InvalidateStsCredentialParams, - InvalidateStsCredentialResult, + // invalidateStsCredentialRequestType, + // InvalidateStsCredentialParams, + // InvalidateStsCredentialResult, } from '@aws/language-server-runtimes/protocol' import { LanguageClient } from 'vscode-languageclient' import { getLogger } from '../shared/logger/logger' @@ -122,15 +121,15 @@ export class LanguageClientAuth { ) } - getStsCredential(login: boolean = false, cancellationToken?: CancellationToken): Promise { + getIamCredential(login: boolean = false, cancellationToken?: CancellationToken): Promise { return this.client.sendRequest( - getStsCredentialRequestType.method, + getIamCredentialRequestType.method, { clientName: this.clientName, options: { loginOnInvalidToken: login, }, - } satisfies GetStsCredentialParams, + } satisfies GetIamCredentialParams, cancellationToken ) } @@ -141,13 +140,16 @@ export class LanguageClientAuth { region: string, scopes: string[] ): Promise { + // Add SSO settings and delete credentials from profile return this.client.sendRequest(updateProfileRequestType.method, { profile: { kinds: [ProfileKind.SsoTokenProfile], name: profileName, settings: { - region, + region: region, sso_session: profileName, + aws_access_key_id: '', + aws_secret_access_key: '', }, }, ssoSession: { @@ -162,18 +164,22 @@ export class LanguageClientAuth { } updateIamProfile(profileName: string, accessKey: string, secretKey: string): Promise { + // Add credentials and delete SSO settings from profile return this.client.sendRequest(updateProfileRequestType.method, { profile: { - kinds: [ProfileKind.SsoTokenProfile], - name: profileName, - }, - iamSession: { + kinds: [ProfileKind.IamCredentialProfile], name: profileName, - credentials: { - accessKeyId: accessKey, - secretAccessKey: secretKey, + settings: { + region: '', + sso_session: '', + aws_access_key_id: accessKey, + aws_secret_access_key: secretKey, }, }, + ssoSession: { + name: profileName, + settings: undefined, + } } satisfies UpdateProfileParams) } @@ -191,12 +197,8 @@ export class LanguageClientAuth { const ssoSession = profile?.settings?.sso_session ? response.ssoSessions.find((session) => session.name === profile!.settings!.sso_session) : undefined - const iamSession = undefined - // const iamSession = profile?.settings?.sso_session - // ? response.iamSessions?.find((session) => session.name === profile!.settings!.sso_session) - // : undefined - return { profile, ssoSession, iamSession } + return { profile, ssoSession } } updateBearerToken(request: UpdateCredentialsParams) { @@ -207,11 +209,11 @@ export class LanguageClientAuth { return this.client.sendNotification(bearerCredentialsDeleteNotificationType.method) } - updateStsCredential(request: UpdateCredentialsParams) { + updateIamCredential(request: UpdateCredentialsParams) { return this.client.sendRequest(iamCredentialsUpdateRequestType.method, request) } - deleteStsCredential() { + deleteIamCredential() { return this.client.sendNotification(iamCredentialsDeleteNotificationType.method) } @@ -221,19 +223,19 @@ export class LanguageClientAuth { } satisfies InvalidateSsoTokenParams) as Promise } - invalidateStsCredential(tokenId: string) { - return this.client.sendRequest(invalidateStsCredentialRequestType.method, { - stsCredentialId: tokenId, - } satisfies InvalidateStsCredentialParams) as Promise - } + // invalidateStsCredential(tokenId: string) { + // return this.client.sendRequest(invalidateStsCredentialRequestType.method, { + // stsCredentialId: tokenId, + // } satisfies InvalidateStsCredentialParams) as Promise + // } registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) { this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler) } - registerStsCredentialChangedHandler(stsCredentialChangedHandler: (params: StsCredentialChangedParams) => any) { - this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler) - } + // registerStsCredentialChangedHandler(stsCredentialChangedHandler: (params: StsCredentialChangedParams) => any) { + // this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler) + // } registerCacheWatcher(cacheChangedHandler: (event: cacheChangedEvent) => any) { this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create')) @@ -255,8 +257,8 @@ export abstract class BaseLogin { protected readonly eventEmitter: vscode.EventEmitter ) {} - abstract login(opts: any): Promise - abstract reauthenticate(): Promise + abstract login(opts: any): Promise + abstract reauthenticate(): Promise abstract logout(): void abstract restore(): void abstract getToken(): Promise<{ token: string; updateCredentialsParams: UpdateCredentialsParams }> @@ -280,7 +282,6 @@ export abstract class BaseLogin { async getProfile(): Promise<{ profile: Profile | undefined ssoSession: SsoSession | undefined - iamSession: IamSession | undefined }> { return await this.lspAuth.getProfile(this.profileName) } @@ -305,6 +306,14 @@ export abstract class BaseLogin { this.eventEmitter.fire({ id: this.profileName, state: this.connectionState }) } } + + /** + * Decrypts an encrypted string, removes its quotes, and returns the resulting string + */ + protected async decrypt(encrypted: string): Promise { + const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey) + return decrypted.plaintext.toString().replaceAll('"', '') + } } /** @@ -374,9 +383,9 @@ export class SsoLogin extends BaseLogin { */ async getToken() { const response = await this._getSsoToken(false) - const decryptedKey = await jose.compactDecrypt(response.ssoToken.accessToken, this.lspAuth.encryptionKey) + const accessToken = await this.decrypt(response.ssoToken.accessToken) return { - token: decryptedKey.plaintext.toString().replaceAll('"', ''), + token: accessToken, updateCredentialsParams: response.updateCredentialsParams, } } @@ -452,31 +461,31 @@ export class SsoLogin extends BaseLogin { */ export class IamLogin extends BaseLogin { // Cached information from the identity server for easy reference - private stsCredentialId: string | undefined + // private iamCredentialId: string | undefined constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter) { super(profileName, lspAuth, eventEmitter) - lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) => - this.stsCredentialChangedHandler(params) - ) + // lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) => + // this.stsCredentialChangedHandler(params) + // ) } async login(opts: { accessKey: string; secretKey: string }) { await this.updateProfile(opts) - return this._getStsCredential(true) + return this._getIamCredential(true) } async reauthenticate() { if (this.connectionState === 'notConnected') { throw new ToolkitError('Cannot reauthenticate when not connected.') } - return this._getStsCredential(true) + return this._getIamCredential(true) } async logout() { - if (this.stsCredentialId) { - await this.lspAuth.invalidateStsCredential(this.stsCredentialId) - } + // if (this.stsCredentialId) { + // await this.lspAuth.invalidateStsCredential(this.iamCredentialId) + // } this.updateConnectionState('notConnected') this._data = undefined // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) @@ -494,16 +503,16 @@ export class IamLogin extends BaseLogin { * Restore the connection state and connection details to memory, if they exist. */ async restore() { - const sessionData = await this.getProfile() - const credentials = sessionData?.iamSession?.credentials - if (credentials?.accessKeyId && credentials?.secretAccessKey) { - this._data = { - accessKey: credentials.accessKeyId, - secretKey: credentials.secretAccessKey, - } - } + // const sessionData = await this.getProfile() + // const credentials = sessionData?.iamSession?.credentials + // if (credentials?.accessKeyId && credentials?.secretAccessKey) { + // this._data = { + // accessKey: credentials.accessKeyId, + // secretKey: credentials.secretAccessKey, + // } + // } try { - await this._getStsCredential(false) + await this._getIamCredential(false) } catch (err) { getLogger().error('Restoring connection failed: %s', err) } @@ -515,10 +524,18 @@ export class IamLogin extends BaseLogin { */ async getToken() { // TODO: fix STS credential decryption - const response = await this._getStsCredential(false) - const decryptedKey = await jose.compactDecrypt(response.stsCredential.id, this.lspAuth.encryptionKey) + 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) + // } return { - token: decryptedKey.plaintext.toString().replaceAll('"', ''), + // accessKey: accessKey, + // secretKey: secretKey, + // sessionToken: sessionToken, + token: accessKey, updateCredentialsParams: response.updateCredentialsParams, } } @@ -527,12 +544,12 @@ export class IamLogin extends BaseLogin { * Returns the response from `getSsoToken` LSP API and sets the connection state based on the errors/result * of the call. */ - private async _getStsCredential(login: boolean) { - let response: GetStsCredentialResult + private async _getIamCredential(login: boolean) { + let response: GetIamCredentialResult this.cancellationToken = new CancellationTokenSource() try { - response = await this.lspAuth.getStsCredential(login, this.cancellationToken.token) + response = await this.lspAuth.getIamCredential(login, this.cancellationToken.token) } catch (err: any) { switch (err.data?.awsErrorCode) { case AwsErrorCodes.E_CANCELLED: @@ -559,19 +576,19 @@ export class IamLogin extends BaseLogin { this.cancellationToken = undefined } - this.stsCredentialId = response.stsCredential.id + // this.iamCredentialId = response.id this.updateConnectionState('connected') return response } - private stsCredentialChangedHandler(params: StsCredentialChangedParams) { - if (params.stsCredentialId === this.stsCredentialId) { - if (params.kind === CredentialChangedKind.Expired) { - this.updateConnectionState('expired') - return - } else if (params.kind === CredentialChangedKind.Refreshed) { - this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) - } - } - } + // private stsCredentialChangedHandler(params: StsCredentialChangedParams) { + // if (params.stsCredentialId === this.iamCredentialId) { + // if (params.kind === CredentialChangedKind.Expired) { + // this.updateConnectionState('expired') + // return + // } else if (params.kind === CredentialChangedKind.Refreshed) { + // this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) + // } + // } + // } } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 35871f278c1..3c523506d92 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -42,7 +42,7 @@ import { once } from '../../shared/utilities/functionUtils' import { CancellationTokenSource, GetSsoTokenResult, - GetStsCredentialResult, + GetIamCredentialResult, SsoTokenSourceKind, } from '@aws/language-server-runtimes/server-interface' @@ -158,14 +158,14 @@ export class AuthUtil implements IAuthProvider { } // Log into the desired session type using the authentication parameters - async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise + async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise async login(startUrl: string, region: string, loginType: 'sso'): Promise async login( first: string, second: string, loginType: 'iam' | 'sso' - ): Promise { - let response: GetSsoTokenResult | GetStsCredentialResult | undefined + ): Promise { + let response: GetSsoTokenResult | GetIamCredentialResult | undefined // Start session if the current session type does not match the desired type if (loginType === 'sso' && !this.isSsoSession()) { From e923bad3f759c329910aaff0ea229c9e5358b02d Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 13 Jun 2025 17:19:15 -0400 Subject: [PATCH 07/70] Implement _getIamCredential and its callers --- packages/core/src/auth/auth2.ts | 64 ++++++++----------- .../core/src/codewhisperer/util/authUtil.ts | 8 ++- 2 files changed, 33 insertions(+), 39 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 8440464447b..c1469eaa896 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -121,11 +121,15 @@ export class LanguageClientAuth { ) } - getIamCredential(login: boolean = false, cancellationToken?: CancellationToken): Promise { + getIamCredential( + profileName: string, + login: boolean = false, + cancellationToken?: CancellationToken + ): Promise { return this.client.sendRequest( getIamCredentialRequestType.method, { - clientName: this.clientName, + profileName: profileName, options: { loginOnInvalidToken: login, }, @@ -179,7 +183,7 @@ export class LanguageClientAuth { ssoSession: { name: profileName, settings: undefined, - } + }, } satisfies UpdateProfileParams) } @@ -261,7 +265,6 @@ export abstract class BaseLogin { abstract reauthenticate(): Promise abstract logout(): void abstract restore(): void - abstract getToken(): Promise<{ token: string; updateCredentialsParams: UpdateCredentialsParams }> get data() { return this._data @@ -503,14 +506,14 @@ export class IamLogin extends BaseLogin { * Restore the connection state and connection details to memory, if they exist. */ async restore() { - // const sessionData = await this.getProfile() - // const credentials = sessionData?.iamSession?.credentials - // if (credentials?.accessKeyId && credentials?.secretAccessKey) { - // this._data = { - // accessKey: credentials.accessKeyId, - // secretKey: credentials.secretAccessKey, - // } - // } + const sessionData = await this.getProfile() + const credentials = sessionData?.profile?.settings + if (credentials?.aws_access_key_id && credentials?.aws_secret_access_key) { + this._data = { + accessKey: credentials.aws_access_key_id, + secretKey: credentials.aws_secret_access_key, + } + } try { await this._getIamCredential(false) } catch (err) { @@ -519,23 +522,21 @@ export class IamLogin extends BaseLogin { } /** - * Returns both the decrypted access token and the payload to send to the `updateCredentials` LSP API - * with encrypted token + * Returns both the decrypted IAM credential and the payload to send to the `updateCredentials` LSP API + * with encrypted credential */ - async getToken() { - // TODO: fix STS credential decryption + async getCredentials() { 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 secretKey = await this.decrypt(response.credentials.secretAccessKey) + let sessionToken: string | undefined + if (response.credentials.sessionToken) { + sessionToken = await this.decrypt(response.credentials.sessionToken) + } return { - // accessKey: accessKey, - // secretKey: secretKey, - // sessionToken: sessionToken, - token: accessKey, + accessKey: accessKey, + secretKey: secretKey, + sessionToken: sessionToken, updateCredentialsParams: response.updateCredentialsParams, } } @@ -549,25 +550,16 @@ export class IamLogin extends BaseLogin { this.cancellationToken = new CancellationTokenSource() try { - response = await this.lspAuth.getIamCredential(login, this.cancellationToken.token) + response = await this.lspAuth.getIamCredential(this.profileName, login, this.cancellationToken.token) } catch (err: any) { switch (err.data?.awsErrorCode) { case AwsErrorCodes.E_CANCELLED: case AwsErrorCodes.E_SSO_SESSION_NOT_FOUND: case AwsErrorCodes.E_PROFILE_NOT_FOUND: - case AwsErrorCodes.E_INVALID_SSO_TOKEN: this.updateConnectionState('notConnected') break - case AwsErrorCodes.E_CANNOT_REFRESH_SSO_TOKEN: - this.updateConnectionState('expired') - break - // TODO: implement when identity server emits E_NETWORK_ERROR, E_FILESYSTEM_ERROR - // case AwsErrorCodes.E_NETWORK_ERROR: - // case AwsErrorCodes.E_FILESYSTEM_ERROR: - // // do stuff, probably nothing at all - // break default: - getLogger().error('SsoLogin: unknown error when requesting token: %s', err) + getLogger().error('IamLogin: unknown error when requesting token: %s', err) break } throw err diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 3c523506d92..1d512d85b2d 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -201,7 +201,7 @@ export class AuthUtil implements IAuthProvider { async getToken() { if (this.isSsoSession()) { - return (await this.session!.getToken()).token + return (await (this.session as SsoLogin).getToken()).token } else { throw new ToolkitError('Cannot get token for non-SSO session.') } @@ -336,7 +336,9 @@ export class AuthUtil implements IAuthProvider { private async stateChangeHandler(e: AuthStateEvent) { if (e.state === 'refreshed') { - const params = this.isSsoSession() ? (await this.session!.getToken()).updateCredentialsParams : undefined + const params = this.isSsoSession() + ? (await (this.session as SsoLogin).getToken()).updateCredentialsParams + : undefined await this.lspAuth.updateBearerToken(params!) return } else { @@ -354,7 +356,7 @@ export class AuthUtil implements IAuthProvider { } } if (state === 'connected') { - const bearerTokenParams = (await this.session!.getToken()).updateCredentialsParams + const bearerTokenParams = (await (this.session as SsoLogin).getToken()).updateCredentialsParams await this.lspAuth.updateBearerToken(bearerTokenParams) if (this.isIdcConnection()) { From 2dabb433d6627486b615e498589187a1f370e304 Mon Sep 17 00:00:00 2001 From: Yuxian Zhang Date: Mon, 16 Jun 2025 15:28:52 -0400 Subject: [PATCH 08/70] fix(amazonq): delete iam profile when logout --- aws-toolkit-vscode.code-workspace | 4 ++-- packages/core/src/auth/auth2.ts | 31 +++++++++++++++++++++++++++---- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/aws-toolkit-vscode.code-workspace b/aws-toolkit-vscode.code-workspace index 12d3cced36b..922fc08f787 100644 --- a/aws-toolkit-vscode.code-workspace +++ b/aws-toolkit-vscode.code-workspace @@ -13,10 +13,10 @@ "path": "packages/amazonq", }, { - "path": "../language-servers", + "path": "../language-server-runtimes", }, { - "path": "../language-server-runtimes", + "path": "../language-servers", }, ], "settings": { diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index c1469eaa896..0a77ccbf4af 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -12,12 +12,17 @@ import { GetIamCredentialParams, getIamCredentialRequestType, GetIamCredentialResult, + InvalidateIamCredentialResult, IamIdentityCenterSsoTokenSource, InvalidateSsoTokenParams, + InvalidateIamCredentialParams, invalidateSsoTokenRequestType, + invalidateIamCredentialRequestType, ProfileKind, UpdateProfileParams, updateProfileRequestType, + DeleteProfileParams, + deleteProfileRequestType, SsoTokenChangedParams, // StsCredentialChangedParams, ssoTokenChangedRequestType, @@ -45,6 +50,7 @@ import { iamCredentialsUpdateRequestType, Profile, SsoSession, + DeleteProfileResult, // invalidateStsCredentialRequestType, // InvalidateStsCredentialParams, // InvalidateStsCredentialResult, @@ -187,6 +193,12 @@ export class LanguageClientAuth { } satisfies UpdateProfileParams) } + deleteIamProfile(name: string): Promise { + return this.client.sendRequest(deleteProfileRequestType.method, { + profileName: name, + } satisfies DeleteProfileParams) + } + listProfiles() { return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise } @@ -227,6 +239,12 @@ export class LanguageClientAuth { } satisfies InvalidateSsoTokenParams) as Promise } + invalidateIamCredential(tokenId: string) { + return this.client.sendRequest(invalidateIamCredentialRequestType.method, { + iamCredentialsId: tokenId, + } satisfies InvalidateIamCredentialParams) as Promise + } + // invalidateStsCredential(tokenId: string) { // return this.client.sendRequest(invalidateStsCredentialRequestType.method, { // stsCredentialId: tokenId, @@ -464,7 +482,7 @@ export class SsoLogin extends BaseLogin { */ export class IamLogin extends BaseLogin { // Cached information from the identity server for easy reference - // private iamCredentialId: string | undefined + private iamCredentialId: string | undefined constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter) { super(profileName, lspAuth, eventEmitter) @@ -486,9 +504,10 @@ export class IamLogin extends BaseLogin { } async logout() { - // if (this.stsCredentialId) { - // await this.lspAuth.invalidateStsCredential(this.iamCredentialId) - // } + if (this.iamCredentialId) { + await this.lspAuth.invalidateIamCredential(this.iamCredentialId) + } + await this.deleteProfile(this.profileName) this.updateConnectionState('notConnected') this._data = undefined // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) @@ -502,6 +521,10 @@ export class IamLogin extends BaseLogin { } } + async deleteProfile(profileName: string) { + await this.lspAuth.deleteIamProfile(profileName) + } + /** * Restore the connection state and connection details to memory, if they exist. */ From ff181f49c7340f51af084a079e1de6b8fee07e0d Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 16 Jun 2025 15:31:07 -0400 Subject: [PATCH 09/70] start modifying auth2 consumers --- .../amazonqFeatureDev/client/featureDev.ts | 8 +- packages/core/src/auth/auth2.ts | 30 +++--- .../src/codewhisperer/client/codewhisperer.ts | 101 ++++++++++++++++-- .../region/regionProfileManager.ts | 37 ++++--- .../core/src/codewhisperer/util/authUtil.ts | 33 +++--- .../shared/clients/codewhispererChatClient.ts | 15 ++- 6 files changed, 173 insertions(+), 51 deletions(-) diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts index 419f6969cc6..9ac57272c9b 100644 --- a/packages/core/src/amazonqFeatureDev/client/featureDev.ts +++ b/packages/core/src/amazonqFeatureDev/client/featureDev.ts @@ -51,7 +51,11 @@ const writeAPIRetryOptions = { // Create a client for featureDev proxy client based off of aws sdk v2 export async function createFeatureDevProxyClient(options?: Partial): Promise { - 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, @@ -59,10 +63,10 @@ export async function createFeatureDevProxyClient(options?: Partial( @@ -217,7 +218,7 @@ export class LanguageClientAuth { return { profile, ssoSession } } - updateBearerToken(request: UpdateCredentialsParams) { + updateBearerToken(request: UpdateCredentialsParams | undefined) { return this.client.sendRequest(bearerCredentialsUpdateRequestType.method, request) } @@ -225,7 +226,7 @@ export class LanguageClientAuth { 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 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 index 81d238c1496..adfc0c3c49d 100644 --- 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 { - 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 { 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 { + // 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 { + // 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 { + // 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 { diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index 2645c573249..59940682a0a 100644 --- 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 { - 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 index 1d512d85b2d..accc102c665 100644 --- 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 + isIamSession(): boolean + getCredential(): Promise 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 index 1210adb7e30..b4be5bb40df 100644 --- a/packages/core/src/shared/clients/codewhispererChatClient.ts +++ b/packages/core/src/shared/clients/codewhispererChatClient.ts @@ -7,17 +7,28 @@ import { ConfiguredRetryStrategy } from '@smithy/util-retry' import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer' import { AuthUtil } from '../../codewhisperer/util/authUtil' import { getUserAgent } from '../telemetry/util' +import { Credentials, Token } from 'aws-sdk' // Create a client for featureDev streaming based off of aws sdk v3 export async function createCodeWhispererChatStreamingClient(): Promise { - const bearerToken = await AuthUtil.instance.getToken() + const credential = await AuthUtil.instance.getCredential() + const authConfig = + typeof credential === 'string' + ? { token: new Token({ token: credential }) } + : { + credentials: new Credentials({ + accessKeyId: credential.accessKeyId, + secretAccessKey: credential.secretAccessKey, + sessionToken: credential.sessionToken, + }), + } const cwsprConfig = getCodewhispererConfig() const streamingClient = new CodeWhispererStreaming({ region: cwsprConfig.region, endpoint: cwsprConfig.endpoint, - token: { token: bearerToken }, customUserAgent: getUserAgent(), retryStrategy: new ConfiguredRetryStrategy(1, (attempt: number) => 500 + attempt ** 10), + ...authConfig, }) return streamingClient } From eb3526a159ceaa0ac15d696cb1031081d505705c Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 17 Jun 2025 16:08:12 -0400 Subject: [PATCH 10/70] undo unnecessary changes --- aws-toolkit-vscode.code-workspace | 6 -- packages/amazonq/.vscode/launch.json | 4 +- .../amazonqFeatureDev/client/featureDev.ts | 8 +- .../src/codewhisperer/client/codewhisperer.ts | 87 +------------------ .../region/regionProfileManager.ts | 48 +++++----- .../core/src/codewhisperer/util/authUtil.ts | 15 +++- .../shared/clients/codewhispererChatClient.ts | 15 +--- 7 files changed, 44 insertions(+), 139 deletions(-) diff --git a/aws-toolkit-vscode.code-workspace b/aws-toolkit-vscode.code-workspace index 922fc08f787..f03aafae2fe 100644 --- a/aws-toolkit-vscode.code-workspace +++ b/aws-toolkit-vscode.code-workspace @@ -12,12 +12,6 @@ { "path": "packages/amazonq", }, - { - "path": "../language-server-runtimes", - }, - { - "path": "../language-servers", - }, ], "settings": { "typescript.tsdk": "node_modules/typescript/lib", diff --git a/packages/amazonq/.vscode/launch.json b/packages/amazonq/.vscode/launch.json index cdeabe152a9..7b7fa1a9150 100644 --- a/packages/amazonq/.vscode/launch.json +++ b/packages/amazonq/.vscode/launch.json @@ -15,8 +15,8 @@ "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010", "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080", // Below allows for overrides used during development - "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", - "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" + // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", + // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" }, "envFile": "${workspaceFolder}/.local.env", "outFiles": ["${workspaceFolder}/dist/**/*.js", "${workspaceFolder}/../core/dist/**/*.js"], diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts index 9ac57272c9b..62e870b51fe 100644 --- a/packages/core/src/amazonqFeatureDev/client/featureDev.ts +++ b/packages/core/src/amazonqFeatureDev/client/featureDev.ts @@ -51,11 +51,7 @@ const writeAPIRetryOptions = { // Create a client for featureDev proxy client based off of aws sdk v2 export async function createFeatureDevProxyClient(options?: Partial): Promise { - 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 bearerToken = await AuthUtil.instance.getBearerToken() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( Service, @@ -63,10 +59,10 @@ export async function createFeatureDevProxyClient(options?: Partial { const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled() session.setFetchCredentialStart() - const credential = await AuthUtil.instance.getCredential() - if (typeof credential !== 'string') { - throw new TypeError('Cannot create user SDK client from IAM credentials') - } + const credential = await AuthUtil.instance.getBearerToken() session.setSdkApiCallStart() const cwsprConfig = getCodewhispererConfig() @@ -161,88 +158,6 @@ export class DefaultCodeWhispererClient { return AuthUtil.instance.isConnected() // TODO: Handle IAM credentials } - // private async createServiceSdkClient(credential: IamCredentials): Promise { - // 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 { - // 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 { - // 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 { diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index 59940682a0a..4bac191158a 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -13,8 +13,8 @@ import { once } from '../../shared/utilities/functionUtils' import CodeWhispererUserClient from '../client/codewhispereruserclient' import { Credentials, HttpRequest, Service } from 'aws-sdk' import { ServiceOptions } from '../../shared/awsClientBuilder' -import tokenApiConfig = require('../client/user-service-2.json') -import iamApiConfig = require('../client/service-2.json') +import userApiConfig = require('../client/user-service-2.json') +import apiConfig = require('../client/service-2.json') import { createConstantMap } from '../../shared/utilities/tsUtils' import { getLogger } from '../../shared/logger/logger' import { pageableToCollection } from '../../shared/utilities/collectionUtils' @@ -426,28 +426,30 @@ export class RegionProfileManager { // Visible for testing only, do not use this directly, please use createQClient(profile) async _createQClient(region: string, endpoint: string): Promise { - 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 + let authConfig: ServiceOptions = {} + if (this.authProvider.isSsoSession()) { + const credential = await this.authProvider.getBearerToken() + authConfig = { + onRequestSetup: [ + (req: any) => { + req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => { + httpRequest.headers['Authorization'] = `Bearer ${credential}` + }) + }, + ], + } + } else if (this.authProvider.isIamSession()) { + const credential = await this.authProvider.getIamCredential() + authConfig = { + credentials: new Credentials({ + accessKeyId: credential.accessKeyId, + secretAccessKey: credential.secretAccessKey, + sessionToken: credential.sessionToken, + }), + } + } const serviceOption: ServiceOptions = { - apiConfig: apiConfig, + apiConfig: this.authProvider.isSsoSession() ? userApiConfig : apiConfig, region: region, endpoint: endpoint, ...authConfig, diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index accc102c665..41389e5e19d 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -61,7 +61,8 @@ export interface IAuthProvider { isIdcConnection(): boolean isSsoSession(): boolean isIamSession(): boolean - getCredential(): Promise + getBearerToken(): Promise + getIamCredential(): Promise readonly profileName: string readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } } @@ -201,9 +202,17 @@ export class AuthUtil implements IAuthProvider { return response } - async getCredential() { + async getBearerToken() { + if (this.isSsoSession()) { + return (await (this.session as SsoLogin).getCredential()).credential + } else { + throw new ToolkitError('Cannot get credential without logging in.') + } + } + + async getIamCredential() { if (this.session) { - return (await this.session.getCredential()).credential + return (await (this.session as IamLogin).getCredential()).credential } else { throw new ToolkitError('Cannot get credential without logging in.') } diff --git a/packages/core/src/shared/clients/codewhispererChatClient.ts b/packages/core/src/shared/clients/codewhispererChatClient.ts index b4be5bb40df..93b7b3bceb4 100644 --- a/packages/core/src/shared/clients/codewhispererChatClient.ts +++ b/packages/core/src/shared/clients/codewhispererChatClient.ts @@ -7,28 +7,17 @@ import { ConfiguredRetryStrategy } from '@smithy/util-retry' import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer' import { AuthUtil } from '../../codewhisperer/util/authUtil' import { getUserAgent } from '../telemetry/util' -import { Credentials, Token } from 'aws-sdk' // Create a client for featureDev streaming based off of aws sdk v3 export async function createCodeWhispererChatStreamingClient(): Promise { - const credential = await AuthUtil.instance.getCredential() - const authConfig = - typeof credential === 'string' - ? { token: new Token({ token: credential }) } - : { - credentials: new Credentials({ - accessKeyId: credential.accessKeyId, - secretAccessKey: credential.secretAccessKey, - sessionToken: credential.sessionToken, - }), - } + const bearerToken = await AuthUtil.instance.getBearerToken() const cwsprConfig = getCodewhispererConfig() const streamingClient = new CodeWhispererStreaming({ region: cwsprConfig.region, endpoint: cwsprConfig.endpoint, + token: { token: bearerToken }, customUserAgent: getUserAgent(), retryStrategy: new ConfiguredRetryStrategy(1, (attempt: number) => 500 + attempt ** 10), - ...authConfig, }) return streamingClient } From 849006b1cb8b69e592c16924531b845eb0e455f8 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 18 Jun 2025 09:48:01 -0400 Subject: [PATCH 11/70] undo more unnecessary changes --- .../tracker/codewhispererTracker.test.ts | 5 ++--- .../amazonqFeatureDev/client/featureDev.ts | 2 +- .../src/codewhisperer/client/codewhisperer.ts | 21 ++++++++++--------- .../region/regionProfileManager.ts | 4 ++-- .../core/src/codewhisperer/util/authUtil.ts | 4 ++-- .../shared/clients/codewhispererChatClient.ts | 2 +- 6 files changed, 19 insertions(+), 19 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts index a7590f18eb4..a43720c81be 100644 --- a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts +++ b/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, }) diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts index 62e870b51fe..419f6969cc6 100644 --- a/packages/core/src/amazonqFeatureDev/client/featureDev.ts +++ b/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): Promise { - const bearerToken = await AuthUtil.instance.getBearerToken() + const bearerToken = await AuthUtil.instance.getToken() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( Service, diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts index 16f75cda46b..bc971a3c0da 100644 --- 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, 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 { 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 { 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', () => { diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index 4bac191158a..d726175e713 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -428,12 +428,12 @@ export class RegionProfileManager { async _createQClient(region: string, endpoint: string): Promise { 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}` }) }, ], diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 41389e5e19d..63de5a0a8f3 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -61,7 +61,7 @@ export interface IAuthProvider { isIdcConnection(): boolean isSsoSession(): boolean isIamSession(): boolean - getBearerToken(): Promise + getToken(): Promise getIamCredential(): Promise 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 { diff --git a/packages/core/src/shared/clients/codewhispererChatClient.ts b/packages/core/src/shared/clients/codewhispererChatClient.ts index 93b7b3bceb4..1210adb7e30 100644 --- a/packages/core/src/shared/clients/codewhispererChatClient.ts +++ b/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 { - const bearerToken = await AuthUtil.instance.getBearerToken() + const bearerToken = await AuthUtil.instance.getToken() const cwsprConfig = getCodewhispererConfig() const streamingClient = new CodeWhispererStreaming({ region: cwsprConfig.region, From 1968d41244921f0037100ee089f487248abe60d3 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 18 Jun 2025 10:02:30 -0400 Subject: [PATCH 12/70] fix logout bug --- packages/core/src/auth/auth2.ts | 1 + packages/core/src/codewhisperer/util/authUtil.ts | 5 ----- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 3add4c70faf..11c572dcf2b 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -367,6 +367,7 @@ export class SsoLogin extends BaseLogin { } async logout() { + this.lspAuth.deleteBearerToken() if (this.ssoTokenId) { await this.lspAuth.invalidateSsoToken(this.ssoTokenId) } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 63de5a0a8f3..46a70dc2a25 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -192,11 +192,6 @@ export class AuthUtil implements IAuthProvider { } logout() { - if (!this.isSsoSession()) { - // Only SSO requires logout - return - } - this.lspAuth.deleteBearerToken() const response = this.session?.logout() this.session = undefined return response From 06c5ad8be6409357920bdb7600b999e83e56aaf2 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 19 Jun 2025 00:02:55 -0400 Subject: [PATCH 13/70] Fix bug where profile failed to be retrieved after signing out and back in --- .../region/regionProfileManager.ts | 29 ++++++++++--------- .../core/src/codewhisperer/util/authUtil.ts | 6 ++-- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index d726175e713..ac1edb0704f 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -11,7 +11,7 @@ import { showConfirmationMessage } from '../../shared/utilities/messages' import globals from '../../shared/extensionGlobals' import { once } from '../../shared/utilities/functionUtils' import CodeWhispererUserClient from '../client/codewhispereruserclient' -import { Credentials, HttpRequest, Service } from 'aws-sdk' +import { Credentials, Service } from 'aws-sdk' import { ServiceOptions } from '../../shared/awsClientBuilder' import userApiConfig = require('../client/user-service-2.json') import apiConfig = require('../client/service-2.json') @@ -426,34 +426,35 @@ export class RegionProfileManager { // Visible for testing only, do not use this directly, please use createQClient(profile) async _createQClient(region: string, endpoint: string): Promise { - let authConfig: ServiceOptions = {} + let serviceOption: ServiceOptions = {} if (this.authProvider.isSsoSession()) { const token = await this.authProvider.getToken() - authConfig = { + serviceOption = { + apiConfig: userApiConfig, + region: region, + endpoint: endpoint, + credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }), onRequestSetup: [ - (req: any) => { - req.on('build', ({ httpRequest }: { httpRequest: HttpRequest }) => { + (req) => { + req.on('build', ({ httpRequest }) => { httpRequest.headers['Authorization'] = `Bearer ${token}` }) }, ], - } + } as ServiceOptions } else if (this.authProvider.isIamSession()) { const credential = await this.authProvider.getIamCredential() - authConfig = { + serviceOption = { + apiConfig: apiConfig, + region: region, + endpoint: endpoint, credentials: new Credentials({ accessKeyId: credential.accessKeyId, secretAccessKey: credential.secretAccessKey, sessionToken: credential.sessionToken, }), - } + } as ServiceOptions } - const serviceOption: ServiceOptions = { - apiConfig: this.authProvider.isSsoSession() ? userApiConfig : apiConfig, - region: region, - endpoint: endpoint, - ...authConfig, - } as ServiceOptions const c = (await globals.sdkClientBuilder.createAwsService( Service, diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 46a70dc2a25..7788c2e3378 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -192,9 +192,7 @@ export class AuthUtil implements IAuthProvider { } logout() { - const response = this.session?.logout() - this.session = undefined - return response + return this.session?.logout() } async getToken() { @@ -361,6 +359,8 @@ export class AuthUtil implements IAuthProvider { await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn) await this.regionProfileManager.clearCache() } + // Session should only be nullified after all methods dependent on session are evaluated + this.session = undefined } if (state === 'connected') { const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined From 60ffc92b6e309efb6cf61f523d9faf6f5445c884 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 19 Jun 2025 10:28:27 -0400 Subject: [PATCH 14/70] Fix region profile selector not triggering --- .../core/src/codewhisperer/client/codewhisperer.ts | 7 ++++--- .../src/codewhisperer/region/regionProfileManager.ts | 8 ++++---- packages/core/src/codewhisperer/util/authUtil.ts | 4 ++-- .../src/login/webview/vue/amazonq/backend_amazonq.ts | 11 +++++------ 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts index bc971a3c0da..0459be92d96 100644 --- a/packages/core/src/codewhisperer/client/codewhisperer.ts +++ b/packages/core/src/codewhisperer/client/codewhisperer.ts @@ -83,15 +83,16 @@ export type Imports = CodeWhispererUserClient.Imports export class DefaultCodeWhispererClient { private async createSdkClient(): Promise { + throw new Error('Do not call this function until IAM is supported by LSP identity server') + 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: new Credentials({ accessKeyId: credential.accessKeyId, secretAccessKey: credential.secretAccessKey }), + credentials: undefined, endpoint: cwsprConfig.endpoint, onRequestSetup: [ (req) => { @@ -127,7 +128,6 @@ export class DefaultCodeWhispererClient { const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled() session.setFetchCredentialStart() const bearerToken = await AuthUtil.instance.getToken() - session.setSdkApiCallStart() const cwsprConfig = getCodewhispererConfig() return (await globals.sdkClientBuilder.createAwsService( @@ -137,6 +137,7 @@ export class DefaultCodeWhispererClient { region: cwsprConfig.region, endpoint: cwsprConfig.endpoint, maxRetries: maxRetries, + credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }), onRequestSetup: [ (req) => { req.on('build', ({ httpRequest }) => { diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index ac1edb0704f..edd2f216fcc 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -145,7 +145,7 @@ export class RegionProfileManager { async listRegionProfile(): Promise { this._profiles = [] - if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { + if (!this.authProvider.isConnected()) { return [] } const availableProfiles: RegionProfile[] = [] @@ -405,7 +405,7 @@ export class RegionProfileManager { if (this.authProvider.isBuilderIdConnection()) { return false } - return this.authProvider.isIdcConnection() && this.activeRegionProfile === undefined + return (this.authProvider.isIdcConnection() || this.authProvider.isIamSession()) && this.activeRegionProfile === undefined } async clearCache() { @@ -414,8 +414,8 @@ export class RegionProfileManager { // TODO: Should maintain sdk client in a better way async createQClient(profile: RegionProfile): Promise { - if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { - throw new Error('No valid SSO connection') + if (!this.authProvider.isConnected()) { + throw new Error('No valid connection') } const endpoint = endpoints.get(profile.region) if (!endpoint) { diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 7788c2e3378..8a596b0f62c 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -355,7 +355,7 @@ export class AuthUtil implements IAuthProvider { private async refreshState(state = this.getAuthState()) { if (state === 'expired' || state === 'notConnected') { this.lspAuth.deleteBearerToken() - if (this.isIdcConnection()) { + if (this.isIdcConnection() || this.isIamSession()) { await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn) await this.regionProfileManager.clearCache() } @@ -370,7 +370,7 @@ export class AuthUtil implements IAuthProvider { await this.lspAuth.updateIamCredential(params) } - if (this.isIdcConnection()) { + if (this.isIdcConnection() || this.isIamSession()) { await this.regionProfileManager.restoreProfileSelection() } } diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index 8923c1ba582..d23132333e5 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -203,13 +203,10 @@ export class AmazonQLoginWebview extends CommonAuthWebview { secretKey: string ): Promise { getLogger().debug(`called startIamCredentialSetup()`) - // Defining separate auth function to emit telemetry before returning from setup + // Defining separate auth function to emit telemetry before returning from this method const runAuth = async (): Promise => { try { await AuthUtil.instance.login(accessKey, secretKey, 'iam') - // Add auth telemetry - this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) - void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Credentials') } catch (e) { getLogger().error('Failed submitting credentials %O', e) return { id: this.id, text: e as string } @@ -217,11 +214,13 @@ export class AmazonQLoginWebview extends CommonAuthWebview { // Enable code suggestions vsCodeState.isFreeTierLimitReached = false await Commands.tryExecute('aws.amazonq.enableCodeSuggestions') + + this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata()) + + void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Credentials') } const result = await runAuth() - - // Emit telemetry this.storeMetricMetadata({ credentialSourceId: 'sharedCredentials', authEnabledFeatures: 'codewhisperer', From e1cfdc3ce444a2c91724036a5ee5565fe229ae3b Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 19 Jun 2025 13:53:57 -0400 Subject: [PATCH 15/70] Add support for IAM credentials to region profile manager --- .../region/regionProfileManager.test.ts | 14 +- .../region/regionProfileManager.ts | 163 ++++++++++++------ .../core/src/codewhisperer/util/authUtil.ts | 13 +- .../codewhisperer/util/customizationUtil.ts | 2 +- .../webview/vue/amazonq/backend_amazonq.ts | 4 - .../test/amazonq/customizationUtil.test.ts | 2 +- 6 files changed, 127 insertions(+), 71 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index a858c3e659e..53c503f5e66 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -62,7 +62,7 @@ describe('RegionProfileManager', async function () { const mockClient = { listAvailableProfiles: listProfilesStub, } - const createClientStub = sinon.stub(regionProfileManager, '_createQClient').resolves(mockClient) + const createClientStub = sinon.stub(regionProfileManager, '_createQUserClient').resolves(mockClient) const profileList = await regionProfileManager.listRegionProfile() @@ -272,11 +272,11 @@ describe('RegionProfileManager', async function () { }) }) - describe('createQClient', function () { + describe('createQUserClient', function () { it(`should configure the endpoint and region from a profile`, async function () { await setupConnection('idc') - const iadClient = await regionProfileManager.createQClient({ + const iadClient = await regionProfileManager.createQUserClient({ name: 'foo', region: 'us-east-1', arn: 'arn', @@ -286,7 +286,7 @@ describe('RegionProfileManager', async function () { assert.deepStrictEqual(iadClient.config.region, 'us-east-1') assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/') - const fraClient = await regionProfileManager.createQClient({ + const fraClient = await regionProfileManager.createQUserClient({ name: 'bar', region: 'eu-central-1', arn: 'arn', @@ -302,7 +302,7 @@ describe('RegionProfileManager', async function () { await assert.rejects( async () => { - await regionProfileManager.createQClient({ + await regionProfileManager.createQUserClient({ name: 'foo', region: 'ap-east-1', arn: 'arn', @@ -314,7 +314,7 @@ describe('RegionProfileManager', async function () { await assert.rejects( async () => { - await regionProfileManager.createQClient({ + await regionProfileManager.createQUserClient({ name: 'foo', region: 'unknown-somewhere', arn: 'arn', @@ -330,7 +330,7 @@ describe('RegionProfileManager', async function () { await regionProfileManager.switchRegionProfile(profileFoo, 'user') assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) - const client = await regionProfileManager._createQClient('eu-central-1', 'https://amazon.com/') + const client = await regionProfileManager._createQUserClient('eu-central-1', 'https://amazon.com/') assert.deepStrictEqual(client.config.region, 'eu-central-1') assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/') diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index edd2f216fcc..491cab28a6e 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -11,6 +11,7 @@ import { showConfirmationMessage } from '../../shared/utilities/messages' import globals from '../../shared/extensionGlobals' import { once } from '../../shared/utilities/functionUtils' import CodeWhispererUserClient from '../client/codewhispereruserclient' +import CodeWhispererClient from '../client/codewhispererclient' import { Credentials, Service } from 'aws-sdk' import { ServiceOptions } from '../../shared/awsClientBuilder' import userApiConfig = require('../client/user-service-2.json') @@ -152,30 +153,63 @@ export class RegionProfileManager { const failedRegions: string[] = [] for (const [region, endpoint] of endpoints.entries()) { - const client = await this._createQClient(region, endpoint) - const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => - client.listAvailableProfiles(request).promise() - const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {} try { - const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') - .flatten() - .promise() - const mappedPfs = profiles.map((it) => { - let accntId = '' - try { - accntId = parse(it.arn).accountId - } catch (e) {} - - return { - name: it.profileName, - region: region, - arn: it.arn, - description: accntId, + // Get region profiles (Q developer profiles) from Q client and authenticate with SSO token + if (this.authProvider.isIdcConnection()) { + const client = await this._createQUserClient(region, endpoint) + const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => { + return client.listAvailableProfiles(request).promise() } - }) + const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {} + const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') + .flatten() + .promise() + const mappedPfs = profiles.map((it) => { + let accntId = '' + try { + accntId = parse(it.arn).accountId + } catch (e) {} + + return { + name: it.profileName, + region: region, + arn: it.arn, + description: accntId, + } + }) - availableProfiles.push(...mappedPfs) - RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) + availableProfiles.push(...mappedPfs) + RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) + } + // Get region profiles (Q developer profiles) from Q client and authenticate with IAM credentials + else if (this.authProvider.isIamSession()) { + const client = await this._createQServiceClient(region, endpoint) + const requester = async (request: CodeWhispererClient.ListProfilesRequest) => { + return client.listProfiles(request).promise() + } + const request: CodeWhispererClient.ListProfilesRequest = {} + const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') + .flatten() + .promise() + const mappedPfs = profiles.map((it) => { + let accntId = '' + try { + accntId = parse(it.arn).accountId + } catch (e) {} + + return { + name: it.profileName, + region: region, + arn: it.arn, + description: accntId, + } + }) + + availableProfiles.push(...mappedPfs) + RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) + } else { + throw new ToolkitError('Failed to list profiles when signed out of identity center and IAM credentials') + } } catch (e) { const logMsg = isAwsError(e) ? `requestId=${e.requestId}; message=${e.message}` : (e as Error).message RegionProfileManager.logger.error(`Failed to list profiles for region ${region}: ${logMsg}`) @@ -201,7 +235,7 @@ export class RegionProfileManager { } async switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) { - if (!this.authProvider.isConnected() || !this.authProvider.isIdcConnection()) { + if (!this.authProvider.isConnected()) { return } @@ -413,7 +447,8 @@ export class RegionProfileManager { } // TODO: Should maintain sdk client in a better way - async createQClient(profile: RegionProfile): Promise { + // Create a Q user client compatible with SSO tokens + async createQUserClient(profile: RegionProfile): Promise { if (!this.authProvider.isConnected()) { throw new Error('No valid connection') } @@ -421,47 +456,63 @@ export class RegionProfileManager { if (!endpoint) { throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`) } - return this._createQClient(profile.region, endpoint) + return this._createQUserClient(profile.region, endpoint) } - // Visible for testing only, do not use this directly, please use createQClient(profile) - async _createQClient(region: string, endpoint: string): Promise { - let serviceOption: ServiceOptions = {} - if (this.authProvider.isSsoSession()) { - const token = await this.authProvider.getToken() - serviceOption = { - apiConfig: userApiConfig, - region: region, - endpoint: endpoint, - credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }), - onRequestSetup: [ - (req) => { - req.on('build', ({ httpRequest }) => { - httpRequest.headers['Authorization'] = `Bearer ${token}` - }) - }, - ], - } as ServiceOptions - } else if (this.authProvider.isIamSession()) { - const credential = await this.authProvider.getIamCredential() - serviceOption = { - apiConfig: apiConfig, - region: region, - endpoint: endpoint, - credentials: new Credentials({ - accessKeyId: credential.accessKeyId, - secretAccessKey: credential.secretAccessKey, - sessionToken: credential.sessionToken, - }), - } as ServiceOptions + // Create a Q service client compatible with IAM credentials + async createQServiceClient(profile: RegionProfile): Promise { + if (!this.authProvider.isConnected()) { + throw new Error('No valid connection') } + const endpoint = endpoints.get(profile.region) + if (!endpoint) { + throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`) + } + return this._createQServiceClient(profile.region, endpoint) + } - const c = (await globals.sdkClientBuilder.createAwsService( + // Visible for testing only, do not use this directly, please use createQUserClient(profile) + async _createQUserClient(region: string, endpoint: string): Promise { + const token = await this.authProvider.getToken() + const serviceOption: ServiceOptions = { + apiConfig: userApiConfig, + region: region, + endpoint: endpoint, + credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }), + onRequestSetup: [ + (req) => { + req.on('build', ({ httpRequest }) => { + httpRequest.headers['Authorization'] = `Bearer ${token}` + }) + }, + ], + } as ServiceOptions + + return (await globals.sdkClientBuilder.createAwsService( Service, serviceOption, undefined )) as CodeWhispererUserClient + } - return c + // Visible for testing only, do not use this directly, please use createQServiceClient(profile) + async _createQServiceClient(region: string, endpoint: string): Promise { + const credential = await this.authProvider.getIamCredential() + const serviceOption: ServiceOptions = { + apiConfig: apiConfig, + region: region, + endpoint: endpoint, + credentials: new Credentials({ + accessKeyId: credential.accessKeyId, + secretAccessKey: credential.secretAccessKey, + sessionToken: credential.sessionToken, + }), + } as ServiceOptions + + return (await globals.sdkClientBuilder.createAwsService( + Service, + serviceOption, + undefined + )) as CodeWhispererClient } } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 8a596b0f62c..927f05ae93f 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -192,12 +192,17 @@ export class AuthUtil implements IAuthProvider { } logout() { + // session will be nullified the next time refreshState() is called return this.session?.logout() } async getToken() { if (this.isSsoSession()) { - return (await (this.session as SsoLogin).getCredential()).credential + const token = (await this.session!.getCredential()).credential + if (typeof token !== 'string') { + throw new ToolkitError('Cannot get token with IAM session') + } + return token } else { throw new ToolkitError('Cannot get credential without logging in.') } @@ -205,7 +210,11 @@ export class AuthUtil implements IAuthProvider { async getIamCredential() { if (this.session) { - return (await (this.session as IamLogin).getCredential()).credential + const credential = (await this.session.getCredential()).credential + if (typeof credential !== 'object') { + throw new ToolkitError('Cannot get token with SSO session') + } + return credential } else { throw new ToolkitError('Cannot get credential without logging in.') } diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts index 04bb85d7a43..ed00663bb26 100644 --- a/packages/core/src/codewhisperer/util/customizationUtil.ts +++ b/packages/core/src/codewhisperer/util/customizationUtil.ts @@ -49,7 +49,7 @@ export class CustomizationProvider { } static async init(profile: RegionProfile): Promise { - const client = await AuthUtil.instance.regionProfileManager.createQClient(profile) + const client = await AuthUtil.instance.regionProfileManager.createQUserClient(profile) return new CustomizationProvider(client, profile) } } diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index d23132333e5..6a0d245849e 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -174,10 +174,6 @@ export class AmazonQLoginWebview extends CommonAuthWebview { @withTelemetryContext({ name: 'signout', class: className }) override async signout(): Promise { - if (!AuthUtil.instance.isSsoSession()) { - throw new ToolkitError(`Cannot signout non-SSO connection`) - } - this.storeMetricMetadata({ authEnabledFeatures: 'codewhisperer', isReAuth: true, diff --git a/packages/core/src/test/amazonq/customizationUtil.test.ts b/packages/core/src/test/amazonq/customizationUtil.test.ts index 19e59b91c03..0ed19cf2197 100644 --- a/packages/core/src/test/amazonq/customizationUtil.test.ts +++ b/packages/core/src/test/amazonq/customizationUtil.test.ts @@ -43,7 +43,7 @@ describe('customizationProvider', function () { regionProfileManager: regionProfileManager, } sinon.stub(AuthUtil, 'instance').get(() => mockAuthUtil) - const createClientStub = sinon.stub(regionProfileManager, 'createQClient') + const createClientStub = sinon.stub(regionProfileManager, 'createQUserClient') const mockProfile = { name: 'foo', region: 'us-east-1', From 422d0856449489a03127838b5887f5ac6da46a4b Mon Sep 17 00:00:00 2001 From: Yuxian Zhang Date: Thu, 19 Jun 2025 17:58:38 -0400 Subject: [PATCH 16/70] feat: remember iam access key --- aws-toolkit-vscode.code-workspace | 6 ++ packages/amazonq/.vscode/launch.json | 4 +- packages/amazonq/src/lsp/client.ts | 11 +++- .../webview/vue/amazonq/backend_amazonq.ts | 11 +++- .../core/src/login/webview/vue/backend.ts | 12 ++++ packages/core/src/login/webview/vue/login.vue | 59 +++++++++++++++---- .../webview/vue/toolkit/backend_toolkit.ts | 8 +++ packages/core/src/shared/globalState.ts | 1 + packages/core/src/shared/settings.ts | 1 + 9 files changed, 98 insertions(+), 15 deletions(-) diff --git a/aws-toolkit-vscode.code-workspace b/aws-toolkit-vscode.code-workspace index f03aafae2fe..922fc08f787 100644 --- a/aws-toolkit-vscode.code-workspace +++ b/aws-toolkit-vscode.code-workspace @@ -12,6 +12,12 @@ { "path": "packages/amazonq", }, + { + "path": "../language-server-runtimes", + }, + { + "path": "../language-servers", + }, ], "settings": { "typescript.tsdk": "node_modules/typescript/lib", diff --git a/packages/amazonq/.vscode/launch.json b/packages/amazonq/.vscode/launch.json index 7b7fa1a9150..cdeabe152a9 100644 --- a/packages/amazonq/.vscode/launch.json +++ b/packages/amazonq/.vscode/launch.json @@ -15,8 +15,8 @@ "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010", "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080", // Below allows for overrides used during development - // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", - // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" + "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js", + "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js" }, "envFile": "${workspaceFolder}/.local.env", "outFiles": ["${workspaceFolder}/dist/**/*.js", "${workspaceFolder}/../core/dist/**/*.js"], diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts index 4395ade9a2c..4c3a1a4d92d 100644 --- a/packages/amazonq/src/lsp/client.ts +++ b/packages/amazonq/src/lsp/client.ts @@ -164,6 +164,9 @@ export async function startLanguageServer( }, credentials: { providesBearerToken: true, + // Add IAM credentials support + providesIamCredentials: true, + supportsAssumeRole: true, }, }, /** @@ -211,9 +214,10 @@ export async function startLanguageServer( /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */ AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => { + const activeProfile = AuthUtil.instance.regionProfileManager.activeRegionProfile void pushConfigUpdate(client, { type: 'profile', - profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn, + profileArn: activeProfile?.arn, }) }) @@ -286,6 +290,11 @@ async function postStartLanguageServer( sso: { startUrl: AuthUtil.instance.connection?.startUrl, }, + // Add IAM credentials metadata + iam: { + region: AuthUtil.instance.connection?.region, + accesskey: AuthUtil.instance.connection?.accessKey, + }, } }) diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index 6a0d245849e..af8a06cd2a7 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ import * as vscode from 'vscode' -import { AwsConnection, SsoConnection } from '../../../../auth/connection' +import { AwsConnection, IamProfile, SsoConnection } from '../../../../auth/connection' import { AuthUtil } from '../../../../codewhisperer/util/authUtil' import { CommonAuthWebview } from '../backend' import { awsIdSignIn } from '../../../../codewhisperer/util/showSsoPrompt' @@ -200,6 +200,9 @@ export class AmazonQLoginWebview extends CommonAuthWebview { ): Promise { getLogger().debug(`called startIamCredentialSetup()`) // Defining separate auth function to emit telemetry before returning from this method + await globals.globalState.update('recentIamKeys', { + accessKey: accessKey, + }) const runAuth = async (): Promise => { try { await AuthUtil.instance.login(accessKey, secretKey, 'iam') @@ -227,6 +230,12 @@ export class AmazonQLoginWebview extends CommonAuthWebview { return result } + async listIamCredentialProfiles(): Promise { + // Amazon Q only supports 1 connection at a time, + // so there isn't a need to de-duplicate connections. + return [] + } + /** If users are unauthenticated in Q/CW, we should always display the auth screen. */ async quitLoginScreen() {} diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts index edb1980a8c0..8bb64c325d4 100644 --- a/packages/core/src/login/webview/vue/backend.ts +++ b/packages/core/src/login/webview/vue/backend.ts @@ -19,6 +19,7 @@ import { scopesCodeWhispererChat, scopesSsoAccountAccess, SsoConnection, + IamProfile, TelemetryMetadata, } from '../../../auth/connection' import { Auth } from '../../../auth/auth' @@ -207,6 +208,8 @@ export abstract class CommonAuthWebview extends VueWebview { abstract listRegionProfiles(): Promise + abstract listIamCredentialProfiles(): Promise + abstract selectRegionProfile(profile: RegionProfile, source: ProfileSwitchIntent): Promise /** @@ -296,6 +299,15 @@ export abstract class CommonAuthWebview extends VueWebview { return globals.globalState.tryGet('recentSso', Object, { startUrl: '', region: 'us-east-1' }) } + getDefaultIamKeys(): { accessKey: string; secretKey: string } { + const devSettings = DevSettings.instance.get('autofillAccessKey', '') + if (devSettings) { + return { accessKey: devSettings, secretKey: '' } + } + + return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '', secretKey: '' }) + } + cancelAuthFlow() { AuthSSOServer.lastInstance?.cancelCurrentFlow() } diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue index 22fcd413972..e6b83772fbd 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -330,6 +330,10 @@ interface ImportedLogin { type: number startUrl: string region: string + // Add IAM credential fields + profileName?: string + accessKey?: string + secretKey?: string // Note: storing secrets has security implications } export default defineComponent({ @@ -349,6 +353,7 @@ export default defineComponent({ data() { return { existingStartUrls: [] as string[], + existingIamAccessKeys: [] as string[], importedLogins: [] as ImportedLogin[], selectedLoginOption: LoginOption.NONE, stage: 'START' as Stage, @@ -366,6 +371,8 @@ export default defineComponent({ }, async created() { const defaultSso = await this.getDefaultSso() + const defaultIamAccessKey = await this.getDefaultIamAccessKey() + this.accessKey = defaultIamAccessKey.accessKey this.startUrl = defaultSso.startUrl this.selectedRegion = defaultSso.region await this.emitUpdate('created') @@ -437,17 +444,44 @@ export default defineComponent({ const selectedConnection = this.importedLogins[this.selectedLoginOption - LoginOption.IMPORTED_LOGINS] - // Imported connections cannot be Builder IDs, they are filtered out in the client. - const error = await client.startEnterpriseSetup( - selectedConnection.startUrl, - selectedConnection.region, - this.app - ) - if (error) { - this.stage = 'START' - void client.errorNotification(error) - } else { - this.stage = 'CONNECTED' + // // Imported connections cannot be Builder IDs, they are filtered out in the client. + // const error = await client.startEnterpriseSetup( + // selectedConnection.startUrl, + // selectedConnection.region, + // this.app + // ) + // if (error) { + // this.stage = 'START' + // void client.errorNotification(error) + // } else { + // this.stage = 'CONNECTED' + // } + // Handle both SSO and IAM imported connections + if (selectedConnection.type === LoginOption.ENTERPRISE_SSO) { + const error = await client.startEnterpriseSetup( + selectedConnection.startUrl, + selectedConnection.region, + this.app + ) + if (error) { + this.stage = 'START' + void client.errorNotification(error) + } else { + this.stage = 'CONNECTED' + } + } else if (selectedConnection.type === LoginOption.IAM_CREDENTIAL) { + // Use stored IAM credentials + const error = await client.startIamCredentialSetup( + selectedConnection.profileName || '', + selectedConnection.accessKey || '', + selectedConnection.secretKey || '' + ) + if (error) { + this.stage = 'START' + void client.errorNotification(error) + } else { + this.stage = 'CONNECTED' + } } } else if (this.selectedLoginOption === LoginOption.IAM_CREDENTIAL) { this.stage = 'AWS_PROFILE' @@ -581,6 +615,9 @@ export default defineComponent({ async getDefaultSso() { return await client.getDefaultSsoProfile() }, + async getDefaultIamAccessKey() { + return await client.getDefaultIamKeys() + }, handleHelpLinkClick() { void client.emitUiClick('auth_helpLink') }, diff --git a/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts b/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts index caec2c764bc..fb108fab8a8 100644 --- a/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts +++ b/packages/core/src/login/webview/vue/toolkit/backend_toolkit.ts @@ -9,6 +9,7 @@ import { getLogger } from '../../../../shared/logger/logger' import { CommonAuthWebview } from '../backend' import { AwsConnection, + IamProfile, SsoConnection, TelemetryMetadata, createSsoProfile, @@ -90,6 +91,9 @@ export class ToolkitLoginWebview extends CommonAuthWebview { secretKey: string ): Promise { getLogger().debug(`called startIamCredentialSetup()`) + await globals.globalState.update('recentIamKeys', { + accessKey: accessKey, + }) // See submitData() in manageCredentials.vue const runAuth = async () => { const data = { aws_access_key_id: accessKey, aws_secret_access_key: secretKey } @@ -157,6 +161,10 @@ export class ToolkitLoginWebview extends CommonAuthWebview { return (await Auth.instance.listConnections()).filter((conn) => isSsoConnection(conn)) as SsoConnection[] } + async listIamCredentialProfiles(): Promise { + return [] + } + override reauthenticateConnection(): Promise { throw new Error('Method not implemented.') } diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts index 65d761412b8..2a11321e5d3 100644 --- a/packages/core/src/shared/globalState.ts +++ b/packages/core/src/shared/globalState.ts @@ -71,6 +71,7 @@ export type globalKey = | 'lastOsStartTime' | 'recentCredentials' | 'recentSso' + | 'recentIamKeys' // List of regions enabled in AWS Explorer. | 'region' // TODO: implement this via `PromptSettings` instead of globalState. diff --git a/packages/core/src/shared/settings.ts b/packages/core/src/shared/settings.ts index 4e3e99f8207..75a6b03780b 100644 --- a/packages/core/src/shared/settings.ts +++ b/packages/core/src/shared/settings.ts @@ -780,6 +780,7 @@ const devSettings = { amazonqWorkspaceLsp: Record(String, String), ssoCacheDirectory: String, autofillStartUrl: String, + autofillAccessKey: String, webAuth: Boolean, notificationsPollInterval: Number, } From 424a1af75b636378cdbae9475bea3baeac86da33 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 20 Jun 2025 17:26:36 -0400 Subject: [PATCH 17/70] Revert unnecessary region profile changes --- .../region/regionProfileManager.test.ts | 14 +- .../region/regionProfileManager.ts | 133 +++++------------- .../core/src/codewhisperer/util/authUtil.ts | 4 +- .../codewhisperer/util/customizationUtil.ts | 2 +- .../test/amazonq/customizationUtil.test.ts | 2 +- 5 files changed, 44 insertions(+), 111 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index 53c503f5e66..a858c3e659e 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -62,7 +62,7 @@ describe('RegionProfileManager', async function () { const mockClient = { listAvailableProfiles: listProfilesStub, } - const createClientStub = sinon.stub(regionProfileManager, '_createQUserClient').resolves(mockClient) + const createClientStub = sinon.stub(regionProfileManager, '_createQClient').resolves(mockClient) const profileList = await regionProfileManager.listRegionProfile() @@ -272,11 +272,11 @@ describe('RegionProfileManager', async function () { }) }) - describe('createQUserClient', function () { + describe('createQClient', function () { it(`should configure the endpoint and region from a profile`, async function () { await setupConnection('idc') - const iadClient = await regionProfileManager.createQUserClient({ + const iadClient = await regionProfileManager.createQClient({ name: 'foo', region: 'us-east-1', arn: 'arn', @@ -286,7 +286,7 @@ describe('RegionProfileManager', async function () { assert.deepStrictEqual(iadClient.config.region, 'us-east-1') assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/') - const fraClient = await regionProfileManager.createQUserClient({ + const fraClient = await regionProfileManager.createQClient({ name: 'bar', region: 'eu-central-1', arn: 'arn', @@ -302,7 +302,7 @@ describe('RegionProfileManager', async function () { await assert.rejects( async () => { - await regionProfileManager.createQUserClient({ + await regionProfileManager.createQClient({ name: 'foo', region: 'ap-east-1', arn: 'arn', @@ -314,7 +314,7 @@ describe('RegionProfileManager', async function () { await assert.rejects( async () => { - await regionProfileManager.createQUserClient({ + await regionProfileManager.createQClient({ name: 'foo', region: 'unknown-somewhere', arn: 'arn', @@ -330,7 +330,7 @@ describe('RegionProfileManager', async function () { await regionProfileManager.switchRegionProfile(profileFoo, 'user') assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo) - const client = await regionProfileManager._createQUserClient('eu-central-1', 'https://amazon.com/') + const client = await regionProfileManager._createQClient('eu-central-1', 'https://amazon.com/') assert.deepStrictEqual(client.config.region, 'eu-central-1') assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/') diff --git a/packages/core/src/codewhisperer/region/regionProfileManager.ts b/packages/core/src/codewhisperer/region/regionProfileManager.ts index 491cab28a6e..2645c573249 100644 --- a/packages/core/src/codewhisperer/region/regionProfileManager.ts +++ b/packages/core/src/codewhisperer/region/regionProfileManager.ts @@ -11,11 +11,9 @@ import { showConfirmationMessage } from '../../shared/utilities/messages' import globals from '../../shared/extensionGlobals' import { once } from '../../shared/utilities/functionUtils' import CodeWhispererUserClient from '../client/codewhispereruserclient' -import CodeWhispererClient from '../client/codewhispererclient' import { Credentials, Service } from 'aws-sdk' import { ServiceOptions } from '../../shared/awsClientBuilder' import userApiConfig = require('../client/user-service-2.json') -import apiConfig = require('../client/service-2.json') import { createConstantMap } from '../../shared/utilities/tsUtils' import { getLogger } from '../../shared/logger/logger' import { pageableToCollection } from '../../shared/utilities/collectionUtils' @@ -146,70 +144,37 @@ export class RegionProfileManager { async listRegionProfile(): Promise { this._profiles = [] - if (!this.authProvider.isConnected()) { + if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { return [] } const availableProfiles: RegionProfile[] = [] const failedRegions: string[] = [] for (const [region, endpoint] of endpoints.entries()) { + const client = await this._createQClient(region, endpoint) + const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => + client.listAvailableProfiles(request).promise() + const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {} try { - // Get region profiles (Q developer profiles) from Q client and authenticate with SSO token - if (this.authProvider.isIdcConnection()) { - const client = await this._createQUserClient(region, endpoint) - const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) => { - return client.listAvailableProfiles(request).promise() + const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') + .flatten() + .promise() + const mappedPfs = profiles.map((it) => { + let accntId = '' + try { + accntId = parse(it.arn).accountId + } catch (e) {} + + return { + name: it.profileName, + region: region, + arn: it.arn, + description: accntId, } - const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {} - const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') - .flatten() - .promise() - const mappedPfs = profiles.map((it) => { - let accntId = '' - try { - accntId = parse(it.arn).accountId - } catch (e) {} - - return { - name: it.profileName, - region: region, - arn: it.arn, - description: accntId, - } - }) - - availableProfiles.push(...mappedPfs) - RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) - } - // Get region profiles (Q developer profiles) from Q client and authenticate with IAM credentials - else if (this.authProvider.isIamSession()) { - const client = await this._createQServiceClient(region, endpoint) - const requester = async (request: CodeWhispererClient.ListProfilesRequest) => { - return client.listProfiles(request).promise() - } - const request: CodeWhispererClient.ListProfilesRequest = {} - const profiles = await pageableToCollection(requester, request, 'nextToken', 'profiles') - .flatten() - .promise() - const mappedPfs = profiles.map((it) => { - let accntId = '' - try { - accntId = parse(it.arn).accountId - } catch (e) {} - - return { - name: it.profileName, - region: region, - arn: it.arn, - description: accntId, - } - }) + }) - availableProfiles.push(...mappedPfs) - RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) - } else { - throw new ToolkitError('Failed to list profiles when signed out of identity center and IAM credentials') - } + availableProfiles.push(...mappedPfs) + RegionProfileManager.logger.debug(`Found ${mappedPfs.length} profiles in region ${region}`) } catch (e) { const logMsg = isAwsError(e) ? `requestId=${e.requestId}; message=${e.message}` : (e as Error).message RegionProfileManager.logger.error(`Failed to list profiles for region ${region}: ${logMsg}`) @@ -235,7 +200,7 @@ export class RegionProfileManager { } async switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) { - if (!this.authProvider.isConnected()) { + if (!this.authProvider.isConnected() || !this.authProvider.isIdcConnection()) { return } @@ -439,7 +404,7 @@ export class RegionProfileManager { if (this.authProvider.isBuilderIdConnection()) { return false } - return (this.authProvider.isIdcConnection() || this.authProvider.isIamSession()) && this.activeRegionProfile === undefined + return this.authProvider.isIdcConnection() && this.activeRegionProfile === undefined } async clearCache() { @@ -447,32 +412,19 @@ export class RegionProfileManager { } // TODO: Should maintain sdk client in a better way - // Create a Q user client compatible with SSO tokens - async createQUserClient(profile: RegionProfile): Promise { - if (!this.authProvider.isConnected()) { - throw new Error('No valid connection') - } - const endpoint = endpoints.get(profile.region) - if (!endpoint) { - throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`) - } - return this._createQUserClient(profile.region, endpoint) - } - - // Create a Q service client compatible with IAM credentials - async createQServiceClient(profile: RegionProfile): Promise { - if (!this.authProvider.isConnected()) { - throw new Error('No valid connection') + async createQClient(profile: RegionProfile): Promise { + if (!this.authProvider.isConnected() || !this.authProvider.isSsoSession()) { + throw new Error('No valid SSO connection') } const endpoint = endpoints.get(profile.region) if (!endpoint) { throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`) } - return this._createQServiceClient(profile.region, endpoint) + return this._createQClient(profile.region, endpoint) } - // Visible for testing only, do not use this directly, please use createQUserClient(profile) - async _createQUserClient(region: string, endpoint: string): Promise { + // Visible for testing only, do not use this directly, please use createQClient(profile) + async _createQClient(region: string, endpoint: string): Promise { const token = await this.authProvider.getToken() const serviceOption: ServiceOptions = { apiConfig: userApiConfig, @@ -487,32 +439,13 @@ export class RegionProfileManager { }, ], } as ServiceOptions - - return (await globals.sdkClientBuilder.createAwsService( + + const c = (await globals.sdkClientBuilder.createAwsService( Service, serviceOption, undefined )) as CodeWhispererUserClient - } - - // Visible for testing only, do not use this directly, please use createQServiceClient(profile) - async _createQServiceClient(region: string, endpoint: string): Promise { - const credential = await this.authProvider.getIamCredential() - const serviceOption: ServiceOptions = { - apiConfig: apiConfig, - region: region, - endpoint: endpoint, - credentials: new Credentials({ - accessKeyId: credential.accessKeyId, - secretAccessKey: credential.secretAccessKey, - sessionToken: credential.sessionToken, - }), - } as ServiceOptions - return (await globals.sdkClientBuilder.createAwsService( - Service, - serviceOption, - undefined - )) as CodeWhispererClient + return c } } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 927f05ae93f..daf8aebe369 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -364,7 +364,7 @@ export class AuthUtil implements IAuthProvider { private async refreshState(state = this.getAuthState()) { if (state === 'expired' || state === 'notConnected') { this.lspAuth.deleteBearerToken() - if (this.isIdcConnection() || this.isIamSession()) { + if (this.isIdcConnection()) { await this.regionProfileManager.invalidateProfile(this.regionProfileManager.activeRegionProfile?.arn) await this.regionProfileManager.clearCache() } @@ -379,7 +379,7 @@ export class AuthUtil implements IAuthProvider { await this.lspAuth.updateIamCredential(params) } - if (this.isIdcConnection() || this.isIamSession()) { + if (this.isIdcConnection()) { await this.regionProfileManager.restoreProfileSelection() } } diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts index ed00663bb26..04bb85d7a43 100644 --- a/packages/core/src/codewhisperer/util/customizationUtil.ts +++ b/packages/core/src/codewhisperer/util/customizationUtil.ts @@ -49,7 +49,7 @@ export class CustomizationProvider { } static async init(profile: RegionProfile): Promise { - const client = await AuthUtil.instance.regionProfileManager.createQUserClient(profile) + const client = await AuthUtil.instance.regionProfileManager.createQClient(profile) return new CustomizationProvider(client, profile) } } diff --git a/packages/core/src/test/amazonq/customizationUtil.test.ts b/packages/core/src/test/amazonq/customizationUtil.test.ts index 0ed19cf2197..19e59b91c03 100644 --- a/packages/core/src/test/amazonq/customizationUtil.test.ts +++ b/packages/core/src/test/amazonq/customizationUtil.test.ts @@ -43,7 +43,7 @@ describe('customizationProvider', function () { regionProfileManager: regionProfileManager, } sinon.stub(AuthUtil, 'instance').get(() => mockAuthUtil) - const createClientStub = sinon.stub(regionProfileManager, 'createQUserClient') + const createClientStub = sinon.stub(regionProfileManager, 'createQClient') const mockProfile = { name: 'foo', region: 'us-east-1', From 9dbe490ef9ee79dec5797f19bf1e43c8a2d7f0bf Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 20 Jun 2025 18:20:57 -0400 Subject: [PATCH 18/70] Add limited IAM support for inline chat --- .../src/inlineChat/provider/inlineChatProvider.ts | 15 +++++++++++++-- packages/core/src/codewhisperer/util/authUtil.ts | 4 ++-- .../src/codewhispererChat/clients/chat/v0/chat.ts | 1 - .../src/shared/clients/qDeveloperChatClient.ts | 5 ++--- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts index 64a67224a2e..64d97c372ac 100644 --- a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts +++ b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts @@ -143,7 +143,7 @@ export class InlineChatProvider { private async generateResponse( triggerPayload: TriggerPayload & { projectContextQueryLatencyMs?: number }, triggerID: string - ) { + ): Promise { const triggerEvent = this.triggerEventsStorage.getTriggerEvent(triggerID) if (triggerEvent === undefined) { return @@ -182,7 +182,18 @@ export class InlineChatProvider { let response: GenerateAssistantResponseCommandOutput | undefined = undefined session.createNewTokenSource() try { - response = await session.chatSso(request) + if (AuthUtil.instance.isSsoSession()) { + response = await session.chatSso(request) + } else { + // Call sendMessage because Q Developer Streaming Client does not have generateAssistantResponse + const { sendMessageResponse, ...rest } = await session.chatIam(request) + // Convert sendMessageCommandOutput to GenerateAssistantResponseCommandOutput + response = { + generateAssistantResponseResponse: sendMessageResponse, + conversationId: session.sessionIdentifier, + ...rest + } + } getLogger().info( `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: %O`, response.$metadata diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index daf8aebe369..7c91a06917d 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -197,8 +197,8 @@ export class AuthUtil implements IAuthProvider { } async getToken() { - if (this.isSsoSession()) { - const token = (await this.session!.getCredential()).credential + if (this.session) { + const token = (await this.session.getCredential()).credential if (typeof token !== 'string') { throw new ToolkitError('Cannot get token with IAM session') } diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index c32f67cdac5..ef0aad6ec25 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts @@ -41,7 +41,6 @@ export class ChatSession { } async chatIam(chatRequest: SendMessageRequest): Promise { const client = await createQDeveloperStreamingClient() - const response = await client.sendMessage(chatRequest) if (!response.sendMessageResponse) { throw new ToolkitError( diff --git a/packages/core/src/shared/clients/qDeveloperChatClient.ts b/packages/core/src/shared/clients/qDeveloperChatClient.ts index d9344b5b406..547591a5faf 100644 --- a/packages/core/src/shared/clients/qDeveloperChatClient.ts +++ b/packages/core/src/shared/clients/qDeveloperChatClient.ts @@ -6,13 +6,12 @@ import { QDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client' import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer' import { getUserAgent } from '../telemetry/util' import { ConfiguredRetryStrategy } from '@smithy/util-retry' +import { AuthUtil } from '../../codewhisperer' // Create a client for featureDev streaming based off of aws sdk v3 export async function createQDeveloperStreamingClient(): Promise { - throw new Error('Do not call this function until IAM is supported by LSP identity server') - const cwsprConfig = getCodewhispererConfig() - const credentials = undefined + const credentials = await AuthUtil.instance.getIamCredential() const streamingClient = new QDeveloperStreaming({ region: cwsprConfig.region, endpoint: cwsprConfig.endpoint, From f044eb6e06f4f5d6c4b5d0345f0da69b04807d30 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 20 Jun 2025 19:22:50 -0400 Subject: [PATCH 19/70] Revert CredentialChangedKind for backwards compatibility --- packages/core/src/auth/auth2.ts | 10 +++++----- packages/core/src/test/credentials/auth2.test.ts | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 11c572dcf2b..d010e599bbe 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -41,7 +41,7 @@ import { iamCredentialsDeleteNotificationType, bearerCredentialsDeleteNotificationType, bearerCredentialsUpdateRequestType, - CredentialChangedKind, + SsoTokenChangedKind, RequestType, ResponseMessage, NotificationType, @@ -473,10 +473,10 @@ export class SsoLogin extends BaseLogin { private ssoTokenChangedHandler(params: SsoTokenChangedParams) { if (params.ssoTokenId === this.ssoTokenId) { - if (params.kind === CredentialChangedKind.Expired) { + if (params.kind === SsoTokenChangedKind.Expired) { this.updateConnectionState('expired') return - } else if (params.kind === CredentialChangedKind.Refreshed) { + } else if (params.kind === SsoTokenChangedKind.Refreshed) { this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) } } @@ -603,10 +603,10 @@ export class IamLogin extends BaseLogin { // private stsCredentialChangedHandler(params: StsCredentialChangedParams) { // if (params.stsCredentialId === this.iamCredentialId) { - // if (params.kind === CredentialChangedKind.Expired) { + // if (params.kind === StsCredentialChangedKind.Expired) { // this.updateConnectionState('expired') // return - // } else if (params.kind === CredentialChangedKind.Refreshed) { + // } else if (params.kind === StsCredentialChangedKind.Refreshed) { // this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) // } // } diff --git a/packages/core/src/test/credentials/auth2.test.ts b/packages/core/src/test/credentials/auth2.test.ts index 85405ef06d5..acd2b1ccfcd 100644 --- a/packages/core/src/test/credentials/auth2.test.ts +++ b/packages/core/src/test/credentials/auth2.test.ts @@ -17,7 +17,7 @@ import { bearerCredentialsUpdateRequestType, bearerCredentialsDeleteNotificationType, ssoTokenChangedRequestType, - CredentialChangedKind, + SsoTokenChangedKind, invalidateSsoTokenRequestType, AwsErrorCodes, } from '@aws/language-server-runtimes/protocol' @@ -180,7 +180,7 @@ describe('LanguageClientAuth', () => { // Simulate a token changed notification const tokenChangedParams: SsoTokenChangedParams = { - kind: CredentialChangedKind.Refreshed, + kind: SsoTokenChangedKind.Refreshed, ssoTokenId: tokenId, } const registeredHandler = client.onNotification.firstCall.args[1] From 7506ecc462be61b89653424e7a1af3861e854805 Mon Sep 17 00:00:00 2001 From: Yuxian Zhang Date: Tue, 24 Jun 2025 23:11:18 -0400 Subject: [PATCH 20/70] fix: fix missing autofill after disabling extension --- packages/core/src/login/webview/vue/backend.ts | 6 +++--- packages/core/src/login/webview/vue/login.vue | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts index 8bb64c325d4..f17aa80acf7 100644 --- a/packages/core/src/login/webview/vue/backend.ts +++ b/packages/core/src/login/webview/vue/backend.ts @@ -299,13 +299,13 @@ export abstract class CommonAuthWebview extends VueWebview { return globals.globalState.tryGet('recentSso', Object, { startUrl: '', region: 'us-east-1' }) } - getDefaultIamKeys(): { accessKey: string; secretKey: string } { + getDefaultIamKeys(): { accessKey: string } { const devSettings = DevSettings.instance.get('autofillAccessKey', '') if (devSettings) { - return { accessKey: devSettings, secretKey: '' } + return { accessKey: devSettings } } - return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '', secretKey: '' }) + return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '' }) } cancelAuthFlow() { diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue index e6b83772fbd..ec32788a030 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -371,10 +371,10 @@ export default defineComponent({ }, async created() { const defaultSso = await this.getDefaultSso() - const defaultIamAccessKey = await this.getDefaultIamAccessKey() - this.accessKey = defaultIamAccessKey.accessKey this.startUrl = defaultSso.startUrl this.selectedRegion = defaultSso.region + const defaultIamAccessKey = await this.getDefaultIamAccessKey() + this.accessKey = defaultIamAccessKey.accessKey await this.emitUpdate('created') }, From ce07bd2c1daf321e7643150a88188e1ab0bb7f26 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 30 Jun 2025 17:10:51 -0400 Subject: [PATCH 21/70] Re-add session restore for IAM --- packages/core/src/codewhisperer/util/authUtil.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 7c91a06917d..7d55c7502c0 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -134,7 +134,7 @@ export class AuthUtil implements IAuthProvider { if (!this.isConnected()) { // Try to restore an IAM session this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) - // await this.session.restore() + await this.session.restore() if (!this.isConnected()) { // If both fail, reset the session this.session = undefined From 4d2b8b0d00952f020dbc2fe495577abc11ef329a Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 1 Jul 2025 09:52:44 -0400 Subject: [PATCH 22/70] Add session token input to login flow --- .../amazonq/test/e2e/amazonq/utils/setup.ts | 2 +- .../region/regionProfileManager.test.ts | 4 +-- .../codewhisperer/util/showSsoPrompt.test.ts | 4 +-- packages/core/src/auth/auth2.ts | 9 ++--- .../codewhisperer/ui/codeWhispererNodes.ts | 2 +- .../core/src/codewhisperer/util/authUtil.ts | 33 ++++++++++--------- .../src/codewhisperer/util/getStartUrl.ts | 2 +- .../src/codewhisperer/util/showSsoPrompt.ts | 2 +- .../webview/vue/amazonq/backend_amazonq.ts | 5 +-- .../core/src/login/webview/vue/backend.ts | 3 +- packages/core/src/login/webview/vue/login.vue | 14 +++++++- 11 files changed, 48 insertions(+), 32 deletions(-) diff --git a/packages/amazonq/test/e2e/amazonq/utils/setup.ts b/packages/amazonq/test/e2e/amazonq/utils/setup.ts index be749fc3e25..5690480b7c0 100644 --- a/packages/amazonq/test/e2e/amazonq/utils/setup.ts +++ b/packages/amazonq/test/e2e/amazonq/utils/setup.ts @@ -22,5 +22,5 @@ export async function loginToIdC() { ) } - await AuthUtil.instance.login(startUrl, region, 'sso') + await AuthUtil.instance.login_sso(startUrl, region) } diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index a858c3e659e..d0fdb6c0bb7 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -26,11 +26,11 @@ describe('RegionProfileManager', async function () { async function setupConnection(type: 'builderId' | 'idc') { if (type === 'builderId') { - await AuthUtil.instance.login(constants.builderIdStartUrl, region, 'sso') + await AuthUtil.instance.login_sso(constants.builderIdStartUrl, region) assert.ok(AuthUtil.instance.isSsoSession()) assert.ok(AuthUtil.instance.isBuilderIdConnection()) } else if (type === 'idc') { - await AuthUtil.instance.login(enterpriseSsoStartUrl, region, 'sso') + await AuthUtil.instance.login_sso(enterpriseSsoStartUrl, region) assert.ok(AuthUtil.instance.isSsoSession()) assert.ok(AuthUtil.instance.isIdcConnection()) } diff --git a/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts b/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts index 1d67db60efc..83684ac44df 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts @@ -28,7 +28,7 @@ describe('showConnectionPrompt', function () { }) it('can select connect to AwsBuilderId', async function () { - sinon.stub(AuthUtil.instance, 'login').resolves() + sinon.stub(AuthUtil.instance, 'login_sso').resolves() getTestWindow().onDidShowQuickPick(async (picker) => { await picker.untilReady() @@ -44,7 +44,7 @@ describe('showConnectionPrompt', function () { it('connectToAwsBuilderId calls AuthUtil login with builderIdStartUrl', async function () { sinon.stub(vscode.commands, 'executeCommand') - const loginStub = sinon.stub(AuthUtil.instance, 'login').resolves() + const loginStub = sinon.stub(AuthUtil.instance, 'login_sso').resolves() await awsIdSignIn() diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index d010e599bbe..fc1a8db2182 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -174,7 +174,7 @@ export class LanguageClientAuth { } satisfies UpdateProfileParams) } - updateIamProfile(profileName: string, accessKey: string, secretKey: string): Promise { + updateIamProfile(profileName: string, accessKey: string, secretKey: string, sessionToken?: string): Promise { // Add credentials and delete SSO settings from profile return this.client.sendRequest(updateProfileRequestType.method, { profile: { @@ -185,6 +185,7 @@ export class LanguageClientAuth { sso_session: '', aws_access_key_id: accessKey, aws_secret_access_key: secretKey, + aws_session_token: sessionToken, }, }, ssoSession: { @@ -497,7 +498,7 @@ export class IamLogin extends BaseLogin { // ) } - async login(opts: { accessKey: string; secretKey: string }) { + async login(opts: { accessKey: string; secretKey: string, sessionToken?: string }) { await this.updateProfile(opts) return this._getIamCredential(true) } @@ -519,8 +520,8 @@ export class IamLogin extends BaseLogin { // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) } - async updateProfile(opts: { accessKey: string; secretKey: string }) { - await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey) + async updateProfile(opts: { accessKey: string; secretKey: string, sessionToken?: string }) { + await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken) this._data = { accessKey: opts.accessKey, secretKey: opts.secretKey, diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts index 207add2f452..0fbc08542d5 100644 --- a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts +++ b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts @@ -271,7 +271,7 @@ export function createSignIn(): DataQuickPickItem<'signIn'> { if (isWeb()) { // TODO: nkomonen, call a Command instead onClick = () => { - void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso') + void AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion) } } diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 7d55c7502c0..52853110dff 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -160,25 +160,26 @@ export class AuthUtil implements IAuthProvider { } } - // Log into the desired session type using the authentication parameters - async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise - async login(startUrl: string, region: string, loginType: 'sso'): Promise - async login( - first: string, - second: string, - loginType: 'iam' | 'sso' - ): Promise { - let response: GetSsoTokenResult | GetIamCredentialResult | undefined - - // Start session if the current session type does not match the desired type - if (loginType === 'sso' && !this.isSsoSession()) { + // Log in using SSO + async login_sso(startUrl: string, region: string): Promise { + let response: GetSsoTokenResult | undefined + // Create SSO login session + if (!this.isSsoSession()) { this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter) - response = await this.session.login({ startUrl: first, region: second, scopes: amazonQScopes }) - } else if (loginType === 'iam' && !this.isIamSession()) { - this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) - response = await this.session.login({ accessKey: first, secretKey: second }) } + response = await (this.session as SsoLogin).login({ startUrl: startUrl, region: region, scopes: amazonQScopes }) + await showAmazonQWalkthroughOnce() + return response + } + // Log in using IAM or STS credentials + async login_iam(accessKey: string, secretKey: string, sessionToken?: string): Promise { + let response: GetIamCredentialResult | undefined + // Create IAM login session + if (!this.isIamSession()) { + this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) + } + response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken }) await showAmazonQWalkthroughOnce() return response } diff --git a/packages/core/src/codewhisperer/util/getStartUrl.ts b/packages/core/src/codewhisperer/util/getStartUrl.ts index 40da222bfb9..76b58db14bb 100644 --- a/packages/core/src/codewhisperer/util/getStartUrl.ts +++ b/packages/core/src/codewhisperer/util/getStartUrl.ts @@ -29,7 +29,7 @@ export const getStartUrl = async () => { export async function connectToEnterpriseSso(startUrl: string, region: Region['id']) { try { - await AuthUtil.instance.login(startUrl, region, 'sso') + await AuthUtil.instance.login_sso(startUrl, region) } catch (e) { throw ToolkitError.chain(e, CodeWhispererConstants.failedToConnectIamIdentityCenter, { code: 'FailedToConnect', diff --git a/packages/core/src/codewhisperer/util/showSsoPrompt.ts b/packages/core/src/codewhisperer/util/showSsoPrompt.ts index 7b1116cf370..f3c0900bc66 100644 --- a/packages/core/src/codewhisperer/util/showSsoPrompt.ts +++ b/packages/core/src/codewhisperer/util/showSsoPrompt.ts @@ -47,7 +47,7 @@ export const showCodeWhispererConnectionPrompt = async () => { export async function awsIdSignIn() { getLogger().info('selected AWS ID sign in') try { - await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso') + await AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion) } catch (e) { throw ToolkitError.chain(e, failedToConnectAwsBuilderId, { code: 'FailedToConnect' }) } diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts index af8a06cd2a7..ba1b5abec4d 100644 --- a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts +++ b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts @@ -196,7 +196,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview { async startIamCredentialSetup( profileName: string, accessKey: string, - secretKey: string + secretKey: string, + sessionToken?: string ): Promise { getLogger().debug(`called startIamCredentialSetup()`) // Defining separate auth function to emit telemetry before returning from this method @@ -205,7 +206,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview { }) const runAuth = async (): Promise => { try { - await AuthUtil.instance.login(accessKey, secretKey, 'iam') + await AuthUtil.instance.login_iam(accessKey, secretKey, sessionToken) } catch (e) { getLogger().error('Failed submitting credentials %O', e) return { id: this.id, text: e as string } diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts index f17aa80acf7..153004f374a 100644 --- a/packages/core/src/login/webview/vue/backend.ts +++ b/packages/core/src/login/webview/vue/backend.ts @@ -174,7 +174,8 @@ export abstract class CommonAuthWebview extends VueWebview { abstract startIamCredentialSetup( profileName: string, accessKey: string, - secretKey: string + secretKey: string, + sessionToken?: string, ): Promise async showResourceExplorer(): Promise { diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue index ec32788a030..2254f6032c1 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -279,6 +279,17 @@ v-model="secretKey" @keydown.enter="handleContinueClick()" /> +
+
Session Token (Optional)
+ +
@@ -367,6 +378,7 @@ export default defineComponent({ profileName: '', accessKey: '', secretKey: '', + sessionToken: '', } }, async created() { @@ -505,7 +517,7 @@ export default defineComponent({ return } this.stage = 'AUTHENTICATING' - const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey) + const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey, this.sessionToken) if (error) { this.stage = 'START' void client.errorNotification(error) From 58307167d52987c383cb27ced9b6e27b04eb2ca4 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 1 Jul 2025 12:00:36 -0400 Subject: [PATCH 23/70] Add session tokens to auth2 logic --- packages/core/src/auth/auth2.ts | 4 +++- packages/core/src/codewhisperer/util/authUtil.ts | 5 +++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index fc1a8db2182..6c1cd40680d 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -273,7 +273,7 @@ export class LanguageClientAuth { export abstract class BaseLogin { protected connectionState: AuthState = 'notConnected' protected cancellationToken: CancellationTokenSource | undefined - protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined + protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string } | undefined constructor( public readonly profileName: string, @@ -525,6 +525,7 @@ export class IamLogin extends BaseLogin { this._data = { accessKey: opts.accessKey, secretKey: opts.secretKey, + sessionToken: opts.sessionToken } } @@ -542,6 +543,7 @@ export class IamLogin extends BaseLogin { this._data = { accessKey: credentials.aws_access_key_id, secretKey: credentials.aws_secret_access_key, + sessionToken: credentials.aws_session_token } } try { diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 52853110dff..29b166554e8 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -64,7 +64,7 @@ export interface IAuthProvider { getToken(): Promise getIamCredential(): Promise readonly profileName: string - readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } + readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string } } /** @@ -132,12 +132,13 @@ export class AuthUtil implements IAuthProvider { this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter) await this.session.restore() if (!this.isConnected()) { + this.session?.logout() // Try to restore an IAM session this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter) await this.session.restore() if (!this.isConnected()) { // If both fail, reset the session - this.session = undefined + this.session?.logout() } } } From ef3745bdeb17dedb2ac23921317ead52d00f3382 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 1 Jul 2025 16:17:14 -0400 Subject: [PATCH 24/70] Replace profile deletion with iam invalidation --- packages/core/src/auth/auth2.ts | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 6c1cd40680d..b65e3715e21 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -21,8 +21,6 @@ import { ProfileKind, UpdateProfileParams, updateProfileRequestType, - DeleteProfileParams, - deleteProfileRequestType, SsoTokenChangedParams, // StsCredentialChangedParams, ssoTokenChangedRequestType, @@ -50,7 +48,6 @@ import { iamCredentialsUpdateRequestType, Profile, SsoSession, - DeleteProfileResult, // invalidateStsCredentialRequestType, // InvalidateStsCredentialParams, // InvalidateStsCredentialResult, @@ -175,10 +172,12 @@ export class LanguageClientAuth { } updateIamProfile(profileName: string, accessKey: string, secretKey: string, sessionToken?: string): Promise { + // Use unknown profile type if invalidating all IAM fields + const kind = !accessKey && !secretKey && !sessionToken ? ProfileKind.EmptyProfile : ProfileKind.IamCredentialProfile // Add credentials and delete SSO settings from profile return this.client.sendRequest(updateProfileRequestType.method, { profile: { - kinds: [ProfileKind.IamCredentialProfile], + kinds: [kind], name: profileName, settings: { region: '', @@ -195,12 +194,6 @@ export class LanguageClientAuth { } satisfies UpdateProfileParams) } - deleteIamProfile(name: string): Promise { - return this.client.sendRequest(deleteProfileRequestType.method, { - profileName: name, - } satisfies DeleteProfileParams) - } - listProfiles() { return this.client.sendRequest(listProfilesRequestType.method, {}) as Promise } @@ -514,7 +507,7 @@ export class IamLogin extends BaseLogin { if (this.iamCredentialId) { await this.lspAuth.invalidateIamCredential(this.iamCredentialId) } - await this.deleteProfile(this.profileName) + await this.lspAuth.updateIamProfile(this.profileName, '', '', '') this.updateConnectionState('notConnected') this._data = undefined // TODO: DeleteProfile api in Identity Service (this doesn't exist yet) @@ -529,10 +522,6 @@ export class IamLogin extends BaseLogin { } } - async deleteProfile(profileName: string) { - await this.lspAuth.deleteIamProfile(profileName) - } - /** * Restore the connection state and connection details to memory, if they exist. */ From 1ec8508b5e587dbf6fd365a7d5e1aee6ee8433f8 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 1 Jul 2025 17:56:29 -0400 Subject: [PATCH 25/70] Rename loginOnInvalidToken --- packages/core/src/auth/auth2.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index b65e3715e21..cb1e3abf528 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -135,7 +135,7 @@ export class LanguageClientAuth { { profileName: profileName, options: { - loginOnInvalidToken: login, + loginOnInvalidStsCredential: login, }, } satisfies GetIamCredentialParams, cancellationToken From 0382876a1ef907315920d8e56a68adb35cef6c04 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 2 Jul 2025 09:52:37 -0400 Subject: [PATCH 26/70] Rename getIamCredential options --- packages/core/src/auth/auth2.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index cb1e3abf528..ccc79a6110b 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -135,7 +135,8 @@ export class LanguageClientAuth { { profileName: profileName, options: { - loginOnInvalidStsCredential: login, + assumeRole: false, + generateOnInvalidStsCredential: login, }, } satisfies GetIamCredentialParams, cancellationToken From f45c92a0d66339f5563dc62186dfdedb20de9fd1 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 2 Jul 2025 10:37:40 -0400 Subject: [PATCH 27/70] Remove assumeRole option --- packages/core/src/auth/auth2.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index ccc79a6110b..c16ebd21200 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -135,7 +135,6 @@ export class LanguageClientAuth { { profileName: profileName, options: { - assumeRole: false, generateOnInvalidStsCredential: login, }, } satisfies GetIamCredentialParams, From b730331846f64aaa6a4e59c804caea03de3ea8ed Mon Sep 17 00:00:00 2001 From: Yuxian Zhang Date: Wed, 2 Jul 2025 17:10:08 -0400 Subject: [PATCH 28/70] feat: enable sts invalidation --- packages/core/src/auth/auth2.ts | 44 +++++++++---------- packages/core/src/login/webview/vue/login.vue | 2 +- 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index c16ebd21200..f941f82557d 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -12,12 +12,12 @@ import { GetIamCredentialParams, getIamCredentialRequestType, GetIamCredentialResult, - InvalidateIamCredentialResult, + InvalidateStsCredentialResult, IamIdentityCenterSsoTokenSource, InvalidateSsoTokenParams, - InvalidateIamCredentialParams, + InvalidateStsCredentialParams, invalidateSsoTokenRequestType, - invalidateIamCredentialRequestType, + invalidateStsCredentialRequestType, ProfileKind, UpdateProfileParams, updateProfileRequestType, @@ -234,18 +234,18 @@ export class LanguageClientAuth { } satisfies InvalidateSsoTokenParams) as Promise } - invalidateIamCredential(tokenId: string) { - return this.client.sendRequest(invalidateIamCredentialRequestType.method, { - iamCredentialsId: tokenId, - } satisfies InvalidateIamCredentialParams) as Promise - } - // invalidateStsCredential(tokenId: string) { // return this.client.sendRequest(invalidateStsCredentialRequestType.method, { - // stsCredentialId: tokenId, + // iamCredentialsId: tokenId, // } satisfies InvalidateStsCredentialParams) as Promise // } + invalidateStsCredential(tokenId: string) { + return this.client.sendRequest(invalidateStsCredentialRequestType.method, { + profileName: tokenId, + } satisfies InvalidateStsCredentialParams) as Promise + } + registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) { this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler) } @@ -505,7 +505,7 @@ export class IamLogin extends BaseLogin { async logout() { if (this.iamCredentialId) { - await this.lspAuth.invalidateIamCredential(this.iamCredentialId) + await this.lspAuth.invalidateStsCredential(this.iamCredentialId) } await this.lspAuth.updateIamProfile(this.profileName, '', '', '') this.updateConnectionState('notConnected') @@ -588,19 +588,19 @@ export class IamLogin extends BaseLogin { this.cancellationToken = undefined } - // this.iamCredentialId = response.id + this.iamCredentialId = response.id this.updateConnectionState('connected') return response } - // private stsCredentialChangedHandler(params: StsCredentialChangedParams) { - // if (params.stsCredentialId === this.iamCredentialId) { - // if (params.kind === StsCredentialChangedKind.Expired) { - // this.updateConnectionState('expired') - // return - // } else if (params.kind === StsCredentialChangedKind.Refreshed) { - // this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) - // } - // } - // } +// private stsCredentialChangedHandler(params: StsCredentialChangedParams) { +// if (params.stsCredentialId === this.iamCredentialId) { +// if (params.kind === StsCredentialChangedKind.Expired) { +// this.updateConnectionState('expired') +// return +// } else if (params.kind === StsCredentialChangedKind.Refreshed) { +// this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' }) +// } +// } +// } } diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue index 2254f6032c1..b6c403da670 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -230,7 +230,7 @@