Attempt #2: Add explicit state for initializing authMachine (#1142) (#1162)

* Revert "Revert "Add explicit state for initializing authMachine (#1142)" (#1148)"

This reverts commit 6b629e0.

* Wait for machine to init before subcomponents render

Author: wlee221

Diff for: .changeset/lucky-numbers-lay.md

@@ -0,0 +1,8 @@
+---
+"@aws-amplify/ui-react": patch
+"@aws-amplify/ui": patch
+"@aws-amplify/ui-vue": patch
+"@aws-amplify/ui-angular": patch
+---
+
+Add explicit `INIT` step for initializing authMachine

Diff for: docs/src/pages/components/authenticator/LabelsAndTextDemo.tsx

@@ -71,15 +71,28 @@ function Screen({ Component }: ScreenProps) {
 }
 
 export function LabelsAndTextDemo({ Component }: ScreenProps) {
+  const OnMachineInit = ({ children }) => {
+    /**
+     * This waits for Authenticator machine to init before its inner components
+     * start consuming machine context.
+     */
+    const { route } = useAuthenticator();
+    if (!route || route === 'idle' || route === 'setup') return null;
+
+    return <>{children}</>;
+  };
+
   return (
     <Authenticator.Provider>
-      <View data-amplify-authenticator="">
-        <View data-amplify-container="">
-          <View data-amplify-body>
-            <Screen Component={Component} />
+      <OnMachineInit>
+        <View data-amplify-authenticator="">
+          <View data-amplify-container="">
+            <View data-amplify-body>
+              <Screen Component={Component} />
+            </View>
           </View>
         </View>
-      </View>
+      </OnMachineInit>
     </Authenticator.Provider>
   );
 }

Diff for: packages/angular/projects/ui-angular/src/lib/services/authenticator.service.ts

@@ -37,18 +37,21 @@ export class AuthenticatorService implements OnDestroy {
     signUpAttributes,
     socialProviders,
   }: AuthenticatorMachineOptions) {
-    const machine = createAuthenticatorMachine({
-      initialState,
-      loginMechanisms,
-      services,
-      signUpAttributes,
-      socialProviders,
+    const machine = createAuthenticatorMachine();
+
+    const authService = interpret(machine).start();
+
+    authService.send({
+      type: 'INIT',
+      data: {
+        initialState,
+        loginMechanisms,
+        socialProviders,
+        signUpAttributes,
+        services,
+      },
     });
 
-    const authService = interpret(machine, {
-      devTools: process.env.NODE_ENV === 'development',
-    }).start();
-
     this._subscription = authService.subscribe((state) => {
       this._authState = state;
       this._facade = getServiceContextFacade(state);

Diff for: packages/react/src/components/Authenticator/Provider/index.tsx

@@ -23,19 +23,20 @@ const useAuthenticatorValue = ({
   signUpAttributes,
   services,
 }: ProviderProps) => {
-  const [state, send] = useMachine(
-    () =>
-      createAuthenticatorMachine({
+  const [state, send] = useMachine(() => createAuthenticatorMachine());
+
+  React.useEffect(() => {
+    send({
+      type: 'INIT',
+      data: {
         initialState,
         loginMechanisms,
-        services,
-        signUpAttributes,
         socialProviders,
-      }),
-    {
-      devTools: process.env.NODE_ENV === 'development',
-    }
-  );
+        signUpAttributes,
+        services,
+      },
+    });
+  }, []);
 
   const components = React.useMemo(
     () => ({ ...defaultComponents, ...customComponents }),

Diff for: packages/react/src/components/Authenticator/Router/index.tsx

@@ -60,6 +60,7 @@ export function Router({
             {(() => {
               switch (route) {
                 case 'idle':
+                case 'setup':
                   return null;
                 case 'confirmSignUp':
                   return <ConfirmSignUp />;

Diff for: packages/ui/src/helpers/auth.ts

@@ -226,6 +226,8 @@ export const getServiceContextFacade = (state: AuthMachineState) => {
     switch (true) {
       case state.matches('idle'):
         return 'idle';
+      case state.matches('setup'):
+        return 'setup';
       case state.matches('signOut'):
         return 'signOut';
       case state.matches('authenticated'):

Diff for: packages/ui/src/machines/authenticator/index.ts

@@ -9,50 +9,54 @@ import { createSignUpMachine } from './signUp';
 const DEFAULT_COUNTRY_CODE = '+1';
 
 export type AuthenticatorMachineOptions = AuthContext['config'] & {
-  initialState?: 'signIn' | 'signUp' | 'resetPassword';
-  services?: Partial<typeof defaultServices>;
+  services?: AuthContext['services'];
 };
 
-export function createAuthenticatorMachine({
-  initialState = 'signIn',
-  loginMechanisms,
-  signUpAttributes,
-  socialProviders,
-  services: customServices,
-}: AuthenticatorMachineOptions) {
-  const services = {
-    ...defaultServices,
-    ...customServices,
-  };
-
+export function createAuthenticatorMachine() {
   return createMachine<AuthContext, AuthEvent>(
     {
       id: 'authenticator',
       initial: 'idle',
       context: {
         user: undefined,
-        config: {
-          loginMechanisms,
-          signUpAttributes,
-          socialProviders,
-        },
+        config: {},
+        services: {},
         actorRef: undefined,
       },
       states: {
         // See: https://xstate.js.org/docs/guides/communication.html#invoking-promises
         idle: {
+          on: {
+            INIT: {
+              target: 'setup',
+              actions: 'configure',
+            },
+          },
+        },
+        setup: {
           invoke: [
             {
               // TODO Wait for Auth to be configured
-              src: 'getCurrentUser',
+              src: (context, _) => context.services.getCurrentUser(),
               onDone: {
                 actions: 'setUser',
                 target: 'authenticated',
               },
-              onError: initialState,
+              onError: [
+                {
+                  target: 'signUp',
+                  cond: (context) => context.config.initialState === 'signUp',
+                },
+                {
+                  target: 'resetPassword',
+                  cond: (context) =>
+                    context.config.initialState === 'resetPassword',
+                },
+                { target: 'signIn' },
+              ],
             },
             {
-              src: 'getAmplifyConfig',
+              src: (context, _) => context.services.getAmplifyConfig(),
               onDone: {
                 actions: 'applyAmplifyConfig',
               },
@@ -87,7 +91,7 @@ export function createAuthenticatorMachine({
           on: {
             SIGN_IN: 'signIn',
             'done.invoke.signUpActor': {
-              target: 'idle',
+              target: 'setup',
               actions: 'setUser',
             },
           },
@@ -159,7 +163,14 @@ export function createAuthenticatorMachine({
               cliLoginMechanisms.push('username');
             }
 
-            // Prefer explicitly configured settings over default CLI values
+            // Prefer explicitly configured settings over default CLI values\
+
+            const {
+              loginMechanisms,
+              signUpAttributes,
+              socialProviders,
+              initialState,
+            } = context.config;
             return {
               loginMechanisms: loginMechanisms ?? cliLoginMechanisms,
               signUpAttributes:
@@ -171,11 +182,13 @@ export function createAuthenticatorMachine({
                   ])
                 ),
               socialProviders: socialProviders ?? cliSocialProviders.sort(),
+              initialState,
             };
           },
         }),
         spawnSignInActor: assign({
           actorRef: (context, event) => {
+            const { services } = context;
             const actor = signInActor({ services }).withContext({
               authAttributes: event.data?.authAttributes,
               user: event.data?.user,
@@ -192,6 +205,7 @@ export function createAuthenticatorMachine({
         }),
         spawnSignUpActor: assign({
           actorRef: (context, event) => {
+            const { services } = context;
             const actor = createSignUpMachine({ services }).withContext({
               authAttributes: event.data?.authAttributes ?? {},
               country_code: DEFAULT_COUNTRY_CODE,
@@ -207,6 +221,7 @@ export function createAuthenticatorMachine({
         }),
         spawnResetPasswordActor: assign({
           actorRef: (context, event) => {
+            const { services } = context;
             const actor = resetPasswordActor({ services }).withContext({
               formValues: {},
               touched: {},
@@ -225,6 +240,13 @@ export function createAuthenticatorMachine({
             return spawn(actor, { name: 'signOutActor' });
           },
         }),
+        configure: assign((_, event) => {
+          const { services: customServices, ...config } = event.data;
+          return {
+            services: { ...defaultServices, ...customServices },
+            config,
+          };
+        }),
       },
       guards: {
         shouldRedirectToSignUp: (_, event): boolean => {
@@ -236,7 +258,6 @@ export function createAuthenticatorMachine({
           return event.data.intent === 'confirmPasswordReset';
         },
       },
-      services,
     }
   );
 }

Diff for: packages/ui/src/types/authMachine.ts

@@ -1,6 +1,7 @@
 import { CognitoUser, CodeDeliveryDetails } from 'amazon-cognito-identity-js';
 import { Interpreter, State } from 'xstate';
 import { ValidationError } from './validator';
+import { defaultServices } from '../machines/authenticator/defaultServices';
 
 export type AuthFormData = Record<string, string>;
 
@@ -10,7 +11,9 @@ export interface AuthContext {
     loginMechanisms?: LoginMechanism[];
     signUpAttributes?: SignUpAttribute[];
     socialProviders?: SocialProvider[];
+    initialState?: 'signIn' | 'signUp' | 'resetPassword';
   };
+  services?: Partial<typeof defaultServices>;
   user?: CognitoUserAmplify;
   username?: string;
   password?: string;
@@ -131,6 +134,7 @@ export type AuthEventTypes =
   | 'SIGN_UP'
   | 'SKIP'
   | 'SUBMIT'
+  | 'INIT'
   | InvokeActorEventTypes;
 
 export enum AuthChallengeNames {

Diff for: packages/vue/src/components/authenticator.vue

@@ -1,6 +1,6 @@
 <script setup lang="ts">
 import { useAuth } from '../composables/useAuth';
-import { ref, computed, useAttrs, watch, Ref } from 'vue';
+import { ref, computed, useAttrs, watch, Ref, onMounted } from 'vue';
 import { useActor, useInterpret } from '@xstate/vue';
 import {
   getActorState,
@@ -58,21 +58,26 @@ const emit = defineEmits([
   'verifyUserSubmit',
   'confirmVerifyUserSubmit',
 ]);
-const machine = createAuthenticatorMachine({
-  initialState,
-  loginMechanisms,
-  services,
-  signUpAttributes,
-  socialProviders,
-});
+const machine = createAuthenticatorMachine();
 
-const service = useInterpret(machine, {
-  devTools: process.env.NODE_ENV === 'development',
-});
+const service = useInterpret(machine);
 
 const { state, send } = useActor(service);
 useAuth(service);
 
+onMounted(() => {
+  send({
+    type: 'INIT',
+    data: {
+      initialState,
+      loginMechanisms,
+      socialProviders,
+      signUpAttributes,
+      services,
+    },
+  });
+});
+
 const actorState = computed(() => getActorState(state.value));
 
 const signInComponent = ref();