create staking manager instead of useMemo to hold in memory

babylonlabs-io · Feb 9, 2025 · f265c5c · f265c5c
1 parent 19f7e44
commit f265c5c
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 37 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -27,7 +27,8 @@
   "dependencies": {
     "@babylonlabs-io/babylon-proto-ts": "0.0.3-canary.5",
     "@babylonlabs-io/bbn-core-ui": "0.9.3",
-    "@babylonlabs-io/bbn-wallet-connect": "0.3.5",
+    "@babylonlabs-io/bbn-wallet-connect": "0.3.6",
+    "@babylonlabs-io/btc-staking-ts": "0.4.0-canary.12",
     "@bitcoin-js/tiny-secp256k1-asmjs": "2.2.3",
     "@bitcoinerlab/secp256k1": "^1.1.1",
     "@cosmjs/proto-signing": "^0.32.4",

diff --git a/src/app/hooks/services/useRegistrationService.ts b/src/app/hooks/services/useRegistrationService.ts
@@ -39,10 +39,10 @@ type RegistrationSigningType = Extract<
 
 const REGISTRATION_STEP_MAP: Record<RegistrationSigningType, RegistrationStep> =
   {
-    "staking-slashing": "registration-staking-slashing",
-    "unbonding-slashing": "registration-unbonding-slashing",
-    "proof-of-possession": "registration-proof-of-possession",
-    "create-btc-delegation-msg": "registration-sign-bbn",
+    [SigningType.STAKING_SLASHING]: "registration-staking-slashing",
+    [SigningType.UNBONDING_SLASHING]: "registration-unbonding-slashing",
+    [SigningType.PROOF_OF_POSSESSION]: "registration-proof-of-possession",
+    [SigningType.CREATE_BTC_DELEGATION_MSG]: "registration-sign-bbn",
   };
 
 export function useRegistrationService() {

diff --git a/src/app/hooks/services/useStakingManagerService.ts b/src/app/hooks/services/useStakingManagerService.ts
@@ -1,5 +1,5 @@
 import { BabylonBtcStakingManager } from "@babylonlabs-io/btc-staking-ts";
-import { useMemo } from "react";
+import { useCallback } from "react";
 
 import { useBTCWallet } from "@/app/context/wallet/BTCWalletProvider";
 import { useCosmosWallet } from "@/app/context/wallet/CosmosWalletProvider";
@@ -20,7 +20,7 @@ export const useStakingManagerService = () => {
 
   const versionedParams = networkInfo?.params.bbnStakingParams?.versions;
 
-  const btcStakingManager = useMemo(() => {
+  const createBtcStakingManager = useCallback(() => {
     if (
       !btcNetwork ||
       !cosmosConnected ||
@@ -61,7 +61,8 @@ export const useStakingManagerService = () => {
     signMessage,
     signBbnTx,
   ]);
+
   return {
-    btcStakingManager,
+    createBtcStakingManager,
   };
 };
diff --git a/src/app/hooks/services/useTransactionService.ts b/src/app/hooks/services/useTransactionService.ts
@@ -49,7 +49,8 @@ export const useTransactionService = () => {
 
   const tipHeight = useMemo(() => tipHeader?.height ?? 0, [tipHeader]);
 
-  const { btcStakingManager } = useStakingManagerService();
+  const { createBtcStakingManager } = useStakingManagerService();
+
   /**
    * Create the delegation EOI
    *
@@ -59,6 +60,8 @@ export const useTransactionService = () => {
    */
   const createDelegationEoi = useCallback(
     async (stakingInput: BtcStakingInputs, feeRate: number) => {
+      const btcStakingManager = createBtcStakingManager();
+
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -84,7 +87,13 @@ export const useTransactionService = () => {
         signedBabylonTx,
       };
     },
-    [availableUTXOs, bech32Address, btcStakingManager, stakerInfo, tipHeight],
+    [
+      availableUTXOs,
+      bech32Address,
+      createBtcStakingManager,
+      stakerInfo,
+      tipHeight,
+    ],
   );
 
   /**
@@ -96,6 +105,7 @@ export const useTransactionService = () => {
    */
   const estimateStakingFee = useCallback(
     (stakingInput: BtcStakingInputs, feeRate: number): number => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -113,7 +123,7 @@ export const useTransactionService = () => {
         feeRate,
       );
     },
-    [btcStakingManager, tipHeight, stakerInfo, availableUTXOs],
+    [createBtcStakingManager, tipHeight, stakerInfo, availableUTXOs],
   );
 
   /**
@@ -129,6 +139,7 @@ export const useTransactionService = () => {
       stakingHeight: number,
       stakingInput: BtcStakingInputs,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -154,7 +165,7 @@ export const useTransactionService = () => {
         signedBabylonTx,
       };
     },
-    [bech32Address, btcStakingManager, stakerInfo, tipHeight],
+    [bech32Address, createBtcStakingManager, stakerInfo, tipHeight],
   );
 
   /**
@@ -172,6 +183,7 @@ export const useTransactionService = () => {
       expectedTxHashHex: string,
       unsignedStakingTxHex: string,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -203,7 +215,7 @@ export const useTransactionService = () => {
     },
     [
       availableUTXOs,
-      btcStakingManager,
+      createBtcStakingManager,
       pushTx,
       refetchUTXOs,
       stakerInfo,
@@ -231,6 +243,7 @@ export const useTransactionService = () => {
         sigHex: string;
       }[],
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -252,7 +265,7 @@ export const useTransactionService = () => {
 
       await pushTx(signedUnbondingTx.toHex());
     },
-    [btcStakingManager, pushTx, stakerInfo, tipHeight],
+    [createBtcStakingManager, pushTx, stakerInfo, tipHeight],
   );
 
   /**
@@ -268,6 +281,7 @@ export const useTransactionService = () => {
       paramVersion: number,
       earlyUnbondingTxHex: string,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -285,7 +299,7 @@ export const useTransactionService = () => {
         );
       await pushTx(signedWithdrawalTx.toHex());
     },
-    [btcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
+    [createBtcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
   );
 
   /**
@@ -301,6 +315,7 @@ export const useTransactionService = () => {
       paramVersion: number,
       stakingTxHex: string,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -318,7 +333,7 @@ export const useTransactionService = () => {
         );
       await pushTx(signedWithdrawalTx.toHex());
     },
-    [btcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
+    [createBtcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
   );
 
   /**
@@ -334,6 +349,7 @@ export const useTransactionService = () => {
       paramVersion: number,
       slashingTxHex: string,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(
         btcStakingManager,
         stakingInput,
@@ -351,7 +367,7 @@ export const useTransactionService = () => {
         );
       await pushTx(signedWithdrawalTx.toHex());
     },
-    [btcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
+    [createBtcStakingManager, defaultFeeRate, pushTx, stakerInfo, tipHeight],
   );
 
   /**
@@ -361,8 +377,9 @@ export const useTransactionService = () => {
    */
   const subscribeToSigningSteps = useCallback(
     (callback: (step: SigningType) => void) => {
+      const btcStakingManager = createBtcStakingManager();
       if (!btcStakingManager) {
-        return () => {};
+        throw new Error("BTC Staking Manager not initialized");
       }
 
       btcStakingManager.on(StakingEventType.SIGNING, callback);
@@ -372,7 +389,7 @@ export const useTransactionService = () => {
         btcStakingManager.off(StakingEventType.SIGNING, callback);
       };
     },
-    [btcStakingManager],
+    [createBtcStakingManager],
   );
 
   return {

diff --git a/src/app/hooks/services/useV1TransactionService.ts b/src/app/hooks/services/useV1TransactionService.ts
@@ -2,6 +2,7 @@ import {
   BabylonBtcStakingManager,
   getUnbondingTxStakerSignature,
   TransactionResult,
+  VersionedStakingParams,
 } from "@babylonlabs-io/btc-staking-ts";
 import { Transaction } from "bitcoinjs-lib";
 import { useCallback, useMemo } from "react";
@@ -43,7 +44,7 @@ export function useV1TransactionService() {
   // The "tag" is not needed for withdrawal or unbonding transactions.
   const versionedParams = networkInfo?.params.bbnStakingParams?.versions;
 
-  const { btcStakingManager } = useStakingManagerService();
+  const { createBtcStakingManager } = useStakingManagerService();
 
   /**
    * Submit the unbonding transaction to babylon API for further processing
@@ -60,7 +61,11 @@ export function useV1TransactionService() {
       stakingHeight: number,
       stakingTxHex: string,
     ) => {
+      const btcStakingManager = createBtcStakingManager();
       validateCommonInputs(btcStakingManager, stakingInput, stakerBtcInfo);
+      if (!versionedParams?.length) {
+        throw new Error("Staking params not loaded");
+      }
 
       const stakingTx = Transaction.fromHex(stakingTxHex);
       // Check if this staking transaction is eligible for unbonding
@@ -73,11 +78,17 @@ export function useV1TransactionService() {
         });
       }
 
+      // Get the param version based on height
+      const { version: paramsVersion } = getBbnParamByBtcHeight(
+        stakingHeight,
+        versionedParams!,
+      );
+
       const { transaction: signedUnbondingTx } =
         await btcStakingManager!.createPartialSignedBtcUnbondingTransaction(
           stakerBtcInfo,
           stakingInput,
-          stakingHeight,
+          paramsVersion,
           stakingTx,
         );
       const stakerSignatureHex =
@@ -93,7 +104,7 @@ export function useV1TransactionService() {
         throw new Error(`Error submitting unbonding transaction: ${error}`);
       }
     },
-    [btcStakingManager, stakerBtcInfo],
+    [createBtcStakingManager, stakerBtcInfo],
   );
 
   /**
@@ -115,14 +126,17 @@ export function useV1TransactionService() {
       stakingTxHex: string,
       earlyUnbondingTxHex?: string,
     ) => {
-      validateCommonInputs(btcStakingManager, stakingInput, stakerBtcInfo);
-      if (!versionedParams?.length) {
-        throw new Error("Staking params not loaded");
-      }
+      const btcStakingManager = createBtcStakingManager();
+      validateCommonInputs(
+        btcStakingManager,
+        stakingInput,
+        stakerBtcInfo,
+        versionedParams,
+      );
       // Get the param version based on height
       const { version: paramVersion } = getBbnParamByBtcHeight(
         stakingHeight,
-        versionedParams,
+        versionedParams!,
       );
 
       validateStakingInput(stakingInput);
@@ -153,7 +167,13 @@ export function useV1TransactionService() {
 
       await pushTx(result.transaction.toHex());
     },
-    [btcStakingManager, defaultFeeRate, pushTx, stakerBtcInfo, versionedParams],
+    [
+      createBtcStakingManager,
+      defaultFeeRate,
+      pushTx,
+      stakerBtcInfo,
+      versionedParams,
+    ],
   );
 
   return {
@@ -172,6 +192,7 @@ const validateCommonInputs = (
   btcStakingManager: BabylonBtcStakingManager | null,
   stakingInput: BtcStakingInputs,
   stakerBtcInfo: { address: string; publicKeyNoCoordHex: string },
+  versionedParams?: VersionedStakingParams[],
 ) => {
   validateStakingInput(stakingInput);
   if (!btcStakingManager) {
@@ -180,4 +201,7 @@ const validateCommonInputs = (
   if (!stakerBtcInfo.address || !stakerBtcInfo.publicKeyNoCoordHex) {
     throw new Error("Staker info not initialized");
   }
+  if (!versionedParams?.length) {
+    throw new Error("Staking params not loaded");
+  }
 };