Multi-step MFA location support

new-ui/src/pages/full/OverviewPage/components/ConnectModal/ConnectModal.tsx

@@ -2,16 +2,18 @@
 import './style.scss';
 
 import { useQuery } from '@tanstack/react-query';
-import { type ReactNode, useEffect } from 'react';
+import { Fragment, type ReactNode, useEffect } from 'react';
 import { useShallow } from 'zustand/shallow';
 import { Modal } from '../../../../../shared/components/Modal/Modal';
 import { useAppData } from '../../../../../shared/providers/AppDataContext';
 import { api } from '../../../../../shared/rust-api/api';
 import { isPresent } from '../../../../../shared/utils/isPresent';
+import { mfaStepCount, mfaToText } from '../../../../../shared/utils/mfa';
 import {
   ConnectModalTitle,
   ConnectModalView,
   type ConnectModalViewValue,
+  mfaMethodToConnectModalView,
 } from './hooks/types';
 import { useConnectModal } from './hooks/useConnectModal';
 import { ConnectModalConnectionError } from './views/ConnectModalConnectionError/ConnectModalConnectionError';
@@ -25,17 +27,26 @@ import { ConnectModalPostureCheckFail } from './views/ConnectModalPostureCheckFa
 export const ConnectModal = () => {
   const reset = useConnectModal((s) => s.reset);
 
-  const [view, visible, location] = useConnectModal(
-    useShallow((s) => [s.view, s.visible, s.location]),
+  const [view, visible, location, stepIndex, stepPlan] = useConnectModal(
+    useShallow((s) => [s.view, s.visible, s.location, s.stepIndex, s.stepPlan]),
   );
 
+  const stepCount = isPresent(location) ? mfaStepCount(location) : 0;
+  const stepMethod = stepPlan[stepIndex];
+  const isOnMfaStepView =
+    isPresent(stepMethod) && view === mfaMethodToConnectModalView(stepMethod);
+  const stepLabel =
+    stepCount > 1 && isOnMfaStepView
+      ? `Step ${stepIndex + 1}/${stepCount}: ${mfaToText(stepMethod)}`
+      : null;
+
   const isOpen = isPresent(view) && isPresent(location) && visible;
 
   return (
     <Modal
       id="connect-modal"
       size="small"
-      title={view ? ConnectModalTitle[view] : ''}
+      title={stepLabel ?? (view ? ConnectModalTitle[view] : '')}
       isOpen={isOpen}
       afterClose={() => {
         reset();
@@ -75,6 +86,7 @@ const ModalContent = () => {
     retry: false,
   });
   const activeView = useConnectModal((s) => s.view);
+  const stepIndex = useConnectModal((s) => s.stepIndex);
 
   // When user completes connection and it's working modal is no longer needed so auto close it
   // biome-ignore lint/correctness/useExhaustiveDependencies: side-effect on connect
@@ -88,5 +100,7 @@ const ModalContent = () => {
 
   if (!activeView) return null;
 
-  return viewContent[activeView];
+  return (
+    <Fragment key={`${activeView}-${stepIndex}`}>{viewContent[activeView]}</Fragment>
+  );
 };

new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModal.tsx

@@ -4,7 +4,9 @@ import {
   MfaMethod,
   type MfaMethodValue,
 } from '../../../../../../shared/rust-api/types';
-import type { ConnectModalViewValue } from './types';
+import { isPresent } from '../../../../../../shared/utils/isPresent';
+import { resolveMfaStepPlan } from '../../../../../../shared/utils/mfa';
+import { type ConnectModalViewValue, mfaMethodToConnectModalView } from './types';
 
 interface StoreValues {
   visible: boolean;
@@ -14,6 +16,9 @@ interface StoreValues {
   postureError: string | null;
   autoStartOpenId: boolean;
   mfaMethod: MfaMethodValue;
+  stepIndex: number;
+  stepPlan: MfaMethodValue[];
+  mfaToken: string | null;
 }
 
 const defaults: StoreValues = {
@@ -24,11 +29,16 @@ const defaults: StoreValues = {
   perviousView: null,
   postureError: null,
   autoStartOpenId: false,
+  stepIndex: 0,
+  stepPlan: [],
+  mfaToken: null,
 } as const;
 
 interface Store extends StoreValues {
   open: (init?: Partial<StoreValues>) => void;
   setView: (view: ConnectModalViewValue, values?: Partial<StoreValues>) => void;
+  setMfaToken: (token: string) => void;
+  goToStep: (stepIndex: number) => void;
   reset: () => void;
 }
 
@@ -38,7 +48,16 @@ export const useConnectModal = create<Store>((set, get) => ({
     set(defaults);
   },
   open: (init) => {
-    set({ ...defaults, ...init, visible: true });
+    const location = init?.location ?? null;
+    const stepPlan = isPresent(location) ? resolveMfaStepPlan(location) : [];
+    set({ ...defaults, ...init, stepPlan, visible: true });
+  },
+  setMfaToken: (token) => {
+    set({ mfaToken: token });
+  },
+  goToStep: (stepIndex) => {
+    const { stepPlan, setView } = get();
+    setView(mfaMethodToConnectModalView(stepPlan[stepIndex]), { stepIndex });
   },
   setView: (view, vals) => {
     const pervious = get().view ?? null;

new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useConnectModalMfaOidc.ts

@@ -29,7 +29,9 @@ export const useConnectModalMfaOidc = ({
   onSessionExpired,
   onServiceUnavailable,
 }: Options = {}) => {
-  const location = useConnectModal(useShallow((s) => s.location));
+  const [location, stepPlan, mfaToken, setMfaToken] = useConnectModal(
+    useShallow((s) => [s.location, s.stepPlan, s.mfaToken, s.setMfaToken]),
+  );
 
   const [isStarting, setIsStarting] = useState(false);
   const [startError, setStartError] = useState<string | null>(null);
@@ -71,13 +73,21 @@ export const useConnectModalMfaOidc = ({
     cleanup();
 
     try {
-      const info = await api.mfaStart(instance.id, location.id, MfaMethod.Oidc);
-      await api.openLink(`${instance.proxy_url}openid/mfa?token=${info.token}`);
+      const session = await api.startMfaStep(
+        instance.id,
+        location.id,
+        MfaMethod.Oidc,
+        stepPlan,
+        mfaToken,
+      );
+      setMfaToken(session.token);
+
+      await api.openLink(`${instance.proxy_url}openid/mfa?token=${session.token}`);
 
       setIsStarting(false);
       setIsPolling(true);
 
-      const taskId = await api.mfaPollOpenId(instance.id, location.id, info.token);
+      const taskId = await api.mfaPollOpenId(instance.id, location.id, session.token);
       taskIdRef.current = taskId;
 
       // The backend brings up the connection itself; completion means connected.
@@ -126,6 +136,9 @@ export const useConnectModalMfaOidc = ({
   }, [
     instance,
     location,
+    stepPlan,
+    mfaToken,
+    setMfaToken,
     cleanup,
     onPostureError,
     onSessionExpired,

new-ui/src/pages/full/OverviewPage/components/ConnectModal/hooks/useMfaStep.ts

@@ -0,0 +1,29 @@
+import { useShallow } from 'zustand/shallow';
+import { isPresent } from '../../../../../../shared/utils/isPresent';
+import { usableMfaMethods } from '../../../../../../shared/utils/mfa';
+import { useConnectModal } from './useConnectModal';
+
+export const useMfaStep = () => {
+  const [location, stepIndex, stepPlan, mfaToken, setMfaToken, goToStep] =
+    useConnectModal(
+      useShallow((s) => [
+        s.location,
+        s.stepIndex,
+        s.stepPlan,
+        s.mfaToken,
+        s.setMfaToken,
+        s.goToStep,
+      ]),
+    );
+
+  const currentStep = location?.mfa_steps[stepIndex];
+
+  return {
+    canPickOtherMethod:
+      isPresent(currentStep) && usableMfaMethods(currentStep).length > 1,
+    stepPlan,
+    mfaToken,
+    setMfaToken,
+    goToStep,
+  };
+};

new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaEmail/ConnectModalMfaEmail.tsx

@@ -10,6 +10,7 @@ import { isPresent } from '../../../../../../../shared/utils/isPresent';
 import { ConnectModalPostureCheckLoading } from '../../components/ConnectModalPostureCheckLoading/ConnectModalPostureCheckLoading';
 import { ConnectModalView } from '../../hooks/types';
 import { useConnectModal } from '../../hooks/useConnectModal';
+import { useMfaStep } from '../../hooks/useMfaStep';
 
 const MIN_POSTURE_LOADER_MS = 500;
 
@@ -18,10 +19,16 @@ export const ConnectModalMfaEmail = () => {
     useShallow((s) => [s.perviousView, s.location]),
   );
 
+  const { canPickOtherMethod, stepPlan, mfaToken, setMfaToken, goToStep } = useMfaStep();
+
   const { verifyCode, isVerifying, verifyError, isStarting, startError } = useMfaConnect(
     location as LocationInfo,
     MfaMethod.Email,
     {
+      stepPlan,
+      mfaToken,
+      setMfaToken,
+      onStepAdvanced: goToStep,
       debounceMs: location?.posture_check_required ? MIN_POSTURE_LOADER_MS : 0,
       onSessionExpired: () =>
         useConnectModal.getState().setView(perviousView ?? ConnectModalView.MfaSettings),
@@ -87,13 +94,15 @@ export const ConnectModalMfaEmail = () => {
         }}
       />
       <Controls>
-        <Button
-          variant={ButtonVariant.Secondary}
-          text="Use different MFA"
-          onClick={() => {
-            useConnectModal.getState().setView(ConnectModalView.MfaSettings);
-          }}
-        />
+        {canPickOtherMethod && (
+          <Button
+            variant={ButtonVariant.Secondary}
+            text="Other methods"
+            onClick={() => {
+              useConnectModal.getState().setView(ConnectModalView.MfaSettings);
+            }}
+          />
+        )}
         <div className="right">
           <Button
             text="Verify"

new-ui/src/pages/full/OverviewPage/components/ConnectModal/views/ConnectModalMfaMobile/ConnectModalMfaMobile.tsx

@@ -10,6 +10,7 @@ import type { LocationInfo } from '../../../../../../../shared/rust-api/types';
 import { ConnectModalPostureCheckLoading } from '../../components/ConnectModalPostureCheckLoading/ConnectModalPostureCheckLoading';
 import { ConnectModalView } from '../../hooks/types';
 import { useConnectModal } from '../../hooks/useConnectModal';
+import { useMfaStep } from '../../hooks/useMfaStep';
 
 type Screen = 'loading' | 'qr' | 'error';
 
@@ -18,9 +19,14 @@ export const ConnectModalMfaMobile = () => {
     useShallow((s) => [s.perviousView, s.location]),
   );
 
+  const { canPickOtherMethod, stepPlan, mfaToken, setMfaToken } = useMfaStep();
+
   const { start, isStarting, startError, qrValue, connectionError } = useMfaMobileConnect(
     location as LocationInfo,
     {
+      stepPlan,
+      mfaToken,
+      setMfaToken,
       onPostureError: (msg) => {
         useConnectModal.setState({ postureError: msg });
         useConnectModal.getState().setView(ConnectModalView.PostureCheckFail);
@@ -68,11 +74,11 @@ export const ConnectModalMfaMobile = () => {
         </div>
       )}
       <Controls>
-        {screen === 'qr' && (
+        {screen === 'qr' && canPickOtherMethod && (
           <Button
             containerProps={{ className: 'full' }}
             variant={ButtonVariant.Secondary}
-            text="Use different MFA"
+            text="Other methods"
             onClick={() => {
               useConnectModal.getState().setView(ConnectModalView.MfaSettings);
             }}