fix(token-select): decouple sortByBalance from balance fetching

Author: gabitoesmiapodo

src/components/sharedComponents/TokenSelect/index.tsx

@@ -130,6 +130,7 @@ const TokenSelect = withSuspenseAndRetry<Props>(
     const { isLoadingBalances, tokensByChainId } = useTokens({
       chainId,
       withBalance: showBalance || resolvedSortByBalance,
+      sortByBalance: resolvedSortByBalance,
     })
 
     const { searchResult, searchTerm, setSearchTerm } = useTokenSearch(

src/hooks/useTokens.test.ts

@@ -117,6 +117,31 @@ describe('updateTokensBalances', () => {
     expect(tokens[1].symbol).toBe('USDC')
     expect(tokens[2].symbol).toBe('DAI')
   })
+
+  it('preserves source order when sortByBalance is false', () => {
+    const prices: TokensResponse = {
+      tokens: {
+        1: [
+          makeLifiToken(LOCAL_NATIVE, 'ETH', 18, '2300'),
+          makeLifiToken(usdcAddress, 'USDC', 6, '1'),
+          makeLifiToken(daiAddress, 'DAI', 18, '1'),
+        ],
+      },
+    }
+    const balances: TokenAmount[] = [
+      { ...makeLifiToken(LOCAL_NATIVE, 'ETH', 18, '2300'), amount: 0n },
+      { ...makeLifiToken(usdcAddress, 'USDC', 6, '1'), amount: 5_000_000n },
+      { ...makeLifiToken(daiAddress, 'DAI', 18, '1'), amount: 0n },
+    ]
+
+    const { tokens } = updateTokensBalances(threeTokens, [balances, prices], {
+      sortByBalance: false,
+    })
+
+    expect(tokens[0].symbol).toBe('ETH')
+    expect(tokens[1].symbol).toBe('USDC')
+    expect(tokens[2].symbol).toBe('DAI')
+  })
 })
 
 describe('updateTokensWithRawBalances', () => {
@@ -172,4 +197,24 @@ describe('updateTokensWithRawBalances', () => {
     expect(tokens[0].extensions?.balance).toBe(500n)
     expect(tokens[0].extensions?.priceUSD).toBeUndefined()
   })
+
+  it('preserves source order when sortByBalance is false', () => {
+    const rawBalances: Record<number, Record<string, bigint>> = {
+      11155111: {
+        [LOCAL_NATIVE]: 0n,
+        [usdcAddress]: 1_000_000n,
+      },
+    }
+    const sepoliaTokens: Tokens = [
+      { chainId: 11155111, address: LOCAL_NATIVE, name: 'Ether', symbol: 'ETH', decimals: 18 },
+      { chainId: 11155111, address: usdcAddress, name: 'USD Coin', symbol: 'USDC', decimals: 6 },
+    ]
+
+    const { tokens } = updateTokensWithRawBalances(sepoliaTokens, rawBalances, {
+      sortByBalance: false,
+    })
+
+    expect(tokens[0].symbol).toBe('ETH')
+    expect(tokens[1].symbol).toBe('USDC')
+  })
 })

src/hooks/useTokens.ts

@@ -49,6 +49,7 @@ export const lifiConfig = createConfig({
  * @param {Address} [params.account] - Account address for balance fetching (defaults to connected wallet)
  * @param {Chain['id']} [params.chainId] - Specific chain ID to filter tokens (defaults to all supported chains)
  * @param {boolean} [params.withBalance=true] - Whether to fetch token balances
+ * @param {boolean} [params.sortByBalance=true] - Whether to sort tokens by balance. When false, source order is preserved even if balances are fetched.
  *
  * @returns {Object} Token data and loading state
  * @returns {Token[]} returns.tokens - Array of tokens with price and balance information
@@ -78,12 +79,15 @@ export const useTokens = (
     account,
     chainId,
     withBalance,
+    sortByBalance = true,
   }: {
     account?: Address
     chainId?: Chain['id']
     withBalance?: boolean
+    sortByBalance?: boolean
   } = {
     withBalance: true,
+    sortByBalance: true,
   },
 ) => {
   const { address } = useWeb3Status()
@@ -188,10 +192,16 @@ export const useTokens = (
   const cache = useMemo(() => {
     if (withBalance && account) {
       if (!isLoadingPrices && !isLoadingBalances && tokensBalances && tokensPricesByChain) {
-        return updateTokensBalances(tokensData.tokens, [tokensBalances, tokensPricesByChain])
+        return updateTokensBalances(tokensData.tokens, [tokensBalances, tokensPricesByChain], {
+          sortByBalance,
+        })
       }
       if (useOnchainFallback && !isLoadingOnchainBalances && onchainBalances && chainId) {
-        return updateTokensWithRawBalances(tokensData.tokens, { [chainId]: onchainBalances })
+        return updateTokensWithRawBalances(
+          tokensData.tokens,
+          { [chainId]: onchainBalances },
+          { sortByBalance },
+        )
       }
     }
     return tokensData
@@ -202,6 +212,7 @@ export const useTokens = (
     isLoadingOnchainBalances,
     isLoadingPrices,
     onchainBalances,
+    sortByBalance,
     tokensBalances,
     tokensData,
     tokensPricesByChain,
@@ -228,6 +239,7 @@ export const useTokens = (
 export function updateTokensBalances(
   tokens: Tokens,
   results: [Array<TokenAmount>, TokensResponse],
+  { sortByBalance = true }: { sortByBalance?: boolean } = {},
 ) {
   const [balanceTokens, prices] = results
 
@@ -273,9 +285,11 @@ export function updateTokensBalances(
   })
   logger.timeEnd('extending tokens with balance info')
 
-  logger.time('sorting tokens by balance')
-  tokensWithBalances.sort(sortFn)
-  logger.timeEnd('sorting tokens by balance')
+  if (sortByBalance) {
+    logger.time('sorting tokens by balance')
+    tokensWithBalances.sort(sortFn)
+    logger.timeEnd('sorting tokens by balance')
+  }
 
   logger.time('updating tokens cache')
   const tokensByChain = tokensWithBalances.reduce(
@@ -306,6 +320,7 @@ export function updateTokensBalances(
 export function updateTokensWithRawBalances(
   tokens: Tokens,
   rawBalances: Record<number, Record<string, bigint>>,
+  { sortByBalance = true }: { sortByBalance?: boolean } = {},
 ) {
   const tokensWithBalances = tokens.map(
     (token): Token => ({
@@ -316,7 +331,9 @@ export function updateTokensWithRawBalances(
     }),
   )
 
-  tokensWithBalances.sort(sortByBalancePresenceFn)
+  if (sortByBalance) {
+    tokensWithBalances.sort(sortByBalancePresenceFn)
+  }
 
   const tokensByChainId = tokensWithBalances.reduce(
     (acc, token) => {