Add request token and wait has energy call

src/components/substrate/SubstrateTxButton.tsx

@@ -14,8 +14,10 @@ import { isFunction } from '@polkadot/util'
 import type { Signer } from '@polkadot/api/types'
 import { VoidFn } from '@polkadot/api/types'
 import { isEmptyStr, newLogger } from '@subsocial/utils'
+import { ESTIMATED_ENERGY_FOR_ONE_TX } from 'src/config/constants'
 import useExternalStorage from 'src/hooks/useExternalStorage'
 import useSignerExternalStorage from 'src/hooks/useSignerExternalStorage'
+import useWaitHasEnergy from 'src/hooks/useWaitHasEnergy'
 import messages from 'src/messages'
 import { useBuildSendEvent } from 'src/providers/AnalyticContext'
 import { useOpenCloseOnBoardingModal } from 'src/rtk/features/onBoarding/onBoardingHooks'
@@ -35,6 +37,7 @@ import {
   SIGNER_REFRESH_TOKEN_KEY,
   SIGNER_TOKEN_KEY,
 } from '../utils/OffchainSigner/ExternalStorage'
+import { requestToken } from '../utils/OffchainUtils'
 import { getWalletBySource } from '../wallets/supportedWallets'
 import styles from './SubstrateTxButton.module.sass'
 import useToggle from './useToggle'
@@ -102,6 +105,8 @@ function TxButton({
   const { api: subsocialApi } = useSubstrate()
   const openOnBoardingModal = useOpenCloseOnBoardingModal()
   const [isSending, , setIsSending] = useToggle(false)
+  const hasEnoughEnergy = useMyAccount(state => (state.energy ?? 0) >= ESTIMATED_ENERGY_FOR_ONE_TX)
+  const waitHasEnergy = useWaitHasEnergy()
 
   const { isMobile } = useResponsiveSize()
   const {
@@ -299,6 +304,12 @@ function TxButton({
         const keypairSigner = useMyAccount.getState().signer
         if (!keypairSigner) throw new Error('No account signer provided')
         account = keypairSigner
+
+        if (!hasEnoughEnergy) {
+          await requestToken({ address: accountId })
+          console.log('Waiting energy...')
+          await waitHasEnergy()
+        }
       } else {
         // use extension signer
         if (isMobile) {

src/components/utils/OffchainUtils.ts

@@ -173,3 +173,14 @@ export const updatePendingOrder = async ({
 
   return res.data
 }
+
+export async function requestToken({ address }: { address: string }) {
+  // make request token as pending transaction so websocket won't disconnect for 10 secs after request token
+  // this is to make energy subscription work
+  const res = await axios.post('/c/api/request-token', {
+    address,
+  })
+  const data = res.data
+  if (!data.success) throw new Error(data.message)
+  return res
+}

src/hooks/useWaitHasEnergy.ts

@@ -0,0 +1,67 @@
+import { useCallback, useEffect, useRef } from 'react'
+import { useMyAccount } from 'src/stores/my-account'
+import useWrapInRef from './useWrapInRef'
+
+const DEFAULT_TIMEOUT = 30_000
+export default function useWaitHasEnergy(
+  isUsingConnectedWallet?: boolean,
+  timeout = DEFAULT_TIMEOUT,
+) {
+  const hasEnergyResolvers = useRef<(() => void)[]>([])
+  const { address, energy, resubscribeEnergy } = useAccountSwitch(isUsingConnectedWallet)
+  const energyRef = useWrapInRef(energy)
+
+  const generateNewPromise = useCallback(() => {
+    return new Promise<void>((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        reject(new Error("You don't have enough energy to perform this action. Please try again"))
+        resubscribeEnergy()
+      }, timeout)
+      hasEnergyResolvers.current.push(() => {
+        clearTimeout(timeoutId)
+        resolve()
+      })
+    })
+  }, [timeout, resubscribeEnergy])
+
+  useEffect(() => {
+    if (!energy || energy <= 0 || hasEnergyResolvers.current.length === 0) return
+    function resolveAllPending() {
+      hasEnergyResolvers.current.forEach(resolve => resolve())
+      hasEnergyResolvers.current = []
+    }
+    resolveAllPending()
+  }, [energy, generateNewPromise])
+
+  useEffect(() => {
+    return () => {
+      hasEnergyResolvers.current.forEach(resolve => resolve())
+      hasEnergyResolvers.current = []
+    }
+  }, [address])
+
+  return () => {
+    // need to use ref because if not it can have stale energy value
+    return !energyRef.current ? generateNewPromise() : Promise.resolve()
+  }
+}
+
+function useAccountSwitch(isUsingConnectedWallet = false) {
+  const address = useMyAccount(state => state.address)
+  const energy = useMyAccount(state => state.energy)
+  const resubscribeEnergy = useMyAccount(state => state._subscribeEnergy)
+
+  const connectedWallet = useMyAccount(state => state.connectedWallet)
+  const resubscribeConnectedWalletEnergy = useMyAccount(
+    state => state._subscribeConnectedWalletEnergy,
+  )
+
+  const usedData = { address, energy, resubscribeEnergy }
+  if (isUsingConnectedWallet) {
+    usedData.address = connectedWallet?.address ?? null
+    usedData.energy = connectedWallet?.energy ?? 0
+    usedData.resubscribeEnergy = resubscribeConnectedWalletEnergy
+  }
+
+  return usedData
+}