refactor: remove user balance store dependency in deposit gauge reward feature

apps/main/src/dex/components/PagePool/components/TransferActions.tsx

@@ -3,7 +3,6 @@ import { useConnection } from 'wagmi'
 import FormConnectWallet from '@/dex/components/FormConnectWallet'
 import AlertSeedAmounts from '@/dex/components/PagePool/components/AlertSeedAmounts'
 import type { TransferProps } from '@/dex/components/PagePool/types'
-import { useSignerAddress } from '@/dex/entities/signer'
 import { usePoolIdByAddressOrId } from '@/dex/hooks/usePoolIdByAddressOrId'
 import { usePoolTokenBalances } from '@/dex/hooks/usePoolTokenBalances'
 import useTokenAlert from '@/dex/hooks/useTokenAlert'
@@ -23,7 +22,7 @@ const TransferActions = ({
   loading?: boolean
   children: ReactNode
 } & Pick<TransferProps, 'poolData' | 'poolDataCacheOrApi' | 'routerParams' | 'seed'>) => {
-  const { data: signerAddress } = useSignerAddress()
+  const { address: signerAddress } = useConnection()
   const { rChainId, rPoolIdOrAddress } = routerParams
   const poolId = usePoolIdByAddressOrId({ chainId: rChainId, poolIdOrAddress: rPoolIdOrAddress })
   const alert = useTokenAlert(poolData?.tokenAddressesAll ?? [])

apps/main/src/dex/entities/gauge/model/estimate-gas-options.ts

@@ -10,13 +10,14 @@ import { gaugeKeys } from './query-keys'
 import { depositRewardAvailable, depositRewardIsApproved } from './query-options'
 
 export const estimateGasDepositRewardApprove = queryFactory({
-  queryKey: ({ rewardTokenId, amount, ...gaugeParams }: DepositRewardApproveParams) =>
+  queryKey: ({ rewardTokenId, amount, userBalance, ...gaugeParams }: DepositRewardApproveParams) =>
     [
       ...rootKeys.gauge({ ...gaugeParams }),
       ...gaugeKeys.estimateGas(),
       'depositRewardApprove',
       { rewardTokenId },
       { amount },
+      { userBalance },
     ] as const,
   queryFn: api.queryEstimateGasDepositRewardApprove,
   refetchInterval: '1m',
@@ -43,14 +44,15 @@ export const estimateGasAddRewardToken = queryFactory({
 })
 
 export const estimateGasDepositReward = queryFactory({
-  queryKey: ({ rewardTokenId, amount, epoch, ...gaugeParams }: DepositRewardParams) =>
+  queryKey: ({ rewardTokenId, amount, epoch, userBalance, ...gaugeParams }: DepositRewardParams) =>
     [
       ...rootKeys.gauge({ ...gaugeParams }),
       ...gaugeKeys.estimateGas(),
       'depositReward',
       { rewardTokenId },
       { amount },
       { epoch },
+      { userBalance },
     ] as const,
   queryFn: api.queryEstimateGasDepositReward,
   refetchInterval: '1m',

apps/main/src/dex/entities/gauge/model/gauge-validation.ts

@@ -1,5 +1,4 @@
 import { enforce, group, test } from 'vest'
-import useStore from '@/dex/store/useStore'
 import { formatNumber } from '@ui/utils'
 import { t } from '@ui-kit/lib/i18n'
 import { TIME_FRAMES } from '@ui-kit/lib/model'
@@ -15,20 +14,23 @@ import { AddRewardParams, DepositRewardApproveParams, DepositRewardParams } from
 export const gaugeAddRewardValidationGroup = ({ distributorId, rewardTokenId }: AddRewardParams) =>
   group('gaugeAddRewardValidationGroup', () => {
     test('distributorId', () => addressValidationFn(distributorId))
-
     test('rewardTokenId', () => tokenIdValidationFn(rewardTokenId))
   })
 
-export const gaugeDepositRewardApproveValidationGroup = ({ rewardTokenId, amount }: DepositRewardApproveParams) =>
+export const gaugeDepositRewardApproveValidationGroup = ({
+  rewardTokenId,
+  amount,
+  userBalance,
+}: DepositRewardApproveParams) =>
   group('gaugeDepositRewardApproveValidationGroup', () => {
     test('rewardTokenId', () => tokenIdValidationFn(rewardTokenId))
-    test('amount', () => validateAmount({ rewardTokenId, amount }))
+    test('amount', () => validateAmount({ rewardTokenId, amount, userBalance }))
   })
 
-export const gaugeDepositRewardValidationGroup = ({ rewardTokenId, amount, epoch }: DepositRewardParams) =>
+export const gaugeDepositRewardValidationGroup = ({ rewardTokenId, amount, epoch, userBalance }: DepositRewardParams) =>
   group('gaugeDepositRewardValidationGroup', () => {
     test('rewardTokenId', () => tokenIdValidationFn(rewardTokenId))
-    test('amount', () => validateAmount({ rewardTokenId, amount }))
+    test('amount', () => validateAmount({ rewardTokenId, amount, userBalance }))
     test('epoch', () => {
       enforce(epoch)
         .message(t`Epoch is required`)
@@ -55,18 +57,17 @@ export const gaugeDepositRewardValidationSuite = createValidationSuite((data: De
   gaugeDepositRewardValidationGroup(data)
 })
 
-function validateAmount({ rewardTokenId, amount }: DepositRewardApproveParams) {
+function validateAmount({ rewardTokenId, amount, userBalance }: DepositRewardApproveParams) {
   amountValidationFn(amount)
   if (!rewardTokenId || !amount) return
 
-  const state = useStore.getState()
-  const userBalancesMapper = state.userBalances.userBalancesMapper
-  const tokenBalance = userBalancesMapper[rewardTokenId]
-
-  if (!tokenBalance) return
+  enforce(userBalance).condition((userBalance) => ({
+    pass: userBalance != null && +userBalance > 0,
+    message: t`Wallet balance is zero for the selected reward token`,
+  }))
 
   enforce(amount).condition((amount) => ({
-    pass: +amount < +tokenBalance,
-    message: t`Amount ${formatNumber(amount, { decimals: 5 })} > wallet balance ${formatNumber(tokenBalance, { decimals: 5 })}`,
+    pass: +amount <= +userBalance!,
+    message: t`Amount ${formatNumber(amount, { decimals: 5 })} > wallet balance ${formatNumber(userBalance, { decimals: 5 })}`,
   }))
 }

apps/main/src/dex/entities/gauge/types.ts

@@ -1,6 +1,7 @@
 import type { Address } from 'viem'
 import { GaugeQuery } from '@ui-kit/lib/model/query'
 import { FieldsOf } from '@ui-kit/lib/validation'
+import type { Decimal } from '@ui-kit/utils'
 
 export type AddReward = {
   rewardTokenId: Address
@@ -13,6 +14,7 @@ export type AddRewardMutation = FieldsOf<AddReward>
 export type DepositRewardApprove = {
   rewardTokenId: Address
   amount: number | string
+  userBalance?: Decimal
 }
 export type DepositRewardApproveQuery = GaugeQuery & DepositRewardApprove
 export type DepositRewardApproveParams = FieldsOf<DepositRewardApproveQuery>

apps/main/src/dex/features/add-gauge-reward-token/ui/AddRewardToken.tsx

@@ -1,10 +1,10 @@
 import { useCallback } from 'react'
 import { FormProvider, useForm } from 'react-hook-form'
 import { zeroAddress } from 'viem'
+import { useConnection } from 'wagmi'
 import AlertFormError from '@/dex/components/AlertFormError'
 import { useAddRewardToken, useGaugeRewardsDistributors, useIsDepositRewardAvailable } from '@/dex/entities/gauge'
 import { useNetworkByChain } from '@/dex/entities/networks'
-import { useSignerAddress } from '@/dex/entities/signer'
 import { addGaugeRewardTokenValidationSuite } from '@/dex/features/add-gauge-reward-token/model'
 import type { AddRewardFormValues, AddRewardTokenProps } from '@/dex/features/add-gauge-reward-token/types'
 import {
@@ -22,7 +22,7 @@ import { t } from '@ui-kit/lib/i18n'
 import { formDefaultOptions } from '@ui-kit/lib/model/form'
 
 export const AddRewardToken = ({ chainId, poolId }: AddRewardTokenProps) => {
-  const { data: signerAddress } = useSignerAddress()
+  const { address: signerAddress } = useConnection()
 
   const { isFetching: isFetchingGaugeRewardsDistributors } = useGaugeRewardsDistributors({ chainId, poolId })
 

apps/main/src/dex/features/deposit-gauge-reward/model/default-values.ts

@@ -5,6 +5,7 @@ import { TIME_FRAMES } from '@ui-kit/lib/model'
 export const DepositRewardDefaultValues: DepositRewardFormValues = {
   rewardTokenId: zeroAddress,
   amount: '',
+  userBalance: undefined,
   epoch: TIME_FRAMES.WEEK,
   step: DepositRewardStep.APPROVAL,
 } as const

apps/main/src/dex/features/deposit-gauge-reward/types/index.ts

@@ -1,4 +1,5 @@
-import { Address } from 'viem'
+import type { Address } from 'viem'
+import type { Decimal } from '@ui-kit/utils'
 
 export enum DepositRewardStep {
   APPROVAL = 'APPROVAL',
@@ -9,6 +10,7 @@ export enum DepositRewardStep {
 export interface DepositRewardFormValues {
   rewardTokenId?: Address
   amount?: string
+  userBalance?: Decimal
   epoch?: number
   step: DepositRewardStep
 }

apps/main/src/dex/features/deposit-gauge-reward/ui/ActionsStepper.tsx

@@ -32,6 +32,7 @@ export const DepositStepper = ({ chainId, poolId }: { chainId: ChainId; poolId:
   const amount = watch('amount')
   const rewardTokenId = watch('rewardTokenId')
   const step = watch('step')
+  const userBalance = watch('userBalance')
 
   const { mutate: depositRewardApprove, isPending: isPendingDepositRewardApprove } = useDepositRewardApprove({
     chainId,
@@ -59,6 +60,7 @@ export const DepositStepper = ({ chainId, poolId }: { chainId: ChainId; poolId:
       {
         rewardTokenId: getValues('rewardTokenId'),
         amount: getValues('amount'),
+        userBalance: getValues('userBalance'),
       },
       { onSuccess: onApproveSuccess, onError: onApproveError },
     )
@@ -70,6 +72,7 @@ export const DepositStepper = ({ chainId, poolId }: { chainId: ChainId; poolId:
         rewardTokenId: getValues('rewardTokenId'),
         amount: getValues('amount'),
         epoch: getValues('epoch'),
+        userBalance: getValues('userBalance'),
       },
       {
         onSuccess: (data: string) => {
@@ -91,6 +94,7 @@ export const DepositStepper = ({ chainId, poolId }: { chainId: ChainId; poolId:
     poolId,
     rewardTokenId,
     amount,
+    userBalance,
   })
 
   useLayoutEffect(() => {

apps/main/src/dex/features/deposit-gauge-reward/ui/AmountTokenInput.tsx

@@ -2,13 +2,13 @@ import { MouseEvent, useCallback, useMemo } from 'react'
 import { useFormContext } from 'react-hook-form'
 import { Address, isAddressEqual } from 'viem'
 import { ethAddress } from 'viem'
+import { useConnection } from 'wagmi'
 import {
   useDepositRewardApproveIsMutating,
   useDepositRewardIsMutating,
   useGaugeRewardsDistributors,
 } from '@/dex/entities/gauge'
 import { useNetworkByChain } from '@/dex/entities/networks'
-import { useIsSignerConnected, useSignerAddress, useTokensBalances } from '@/dex/entities/signer'
 import { type DepositRewardFormValues, DepositRewardStep } from '@/dex/features/deposit-gauge-reward/types'
 import {
   FlexItemAmount,
@@ -26,26 +26,20 @@ import { formatNumber } from '@ui/utils'
 import { type TokenOption, TokenSelector } from '@ui-kit/features/select-token'
 import { t } from '@ui-kit/lib/i18n'
 import { useTokenUsdRates } from '@ui-kit/lib/model/entities/token-usd-rate'
+import { useTokenBalances } from '@ui-kit/queries/token-balance.query'
 
 export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId: string }) => {
   const { setValue, getValues, formState, watch } = useFormContext<DepositRewardFormValues>()
   const rewardTokenId = watch('rewardTokenId')
   const amount = watch('amount')
   const epoch = watch('epoch')
 
-  const { data: signerAddress } = useSignerAddress()
-  const { data: haveSigner } = useIsSignerConnected()
+  const { address: signerAddress } = useConnection()
   const isMaxLoading = useStore((state) => state.quickSwap.isMaxLoading)
   const {
     data: { networkId },
   } = useNetworkByChain({ chainId })
 
-  const userBalancesMapper = useStore((state) => state.userBalances.userBalancesMapper)
-  const userTokens = Object.entries(userBalancesMapper)
-    .filter(([, balance]) => parseFloat(balance ?? '0') > 0)
-    .map(([address]) => address)
-  const { data: tokenPrices } = useTokenUsdRates({ chainId, tokenAddresses: userTokens })
-
   const { tokensMapper } = useTokensMapper(chainId)
 
   const { data: rewardDistributors, isPending: isPendingRewardDistributors } = useGaugeRewardsDistributors({
@@ -56,11 +50,6 @@ export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId
   const isMutatingDepositRewardApprove = useDepositRewardApproveIsMutating({ chainId, poolId, rewardTokenId, amount })
   const isMutatingDepositReward = useDepositRewardIsMutating({ chainId, poolId, rewardTokenId, amount, epoch })
 
-  const {
-    data: [tokenBalance],
-    isLoading: isTokenBalancesLoading,
-  } = useTokensBalances([rewardTokenId])
-
   const filteredTokens = useMemo<TokenOption[]>(() => {
     if (isPendingRewardDistributors || !rewardDistributors || !signerAddress) return []
 
@@ -89,6 +78,19 @@ export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId
   }, [isPendingRewardDistributors, rewardDistributors, signerAddress, tokensMapper, getValues, networkId, setValue])
 
   const token = filteredTokens.find((x) => x.address === rewardTokenId)
+  const tokenAddresses = filteredTokens.map((t) => t.address)
+
+  const { data: tokenPrices } = useTokenUsdRates({ chainId, tokenAddresses })
+  const { data: tokenBalances, isLoading: isTokenBalancesLoading } = useTokenBalances({
+    chainId,
+    userAddress: signerAddress,
+    tokenAddresses,
+  })
+
+  const rewardTokenBalance = useMemo(
+    () => rewardTokenId && tokenBalances && tokenBalances[rewardTokenId],
+    [rewardTokenId, tokenBalances],
+  )
 
   const onChangeAmount = useCallback(
     (amount: string) => {
@@ -109,10 +111,10 @@ export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId
   const onMaxButtonClick = useCallback(
     (e?: MouseEvent<HTMLButtonElement>) => {
       e?.preventDefault()
-      if (!tokenBalance) return
-      setValue('amount', tokenBalance, { shouldValidate: true })
+      if (!rewardTokenBalance) return
+      setValue('amount', rewardTokenBalance, { shouldValidate: true })
     },
-    [tokenBalance, setValue],
+    [rewardTokenBalance, setValue],
   )
 
   const isDisabled = isMutatingDepositReward || isMutatingDepositRewardApprove
@@ -129,10 +131,10 @@ export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId
             id="deposit-amount"
             type="number"
             labelProps={
-              haveSigner && {
+              signerAddress && {
                 label: t`Avail.`,
                 descriptionLoading: isTokenBalancesLoading,
-                description: formatNumber(tokenBalance, { decimals: 5 }),
+                description: formatNumber(rewardTokenBalance),
               }
             }
             testId="deposit-amount"
@@ -154,7 +156,7 @@ export const AmountTokenInput = ({ chainId, poolId }: { chainId: ChainId; poolId
             selectedToken={token}
             tokens={filteredTokens}
             disabled={isDisabled}
-            balances={userBalancesMapper}
+            balances={tokenBalances}
             tokenPrices={tokenPrices}
             onToken={onChangeToken}
           />

apps/main/src/dex/features/deposit-gauge-reward/ui/DepositReward.tsx

@@ -1,4 +1,6 @@
+import { useEffect } from 'react'
 import { FormProvider, useForm } from 'react-hook-form'
+import { useConnection } from 'wagmi'
 import AlertFormError from '@/dex/components/AlertFormError'
 import { useGaugeRewardsDistributors } from '@/dex/entities/gauge'
 import { DepositRewardDefaultValues, depositRewardValidationSuite } from '@/dex/features/deposit-gauge-reward/model'
@@ -16,6 +18,7 @@ import { FormErrorsDisplay } from '@ui/FormErrorsDisplay'
 import { BlockSkeleton } from '@ui/skeleton'
 import { FormContainer, FormFieldsContainer, GroupedFieldsContainer } from '@ui/styled-containers'
 import { formDefaultOptions } from '@ui-kit/lib/model/form'
+import { useTokenBalance } from '@ui-kit/queries/token-balance.query'
 
 export const DepositReward = ({ chainId, poolId }: { chainId: ChainId; poolId: string }) => {
   const { isPending: isPendingRewardDistributors } = useGaugeRewardsDistributors({
@@ -29,6 +32,15 @@ export const DepositReward = ({ chainId, poolId }: { chainId: ChainId; poolId: s
     defaultValues: DepositRewardDefaultValues,
   })
 
+  const rewardTokenId = methods.watch('rewardTokenId')
+  const { address: userAddress } = useConnection()
+  const { data: userBalance } = useTokenBalance({ chainId, userAddress, tokenAddress: rewardTokenId })
+
+  // Sync userBalance from query into form for validation
+  useEffect(() => {
+    methods.setValue('userBalance', userBalance, { shouldValidate: true })
+  }, [userBalance, methods])
+
   if (isPendingRewardDistributors) {
     return <BlockSkeleton height={440} />
   }