Merge pull request #73 from Cardinal-Cryptography/sd-78-fix-modal-displayed-on-disconnect-click

SD-78: Fix modal displayed on disconnect click

Author: rAskVAL

src/domains/chains/components/WalletProvider.tsx

@@ -3,13 +3,12 @@ import { useQueryClient } from '@tanstack/react-query';
 import {
   createContext,
   useContext,
-  ReactNode,
+  ReactNode, useEffect,
 } from 'react';
 import { Address } from 'viem';
 import { useAccount } from 'wagmi';
 
-import useConnectedChainNetworkEnvironment from 'src/domains/chains/utils/useConnectedChainNetworkEnvironment';
-import getQueryKey from 'src/domains/misc/utils/getQueryKey';
+import { QUERY_KEYS } from 'src/domains/misc/utils/getQueryKey';
 import { clearShielderIndexedDB } from 'src/domains/shielder/stores/getShielderIndexedDB';
 import useShielderStore from 'src/domains/shielder/stores/shielder';
 import useShielderPrivateKey from 'src/domains/shielder/utils/useShielderPrivateKey';
@@ -27,37 +26,36 @@ const WalletContext = createContext<WalletContextType | undefined>(undefined);
 
 const WalletProvider = ({ children }: { children: ReactNode }) => {
   const { open: openModal } = useAppKit();
-  const { address, isConnected, status } = useAccount();
+  const { address, isReconnecting, isConnecting, isConnected, isDisconnected } = useAccount();
   const { disconnect } = useDisconnect();
   const { shielderPrivateKey } = useShielderPrivateKey(address);
   const clearShielderPrivateKeys = useShielderStore(store => store.clearShielderPrivateKeys);
 
   const queryClient = useQueryClient();
 
-  const networkEnvironment = useConnectedChainNetworkEnvironment();
+  useEffect(() => {
+    if (!isDisconnected) return;
+
+    queryClient.removeQueries({
+      predicate: query =>
+        query.queryKey[0] === QUERY_KEYS.shielderPrivateKey,
+    });
 
-  const handleDisconnect = async () => {
-    if(address && networkEnvironment) {
-      queryClient.removeQueries({
-        queryKey: getQueryKey.shielderPrivateKey(address, networkEnvironment),
-      });
-    }
     void queryClient.removeQueries({
       predicate: query =>
-        query.queryKey[0] === 'shielder-client',
+        query.queryKey[0] === QUERY_KEYS.shielderClient,
     });
 
     clearShielderPrivateKeys();
-    await clearShielderIndexedDB();
-    await disconnect();
-  };
+    void clearShielderIndexedDB();
+  }, [address, queryClient, clearShielderPrivateKeys, isDisconnected]);
 
   const value = {
-    disconnect: handleDisconnect,
+    disconnect,
     openModal,
     address,
     isConnected,
-    isLoading: status !== 'connected' && status !== 'disconnected',
+    isLoading: isReconnecting || isConnecting,
     privateKey: shielderPrivateKey,
   };
 

src/domains/misc/utils/getQueryKey.ts

@@ -1,36 +1,91 @@
+export const QUERY_KEYS = {
+  shielderFees: 'shielderFees',
+  shielderClient: 'shielder-client',
+  shielderTransactions: 'shielder-transactions',
+  tokenDecimals: 'token-decimals',
+  tokenName: 'token-name',
+  tokenSymbol: 'token-symbol',
+  tokenPublicBalance: 'token-public-balance',
+  tokenShieldedBalance: 'token-shielded-balance',
+  shielderPrivateKey: 'shielderPrivateKey',
+  wasmCryptoClient: 'wasm-crypto-client',
+} as const;
+
 import type { Address } from 'viem';
 
 import { NetworkEnvironment } from 'src/domains/chains/types/misc';
 
 const getQueryKey = {
-  shielderFees: (walletAddress: Address, chainId: string) => ['shielderFees', walletAddress, chainId] ,
+  shielderFees: (walletAddress: Address, chainId: string) => [
+    QUERY_KEYS.shielderFees,
+    walletAddress,
+    chainId,
+  ],
   shielderClient: (chainId: number, shielderPrivateKey: Address) => [
-    'shielder-client',
+    QUERY_KEYS.shielderClient,
     chainId,
     shielderPrivateKey,
-  ] ,
+  ],
   shielderTransactions: (accountAddress: Address, chainId: number) => [
-    'shielder-transactions',
+    QUERY_KEYS.shielderTransactions,
     accountAddress,
     chainId,
-  ] ,
-  tokenDecimals: (tokenAddress: Address | 'native', chainId: string) => ['token-decimals', chainId, tokenAddress],
-  tokenName: (tokenAddress: Address | 'native', chainId: string) => ['token-name', chainId, tokenAddress] ,
-  tokenSymbol: (tokenAddress: Address | 'native', chainId: string) => ['token-symbol', chainId, tokenAddress] ,
-  tokenPublicBalance: (tokenAddress: Address | 'native', chainId: string, walletAddress: Address) => [
-    'token-public-balance',
+  ],
+  tokenDecimals: (
+    tokenAddress: Address | 'native',
+    chainId: string
+  ) => [
+    QUERY_KEYS.tokenDecimals,
+    chainId,
+    tokenAddress,
+  ],
+  tokenName: (
+    tokenAddress: Address | 'native',
+    chainId: string
+  ) => [
+    QUERY_KEYS.tokenName,
+    chainId,
+    tokenAddress,
+  ],
+  tokenSymbol: (
+    tokenAddress: Address | 'native',
+    chainId: string
+  ) => [
+    QUERY_KEYS.tokenSymbol,
+    chainId,
+    tokenAddress,
+  ],
+  tokenPublicBalance: (
+    tokenAddress: Address | 'native',
+    chainId: string,
+    walletAddress: Address
+  ) => [
+    QUERY_KEYS.tokenPublicBalance,
     tokenAddress,
     chainId,
     walletAddress,
-  ] ,
-  tokenShieldedBalance: (tokenAddress: Address | 'native', chainId: string, walletAddress: Address) => [
-    'token-shielded-balance',
+  ],
+  tokenShieldedBalance: (
+    tokenAddress: Address | 'native',
+    chainId: string,
+    walletAddress: Address
+  ) => [
+    QUERY_KEYS.tokenShieldedBalance,
     tokenAddress,
     chainId,
     walletAddress,
-  ] ,
-  shielderPrivateKey: (walletAddress: Address, networkEnvironment: NetworkEnvironment) => ['shielderPrivateKey', networkEnvironment, walletAddress] ,
-  wasmCryptoClient: () => ['wasm-crypto-client'] ,
+  ],
+  shielderPrivateKey: (
+    walletAddress: Address,
+    networkEnvironment: NetworkEnvironment
+  ) => [
+    QUERY_KEYS.shielderPrivateKey,
+    networkEnvironment,
+    walletAddress,
+  ],
+  wasmCryptoClient: () => [
+    QUERY_KEYS.wasmCryptoClient,
+  ],
 };
 
 export default getQueryKey;