diff --git a/packages/common/src/api/tan-query/wallets/useAudioBalance.ts b/packages/common/src/api/tan-query/wallets/useAudioBalance.ts index d7b5b013c15..94cabd077ed 100644 --- a/packages/common/src/api/tan-query/wallets/useAudioBalance.ts +++ b/packages/common/src/api/tan-query/wallets/useAudioBalance.ts @@ -4,7 +4,8 @@ import { queryOptions, useQueries, useQuery, - type QueryFunctionContext + type QueryFunctionContext, + useQueryClient } from '@tanstack/react-query' import { call, getContext } from 'typed-redux-saga' import { getAddress } from 'viem' @@ -449,3 +450,47 @@ export const invalidateAudioBalance = ({ }) queryClient.invalidateQueries({ queryKey }) } + +/** + * Helper function to poll the audio balance until it changes from the original value. + * @param queryClient + * @param splWallet - The Solana wallet address + * @param maxAttempts - Maximum number of polling attempts (default: 10) + * @param delayMs - Delay between polling attempts in milliseconds (default: 300) + */ +export const pollUntilAudioBalanceChanges = async ( + queryClient: ReturnType, + splWallet: string, + maxAttempts = 10, + delayMs = 300 +): Promise => { + const queryKey = getWalletAudioBalanceQueryKey({ + address: splWallet, + chain: Chain.Sol, + includeStaked: false + }) + const originalBalance = queryClient.getQueryData(queryKey) + + for (let attempt = 0; attempt < maxAttempts; attempt++) { + // Wait before polling (except on first attempt where we check immediately) + if (attempt > 0) { + await new Promise((resolve) => setTimeout(resolve, delayMs)) + } + + // Invalidate and refetch the balance + await queryClient.invalidateQueries({ queryKey }) + await queryClient.refetchQueries({ queryKey }) + + // Check if the balance has changed + const newBalance = queryClient.getQueryData(queryKey) + if (newBalance !== originalBalance) { + return + } + } + + // If we've exhausted all attempts, log a warning but don't throw + // The balance will eventually update on the next stale query refetch + console.warn( + `Audio balance polled viia pollUntilAudioBalanceChanges but the value was not changed after ${maxAttempts} attempts` + ) +} diff --git a/packages/web/src/hooks/useClaimFees.ts b/packages/web/src/hooks/useClaimFees.ts index 3a15965c5a9..b622d0a8104 100644 --- a/packages/web/src/hooks/useClaimFees.ts +++ b/packages/web/src/hooks/useClaimFees.ts @@ -3,7 +3,7 @@ import { getArtistCoinQueryKey, useCurrentAccountUser, useQueryContext, - QUERY_KEYS + pollUntilAudioBalanceChanges } from '@audius/common/api' import { Feature } from '@audius/common/models' import { createUserBankIfNeeded } from '@audius/common/services' @@ -105,7 +105,7 @@ export const useClaimFees = ( } }) }, - onSuccess: (data, variables, context) => { + onSuccess: async (data, variables, context) => { // Optimistically update the unclaimed fees data const queryKey = getArtistCoinQueryKey(variables.tokenMint) queryClient.setQueryData(queryKey, (existingCoin) => { @@ -118,14 +118,14 @@ export const useClaimFees = ( } } }) - - // Invalidate audio balance queries to refresh user's AUDIO balance - queryClient.invalidateQueries({ - queryKey: [QUERY_KEYS.audioBalance] - }) - // Call the original onSuccess if provided options?.onSuccess?.(data, variables, context) + + // Poll audio balance queries until the balance actually changes + // The reason we want to do polling logic here is we dont actually have the value + if (currentUser?.spl_wallet) { + await pollUntilAudioBalanceChanges(queryClient, currentUser.spl_wallet) + } } }) }