UI fixes

app/me/layout.tsx

@@ -4,7 +4,7 @@ import { AppSidebar } from '@/components/app-sidebar';
 import { SiteHeader } from '@/components/site-header';
 import { SidebarInset, SidebarProvider } from '@/components/ui/sidebar';
 import { useAuthStatus } from '@/hooks/use-auth';
-import React, { useMemo } from 'react';
+import React, { useMemo, useRef, useEffect } from 'react';
 import LoadingSpinner from '@/components/LoadingSpinner';
 
 interface MeLayoutProps {
@@ -33,6 +33,13 @@ const getId = (item: ProfileItemWithId): string | undefined =>
 
 const MeLayout = ({ children }: MeLayoutProps): React.ReactElement => {
   const { user, isLoading } = useAuthStatus();
+  // Track whether we've completed the very first load.
+  // This prevents children from unmounting during background session refetches
+  // (e.g. on window focus), which would destroy component state like 2FA steps.
+  const hasLoadedOnce = useRef(false);
+  useEffect(() => {
+    if (!isLoading) hasLoadedOnce.current = true;
+  }, [isLoading]);
 
   const { name = '', email = '', profile, image: userImage = '' } = user || {};
   const typedProfile = profile as MeLayoutProfile | null | undefined;
@@ -62,7 +69,8 @@ const MeLayout = ({ children }: MeLayoutProps): React.ReactElement => {
     }).length;
   }, [typedProfile]);
 
