feat(wallet): save selected wallet type local storage

(#570)

apps/deploy-web/src/components/layout/WalletStatus.tsx

@@ -21,13 +21,15 @@ const withBilling = browserEnvConfig.NEXT_PUBLIC_BILLING_ENABLED;
 
 export function WalletStatus() {
   const { walletName, isWalletLoaded, isWalletConnected, isManaged, isWalletLoading, isTrialing } = useWallet();
-  const { balance: walletBalance } = useWalletBalance();
+  const { balance: walletBalance, isLoading: isWalletBalanceLoading } = useWalletBalance();
   const [isSignedInWithTrial] = useAtom(walletStore.isSignedInWithTrial);
   const [open, setOpen] = useState(false);
+  const isLoadingBalance = isWalletBalanceLoading && !walletBalance;
+  const isInit = isWalletLoaded && !isWalletLoading && !isLoadingBalance;
 
   return (
     <>
-      {isWalletLoaded && !isWalletLoading ? (
+      {isInit ? (
         isWalletConnected ? (
           <div className="flex w-full items-center">
             <div className="w-full py-2">

apps/deploy-web/src/components/wallet/ManagedWalletPopup.tsx

@@ -4,6 +4,7 @@ import { Button, Separator } from "@akashnetwork/ui/components";
 import { CoinsSwap, HandCard } from "iconoir-react";
 
 import { TopUpAmountPicker } from "@src/components/top-up-amount-picker/TopUpAmountPicker";
+import { useSelectedChain } from "@src/context/CustomChainProvider";
 import { useWallet } from "@src/context/WalletProvider";
 import { useLoginRequiredEventHandler } from "@src/hooks/useLoginRequiredEventHandler";
 import { useManagedEscrowFaqModal } from "@src/hooks/useManagedEscrowFaqModal";
@@ -18,9 +19,10 @@ interface ManagedWalletPopupProps extends React.PropsWithChildren {
 
 export const ManagedWalletPopup: React.FC<ManagedWalletPopupProps> = ({ walletBalance }) => {
   const isEmailVerified = useIsEmailVerified();
-  const { switchWalletType, isManaged, isTrialing } = useWallet();
+  const { isManaged, isTrialing, switchWalletType } = useWallet();
   const whenLoggedIn = useLoginRequiredEventHandler();
   const { showManagedEscrowFaqModal } = useManagedEscrowFaqModal();
+  const { connect, isWalletConnected } = useSelectedChain();
 
   const goToCheckout = () => {
     window.location.href = "/api/proxy/v1/checkout";
@@ -77,12 +79,17 @@ export const ManagedWalletPopup: React.FC<ManagedWalletPopupProps> = ({ walletBa
         <VerifyEmail />
 
         <TopUpAmountPicker mdMode="click" className="w-full">
-          <Button onClick={whenLoggedIn(goToCheckout, "Sign In or Sign Up to add funds")} variant="outline" className="w-full space-x-2" disabled={!isEmailVerified}>
+          <Button
+            onClick={whenLoggedIn(goToCheckout, "Sign In or Sign Up to add funds")}
+            variant="outline"
+            className="w-full space-x-2"
+            disabled={!isEmailVerified}
+          >
             <HandCard />
             <span>Add Funds</span>
           </Button>
         </TopUpAmountPicker>
-        <Button onClick={switchWalletType} variant="outline" className="w-full space-x-2">
+        <Button onClick={isWalletConnected ? switchWalletType : connect} variant="outline" className="w-full space-x-2">
           <CoinsSwap />
           <span>Switch to Wallet Payments</span>
         </Button>

apps/deploy-web/src/context/CustomChainProvider/CustomChainProvider.tsx

@@ -69,11 +69,17 @@ export function useSelectedChain() {
 }
 
 const ModalWrapper = (props: WalletModalProps) => {
-  const [, setIsWalletModalOpen] = useAtom(walletStore.isWalletModalOpen);
+  const { isWalletConnected } = useSelectedChain();
+  const [isWalletModalOpen, setIsWalletModalOpen] = useAtom(walletStore.isWalletModalOpen);
+  const [, setSelectedWalletType] = useAtom(walletStore.selectedWalletType);
 
   useEffect(() => {
     setIsWalletModalOpen(props.isOpen);
-  }, [props.isOpen]);
+
+    if (isWalletModalOpen && !props.isOpen && isWalletConnected) {
+      setSelectedWalletType("custodial");
+    }
+  }, [isWalletModalOpen, props.isOpen, isWalletConnected]);
 
   return <DefaultModal {...props} isOpen={props.isOpen} setOpen={props.setOpen} />;
 };

apps/deploy-web/src/context/WalletProvider/WalletProvider.tsx

@@ -24,7 +24,7 @@ import networkStore from "@src/store/networkStore";
 import walletStore from "@src/store/walletStore";
 import { AnalyticsCategory, AnalyticsEvents } from "@src/types/analytics";
 import { UrlService } from "@src/utils/urlUtils";
-import { getSelectedStorageWallet, getStorageWallets, updateStorageManagedWallet, updateStorageWallets } from "@src/utils/walletUtils";
+import { getStorageWallets, updateStorageManagedWallet, updateStorageWallets } from "@src/utils/walletUtils";
 import { useSelectedChain } from "../CustomChainProvider";
 import { useSettings } from "../SettingsProvider";
 
@@ -66,8 +66,6 @@ const MESSAGE_STATES: Record<string, LoadingState> = {
   "/akash.deployment.v1beta3.MsgDepositDeployment": "depositingDeployment"
 };
 
-const initialWallet = getSelectedStorageWallet();
-
 export const WalletProvider = ({ children }) => {
   const [isWalletLoaded, setIsWalletLoaded] = useState<boolean>(true);
   const [loadingState, setLoadingState] = useState<LoadingState | undefined>(undefined);
@@ -76,11 +74,9 @@ export const WalletProvider = ({ children }) => {
   const { settings } = useSettings();
   const user = useUser();
   const userWallet = useSelectedChain();
-  const { wallet: managedWallet, isLoading, create: createManagedWallet } = useManagedWallet();
-  const [isWalletModalOpen, setIsWalletModelOpen] = useAtom(walletStore.isWalletModalOpen);
-  const [selectedWalletType, selectWalletType] = useState<"managed" | "custodial">(
-    initialWallet?.selected && initialWallet?.isManaged ? "managed" : "custodial"
-  );
+  const { wallet: managedWallet, isLoading: isManagedWalletLoading, create: createManagedWallet } = useManagedWallet();
+  const [, setIsWalletModelOpen] = useAtom(walletStore.isWalletModalOpen);
+  const [selectedWalletType, setSelectedWalletType] = useAtom(walletStore.selectedWalletType);
   const {
     address: walletAddress,
     username,
@@ -93,11 +89,9 @@ export const WalletProvider = ({ children }) => {
     fee: { default: feeGranter }
   } = useAllowance(walletAddress as string, isManaged);
   const [selectedNetworkId, setSelectedNetworkId] = networkStore.useSelectedNetworkIdStore({ reloadOnChange: true });
-  const shouldAutoConnectManagedWallet =
-    !isWalletConnected && selectedWalletType === "custodial" && !!managedWallet && !isWalletModalOpen && !userWallet.isWalletConnecting;
+  const isLoading = (selectedWalletType === "managed" && isManagedWalletLoading) || (selectedWalletType === "custodial" && userWallet.isWalletConnecting);
 
   useWhen(walletAddress, loadWallet);
-  useWhen(shouldAutoConnectManagedWallet, switchWalletType);
 
   useEffect(() => {
     if (!settings.apiEndpoint || !settings.rpcEndpoint) return;
@@ -126,14 +120,14 @@ export const WalletProvider = ({ children }) => {
       });
     }
 
-    selectWalletType(prev => (prev === "custodial" ? "managed" : "custodial"));
+    setSelectedWalletType(prev => (prev === "custodial" ? "managed" : "custodial"));
   }
 
   function connectManagedWallet() {
     if (!managedWallet) {
       createManagedWallet();
     }
-    selectWalletType("managed");
+    setSelectedWalletType("managed");
   }
 
   function logout() {
@@ -145,6 +139,10 @@ export const WalletProvider = ({ children }) => {
     });
 
     router.push(UrlService.home());
+
+    if (managedWallet) {
+      setSelectedWalletType("managed");
+    }
   }
 
   async function connectWallet() {

apps/deploy-web/src/hooks/useManagedWallet.ts

@@ -2,6 +2,7 @@ import { useEffect, useMemo } from "react";
 import { useAtom } from "jotai";
 
 import { browserEnvConfig } from "@src/config/browser-env.config";
+import { useSelectedChain } from "@src/context/CustomChainProvider";
 import { useUser } from "@src/hooks/useUser";
 import { useCreateManagedWalletMutation, useManagedWalletQuery } from "@src/queries/useManagedWalletQuery";
 import walletStore from "@src/store/walletStore";
@@ -14,7 +15,20 @@ const isBillingEnabled = NEXT_PUBLIC_BILLING_ENABLED;
 export const useManagedWallet = () => {
   const user = useUser();
   const { user: signedInUser } = useCustomUser();
-  const { data: queried, isFetched, isLoading: isFetching, refetch } = useManagedWalletQuery(isBillingEnabled ? user?.id : undefined);
+  const userWallet = useSelectedChain();
+  const [selectedWalletType, setSelectedWalletType] = useAtom(walletStore.selectedWalletType);
+  const {
+    data: queried,
+    isFetched,
+    isLoading: isFetching,
+    refetch
+  } = useManagedWalletQuery(isBillingEnabled ? user?.id : undefined, {
+    onSuccess: wallet => {
+      if (selectedWalletType === "custodial" && wallet && !userWallet.isWalletConnected && !userWallet.isWalletConnecting) {
+        setSelectedWalletType("managed");
+      }
+    }
+  });
   const { mutate: create, data: created, isLoading: isCreating, isSuccess: isCreated } = useCreateManagedWalletMutation();
   const wallet = useMemo(() => queried || created, [queried, created]);
   const isLoading = isFetching || isCreating;

apps/deploy-web/src/queries/useManagedWalletQuery.ts

@@ -1,10 +1,11 @@
-import { useMutation, useQuery, useQueryClient } from "react-query";
+import { QueryKey, useMutation, useQuery, useQueryClient, UseQueryOptions } from "react-query";
+import { ApiWalletOutput } from "@akashnetwork/http-sdk";
 
 import { managedWalletHttpService } from "@src/services/managed-wallet-http/managed-wallet-http.service";
 
 const MANAGED_WALLET = "MANAGED_WALLET";
 
-export function useManagedWalletQuery(userId?: string) {
+export function useManagedWalletQuery(userId?: string, options?: Omit<UseQueryOptions<ApiWalletOutput, Error, any, QueryKey>, "queryKey" | "queryFn">) {
   return useQuery(
     [MANAGED_WALLET, userId],
     async () => {
@@ -14,7 +15,8 @@ export function useManagedWalletQuery(userId?: string) {
     },
     {
       enabled: !!userId,
-      staleTime: Infinity
+      staleTime: Infinity,
+      ...options
     }
   );
 }

apps/deploy-web/src/store/walletStore.ts

@@ -2,9 +2,11 @@ import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 
 const isSignedInWithTrial = atomWithStorage<boolean>("isSignedInWithTrial", false);
+const selectedWalletType = atomWithStorage<"managed" | "custodial">("selectedWalletType", "custodial");
 const isWalletModalOpen = atom<boolean>(false);
 
 export default {
   isSignedInWithTrial,
+  selectedWalletType,
   isWalletModalOpen
 };