From 153651df35316704bb7c5a5650b2bb4eca4029e1 Mon Sep 17 00:00:00 2001
From: MananTank <manantankm@gmail.com>
Date: Mon, 27 Jan 2025 19:49:51 +0000
Subject: [PATCH] [TOOL-3241] Fix previous account shown after logging out in
 onboarding flow (#6061)

---
 apps/dashboard/src/app/login/LoginPage.tsx    | 49 ++++++++++++-------
 .../src/app/login/onboarding/General.tsx      |  3 ++
 .../onboarding/on-boarding-ui.client.tsx      |  2 +
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/apps/dashboard/src/app/login/LoginPage.tsx b/apps/dashboard/src/app/login/LoginPage.tsx
index 0fc410aa485..8eaf5f30737 100644
--- a/apps/dashboard/src/app/login/LoginPage.tsx
+++ b/apps/dashboard/src/app/login/LoginPage.tsx
@@ -9,7 +9,7 @@ import { useDashboardRouter } from "@/lib/DashboardRouter";
 import type { Account } from "@3rdweb-sdk/react/hooks/useApi";
 import { useTheme } from "next-themes";
 import Link from "next/link";
-import { Suspense, lazy, useState } from "react";
+import { Suspense, lazy, useEffect, useState } from "react";
 import { ConnectEmbed, useActiveWalletConnectionStatus } from "thirdweb/react";
 import { createWallet, inAppWallet } from "thirdweb/wallets";
 import { ClientOnly } from "../../components/ClientOnly/ClientOnly";
@@ -154,6 +154,32 @@ function PageContent(props: {
     router.replace(props.redirectPath);
   }
 
+  async function onLogin() {
+    const account = await getRawAccountAction();
+
+    // shouldn't happen - but if account is not found, stay on login page
+    if (!account) {
+      return;
+    }
+
+    if (!isOnboardingComplete(account)) {
+      setScreen({
+        id: "onboarding",
+        account,
+      });
+    } else {
+      onComplete();
+    }
+  }
+
+  // eslint-disable-next-line no-restricted-syntax
+  useEffect(() => {
+    // if suddenly disconnected
+    if (connectionStatus !== "connected" && screen.id !== "login") {
+      setScreen({ id: "login" });
+    }
+  }, [connectionStatus, screen.id]);
+
   if (connectionStatus === "connecting") {
     return <LoadingCard />;
   }
@@ -170,29 +196,14 @@ function PageContent(props: {
           onComplete={onComplete}
           redirectPath={props.redirectPath}
           redirectToCheckout={redirectToCheckout}
+          onLogout={() => {
+            setScreen({ id: "login" });
+          }}
         />
       </Suspense>
     );
   }
 
-  async function onLogin() {
-    const account = await getRawAccountAction();
-
-    // shouldn't happen - but if account is not found, stay on login page
-    if (!account) {
-      return;
-    }
-
-    if (!isOnboardingComplete(account)) {
-      setScreen({
-        id: "onboarding",
-        account,
-      });
-    } else {
-      onComplete();
-    }
-  }
-
   return <LoadingCard />;
 }
 
diff --git a/apps/dashboard/src/app/login/onboarding/General.tsx b/apps/dashboard/src/app/login/onboarding/General.tsx
index c34833d2648..9eaadf20cc5 100644
--- a/apps/dashboard/src/app/login/onboarding/General.tsx
+++ b/apps/dashboard/src/app/login/onboarding/General.tsx
@@ -14,12 +14,14 @@ type OnboardingGeneralProps = {
   account: Account;
   onSave: (email: string) => void;
   onDuplicate: (email: string) => void;
+  onLogout: () => void;
 };
 
 export const OnboardingGeneral: React.FC<OnboardingGeneralProps> = ({
   account,
   onSave,
   onDuplicate,
+  onLogout,
 }) => {
   const [existing, setExisting] = useState(false);
   const activeWallet = useActiveWallet();
@@ -27,6 +29,7 @@ export const OnboardingGeneral: React.FC<OnboardingGeneralProps> = ({
 
   async function handleLogout() {
     await doLogout();
+    onLogout();
     if (activeWallet) {
       disconnect(activeWallet);
     }
diff --git a/apps/dashboard/src/app/login/onboarding/on-boarding-ui.client.tsx b/apps/dashboard/src/app/login/onboarding/on-boarding-ui.client.tsx
index ba4fb375337..85467d12253 100644
--- a/apps/dashboard/src/app/login/onboarding/on-boarding-ui.client.tsx
+++ b/apps/dashboard/src/app/login/onboarding/on-boarding-ui.client.tsx
@@ -21,6 +21,7 @@ type OnboardingScreen =
 function OnboardingUI(props: {
   account: Account;
   onComplete: () => void;
+  onLogout: () => void;
   // path to redirect from stripe
   redirectPath: string;
   redirectToCheckout: RedirectBillingCheckoutAction;
@@ -68,6 +69,7 @@ function OnboardingUI(props: {
       {screen.id === "onboarding" && (
         <OnboardingGeneral
           account={account}
+          onLogout={props.onLogout}
           onSave={(email) => {
             setUpdatedEmail(email);
             setScreen({