Fix type checks for AuthUtil sessions

Author: liramon1

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.')
             }
         })
     })

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'

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,36 +213,31 @@ 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.')
         }
     }
 
     get connection() {
         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()) {