feat: use reward state for reward components (#1492)

* chore(BABY): Keep the useRewardState to reward components

---------

Co-authored-by: wjrjerome <[email protected]>

Author: jonybur

src/ui/baby/components/RewardPreviewModal/index.tsx

@@ -5,9 +5,10 @@ import {
   DialogHeader,
   Text,
 } from "@babylonlabs-io/core-ui";
-import { PropsWithChildren, useEffect, useState } from "react";
+import { useEffect, useState } from "react";
 
 import babylon from "@/infrastructure/babylon";
+import { useRewardState } from "@/ui/baby/state/RewardState";
 import { ResponsiveDialog } from "@/ui/common/components/Modals/ResponsiveDialog";
 import { getNetworkConfigBBN } from "@/ui/common/config/network/bbn";
 import { useCosmosWallet } from "@/ui/common/context/wallet/CosmosWalletProvider";
@@ -19,27 +20,17 @@ import { maxDecimals } from "@/ui/common/utils/maxDecimals";
 
 const { coinSymbol, displayUSD } = getNetworkConfigBBN();
 
-interface PreviewModalProps {
-  open: boolean;
-  processing?: boolean;
-  title: string;
-  onClose: () => void;
-  onProceed: () => void;
-  rewards: Array<{ validatorAddress: string }>;
-  totalReward: bigint;
-}
-
-export const RewardsPreviewModal = ({
-  open,
-  processing = false,
-  title,
-  onClose,
-  onProceed,
-  rewards,
-  totalReward,
-}: PropsWithChildren<PreviewModalProps>) => {
+export const RewardsPreviewModal = () => {
   const { bech32Address } = useCosmosWallet();
   const { estimateBbnGasFee } = useBbnTransaction();
+  const {
+    totalReward,
+    rewards,
+    showClaimModal: open,
+    loading: processing,
+    closeClaimModal: onClose,
+    claimAll: onProceed,
+  } = useRewardState();
   const babyPrice = usePrice(coinSymbol);
   const [feeAmount, setFeeAmount] = useState<number>(0);
   const [feeLoading, setFeeLoading] = useState<boolean>(false);
@@ -51,7 +42,7 @@ export const RewardsPreviewModal = ({
 
       setFeeLoading(true);
       try {
-        const msgs = rewards.map((reward) =>
+        const msgs = rewards.map((reward: { validatorAddress: string }) =>
           babylon.txs.baby.createClaimRewardMsg({
             validatorAddress: reward.validatorAddress,
             delegatorAddress: bech32Address,
@@ -85,7 +76,7 @@ export const RewardsPreviewModal = ({
     : undefined;
   return (
     <ResponsiveDialog open={open} onClose={onClose}>
-      <DialogHeader title={title} className="text-accent-primary" />
+      <DialogHeader title="Claim Rewards" className="text-accent-primary" />
 
       <DialogBody className="no-scrollbar mb-[40px] mt-8 flex max-h-[calc(100vh-12rem)] flex-col gap-[40px] overflow-y-auto text-accent-primary">
         <div className="flex flex-col gap-4">

src/ui/baby/hooks/api/useEpochPolling.ts

@@ -1,4 +1,4 @@
-import { useQuery, useQueryClient } from "@tanstack/react-query";
+import { useQueryClient } from "@tanstack/react-query";
 import { useEffect, useRef } from "react";
 
 import babylon from "@/infrastructure/babylon";
@@ -9,54 +9,52 @@ import { usePendingOperationsService } from "../services/usePendingOperationsSer
 import { BABY_DELEGATIONS_KEY } from "./useDelegations";
 import { BABY_UNBONDING_DELEGATIONS_KEY } from "./useUnbondingDelegations";
 
-const BABYLON_CURRENT_EPOCH_KEY = "BABYLON_CURRENT_EPOCH";
-
 export function useEpochPolling(address?: string) {
   const queryClient = useQueryClient();
   const previousEpochRef = useRef<number | undefined>(undefined);
   const { cleanupAllPendingOperationsFromStorage } =
     usePendingOperationsService();
-
-  const { data: currentEpoch } = useQuery<number, Error>({
-    queryKey: [BABYLON_CURRENT_EPOCH_KEY],
-    queryFn: async () => {
-      const client = await babylon.client();
-      const { currentEpoch } = await client.baby.getCurrentEpoch();
-      return currentEpoch;
-    },
-    refetchInterval: ONE_MINUTE,
-    // Only poll if we have an address (user is connected)
-    enabled: Boolean(address),
-  });
+  const cleanupRef = useRef(cleanupAllPendingOperationsFromStorage);
+  cleanupRef.current = cleanupAllPendingOperationsFromStorage;
 
   useEffect(() => {
-    if (currentEpoch === undefined) return;
-
-    if (previousEpochRef.current === undefined) {
-      previousEpochRef.current = currentEpoch;
-      return;
-    }
-
-    if (currentEpoch !== previousEpochRef.current) {
-      // Clean up all pending operations from localStorage (everything from
-      // previous epochs is finalized)
-      cleanupAllPendingOperationsFromStorage();
-
-      // Invalidate all delegation queries since epoch change affects entire
-      // blockchain state
-      queryClient.invalidateQueries({
-        queryKey: [BABY_DELEGATIONS_KEY],
-      });
-      // Also invalidate unbonding delegations since they're epoch-dependent
-      queryClient.invalidateQueries({
-        queryKey: [BABY_UNBONDING_DELEGATIONS_KEY],
-      });
-      previousEpochRef.current = currentEpoch;
-    }
-  }, [
-    currentEpoch,
-    queryClient,
-    address,
-    cleanupAllPendingOperationsFromStorage,
-  ]);
+    if (!address) return;
+
+    let cancelled = false;
+
+    const checkEpoch = async () => {
+      try {
+        const client = await babylon.client();
+        const { currentEpoch } = await client.baby.getCurrentEpoch();
+
+        if (previousEpochRef.current === undefined) {
+          previousEpochRef.current = currentEpoch;
+          return;
+        }
+
+        if (!cancelled && currentEpoch !== previousEpochRef.current) {
+          cleanupRef.current?.();
+
+          queryClient.invalidateQueries({
+            queryKey: [BABY_DELEGATIONS_KEY],
+            refetchType: "active",
+          });
+          queryClient.invalidateQueries({
+            queryKey: [BABY_UNBONDING_DELEGATIONS_KEY],
+            refetchType: "active",
+          });
+          previousEpochRef.current = currentEpoch;
+        }
+      } catch {
+        // ignore transient errors
+      }
+    };
+
+    checkEpoch();
+    const id = setInterval(checkEpoch, ONE_MINUTE);
+    return () => {
+      cancelled = true;
+      clearInterval(id);
+    };
+  }, [address, queryClient]);
 }

src/ui/baby/layout.tsx

@@ -2,7 +2,7 @@ import { useWalletConnect } from "@babylonlabs-io/wallet-connector";
 import { useEffect, useState } from "react";
 
 import { DelegationState } from "@/ui/baby/state/DelegationState";
-import { RewardState, useRewardState } from "@/ui/baby/state/RewardState";
+import { RewardState } from "@/ui/baby/state/RewardState";
 import { StakingState } from "@/ui/baby/state/StakingState";
 import { ValidatorState } from "@/ui/baby/state/ValidatorState";
 import { AuthGuard } from "@/ui/common/components/Common/AuthGuard";
@@ -53,28 +53,11 @@ function BabyLayoutContent() {
   useEpochPolling(bech32Address);
 
   const RewardsTab: React.FC = () => {
-    const {
-      showClaimModal,
-      closeClaimModal,
-      claimAll,
-      loading: rewardLoading,
-      totalReward,
-      rewards,
-    } = useRewardState();
-
     return (
       <Section>
         <div className="h-[500px]">
           <RewardCard />
-          <RewardsPreviewModal
-            open={showClaimModal}
-            processing={rewardLoading}
-            title="Claim Rewards"
-            totalReward={totalReward}
-            rewards={rewards}
-            onClose={closeClaimModal}
-            onProceed={claimAll}
-          />
+          <RewardsPreviewModal />
         </div>
       </Section>
     );