fix: refetch token balances after dust success

Author: mmioana

src/components/composite/DustFlow/hooks/useDustModalFlow.tsx

@@ -46,7 +46,7 @@ export const useDustModalFlow = ({
   isOpen,
 }: UseDustModalFlowOptions) => {
   const { t } = useTranslation();
-  const { refresh: refreshPortfolio } = usePortfolioState();
+  const { refreshForTokens } = usePortfolioState();
   const { nonNativeBalances, chains, nativeExtendedTokens } = useDustBalances();
   const fallbackNativeToken = useFallbackNativeToken(nativeExtendedTokens);
   const formFields = useDustFormFields({
@@ -67,6 +67,20 @@ export const useDustModalFlow = ({
     prevChainId: number | undefined;
   }>({ prevThreshold: undefined, prevChainId: undefined });
 
+  const completedDustSummaryRef = useRef<DustSummaryValue | null>(null);
+
+  const refreshCompletedDustTokens = useCallback(() => {
+    const summary = completedDustSummaryRef.current;
+    if (!summary) {
+      return;
+    }
+    completedDustSummaryRef.current = null;
+    void refreshForTokens(summary.address, [
+      ...summary.selectedBalances.map((b) => b.token),
+      summary.nativeToken,
+    ]);
+  }, [refreshForTokens]);
+
   useEffect(() => {
     if (isOpen) {
       dustFieldSyncRef.current = {
@@ -239,22 +253,27 @@ export const useDustModalFlow = ({
     requiresConfirmation: false,
     executorType: supportsBatchTransactions ? 'batch' : 'single',
     fetchCallData,
+    onSuccess: () => {
+      if (dustSummary) {
+        completedDustSummaryRef.current = dustSummary;
+      }
+    },
   });
 
   const statusSheet = useDustConversionStatusSheet({
     transactionForm,
     toTokenBalance: nativeTokenBalance,
     onSuccess: () => {
+      refreshCompletedDustTokens();
       setDustSummary(null);
-      refreshPortfolio();
       widgetNav?.resetForm();
     },
   });
 
   const handleModalClose = () => {
     onClose();
     transactionForm.resetForm();
-    refreshPortfolio();
+    refreshCompletedDustTokens();
   };
 
   const views = useMemo(

src/providers/PortfolioProvider/PortfolioContext.ts

@@ -98,6 +98,7 @@ const defaultOrchestrationState: OrchestrationState = {
   },
   refresh: () => {},
   refreshByAddress: () => {},
+  refreshForTokens: () => Promise.resolve(),
 };
 
 export const PortfolioContext = createContext<PortfolioContextValue>({

src/providers/PortfolioProvider/hooks/useBalancesData.ts

@@ -2,9 +2,12 @@ import { useCallback, useMemo } from 'react';
 import { useQueries, useQueryClient } from '@tanstack/react-query';
 import type { Account } from '@lifi/widget-provider';
 import { useAccount } from '@lifi/wallet-management';
-import type { ChainType } from '@lifi/sdk';
-import { compact, max } from 'lodash';
-import { fetchBalancesForAddress } from '../lib/fetchBalancesForAddresses';
+import type { ChainType, Token as LifiToken } from '@lifi/sdk';
+import { compact, differenceWith, max } from 'lodash';
+import {
+  fetchBalancesForAddress,
+  fetchTokenBalancesForAddress,
+} from '../lib/fetchBalancesForAddresses';
 import { usePortfolioCacheStore } from '@/stores/portfolio/PortfolioCacheStore';
 import type { TokenBalance } from '@/types/tokens';
 
@@ -26,6 +29,7 @@ export interface UseTokensDataResult extends ResultState {
   refetch: () => void;
   cancel: () => void;
   refetchForAddress: (address: string) => void;
+  refetchForTokens: (address: string, tokens: LifiToken[]) => Promise<void>;
   stateByAddress: Record<string, ResultState>;
 }
 
@@ -141,6 +145,35 @@ export const useBalancesData = (): UseTokensDataResult => {
     [queries, addresses, queryClient],
   );
 
+  const refetchForTokens = useCallback(
+    async (address: string, tokens: LifiToken[]) => {
+      const freshBalances = await fetchTokenBalancesForAddress(address, tokens);
+
+      queryClient.setQueryData<TokenQueryData>(
+        ['portfolio-tokens', address],
+        (prev) => {
+          if (!prev) {
+            return prev;
+          }
+
+          const kept = differenceWith(
+            prev.balances,
+            tokens,
+            (b, t) =>
+              b.token.chainId === t.chainId &&
+              b.token.address.toLowerCase() === t.address.toLowerCase(),
+          );
+
+          return {
+            ...prev,
+            balances: [...kept, ...freshBalances],
+          };
+        },
+      );
+    },
+    [queryClient],
+  );
+
   const balancesByAddress = useMemo(() => {
     return queries.reduce(
       (acc, query, index) => {
@@ -213,5 +246,6 @@ export const useBalancesData = (): UseTokensDataResult => {
     refetch,
     cancel,
     refetchForAddress,
+    refetchForTokens,
   };
 };

src/providers/PortfolioProvider/hooks/useOrchestrationState.ts

@@ -133,6 +133,16 @@ export const useOrchestrationState = (
     [balances.refetchForAddress],
   );
 
+  const refreshForTokens = useCallback(
+    (
+      address: string,
+      tokens: Parameters<typeof balances.refetchForTokens>[1],
+    ) => {
+      return balances.refetchForTokens(address, tokens);
+    },
+    [balances.refetchForTokens],
+  );
+
   return useMemo(
     () => ({
       isEmpty,
@@ -150,6 +160,7 @@ export const useOrchestrationState = (
       },
       refresh,
       refreshByAddress,
+      refreshForTokens,
     }),
     [
       isEmpty,
@@ -165,6 +176,7 @@ export const useOrchestrationState = (
       pricesSource,
       refresh,
       refreshByAddress,
+      refreshForTokens,
     ],
   );
 };

src/providers/PortfolioProvider/lib/fetchBalancesForAddresses.ts

@@ -5,7 +5,7 @@ import {
   getWalletBalances,
 } from '@lifi/sdk';
 import type { Token, TokenAmount } from '@lifi/sdk';
-import { flatMap } from 'lodash';
+import { flatMap, groupBy } from 'lodash';
 import { createBatchFetcher } from '@/utils/batches/fetcher';
 import { sdkClient } from '@/utils/instrumentation/lifiSdkConfig';
 import {
@@ -22,7 +22,7 @@ const createTokensBalanceFromPlain = <
   return createTokenBalance(createExtendedToken(token), token.amount);
 };
 
-const createTokenBalancesFromPlainArray = <
+export const createTokenBalancesFromPlainArray = <
   T extends Omit<TokenAmount, 'amount'> & {
     amount?: string | bigint | undefined;
   },
@@ -46,6 +46,24 @@ export interface FetchBalancesParams {
   onComplete?: (tokens: TokenBalance[]) => void;
 }
 
+/** Fetch fresh balances for a specific list of tokens — bypasses getWalletBalances indexer lag */
+export const fetchTokenBalancesForAddress = async (
+  address: string,
+  tokens: Token[],
+): Promise<TokenBalance[]> => {
+  const grouped = groupBy(tokens, (t) => t.chainId);
+
+  const { results } = createBatchFetcher<Token, TokenAmount>(
+    grouped,
+    async (_chainId, tokenBatch) => {
+      const balances = await getTokenBalances(sdkClient, address, tokenBatch);
+      return balances.filter((t) => t.amount && t.amount > BigInt(0));
+    },
+  );
+
+  return createTokenBalancesFromPlainArray(await results);
+};
+
 /** Fetch tokens for EVM wallet (no batching needed) */
 export const fetchBalancesForEVMAddress = async (address: string) => {
   const walletBalances = await getWalletBalances(sdkClient, address);

src/providers/PortfolioProvider/types.ts

@@ -4,6 +4,7 @@ import type {
   WalletToken,
 } from '@/types/tokens';
 import type { App, Chain, Protocol } from '@/types/jumper-backend';
+import type { Token as LifiToken } from '@lifi/sdk';
 import type { DefiPosition } from '@/utils/positions/type-guards';
 
 /**
@@ -123,4 +124,7 @@ export interface OrchestrationState {
 
   /** Refresh data source by address */
   refreshByAddress: (address: string) => void;
+
+  /** Fetch fresh balances for specific tokens and merge into cache, bypassing getWalletBalances indexer lag */
+  refreshForTokens: (address: string, tokens: LifiToken[]) => Promise<void>;
 }