From c0c47641a542ef24cde765488174de93d740ddf2 Mon Sep 17 00:00:00 2001
From: Yuxian Zhang <yuxianrz@amazon.com>
Date: Fri, 1 Aug 2025 14:08:22 -0400
Subject: [PATCH 1/5] disable inline chat, fix other bugs, autofill access key

---
 packages/amazonq/src/extensionNode.ts         |  4 +-
 .../inlineChat/provider/inlineChatProvider.ts |  9 ++-
 .../unit/codewhisperer/util/authUtil.test.ts  |  2 +-
 .../codewhisperer/ui/codeWhispererNodes.ts    |  6 +-
 .../core/src/codewhisperer/util/authUtil.ts   |  2 +-
 .../core/src/login/webview/vue/backend.ts     | 12 +++-
 packages/core/src/login/webview/vue/login.vue | 55 ++++++++++++++-----
 packages/core/src/shared/featureConfig.ts     |  2 +-
 packages/core/src/shared/globalState.ts       |  1 +
 packages/core/src/shared/settings.ts          |  1 +
 packages/core/src/test/amazonq/utils.ts       | 14 ++---
 11 files changed, 79 insertions(+), 29 deletions(-)

diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts
index 576757c36e2..dbc669909c8 100644
--- a/packages/amazonq/src/extensionNode.ts
+++ b/packages/amazonq/src/extensionNode.ts
@@ -100,8 +100,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
     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<GenerateAssistantResponseCommandOutput | undefined> {
         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..62ca4002f43 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 with SSO session')
             }
         })
 
diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
index 1d0dc8c51f0..ed50a8ba09c 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.isIamConnection()
+        ? '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..68d01e51e15 100644
--- a/packages/core/src/codewhisperer/util/authUtil.ts
+++ b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -219,7 +219,7 @@ export class AuthUtil implements IAuthProvider {
         if (this.session) {
             const credential = (await this.session.getCredential()).credential
             if (typeof credential !== 'object') {
-                throw new ToolkitError('Cannot get token with SSO session')
+                throw new ToolkitError('Cannot get credential with SSO session')
             }
             return credential
         } else {
diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts
index edb1980a8c0..bbc33ebb544 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<AwsConnection[] | undefined>
 
     async errorNotification(e: AuthError) {
-        void vscode.window.showInformationMessage(`${e.text}`)
+        showMessage('error', e.text)
     }
 
     abstract quitLoginScreen(): Promise<void>
@@ -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 @@
 
         <template v-if="stage === 'AUTHENTICATING'">
             <div class="auth-container-section">
-                <div v-if="app === 'TOOLKIT' && profileName.length > 0" class="header bottomMargin">
+                <div v-if="selectedLoginOption === LoginOption.IAM_CREDENTIAL" class="header bottomMargin">
                     Connecting to IAM...
                 </div>
                 <div v-else class="header bottomMargin">Authenticating in browser...</div>
@@ -273,7 +273,7 @@
             <div class="title">Secret Access Key</div>
             <input
                 class="iamInput bottomMargin"
-                type="text"
+                type="password"
                 id="secretKey"
                 name="secretKey"
                 v-model="secretKey"
@@ -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,
@@ -368,6 +373,8 @@ export default defineComponent({
         const defaultSso = await this.getDefaultSso()
         this.startUrl = defaultSso.startUrl
         this.selectedRegion = defaultSso.region
+        const defaultIamAccessKey = await this.getDefaultIamAccessKey()
+        this.accessKey = defaultIamAccessKey.accessKey
         await this.emitUpdate('created')
     },
 
@@ -397,6 +404,10 @@ export default defineComponent({
             }
         },
         handleDocumentClick(event: any) {
+            // Only reset selection when in START stage to avoid clearing during authentication
+            if (this.stage !== 'START') {
+                return
+            }
             const isClickInsideSelectableItems = event.target.closest('.selectable-item')
             if (!isClickInsideSelectableItems) {
                 this.selectedLoginOption = 0
@@ -437,17 +448,32 @@ 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'
+                    // 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 +607,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/shared/featureConfig.ts b/packages/core/src/shared/featureConfig.ts
index 7cc6a9cbfc7..42ffb15ad66 100644
--- a/packages/core/src/shared/featureConfig.ts
+++ b/packages/core/src/shared/featureConfig.ts
@@ -118,7 +118,7 @@ export class FeatureConfigProvider {
     }
 
     async fetchFeatureConfigs(): Promise<void> {
-        if (AuthUtil.instance.isConnectionExpired()) {
+        if (AuthUtil.instance.isConnectionExpired() || AuthUtil.instance.isIamSession()) {
             return
         }
 
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,
 }
diff --git a/packages/core/src/test/amazonq/utils.ts b/packages/core/src/test/amazonq/utils.ts
index ec2e7020e4e..f7fc5640e2f 100644
--- a/packages/core/src/test/amazonq/utils.ts
+++ b/packages/core/src/test/amazonq/utils.ts
@@ -73,12 +73,12 @@ export async function createSession({
 
     const client = sinon.createStubInstance(FeatureDevClient)
     client.createConversation.resolves(conversationID)
-    const session = new Session(sessionConfig, messenger, tabID, sessionState, client)
+    const sessionNew = new Session(sessionConfig, messenger, tabID, sessionState, client)
 
-    sinon.stub(session, 'conversationId').get(() => conversationID)
-    sinon.stub(session, 'uploadId').get(() => uploadID)
+    sinon.stub(sessionNew, 'conversationId').get(() => conversationID)
+    sinon.stub(sessionNew, 'uploadId').get(() => uploadID)
 
-    return session
+    return sessionNew
 }
 
 export async function sessionRegisterProvider(session: Session, uri: vscode.Uri, fileContents: Uint8Array) {
@@ -86,9 +86,9 @@ export async function sessionRegisterProvider(session: Session, uri: vscode.Uri,
 }
 
 export function generateVirtualMemoryUri(uploadID: string, filePath: string, scheme: string) {
-    const generationFilePath = path.join(uploadID, filePath)
-    const uri = vscode.Uri.from({ scheme, path: generationFilePath })
-    return uri
+    const generationFilePathNew = path.join(uploadID, filePath)
+    const uriNew = vscode.Uri.from({ scheme, path: generationFilePathNew })
+    return uriNew
 }
 
 export async function sessionWriteFile(session: Session, uri: vscode.Uri, encodedContent: Uint8Array) {

From 69915c167117d07223e549a8c419edb7e5c40281 Mon Sep 17 00:00:00 2001
From: Yuxian Zhang <yuxianrz@amazon.com>
Date: Fri, 1 Aug 2025 14:20:27 -0400
Subject: [PATCH 2/5] fix: add codeWhispererNodes error

---
 packages/core/src/codewhisperer/util/authUtil.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
index 68d01e51e15..fae8552c388 100644
--- a/packages/core/src/codewhisperer/util/authUtil.ts
+++ b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -255,6 +255,10 @@ export class AuthUtil implements IAuthProvider {
         return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl)
     }
 
+    isIamConnection() {
+        return Boolean(this.connection?.accessKey && this.connection?.secretKey)
+    }
+
     isInternalAmazonUser(): boolean {
         return this.isConnected() && this.connection?.startUrl === internalStartUrl
     }

From d190d0187fbe58f1cfcacf7dec118f266f2ccb78 Mon Sep 17 00:00:00 2001
From: Yuxian Zhang <yuxianrz@amazon.com>
Date: Fri, 1 Aug 2025 14:34:33 -0400
Subject: [PATCH 3/5] fix: lint

---
 packages/amazonq/src/extensionNode.ts          | 5 ++++-
 packages/core/src/login/webview/vue/backend.ts | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts
index dbc669909c8..d7f22c2996d 100644
--- a/packages/amazonq/src/extensionNode.ts
+++ b/packages/amazonq/src/extensionNode.ts
@@ -100,7 +100,10 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
     const state = AuthUtil.instance.getAuthState()
 
-    if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())) {
+    if (
+        AuthUtil.instance.isConnected() &&
+        !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())
+    ) {
         getLogger().error('Current Amazon Q connection is not SSO nor IAM')
     }
 
diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts
index bbc33ebb544..b455c205ec8 100644
--- a/packages/core/src/login/webview/vue/backend.ts
+++ b/packages/core/src/login/webview/vue/backend.ts
@@ -184,7 +184,7 @@ export abstract class CommonAuthWebview extends VueWebview {
     abstract fetchConnections(): Promise<AwsConnection[] | undefined>
 
     async errorNotification(e: AuthError) {
-        showMessage('error', e.text)
+        await showMessage('error', e.text)
     }
 
     abstract quitLoginScreen(): Promise<void>

From 6c5492877ed2e5cd1177cbe7c8cf570a9c348ac8 Mon Sep 17 00:00:00 2001
From: Ramon Li <liramon@amazon.com>
Date: Fri, 1 Aug 2025 15:38:10 -0400
Subject: [PATCH 4/5] fix: unit tests

---
 packages/core/src/auth/auth2.ts        |  9 +++++++--
 packages/core/src/test/testAuthUtil.ts | 17 ++++++++++++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

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<string> {
-        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/test/testAuthUtil.ts b/packages/core/src/test/testAuthUtil.ts
index 661b2c190c5..929eca99b65 100644
--- a/packages/core/src/test/testAuthUtil.ts
+++ b/packages/core/src/test/testAuthUtil.ts
@@ -26,15 +26,22 @@ export async function createTestAuthUtil() {
         },
     }
 
+    const fakeCredential = {
+        credentials: {
+            accessKeyId: 'fake-access-key-id',
+            secretAccessKey: 'fake-secret-access-key',
+            sessionToken: 'fake-session-token',
+        },
+        updateCredentialsParams: {
+            data: 'fake-data',
+        },
+    }
+
     const mockLspAuth: Partial<LanguageClientAuth> = {
         registerSsoTokenChangedHandler: sinon.stub().resolves(),
         updateSsoProfile: sinon.stub().resolves(),
         getSsoToken: sinon.stub().resolves(fakeToken),
-        getIamCredential: sinon.stub().resolves({
-            accessKeyId: 'fake-access-key-id',
-            secretAccessKey: 'fake-secret-access-key',
-            sessionToken: 'fake-session-token',
-        }),
+        getIamCredential: sinon.stub().resolves(fakeCredential),
         getProfile: sinon.stub().resolves({
             sso_registration_scopes: ['codewhisperer'],
         }),

From bcf7d39e8c1f26230ec7d0eeb5b81369b3c4d173 Mon Sep 17 00:00:00 2001
From: Ramon Li <liramon@amazon.com>
Date: Mon, 4 Aug 2025 11:24:20 -0400
Subject: [PATCH 5/5] Fix type checks for AuthUtil sessions

---
 .../unit/codewhisperer/util/authUtil.test.ts  |  4 +-
 .../codewhisperer/ui/codeWhispererNodes.ts    |  2 +-
 .../core/src/codewhisperer/util/authUtil.ts   | 62 ++++++++++---------
 3 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts
index 62ca4002f43..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 credential 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/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
index ed50a8ba09c..70facac1ee0 100644
--- a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
+++ b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
@@ -192,7 +192,7 @@ 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.isIamConnection()
+    const connection = AuthUtil.instance.isIamSession()
         ? 'IAM Credentials'
         : AuthUtil.instance.isBuilderIdConnection()
           ? 'AWS Builder ID'
diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
index fae8552c388..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 credential 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'
         }
@@ -255,10 +259,6 @@ export class AuthUtil implements IAuthProvider {
         return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl)
     }
 
-    isIamConnection() {
-        return Boolean(this.connection?.accessKey && this.connection?.secretKey)
-    }
-
     isInternalAmazonUser(): boolean {
         return this.isConnected() && this.connection?.startUrl === internalStartUrl
     }
@@ -360,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}`)
@@ -387,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()) {