feat(mfi-v2-ui):store login provider

apps/marginfi-v2-ui/src/components/common/Wallet/WalletButton.tsx

@@ -1,25 +1,96 @@
-import React from "react";
+import React, { useCallback, useMemo } from "react";
+import Image from "next/image";
+import { useWallet } from "@solana/wallet-adapter-react";
 
 import { useUiStore } from "~/store";
-import { useWalletContext } from "~/hooks/useWalletContext";
+import { WalletInfo, Web3AuthProvider, useWalletContext } from "~/hooks/useWalletContext";
 
 import { Wallet } from "~/components/common/Wallet";
-
+import { IconChevronDown, IconBrandGoogle, IconBrandX, IconBrandApple, IconMrgn } from "~/components/ui/icons";
 import { Button } from "~/components/ui/button";
 
+const web3AuthIconMap: { [key in Web3AuthProvider]: { icon: JSX.Element } } = {
+  google: {
+    icon: <IconBrandGoogle size={20} />,
+  },
+  twitter: {
+    icon: <IconBrandX size={20} />,
+  },
+  apple: {
+    icon: <IconBrandApple size={20} />,
+  },
+  email_passwordless: {
+    icon: <IconMrgn size={20} />,
+  },
+};
+
 export const WalletButton = () => {
-  const { connected } = useWalletContext();
+  const { select } = useWallet();
+  const { connected, connecting, isLoading, loginWeb3Auth } = useWalletContext();
   const [setIsWalletAuthDialogOpen] = useUiStore((state) => [state.setIsWalletAuthDialogOpen]);
 
+  const walletInfo = useMemo(
+    () => (connected ? null : (JSON.parse(localStorage.getItem("walletInfo") ?? "null") as WalletInfo | null)),
+    [connected]
+  );
+
+  const WalletIcon = useMemo(() => {
+    if (walletInfo?.icon) {
+      const iconSrc = walletInfo?.icon;
+      return function WalletIconComp() {
+        return <Image src={iconSrc} alt="wallet_icon" width={20} height={20} />;
+      };
+    } else {
+      return function WalletIconComp() {
+        return walletInfo ? (
+          web3AuthIconMap[walletInfo.name as Web3AuthProvider]?.icon || <IconMrgn size={20} />
+        ) : (
+          <IconMrgn size={20} />
+        );
+      };
+    }
+  }, [walletInfo]);
+
+  const handleWalletConnect = useCallback(() => {
+    try {
+      if (!walletInfo) throw new Error("No local storage");
+      if (walletInfo.web3Auth) {
+        if (walletInfo.name === "email_passwordless") {
+          loginWeb3Auth("email_passwordless", { login_hint: walletInfo.email });
+        } else {
+          loginWeb3Auth(walletInfo.name);
+        }
+      } else {
+        select(walletInfo.name as any);
+      }
+    } catch (error) {
+      setIsWalletAuthDialogOpen(true);
+    }
+  }, [walletInfo, setIsWalletAuthDialogOpen]);
+
   return (
     <>
-      {!connected ? (
-        <Button onClick={() => setIsWalletAuthDialogOpen(true)} className="gap-1.5">
-          Sign In
-        </Button>
-      ) : (
-        <Wallet />
-      )}
+      {!isLoading &&
+        (!connected ? (
+          <Button className="gap-1.5 py-0 pr-2">
+            <div className="flex flex-row relative h-full gap-4">
+              <div onClick={() => handleWalletConnect()} className="inline-flex items-center gap-2">
+                Sign in with
+                {walletInfo && <WalletIcon />}
+              </div>
+              {walletInfo && (
+                <div
+                  onClick={() => setIsWalletAuthDialogOpen(true)}
+                  className="pl-2 border-l border-border inline-flex items-center"
+                >
+                  <IconChevronDown size={20} />
+                </div>
+              )}
+            </div>
+          </Button>
+        ) : (
+          <Wallet />
+        ))}
     </>
   );
 };

apps/marginfi-v2-ui/src/hooks/useWalletContext.tsx

@@ -13,6 +13,8 @@ import { SolanaWallet, SolanaPrivateKeyProvider } from "@web3auth/solana-provide
 import base58 from "bs58";
 import { Wallet } from "@mrgnlabs/mrgn-common";
 
