Add polling for audio balance changes

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<typeof useQueryClient>,
+  splWallet: string,
+  maxAttempts = 10,
+  delayMs = 300
+): Promise<void> => {
+  const queryKey = getWalletAudioBalanceQueryKey({
+    address: splWallet,
+    chain: Chain.Sol,
+    includeStaked: false
+  })
+  const originalBalance = queryClient.getQueryData<AudioWei>(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<AudioWei>(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`
+  )
+}

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<Coin>(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)
+      }
     }
   })
 }