feat: enable mfa serial input in pop up window

Author: yuxianrz

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

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 code 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
  *

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'

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.