From adcf720b1c14954aa23986c8d3be1c24e7cbd410 Mon Sep 17 00:00:00 2001
From: quangdz1704 <Ntq.1704@gmail.com>
Date: Mon, 22 Jul 2024 16:34:40 +0700
Subject: [PATCH] fix: reset balance when dissconnect

---
 components/layout/connectButton/index.tsx | 32 ++++++++---------------
 stores/token/useTokenStore.ts             | 12 +++++++++
 2 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/components/layout/connectButton/index.tsx b/components/layout/connectButton/index.tsx
index a6afda7..4bb0fb2 100644
--- a/components/layout/connectButton/index.tsx
+++ b/components/layout/connectButton/index.tsx
@@ -33,6 +33,7 @@ import classNames from "classnames";
 import { FC, useCallback, useEffect, useRef, useState } from "react";
 import ConnectedInfo from "../connectedInfo";
 import styles from "./index.module.scss";
+import { useTokenActions } from "@/stores/token/selector";
 
 export type ConnectStatus =
   | "init"
@@ -55,6 +56,8 @@ const ConnectButton: FC<{ fullWidth?: boolean }> = ({ fullWidth }) => {
     handleSetTonAddress,
     handleSetTonWallet,
   } = useAuthenticationActions();
+  const { handleResetAmountsCache, handleResetTonAmountsCache } =
+    useTokenActions();
   const [connectStatus, setConnectStatus] = useState<
     OraiWallet | TonWallet | "init"
   >("init");
@@ -100,20 +103,9 @@ const ConnectButton: FC<{ fullWidth?: boolean }> = ({ fullWidth }) => {
     try {
       setConnectStatus(walletType);
 
-      const walletsList = await connector.getWallets(); // or use `walletsList` fetched before
+      const _walletsList = await connector.getWallets(); // or use `walletsList` fetched before
 
-      const embeddedWallet = walletsList.find(
-        isWalletInfoCurrentlyEmbedded
-      ) as WalletInfoCurrentlyEmbedded;
-
-      // if (embeddedWallet) {
-      //   connector.connect({ jsBridgeKey: embeddedWallet.jsBridgeKey });
-      //   return;
-      // }
-
-      // await connector.disconnect();
-
-      const addressConnected = connector.connect({
+      const _addressConnected = connector.connect({
         jsBridgeKey: walletType || "tonkeeper",
       });
       console.log("addressConnected", connector.account);
@@ -128,8 +120,10 @@ const ConnectButton: FC<{ fullWidth?: boolean }> = ({ fullWidth }) => {
 
   const handleDisconnectOraichain = (walletType: OraiWallet) => {
     if (oraiAddress && walletType === oraiWallet) {
-      handleSetOraiAddress({ oraiAddress: undefined }),
-        handleSetOraiWallet({ oraiWallet: undefined });
+      handleSetOraiAddress({ oraiAddress: undefined });
+      handleSetOraiWallet({ oraiWallet: undefined });
+      handleResetAmountsCache();
+      handleResetTonAmountsCache();
     }
   };
 
@@ -142,6 +136,8 @@ const ConnectButton: FC<{ fullWidth?: boolean }> = ({ fullWidth }) => {
       if (tonAddress && walletType === tonWallet) {
         handleSetTonAddress({ tonAddress: undefined });
         handleSetTonWallet({ tonWallet: undefined });
+        handleResetTonAmountsCache();
+        handleResetAmountsCache();
       }
     } catch (error) {
       console.log("error disconnect TON :>>", error);
@@ -181,12 +177,6 @@ const ConnectButton: FC<{ fullWidth?: boolean }> = ({ fullWidth }) => {
     }
   }, [tonAddress, tonWallet]);
 
-  // useEffect(() => {
-  //   if (open && oraiAddress) {
-  //     setStep(2);
-  //   }
-  // }, [open]);
-
   // @ts-ignore
   const isCheckOwallet =
     typeof window !== "undefined" && window?.owallet?.["isOwallet"];
diff --git a/stores/token/useTokenStore.ts b/stores/token/useTokenStore.ts
index 7e23869..32f10a7 100644
--- a/stores/token/useTokenStore.ts
+++ b/stores/token/useTokenStore.ts
@@ -13,6 +13,8 @@ export interface TokenActions {
   handleSetPricesCache: (prices: CoinGeckoPrices<string>) => void;
   handleSetAmountsCache: (amounts: AmountDetails) => void;
   handleSetTonAmountsCache: (amountsTon: AmountDetails) => void;
+  handleResetAmountsCache: () => void;
+  handleResetTonAmountsCache: () => void;
 }
 
 const initialState: IToken = {
@@ -47,6 +49,16 @@ const useTokenStore = create<IToken & { actions: TokenActions }>()(
               ...amountsTon,
             };
           }),
+        handleResetTonAmountsCache() {
+          set((state) => {
+            state.amountsTon = {};
+          });
+        },
+        handleResetAmountsCache() {
+          set((state) => {
+            state.amounts = {};
+          });
+        },
       },
     })),
     {