Namadillo: Optimistically updating staking rewards balance (#1115)

* feat(namadillo): adding optimistic updates to staking rewards claiming

* refactor(namadillo): making claim rewards object clearer

* refactor(namadillo): using account?.address as a dependency instead of account only

Author: pedrorezende

apps/namadillo/src/atoms/staking/atoms.ts

@@ -96,21 +96,6 @@ export const createWithdrawTxAtomFamily = atomFamily((id: string) => {
   });
 });
 
-export const claimRewardsAtom = atomWithMutation((get) => {
-  const chain = get(chainAtom);
-  return {
-    mutationKey: ["create-claim-tx"],
-    enabled: chain.isSuccess,
-    mutationFn: async ({
-      params,
-      gasConfig,
-      account,
-    }: BuildTxAtomParams<ClaimRewardsMsgValue>) => {
-      return createClaimTx(chain.data!, account, params, gasConfig);
-    },
-  };
-});
-
 export const claimableRewardsAtom = atomWithQuery<AddressBalance>((get) => {
   const account = get(defaultAccountAtom);
   const chainParameters = get(chainParametersAtom);
@@ -134,6 +119,21 @@ export const claimableRewardsAtom = atomWithQuery<AddressBalance>((get) => {
   };
 });
 
+export const claimRewardsAtom = atomWithMutation((get) => {
+  const chain = get(chainAtom);
+  return {
+    mutationKey: ["create-claim-tx"],
+    enabled: chain.isSuccess,
+    mutationFn: async ({
+      params,
+      gasConfig,
+      account,
+    }: BuildTxAtomParams<ClaimRewardsMsgValue>) => {
+      return createClaimTx(chain.data!, account, params, gasConfig);
+    },
+  };
+});
+
 export const claimAndStakeRewardsAtom = atomWithMutation((get) => {
   const chain = get(chainAtom);
   const claimableRewards = get(claimableRewardsAtom);

apps/namadillo/src/atoms/staking/services.ts

@@ -12,6 +12,7 @@ import {
   WithdrawProps,
   WrapperTxProps,
 } from "@namada/types";
+import { queryClient } from "App/Common/QueryProvider";
 import { getSdkInstance } from "hooks";
 import { TransactionPair, buildTxPair } from "lib/query";
 import { Address, AddressBalance, ChainSettings, GasConfig } from "types";
@@ -159,3 +160,11 @@ export const createClaimAndStakeTx = async (
     account.address
   );
 };
+
+export const clearClaimRewards = (accountAddress: string): void => {
+  const emptyClaimRewards = {};
+  queryClient.setQueryData(
+    ["claim-rewards", accountAddress],
+    () => emptyClaimRewards
+  );
+};

apps/namadillo/src/hooks/useTransactionCallbacks.tsx

@@ -1,29 +1,36 @@
-import { accountBalanceAtom } from "atoms/accounts";
+import { accountBalanceAtom, defaultAccountAtom } from "atoms/accounts";
 import { shouldUpdateBalanceAtom, shouldUpdateProposalAtom } from "atoms/etc";
-import { claimableRewardsAtom } from "atoms/staking";
+import { claimableRewardsAtom, clearClaimRewards } from "atoms/staking";
 import { useAtomValue, useSetAtom } from "jotai";
 import { useTransactionEventListener } from "utils";
 
 export const useTransactionCallback = (): void => {
   const { refetch: refetchBalances } = useAtomValue(accountBalanceAtom);
   const { refetch: refetchRewards } = useAtomValue(claimableRewardsAtom);
+  const { data: account } = useAtomValue(defaultAccountAtom);
   const shouldUpdateBalance = useSetAtom(shouldUpdateBalanceAtom);
 
   const onBalanceUpdate = (): void => {
     // TODO: refactor this after event subscription is enabled on indexer
     shouldUpdateBalance(true);
     refetchBalances();
-    refetchRewards();
 
     const timePolling = 6 * 1000;
     setTimeout(() => shouldUpdateBalance(false), timePolling);
+
+    if (account?.address) {
+      clearClaimRewards(account.address);
+      setTimeout(() => refetchRewards(), timePolling);
+    }
   };
 
   useTransactionEventListener("Bond.Success", onBalanceUpdate);
   useTransactionEventListener("Unbond.Success", onBalanceUpdate);
   useTransactionEventListener("Withdraw.Success", onBalanceUpdate);
   useTransactionEventListener("Redelegate.Success", onBalanceUpdate);
-  useTransactionEventListener("ClaimRewards.Success", onBalanceUpdate);
+  useTransactionEventListener("ClaimRewards.Success", onBalanceUpdate, [
+    account?.address,
+  ]);
 
   const shouldUpdateProposal = useSetAtom(shouldUpdateProposalAtom);
 

apps/namadillo/src/utils/index.ts

@@ -36,14 +36,15 @@ export const proposalIdToString = (proposalId: bigint): string =>
 
 export const useTransactionEventListener = <T extends keyof WindowEventMap>(
   event: T,
-  handler: (this: Window, ev: WindowEventMap[T]) => void
+  handler: (this: Window, ev: WindowEventMap[T]) => void,
+  deps: React.DependencyList = []
 ): void => {
   useEffect(() => {
     window.addEventListener(event, handler);
     return () => {
       window.removeEventListener(event, handler);
     };
-  }, []);
+  }, deps);
 };
 
 const secondsToDateTime = (seconds: bigint): DateTime =>