fix: shield plan token approval amount incorrect when plan change

Author: tuna1207

ui/hooks/subscription/useSubscriptionPricing.ts

@@ -43,7 +43,11 @@ export const useAvailableTokenBalances = (params: {
   paymentChains?: ChainPaymentInfo[];
   price?: ProductPrice;
   productType: ProductType;
-}): TokenWithApprovalAmount[] => {
+}): {
+  availableTokenBalances: TokenWithApprovalAmount[];
+  pending: boolean;
+  error: Error | undefined;
+} => {
   const { paymentChains, price, productType } = params;
 
   const paymentChainIds = useMemo(
@@ -70,10 +74,6 @@ export const useAvailableTokenBalances = (params: {
   // Poll and update evm balances for payment chains
   pollAndUpdateEvmBalances({ chainIds: paymentChainIds });
 
-  const [availableTokenBalances, setAvailableTokenBalances] = useState<
-    TokenWithApprovalAmount[]
-  >([]);
-
   const validTokenBalances = useMemo(() => {
     return evmBalances.filter((token) => {
       const supportedTokensForChain =
@@ -96,72 +96,74 @@ export const useAvailableTokenBalances = (params: {
     });
   }, [evmBalances, paymentChainTokenMap]);
 
-  useEffect(() => {
+  const {
+    value: availableTokenBalances,
+    pending,
+    error,
+  } = useAsyncResult(async (): Promise<TokenWithApprovalAmount[]> => {
     if (!price || !paymentChainTokenMap) {
-      return;
+      return [];
     }
 
-    const getAvailableTokenBalances = async () => {
-      const availableTokens: TokenWithApprovalAmount[] = [];
-
-      const cryptoApprovalAmounts = await Promise.all(
-        validTokenBalances.map((token) => {
-          const tokenPaymentInfo = paymentChainTokenMap?.[
-            token.chainId as Hex
-          ]?.find(
-            (t) => t.address.toLowerCase() === token.address.toLowerCase(),
+    const availableTokens: TokenWithApprovalAmount[] = [];
+
+    const cryptoApprovalAmounts = await Promise.all(
+      validTokenBalances.map((token) => {
+        const tokenPaymentInfo = paymentChainTokenMap?.[
+          token.chainId as Hex
+        ]?.find((t) => t.address.toLowerCase() === token.address.toLowerCase());
+        if (!tokenPaymentInfo) {
+          log.error(
+            '[useAvailableTokenBalances] tokenPaymentInfo not found',
+            token,
           );
-          if (!tokenPaymentInfo) {
-            log.error(
-              '[useAvailableTokenBalances] tokenPaymentInfo not found',
-              token,
-            );
-            return null;
-          }
-          return getSubscriptionCryptoApprovalAmount({
-            chainId: token.chainId as Hex,
-            paymentTokenAddress: token.address as Hex,
-            productType,
-            interval: price.interval,
-          });
-        }),
-      );
-
-      cryptoApprovalAmounts.forEach((amount, index) => {
-        const token = validTokenBalances[index];
-        if (!token.balance) {
-          return;
-        }
-        // NOTE: we are using stable coin for subscription atm, so we need to scale the balance by the decimals
-        const scaledFactor = 10n ** 6n;
-        const scaledBalance =
-          BigInt(Math.round(Number(token.balance) * Number(scaledFactor))) /
-          scaledFactor;
-        const tokenHasEnoughBalance =
-          amount &&
-          scaledBalance * BigInt(10 ** token.decimals) >=
-            BigInt(amount.approveAmount);
-        if (tokenHasEnoughBalance) {
-          availableTokens.push({
-            ...token,
-            approvalAmount: {
-              approveAmount: amount.approveAmount,
-              chainId: token.chainId as Hex,
-              paymentAddress: amount.paymentAddress,
-              paymentTokenAddress: amount.paymentTokenAddress,
-            },
-            type: token.isNative ? AssetType.native : AssetType.token,
-          } as TokenWithApprovalAmount);
+          return null;
         }
-      });
-
-      setAvailableTokenBalances(availableTokens);
-    };
+        return getSubscriptionCryptoApprovalAmount({
+          chainId: token.chainId as Hex,
+          paymentTokenAddress: token.address as Hex,
+          productType,
+          interval: price.interval,
+        });
+      }),
+    );
+
+    cryptoApprovalAmounts.forEach((amount, index) => {
+      const token = validTokenBalances[index];
+      if (!token.balance) {
+        return;
+      }
+      // NOTE: we are using stable coin for subscription atm, so we need to scale the balance by the decimals
+      const scaledFactor = 10n ** 6n;
+      const scaledBalance =
+        BigInt(Math.round(Number(token.balance) * Number(scaledFactor))) /
+        scaledFactor;
+      const tokenHasEnoughBalance =
+        amount &&
+        scaledBalance * BigInt(10 ** token.decimals) >=
+          BigInt(amount.approveAmount);
+      if (tokenHasEnoughBalance) {
+        availableTokens.push({
+          ...token,
+          approvalAmount: {
+            approveAmount: amount.approveAmount,
+            chainId: token.chainId as Hex,
+            paymentAddress: amount.paymentAddress,
+            paymentTokenAddress: amount.paymentTokenAddress,
+          },
+          type: token.isNative ? AssetType.native : AssetType.token,
+        } as TokenWithApprovalAmount);
+      }
+    });
 
-    getAvailableTokenBalances();
-  }, [price, productType, paymentChainTokenMap, validTokenBalances]);
+    return availableTokens;
+  }, [price, paymentChainTokenMap, validTokenBalances]);
 
-  return availableTokenBalances;
+  return {
+    availableTokenBalances: availableTokenBalances ?? [],
+    pending,
+    error,
+  };
 };
 
 /**

ui/pages/shield-plan/shield-plan.tsx

@@ -135,11 +135,12 @@ const ShieldPlan = () => {
     return pricingPlans?.find((plan) => plan.interval === selectedPlan);
   }, [pricingPlans, selectedPlan]);
 
-  const availableTokenBalances = useAvailableTokenBalances({
-    paymentChains: cryptoPaymentMethod?.chains,
-    price: selectedProductPrice,
-    productType: PRODUCT_TYPES.SHIELD,
-  });
+  const { availableTokenBalances, pending: pendingAvailableTokenBalances } =
+    useAvailableTokenBalances({
+      paymentChains: cryptoPaymentMethod?.chains,
+      price: selectedProductPrice,
+      productType: PRODUCT_TYPES.SHIELD,
+    });
   const hasAvailableToken = availableTokenBalances.length > 0;
 
   const [selectedPaymentMethod, setSelectedPaymentMethod] =
@@ -165,17 +166,23 @@ const ShieldPlan = () => {
 
   // set selected token to the first available token if no token is selected
   useEffect(() => {
-    if (selectedToken || availableTokenBalances.length === 0) {
+    if (
+      pendingAvailableTokenBalances ||
+      selectedToken ||
+      availableTokenBalances.length === 0
+    ) {
       return;
     }
 
     const lastUsedPaymentToken = lastUsedPaymentDetails?.paymentTokenAddress;
     const lastUsedPaymentMethod = lastUsedPaymentDetails?.type;
+    const lastUsedPaymentPlan = lastUsedPaymentDetails?.plan;
 
     let lastUsedSelectedToken = availableTokenBalances[0];
     if (
       lastUsedPaymentToken &&
-      lastUsedPaymentMethod === PAYMENT_TYPES.byCrypto
+      lastUsedPaymentMethod === PAYMENT_TYPES.byCrypto &&
+      lastUsedPaymentPlan === selectedPlan
     ) {
       lastUsedSelectedToken =
         availableTokenBalances.find(
@@ -185,12 +192,19 @@ const ShieldPlan = () => {
 
     setSelectedToken(lastUsedSelectedToken);
   }, [
+    pendingAvailableTokenBalances,
     availableTokenBalances,
     selectedToken,
     setSelectedToken,
     lastUsedPaymentDetails,
+    selectedPlan,
   ]);
 
+  // reset selected token if selected plan changes
+  useEffect(() => {
+    setSelectedToken(undefined);
+  }, [selectedPlan, setSelectedToken]);
+
   // set default selected payment method to crypto if selected token available
   useEffect(() => {
     if (selectedToken) {