fix(auth): fix UI and error message bug, disable inline chat for IAM, autofill access key id

packages/amazonq/src/extensionNode.ts

@@ -100,8 +100,11 @@ 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 {

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

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

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',

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 {
@@ -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
     }

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}`)
+        await 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()
     }

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

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
         }
 

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.

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,
 }

packages/core/src/test/amazonq/utils.ts

@@ -73,22 +73,22 @@ 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) {
     session.config.fs.registerProvider(uri, new VirtualMemoryFile(fileContents))
 }
 
 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) {