feat: add mfa serial input, telemetry for iam and sync with small fix in flare-iam

yuxianrz · yuxianrz · commit 55b575d0ecc7 · 2025-08-01T12:23:55.000-04:00
diff --git a/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts b/packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts
@@ -16,6 +16,7 @@ import {
     ChatTriggerType,
     EditorContextExtractor,
     PromptMessage,
+    TriggerEvent,
     TriggerEventsStorage,
     TriggerPayload,
     triggerPayloadToChatRequest,
@@ -30,6 +31,7 @@ import { extractAuthFollowUp } from 'aws-core-vscode/amazonq'
 import { InlineChatParams, InlineChatResult } from '@aws/language-server-runtimes-types'
 import { decryptResponse, encryptRequest } from '../../lsp/encryption'
 import { getCursorState } from '../../lsp/utils'
+import { CwsprChatTriggerInteraction, telemetry } from 'aws-core-vscode/telemetry'
 
 export class InlineChatProvider {
     private readonly editorContextExtractor: EditorContextExtractor
@@ -68,7 +70,34 @@ export class InlineChatProvider {
         }
     }
 
+    private getTriggerInteractionFromTriggerEvent(triggerEvent: TriggerEvent | undefined): CwsprChatTriggerInteraction {
+        switch (triggerEvent?.type) {
+            case 'editor_context_command':
+                return triggerEvent.command?.triggerType === 'keybinding' ? 'hotkeys' : 'contextMenu'
+            case 'follow_up':
+            case 'chat_message':
+            default:
+                return 'click'
+        }
+    }
+
     public async processPromptMessageLSP(message: PromptMessage): Promise<InlineChatResult> {
+        const triggerInteraction = this.getTriggerInteractionFromTriggerEvent(
+            this.triggerEventsStorage.getLastTriggerEventByTabID(message.tabID)
+        )
+        if (!AuthUtil.instance.isSsoSession()) {
+            telemetry.amazonq_messageResponseError.emit({
+                result: 'Failed',
+                cwsprChatConversationType: 'Chat',
+                cwsprChatRequestLength: message.message?.length ?? 0,
+                cwsprChatResponseCode: 401,
+                cwsprChatTriggerInteraction: triggerInteraction,
+                reason: 'AuthenticationError',
+                reasonDesc: 'Inline chat requires SSO authentication, but current session is not',
+            })
+            throw new ToolkitError('Inline chat is only available with SSO authentication')
+        }
+
         // TODO: handle partial responses.
         getLogger().info('Making inline chat request with message %O', message)
         const params = this.getCurrentEditorParams(message.message ?? '')
@@ -83,6 +112,23 @@ export class InlineChatProvider {
 
     // TODO: remove in favor of LSP implementation.
     public async processPromptMessage(message: PromptMessage) {
+        const triggerInteraction = this.getTriggerInteractionFromTriggerEvent(
+            this.triggerEventsStorage.getLastTriggerEventByTabID(message.tabID)
+        )
+        if (!AuthUtil.instance.isSsoSession()) {
+            telemetry.amazonq_messageResponseError.emit({
+                result: 'Failed',
+                cwsprChatConversationType: 'Chat',
+                cwsprChatRequestLength: message.message?.length ?? 0,
+                cwsprChatResponseCode: 401,
+                cwsprChatTriggerInteraction: triggerInteraction,
+                reason: 'AuthenticationError',
+                reasonDesc: 'Inline chat requires SSO authentication, but current session is not',
+                credentialStartUrl: AuthUtil.instance.connection?.startUrl,
+            })
+            throw new ToolkitError('Inline chat is only available with SSO authentication')
+        }
+
         return this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then((context) => {
@@ -143,7 +189,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 +228,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/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts
@@ -58,7 +58,7 @@ import { processUtils } from 'aws-core-vscode/shared'
 import { activate as activateChat } from './chat/activation'
 import { activate as activeInlineChat } from '../inlineChat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
-import { auth2, getMfaTokenFromUser } from 'aws-core-vscode/auth'
+import { auth2, getMfaTokenFromUser, getMfaSerialFromUser } from 'aws-core-vscode/auth'
 import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { telemetry } from 'aws-core-vscode/telemetry'
 import { SessionManager } from '../app/inline/sessionManager'
@@ -343,8 +343,17 @@ async function postStartLanguageServer(
     client.onRequest(
         auth2.notificationTypes.getMfaCode.method,
         async (params: GetMfaCodeParams): Promise<GetMfaCodeResult> => {
-            const mfaCode = await getMfaTokenFromUser(params.mfaSerial, params.profileName)
-            return { code: mfaCode ?? '' }
+            if (params.mfaSerial) {
+                globals.globalState.update('recentMfaSerial', { mfaSerial: params.mfaSerial })
+            }
+            const defaultMfaSerial = globals.globalState.tryGet('recentMfaSerial', Object, {
+                mfaSerial: '',
+            }).mfaSerial
+            let mfaSerial = await getMfaSerialFromUser(defaultMfaSerial, params.profileName)
+            mfaSerial = mfaSerial.trim()
+            globals.globalState.update('recentMfaSerial', { mfaSerial: mfaSerial })
+            const mfaCode = await getMfaTokenFromUser(mfaSerial, params.profileName)
+            return { code: mfaCode ?? '', mfaSerial: mfaSerial ?? '' }
         }
     )
 
diff --git a/packages/core/src/auth/credentials/utils.ts b/packages/core/src/auth/credentials/utils.ts
@@ -102,6 +102,34 @@ export class CredentialsSettings extends fromExtensionManifest('aws', { profile:
 
 const errorMessageUserCancelled = localize('AWS.error.mfa.userCancelled', 'User cancelled entering authentication code')
 
+/**
+ * @description Prompts user for MFA serial number
+ *
+ * Entered token is passed to the callback.
+ * If user cancels out, the callback is passed an error with a fixed message string.
+ *
+ * @param profileName Name of Credentials profile we are asking an MFA Token for
+ * @param callback tokens/errors are passed through here
+ */
+export async function getMfaSerialFromUser(defaultSerial: string, profileName: string): Promise<string> {
+    const inputBox = createInputBox({
+        ignoreFocusOut: true,
+        placeholder: localize('AWS.prompt.mfa.enterCode.placeholder', 'Enter mfaSerial Number Here'),
+        title: localize('AWS.prompt.mfa.enterCode.title', 'MFA Challenge for {0}', profileName),
+        prompt: localize('AWS.prompt.mfa.enterCode.prompt', 'Enter Serial Number for MFA device', defaultSerial),
+        value: defaultSerial, // Pre-fill with default value
+    })
+
+    const token = await inputBox.prompt()
+
+    // Distinguish user cancel vs code entry issues with the error message
+    if (!isValidResponse(token)) {
+        throw new Error(errorMessageUserCancelled)
+    }
+
+    return token
+}
+
 /**
  * @description Prompts user for MFA token
  *
diff --git a/packages/core/src/auth/index.ts b/packages/core/src/auth/index.ts
@@ -22,7 +22,7 @@ export {
 } from './connection'
 export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'
-export { getMfaTokenFromUser } from './credentials/utils'
+export { getMfaTokenFromUser, getMfaSerialFromUser } from './credentials/utils'
 export { LoginManager } from './deprecated/loginManager'
 export * as constants from './sso/constants'
 export * as cache from './sso/cache'
diff --git 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
@@ -277,6 +277,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
     }
diff --git a/packages/core/src/codewhispererChat/index.ts b/packages/core/src/codewhispererChat/index.ts
@@ -12,6 +12,6 @@ export { ChatTriggerType, PromptMessage, TriggerPayload } from './controllers/ch
 export { UserIntentRecognizer } from './controllers/chat/userIntent/userIntentRecognizer'
 export { EditorContextExtractor } from './editor/context/extractor'
 export { ChatSessionStorage } from './storages/chatSession'
-export { TriggerEventsStorage } from './storages/triggerEvents'
+export { TriggerEventsStorage, TriggerEvent } from './storages/triggerEvents'
 export { ReferenceLogController } from './view/messages/referenceLogController'
 export { extractLanguageNameFromFile } from './editor/context/file/languages'
diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
@@ -20,7 +20,7 @@ import { builderIdStartUrl } from '../../../../auth/sso/constants'
 import { RegionProfile, vsCodeState } from '../../../../codewhisperer/models/model'
 import { randomUUID } from '../../../../shared/crypto'
 import globals from '../../../../shared/extensionGlobals'
-import { telemetry } from '../../../../shared/telemetry/telemetry'
+import { CredentialType, telemetry } from '../../../../shared/telemetry/telemetry'
 import { ProfileSwitchIntent } from '../../../../codewhisperer/region/regionProfileManager'
 
 const className = 'AmazonQLoginWebview'
@@ -204,12 +204,21 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         // Defining separate auth function to emit telemetry before returning from this method
         await globals.globalState.update('recentIamKeys', { accessKey: accessKey })
         await globals.globalState.update('recentRoleArn', { roleArn: roleArn })
+        let credentialsType: CredentialType | undefined
+        if (!sessionToken && !roleArn) {
+            credentialsType = 'staticProfile'
+        } else if (roleArn) {
+            credentialsType = 'assumeRoleProfile'
+        } else {
+            credentialsType = 'staticSessionProfile'
+        }
+
         const runAuth = async (): Promise<AuthError | undefined> => {
             try {
                 await AuthUtil.instance.loginIam(accessKey, secretKey, sessionToken, roleArn)
             } catch (e) {
                 getLogger().error('Failed submitting credentials %O', e)
-                const message = e instanceof Error ? e.message : e as string
+                const message = e instanceof Error ? e.message : (e as string)
                 return { id: this.id, text: message }
             }
             // Enable code suggestions
@@ -224,7 +233,10 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         const result = await runAuth()
         this.storeMetricMetadata({
             credentialSourceId: 'sharedCredentials',
-            authEnabledFeatures: 'codewhisperer',
+            featureId: 'codewhisperer',
+            credentialType: credentialsType,
+            isReAuth: false,
+            isAggregated: false,
             ...this.getResultForMetrics(result),
         })
         this.emitAuthMetric()
diff --git a/packages/core/src/login/webview/vue/backend.ts b/packages/core/src/login/webview/vue/backend.ts
@@ -177,7 +177,7 @@ export abstract class CommonAuthWebview extends VueWebview {
         accessKey: string,
         secretKey: string,
         sessionToken?: string,
-        role_arn?: string,
+        roleArn?: string
     ): Promise<AuthError | undefined>
 
     async showResourceExplorer(): Promise<void> {
diff --git a/packages/core/src/login/webview/vue/login.vue b/packages/core/src/login/webview/vue/login.vue
@@ -512,6 +512,9 @@ export default defineComponent({
                         }
                     }
                 } else if (this.selectedLoginOption === LoginOption.IAM_CREDENTIAL) {
+                    // Emit telemetry when IAM Credentials option is selected and Continue is clicked
+                    void client.emitUiClick('auth_credentialsOption')
+
                     this.stage = 'AWS_PROFILE'
                     this.$nextTick(() => document.getElementById('profileName')!.focus())
                 }
@@ -533,7 +536,13 @@ export default defineComponent({
                     return
                 }
                 this.stage = 'AUTHENTICATING'
-                const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey, this.sessionToken, this.roleArn)
+                const error = await client.startIamCredentialSetup(
+                    this.profileName,
+                    this.accessKey,
+                    this.secretKey,
+                    this.sessionToken,
+                    this.roleArn
+                )
                 if (error) {
                     this.stage = 'START'
                     void client.errorNotification(error)
diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
@@ -73,6 +73,7 @@ export type globalKey =
     | 'recentSso'
     | 'recentIamKeys'
     | 'recentRoleArn'
+    | 'recentMfaSerial'
     // List of regions enabled in AWS Explorer.
     | 'region'
     // TODO: implement this via `PromptSettings` instead of globalState.

Original file line number	Diff line number	Diff line change
`@@ -277,6 +277,10 @@ export class AuthUtil implements IAuthProvider {`
`277`	`277`	`return Boolean(this.connection?.startUrl && this.connection?.startUrl !== builderIdStartUrl)`
`278`	`278`	`}`
`279`	`279`
	`280`	`+ isIamConnection() {`
	`281`	`+ return Boolean(this.connection?.accessKey && this.connection?.secretKey)`
	`282`	`+ }`
	`283`	`+`
`280`	`284`	`isInternalAmazonUser(): boolean {`
`281`	`285`	`return this.isConnected() && this.connection?.startUrl === internalStartUrl`
`282`	`286`	`}`
Original file line number	Diff line number	Diff line change
`@@ -512,6 +512,9 @@ export default defineComponent({`
`512`	`512`	`}`
`513`	`513`	`}`
`514`	`514`	`} else if (this.selectedLoginOption === LoginOption.IAM_CREDENTIAL) {`
	`515`	`+ // Emit telemetry when IAM Credentials option is selected and Continue is clicked`
	`516`	`+ void client.emitUiClick('auth_credentialsOption')`
	`517`	`+`
`515`	`518`	`this.stage = 'AWS_PROFILE'`
`516`	`519`	`this.$nextTick(() => document.getElementById('profileName')!.focus())`
`517`	`520`	`}`
`@@ -533,7 +536,13 @@ export default defineComponent({`
`533`	`536`	`return`
`534`	`537`	`}`
`535`	`538`	`this.stage = 'AUTHENTICATING'`
`536`		`- const error = await client.startIamCredentialSetup(this.profileName, this.accessKey, this.secretKey, this.sessionToken, this.roleArn)`
	`539`	`+ const error = await client.startIamCredentialSetup(`
	`540`	`+ this.profileName,`
	`541`	`+ this.accessKey,`
	`542`	`+ this.secretKey,`
	`543`	`+ this.sessionToken,`
	`544`	`+ this.roleArn`
	`545`	`+ )`
`537`	`546`	`if (error) {`
`538`	`547`	`this.stage = 'START'`
`539`	`548`	`void client.errorNotification(error)`