fix: account switch on wagmi, multi-chain related params for account …

…controller methods
reown-com · Aug 9, 2024 · b9d0e18 · b9d0e18
1 parent a3aba6f
commit b9d0e18
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 24 deletions.
diff --git a/packages/core/src/controllers/AccountController.ts b/packages/core/src/controllers/AccountController.ts
@@ -129,8 +129,9 @@ export const AccountController = {
       ChainController.setAccountProp('tokenBalance', tokenBalance, chain)
     }
   },
-  setShouldUpdateToAddress(address: string) {
-    ChainController.setAccountProp('shouldUpdateToAddress', address)
+
+  setShouldUpdateToAddress(address: string, chain?: Chain) {
+    ChainController.setAccountProp('shouldUpdateToAddress', address, chain)
   },
 
   setAllAccounts(accounts: AccountType[], chain?: Chain) {
@@ -140,13 +141,13 @@ export const AccountController = {
   addAddressLabel(address: string, label: string) {
     const map = ChainController.getAccountProp('addressLabels') || new Map()
     map.set(address, label)
-    ChainController.setAccountProp('addressLabels', map)
+    ChainController.setAccountProp('addressLabels', map, ChainController.state.activeChain)
   },
 
   removeAddressLabel(address: string) {
     const map = ChainController.getAccountProp('addressLabels') || new Map()
     map.delete(address)
-    ChainController.setAccountProp('addressLabels', map)
+    ChainController.setAccountProp('addressLabels', map, ChainController.state.activeChain)
   },
 
   setConnectedWalletInfo(

diff --git a/packages/ethers/src/client.ts b/packages/ethers/src/client.ts
@@ -811,7 +811,10 @@ export class Web3Modal extends Web3ModalScaffold {
       EthersStoreUtil.setProvider(WalletConnectProvider as unknown as Provider)
       EthersStoreUtil.setStatus('connected')
       EthersStoreUtil.setIsConnected(true)
-      this.setAllAccounts(WalletConnectProvider.accounts.map(address => ({ address, type: 'eoa' })))
+      this.setAllAccounts(
+        WalletConnectProvider.accounts.map(address => ({ address, type: 'eoa' })),
+        this.chain
+      )
       const session = WalletConnectProvider.signer?.session
       for (const address of WalletConnectProvider.accounts) {
         const label = session?.sessionProperties?.[address]
@@ -836,7 +839,10 @@ export class Web3Modal extends Web3ModalScaffold {
         EthersStoreUtil.setProvider(config.injected)
         EthersStoreUtil.setStatus('connected')
         EthersStoreUtil.setIsConnected(true)
-        this.setAllAccounts(addresses.map(address => ({ address, type: 'eoa' })))
+        this.setAllAccounts(
+          addresses.map(address => ({ address, type: 'eoa' })),
+          this.chain
+        )
         this.setAddress(addresses[0])
         this.watchCoinbase(config)
       }
@@ -854,7 +860,10 @@ export class Web3Modal extends Web3ModalScaffold {
         EthersStoreUtil.setProvider(provider)
         EthersStoreUtil.setStatus('connected')
         EthersStoreUtil.setIsConnected(true)
-        this.setAllAccounts(addresses.map(address => ({ address, type: 'eoa' })))
+        this.setAllAccounts(
+          addresses.map(address => ({ address, type: 'eoa' })),
+          this.chain
+        )
         this.setAddress(addresses[0])
         this.watchEIP6963(provider)
       }
@@ -875,7 +884,10 @@ export class Web3Modal extends Web3ModalScaffold {
         EthersStoreUtil.setProvider(config.coinbase)
         EthersStoreUtil.setStatus('connected')
         EthersStoreUtil.setIsConnected(true)
-        this.setAllAccounts(addresses.map(address => ({ address, type: 'eoa' })))
+        this.setAllAccounts(
+          addresses.map(address => ({ address, type: 'eoa' })),
+          this.chain
+        )
         this.setAddress(addresses[0])
         this.watchCoinbase(config)
       }
@@ -903,7 +915,8 @@ export class Web3Modal extends Web3ModalScaffold {
         this.setAllAccounts(
           accounts.length > 0
             ? accounts
-            : [{ address, type: preferredAccountType as 'eoa' | 'smartAccount' }]
+            : [{ address, type: preferredAccountType as 'eoa' | 'smartAccount' }],
+          this.chain
         )
         EthersStoreUtil.setChainId(chainId)
         EthersStoreUtil.setProviderType(ConstantsUtil.AUTH_CONNECTOR_ID as 'w3mAuth')
@@ -1003,9 +1016,12 @@ export class Web3Modal extends Web3ModalScaffold {
       const currentAccount = accounts?.[0]
       if (currentAccount) {
         EthersStoreUtil.setAddress(getOriginalAddress(currentAccount) as Address)
-        this.setAllAccounts(accounts.map(address => ({ address, type: 'eoa' })))
+        this.setAllAccounts(
+          accounts.map(address => ({ address, type: 'eoa' })),
+          this.chain
+        )
       } else {
-        this.setAllAccounts([])
+        this.setAllAccounts([], this.chain)
         localStorage.removeItem(EthersConstantsUtil.WALLET_ID)
         EthersStoreUtil.reset()
       }
@@ -1179,7 +1195,7 @@ export class Web3Modal extends Web3ModalScaffold {
     } else if (!isConnected && this.hasSyncedConnectedAccount) {
       this.resetWcConnection()
       this.resetNetwork()
-      this.setAllAccounts([])
+      this.setAllAccounts([], this.chain)
     }
   }
 

diff --git a/packages/scaffold-ui/src/views/w3m-switch-address-view/index.ts b/packages/scaffold-ui/src/views/w3m-switch-address-view/index.ts
@@ -1,6 +1,7 @@
 import {
   AccountController,
   BlockchainApiController,
+  ChainController,
   ModalController,
   NetworkController,
   OptionsController,
@@ -139,7 +140,7 @@ export class W3mSwitchAddressView extends LitElement {
   }
 
   private onSwitchAddress(address: string) {
-    AccountController.setShouldUpdateToAddress(address)
+    AccountController.setShouldUpdateToAddress(address, ChainController.state.activeChain)
     ModalController.close()
   }
 }

diff --git a/packages/scaffold/src/client.ts b/packages/scaffold/src/client.ts
@@ -192,8 +192,11 @@ export class Web3ModalScaffold {
 
   protected getIsConnectedState = () => AccountController.state.isConnected
 
-  protected setAllAccounts: (typeof AccountController)['setAllAccounts'] = (addresses = []) => {
-    AccountController.setAllAccounts(addresses)
+  protected setAllAccounts: (typeof AccountController)['setAllAccounts'] = (
+    addresses = [],
+    chain
+  ) => {
+    AccountController.setAllAccounts(addresses, chain)
     OptionsController.setHasMultipleAddresses(addresses?.length > 1)
   }
 

diff --git a/packages/wagmi/src/client.ts b/packages/wagmi/src/client.ts
@@ -409,7 +409,9 @@ export class Web3Modal extends Web3ModalScaffold {
               isConnected: true,
               addresses: response.accounts,
               connector,
-              chainId: response.chainId
+              chainId: response.chainId,
+              status: 'connected',
+              resetAccount: false
             })
           )
         }
@@ -460,7 +462,8 @@ export class Web3Modal extends Web3ModalScaffold {
     chainId,
     connector,
     addresses,
-    status
+    status,
+    resetAccount = true
   }: Partial<
     Pick<
       GetAccountReturnType,
@@ -472,7 +475,7 @@ export class Web3Modal extends Web3ModalScaffold {
       | 'addresses'
       | 'status'
     >
-  >) {
+  > & { resetAccount?: boolean }) {
     const caipAddress: CaipAddress = `${ConstantsUtil.EIP155}:${chainId}:${address}`
 
     if (this.getCaipAddress() === caipAddress) {
@@ -482,7 +485,9 @@ export class Web3Modal extends Web3ModalScaffold {
     const connected = isConnected && status === 'connected'
 
     if (connected && address && chainId) {
-      this.resetAccount()
+      if (resetAccount) {
+        this.resetAccount()
+      }
       this.syncNetwork(address, chainId, isConnected)
       this.setIsConnected(isConnected)
       this.setCaipAddress(caipAddress)
@@ -498,15 +503,20 @@ export class Web3Modal extends Web3ModalScaffold {
       // Set by authConnector.onIsConnectedHandler as we need the account type
       const isAuthConnector = connector?.id === ConstantsUtil.AUTH_CONNECTOR_ID
       if (!isAuthConnector && addresses?.length) {
-        this.setAllAccounts(addresses.map(addr => ({ address: addr, type: 'eoa' })))
+        this.setAllAccounts(
+          addresses.map(addr => ({ address: addr, type: 'eoa' })),
+          this.chain
+        )
       }
 
       this.hasSyncedConnectedAccount = true
     } else if (isDisconnected && this.hasSyncedConnectedAccount) {
-      this.resetAccount()
+      if (resetAccount) {
+        this.resetAccount()
+      }
       this.resetWcConnection()
       this.resetNetwork()
-      this.setAllAccounts([])
+      this.setAllAccounts([], this.chain)
 
       this.hasSyncedConnectedAccount = false
     }
@@ -790,7 +800,8 @@ export class Web3Modal extends Web3ModalScaffold {
               address: req.address,
               type: (req.preferredAccountType || 'eoa') as W3mFrameTypes.AccountType
             }
-          ]
+          ],
+          this.chain
         )
       })
 
@@ -807,7 +818,8 @@ export class Web3Modal extends Web3ModalScaffold {
           address: address as `0x${string}`,
           isConnected: true,
           chainId: NetworkUtil.caipNetworkIdToNumber(this.getCaipNetwork()?.id),
-          connector
+          connector,
+          status: 'connected'
         })
       })
     }