@azure/msal-browser 5.0.0-alpha.0 → 5.0.1

package/src/custom_auth/reset_password/auth_flow/state/ResetPasswordCodeRequiredState.ts

@@ -8,11 +8,17 @@ import { ResetPasswordSubmitCodeResult } from "../result/ResetPasswordSubmitCode
 import { ResetPasswordCodeRequiredStateParameters } from "./ResetPasswordStateParameters.js";
 import { ResetPasswordState } from "./ResetPasswordState.js";
 import { ResetPasswordPasswordRequiredState } from "./ResetPasswordPasswordRequiredState.js";
+import { RESET_PASSWORD_CODE_REQUIRED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Reset password code required state.
  */
 export class ResetPasswordCodeRequiredState extends ResetPasswordState<ResetPasswordCodeRequiredStateParameters> {
+    /**
+     * The type of the state.
+     */
+    stateType = RESET_PASSWORD_CODE_REQUIRED_STATE_TYPE;
+
     /**
      * Submits a one-time passcode that the customer user received in their email in order to continue password reset flow.
      * @param {string} code - The code to submit.
@@ -56,6 +62,7 @@ export class ResetPasswordCodeRequiredState extends ResetPasswordState<ResetPass
                     signInClient: this.stateParameters.signInClient,
                     cacheClient: this.stateParameters.cacheClient,
                     jitClient: this.stateParameters.jitClient,
+                    mfaClient: this.stateParameters.mfaClient,
                     username: this.stateParameters.username,
                 })
             );
@@ -108,6 +115,7 @@ export class ResetPasswordCodeRequiredState extends ResetPasswordState<ResetPass
                     signInClient: this.stateParameters.signInClient,
                     cacheClient: this.stateParameters.cacheClient,
                     jitClient: this.stateParameters.jitClient,
+                    mfaClient: this.stateParameters.mfaClient,
                     username: this.stateParameters.username,
                     codeLength: result.codeLength,
                 })

package/src/custom_auth/reset_password/auth_flow/state/ResetPasswordCompletedState.ts

@@ -4,8 +4,14 @@
  */
 
 import { SignInContinuationState } from "../../../sign_in/auth_flow/state/SignInContinuationState.js";
+import { RESET_PASSWORD_COMPLETED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /**
  * Represents the state that indicates the successful completion of a password reset operation.
  */
-export class ResetPasswordCompletedState extends SignInContinuationState {}
+export class ResetPasswordCompletedState extends SignInContinuationState {
+    /**
+     * The type of the state.
+     */
+    stateType = RESET_PASSWORD_COMPLETED_STATE_TYPE;
+}

package/src/custom_auth/reset_password/auth_flow/state/ResetPasswordFailedState.ts

@@ -4,8 +4,14 @@
  */
 
 import { AuthFlowStateBase } from "../../../core/auth_flow/AuthFlowState.js";
+import { RESET_PASSWORD_FAILED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /**
  * State of a reset password operation that has failed.
  */
-export class ResetPasswordFailedState extends AuthFlowStateBase {}
+export class ResetPasswordFailedState extends AuthFlowStateBase {
+    /**
+     * The type of the state.
+     */
+    stateType = RESET_PASSWORD_FAILED_STATE_TYPE;
+}

package/src/custom_auth/reset_password/auth_flow/state/ResetPasswordPasswordRequiredState.ts

@@ -8,11 +8,17 @@ import { ResetPasswordState } from "./ResetPasswordState.js";
 import { ResetPasswordPasswordRequiredStateParameters } from "./ResetPasswordStateParameters.js";
 import { ResetPasswordCompletedState } from "./ResetPasswordCompletedState.js";
 import { SignInScenario } from "../../../sign_in/auth_flow/SignInScenario.js";
+import { RESET_PASSWORD_PASSWORD_REQUIRED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Reset password password required state.
  */
 export class ResetPasswordPasswordRequiredState extends ResetPasswordState<ResetPasswordPasswordRequiredStateParameters> {
+    /**
+     * The type of the state.
+     */
+    stateType = RESET_PASSWORD_PASSWORD_REQUIRED_STATE_TYPE;
+
     /**
      * Submits a new password for reset password flow.
      * @param {string} password - The password to submit.
@@ -59,6 +65,7 @@ export class ResetPasswordPasswordRequiredState extends ResetPasswordState<Reset
                     signInClient: this.stateParameters.signInClient,
                     cacheClient: this.stateParameters.cacheClient,
                     jitClient: this.stateParameters.jitClient,
+                    mfaClient: this.stateParameters.mfaClient,
                     signInScenario: SignInScenario.SignInAfterPasswordReset,
                 })
             );

package/src/custom_auth/reset_password/auth_flow/state/ResetPasswordStateParameters.ts

@@ -8,6 +8,7 @@ import { SignInClient } from "../../../sign_in/interaction_client/SignInClient.j
 import { CustomAuthSilentCacheClient } from "../../../get_account/interaction_client/CustomAuthSilentCacheClient.js";
 import { AuthFlowActionRequiredStateParameters } from "../../../core/auth_flow/AuthFlowState.js";
 import { JitClient } from "../../../core/interaction_client/jit/JitClient.js";
+import { MfaClient } from "../../../core/interaction_client/mfa/MfaClient.js";
 
 export interface ResetPasswordStateParameters
     extends AuthFlowActionRequiredStateParameters {
@@ -16,6 +17,7 @@ export interface ResetPasswordStateParameters
     signInClient: SignInClient;
     cacheClient: CustomAuthSilentCacheClient;
     jitClient: JitClient;
+    mfaClient: MfaClient;
 }
 
 export type ResetPasswordPasswordRequiredStateParameters =

package/src/custom_auth/reset_password/interaction_client/ResetPasswordClient.ts

@@ -44,12 +44,12 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
     async start(
         parameters: ResetPasswordStartParams
     ): Promise<ResetPasswordCodeRequiredResult> {
-        const correlationId = parameters.correlationId;
+        const correlationId = parameters.correlationId || this.correlationId;
         const apiId = PublicApiId.PASSWORD_RESET_START;
         const telemetryManager = initializeServerTelemetryManager(
             apiId,
             this.config.auth.clientId,
-            this.correlationId,
+            correlationId,
             this.browserStorage,
             this.logger
         );
@@ -77,7 +77,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
         const challengeRequest: ResetPasswordChallengeRequest = {
             continuation_token: startResponse.continuation_token ?? "",
             challenge_type: this.getChallengeTypes(parameters.challengeType),
-            correlationId: correlationId,
+            correlationId: startResponse.correlation_id || correlationId,
             telemetryManager: telemetryManager,
         };
 
@@ -92,7 +92,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
     async submitCode(
         parameters: ResetPasswordSubmitCodeParams
     ): Promise<ResetPasswordPasswordRequiredResult> {
-        const correlationId = parameters.correlationId;
+        const correlationId = parameters.correlationId || this.correlationId;
         ensureArgumentIsNotEmptyString(
             "parameters.code",
             parameters.code,
@@ -103,7 +103,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
         const telemetryManager = initializeServerTelemetryManager(
             apiId,
             this.config.auth.clientId,
-            this.correlationId,
+            correlationId,
             this.browserStorage,
             this.logger
         );
@@ -131,7 +131,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
         );
 
         return {
-            correlationId: response.correlation_id,
+            correlationId: response.correlation_id || correlationId,
             continuationToken: response.continuation_token ?? "",
         };
     }
@@ -144,11 +144,12 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
     async resendCode(
         parameters: ResetPasswordResendCodeParams
     ): Promise<ResetPasswordCodeRequiredResult> {
+        const correlationId = parameters.correlationId || this.correlationId;
         const apiId = PublicApiId.PASSWORD_RESET_RESEND_CODE;
         const telemetryManager = initializeServerTelemetryManager(
             apiId,
             this.config.auth.clientId,
-            this.correlationId,
+            correlationId,
             this.browserStorage,
             this.logger
         );
@@ -156,7 +157,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
         const challengeRequest: ResetPasswordChallengeRequest = {
             continuation_token: parameters.continuationToken,
             challenge_type: this.getChallengeTypes(parameters.challengeType),
-            correlationId: parameters.correlationId,
+            correlationId: correlationId,
             telemetryManager: telemetryManager,
         };
 
@@ -171,7 +172,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
     async submitNewPassword(
         parameters: ResetPasswordSubmitNewPasswordParams
     ): Promise<ResetPasswordCompletedResult> {
-        const correlationId = parameters.correlationId;
+        const correlationId = parameters.correlationId || this.correlationId;
 
         ensureArgumentIsNotEmptyString(
             "parameters.newPassword",
@@ -183,7 +184,7 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
         const telemetryManager = initializeServerTelemetryManager(
             apiId,
             this.config.auth.clientId,
-            this.correlationId,
+            correlationId,
             this.browserStorage,
             this.logger
         );
@@ -207,13 +208,13 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
 
         this.logger.verbose(
             "Submit endpoint called successfully with new password for password reset.",
-            correlationId
+            submitResponse.correlation_id || correlationId
         );
 
         return this.performPollCompletionRequest(
             submitResponse.continuation_token ?? "",
             submitResponse.poll_interval,
-            correlationId,
+            submitResponse.correlation_id || correlationId,
             telemetryManager
         );
     }
@@ -234,18 +235,18 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
 
         this.logger.verbose(
             "Challenge endpoint for password reset returned successfully.",
-            correlationId
+            response.correlation_id || correlationId
         );
 
         if (response.challenge_type === ChallengeType.OOB) {
             // Code is required
             this.logger.verbose(
                 "Code is required for password reset flow.",
-                correlationId
+                response.correlation_id || correlationId
             );
 
             return {
-                correlationId: response.correlation_id,
+                correlationId: response.correlation_id || correlationId,
                 continuationToken: response.continuation_token ?? "",
                 challengeChannel: response.challenge_channel ?? "",
                 challengeTargetLabel: response.challenge_target_label ?? "",
@@ -257,13 +258,13 @@ export class ResetPasswordClient extends CustomAuthInteractionClientBase {
 
         this.logger.error(
             `Unsupported challenge type '${response.challenge_type}' returned from challenge endpoint for password reset.`,
-            correlationId
+            response.correlation_id || correlationId
         );
 
         throw new CustomAuthApiError(
             CustomAuthApiErrorCode.UNSUPPORTED_CHALLENGE_TYPE,
             `Unsupported challenge type '${response.challenge_type}'.`,
-            correlationId
+            response.correlation_id || correlationId
         );
     }
 

package/src/custom_auth/sign_in/auth_flow/result/SignInResendCodeResult.ts

@@ -4,6 +4,10 @@
  */
 
 import { AuthFlowResultBase } from "../../../core/auth_flow/AuthFlowResultBase.js";
+import {
+    SIGN_IN_CODE_REQUIRED_STATE_TYPE,
+    SIGN_IN_FAILED_STATE_TYPE,
+} from "../../../core/auth_flow/AuthFlowStateTypes.js";
 import { SignInResendCodeError } from "../error_type/SignInError.js";
 import type { SignInCodeRequiredState } from "../state/SignInCodeRequiredState.js";
 import { SignInFailedState } from "../state/SignInFailedState.js";
@@ -39,7 +43,7 @@ export class SignInResendCodeResult extends AuthFlowResultBase<
      * Checks if the result is in a failed state.
      */
     isFailed(): this is SignInResendCodeResult & { state: SignInFailedState } {
-        return this.state instanceof SignInFailedState;
+        return this.state.stateType === SIGN_IN_FAILED_STATE_TYPE;
     }
 
     /**
@@ -52,7 +56,7 @@ export class SignInResendCodeResult extends AuthFlowResultBase<
          * The instanceof operator couldn't be used here to check the state type since the circular dependency issue.
          * So we are using the constructor name to check the state type.
          */
-        return this.state.constructor?.name === "SignInCodeRequiredState";
+        return this.state.stateType === SIGN_IN_CODE_REQUIRED_STATE_TYPE;
     }
 }
 

package/src/custom_auth/sign_in/auth_flow/result/SignInResult.ts

@@ -11,6 +11,15 @@ import { SignInPasswordRequiredState } from "../state/SignInPasswordRequiredStat
 import { SignInFailedState } from "../state/SignInFailedState.js";
 import { SignInCompletedState } from "../state/SignInCompletedState.js";
 import { AuthMethodRegistrationRequiredState } from "../../../core/auth_flow/jit/state/AuthMethodRegistrationState.js";
+import { MfaAwaitingState } from "../../../core/auth_flow/mfa/state/MfaState.js";
+import {
+    SIGN_IN_CODE_REQUIRED_STATE_TYPE,
+    SIGN_IN_PASSWORD_REQUIRED_STATE_TYPE,
+    SIGN_IN_FAILED_STATE_TYPE,
+    SIGN_IN_COMPLETED_STATE_TYPE,
+    AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE,
+    MFA_AWAITING_STATE_TYPE,
+} from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Result of a sign-in operation.
@@ -44,7 +53,7 @@ export class SignInResult extends AuthFlowResultBase<
      * Checks if the result is in a failed state.
      */
     isFailed(): this is SignInResult & { state: SignInFailedState } {
-        return this.state instanceof SignInFailedState;
+        return this.state.stateType === SIGN_IN_FAILED_STATE_TYPE;
     }
 
     /**
@@ -53,7 +62,7 @@ export class SignInResult extends AuthFlowResultBase<
     isCodeRequired(): this is SignInResult & {
         state: SignInCodeRequiredState;
     } {
-        return this.state instanceof SignInCodeRequiredState;
+        return this.state.stateType === SIGN_IN_CODE_REQUIRED_STATE_TYPE;
     }
 
     /**
@@ -62,14 +71,14 @@ export class SignInResult extends AuthFlowResultBase<
     isPasswordRequired(): this is SignInResult & {
         state: SignInPasswordRequiredState;
     } {
-        return this.state instanceof SignInPasswordRequiredState;
+        return this.state.stateType === SIGN_IN_PASSWORD_REQUIRED_STATE_TYPE;
     }
 
     /**
      * Checks if the result is in a completed state.
      */
     isCompleted(): this is SignInResult & { state: SignInCompletedState } {
-        return this.state instanceof SignInCompletedState;
+        return this.state.stateType === SIGN_IN_COMPLETED_STATE_TYPE;
     }
 
     /**
@@ -79,7 +88,18 @@ export class SignInResult extends AuthFlowResultBase<
     isAuthMethodRegistrationRequired(): this is SignInResult & {
         state: AuthMethodRegistrationRequiredState;
     } {
-        return this.state instanceof AuthMethodRegistrationRequiredState;
+        return (
+            this.state.stateType ===
+            AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE
+        );
+    }
+
+    /**
+     * Checks if the result requires MFA.
+     * @warning This API is experimental. It may be changed in the future without notice. Do not use in production applications.
+     */
+    isMfaRequired(): this is SignInResult & { state: MfaAwaitingState } {
+        return this.state.stateType === MFA_AWAITING_STATE_TYPE;
     }
 }
 
@@ -91,10 +111,12 @@ export class SignInResult extends AuthFlowResultBase<
  * - SignInFailedState: The sign-in process has failed.
  * - SignInCompletedState: The sign-in process is completed.
  * - AuthMethodRegistrationRequiredState: The sign-in process requires authentication method registration.
+ * - MfaAwaitingState: The sign-in process requires MFA.
  */
 export type SignInResultState =
     | SignInCodeRequiredState
     | SignInPasswordRequiredState
     | SignInFailedState
     | SignInCompletedState
-    | AuthMethodRegistrationRequiredState;
+    | AuthMethodRegistrationRequiredState
+    | MfaAwaitingState;

package/src/custom_auth/sign_in/auth_flow/result/SignInSubmitCodeResult.ts

@@ -8,6 +8,14 @@ import { SignInCompletedState } from "../state/SignInCompletedState.js";
 import { SignInFailedState } from "../state/SignInFailedState.js";
 import { AuthFlowResultBase } from "../../../core/auth_flow/AuthFlowResultBase.js";
 import { CustomAuthAccountData } from "../../../get_account/auth_flow/CustomAuthAccountData.js";
+import { AuthMethodRegistrationRequiredState } from "../../../core/auth_flow/jit/state/AuthMethodRegistrationState.js";
+import { MfaAwaitingState } from "../../../core/auth_flow/mfa/state/MfaState.js";
+import {
+    SIGN_IN_FAILED_STATE_TYPE,
+    SIGN_IN_COMPLETED_STATE_TYPE,
+    AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE,
+    MFA_AWAITING_STATE_TYPE,
+} from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Result of a sign-in submit code operation.
@@ -35,7 +43,7 @@ export class SignInSubmitCodeResult extends AuthFlowResultBase<
      * Checks if the result is in a failed state.
      */
     isFailed(): this is SignInSubmitCodeResult & { state: SignInFailedState } {
-        return this.state instanceof SignInFailedState;
+        return this.state.stateType === SIGN_IN_FAILED_STATE_TYPE;
     }
 
     /**
@@ -44,7 +52,30 @@ export class SignInSubmitCodeResult extends AuthFlowResultBase<
     isCompleted(): this is SignInSubmitCodeResult & {
         state: SignInCompletedState;
     } {
-        return this.state instanceof SignInCompletedState;
+        return this.state.stateType === SIGN_IN_COMPLETED_STATE_TYPE;
+    }
+
+    /**
+     * Checks if the result requires authentication method registration.
+     * @warning This API is experimental. It may be changed in the future without notice. Do not use in production applications.
+     */
+    isAuthMethodRegistrationRequired(): this is SignInSubmitCodeResult & {
+        state: AuthMethodRegistrationRequiredState;
+    } {
+        return (
+            this.state.stateType ===
+            AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE
+        );
+    }
+
+    /**
+     * Checks if the result requires MFA.
+     * @warning This API is experimental. It may be changed in the future without notice. Do not use in production applications.
+     */
+    isMfaRequired(): this is SignInSubmitCodeResult & {
+        state: MfaAwaitingState;
+    } {
+        return this.state.stateType === MFA_AWAITING_STATE_TYPE;
     }
 }
 
@@ -53,7 +84,11 @@ export class SignInSubmitCodeResult extends AuthFlowResultBase<
  * This includes:
  * - SignInCompletedState: The sign-in process has completed successfully.
  * - SignInFailedState: The sign-in process has failed.
+ * - AuthMethodRegistrationRequiredState: The user needs to register an authentication method.
+ * - MfaAwaitingState: The user is in a multi-factor authentication (MFA) waiting state.
  */
 export type SignInSubmitCodeResultState =
     | SignInCompletedState
-    | SignInFailedState;
+    | SignInFailedState
+    | AuthMethodRegistrationRequiredState
+    | MfaAwaitingState;

package/src/custom_auth/sign_in/auth_flow/result/SignInSubmitPasswordResult.ts

@@ -9,6 +9,13 @@ import { SignInFailedState } from "../state/SignInFailedState.js";
 import { AuthFlowResultBase } from "../../../core/auth_flow/AuthFlowResultBase.js";
 import { CustomAuthAccountData } from "../../../get_account/auth_flow/CustomAuthAccountData.js";
 import { AuthMethodRegistrationRequiredState } from "../../../core/auth_flow/jit/state/AuthMethodRegistrationState.js";
+import { MfaAwaitingState } from "../../../core/auth_flow/mfa/state/MfaState.js";
+import {
+    SIGN_IN_FAILED_STATE_TYPE,
+    SIGN_IN_COMPLETED_STATE_TYPE,
+    AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE,
+    MFA_AWAITING_STATE_TYPE,
+} from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Result of a sign-in submit password operation.
@@ -33,7 +40,7 @@ export class SignInSubmitPasswordResult extends AuthFlowResultBase<
     isFailed(): this is SignInSubmitPasswordResult & {
         state: SignInFailedState;
     } {
-        return this.state instanceof SignInFailedState;
+        return this.state.stateType === SIGN_IN_FAILED_STATE_TYPE;
     }
 
     /**
@@ -42,7 +49,7 @@ export class SignInSubmitPasswordResult extends AuthFlowResultBase<
     isCompleted(): this is SignInSubmitPasswordResult & {
         state: SignInCompletedState;
     } {
-        return this.state instanceof SignInCompletedState;
+        return this.state.stateType === SIGN_IN_COMPLETED_STATE_TYPE;
     }
 
     /**
@@ -52,7 +59,20 @@ export class SignInSubmitPasswordResult extends AuthFlowResultBase<
     isAuthMethodRegistrationRequired(): this is SignInSubmitPasswordResult & {
         state: AuthMethodRegistrationRequiredState;
     } {
-        return this.state instanceof AuthMethodRegistrationRequiredState;
+        return (
+            this.state.stateType ===
+            AUTH_METHOD_REGISTRATION_REQUIRED_STATE_TYPE
+        );
+    }
+
+    /**
+     * Checks if the result requires MFA.
+     * @warning This API is experimental. It may be changed in the future without notice. Do not use in production applications.
+     */
+    isMfaRequired(): this is SignInSubmitPasswordResult & {
+        state: MfaAwaitingState;
+    } {
+        return this.state.stateType === MFA_AWAITING_STATE_TYPE;
     }
 }
 
@@ -62,8 +82,10 @@ export class SignInSubmitPasswordResult extends AuthFlowResultBase<
  * - SignInCompletedState: The sign-in process has completed successfully.
  * - SignInFailedState: The sign-in process has failed.
  * - AuthMethodRegistrationRequiredState: The sign-in process requires authentication method registration.
+ * - MfaAwaitingState: The sign-in process requires MFA.
  */
 export type SignInSubmitPasswordResultState =
     | SignInCompletedState
     | SignInFailedState
-    | AuthMethodRegistrationRequiredState;
+    | AuthMethodRegistrationRequiredState
+    | MfaAwaitingState;

package/src/custom_auth/sign_in/auth_flow/state/SignInCodeRequiredState.ts

@@ -3,7 +3,6 @@
  * Licensed under the MIT License.
  */
 
-import { CustomAuthAccountData } from "../../../get_account/auth_flow/CustomAuthAccountData.js";
 import {
     SignInResendCodeParams,
     SignInSubmitCodeParams,
@@ -12,12 +11,17 @@ import { SignInResendCodeResult } from "../result/SignInResendCodeResult.js";
 import { SignInSubmitCodeResult } from "../result/SignInSubmitCodeResult.js";
 import { SignInCodeRequiredStateParameters } from "./SignInStateParameters.js";
 import { SignInState } from "./SignInState.js";
-import { SignInCompletedState } from "./SignInCompletedState.js";
+import { SIGN_IN_CODE_REQUIRED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /*
  * Sign-in code required state.
  */
 export class SignInCodeRequiredState extends SignInState<SignInCodeRequiredStateParameters> {
+    /**
+     * The type of the state.
+     */
+    stateType = SIGN_IN_CODE_REQUIRED_STATE_TYPE;
+
     /**
      * Once user configures email one-time passcode as a authentication method in Microsoft Entra, a one-time passcode will be sent to the user’s email.
      * Submit this one-time passcode to continue sign-in flow.
@@ -45,7 +49,7 @@ export class SignInCodeRequiredState extends SignInState<SignInCodeRequiredState
                 this.stateParameters.correlationId
             );
 
-            const completedResult =
+            const submitCodeResult =
                 await this.stateParameters.signInClient.submitCode(
                     submitCodeParams
                 );
@@ -55,17 +59,18 @@ export class SignInCodeRequiredState extends SignInState<SignInCodeRequiredState
                 this.stateParameters.correlationId
             );
 
-            const accountInfo = new CustomAuthAccountData(
-                completedResult.authenticationResult.account,
-                this.stateParameters.config,
-                this.stateParameters.cacheClient,
-                this.stateParameters.logger,
-                this.stateParameters.correlationId
+            const nextState = this.handleSignInResult(
+                submitCodeResult,
+                this.stateParameters.scopes
             );
 
+            if (nextState.error) {
+                return SignInSubmitCodeResult.createWithError(nextState.error);
+            }
+
             return new SignInSubmitCodeResult(
-                new SignInCompletedState(),
-                accountInfo
+                nextState.state,
+                nextState.accountInfo
             );
         } catch (error) {
             this.stateParameters.logger.errorPii(
@@ -115,6 +120,7 @@ export class SignInCodeRequiredState extends SignInState<SignInCodeRequiredState
                     signInClient: this.stateParameters.signInClient,
                     cacheClient: this.stateParameters.cacheClient,
                     jitClient: this.stateParameters.jitClient,
+                    mfaClient: this.stateParameters.mfaClient,
                     username: this.stateParameters.username,
                     codeLength: result.codeLength,
                     scopes: this.stateParameters.scopes,

package/src/custom_auth/sign_in/auth_flow/state/SignInCompletedState.ts

@@ -4,9 +4,15 @@
  */
 
 import { AuthFlowStateBase } from "../../../core/auth_flow/AuthFlowState.js";
+import { SIGN_IN_COMPLETED_STATE_TYPE } from "../../../core/auth_flow/AuthFlowStateTypes.js";
 
 /**
  * Represents the completed state of the sign-in operation.
  * This state indicates that the sign-in process has finished successfully.
  */
-export class SignInCompletedState extends AuthFlowStateBase {}
+export class SignInCompletedState extends AuthFlowStateBase {
+    /**
+     * The type of the state.
+     */
+    stateType = SIGN_IN_COMPLETED_STATE_TYPE;
+}