From d06997453da3a5cfb57c5d2a620ffd1020dcfca3 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 11:26:08 +0000 Subject: [PATCH 01/10] Return unprocessed ERC-20 tokens for RPC fallback Co-authored-by: Prithpal Sooriya --- eslint-suppressions.json | 3 - packages/assets-controllers/CHANGELOG.md | 1 + .../src/TokenBalancesController.test.ts | 89 +++++++++++++++ .../src/TokenBalancesController.ts | 87 ++++++++++++++ .../src/TokenDetectionController.test.ts | 11 +- .../src/TokenDetectionController.ts | 6 - .../api-balance-fetcher.test.ts | 81 +++++-------- .../api-balance-fetcher.ts | 46 ++++++-- .../rpc-service/rpc-balance-fetcher.test.ts | 60 ++++++++++ .../src/rpc-service/rpc-balance-fetcher.ts | 107 +++++++++++++----- 10 files changed, 386 insertions(+), 105 deletions(-) diff --git a/eslint-suppressions.json b/eslint-suppressions.json index ae9bd334a0e..e80df9da18d 100644 --- a/eslint-suppressions.json +++ b/eslint-suppressions.json @@ -477,9 +477,6 @@ } }, "packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts": { - "@typescript-eslint/explicit-function-return-type": { - "count": 1 - }, "@typescript-eslint/prefer-nullish-coalescing": { "count": 2 }, diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index c7c860f01e7..b8ba12370bd 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/packages/assets-controllers/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Gate `TokenListController` polling on controller initialization to avoid duplicate token list API requests during startup races ([#8113](https://github.com/MetaMask/core/pull/8113)) +- Update token balance fallback behavior so missing ERC-20 balances from `AccountsApiBalanceFetcher` are returned as `unprocessedTokens` and fetched through RPC fallback, rather than being forcibly set to zero ([#0000](https://github.com/MetaMask/core/pull/0000)) ## [100.1.0] diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index e123a0c3f3c..6301e1be6a4 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/packages/assets-controllers/src/TokenBalancesController.test.ts @@ -17,6 +17,7 @@ import BN from 'bn.js'; import type nock from 'nock'; import { mockAPI_accountsAPI_MultichainAccountBalances as mockAPIAccountsAPIMultichainAccountBalancesCamelCase } from './__fixtures__/account-api-v4-mocks'; +import { AccountsApiBalanceFetcher } from './multi-chain-accounts-service/api-balance-fetcher'; import * as multicall from './multicall'; import { RpcBalanceFetcher } from './rpc-service/rpc-balance-fetcher'; import type { @@ -6680,6 +6681,94 @@ describe('TokenBalancesController', () => { messengerCallSpy.mockRestore(); }); + it('should forward unprocessed token fallbacks from API fetcher to RPC fetcher', async () => { + const chainId = '0x1' as ChainIdHex; + const accountAddress = '0x0000000000000000000000000000000000000000'; + const token1 = '0x1111111111111111111111111111111111111111'; + + const tokens = { + allDetectedTokens: {}, + allTokens: { + [chainId]: { + [accountAddress]: [ + { address: token1, symbol: 'TK1', decimals: 18 }, + ], + }, + }, + }; + + const selectedAccount = createMockInternalAccount({ + address: accountAddress, + }); + + const apiFetchSpy = jest + .spyOn(AccountsApiBalanceFetcher.prototype, 'fetch') + .mockResolvedValue({ + balances: [ + { + success: true, + value: new BN(1), + account: accountAddress, + token: NATIVE_TOKEN_ADDRESS as Hex, + chainId, + }, + ], + unprocessedTokens: { + [chainId]: { + [accountAddress]: [token1], + }, + }, + }); + + const { controller } = setupController({ + tokens, + listAccounts: [selectedAccount], + config: { + accountsApiChainIds: () => [chainId], + }, + }); + + const rpcFetchSpy = jest + .spyOn(RpcBalanceFetcher.prototype, 'fetch') + .mockResolvedValue({ + balances: [ + { + success: true, + value: new BN(200), + account: accountAddress as ChecksumAddress, + token: token1 as Hex, + chainId, + }, + ], + }); + + await controller.updateBalances({ + chainIds: [chainId], + queryAllAccounts: true, + }); + + expect(apiFetchSpy).toHaveBeenCalled(); + expect(rpcFetchSpy).toHaveBeenCalledWith( + expect.objectContaining({ + chainIds: [chainId], + unprocessedTokens: { + [chainId]: { + [accountAddress]: [token1], + }, + }, + }), + ); + + expect( + controller.state.tokenBalances[accountAddress as ChecksumAddress]?.[ + chainId + ]?.[toChecksumHexAddress(token1) as ChecksumAddress], + ).toStrictEqual(toHex(200)); + + apiFetchSpy.mockRestore(); + rpcFetchSpy.mockRestore(); + }); + it('should handle fetcher throwing error (lines 868-880)', async () => { const chainId = '0x1'; const accountAddress = '0x0000000000000000000000000000000000000000'; diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 2c8f0e91a43..c158ccd7d12 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -66,6 +66,7 @@ import { AccountsApiBalanceFetcher } from './multi-chain-accounts-service/api-ba import type { BalanceFetcher, ProcessedBalance, + UnprocessedTokens, } from './multi-chain-accounts-service/api-balance-fetcher'; import { RpcBalanceFetcher } from './rpc-service/rpc-balance-fetcher'; import type { @@ -818,6 +819,68 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ }): Promise { const aggregated: ProcessedBalance[] = []; let remainingChains = [...targetChains]; + const unprocessedTokens: UnprocessedTokens = {}; + + const getUnprocessedTokensForChains = ( + chains: ChainIdHex[], + ): UnprocessedTokens | undefined => { + const filteredUnprocessedTokens = chains.reduce( + (accumulator, chainId) => { + if (unprocessedTokens[chainId]) { + accumulator[chainId] = unprocessedTokens[chainId]; + } + return accumulator; + }, + {}, + ); + + return Object.keys(filteredUnprocessedTokens).length > 0 + ? filteredUnprocessedTokens + : undefined; + }; + + const mergeUnprocessedTokens = ( + incomingUnprocessedTokens?: UnprocessedTokens, + ): void => { + if (!incomingUnprocessedTokens) { + return; + } + + Object.entries(incomingUnprocessedTokens).forEach( + ([chainId, accountsWithTokens]) => { + if (!accountsWithTokens) { + return; + } + + const chainIdHex = chainId as ChainIdHex; + unprocessedTokens[chainIdHex] ??= {}; + const currentChainTokens = unprocessedTokens[chainIdHex] as Record< + string, + string[] + >; + + Object.entries(accountsWithTokens).forEach(([account, tokens]) => { + if (!tokens.length) { + return; + } + + currentChainTokens[account] ??= []; + const currentAccountTokens = currentChainTokens[account]; + const existingTokenSet = new Set( + currentAccountTokens.map((token) => token.toLowerCase()), + ); + + tokens.forEach((token) => { + const lowerCaseToken = token.toLowerCase(); + if (!existingTokenSet.has(lowerCaseToken)) { + currentAccountTokens.push(token); + existingTokenSet.add(lowerCaseToken); + } + }); + }); + }, + ); + }; for (const fetcher of this.#balanceFetchers) { const supportedChains = remainingChains.filter((chain) => @@ -834,6 +897,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ selectedAccount, allAccounts, jwtToken, + unprocessedTokens: getUnprocessedTokensForChains(supportedChains), }); if (result.balances?.length) { @@ -843,6 +907,10 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ remainingChains = remainingChains.filter( (chain) => !processed.has(chain), ); + + processed.forEach((chainId) => { + delete unprocessedTokens[chainId]; + }); } if (result.unprocessedChainIds?.length) { @@ -862,6 +930,25 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ .catch(() => { // Silently handle token detection errors }); + + result.unprocessedChainIds.forEach((chainId) => { + delete unprocessedTokens[chainId]; + }); + } + + if (result.unprocessedTokens) { + mergeUnprocessedTokens(result.unprocessedTokens); + + const unprocessedTokenChainIds = Object.keys( + result.unprocessedTokens, + ) as ChainIdHex[]; + const currentRemaining = [...remainingChains]; + const chainsToAdd = unprocessedTokenChainIds.filter( + (chainId) => + supportedChains.includes(chainId) && + !currentRemaining.includes(chainId), + ); + remainingChains.push(...chainsToAdd); } } catch (error) { console.warn( diff --git a/packages/assets-controllers/src/TokenDetectionController.test.ts b/packages/assets-controllers/src/TokenDetectionController.test.ts index 55e92667b66..e7706a7f917 100644 --- a/packages/assets-controllers/src/TokenDetectionController.test.ts +++ b/packages/assets-controllers/src/TokenDetectionController.test.ts @@ -3408,12 +3408,10 @@ describe('TokenDetectionController', () => { ); }); - it('should skip tokens not found in cache and log warning', async () => { + it('should skip tokens not found in cache', async () => { const mockTokenAddress = '0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48'; const chainId = '0xa86a'; - const consoleSpy = jest.spyOn(console, 'warn').mockImplementation(); - await withController( { options: { @@ -3434,19 +3432,12 @@ describe('TokenDetectionController', () => { chainId: chainId as Hex, }); - // Should log warning about missing token metadata - expect(consoleSpy).toHaveBeenCalledWith( - expect.stringContaining('Token metadata not found in cache'), - ); - // Should not call addTokens if no tokens have metadata expect(callActionSpy).not.toHaveBeenCalledWith( 'TokensController:addTokens', expect.anything(), expect.anything(), ); - - consoleSpy.mockRestore(); }, ); }); diff --git a/packages/assets-controllers/src/TokenDetectionController.ts b/packages/assets-controllers/src/TokenDetectionController.ts index d3c2f312806..de7ce68be4f 100644 --- a/packages/assets-controllers/src/TokenDetectionController.ts +++ b/packages/assets-controllers/src/TokenDetectionController.ts @@ -838,9 +838,6 @@ export class TokenDetectionController extends StaticIntervalPollingController { }); }); - describe('erc20 token zero balance guarantee', () => { + describe('erc20 token unprocessed handling', () => { + const trackedToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; + const arrangeBalanceFetcher = (): AccountsApiBalanceFetcher => { const responseWithoutErc20: GetBalancesResponse = { count: 1, @@ -820,7 +822,7 @@ describe('AccountsApiBalanceFetcher', () => { [MOCK_ADDRESS_1]: { '0x1': { [ZERO_ADDRESS]: {}, - '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831': '0x814a20', // previously had balance, should be zero now if api doesn't return it + [trackedToken]: '0x814a20', }, }, }), @@ -829,7 +831,7 @@ describe('AccountsApiBalanceFetcher', () => { return balanceFetcher; }; - it('should include erc20 token entry for addresses even when API does not return erc20 balance', async () => { + it('should return missing erc20 balances as unprocessed tokens', async () => { balanceFetcher = arrangeBalanceFetcher(); const result = await balanceFetcher.fetch({ @@ -839,17 +841,17 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - expect(result.balances).toHaveLength(2); + expect(result.balances).toHaveLength(1); expect(result.balances[0].token).toStrictEqual(ZERO_ADDRESS); - expect(result.balances[1].token).toBe( - '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'.toLowerCase(), - ); - expect(result.balances[1].value).toStrictEqual(new BN('0')); // balance is zero now since API did not return a value for this token + expect(result.unprocessedTokens).toStrictEqual({ + '0x1': { + [MOCK_ADDRESS_1]: [trackedToken.toLowerCase()], + }, + }); }); - it('should not zero out erc20 balances for accounts excluded from selected-account requests', async () => { - const selectedAccountToken = - '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; + it('should not include erc20 tokens for accounts excluded from selected-account requests', async () => { + const selectedAccountToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; const excludedAccountToken = '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E'; mockFetchMultiChainBalancesV4.mockResolvedValue({ @@ -884,32 +886,19 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - const zeroedSelectedAccountToken = result.balances.find( - (balance) => - balance.account === MOCK_ADDRESS_1 && - balance.token === selectedAccountToken.toLowerCase(), - ); - expect(zeroedSelectedAccountToken).toStrictEqual( + expect(result.unprocessedTokens?.['0x1']).toStrictEqual( expect.objectContaining({ - success: true, - value: new BN('0'), - account: MOCK_ADDRESS_1, - token: selectedAccountToken.toLowerCase(), - chainId: '0x1', + [MOCK_ADDRESS_1]: [selectedAccountToken.toLowerCase()], }), ); - const zeroedExcludedAccountToken = result.balances.find( - (balance) => - balance.account === MOCK_ADDRESS_2 && - balance.token === excludedAccountToken.toLowerCase(), - ); - expect(zeroedExcludedAccountToken).toBeUndefined(); + const excludedAccountUnprocessedTokens = + result.unprocessedTokens?.['0x1']?.[MOCK_ADDRESS_2]; + expect(excludedAccountUnprocessedTokens).toBeUndefined(); }); - it('should not zero out erc20 balances for accounts excluded from all-accounts requests', async () => { - const includedAccountToken = - '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; + it('should not include erc20 tokens for accounts excluded from all-accounts requests', async () => { + const includedAccountToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; const excludedAccount = '0x1111111111111111111111111111111111111111'; const excludedAccountToken = '0xA0b86a33E6441c86c33E1C6B9cD964c0BA2A86B'; @@ -953,30 +942,18 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - const zeroedIncludedAccountToken = result.balances.find( - (balance) => - balance.account === MOCK_ADDRESS_1 && - balance.token === includedAccountToken.toLowerCase(), - ); - expect(zeroedIncludedAccountToken).toStrictEqual( + expect(result.unprocessedTokens?.['0x1']).toStrictEqual( expect.objectContaining({ - success: true, - value: new BN('0'), - account: MOCK_ADDRESS_1, - token: includedAccountToken.toLowerCase(), - chainId: '0x1', + [MOCK_ADDRESS_1]: [includedAccountToken.toLowerCase()], }), ); - const zeroedExcludedAccountToken = result.balances.find( - (balance) => - balance.account === excludedAccount && - balance.token === excludedAccountToken.toLowerCase(), - ); - expect(zeroedExcludedAccountToken).toBeUndefined(); + const excludedAccountUnprocessedTokens = + result.unprocessedTokens?.['0x1']?.[excludedAccount]; + expect(excludedAccountUnprocessedTokens).toBeUndefined(); }); - it('should not include erc20 token entry for chains that are not supported by account API', async () => { + it('should not include erc20 unprocessed tokens for chains not supported by account API', async () => { balanceFetcher = arrangeBalanceFetcher(); balanceFetcher = new AccountsApiBalanceFetcher( @@ -987,10 +964,11 @@ describe('AccountsApiBalanceFetcher', () => { '0x1': { [ZERO_ADDRESS]: {}, }, - // Avalanche is not a supported chain, so balances should not be zeroed out + // Avalanche is not a supported chain, so this token should not be included + // in the unprocessed token response. '0xa86a': { [ZERO_ADDRESS]: {}, - '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E': '0x814a20', // USDC AVAX has balance, should not be zeroed out + '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E': '0x814a20', }, }, }), @@ -1011,6 +989,7 @@ describe('AccountsApiBalanceFetcher', () => { value: expect.any(BN), }), ); + expect(result.unprocessedTokens).toBeUndefined(); }); }); diff --git a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts index bd0386960cb..7694bdb4320 100644 --- a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts +++ b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts @@ -39,9 +39,14 @@ export type ProcessedBalance = { chainId: ChainIdHex; }; +export type UnprocessedTokens = Partial< + Record> +>; + export type BalanceFetchResult = { balances: ProcessedBalance[]; unprocessedChainIds?: ChainIdHex[]; + unprocessedTokens?: UnprocessedTokens; }; export type BalanceFetcher = { @@ -52,6 +57,7 @@ export type BalanceFetcher = { selectedAccount: ChecksumAddress; allAccounts: InternalAccount[]; jwtToken?: string; + unprocessedTokens?: UnprocessedTokens; }): Promise; }; @@ -415,6 +421,26 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { ) : selectedAccount.toLowerCase() === address.toLowerCase(); + const unprocessedTokens: UnprocessedTokens = {}; + + const addUnprocessedToken = ( + chainId: ChainIdHex, + account: string, + tokenAddress: string, + ): void => { + unprocessedTokens[chainId] ??= {}; + const chainUnprocessedTokens = unprocessedTokens[chainId] as Record< + string, + string[] + >; + + chainUnprocessedTokens[account] ??= []; + const accountUnprocessedTokens = chainUnprocessedTokens[account]; + if (!accountUnprocessedTokens.includes(tokenAddress)) { + accountUnprocessedTokens.push(tokenAddress); + } + }; + // Add zero native balance entries for addresses that API didn't return addressChainMap.forEach((chains, address) => { chains.forEach((chainId) => { @@ -442,7 +468,9 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { }); }); - // Add zero erc-20 balance entries for addresses that API didn't return + // Track ERC-20 balances that were not returned by Accounts API. + // These can then be fetched by a fallback fetcher (RPC) without + // overwriting potentially stale balances with zero values. if (this.#getUserTokens) { const userTokens = this.#getUserTokens(); Object.entries(userTokens).forEach(([account, chains]) => { @@ -462,13 +490,11 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { isAccountIncluded; if (isERC && shouldZeroOutBalance) { - results.push({ - success: true, - value: new BN('0'), - account: account as ChecksumAddress, - token: tokenLowerCase as ChecksumAddress, - chainId: chainId as ChainIdHex, - }); + addUnprocessedToken( + chainId as ChainIdHex, + account, + tokenLowerCase, + ); } }); }); @@ -481,6 +507,10 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { return { balances: results, unprocessedChainIds, + unprocessedTokens: + Object.keys(unprocessedTokens).length > 0 + ? unprocessedTokens + : undefined, }; } } diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts index db0f93f3e9c..184c7c37723 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts @@ -322,6 +322,66 @@ describe('RpcBalanceFetcher', () => { expect(result.balances.length).toBeGreaterThan(0); }); + it('should fetch only unprocessed tokens without native or staked balances', async () => { + mockGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({ + tokenBalances: { + [MOCK_TOKEN_ADDRESS_1]: { + [MOCK_ADDRESS_1]: new BN('42'), + }, + [ZERO_ADDRESS]: { + [MOCK_ADDRESS_1]: new BN('123'), + }, + }, + stakedBalances: { + [MOCK_ADDRESS_1]: new BN('999'), + }, + }); + + const result = await rpcBalanceFetcher.fetch({ + chainIds: [MOCK_CHAIN_ID], + queryAllAccounts: false, + selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, + allAccounts: MOCK_INTERNAL_ACCOUNTS, + unprocessedTokens: { + [MOCK_CHAIN_ID]: { + [MOCK_ADDRESS_1]: [MOCK_TOKEN_ADDRESS_1], + }, + }, + }); + + expect(mockGetTokensState).not.toHaveBeenCalled(); + expect(mockGetTokenBalancesForMultipleAddresses).toHaveBeenCalledWith( + [ + { + accountAddress: MOCK_ADDRESS_1, + tokenAddresses: [MOCK_TOKEN_ADDRESS_1], + }, + ], + MOCK_CHAIN_ID, + mockProvider, + false, + false, + ); + + expect( + result.balances.some((balance) => balance.token === ZERO_ADDRESS), + ).toBe(false); + expect( + result.balances.some( + (balance) => balance.token === STAKING_CONTRACT_ADDRESS, + ), + ).toBe(false); + expect(result.balances).toStrictEqual([ + { + success: true, + value: new BN('42'), + account: MOCK_ADDRESS_1, + token: MOCK_TOKEN_ADDRESS_1, + chainId: MOCK_CHAIN_ID, + }, + ]); + }); + it('should handle multiple chain IDs', async () => { await rpcBalanceFetcher.fetch({ chainIds: [MOCK_CHAIN_ID, MOCK_CHAIN_ID_2], diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts index 8cfa0936a83..85b944dddad 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts @@ -9,6 +9,7 @@ import type { Hex } from '@metamask/utils'; import BN from 'bn.js'; import { STAKING_CONTRACT_ADDRESS_BY_CHAINID } from '../AssetsContractController'; +import type { UnprocessedTokens } from '../multi-chain-accounts-service/api-balance-fetcher'; import { getTokenBalancesForMultipleAddresses } from '../multicall'; import type { TokensControllerState } from '../TokensController'; @@ -28,6 +29,7 @@ export type ProcessedBalance = { export type BalanceFetchResult = { balances: ProcessedBalance[]; unprocessedChainIds?: ChainIdHex[]; + unprocessedTokens?: UnprocessedTokens; }; export type BalanceFetcher = { @@ -37,6 +39,7 @@ export type BalanceFetcher = { queryAllAccounts: boolean; selectedAccount: ChecksumAddress; allAccounts: InternalAccount[]; + unprocessedTokens?: UnprocessedTokens; }): Promise; }; @@ -82,18 +85,39 @@ export class RpcBalanceFetcher implements BalanceFetcher { queryAllAccounts, selectedAccount, allAccounts, + unprocessedTokens, }: Parameters[0]): Promise { // Process all chains in parallel for better performance const chainProcessingPromises = chainIds.map(async (chainId) => { - const tokensState = this.#getTokensState(); - const accountTokenGroups = buildAccountTokenGroupsStatic( - chainId, - queryAllAccounts, - selectedAccount, - allAccounts, - tokensState.allTokens, - tokensState.allDetectedTokens, + const chainUnprocessedTokens = unprocessedTokens?.[chainId]; + const hasUnprocessedTokensForChain = Boolean( + chainUnprocessedTokens && + Object.keys(chainUnprocessedTokens).length > 0, ); + + let accountTokenGroups: { + accountAddress: ChecksumAddress; + tokenAddresses: ChecksumAddress[]; + }[]; + + if (hasUnprocessedTokensForChain) { + accountTokenGroups = buildUnprocessedAccountTokenGroupsStatic( + chainUnprocessedTokens as Record, + queryAllAccounts, + selectedAccount, + allAccounts, + ); + } else { + const tokensState = this.#getTokensState(); + accountTokenGroups = buildAccountTokenGroupsStatic( + chainId, + queryAllAccounts, + selectedAccount, + allAccounts, + tokensState.allTokens, + tokensState.allDetectedTokens, + ); + } if (!accountTokenGroups.length) { return []; } @@ -101,14 +125,16 @@ export class RpcBalanceFetcher implements BalanceFetcher { const provider = this.#getProvider(chainId); await this.#ensureFreshBlockData(chainId); + const includeNativeAndStaked = !hasUnprocessedTokensForChain; + const balanceResult = await safelyExecuteWithTimeout( async () => { return await getTokenBalancesForMultipleAddresses( accountTokenGroups, chainId, provider, - true, // include native - true, // include staked + includeNativeAndStaked, + includeNativeAndStaked, ); }, true, @@ -123,23 +149,25 @@ export class RpcBalanceFetcher implements BalanceFetcher { const { tokenBalances, stakedBalances } = balanceResult; const chainResults: ProcessedBalance[] = []; - // Add native token entries for all addresses being processed - const allAddressesForNative = new Set(); - accountTokenGroups.forEach((group) => { - allAddressesForNative.add(group.accountAddress); - }); + if (includeNativeAndStaked) { + // Add native token entries for all addresses being processed + const allAddressesForNative = new Set(); + accountTokenGroups.forEach((group) => { + allAddressesForNative.add(group.accountAddress); + }); - // Ensure native token entries exist for all addresses - allAddressesForNative.forEach((address) => { - const nativeBalance = tokenBalances[ZERO_ADDRESS]?.[address] || null; - chainResults.push({ - success: true, - value: nativeBalance || new BN('0'), - account: address as ChecksumAddress, - token: ZERO_ADDRESS, - chainId, + // Ensure native token entries exist for all addresses + allAddressesForNative.forEach((address) => { + const nativeBalance = tokenBalances[ZERO_ADDRESS]?.[address] || null; + chainResults.push({ + success: true, + value: nativeBalance || new BN('0'), + account: address as ChecksumAddress, + token: ZERO_ADDRESS, + chainId, + }); }); - }); + } // Add other token balances Object.entries(tokenBalances).forEach(([tokenAddr, balances]) => { @@ -160,7 +188,7 @@ export class RpcBalanceFetcher implements BalanceFetcher { // Add staked balances for all addresses being processed const stakingContractAddress = this.#getStakingContractAddress(chainId); - if (stakingContractAddress) { + if (includeNativeAndStaked && stakingContractAddress) { // Get all unique addresses being processed for this chain const allAddresses = new Set(); accountTokenGroups.forEach((group) => { @@ -237,7 +265,7 @@ function buildAccountTokenGroupsStatic( tokenAddress: ChecksumAddress; }[] = []; - const add = ([account, tokens]: [string, unknown[]]) => { + const add = ([account, tokens]: [string, unknown[]]): void => { const shouldInclude = queryAllAccounts || checksum(account) === checksum(selectedAccount); if (!shouldInclude) { @@ -294,3 +322,28 @@ function buildAccountTokenGroupsStatic( tokenAddresses, })); } + +function buildUnprocessedAccountTokenGroupsStatic( + chainUnprocessedTokens: Record, + queryAllAccounts: boolean, + selectedAccount: ChecksumAddress, + allAccounts: InternalAccount[], +): { accountAddress: ChecksumAddress; tokenAddresses: ChecksumAddress[] }[] { + const includedAccounts = new Set( + queryAllAccounts + ? allAccounts.map((account) => account.address.toLowerCase()) + : [selectedAccount.toLowerCase()], + ); + + return Object.entries(chainUnprocessedTokens) + .filter(([accountAddress, tokenAddresses]) => { + return ( + includedAccounts.has(accountAddress.toLowerCase()) && + tokenAddresses.length > 0 + ); + }) + .map(([accountAddress, tokenAddresses]) => ({ + accountAddress: accountAddress as ChecksumAddress, + tokenAddresses: tokenAddresses as ChecksumAddress[], + })); +} From 6e810b14096836db6647cca11317f5510b48611b Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 11:31:12 +0000 Subject: [PATCH 02/10] Fix changelog entry PR link for check Co-authored-by: Prithpal Sooriya --- packages/assets-controllers/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/assets-controllers/CHANGELOG.md b/packages/assets-controllers/CHANGELOG.md index b8ba12370bd..4da421eb912 100644 --- a/packages/assets-controllers/CHANGELOG.md +++ b/packages/assets-controllers/CHANGELOG.md @@ -14,7 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Gate `TokenListController` polling on controller initialization to avoid duplicate token list API requests during startup races ([#8113](https://github.com/MetaMask/core/pull/8113)) -- Update token balance fallback behavior so missing ERC-20 balances from `AccountsApiBalanceFetcher` are returned as `unprocessedTokens` and fetched through RPC fallback, rather than being forcibly set to zero ([#0000](https://github.com/MetaMask/core/pull/0000)) +- Update token balance fallback behavior so missing ERC-20 balances from `AccountsApiBalanceFetcher` are returned as `unprocessedTokens` and fetched through RPC fallback, rather than being forcibly set to zero ([#8132](https://github.com/MetaMask/core/pull/8132)) ## [100.1.0] From 09111320acb9978cd83890f7ba17b248ecafe100 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 13:53:06 +0000 Subject: [PATCH 03/10] Reshape unprocessed tokens to account-first mapping Co-authored-by: Prithpal Sooriya --- .../src/TokenBalancesController.test.ts | 8 +-- .../src/TokenBalancesController.ts | 70 ++++++++++++------- .../api-balance-fetcher.test.ts | 16 ++--- .../api-balance-fetcher.ts | 27 +++---- .../rpc-service/rpc-balance-fetcher.test.ts | 4 +- .../src/rpc-service/rpc-balance-fetcher.ts | 24 ++++--- 6 files changed, 87 insertions(+), 62 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.test.ts b/packages/assets-controllers/src/TokenBalancesController.test.ts index 6301e1be6a4..950258c0858 100644 --- a/packages/assets-controllers/src/TokenBalancesController.test.ts +++ b/packages/assets-controllers/src/TokenBalancesController.test.ts @@ -6714,8 +6714,8 @@ describe('TokenBalancesController', () => { }, ], unprocessedTokens: { - [chainId]: { - [accountAddress]: [token1], + [accountAddress]: { + [chainId]: [token1], }, }, }); @@ -6752,8 +6752,8 @@ describe('TokenBalancesController', () => { expect.objectContaining({ chainIds: [chainId], unprocessedTokens: { - [chainId]: { - [accountAddress]: [token1], + [accountAddress]: { + [chainId]: [token1], }, }, }), diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index c158ccd7d12..f616f23986b 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -824,14 +824,26 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ const getUnprocessedTokensForChains = ( chains: ChainIdHex[], ): UnprocessedTokens | undefined => { - const filteredUnprocessedTokens = chains.reduce( - (accumulator, chainId) => { - if (unprocessedTokens[chainId]) { - accumulator[chainId] = unprocessedTokens[chainId]; + const chainSet = new Set(chains); + const filteredUnprocessedTokens: UnprocessedTokens = {}; + + Object.entries(unprocessedTokens).forEach( + ([account, unprocessedTokensByChain]) => { + const filteredByChain: Record = {}; + + Object.entries(unprocessedTokensByChain).forEach( + ([chainId, tokenAddresses]) => { + const chainIdHex = chainId as ChainIdHex; + if (chainSet.has(chainIdHex) && tokenAddresses.length > 0) { + filteredByChain[chainIdHex] = tokenAddresses; + } + }, + ); + + if (Object.keys(filteredByChain).length > 0) { + filteredUnprocessedTokens[account] = filteredByChain; } - return accumulator; }, - {}, ); return Object.keys(filteredUnprocessedTokens).length > 0 @@ -847,25 +859,18 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ } Object.entries(incomingUnprocessedTokens).forEach( - ([chainId, accountsWithTokens]) => { - if (!accountsWithTokens) { - return; - } - - const chainIdHex = chainId as ChainIdHex; - unprocessedTokens[chainIdHex] ??= {}; - const currentChainTokens = unprocessedTokens[chainIdHex] as Record< - string, - string[] - >; + ([account, tokensByChainId]) => { + unprocessedTokens[account] ??= {}; + const currentTokensByChainId = unprocessedTokens[account]; - Object.entries(accountsWithTokens).forEach(([account, tokens]) => { + Object.entries(tokensByChainId).forEach(([chainId, tokens]) => { if (!tokens.length) { return; } - currentChainTokens[account] ??= []; - const currentAccountTokens = currentChainTokens[account]; + const chainIdHex = chainId as ChainIdHex; + currentTokensByChainId[chainIdHex] ??= []; + const currentAccountTokens = currentTokensByChainId[chainIdHex]; const existingTokenSet = new Set( currentAccountTokens.map((token) => token.toLowerCase()), ); @@ -882,6 +887,15 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ ); }; + const removeChainFromUnprocessedTokens = (chainId: ChainIdHex): void => { + Object.keys(unprocessedTokens).forEach((account) => { + delete unprocessedTokens[account][chainId]; + if (Object.keys(unprocessedTokens[account]).length === 0) { + delete unprocessedTokens[account]; + } + }); + }; + for (const fetcher of this.#balanceFetchers) { const supportedChains = remainingChains.filter((chain) => fetcher.supports(chain), @@ -909,7 +923,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ ); processed.forEach((chainId) => { - delete unprocessedTokens[chainId]; + removeChainFromUnprocessedTokens(chainId); }); } @@ -932,16 +946,22 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ }); result.unprocessedChainIds.forEach((chainId) => { - delete unprocessedTokens[chainId]; + removeChainFromUnprocessedTokens(chainId); }); } if (result.unprocessedTokens) { mergeUnprocessedTokens(result.unprocessedTokens); - const unprocessedTokenChainIds = Object.keys( - result.unprocessedTokens, - ) as ChainIdHex[]; + const unprocessedTokenChainIdsSet = new Set(); + Object.values(result.unprocessedTokens).forEach((tokensByChainId) => { + Object.keys(tokensByChainId).forEach((chainId) => { + unprocessedTokenChainIdsSet.add(chainId as ChainIdHex); + }); + }); + const unprocessedTokenChainIds = Array.from( + unprocessedTokenChainIdsSet, + ); const currentRemaining = [...remainingChains]; const chainsToAdd = unprocessedTokenChainIds.filter( (chainId) => diff --git a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts index b286a93c4e8..8217bb6b4a2 100644 --- a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts @@ -844,8 +844,8 @@ describe('AccountsApiBalanceFetcher', () => { expect(result.balances).toHaveLength(1); expect(result.balances[0].token).toStrictEqual(ZERO_ADDRESS); expect(result.unprocessedTokens).toStrictEqual({ - '0x1': { - [MOCK_ADDRESS_1]: [trackedToken.toLowerCase()], + [MOCK_ADDRESS_1]: { + '0x1': [trackedToken.toLowerCase()], }, }); }); @@ -886,14 +886,14 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - expect(result.unprocessedTokens?.['0x1']).toStrictEqual( + expect(result.unprocessedTokens?.[MOCK_ADDRESS_1]).toStrictEqual( expect.objectContaining({ - [MOCK_ADDRESS_1]: [selectedAccountToken.toLowerCase()], + '0x1': [selectedAccountToken.toLowerCase()], }), ); const excludedAccountUnprocessedTokens = - result.unprocessedTokens?.['0x1']?.[MOCK_ADDRESS_2]; + result.unprocessedTokens?.[MOCK_ADDRESS_2]?.['0x1']; expect(excludedAccountUnprocessedTokens).toBeUndefined(); }); @@ -942,14 +942,14 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - expect(result.unprocessedTokens?.['0x1']).toStrictEqual( + expect(result.unprocessedTokens?.[MOCK_ADDRESS_1]).toStrictEqual( expect.objectContaining({ - [MOCK_ADDRESS_1]: [includedAccountToken.toLowerCase()], + '0x1': [includedAccountToken.toLowerCase()], }), ); const excludedAccountUnprocessedTokens = - result.unprocessedTokens?.['0x1']?.[excludedAccount]; + result.unprocessedTokens?.[excludedAccount]?.['0x1']; expect(excludedAccountUnprocessedTokens).toBeUndefined(); }); diff --git a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts index 7694bdb4320..23f2c8a6952 100644 --- a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts +++ b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts @@ -39,9 +39,14 @@ export type ProcessedBalance = { chainId: ChainIdHex; }; -export type UnprocessedTokens = Partial< - Record> ->; +/** + * Account -> ChainId -> TokenAddress[] + */ +export type UnprocessedTokens = { + [account: string]: { + [chainId: ChainIdHex]: string[]; + }; +}; export type BalanceFetchResult = { balances: ProcessedBalance[]; @@ -424,18 +429,14 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { const unprocessedTokens: UnprocessedTokens = {}; const addUnprocessedToken = ( - chainId: ChainIdHex, account: string, + chainId: ChainIdHex, tokenAddress: string, ): void => { - unprocessedTokens[chainId] ??= {}; - const chainUnprocessedTokens = unprocessedTokens[chainId] as Record< - string, - string[] - >; - - chainUnprocessedTokens[account] ??= []; - const accountUnprocessedTokens = chainUnprocessedTokens[account]; + unprocessedTokens[account] ??= {}; + const accountUnprocessedTokensByChain = unprocessedTokens[account]; + accountUnprocessedTokensByChain[chainId] ??= []; + const accountUnprocessedTokens = accountUnprocessedTokensByChain[chainId]; if (!accountUnprocessedTokens.includes(tokenAddress)) { accountUnprocessedTokens.push(tokenAddress); } @@ -491,8 +492,8 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { if (isERC && shouldZeroOutBalance) { addUnprocessedToken( - chainId as ChainIdHex, account, + chainId as ChainIdHex, tokenLowerCase, ); } diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts index 184c7c37723..bd3ff1fb90b 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts @@ -343,8 +343,8 @@ describe('RpcBalanceFetcher', () => { selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, allAccounts: MOCK_INTERNAL_ACCOUNTS, unprocessedTokens: { - [MOCK_CHAIN_ID]: { - [MOCK_ADDRESS_1]: [MOCK_TOKEN_ADDRESS_1], + [MOCK_ADDRESS_1]: { + [MOCK_CHAIN_ID]: [MOCK_TOKEN_ADDRESS_1], }, }, }); diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts index 85b944dddad..b6ff0c45a21 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts @@ -89,10 +89,11 @@ export class RpcBalanceFetcher implements BalanceFetcher { }: Parameters[0]): Promise { // Process all chains in parallel for better performance const chainProcessingPromises = chainIds.map(async (chainId) => { - const chainUnprocessedTokens = unprocessedTokens?.[chainId]; const hasUnprocessedTokensForChain = Boolean( - chainUnprocessedTokens && - Object.keys(chainUnprocessedTokens).length > 0, + unprocessedTokens && + Object.values(unprocessedTokens).some( + (tokensByChain) => (tokensByChain[chainId]?.length ?? 0) > 0, + ), ); let accountTokenGroups: { @@ -102,7 +103,8 @@ export class RpcBalanceFetcher implements BalanceFetcher { if (hasUnprocessedTokensForChain) { accountTokenGroups = buildUnprocessedAccountTokenGroupsStatic( - chainUnprocessedTokens as Record, + unprocessedTokens as UnprocessedTokens, + chainId, queryAllAccounts, selectedAccount, allAccounts, @@ -324,7 +326,8 @@ function buildAccountTokenGroupsStatic( } function buildUnprocessedAccountTokenGroupsStatic( - chainUnprocessedTokens: Record, + unprocessedTokens: UnprocessedTokens, + chainId: ChainIdHex, queryAllAccounts: boolean, selectedAccount: ChecksumAddress, allAccounts: InternalAccount[], @@ -335,15 +338,16 @@ function buildUnprocessedAccountTokenGroupsStatic( : [selectedAccount.toLowerCase()], ); - return Object.entries(chainUnprocessedTokens) - .filter(([accountAddress, tokenAddresses]) => { + return Object.entries(unprocessedTokens) + .filter(([accountAddress, tokensByChain]) => { + const tokenAddresses = tokensByChain[chainId]; return ( includedAccounts.has(accountAddress.toLowerCase()) && - tokenAddresses.length > 0 + (tokenAddresses?.length ?? 0) > 0 ); }) - .map(([accountAddress, tokenAddresses]) => ({ + .map(([accountAddress, tokensByChain]) => ({ accountAddress: accountAddress as ChecksumAddress, - tokenAddresses: tokenAddresses as ChecksumAddress[], + tokenAddresses: tokensByChain[chainId] as ChecksumAddress[], })); } From 846893275091a1918e9765343b12c738edd46219 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 15:35:05 +0000 Subject: [PATCH 04/10] Refactor TokenBalancesController to enhance balance fetcher structure and improve unprocessed token handling - Updated #balanceFetchers to include fetcher names for better traceability. - Simplified unprocessed token management by introducing previous state tracking. - Enhanced error reporting for unprocessed tokens during balance fetching. - Refactored balance fetching logic to streamline processing across multiple chains. --- .../src/TokenBalancesController.ts | 196 ++++++-------- .../api-balance-fetcher.ts | 11 +- .../src/rpc-service/rpc-balance-fetcher.ts | 240 ++++++++++-------- 3 files changed, 222 insertions(+), 225 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index f616f23986b..fa7c807926a 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -279,7 +279,7 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ readonly #isOnboarded: () => boolean; - readonly #balanceFetchers: BalanceFetcher[]; + readonly #balanceFetchers: { fetcher: BalanceFetcher; name: string }[]; #allTokens: TokensControllerState['allTokens'] = {}; @@ -349,11 +349,21 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ // Always include AccountsApiFetcher - it dynamically checks allowExternalServices() in supports() this.#balanceFetchers = [ - this.#createAccountsApiFetcher(), - new RpcBalanceFetcher(this.#getProvider, this.#getNetworkClient, () => ({ - allTokens: this.#allTokens, - allDetectedTokens: this.#detectedTokens, - })), + { + fetcher: this.#createAccountsApiFetcher(), + name: 'AccountsApiFetcher', + }, + { + fetcher: new RpcBalanceFetcher( + this.#getProvider, + this.#getNetworkClient, + () => ({ + allTokens: this.#allTokens, + allDetectedTokens: this.#detectedTokens, + }), + ), + name: 'RpcFetcher', + }, ]; this.setIntervalLength(interval); @@ -819,84 +829,10 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ }): Promise { const aggregated: ProcessedBalance[] = []; let remainingChains = [...targetChains]; - const unprocessedTokens: UnprocessedTokens = {}; - - const getUnprocessedTokensForChains = ( - chains: ChainIdHex[], - ): UnprocessedTokens | undefined => { - const chainSet = new Set(chains); - const filteredUnprocessedTokens: UnprocessedTokens = {}; - - Object.entries(unprocessedTokens).forEach( - ([account, unprocessedTokensByChain]) => { - const filteredByChain: Record = {}; - - Object.entries(unprocessedTokensByChain).forEach( - ([chainId, tokenAddresses]) => { - const chainIdHex = chainId as ChainIdHex; - if (chainSet.has(chainIdHex) && tokenAddresses.length > 0) { - filteredByChain[chainIdHex] = tokenAddresses; - } - }, - ); - - if (Object.keys(filteredByChain).length > 0) { - filteredUnprocessedTokens[account] = filteredByChain; - } - }, - ); - - return Object.keys(filteredUnprocessedTokens).length > 0 - ? filteredUnprocessedTokens - : undefined; - }; - - const mergeUnprocessedTokens = ( - incomingUnprocessedTokens?: UnprocessedTokens, - ): void => { - if (!incomingUnprocessedTokens) { - return; - } - - Object.entries(incomingUnprocessedTokens).forEach( - ([account, tokensByChainId]) => { - unprocessedTokens[account] ??= {}; - const currentTokensByChainId = unprocessedTokens[account]; + let previousUnprocessedTokens: UnprocessedTokens | undefined; + let previousFetcherName: string | undefined; - Object.entries(tokensByChainId).forEach(([chainId, tokens]) => { - if (!tokens.length) { - return; - } - - const chainIdHex = chainId as ChainIdHex; - currentTokensByChainId[chainIdHex] ??= []; - const currentAccountTokens = currentTokensByChainId[chainIdHex]; - const existingTokenSet = new Set( - currentAccountTokens.map((token) => token.toLowerCase()), - ); - - tokens.forEach((token) => { - const lowerCaseToken = token.toLowerCase(); - if (!existingTokenSet.has(lowerCaseToken)) { - currentAccountTokens.push(token); - existingTokenSet.add(lowerCaseToken); - } - }); - }); - }, - ); - }; - - const removeChainFromUnprocessedTokens = (chainId: ChainIdHex): void => { - Object.keys(unprocessedTokens).forEach((account) => { - delete unprocessedTokens[account][chainId]; - if (Object.keys(unprocessedTokens[account]).length === 0) { - delete unprocessedTokens[account]; - } - }); - }; - - for (const fetcher of this.#balanceFetchers) { + for (const { fetcher, name: fetcherName } of this.#balanceFetchers) { const supportedChains = remainingChains.filter((chain) => fetcher.supports(chain), ); @@ -911,9 +847,52 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ selectedAccount, allAccounts, jwtToken, - unprocessedTokens: getUnprocessedTokensForChains(supportedChains), + unprocessedTokens: previousUnprocessedTokens, }); + // Balance Error Reporting - for unprocessed tokens from last fetcher, if balances are retrieved + const unprocessedTokensForReporting = previousUnprocessedTokens; + if (unprocessedTokensForReporting && result.balances?.length) { + const confirmedUnprocessedTokens: { + chainId: string; + tokenAddress: string; + }[] = []; + + // Capture balances that were found (> 0 balance), and was unprocessed + result.balances.forEach((bal) => { + const lowercaseAccount = bal.account.toLowerCase(); + const lowercaseTokenAddress = bal.token.toLowerCase(); + + const hasResultBalance = + bal.success && bal.token && bal.value && !bal.value.isZero(); + const isUnprocessed = unprocessedTokensForReporting?.[ + lowercaseAccount + ]?.[bal.chainId]?.includes(lowercaseTokenAddress); + + if (hasResultBalance && isUnprocessed) { + confirmedUnprocessedTokens.push({ + chainId: bal.chainId, + tokenAddress: lowercaseTokenAddress, + }); + } + }); + + const confirmedUnprocessedTokenStrings = + confirmedUnprocessedTokens.map( + (token) => `${token.chainId}:${token.tokenAddress}`, + ); + if (confirmedUnprocessedTokens.length) { + console.warn( + `TokenBalanceController: fetcher ${previousFetcherName} did not process tokens (instead handled by fetcher ${fetcherName}): ${confirmedUnprocessedTokenStrings.join(', ')}`, + ); + } + } + + // Set new previous fields + previousUnprocessedTokens = result.unprocessedTokens; + previousFetcherName = fetcherName; + + // Add balances, and removed processed chains if (result.balances?.length) { aggregated.push(...result.balances); @@ -921,20 +900,22 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ remainingChains = remainingChains.filter( (chain) => !processed.has(chain), ); - - processed.forEach((chainId) => { - removeChainFromUnprocessedTokens(chainId); - }); } - if (result.unprocessedChainIds?.length) { - const currentRemaining = [...remainingChains]; - const chainsToAdd = result.unprocessedChainIds.filter( - (chainId) => - supportedChains.includes(chainId) && - !currentRemaining.includes(chainId), + // Add unprocessed chains (from missing chains or missing tokens) + if (result.unprocessedChainIds || result.unprocessedTokens) { + const resultUnprocessedChains = result.unprocessedChainIds ?? []; + const resultUnsupportedTokenChains = Object.entries( + result.unprocessedTokens ?? {}, + ).flatMap(([_account, chainMap]) => Object.keys(chainMap)) as Hex[]; + + remainingChains = Array.from( + new Set([ + ...remainingChains, + ...resultUnprocessedChains, + ...resultUnsupportedTokenChains, + ]), ); - remainingChains.push(...chainsToAdd); this.messenger .call('TokenDetectionController:detectTokens', { @@ -944,31 +925,6 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ .catch(() => { // Silently handle token detection errors }); - - result.unprocessedChainIds.forEach((chainId) => { - removeChainFromUnprocessedTokens(chainId); - }); - } - - if (result.unprocessedTokens) { - mergeUnprocessedTokens(result.unprocessedTokens); - - const unprocessedTokenChainIdsSet = new Set(); - Object.values(result.unprocessedTokens).forEach((tokensByChainId) => { - Object.keys(tokensByChainId).forEach((chainId) => { - unprocessedTokenChainIdsSet.add(chainId as ChainIdHex); - }); - }); - const unprocessedTokenChainIds = Array.from( - unprocessedTokenChainIdsSet, - ); - const currentRemaining = [...remainingChains]; - const chainsToAdd = unprocessedTokenChainIds.filter( - (chainId) => - supportedChains.includes(chainId) && - !currentRemaining.includes(chainId), - ); - remainingChains.push(...chainsToAdd); } } catch (error) { console.warn( diff --git a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts index 23f2c8a6952..8cc9bd889b6 100644 --- a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts +++ b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts @@ -62,7 +62,7 @@ export type BalanceFetcher = { selectedAccount: ChecksumAddress; allAccounts: InternalAccount[]; jwtToken?: string; - unprocessedTokens?: UnprocessedTokens; + unprocessedTokens?: UnprocessedTokens; // API Balance Fetcher does not process unprocessed tokens }): Promise; }; @@ -234,7 +234,10 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { return results; } - async #fetchBalances(addrs: CaipAccountAddress[], jwtToken?: string) { + async #fetchBalances( + addrs: CaipAccountAddress[], + jwtToken?: string, + ): Promise { // If we have fewer than or equal to the batch size, make a single request if (addrs.length <= ACCOUNTS_API_BATCH_SIZE) { return await fetchMultiChainBalancesV4( @@ -290,7 +293,7 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { }: Parameters[0]): Promise { const caipAddrs: CaipAccountAddress[] = []; - for (const chainId of chainIds.filter((c) => this.supports(c))) { + for (const chainId of chainIds.filter((chain) => this.supports(chain))) { if (queryAllAccounts) { allAccounts.forEach((a) => caipAddrs.push(toCaipAccount(chainId, a.address as ChecksumAddress)), @@ -492,7 +495,7 @@ export class AccountsApiBalanceFetcher implements BalanceFetcher { if (isERC && shouldZeroOutBalance) { addUnprocessedToken( - account, + account.toLowerCase(), chainId as ChainIdHex, tokenLowerCase, ); diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts index b6ff0c45a21..3b9a315c372 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts @@ -89,37 +89,36 @@ export class RpcBalanceFetcher implements BalanceFetcher { }: Parameters[0]): Promise { // Process all chains in parallel for better performance const chainProcessingPromises = chainIds.map(async (chainId) => { - const hasUnprocessedTokensForChain = Boolean( - unprocessedTokens && - Object.values(unprocessedTokens).some( - (tokensByChain) => (tokensByChain[chainId]?.length ?? 0) > 0, - ), - ); + // if there are unprocessed tokens for a chain, it means the chain was partially processed. + // because of this, we need to build distinct account <-> token groups to process + const hasUnprocessedTokensForChain = queryAllAccounts + ? Object.values(unprocessedTokens ?? {}).some((chainMap) => + Boolean(chainMap[chainId] && chainMap[chainId].length > 0), + ) + : Boolean( + unprocessedTokens?.[selectedAccount.toLowerCase()]?.[chainId] && + unprocessedTokens[selectedAccount.toLowerCase()][chainId].length > + 0, + ); + + const tokensState = this.#getTokensState(); + const { accountTokenGroups, includeNativeAndStaked } = + hasUnprocessedTokensForChain + ? buildUnprocessedAccountTokenGroupsStatic( + chainId, + queryAllAccounts, + selectedAccount, + unprocessedTokens as UnprocessedTokens, + ) + : buildAccountTokenGroupsStatic( + chainId, + queryAllAccounts, + selectedAccount, + allAccounts, + tokensState.allTokens, + tokensState.allDetectedTokens, + ); - let accountTokenGroups: { - accountAddress: ChecksumAddress; - tokenAddresses: ChecksumAddress[]; - }[]; - - if (hasUnprocessedTokensForChain) { - accountTokenGroups = buildUnprocessedAccountTokenGroupsStatic( - unprocessedTokens as UnprocessedTokens, - chainId, - queryAllAccounts, - selectedAccount, - allAccounts, - ); - } else { - const tokensState = this.#getTokensState(); - accountTokenGroups = buildAccountTokenGroupsStatic( - chainId, - queryAllAccounts, - selectedAccount, - allAccounts, - tokensState.allTokens, - tokensState.allDetectedTokens, - ); - } if (!accountTokenGroups.length) { return []; } @@ -127,8 +126,6 @@ export class RpcBalanceFetcher implements BalanceFetcher { const provider = this.#getProvider(chainId); await this.#ensureFreshBlockData(chainId); - const includeNativeAndStaked = !hasUnprocessedTokensForChain; - const balanceResult = await safelyExecuteWithTimeout( async () => { return await getTokenBalancesForMultipleAddresses( @@ -200,10 +197,10 @@ export class RpcBalanceFetcher implements BalanceFetcher { // Add staked balance entry for each address const checksummedStakingAddress = checksum(stakingContractAddress); allAddresses.forEach((address) => { - const stakedBalance = stakedBalances?.[address] || null; + const stakedBalance = stakedBalances?.[address] ?? null; chainResults.push({ success: true, - value: stakedBalance || new BN('0'), + value: stakedBalance ?? new BN('0'), account: address as ChecksumAddress, token: checksummedStakingAddress, chainId, @@ -242,70 +239,37 @@ export class RpcBalanceFetcher implements BalanceFetcher { } } -/** - * Merges imported & detected tokens for the requested chain and returns a list - * of `{ accountAddress, tokenAddresses[] }` suitable for getTokenBalancesForMultipleAddresses. - * - * @param chainId - The chain ID to build account token groups for - * @param queryAllAccounts - Whether to query all accounts or just the selected one - * @param selectedAccount - The currently selected account - * @param allAccounts - All available accounts - * @param allTokens - All tokens from TokensController - * @param allDetectedTokens - All detected tokens from TokensController - * @returns Array of account/token groups for multicall - */ -function buildAccountTokenGroupsStatic( - chainId: ChainIdHex, +type AccountTokenGroup = { + accountAddress: ChecksumAddress; + tokenAddresses: ChecksumAddress[]; +}; + +function buildAccountTokenGroups( queryAllAccounts: boolean, selectedAccount: ChecksumAddress, - allAccounts: InternalAccount[], - allTokens: TokensControllerState['allTokens'], - allDetectedTokens: TokensControllerState['allDetectedTokens'], -): { accountAddress: ChecksumAddress; tokenAddresses: ChecksumAddress[] }[] { + accountTokenMap: { [account: string]: string[] }, +): AccountTokenGroup[] { const pairs: { accountAddress: ChecksumAddress; tokenAddress: ChecksumAddress; }[] = []; - const add = ([account, tokens]: [string, unknown[]]): void => { + const add = ([account, tokens]: [string, string[]]): void => { + const checksumAccount = checksum(account); const shouldInclude = - queryAllAccounts || checksum(account) === checksum(selectedAccount); + queryAllAccounts || checksumAccount === checksum(selectedAccount); if (!shouldInclude) { return; } - tokens.forEach((t: unknown) => + tokens.forEach((token: string) => pairs.push({ - accountAddress: account as ChecksumAddress, - tokenAddress: checksum((t as { address: string }).address), + accountAddress: checksumAccount, + tokenAddress: checksum(token), }), ); }; - Object.entries(allTokens[chainId] ?? {}).forEach( - add as (entry: [string, unknown]) => void, - ); - Object.entries(allDetectedTokens[chainId] ?? {}).forEach( - add as (entry: [string, unknown]) => void, - ); - - // Always include native token for relevant accounts - if (queryAllAccounts) { - allAccounts.forEach((a) => { - pairs.push({ - accountAddress: a.address as ChecksumAddress, - tokenAddress: ZERO_ADDRESS, - }); - }); - } else { - pairs.push({ - accountAddress: selectedAccount, - tokenAddress: ZERO_ADDRESS, - }); - } - - if (!pairs.length) { - return []; - } + Object.entries(accountTokenMap).forEach(add); // group by account const map = new Map(); @@ -325,29 +289,103 @@ function buildAccountTokenGroupsStatic( })); } -function buildUnprocessedAccountTokenGroupsStatic( - unprocessedTokens: UnprocessedTokens, +/** + * Merges imported & detected tokens for the requested chain and returns a list + * of `{ accountAddress, tokenAddresses[] }` suitable for getTokenBalancesForMultipleAddresses. + * + * @param chainId - The chain ID to build account token groups for + * @param queryAllAccounts - Whether to query all accounts or just the selected one + * @param selectedAccount - The currently selected account + * @param allAccounts - All available accounts + * @param allTokens - All tokens from TokensController + * @param allDetectedTokens - All detected tokens from TokensController + * @returns Array of account/token groups for multicall + */ +function buildAccountTokenGroupsStatic( chainId: ChainIdHex, queryAllAccounts: boolean, selectedAccount: ChecksumAddress, allAccounts: InternalAccount[], -): { accountAddress: ChecksumAddress; tokenAddresses: ChecksumAddress[] }[] { - const includedAccounts = new Set( - queryAllAccounts - ? allAccounts.map((account) => account.address.toLowerCase()) - : [selectedAccount.toLowerCase()], - ); + allTokens: TokensControllerState['allTokens'], + allDetectedTokens: TokensControllerState['allDetectedTokens'], +): { + accountTokenGroups: AccountTokenGroup[]; + includeNativeAndStaked: true; +} { + const accountTokenMap: { [account: string]: string[] } = {}; + + // Add all tokens + Object.entries(allTokens[chainId] ?? {}).forEach(([account, tokens]) => { + const lowercaseAccount = account.toLowerCase(); + accountTokenMap[lowercaseAccount] = tokens.map((token) => + token.address.toLowerCase(), + ); + }); - return Object.entries(unprocessedTokens) - .filter(([accountAddress, tokensByChain]) => { - const tokenAddresses = tokensByChain[chainId]; - return ( - includedAccounts.has(accountAddress.toLowerCase()) && - (tokenAddresses?.length ?? 0) > 0 + // Add all detected tokens + Object.entries(allDetectedTokens[chainId] ?? {}).forEach( + ([account, tokens]) => { + const lowercaseAccount = account.toLowerCase(); + if (!accountTokenMap[lowercaseAccount]) { + accountTokenMap[lowercaseAccount] = []; + } + accountTokenMap[lowercaseAccount] = Array.from( + new Set([ + ...accountTokenMap[lowercaseAccount], + ...tokens.map((token) => token.address.toLowerCase()), + ]), ); - }) - .map(([accountAddress, tokensByChain]) => ({ - accountAddress: accountAddress as ChecksumAddress, - tokenAddresses: tokensByChain[chainId] as ChecksumAddress[], - })); + }, + ); + + // Add native tokens + if (queryAllAccounts) { + allAccounts.forEach((a) => { + accountTokenMap[a.address.toLowerCase()] = [ZERO_ADDRESS]; + }); + } else { + accountTokenMap[selectedAccount.toLowerCase()] = [ZERO_ADDRESS]; + } + + return { + accountTokenGroups: buildAccountTokenGroups( + queryAllAccounts, + selectedAccount, + accountTokenMap, + ), + includeNativeAndStaked: true, + }; +} + +function buildUnprocessedAccountTokenGroupsStatic( + chainId: ChainIdHex, + queryAllAccounts: boolean, + selectedAccount: ChecksumAddress, + unprocessedTokens: UnprocessedTokens, +): { + accountTokenGroups: AccountTokenGroup[]; + includeNativeAndStaked: false; +} { + const accountTokenMap: { [account: string]: string[] } = {}; + Object.entries(unprocessedTokens).forEach(([account, tokens]) => { + const lowercaseAccount = account.toLowerCase(); + if ( + queryAllAccounts || + lowercaseAccount === selectedAccount.toLowerCase() + ) { + const tokenAddresses = + tokens?.[chainId]?.map((tokenAddress) => tokenAddress.toLowerCase()) ?? + []; + accountTokenMap[lowercaseAccount] = tokenAddresses; + } + }); + + return { + accountTokenGroups: buildAccountTokenGroups( + queryAllAccounts, + selectedAccount, + accountTokenMap, + ), + includeNativeAndStaked: false, + }; } From 8ee87afcbd2e14ecdcd8f90d780e08bca31eb7d4 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 15:42:33 +0000 Subject: [PATCH 05/10] Enhance TokenBalancesController to improve balance aggregation and unprocessed chain handling - Added logic to aggregate balances and filter out processed chains. - Implemented handling for unprocessed chains and tokens, ensuring they are detected and processed correctly. - Introduced error handling for token detection calls to prevent application crashes on failures. --- .../src/TokenBalancesController.ts | 70 +++++++++---------- 1 file changed, 35 insertions(+), 35 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index fa7c807926a..9a307c26f8e 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -850,6 +850,41 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ unprocessedTokens: previousUnprocessedTokens, }); + // Add balances, and removed processed chains + if (result.balances?.length) { + aggregated.push(...result.balances); + + const processed = new Set(result.balances.map((b) => b.chainId)); + remainingChains = remainingChains.filter( + (chain) => !processed.has(chain), + ); + } + + // Add unprocessed chains (from missing chains or missing tokens) + if (result.unprocessedChainIds || result.unprocessedTokens) { + const resultUnprocessedChains = result.unprocessedChainIds ?? []; + const resultUnsupportedTokenChains = Object.entries( + result.unprocessedTokens ?? {}, + ).flatMap(([_account, chainMap]) => Object.keys(chainMap)) as Hex[]; + + remainingChains = Array.from( + new Set([ + ...remainingChains, + ...resultUnprocessedChains, + ...resultUnsupportedTokenChains, + ]), + ); + + this.messenger + .call('TokenDetectionController:detectTokens', { + chainIds: result.unprocessedChainIds, + forceRpc: true, + }) + .catch(() => { + // Silently handle token detection errors + }); + } + // Balance Error Reporting - for unprocessed tokens from last fetcher, if balances are retrieved const unprocessedTokensForReporting = previousUnprocessedTokens; if (unprocessedTokensForReporting && result.balances?.length) { @@ -891,41 +926,6 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ // Set new previous fields previousUnprocessedTokens = result.unprocessedTokens; previousFetcherName = fetcherName; - - // Add balances, and removed processed chains - if (result.balances?.length) { - aggregated.push(...result.balances); - - const processed = new Set(result.balances.map((b) => b.chainId)); - remainingChains = remainingChains.filter( - (chain) => !processed.has(chain), - ); - } - - // Add unprocessed chains (from missing chains or missing tokens) - if (result.unprocessedChainIds || result.unprocessedTokens) { - const resultUnprocessedChains = result.unprocessedChainIds ?? []; - const resultUnsupportedTokenChains = Object.entries( - result.unprocessedTokens ?? {}, - ).flatMap(([_account, chainMap]) => Object.keys(chainMap)) as Hex[]; - - remainingChains = Array.from( - new Set([ - ...remainingChains, - ...resultUnprocessedChains, - ...resultUnsupportedTokenChains, - ]), - ); - - this.messenger - .call('TokenDetectionController:detectTokens', { - chainIds: result.unprocessedChainIds, - forceRpc: true, - }) - .catch(() => { - // Silently handle token detection errors - }); - } } catch (error) { console.warn( `Balance fetcher failed for chains ${supportedChains.join(', ')}: ${String(error)}`, From 9a3ea4e6e0f9994324aefd09c3558571ff1ffae7 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 16:48:48 +0000 Subject: [PATCH 06/10] Refactor ERC-20 token handling in balance fetcher tests - Updated test descriptions for clarity on unprocessed token handling. - Improved token address handling by ensuring consistent formatting. - Enhanced logic to correctly manage unprocessed tokens for selected and excluded accounts. - Removed redundant test cases to streamline the testing process. --- .../api-balance-fetcher.test.ts | 63 ++++++++--------- .../rpc-service/rpc-balance-fetcher.test.ts | 69 ++----------------- .../src/rpc-service/rpc-balance-fetcher.ts | 6 +- 3 files changed, 38 insertions(+), 100 deletions(-) diff --git a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts index 8217bb6b4a2..ce4d5ec9733 100644 --- a/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.test.ts @@ -802,9 +802,7 @@ describe('AccountsApiBalanceFetcher', () => { }); }); - describe('erc20 token unprocessed handling', () => { - const trackedToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; - + describe('erc20 token unprocessed token handling', () => { const arrangeBalanceFetcher = (): AccountsApiBalanceFetcher => { const responseWithoutErc20: GetBalancesResponse = { count: 1, @@ -822,7 +820,7 @@ describe('AccountsApiBalanceFetcher', () => { [MOCK_ADDRESS_1]: { '0x1': { [ZERO_ADDRESS]: {}, - [trackedToken]: '0x814a20', + '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831': '0x814a20', // previously had balance, should be zero now if api doesn't return it }, }, }), @@ -831,7 +829,7 @@ describe('AccountsApiBalanceFetcher', () => { return balanceFetcher; }; - it('should return missing erc20 balances as unprocessed tokens', async () => { + it('includes unprocessed tokens for missing erc20 balances for selected account', async () => { balanceFetcher = arrangeBalanceFetcher(); const result = await balanceFetcher.fetch({ @@ -844,14 +842,15 @@ describe('AccountsApiBalanceFetcher', () => { expect(result.balances).toHaveLength(1); expect(result.balances[0].token).toStrictEqual(ZERO_ADDRESS); expect(result.unprocessedTokens).toStrictEqual({ - [MOCK_ADDRESS_1]: { - '0x1': [trackedToken.toLowerCase()], + [MOCK_ADDRESS_1.toLowerCase()]: { + '0x1': ['0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'.toLowerCase()], }, }); }); - it('should not include erc20 tokens for accounts excluded from selected-account requests', async () => { - const selectedAccountToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; + it('does not include unprocessed tokens for non selected accounts', async () => { + const selectedAccountToken = + '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; const excludedAccountToken = '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E'; mockFetchMultiChainBalancesV4.mockResolvedValue({ @@ -886,20 +885,21 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - expect(result.unprocessedTokens?.[MOCK_ADDRESS_1]).toStrictEqual( - expect.objectContaining({ + expect(result.unprocessedTokens).toStrictEqual({ + // Does not include non-selected accounts + [MOCK_ADDRESS_1.toLowerCase()]: { '0x1': [selectedAccountToken.toLowerCase()], - }), - ); + }, + }); - const excludedAccountUnprocessedTokens = - result.unprocessedTokens?.[MOCK_ADDRESS_2]?.['0x1']; - expect(excludedAccountUnprocessedTokens).toBeUndefined(); + expect( + result.unprocessedTokens?.[MOCK_ADDRESS_2.toLowerCase()], + ).toBeUndefined(); }); - it('should not include erc20 tokens for accounts excluded from all-accounts requests', async () => { - const includedAccountToken = '0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; - const excludedAccount = '0x1111111111111111111111111111111111111111'; + it('includes unprocessed tokens for missing erc20 balances for all accounts', async () => { + const includedAccountToken = + '0x0xaf88d065e77c8cC2239327C5EDb3A432268e5831'; const excludedAccountToken = '0xA0b86a33E6441c86c33E1C6B9cD964c0BA2A86B'; mockFetchMultiChainBalancesV4.mockResolvedValue({ @@ -926,7 +926,7 @@ describe('AccountsApiBalanceFetcher', () => { [includedAccountToken]: '0x814a20', }, }, - [excludedAccount]: { + [MOCK_ADDRESS_2]: { '0x1': { [ZERO_ADDRESS]: {}, [excludedAccountToken]: '0x814a20', @@ -942,18 +942,17 @@ describe('AccountsApiBalanceFetcher', () => { allAccounts: MOCK_INTERNAL_ACCOUNTS, }); - expect(result.unprocessedTokens?.[MOCK_ADDRESS_1]).toStrictEqual( - expect.objectContaining({ + expect(result.unprocessedTokens).toStrictEqual({ + [MOCK_ADDRESS_1.toLowerCase()]: { '0x1': [includedAccountToken.toLowerCase()], - }), - ); - - const excludedAccountUnprocessedTokens = - result.unprocessedTokens?.[excludedAccount]?.['0x1']; - expect(excludedAccountUnprocessedTokens).toBeUndefined(); + }, + [MOCK_ADDRESS_2.toLowerCase()]: { + '0x1': [excludedAccountToken.toLowerCase()], + }, + }); }); - it('should not include erc20 unprocessed tokens for chains not supported by account API', async () => { + it('should not include erc20 token entry for chains that are not supported by account API', async () => { balanceFetcher = arrangeBalanceFetcher(); balanceFetcher = new AccountsApiBalanceFetcher( @@ -964,11 +963,10 @@ describe('AccountsApiBalanceFetcher', () => { '0x1': { [ZERO_ADDRESS]: {}, }, - // Avalanche is not a supported chain, so this token should not be included - // in the unprocessed token response. + // Avalanche is not a supported chain, so balances should not be zeroed out '0xa86a': { [ZERO_ADDRESS]: {}, - '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E': '0x814a20', + '0xB97EF9Ef8734C71904D8002F8b6Bc66Dd9c48a6E': '0x814a20', // USDC AVAX has balance, should not be zeroed out }, }, }), @@ -989,7 +987,6 @@ describe('AccountsApiBalanceFetcher', () => { value: expect.any(BN), }), ); - expect(result.unprocessedTokens).toBeUndefined(); }); }); diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts index bd3ff1fb90b..9e7b9c4b20c 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts @@ -8,9 +8,9 @@ import type { ChainIdHex, ChecksumAddress } from './rpc-balance-fetcher'; import type { TokensControllerState } from '../TokensController'; const MOCK_ADDRESS_1 = '0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045'; -const MOCK_ADDRESS_2 = '0x742d35cc6675c4f17f41140100aa83a4b1fa4c82'; +const MOCK_ADDRESS_2 = '0x1c5E29b17822CC96B834092Ec056a9cd4e833C09'; const MOCK_TOKEN_ADDRESS_1 = '0x6B175474E89094C44Da98b954EedeAC495271d0F'; -const MOCK_TOKEN_ADDRESS_2 = '0xA0b86a33E6441c86c33E1C6B9cD964c0BA2A86B'; +const MOCK_TOKEN_ADDRESS_2 = '0x00000000000C2E074eC69A0dFb2997BA6C7d2e1e'; const MOCK_CHAIN_ID = '0x1' as ChainIdHex; const MOCK_CHAIN_ID_2 = '0x89' as ChainIdHex; const ZERO_ADDRESS = @@ -322,66 +322,6 @@ describe('RpcBalanceFetcher', () => { expect(result.balances.length).toBeGreaterThan(0); }); - it('should fetch only unprocessed tokens without native or staked balances', async () => { - mockGetTokenBalancesForMultipleAddresses.mockResolvedValueOnce({ - tokenBalances: { - [MOCK_TOKEN_ADDRESS_1]: { - [MOCK_ADDRESS_1]: new BN('42'), - }, - [ZERO_ADDRESS]: { - [MOCK_ADDRESS_1]: new BN('123'), - }, - }, - stakedBalances: { - [MOCK_ADDRESS_1]: new BN('999'), - }, - }); - - const result = await rpcBalanceFetcher.fetch({ - chainIds: [MOCK_CHAIN_ID], - queryAllAccounts: false, - selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, - allAccounts: MOCK_INTERNAL_ACCOUNTS, - unprocessedTokens: { - [MOCK_ADDRESS_1]: { - [MOCK_CHAIN_ID]: [MOCK_TOKEN_ADDRESS_1], - }, - }, - }); - - expect(mockGetTokensState).not.toHaveBeenCalled(); - expect(mockGetTokenBalancesForMultipleAddresses).toHaveBeenCalledWith( - [ - { - accountAddress: MOCK_ADDRESS_1, - tokenAddresses: [MOCK_TOKEN_ADDRESS_1], - }, - ], - MOCK_CHAIN_ID, - mockProvider, - false, - false, - ); - - expect( - result.balances.some((balance) => balance.token === ZERO_ADDRESS), - ).toBe(false); - expect( - result.balances.some( - (balance) => balance.token === STAKING_CONTRACT_ADDRESS, - ), - ).toBe(false); - expect(result.balances).toStrictEqual([ - { - success: true, - value: new BN('42'), - account: MOCK_ADDRESS_1, - token: MOCK_TOKEN_ADDRESS_1, - chainId: MOCK_CHAIN_ID, - }, - ]); - }); - it('should handle multiple chain IDs', async () => { await rpcBalanceFetcher.fetch({ chainIds: [MOCK_CHAIN_ID, MOCK_CHAIN_ID_2], @@ -673,7 +613,7 @@ describe('RpcBalanceFetcher', () => { ); }); - it('should handle duplicate tokens in the same group', async () => { + it('removes duplicates in the same group', async () => { const tokensStateWithDuplicates = { allTokens: { [MOCK_CHAIN_ID]: { @@ -714,8 +654,7 @@ describe('RpcBalanceFetcher', () => { { accountAddress: MOCK_ADDRESS_1, tokenAddresses: [ - MOCK_TOKEN_ADDRESS_1, - MOCK_TOKEN_ADDRESS_1, + MOCK_TOKEN_ADDRESS_1, // we do not have duplicates addresses in request! ZERO_ADDRESS, ], }, diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts index 3b9a315c372..2ddca2a0ddc 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts @@ -341,10 +341,12 @@ function buildAccountTokenGroupsStatic( // Add native tokens if (queryAllAccounts) { allAccounts.forEach((a) => { - accountTokenMap[a.address.toLowerCase()] = [ZERO_ADDRESS]; + accountTokenMap[a.address.toLowerCase()] ??= []; + accountTokenMap[a.address.toLowerCase()].push(ZERO_ADDRESS); }); } else { - accountTokenMap[selectedAccount.toLowerCase()] = [ZERO_ADDRESS]; + accountTokenMap[selectedAccount.toLowerCase()] ??= []; + accountTokenMap[selectedAccount.toLowerCase()].push(ZERO_ADDRESS); } return { From 16c9121ed00c2d61c196755b3690c26f7a5cc026 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 16:51:42 +0000 Subject: [PATCH 07/10] Remove unused ESLint suppressions for specific files in assets-controllers --- eslint-suppressions.json | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/eslint-suppressions.json b/eslint-suppressions.json index e80df9da18d..6965155c0ba 100644 --- a/eslint-suppressions.json +++ b/eslint-suppressions.json @@ -440,14 +440,6 @@ "count": 4 } }, - "packages/assets-controllers/src/multi-chain-accounts-service/api-balance-fetcher.ts": { - "@typescript-eslint/explicit-function-return-type": { - "count": 1 - }, - "id-length": { - "count": 1 - } - }, "packages/assets-controllers/src/multi-chain-accounts-service/multi-chain-accounts.test.ts": { "@typescript-eslint/explicit-function-return-type": { "count": 3 @@ -476,14 +468,6 @@ "count": 17 } }, - "packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts": { - "@typescript-eslint/prefer-nullish-coalescing": { - "count": 2 - }, - "id-length": { - "count": 1 - } - }, "packages/assets-controllers/src/selectors/stringify-balance.ts": { "@typescript-eslint/explicit-function-return-type": { "count": 1 From fdf8344846ec2aed2117dc7cc21c41fc0bb4681c Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 17:29:27 +0000 Subject: [PATCH 08/10] Refactor account address handling in balance fetcher - Updated account address handling to use ChecksumAddress type for improved type safety. - Simplified token address mapping by removing unnecessary lowercase transformations. - Enhanced logic for managing account token groups to ensure consistency across different account formats. --- .../rpc-service/rpc-balance-fetcher.test.ts | 4 +-- .../src/rpc-service/rpc-balance-fetcher.ts | 26 ++++++++----------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts index 9e7b9c4b20c..1ff8b064c15 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts @@ -8,9 +8,9 @@ import type { ChainIdHex, ChecksumAddress } from './rpc-balance-fetcher'; import type { TokensControllerState } from '../TokensController'; const MOCK_ADDRESS_1 = '0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045'; -const MOCK_ADDRESS_2 = '0x1c5E29b17822CC96B834092Ec056a9cd4e833C09'; +const MOCK_ADDRESS_2 = '0x742d35cc6675c4f17f41140100aa83a4b1fa4c82'; const MOCK_TOKEN_ADDRESS_1 = '0x6B175474E89094C44Da98b954EedeAC495271d0F'; -const MOCK_TOKEN_ADDRESS_2 = '0x00000000000C2E074eC69A0dFb2997BA6C7d2e1e'; +const MOCK_TOKEN_ADDRESS_2 = '0xA0b86a33E6441c86c33E1C6B9cD964c0BA2A86B'; const MOCK_CHAIN_ID = '0x1' as ChainIdHex; const MOCK_CHAIN_ID_2 = '0x89' as ChainIdHex; const ZERO_ADDRESS = diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts index 2ddca2a0ddc..49c369f8a4d 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.ts @@ -263,7 +263,7 @@ function buildAccountTokenGroups( } tokens.forEach((token: string) => pairs.push({ - accountAddress: checksumAccount, + accountAddress: account as ChecksumAddress, tokenAddress: checksum(token), }), ); @@ -316,23 +316,19 @@ function buildAccountTokenGroupsStatic( // Add all tokens Object.entries(allTokens[chainId] ?? {}).forEach(([account, tokens]) => { - const lowercaseAccount = account.toLowerCase(); - accountTokenMap[lowercaseAccount] = tokens.map((token) => - token.address.toLowerCase(), - ); + accountTokenMap[account] = tokens.map((token) => token.address); }); // Add all detected tokens Object.entries(allDetectedTokens[chainId] ?? {}).forEach( ([account, tokens]) => { - const lowercaseAccount = account.toLowerCase(); - if (!accountTokenMap[lowercaseAccount]) { - accountTokenMap[lowercaseAccount] = []; + if (!accountTokenMap[account]) { + accountTokenMap[account] = []; } - accountTokenMap[lowercaseAccount] = Array.from( + accountTokenMap[account] = Array.from( new Set([ - ...accountTokenMap[lowercaseAccount], - ...tokens.map((token) => token.address.toLowerCase()), + ...accountTokenMap[account], + ...tokens.map((token) => token.address), ]), ); }, @@ -341,12 +337,12 @@ function buildAccountTokenGroupsStatic( // Add native tokens if (queryAllAccounts) { allAccounts.forEach((a) => { - accountTokenMap[a.address.toLowerCase()] ??= []; - accountTokenMap[a.address.toLowerCase()].push(ZERO_ADDRESS); + accountTokenMap[a.address] ??= []; + accountTokenMap[a.address].push(ZERO_ADDRESS); }); } else { - accountTokenMap[selectedAccount.toLowerCase()] ??= []; - accountTokenMap[selectedAccount.toLowerCase()].push(ZERO_ADDRESS); + accountTokenMap[selectedAccount] ??= []; + accountTokenMap[selectedAccount].push(ZERO_ADDRESS); } return { From 365b654206c5fc5eec411c8e3d0c22b03f198083 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 18:00:03 +0000 Subject: [PATCH 09/10] test: add rpc unprocessed token routing coverage Co-authored-by: Prithpal Sooriya --- .../rpc-service/rpc-balance-fetcher.test.ts | 135 ++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts index 1ff8b064c15..a0a61bb666d 100644 --- a/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts +++ b/packages/assets-controllers/src/rpc-service/rpc-balance-fetcher.test.ts @@ -5,6 +5,7 @@ import BN from 'bn.js'; import { RpcBalanceFetcher } from './rpc-balance-fetcher'; import type { ChainIdHex, ChecksumAddress } from './rpc-balance-fetcher'; +import type { UnprocessedTokens } from '../multi-chain-accounts-service/api-balance-fetcher'; import type { TokensControllerState } from '../TokensController'; const MOCK_ADDRESS_1 = '0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045'; @@ -462,6 +463,140 @@ describe('RpcBalanceFetcher', () => { true, ); }); + + it('uses unprocessed tokens for selected account and skips native/staked fetches', async () => { + const unprocessedTokens: UnprocessedTokens = { + [MOCK_ADDRESS_1.toLowerCase()]: { + [MOCK_CHAIN_ID]: [MOCK_TOKEN_ADDRESS_1], + }, + }; + + mockGetTokenBalancesForMultipleAddresses.mockResolvedValue({ + tokenBalances: { + [MOCK_TOKEN_ADDRESS_1]: { + [MOCK_ADDRESS_1.toLowerCase()]: new BN('123'), + }, + }, + stakedBalances: { + [MOCK_ADDRESS_1.toLowerCase()]: new BN('999'), + }, + }); + + const result = await rpcBalanceFetcher.fetch({ + chainIds: [MOCK_CHAIN_ID], + queryAllAccounts: false, + selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, + allAccounts: MOCK_INTERNAL_ACCOUNTS, + unprocessedTokens, + }); + + expect(mockGetTokenBalancesForMultipleAddresses).toHaveBeenCalledWith( + [ + { + accountAddress: MOCK_ADDRESS_1.toLowerCase(), + tokenAddresses: [MOCK_TOKEN_ADDRESS_1], + }, + ], + MOCK_CHAIN_ID, + mockProvider, + false, + false, + ); + expect(result.balances).toHaveLength(1); + expect(result.balances[0]).toMatchObject({ + account: MOCK_ADDRESS_1.toLowerCase(), + chainId: MOCK_CHAIN_ID, + }); + expect( + result.balances.some((balance) => balance.token === ZERO_ADDRESS), + ).toBe(false); + expect( + result.balances.some( + (balance) => balance.token === STAKING_CONTRACT_ADDRESS, + ), + ).toBe(false); + }); + + it('uses unprocessed tokens per-chain and falls back to regular mode for other chains', async () => { + const unprocessedTokens: UnprocessedTokens = { + [MOCK_ADDRESS_1.toLowerCase()]: { + [MOCK_CHAIN_ID]: [MOCK_TOKEN_ADDRESS_1], + }, + }; + + await rpcBalanceFetcher.fetch({ + chainIds: [MOCK_CHAIN_ID, MOCK_CHAIN_ID_2], + queryAllAccounts: true, + selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, + allAccounts: MOCK_INTERNAL_ACCOUNTS, + unprocessedTokens, + }); + + const chain1Call = + mockGetTokenBalancesForMultipleAddresses.mock.calls.find( + ([, chainId]) => chainId === MOCK_CHAIN_ID, + ); + expect(chain1Call).toBeDefined(); + expect(chain1Call?.[0]).toStrictEqual([ + { + accountAddress: MOCK_ADDRESS_1.toLowerCase(), + tokenAddresses: [MOCK_TOKEN_ADDRESS_1], + }, + ]); + expect(chain1Call?.[3]).toBe(false); + expect(chain1Call?.[4]).toBe(false); + + const chain2Call = + mockGetTokenBalancesForMultipleAddresses.mock.calls.find( + ([, chainId]) => chainId === MOCK_CHAIN_ID_2, + ); + expect(chain2Call).toBeDefined(); + expect(chain2Call?.[0]).toStrictEqual([ + { + accountAddress: MOCK_ADDRESS_1, + tokenAddresses: [MOCK_TOKEN_ADDRESS_1, ZERO_ADDRESS], + }, + { + accountAddress: MOCK_ADDRESS_2, + tokenAddresses: [ZERO_ADDRESS], + }, + ]); + expect(chain2Call?.[3]).toBe(true); + expect(chain2Call?.[4]).toBe(true); + }); + + it('ignores unprocessed tokens from non-selected accounts when queryAllAccounts is false', async () => { + const unprocessedTokens: UnprocessedTokens = { + [MOCK_ADDRESS_2.toLowerCase()]: { + [MOCK_CHAIN_ID]: [MOCK_TOKEN_ADDRESS_2], + }, + }; + + await rpcBalanceFetcher.fetch({ + chainIds: [MOCK_CHAIN_ID], + queryAllAccounts: false, + selectedAccount: MOCK_ADDRESS_1 as ChecksumAddress, + allAccounts: MOCK_INTERNAL_ACCOUNTS, + unprocessedTokens, + }); + + expect(mockGetTokenBalancesForMultipleAddresses).toHaveBeenCalledWith( + [ + { + accountAddress: MOCK_ADDRESS_1, + tokenAddresses: [ + MOCK_TOKEN_ADDRESS_1, + MOCK_TOKEN_ADDRESS_2, + ZERO_ADDRESS, + ], + }, + ], + MOCK_CHAIN_ID, + mockProvider, + true, + true, + ); + }); }); describe('Token grouping integration (via fetch)', () => { From 3240c4dc62d12eb55874e4b3fc0184c121f6373e Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Fri, 6 Mar 2026 19:28:21 +0000 Subject: [PATCH 10/10] Refactor TokenBalancesController to improve unprocessed chain handling - Introduced a new variable for unprocessed chain IDs to enhance clarity. - Updated logic to aggregate remaining chains more efficiently by combining existing and unprocessed chains. - Improved readability and maintainability of the code structure. --- .../assets-controllers/src/TokenBalancesController.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/packages/assets-controllers/src/TokenBalancesController.ts b/packages/assets-controllers/src/TokenBalancesController.ts index 9a307c26f8e..316c8e4f7ab 100644 --- a/packages/assets-controllers/src/TokenBalancesController.ts +++ b/packages/assets-controllers/src/TokenBalancesController.ts @@ -866,18 +866,20 @@ export class TokenBalancesController extends StaticIntervalPollingController<{ const resultUnsupportedTokenChains = Object.entries( result.unprocessedTokens ?? {}, ).flatMap(([_account, chainMap]) => Object.keys(chainMap)) as Hex[]; - - remainingChains = Array.from( + const unprocessedChainIds = Array.from( new Set([ - ...remainingChains, ...resultUnprocessedChains, ...resultUnsupportedTokenChains, ]), ); + remainingChains = Array.from( + new Set([...remainingChains, ...unprocessedChainIds]), + ); + this.messenger .call('TokenDetectionController:detectTokens', { - chainIds: result.unprocessedChainIds, + chainIds: unprocessedChainIds, forceRpc: true, }) .catch(() => {