From 4f860c6de38fae5a24688ca5335df24bdf598521 Mon Sep 17 00:00:00 2001
From: noah <noahsaso@gmail.com>
Date: Thu, 3 Aug 2023 12:37:17 -0700
Subject: [PATCH] Improve wallet connection (#1321)

---
 .../stateful/components/WalletProvider.tsx    | 37 ++++++++++++-------
 .../components/wallet/WalletUiConnected.tsx   |  3 +-
 packages/stateful/hooks/useWallet.ts          | 29 ++++++++++++---
 3 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/packages/stateful/components/WalletProvider.tsx b/packages/stateful/components/WalletProvider.tsx
index fdfa28cfb8..fa36410d9e 100644
--- a/packages/stateful/components/WalletProvider.tsx
+++ b/packages/stateful/components/WalletProvider.tsx
@@ -22,7 +22,9 @@ import {
   useRef,
 } from 'react'
 import { useTranslation } from 'react-i18next'
+import { useRecoilValue } from 'recoil'
 
+import { mountedInBrowserAtom } from '@dao-dao/state/recoil'
 import { Web3AuthPrompt } from '@dao-dao/types'
 import {
   CHAIN_ENDPOINTS,
@@ -206,32 +208,41 @@ const InnerWalletProvider = ({ children }: PropsWithChildren<{}>) => {
   }, [chain.chain_id, walletManager])
 
   // Auto-connect to Keplr mobile web if in that context.
-  const attemptedKeplrMobileWebConnectionRef = useRef(false)
+  const mountedInBrowser = useRecoilValue(mountedInBrowserAtom)
+  const connectingKeplrMobileWebConnectionRef = useRef(false)
   useEffect(() => {
     if (
       typeof window === 'undefined' ||
+      !mountedInBrowser ||
       !isWalletDisconnected ||
-      attemptedKeplrMobileWebConnectionRef.current
+      connectingKeplrMobileWebConnectionRef.current
     ) {
       return
     }
 
-    // Don't try again.
-    attemptedKeplrMobileWebConnectionRef.current = true
+    // Only try once at a time.
+    connectingKeplrMobileWebConnectionRef.current = true
 
     // Connect.
     ;(async () => {
-      const keplr = await (
-        await import('@keplr-wallet/stores')
-      ).getKeplrFromWindow()
-      const keplrWallet = walletRepo.wallets.find(
-        (wallet) => wallet.walletInfo.name === 'keplr-extension'
-      )
-      if (keplr && keplrWallet && keplr.mode === 'mobile-web') {
-        keplrWallet.connect()
+      try {
+        const keplr = await (
+          await import('@keplr-wallet/stores')
+        ).getKeplrFromWindow()
+        const keplrWallet = walletRepo.wallets.find(
+          (wallet) =>
+            wallet.walletInfo.name === keplrExtensionWallets[0].walletInfo.name
+        )
+        if (keplr && keplrWallet && keplr.mode === 'mobile-web') {
+          await keplrWallet.connect()
+        }
+      } catch (err) {
+        console.error('Keplr mobile web connection failed', err)
+      } finally {
+        connectingKeplrMobileWebConnectionRef.current = false
       }
     })()
-  }, [isWalletDisconnected, walletRepo])
+  }, [isWalletDisconnected, mountedInBrowser, walletRepo])
 
   return <>{children}</>
 }
diff --git a/packages/stateful/components/wallet/WalletUiConnected.tsx b/packages/stateful/components/wallet/WalletUiConnected.tsx
index 828fe3f21a..7f04c097fc 100644
--- a/packages/stateful/components/wallet/WalletUiConnected.tsx
+++ b/packages/stateful/components/wallet/WalletUiConnected.tsx
@@ -89,7 +89,8 @@ export const WalletUiConnected = ({
       </div>
 
       {/* In Keplr mobile web, the wallet is force connected and cannot be logged out of, so only show the log out button for all other options. */}
-      {(walletRepo?.current.walletName !== keplrExtensionWallet.walletName ||
+      {(walletRepo?.current.walletInfo.name !==
+        keplrExtensionWallet.walletInfo.name ||
         !(walletRepo?.current.client instanceof KeplrClient) ||
         walletRepo?.current.client.client.mode !== 'mobile-web') && (
         <Button
diff --git a/packages/stateful/hooks/useWallet.ts b/packages/stateful/hooks/useWallet.ts
index c4bc67bb8a..1bd25a7638 100644
--- a/packages/stateful/hooks/useWallet.ts
+++ b/packages/stateful/hooks/useWallet.ts
@@ -5,8 +5,14 @@ import { useChain as useWalletChain } from '@cosmos-kit/react-lite'
 import { useEffect, useMemo, useRef, useState } from 'react'
 import { useRecoilValue } from 'recoil'
 
-import { walletChainIdAtom } from '@dao-dao/state/recoil'
-import { useChainContextIfAvailable } from '@dao-dao/stateless'
+import {
+  walletChainIdAtom,
+  walletHexPublicKeySelector,
+} from '@dao-dao/state/recoil'
+import {
+  useCachedLoading,
+  useChainContextIfAvailable,
+} from '@dao-dao/stateless'
 import { LoadingData } from '@dao-dao/types'
 import { getChainForChainId } from '@dao-dao/utils'
 
@@ -47,6 +53,16 @@ export const useWallet = ({
   const [account, setAccount] = useState<WalletAccount>()
   const [hexPublicKeyData, setHexPublicKeyData] = useState<string>()
 
+  const hexPublicKeyFromChain = useCachedLoading(
+    _walletChain.address && loadAccount
+      ? walletHexPublicKeySelector({
+          walletAddress: _walletChain.address,
+          chainId: _walletChain.chain.chain_id,
+        })
+      : undefined,
+    undefined
+  )
+
   useEffect(() => {
     if (!loadAccount) {
       return
@@ -84,9 +100,11 @@ export const useWallet = ({
       // Use chain from our version of the chain-registry.
       chain: getChainForChainId(walletChainRef.current.chain.chain_id),
       account,
-      hexPublicKey: !hexPublicKeyData
-        ? { loading: true }
-        : { loading: false, data: hexPublicKeyData },
+      hexPublicKey: hexPublicKeyData
+        ? { loading: false, data: hexPublicKeyData }
+        : !hexPublicKeyFromChain.loading && hexPublicKeyFromChain.data
+        ? { loading: false, data: hexPublicKeyFromChain.data }
+        : { loading: true },
     }),
     // eslint-disable-next-line react-hooks/exhaustive-deps
     [
@@ -95,6 +113,7 @@ export const useWallet = ({
       walletChainRef.current.address,
       walletChainRef.current.chain.chain_id,
       walletChainRef.current.status,
+      hexPublicKeyFromChain,
     ]
   )