[PAY-2238] Improve balance display with recovery process

packages/common/src/hooks/useUSDCBalance.ts

@@ -55,6 +55,7 @@ export const useUSDCBalance = ({
     }
   }, [audiusBackend, setData])
 
+  // Refresh balance on mount
   useEffect(() => {
     refresh()
   }, [refresh])
@@ -69,5 +70,13 @@ export const useUSDCBalance = ({
     clearInterval(id)
   }, [id])
 
-  return { balanceStatus, recoveryStatus, data, refresh, cancelPolling }
+  // If we haven't loaded the balance yet for the first time or we're
+  // actively recovering, then we will be in loading state.
+  const status =
+    balanceStatus === Status.IDLE ||
+    (balanceStatus === Status.LOADING && data === null) ||
+    recoveryStatus === Status.LOADING
+      ? Status.LOADING
+      : balanceStatus
+  return { status, data, refresh, cancelPolling }
 }

packages/common/src/store/buy-usdc/sagas.ts

@@ -46,6 +46,8 @@ import {
 } from './slice'
 import { BuyUSDCError, BuyUSDCErrorCode } from './types'
 import { getBuyUSDCRemoteConfig, getUSDCUserBank } from './utils'
+import { setUSDCBalance } from 'store/wallet/slice'
+import { StringUSDC } from 'models/Wallet'
 
 type PurchaseStepParams = {
   desiredAmount: number
@@ -424,6 +426,17 @@ function* recoverPurchaseIfNecessary() {
       return
     }
 
+    const userBankAccountInfo = yield* call(
+      getTokenAccountInfo,
+      audiusBackendInstance,
+      {
+        tokenAccount: userBank,
+        mint: 'usdc'
+      }
+    )
+
+    const userBankInitialBalance = userBankAccountInfo?.amount ?? BigInt(0)
+
     const userBankAddress = userBank.toBase58()
 
     // Transfer all USDC from the from the root wallet to the user bank
@@ -480,6 +493,21 @@ function* recoverPurchaseIfNecessary() {
       }
     )
 