-  if (isLoading) {
+  // Only show full-screen spinner on first load, not on background refetches
+  if (isLoading && !hasLoadedOnce.current) {
     return (
       <div className='flex h-screen items-center justify-center'>
         <LoadingSpinner size='xl' color='white' />

app/me/settings/SettingsContent.tsx

@@ -8,15 +8,20 @@ import { getMe } from '@/lib/api/auth';
 import { GetMeResponse } from '@/lib/api/types';
 import { Skeleton } from '@/components/ui/skeleton';
 import Settings from '@/components/profile/update/Settings';
-import { SecuritySettingsTab } from '@/components/profile/update/SecuritySettingsTab';
+import TwoFactorTab from '@/components/profile/update/TwoFactorTab';
+import SecurityTab from '@/components/profile/update/SecurityTab';
 import { IdentityVerificationSection } from '@/components/didit/IdentityVerificationSection';
 import { invalidateAuthProfileCache } from '@/hooks/use-auth';
+import { useRef } from 'react';
+import { Loader2 } from 'lucide-react';
 
 const SettingsContent = () => {
   const searchParams = useSearchParams();
   const fromVerification = searchParams.get('verification') === 'complete';
   const [userData, setUserData] = useState<GetMeResponse | null>(null);
   const [isLoading, setIsLoading] = useState(true);
+  // Prevent unmounting tabs on background refetches (e.g. after 2FA enable)
+  const hasLoadedOnce = useRef(false);
 
   const fetchUserData = useCallback(async () => {
     try {
@@ -26,11 +31,15 @@ const SettingsContent = () => {
       setUserData(null);
     } finally {
       setIsLoading(false);
+      hasLoadedOnce.current = true;
     }
   }, []);
 
   useEffect(() => {
-    setIsLoading(true);
+    // Only set isLoading true on the very first fetch
+    if (!hasLoadedOnce.current) {
+      setIsLoading(true);
+    }
     fetchUserData();
   }, [fetchUserData]);
 
@@ -39,7 +48,8 @@ const SettingsContent = () => {
     invalidateAuthProfileCache();
   }, [fetchUserData]);
 
-  if (isLoading) {
+  // Only show skeleton on first load — not on background refetches
+  if (isLoading && !hasLoadedOnce.current) {
     return (
       <div>
         <Skeleton className='h-full w-full' />
@@ -113,7 +123,14 @@ const SettingsContent = () => {
             </TabsTrigger>
           </TabsList>
           <TabsContent value='profile'>
-            <Profile user={userData?.user as User} />
+            {userData?.user ? (
+              <Profile user={userData.user as User} />
+            ) : (
+              <div className='flex items-center justify-center p-12'>
+                <Loader2 className='mr-2 h-8 w-8 animate-spin text-[#a7f950]' />
+                <span className='text-zinc-500'>Loading profile...</span>
+              </div>
+            )}
           </TabsContent>
           <TabsContent value='settings'>
             <Settings />
@@ -127,8 +144,30 @@ const SettingsContent = () => {
           <TabsContent value='preferences' className='space-y-6'>
             <Settings visibleSections={['appearance', 'preferences']} />
           </TabsContent>
-          <TabsContent value='security' className='space-y-6'>
-            <SecuritySettingsTab />
+          <TabsContent value='security'>
+            {userData?.user ? (
+              <SecurityTab user={userData.user as User} />
+            ) : (
+              <div className='flex items-center justify-center p-12'>
+                <Loader2 className='mr-2 h-8 w-8 animate-spin text-[#a7f950]' />
+                <span className='text-zinc-500'>
+                  Loading security settings...
+                </span>
+              </div>
+            )}
+          </TabsContent>
+          <TabsContent value='2fa'>
+            {userData?.user ? (
+              <TwoFactorTab
+                user={userData.user as User}
+                onStatusChange={fetchUserData}
+              />
+            ) : (
+              <div className='flex items-center justify-center p-12'>
+                <Loader2 className='mr-2 h-8 w-8 animate-spin text-[#a7f950]' />
+                <span className='text-zinc-500'>Loading 2FA settings...</span>
+              </div>
+            )}
           </TabsContent>
           <TabsContent value='identity' className='space-y-6'>
             <IdentityVerificationSection

components/auth/LoginWrapper.tsx

@@ -7,6 +7,7 @@ import { useForm } from 'react-hook-form';
 import { toast } from 'sonner';
 import z from 'zod';
 import LoginForm from './LoginForm';
+import TwoFactorVerify from './TwoFactorVerify';
 import { authClient } from '@/lib/auth-client';
 import { useAuthStore } from '@/lib/stores/auth-store';
 
@@ -32,6 +33,7 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
   const [showPassword, setShowPassword] = useState(false);
   const [isLoading, setIsLoading] = useState(false);
   const [lastMethod, setLastMethod] = useState<string | null>(null);
+  const [twoFactorRequired, setTwoFactorRequired] = useState(false);
 
   const callbackUrl = searchParams.get('callbackUrl')
     ? decodeURIComponent(searchParams.get('callbackUrl')!)
@@ -42,16 +44,6 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
     setLastMethod(method);
   }, []);
 
-  useEffect(() => {
-    const method = authClient.getLastUsedLoginMethod();
-    setLastMethod(method);
-  }, []);
-
-  useEffect(() => {
-    const method = authClient.getLastUsedLoginMethod();
-    setLastMethod(method);
-  }, []);
-
   const form = useForm<FormData>({
     resolver: zodResolver(formSchema),
     defaultValues: {
@@ -89,6 +81,17 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
       const errorStatus = error.status;
       const errorCode = error.code;
 
+      // Check for 2FA requirement
+      if (
+        errorStatus === 403 &&
+        (errorMessage === 'two_factor_required' ||
+          errorMessage?.includes('two_factor') ||
+          errorCode === 'TWO_FACTOR_REQUIRED')
+      ) {
+        setTwoFactorRequired(true);
+        return;
+      }
+
       if (errorStatus === 403 || errorCode === 'FORBIDDEN') {
         const message =
           'Please verify your email before signing in. Check your inbox for a verification link.';
@@ -171,8 +174,36 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
       setIsLoading(true);
       setLoadingState(true);
 
+      const syncSession = async () => {
+        const session = await authClient.getSession();
+        if (session && typeof session === 'object' && 'user' in session) {
+          const sessionUser = session.user as
+            | {
+                id: string;
+                email: string;
+                name?: string | null;
+                image?: string | null;
+              }
+            | null
+            | undefined;
+
+          if (sessionUser && sessionUser.id && sessionUser.email) {
+            const authStore = useAuthStore.getState();
+            await authStore.syncWithSession({
+              id: sessionUser.id,
+              email: sessionUser.email,
+              name: sessionUser.name || undefined,
+              image: sessionUser.image || undefined,
+              role: 'USER',
+              username: undefined,
+              accessToken: undefined,
+            });
+          }
+        }
+      };
+
       try {
-        const { error } = await authClient.signIn.email(
+        const { data, error } = await authClient.signIn.email(
           {
             email: values.email,
             password: values.password,
@@ -184,42 +215,26 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
               setIsLoading(true);
               setLoadingState(true);
             },
-            onSuccess: async () => {
-              await new Promise(resolve => setTimeout(resolve, 200));
+            onSuccess: async ctx => {
+              const resData = ctx?.data as
+                | {
+                    twoFactorRequired?: boolean;
+                    twoFactorRedirect?: boolean;
+                  }
+                | undefined;
 
-              const session = await authClient.getSession();
-
-              if (session && typeof session === 'object' && 'user' in session) {
-                const sessionUser = session.user as
-                  | {
-                      id: string;
-                      email: string;
-                      name?: string | null;
-                      image?: string | null;
-                    }
-                  | null
-                  | undefined;
-
-                if (sessionUser && sessionUser.id && sessionUser.email) {
-                  const authStore = useAuthStore.getState();
-                  await authStore.syncWithSession({
-                    id: sessionUser.id,
-                    email: sessionUser.email,
-                    name: sessionUser.name || undefined,
-                    image: sessionUser.image || undefined,
-                    role: 'USER',
-                    username: undefined,
-                    accessToken: undefined,
-                  });
-                }
+              if (resData?.twoFactorRequired || resData?.twoFactorRedirect) {
+                setTwoFactorRequired(true);
+                setIsLoading(false);
+                setLoadingState(false);
+                return;
               }
 
-              // Keep loading state active during redirect
-              // The page will unmount when redirecting, so no need to set false
+              await new Promise(resolve => setTimeout(resolve, 200));
+              await syncSession();
               window.location.href = callbackUrl;
             },
             onError: ctx => {
-              // Handle error from Better Auth callback
               const errorObj = ctx.error || ctx;
               handleAuthError(
                 typeof errorObj === 'object'
@@ -233,14 +248,21 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
           }
         );
 
-        // Handle error from return value
         if (error) {
           handleAuthError(error, values);
           setIsLoading(false);
           setLoadingState(false);
+        } else if (
+          (data as { twoFactorRequired?: boolean; twoFactorRedirect?: boolean })
+            ?.twoFactorRequired ||
+          (data as { twoFactorRequired?: boolean; twoFactorRedirect?: boolean })
+            ?.twoFactorRedirect
+        ) {
+          setTwoFactorRequired(true);
+          setIsLoading(false);
+          setLoadingState(false);
         }
       } catch (error) {
-        // Handle unexpected errors
         const errorObj =
           error instanceof Error
             ? { message: error.message, status: undefined, code: undefined }
@@ -254,6 +276,17 @@ const LoginWrapper = ({ setLoadingState }: LoginWrapperProps) => {
     [handleAuthError, setLoadingState, callbackUrl]
   );
 
+  if (twoFactorRequired) {
+    return (
+      <TwoFactorVerify
+        onSuccess={async () => {
+          window.location.href = callbackUrl;
+        }}
+        onCancel={() => setTwoFactorRequired(false)}
+      />
+    );
+  }
+
   return (
     <LoginForm
       form={form}