CP - 13365 - fix for imported PK loading balance issues (#3581)

alexnicolae-ava · web-flow · commit eb416a905776 · 2026-02-13T22:35:00.000+02:00
diff --git a/packages/core-mobile/app/new/features/portfolio/hooks/useAllBalances.ts b/packages/core-mobile/app/new/features/portfolio/hooks/useAllBalances.ts
@@ -1,7 +1,7 @@
 import { QueryObserverResult } from '@tanstack/react-query'
 import { useSelector } from 'react-redux'
 import { AdjustedNormalizedBalancesForAccounts } from 'services/balance/types'
-import { selectAccounts } from 'store/account'
+import { selectAccountsArray } from 'store/account'
 import { useAccountsBalances } from './useAccountsBalances'
 
 /**
@@ -20,7 +20,11 @@ export function useAllBalances(options?: {
     QueryObserverResult<AdjustedNormalizedBalancesForAccounts, Error>
   >
 } {
-  const allAccounts = useSelector(selectAccounts)
+  // Use the memoized selector so the array reference is stable between
+  // renders.  Object.values(selectAccounts) created a new array every render,
+  // which made the useEffect in useAccountsBalances fire on every render,
+  // flooding queryClient.setQueryData and causing a render storm / freeze.
+  const allAccounts = useSelector(selectAccountsArray)
 
-  return useAccountsBalances(Object.values(allAccounts), options)
+  return useAccountsBalances(allAccounts, options)
 }
diff --git a/packages/core-mobile/app/new/features/portfolio/hooks/useIsLoadingBalancesForWallet.ts b/packages/core-mobile/app/new/features/portfolio/hooks/useIsLoadingBalancesForWallet.ts
@@ -1,17 +1,27 @@
 import { useSelector } from 'react-redux'
-import { selectAccountsByWalletId } from 'store/account'
+import { selectAccountsByWalletId, selectImportedAccounts } from 'store/account'
 import { RootState } from 'store/types'
 import { Wallet } from 'store/wallet/types'
+import { IMPORTED_ACCOUNTS_VIRTUAL_WALLET_ID } from 'features/wallets/consts'
 import { useAllBalances } from './useAllBalances'
 
 /**
  * Returns true if balances are currently loading for the given wallet.
  * Checks if any account within the wallet is missing balance data.
  */
 export const useIsLoadingBalancesForWallet = (wallet?: Wallet): boolean => {
-  const accounts = useSelector((state: RootState) =>
+  const accountsForWallet = useSelector((state: RootState) =>
     selectAccountsByWalletId(state, wallet?.id ?? '')
   )
+  const importedAccounts = useSelector(selectImportedAccounts)
+
+  // The virtual "Imported" wallet groups all private-key accounts under one
+  // card, but those accounts have real wallet UUIDs — not the virtual ID.
+  const accounts =
+    wallet?.id === IMPORTED_ACCOUNTS_VIRTUAL_WALLET_ID
+      ? importedAccounts
+      : accountsForWallet
+
   const { data } = useAllBalances()
 
   if (!wallet || accounts.length === 0) return true
diff --git a/packages/core-mobile/app/new/features/portfolio/hooks/useWalletBalances.ts b/packages/core-mobile/app/new/features/portfolio/hooks/useWalletBalances.ts
@@ -2,9 +2,10 @@ import { QueryObserverResult } from '@tanstack/react-query'
 import { useMemo } from 'react'
 import { useSelector } from 'react-redux'
 import { AdjustedNormalizedBalancesForAccounts } from 'services/balance/types'
-import { selectAccountsByWalletId } from 'store/account'
+import { selectAccountsByWalletId, selectImportedAccounts } from 'store/account'
 import { RootState } from 'store/types'
 import { Wallet } from 'store/wallet/types'
+import { IMPORTED_ACCOUNTS_VIRTUAL_WALLET_ID } from 'new/features/wallets/consts'
 import { useAllBalances } from './useAllBalances'
 
 /**
@@ -25,9 +26,19 @@ export const useWalletBalances = (
     QueryObserverResult<AdjustedNormalizedBalancesForAccounts, Error>
   >
 } => {
-  const accounts = useSelector((state: RootState) =>
+  const isVirtualImportedWallet =
+    wallet?.id === IMPORTED_ACCOUNTS_VIRTUAL_WALLET_ID
+
+  // For the virtual "Imported" wallet, imported accounts have their own
+  // real walletId — not the virtual ID — so selectAccountsByWalletId
+  // would return an empty array. Fall back to selectImportedAccounts.
+  const accountsByWalletId = useSelector((state: RootState) =>
     selectAccountsByWalletId(state, wallet?.id ?? '')
   )
+  const importedAccounts = useSelector(selectImportedAccounts)
+  const accounts = isVirtualImportedWallet
+    ? importedAccounts
+    : accountsByWalletId
 
   const {
     data: allAccountsBalances,
diff --git a/packages/core-mobile/app/services/balance/BalanceService.test.ts b/packages/core-mobile/app/services/balance/BalanceService.test.ts
@@ -125,6 +125,22 @@ const testAccount2: Account = {
   addressPVM: 'P-avax1bbhxdv3wqxd42rxdalvp2knxs244r06wrxmvlg'
 }
 
+// Imported PK account that shares the same EVM address as testAccount
+// (simulates exporting PK from mnemonic and re-importing it)
+const importedAccount: Account = {
+  name: 'Imported Account',
+  id: 'a1b2c3d4-e5f6-7890-abcd-ef1234567890',
+  walletId: 'imported-wallet-id',
+  index: 0,
+  type: CoreAccountType.IMPORTED as CoreAccountType.IMPORTED,
+  addressC: '0x066b2322a30d7C5838035112F3b816b46D639bBC', // same as testAccount
+  addressCoreEth: '0x066b2322a30d7C5838035112F3b816b46D639bBC',
+  addressBTC: 'bc1qdifferentbtcaddressforimportedpkaccount',
+  addressSVM: 'DifferentSolAddressForImportedPkAccount1234567',
+  addressAVM: 'X-avax1differentavmaddressforimported',
+  addressPVM: 'P-avax1differentpvmaddressforimported'
+}
+
 const cChainNetwork = {
   chainId: 43114,
   chainName: 'Avalanche C-Chain',
@@ -688,6 +704,266 @@ describe('BalanceService', () => {
     })
   })
 
+  describe('duplicate address handling (imported PK accounts)', () => {
+    describe('getBalancesForAccounts', () => {
+      it('should assign balance to both accounts when they share the same EVM address', async () => {
+        // API returns one balance for the shared EVM address
+        const mockStreamResponse = [
+          {
+            caip2Id: 'eip155:43114',
+            networkType: 'evm',
+            id: testAccount.addressC // shared address
+          }
+        ]
+
+        mockGetBalancesStream.mockReturnValue(
+          createAsyncIterator(mockStreamResponse)
+        )
+
+        const mockNormalizedBalance = {
+          accountId: 'PLACEHOLDER',
+          chainId: 43114,
+          tokens: [{ symbol: 'AVAX', balance: 1000n }],
+          dataAccurate: true,
+          error: null
+        }
+
+        // mapBalanceResponseToLegacy is called once per matched account
+        mockMapBalanceResponseToLegacy
+          .mockReturnValueOnce({
+            ...mockNormalizedBalance,
+            accountId: testAccount.id
+          })
+          .mockReturnValueOnce({
+            ...mockNormalizedBalance,
+            accountId: importedAccount.id
+          })
+
+        const onBalanceLoaded = jest.fn()
+
+        const result = await balanceService.getBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount],
+          currency: 'usd',
+          onBalanceLoaded,
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, ['avax1test1', 'avax1test2']],
+            [importedAccount.id, []]
+          ]),
+          xpubByAccountId: new Map()
+        })
+
+        // Both accounts should have received the balance
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(1)
+        expect(result[testAccount.id]?.[0]?.accountId).toBe(testAccount.id)
+        expect(result[importedAccount.id]?.[0]?.accountId).toBe(
+          importedAccount.id
+        )
+        // onBalanceLoaded should fire for both
+        expect(onBalanceLoaded).toHaveBeenCalledTimes(2)
+      })
+
+      it('should match address case-insensitively (API returns lowercase, account stores checksummed)', async () => {
+        // API returns the address in LOWERCASE
+        const lowercaseAddress =
+          testAccount.addressC?.toLowerCase() ?? 'missing'
+
+        const mockStreamResponse = [
+          {
+            caip2Id: 'eip155:43114',
+            networkType: 'evm',
+            id: lowercaseAddress // lowercase from API
+          }
+        ]
+
+        mockGetBalancesStream.mockReturnValue(
+          createAsyncIterator(mockStreamResponse)
+        )
+
+        mockMapBalanceResponseToLegacy
+          .mockReturnValueOnce({
+            accountId: testAccount.id,
+            chainId: 43114,
+            tokens: [{ symbol: 'AVAX', balance: 500n }],
+            dataAccurate: true,
+            error: null
+          })
+          .mockReturnValueOnce({
+            accountId: importedAccount.id,
+            chainId: 43114,
+            tokens: [{ symbol: 'AVAX', balance: 500n }],
+            dataAccurate: true,
+            error: null
+          })
+
+        const result = await balanceService.getBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount],
+          currency: 'usd',
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, ['avax1test1', 'avax1test2']],
+            [importedAccount.id, []]
+          ]),
+          xpubByAccountId: new Map()
+        })
+
+        // Both should match despite case difference
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(1)
+      })
+
+      it('should not duplicate when account id matches directly in accountById', async () => {
+        // If the API returns the account.id directly as the balance id,
+        // it should resolve via accountById (direct lookup) — not the address map.
+        const mockStreamResponse = [
+          {
+            caip2Id: 'eip155:43114',
+            networkType: 'evm',
+            id: testAccount.id // matches account id directly
+          }
+        ]
+
+        mockGetBalancesStream.mockReturnValue(
+          createAsyncIterator(mockStreamResponse)
+        )
+
+        mockMapBalanceResponseToLegacy.mockReturnValue({
+          accountId: testAccount.id,
+          chainId: 43114,
+          tokens: [],
+          dataAccurate: true,
+          error: null
+        })
+
+        const result = await balanceService.getBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount],
+          currency: 'usd',
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, []],
+            [importedAccount.id, []]
+          ]),
+          xpubByAccountId: new Map()
+        })
+
+        // Only testAccount should match (by id), not importedAccount
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(0)
+      })
+    })
+
+    describe('getVMBalancesForAccounts', () => {
+      it('should assign VM balance to both accounts sharing the same EVM address', async () => {
+        const mockModule = {
+          getBalances: jest.fn().mockResolvedValue({
+            [testAccount.addressC!]: {
+              avax: {
+                name: 'Avalanche',
+                symbol: 'AVAX',
+                decimals: 18,
+                balance: '1000000000000000000',
+                type: TokenType.NATIVE
+              }
+            }
+          })
+        }
+        mockLoadModuleByNetwork.mockResolvedValue(mockModule)
+
+        const onBalanceLoaded = jest.fn()
+
+        const result = await balanceService.getVMBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount],
+          currency: 'usd',
+          customTokens: {},
+          onBalanceLoaded,
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, ['avax1test1', 'avax1test2']],
+            [importedAccount.id, []]
+          ])
+        })
+
+        // Both accounts should have received the VM balance
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(1)
+        expect(result[testAccount.id]?.[0]?.accountId).toBe(testAccount.id)
+        expect(result[importedAccount.id]?.[0]?.accountId).toBe(
+          importedAccount.id
+        )
+      })
+
+      it('should handle VM error for both accounts sharing an address', async () => {
+        const mockModule = {
+          getBalances: jest.fn().mockResolvedValue({
+            [testAccount.addressC!]: {
+              error: new Error('VM fetch failed')
+            }
+          })
+        }
+        mockLoadModuleByNetwork.mockResolvedValue(mockModule)
+
+        const result = await balanceService.getVMBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount],
+          currency: 'usd',
+          customTokens: {},
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, ['avax1test1', 'avax1test2']],
+            [importedAccount.id, []]
+          ])
+        })
+
+        // Both accounts should have received an error entry
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(1)
+        expect(result[testAccount.id]?.[0]?.dataAccurate).toBe(false)
+        expect(result[importedAccount.id]?.[0]?.dataAccurate).toBe(false)
+      })
+
+      it('should not add duplicate account entries for the same address', async () => {
+        // Create a second imported account with the exact same address
+        const importedAccount2: Account = {
+          ...importedAccount,
+          id: 'duplicate-imported-id',
+          name: 'Imported Account 2'
+        }
+
+        const mockModule = {
+          getBalances: jest.fn().mockResolvedValue({
+            [testAccount.addressC!]: {
+              avax: {
+                name: 'Avalanche',
+                symbol: 'AVAX',
+                decimals: 18,
+                balance: '1000000000000000000',
+                type: TokenType.NATIVE
+              }
+            }
+          })
+        }
+        mockLoadModuleByNetwork.mockResolvedValue(mockModule)
+
+        const result = await balanceService.getVMBalancesForAccounts({
+          networks: [cChainNetwork as any],
+          accounts: [testAccount, importedAccount, importedAccount2],
+          currency: 'usd',
+          customTokens: {},
+          xpAddressesByAccountId: new Map([
+            [testAccount.id, ['avax1test1', 'avax1test2']],
+            [importedAccount.id, []],
+            [importedAccount2.id, []]
+          ])
+        })
+
+        // All three accounts should get exactly 1 balance entry
+        expect(result[testAccount.id]).toHaveLength(1)
+        expect(result[importedAccount.id]).toHaveLength(1)
+        expect(result[importedAccount2.id]).toHaveLength(1)
+      })
+    })
+  })
+
   describe('edge cases', () => {
     it('should handle empty networks array', async () => {
       const result = await balanceService.getBalancesForAccount({
diff --git a/packages/core-mobile/app/services/balance/BalanceService.ts b/packages/core-mobile/app/services/balance/BalanceService.ts