+    const updatedBalance = yield* call(
+      pollForTokenBalanceChange,
+      audiusBackendInstance,
+      {
+        tokenAccount: userBank,
+        mint: 'usdc',
+        initialBalance: userBankInitialBalance,
+        maxRetryCount: TRANSACTION_RETRY_COUNT
+      }
+    )
+
+    yield* put(
+      setUSDCBalance({ amount: updatedBalance.toString() as StringUSDC })
+    )
+
     yield* put(recoveryStatusChanged({ status: Status.SUCCESS }))
     yield* call(
       track,

packages/mobile/src/components/premium-track-purchase-drawer/hooks/usePurchaseContentFormState.ts

@@ -1,10 +1,7 @@
-import { useEffect } from 'react'
-
 import {
-  purchaseContentSelectors,
   isContentPurchaseInProgress,
-  useUSDCBalance,
-  Status
+  purchaseContentSelectors,
+  useUSDCBalance
 } from '@audius/common'
 import { useSelector } from 'react-redux'
 
@@ -22,14 +19,7 @@ export const usePurchaseContentFormState = ({ price }: { price: number }) => {
   const error = useSelector(getPurchaseContentError)
   const isUnlocking = !error && isContentPurchaseInProgress(stage)
 
-  const { data: currentBalance, recoveryStatus, refresh } = useUSDCBalance()
-
-  // Refresh balance on successful recovery
-  useEffect(() => {
-    if (recoveryStatus === Status.SUCCESS) {
-      refresh()
-    }
-  }, [recoveryStatus, refresh])
+  const { data: currentBalance } = useUSDCBalance()
 
   const purchaseSummaryValues = usePurchaseSummaryValues({
     price,

packages/mobile/src/components/usdc-balance-pill/USDCBalancePill.tsx

@@ -33,18 +33,15 @@ const useStyles = makeStyles(({ spacing, palette }) => ({
 
 export const USDCBalancePill = () => {
   const styles = useStyles()
-  const { data: usdcBalance, balanceStatus: usdcBalanceStatus } =
-    useUSDCBalance()
-  const isUsdcBalanceLoading =
-    usdcBalance === null || usdcBalanceStatus === Status.LOADING
+  const { data: usdcBalance, status: usdcBalanceStatus } = useUSDCBalance()
   const balanceCents = formatUSDCWeiToFloorCentsNumber(
     (usdcBalance ?? new BN(0)) as BNUSDC
   )
   const usdcBalanceFormatted = formatCurrencyBalance(balanceCents / 100)
   return (
     <View style={styles.root}>
       <LogoUSDC height={spacing(5)} width={spacing(5)} />
-      {isUsdcBalanceLoading ? (
+      {usdcBalanceStatus === Status.LOADING ? (
         <Skeleton
           style={styles.amount}
           height={spacing(4.5)}

packages/web/src/components/premium-content-purchase-modal/hooks/usePurchaseContentFormState.ts

@@ -1,10 +1,7 @@
-import { useEffect } from 'react'
-
 import {
   purchaseContentSelectors,
   isContentPurchaseInProgress,
-  useUSDCBalance,
-  Status
+  useUSDCBalance
 } from '@audius/common'
 import { useSelector } from 'react-redux'
 
@@ -22,14 +19,7 @@ export const usePurchaseContentFormState = ({ price }: { price: number }) => {
   const error = useSelector(getPurchaseContentError)
   const isUnlocking = !error && isContentPurchaseInProgress(stage)
 
-  const { data: currentBalance, recoveryStatus, refresh } = useUSDCBalance()
-
-  // Refresh balance on successful recovery
-  useEffect(() => {
-    if (recoveryStatus === Status.SUCCESS) {
-      refresh()
-    }
-  }, [recoveryStatus, refresh])
+  const { data: currentBalance } = useUSDCBalance()
 
   const purchaseSummaryValues = usePurchaseSummaryValues({
     price,

packages/web/src/components/usdc-balance-pill/USDCBalancePill.tsx

@@ -19,8 +19,7 @@ type USDCPillProps = {
 }
 
 export const USDCBalancePill = ({ className }: USDCPillProps) => {
-  const { data: balance, balanceStatus: usdcBalanceStatus } = useUSDCBalance()
-  const isLoading = balance === null || usdcBalanceStatus === Status.LOADING
+  const { data: balance, status: usdcBalanceStatus } = useUSDCBalance()
   const balanceCents = formatUSDCWeiToFloorCentsNumber(
     (balance ?? new BN(0)) as BNUSDC
   )
@@ -29,7 +28,7 @@ export const USDCBalancePill = ({ className }: USDCPillProps) => {
   return (
     <div className={cn(styles.container, className)}>
       <Icon className={styles.icon} icon={LogoUSDC} size='medium' />
-      {isLoading ? (
+      {usdcBalanceStatus === Status.LOADING ? (
         <Skeleton className={styles.skeleton} />
       ) : (
         <span className={styles.amount}>${balanceFormatted}</span>

packages/web/src/components/withdraw-usdc-modal/components/TransferSuccessful.tsx

@@ -60,7 +60,7 @@ export const TransferSuccessful = ({
   priorBalanceCents: number
   onClickDone: () => void
 }) => {
-  const { data: balance, balanceStatus } = useUSDCBalance()
+  const { data: balance, status: balanceStatus } = useUSDCBalance()
   const signature = useSelector(getWithdrawTransaction)
   const balanceNumber = formatUSDCWeiToFloorCentsNumber(
     (balance ?? new BN(0)) as BNUSDC

packages/web/src/pages/dashboard-page/DashboardPage.tsx

@@ -1,13 +1,6 @@
 import { useState, Suspense, ReactNode, useEffect, useCallback } from 'react'
 
-import {
-  Status,
-  Track,
-  formatCount,
-  themeSelectors,
-  combineStatuses,
-  useUSDCBalance
-} from '@audius/common'
+import { Status, Track, formatCount, themeSelectors } from '@audius/common'
 import cn from 'classnames'
 import { each } from 'lodash'
 import moment, { Moment } from 'moment'
@@ -77,12 +70,6 @@ export const DashboardPage = () => {
   const listenData = useSelector(getDashboardListenData)
   const dashboardStatus = useSelector(getDashboardStatus)
   const theme = useSelector(getTheme)
-  const { data: balance, balanceStatus } = useUSDCBalance()
-  const statuses = [dashboardStatus]
-  if (balance === null) {
-    statuses.push(balanceStatus)
-  }
-  const status = combineStatuses(statuses)
 
   const header = <Header primary={messages.title} />
 
@@ -187,10 +174,7 @@ export const DashboardPage = () => {
       contentClassName={styles.pageContainer}
       header={header}
     >
-      {!account ||
-      balance === null ||
-      !listenData ||
-      status === Status.LOADING ? (
+      {!account || !listenData || dashboardStatus === Status.LOADING ? (
         <LoadingSpinner className={styles.spinner} />
       ) : (
         <>

packages/web/src/pages/pay-and-earn-page/PayAndEarnPage.tsx

@@ -1,6 +1,6 @@
 import { useEffect } from 'react'
 
-import { Status, buyUSDCActions, useUSDCBalance } from '@audius/common'
+import { buyUSDCActions } from '@audius/common'
 import { useDispatch } from 'react-redux'
 
 import { useIsMobile } from 'utils/clientUtil'
@@ -12,14 +12,6 @@ import { PayAndEarnPageProps } from './types'
 export const PayAndEarnPage = (props: PayAndEarnPageProps) => {
   const isMobile = useIsMobile()
   const dispatch = useDispatch()
-  const { recoveryStatus, refresh } = useUSDCBalance()
-
-  // Refresh balance on successful recovery
-  useEffect(() => {
-    if (recoveryStatus === Status.SUCCESS) {
-      refresh()
-    }
-  }, [recoveryStatus, refresh])
 
   // Always check for recoverable USDC on page load
   useEffect(() => {

packages/web/src/pages/pay-and-earn-page/components/USDCCard.module.css

@@ -48,8 +48,8 @@
 }
 
 .spinner {
-  height: var(--unit-5);
-  width: var(--unit-5);
+  height: var(--unit-8);
+  width: var(--unit-8);
   align-self: center;
 }
 

packages/web/src/pages/pay-and-earn-page/components/USDCCard.tsx

@@ -7,7 +7,9 @@ import {
   formatCurrencyBalance,
   formatUSDCWeiToFloorCentsNumber,
   useWithdrawUSDCModal,
-  useAddFundsModal
+  useAddFundsModal,
+  useUSDCBalance,
+  Status
 } from '@audius/common'
 import { Button, PlainButton, IconQuestionCircle, Flex } from '@audius/harmony'
 import { LogoUSDC } from '@audius/stems'
@@ -34,15 +36,10 @@ const messages = {
   withdrawalHistory: 'Withdrawal History'
 }
 
-export const USDCCard = ({
-  balance,
-  refreshing = false
-}: {
-  balance: BNUSDC
-  refreshing?: boolean
-}) => {
+export const USDCCard = () => {
   const { onOpen: openWithdrawUSDCModal } = useWithdrawUSDCModal()
   const { onOpen: openAddFundsModal } = useAddFundsModal()
+  const { data: balance, status: balanceStatus } = useUSDCBalance()
 
   const balanceCents = formatUSDCWeiToFloorCentsNumber(
     (balance ?? new BN(0)) as BNUSDC
@@ -93,15 +90,18 @@ export const USDCCard = ({
           </div>
 
           <Flex gap='m'>
-            {refreshing && <LoadingSpinner className={styles.spinner} />}
-            <Text
-              variant='heading'
-              color='staticWhite'
-              strength='strong'
-              size='xxLarge'
-            >
-              ${balanceFormatted}
-            </Text>
+            {balanceStatus === Status.LOADING ? (
+              <LoadingSpinner className={styles.spinner} />
+            ) : (
+              <Text
+                variant='heading'
+                color='staticWhite'
+                strength='strong'
+                size='xxLarge'
+              >
+                ${balanceFormatted}
+              </Text>
+            )}
           </Flex>
         </div>
         <div className={styles.usdcInfo}>
@@ -117,12 +117,22 @@ export const USDCCard = ({
       </div>
       <div className={styles.withdrawContainer}>
         <div className={styles.addFundsButton}>
-          <Button variant='secondary' fullWidth onClick={handleAddFunds}>
+          <Button
+            variant='secondary'
+            fullWidth
+            onClick={handleAddFunds}
+            disabled={balanceStatus === Status.LOADING}
+          >
             {messages.addFunds}
           </Button>
         </div>
         <div className={styles.withdrawButton}>
-          <Button variant='secondary' fullWidth onClick={handleWithdraw}>
+          <Button
+            variant='secondary'
+            fullWidth
+            onClick={handleWithdraw}
+            disabled={balanceStatus === Status.LOADING}
+          >
             {messages.withdraw}
           </Button>
         </div>

packages/web/src/pages/pay-and-earn-page/desktop/PayAndEarnPage.tsx

@@ -1,6 +1,6 @@
 import { useCallback, useEffect, useState } from 'react'
 
-import { useUSDCBalance, accountSelectors, Status } from '@audius/common'
+import { accountSelectors } from '@audius/common'
 import {
   Button,
   Flex,
@@ -42,7 +42,6 @@ type TableMetadata = {
 
 export const PayAndEarnPage = ({ tableView }: PayAndEarnPageProps) => {
   const dispatch = useDispatch()
-  const { data: balance, balanceStatus, recoveryStatus } = useUSDCBalance()
   const accountHasTracks = useSelector(getAccountHasTracks)
 
   const [tableOptions, setTableOptions] = useState<TableType[] | null>(null)
@@ -135,17 +134,11 @@ export const PayAndEarnPage = ({ tableView }: PayAndEarnPageProps) => {
       contentClassName={styles.pageContainer}
       header={header}
     >
-      {!tableOptions || !selectedTable || balance === null ? (
+      {!tableOptions || !selectedTable ? (
         <LoadingSpinner className={styles.spinner} />
       ) : (
         <>
-          <USDCCard
-            balance={balance}
-            refreshing={
-              balanceStatus === Status.LOADING ||
-              recoveryStatus === Status.LOADING
-            }
-          />
+          <USDCCard />
           <Paper w='100%'>
             <Flex direction='column' w='100%'>
               <Flex