fix: session callback on otp register

apps/login/app/(login)/otp/[method]/set/page.tsx

@@ -6,8 +6,10 @@ import {
   registerTOTP,
   server,
 } from "#/lib/zitadel";
+import Alert from "#/ui/Alert";
 import DynamicTheme from "#/ui/DynamicTheme";
 import TOTPRegister from "#/ui/TOTPRegister";
+import UserAvatar from "#/ui/UserAvatar";
 import { getMostRecentCookieWithLoginname } from "#/utils/cookies";
 
 export default async function Page({
@@ -21,27 +23,25 @@ export default async function Page({
   const { method } = params;
 
   const branding = await getBrandingSettings(server, organization);
+  const { session, token } = await loadSession(loginName, organization);
 
-  const totpResponse = await loadSession(loginName, organization).then(
-    ({ session, token }) => {
+  let totpResponse;
   if (session && session.factors?.user?.id) {
     if (method === "time-based") {
       // inconsistency with token: email works with machine token, totp works with session token
-          return registerTOTP(session.factors.user.id, token);
+      totpResponse = await registerTOTP(session.factors.user.id, token);
     } else if (method === "sms") {
       // does not work
-          return addOTPSMS(session.factors.user.id);
+      await addOTPSMS(session.factors.user.id);
     } else if (method === "email") {
       // works
-          return addOTPEmail(session.factors.user.id);
+      await addOTPEmail(session.factors.user.id);
     } else {
       throw new Error("Invalid method");
     }
   } else {
     throw new Error("No session found");
   }
-    }
-  );
 
   async function loadSession(loginName?: string, organization?: string) {
     const recent = await getMostRecentCookieWithLoginname(
@@ -58,6 +58,23 @@ export default async function Page({
     <DynamicTheme branding={branding}>
       <div className="flex flex-col items-center space-y-4">
         <h1>Register 2-factor</h1>
+        {!session && (
+          <div className="py-4">
+            <Alert>
+              Could not get the context of the user. Make sure to enter the
+              username first or provide a loginName as searchParam.
+            </Alert>
+          </div>
+        )}
+
+        {session && (
+          <UserAvatar
+            loginName={loginName ?? session.factors?.user?.loginName}
+            displayName={session.factors?.user?.displayName}
+            showDropdown
+          ></UserAvatar>
+        )}
+
         {totpResponse && "uri" in totpResponse && "secret" in totpResponse ? (
           <>
             <p className="ztdl-p">

apps/login/ui/TOTPRegister.tsx

@@ -44,8 +44,10 @@ export default function TOTPRegister({
   });
 
   async function continueWithCode(values: Inputs) {
+    setLoading(true);
     return verifyTOTP(values.code, loginName, organization)
       .then((response) => {
+        setLoading(false);
         if (authRequestId && sessionId) {
           const params = new URLSearchParams({
             sessionId: sessionId,
@@ -73,6 +75,7 @@ export default function TOTPRegister({
         }
       })
       .catch((e) => {
+        setLoading(false);
         setError(e.message);
       });
   }