From 996c3a3a452a795bb6bb9d3f01fe49e5e61db6b2 Mon Sep 17 00:00:00 2001
From: Brian Wagner <brian.wagner35@gmail.com>
Date: Fri, 12 Jul 2024 19:26:50 -0400
Subject: [PATCH] fix: only reset auth state when using the logout hook

---
 .../react/src/components/auth/context.ts      |  1 +
 account-kit/react/src/context.tsx             | 27 ++++++++++---------
 account-kit/react/src/hooks/useLogout.ts      |  5 +++-
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/account-kit/react/src/components/auth/context.ts b/account-kit/react/src/components/auth/context.ts
index 56d37d6448..c54289849e 100644
--- a/account-kit/react/src/components/auth/context.ts
+++ b/account-kit/react/src/components/auth/context.ts
@@ -16,6 +16,7 @@ export type AuthStep =
 type AuthContextType = {
   authStep: AuthStep;
   setAuthStep: (step: AuthStep) => void;
+  resetAuthStep: () => void;
 };
 
 export const AuthModalContext = createContext<AuthContextType | undefined>(
diff --git a/account-kit/react/src/context.tsx b/account-kit/react/src/context.tsx
index 76c1a06ca7..f661659fc6 100644
--- a/account-kit/react/src/context.tsx
+++ b/account-kit/react/src/context.tsx
@@ -118,6 +118,19 @@ export const AlchemyAccountProvider = (
   const openAuthModal = useCallback(() => setIsModalOpen(true), []);
   const closeAuthModal = useCallback(() => setIsModalOpen(false), []);
 
+  /**
+   * Reset the auth step to the initial state. This also clears the email auth query params from the URL.
+   */
+  const resetAuthStep = useCallback(() => {
+    setAuthStep({ type: "initial" });
+
+    const url = new URL(window.location.href);
+    url.searchParams.delete("orgId");
+    url.searchParams.delete("bundle");
+    url.searchParams.delete(IS_SIGNUP_QP);
+    window.history.replaceState({}, "", url.toString());
+  }, []);
+
   const initialContext = useMemo(
     () => ({
       config,
@@ -151,19 +164,6 @@ export const AlchemyAccountProvider = (
     }
   }, [status, config.ui, openAuthModal]);
 
-  // Force the auth UI back to the initial state when disconnected
-  useEffect(() => {
-    if (status === "DISCONNECTED" && !isAuthenticating) {
-      setAuthStep({ type: "initial" });
-      // Remove the relevant query params (could also do this once connected)
-      const url = new URL(window.location.href);
-      url.searchParams.delete("orgId");
-      url.searchParams.delete("bundle");
-      url.searchParams.delete(IS_SIGNUP_QP);
-      window.history.replaceState({}, "", url.toString());
-    }
-  }, [status, isAuthenticating]);
-
   return (
     <Hydrate {...props}>
       <AlchemyAccountContext.Provider value={initialContext}>
@@ -172,6 +172,7 @@ export const AlchemyAccountProvider = (
             value={{
               authStep,
               setAuthStep,
+              resetAuthStep,
             }}
           >
             {children}
diff --git a/account-kit/react/src/hooks/useLogout.ts b/account-kit/react/src/hooks/useLogout.ts
index dd5841947c..c8811315a9 100644
--- a/account-kit/react/src/hooks/useLogout.ts
+++ b/account-kit/react/src/hooks/useLogout.ts
@@ -5,6 +5,7 @@ import { useDisconnect } from "wagmi";
 import { useAlchemyAccountContext } from "../context.js";
 import type { BaseHookMutationArgs } from "../types.js";
 import { useSigner } from "./useSigner.js";
+import { useAuthContext } from "../components/auth/context.js";
 
 export type UseLogoutMutationArgs = BaseHookMutationArgs<void, void>;
 
@@ -44,6 +45,7 @@ export function useLogout(
   } = useAlchemyAccountContext();
   const signer = useSigner();
   const { disconnectAsync } = useDisconnect({ config: wagmiConfig });
+  const { resetAuthStep } = useAuthContext();
 
   const {
     mutate: logout,
@@ -53,7 +55,8 @@ export function useLogout(
     {
       mutationFn: async () => {
         await disconnectAsync();
-        return signer?.disconnect();
+        await signer?.disconnect();
+        resetAuthStep();
       },
       ...mutationArgs,
     },