add more mfa feature

Author: yuxianrz

packages/amazonq/src/lsp/client.ts

@@ -32,7 +32,6 @@ import {
     updateConfigurationRequestType,
     GetMfaCodeParams,
     GetMfaCodeResult,
-    getMfaCodeRequestType,
 } from '@aws/language-server-runtimes/protocol'
 import {
     AuthUtil,
@@ -59,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 } from 'aws-core-vscode/auth'
+import { auth2, getMfaTokenFromUser } 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'
@@ -346,10 +345,10 @@ async function postStartLanguageServer(
     )
 
     // Handler for when Flare needs to assume a role with MFA code
-    client.onRequest<GetMfaCodeParams, GetMfaCodeResult>(
-        getMfaCodeRequestType.method,
+    client.onRequest(
+        auth2.notificationTypes.getMfaCode.method,
         async (params: GetMfaCodeParams): Promise<GetMfaCodeResult> => {
-            const mfaCode = await vscode.window.showInputBox({ title: 'Enter MFA Code' })
+            const mfaCode = await getMfaTokenFromUser(params.mfaSerial, params.profileName)
             return { code: mfaCode ?? '' }
         }
     )

packages/core/src/auth/auth2.ts

@@ -49,6 +49,9 @@ import {
     iamCredentialsUpdateRequestType,
     Profile,
     SsoSession,
+    GetMfaCodeParams,
+    getMfaCodeRequestType,
+
 } from '@aws/language-server-runtimes/protocol'
 import { LanguageClient } from 'vscode-languageclient'
 import { getLogger } from '../shared/logger/logger'
@@ -70,6 +73,9 @@ export const notificationTypes = {
     getConnectionMetadata: new RequestType<undefined, ConnectionMetadata, Error>(
         getConnectionMetadataRequestType.method
     ),
+    getMfaCode: new RequestType<GetMfaCodeParams, ResponseMessage, Error>(
+        getMfaCodeRequestType.method
+    )
 }
 
 export type AuthState = 'notConnected' | 'connected' | 'expired'
@@ -140,7 +146,7 @@ export class LanguageClientAuth {
             {
                 profileName: profileName,
                 options: {
-                    generateOnInvalidStsCredential: login,
+                    callStsOnInvalidIamCredential: login,
                 },
             } satisfies GetIamCredentialParams,
             cancellationToken
@@ -182,7 +188,7 @@ export class LanguageClientAuth {
         let profile: Profile
         if (roleArn) {
             profile = {
-                kinds: [ProfileKind.IamRoleSourceProfile],
+                kinds: [ProfileKind.IamSourceProfileProfile],
                 name: profileName,
                 settings: {
                     sso_session: '',
@@ -195,7 +201,7 @@ export class LanguageClientAuth {
             }
         } else if (accessKey && secretKey) {
             profile = {
-                kinds: [ProfileKind.IamUserProfile],
+                kinds: [ProfileKind.IamCredentialsProfile],
                 name: profileName,
                 settings: {
                     sso_session: '',

packages/core/src/auth/index.ts

@@ -22,6 +22,7 @@ export {
 } from './connection'
 export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'
+export { getMfaTokenFromUser } from './credentials/utils'
 export { LoginManager } from './deprecated/loginManager'
 export * as constants from './sso/constants'
 export * as cache from './sso/cache'

packages/core/src/test/credentials/auth2.test.ts

@@ -110,7 +110,7 @@ describe('LanguageClientAuth', () => {
             const requestParams = client.sendRequest.firstCall.args[1]
             sinon.assert.match(requestParams.profile, {
                 name: profileName,
-                kinds: [ProfileKind.IamUserProfile],
+                kinds: [ProfileKind.IamCredentialsProfile],
             })
             sinon.assert.match(requestParams.profile.settings, {
                 aws_access_key_id: 'accessKey',