diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts index 576757c36e2..d7f22c2996d 100644 --- a/packages/amazonq/src/extensionNode.ts +++ b/packages/amazonq/src/extensionNode.ts @@ -100,8 +100,11 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) { async function getAuthState(): Promise> { const state = AuthUtil.instance.getAuthState() - if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) { - getLogger().error('Current Amazon Q connection is not SSO') + if ( + AuthUtil.instance.isConnected() && + !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker()) + ) { + getLogger().error('Current Amazon Q connection is not SSO nor IAM') } return { diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts index 64a67224a2e..6c22cb06b84 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,12 @@ 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 + throw new ToolkitError('Inline chat is only available with SSO authentication') + } getLogger().info( `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: %O`, response.$metadata diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index d835427006c..233a7371425 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -483,7 +483,7 @@ describe('AuthUtil', async function () { await auth.getIamCredential() assert.fail('Should have thrown an error') } catch (err) { - assert.strictEqual((err as Error).message, 'Cannot get token with SSO session') + assert.strictEqual((err as Error).message, 'Cannot get credential without logging in with IAM.') } }) @@ -494,7 +494,7 @@ describe('AuthUtil', async function () { await auth.getIamCredential() assert.fail('Should have thrown an error') } catch (err) { - assert.strictEqual((err as Error).message, 'Cannot get credential without logging in.') + assert.strictEqual((err as Error).message, 'Cannot get credential without logging in with IAM.') } }) }) diff --git a/packages/core/src/auth/auth2.ts b/packages/core/src/auth/auth2.ts index 6df5409fceb..e69c6730a9a 100644 --- a/packages/core/src/auth/auth2.ts +++ b/packages/core/src/auth/auth2.ts @@ -344,8 +344,13 @@ export abstract class BaseLogin { * 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('"', '') + try { + const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey) + return decrypted.plaintext.toString().replaceAll('"', '') + } catch (e) { + getLogger().error(`Failed to decrypt: ${encrypted}`) + return encrypted + } } } diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts index 1d0dc8c51f0..70facac1ee0 100644 --- a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts +++ b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts @@ -192,7 +192,11 @@ export function createManageSubscription(): DataQuickPickItem<'manageSubscriptio export function createSignout(): DataQuickPickItem<'signout'> { const label = localize('AWS.codewhisperer.signoutNode.label', 'Sign Out') const icon = getIcon('vscode-export') - const connection = AuthUtil.instance.isBuilderIdConnection() ? 'AWS Builder ID' : 'IAM Identity Center' + const connection = AuthUtil.instance.isIamSession() + ? 'IAM Credentials' + : AuthUtil.instance.isBuilderIdConnection() + ? 'AWS Builder ID' + : 'IAM Identity Center' return { data: 'signout', diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 8745d6739aa..a4a2eff051e 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -30,7 +30,16 @@ 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, Login, SsoLogin, IamLogin } from '../../auth/auth2' +import { + AuthStateEvent, + cacheChangedEvent, + LanguageClientAuth, + Login, + SsoLogin, + IamLogin, + AuthState, + LoginTypes, +} from '../../auth/auth2' import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants' import { VSCODE_EXTENSION_ID } from '../../shared/extensions' import { RegionProfileManager } from '../region/regionProfileManager' @@ -108,11 +117,11 @@ export class AuthUtil implements IAuthProvider { } isSsoSession(): boolean { - return this.session instanceof SsoLogin + return this.session?.loginType === LoginTypes.SSO } isIamSession(): boolean { - return this.session instanceof IamLogin + return this.session?.loginType === LoginTypes.IAM } /** @@ -204,26 +213,20 @@ export class AuthUtil implements IAuthProvider { } async getToken() { - if (this.session) { - const token = (await this.session.getCredential()).credential - if (typeof token !== 'string') { - throw new ToolkitError('Cannot get token with IAM session') - } - return token + if (this.isSsoSession()) { + const response = await this.session!.getCredential() + return response.credential as string } else { - throw new ToolkitError('Cannot get credential without logging in.') + throw new ToolkitError('Cannot get credential without logging in with SSO.') } } async getIamCredential() { - if (this.session) { - const credential = (await this.session.getCredential()).credential - if (typeof credential !== 'object') { - throw new ToolkitError('Cannot get token with SSO session') - } - return credential + if (this.isIamSession()) { + const response = await this.session!.getCredential() + return response.credential as IamCredentials } else { - throw new ToolkitError('Cannot get credential without logging in.') + throw new ToolkitError('Cannot get credential without logging in with IAM.') } } @@ -231,9 +234,10 @@ export class AuthUtil implements IAuthProvider { return this.session?.data } - getAuthState() { - if (this.session) { - return this.session.getConnectionState() + getAuthState(): AuthState { + // Check if getConnectionState exists in case of type casts + if (typeof this.session?.getConnectionState === 'function') { + return this.session!.getConnectionState() } else { return 'notConnected' } @@ -356,11 +360,12 @@ export class AuthUtil implements IAuthProvider { private async stateChangeHandler(e: AuthStateEvent) { if (e.state === 'refreshed') { - const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined if (this.isSsoSession()) { - await this.lspAuth.updateBearerToken(params) + const params = await this.session!.getCredential() + await this.lspAuth.updateBearerToken(params.updateCredentialsParams) } else if (this.isIamSession()) { - await this.lspAuth.updateIamCredential(params) + const params = await this.session!.getCredential() + await this.lspAuth.updateIamCredential(params.updateCredentialsParams) } } else { this.logger.info(`codewhisperer: connection changed to ${e.state}`) @@ -383,11 +388,12 @@ export class AuthUtil implements IAuthProvider { this.session = undefined } if (state === 'connected') { - const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined if (this.isSsoSession()) { - await this.lspAuth.updateBearerToken(params) + const params = await this.session!.getCredential() + await this.lspAuth.updateBearerToken(params.updateCredentialsParams) } else if (this.isIamSession()) { - await this.lspAuth.updateIamCredential(params) + const params = await this.session!.getCredential() + await this.lspAuth.updateIamCredential(params.updateCredentialsParams) } if (this.isIdcConnection()) { diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts index edb1980a8c0..b455c205ec8 100644 --- a/packages/core/src/login/webview/vue/backend.ts +++ b/packages/core/src/login/webview/vue/backend.ts @@ -33,6 +33,7 @@ import { getLogger } from '../../../shared/logger/logger' import { isValidUrl } from '../../../shared/utilities/uriUtils' import { RegionProfile } from '../../../codewhisperer/models/model' import { ProfileSwitchIntent } from '../../../codewhisperer/region/regionProfileManager' +import { showMessage } from '../../../shared/utilities/messages' export abstract class CommonAuthWebview extends VueWebview { private readonly className = 'CommonAuthWebview' @@ -183,7 +184,7 @@ export abstract class CommonAuthWebview extends VueWebview { abstract fetchConnections(): Promise async errorNotification(e: AuthError) { - void vscode.window.showInformationMessage(`${e.text}`) + await showMessage('error', e.text) } abstract quitLoginScreen(): Promise @@ -296,6 +297,15 @@ export abstract class CommonAuthWebview extends VueWebview { return globals.globalState.tryGet('recentSso', Object, { startUrl: '', region: 'us-east-1' }) } + getDefaultIamKeys(): { accessKey: string } { + const devSettings = DevSettings.instance.get('autofillAccessKey', '') + if (devSettings) { + return { accessKey: devSettings } + } + + return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '' }) + } + 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 c61e7c1dabd..90eabc09db4 100644 --- a/packages/core/src/login/webview/vue/login.vue +++ b/packages/core/src/login/webview/vue/login.vue @@ -230,7 +230,7 @@