From f37d5ab84e68f426ca7a5d2da1725da9beaadf4a Mon Sep 17 00:00:00 2001
From: Nicolas Mahe <nicolas@mahe.me>
Date: Mon, 11 Mar 2024 17:53:45 +0700
Subject: [PATCH] open sign modal only when auto-login fails

---
 hooks/useAccount.ts | 11 ++++++++---
 pages/_app.tsx      | 41 +++++++++++++++++++++++++----------------
 2 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/hooks/useAccount.ts b/hooks/useAccount.ts
index e7c9ae7b..c60fb951 100644
--- a/hooks/useAccount.ts
+++ b/hooks/useAccount.ts
@@ -11,7 +11,7 @@ type AccountDetail = {
   address?: string
   jwtToken: string | null
   logout: () => Promise<void>
-  login: (connector: Connector) => Promise<void>
+  login: (connector: Connector) => Promise<boolean>
   sign: (connector: Connector) => Promise<void>
 }
 
@@ -64,8 +64,12 @@ export default function useAccount(): AccountDetail {
       const signer = walletClientToSigner(wallet)
       const currentAddress = (await signer.getAddress()).toLowerCase()
       if (jwt && currentAddress === jwt.address) {
-        return setAuthenticationToken(jwt.token)
+        setAuthenticationToken(jwt.token)
+        // fully login, nothing to do
+        return true
       }
+      // not fully login, need to sign
+      return false
     },
     [jwt, setAuthenticationToken],
   )
@@ -76,6 +80,7 @@ export default function useAccount(): AccountDetail {
       const signer = walletClientToSigner(wallet)
       const { jwtToken } = await authenticate(signer)
 
+      setAuthenticationToken(jwtToken)
       const newJwt = jwtDecode<JwtPayload>(jwtToken)
       setCookie(COOKIE_JWT_TOKEN, jwtToken, {
         ...COOKIE_OPTIONS,
@@ -87,7 +92,7 @@ export default function useAccount(): AccountDetail {
           : {}),
       })
     },
-    [authenticate, setCookie],
+    [authenticate, setAuthenticationToken, setCookie],
   )
 
   // Server side
diff --git a/pages/_app.tsx b/pages/_app.tsx
index 08a91ca1..b061927a 100644
--- a/pages/_app.tsx
+++ b/pages/_app.tsx
@@ -18,12 +18,14 @@ import React, {
   Fragment,
   JSX,
   PropsWithChildren,
+  useCallback,
   useEffect,
   useMemo,
   useState,
 } from 'react'
 import { Cookies, CookiesProvider, useCookies } from 'react-cookie'
 import {
+  Connector,
   WagmiConfig,
   useDisconnect,
   useAccount as useWagmiAccount,
@@ -128,36 +130,43 @@ function AccountProvider({
   const { LITEFLOW_API_KEY, BASE_URL } = useEnvironment()
   const { jwtToken, login, logout } = useAccount()
   const { disconnect } = useDisconnect()
-  const { isOpen, onOpen, onClose } = useDisclosure()
+  const { isOpen, onOpen: openSignModal, onClose } = useDisclosure()
   const [cookies] = useCookies<string, COOKIES>([COOKIE_JWT_TOKEN])
 
   const { connector } = useWagmiAccount({
-    async onConnect({ connector }) {
-      if (!connector) return
-      try {
-        await login(connector)
-        // jwtToken from useAccount is null on refresh, so we need to check the cookies
-        if (!cookies[COOKIE_JWT_TOKEN]) onOpen()
-      } catch (e: any) {
-        disconnect()
-      }
-    },
-    onDisconnect() {
+    onConnect: useCallback(
+      async ({ connector }: { connector?: Connector }) => {
+        if (!connector) return
+        try {
+          const isConnected = await login(connector)
+          if (!isConnected) openSignModal()
+        } catch (e: any) {
+          disconnect()
+        }
+      },
+      [disconnect, login, openSignModal],
+    ),
+    onDisconnect: useCallback(() => {
       void logout()
-    },
+    }, [logout]),
   })
 
   // handle change of account
   useEffect(() => {
     if (!connector) return
-    const handleLogin = () => {
-      login(connector).then(onOpen).catch(onError)
+    const handleLogin = async () => {
+      try {
+        const isConnected = await login(connector)
+        if (!isConnected) openSignModal()
+      } catch (error) {
+        onError(1)
+      }
     }
     connector.on('change', handleLogin)
     return () => {
       connector.off('change', handleLogin)
     }
-  }, [connector, login, onError, onOpen])
+  }, [connector, cookies, login, onError, openSignModal])
 
   const client = useMemo(
     // The client needs to be reset when the jwtToken changes but only on the client as the server will