+import { showErrorToast } from "~/utils/toastUtils";
+
 // wallet adapter context type to override with web3auth data
 // this allows us to pass web3auth wallet to 3rd party services that expect wallet adapter
 type WalletContextOverride = {
@@ -43,6 +45,7 @@ type WalletContextProps = {
   walletAddress: PublicKey;
   walletContextState: WalletContextStateOverride | WalletContextState;
   isOverride: boolean;
+  isLoading: boolean;
   loginWeb3Auth: (
     provider: string,
     extraLoginOptions?: Partial<{
@@ -58,6 +61,12 @@ type WalletContextProps = {
 
 type Web3AuthSocialProvider = "google" | "twitter" | "apple";
 type Web3AuthProvider = "email_passwordless" | Web3AuthSocialProvider;
+type WalletInfo = {
+  name: string;
+  web3Auth: boolean;
+  icon?: string;
+  email?: string;
+};
 
 const web3AuthChainConfig = {
   chainNamespace: CHAIN_NAMESPACES.SOLANA,
@@ -110,7 +119,7 @@ const makeweb3AuthWalletContextState = (wallet: Wallet): WalletContextStateOverr
 const WalletContext = React.createContext<WalletContextProps | undefined>(undefined);
 
 const WalletProvider = ({ children }: { children: React.ReactNode }) => {
-  const { query } = useRouter();
+  const { query, asPath, replace } = useRouter();
   const [web3AuthPkCookie, setWeb3AuthPkCookie] = useCookies(["mrgnPrivateKeyRequested"]);
 
   // default wallet adapter context, overwritten when web3auth is connected
@@ -123,12 +132,21 @@ const WalletProvider = ({ children }: { children: React.ReactNode }) => {
   const [web3AuthLoginType, setWeb3AuthLoginType] = React.useState<string>("");
   const [web3AuthPk, setWeb3AuthPk] = React.useState<string>("");
   const [web3AuthEmail, setWeb3AuthEmail] = React.useState<string>("");
+  const [isLoading, setIsLoading] = React.useState<boolean>(false);
 
   // if web3auth is connected, override wallet adapter context, otherwise use default
   const walletContextState = React.useMemo(() => {
     if (web3Auth?.connected && web3AuthWalletData) {
       return makeweb3AuthWalletContextState(web3AuthWalletData);
     } else {
+      if (walletContextStateDefault.wallet) {
+        const walletInfo: WalletInfo = {
+          name: walletContextStateDefault.wallet.adapter.name,
+          icon: walletContextStateDefault.wallet.adapter.icon,
+          web3Auth: false,
+        };
+        localStorage.setItem("walletInfo", JSON.stringify(walletInfo));
+      }
       return walletContextStateDefault;
     }
   }, [web3Auth?.connected, web3AuthWalletData, walletContextStateDefault.connected]);
@@ -263,7 +281,7 @@ const WalletProvider = ({ children }: { children: React.ReactNode }) => {
   const loginWeb3Auth = React.useCallback(
     async (provider: string, extraLoginOptions: any = {}, cb?: () => void) => {
       if (!web3Auth) {
-        toast.error("Error connecting to web3Auth");
+        showErrorToast("marginfi account not ready.");
         return;
       }
 
@@ -272,6 +290,13 @@ const WalletProvider = ({ children }: { children: React.ReactNode }) => {
           loginProvider: provider,
           extraLoginOptions,
         });
+
+        const walletInfo: WalletInfo = {
+          name: provider!,
+          web3Auth: true,
+          email: extraLoginOptions.login_hint,
+        };
+        localStorage.setItem("walletInfo", JSON.stringify(walletInfo));
       } catch (error) {
         console.error(error);
       }
@@ -285,15 +310,24 @@ const WalletProvider = ({ children }: { children: React.ReactNode }) => {
       await web3Auth.logout();
       setWeb3AuthWalletData(undefined);
     } else {
-      walletContextState?.disconnect();
+      await walletContextState?.disconnect();
+    }
+
+    if (asPath.includes("#")) {
+      // Remove the hash and update the URL
+      const newUrl = asPath.split("#")[0];
+      replace(newUrl);
     }
+    setIsLoading(false);
     setPfp("");
   }, [walletContextState, web3Auth?.connected, walletContextStateDefault]);
 
   // if web3auth is connected, fetch wallet data
   React.useEffect(() => {
     if (!web3Auth?.connected || !web3Auth?.provider || web3AuthWalletData) return;
+    setIsLoading(true);
     makeweb3AuthWalletData(web3Auth.provider);
+    setIsLoading(false);
   }, [web3Auth?.connected, web3Auth?.provider, web3AuthWalletData]);
 
   // initialize web3auth sdk on page load
@@ -324,22 +358,25 @@ const WalletProvider = ({ children }: { children: React.ReactNode }) => {
         setweb3Auth(web3AuthInstance);
       } catch (error) {
         console.error(error);
+      } finally {
+        setIsLoading(false);
       }
     };
-
+    setIsLoading(true);
     init();
   }, []);
 
   return (
     <WalletContext.Provider
       value={{
         connecting: walletContextState?.connecting,
-        connected: Boolean(walletContextState?.connected || web3Auth?.connected),
+        connected: Boolean(walletContextState?.connected),
         web3AuthConncected: web3Auth?.connected,
         wallet,
         walletAddress: wallet?.publicKey as PublicKey,
         walletContextState,
         isOverride,
+        isLoading,
         loginWeb3Auth,
         logout,
         requestPrivateKey,
@@ -360,5 +397,5 @@ const useWalletContext = () => {
   return context;
 };
 
-export type { WalletContextStateOverride, Web3AuthSocialProvider, Web3AuthProvider };
+export type { WalletContextStateOverride, Web3AuthSocialProvider, Web3AuthProvider, WalletInfo };
 export { WalletProvider, useWalletContext };