feat: Babylon SDK Staking (#34)

* feat: baby staking

* chore: getBabyStakingDelegations

* chore: createWithdrawRewardForBTCStakingMsg

* chore: createClaimRewardForBABYStakingMsg

* chore: createStakeBABYMsg

* chore: createUnstakeBABYMsg

* chore: msg withdraw reward for btc staking

* chore: msg stake baby

* chore: msg unstake baby

* chore: msg withdraw reward for baby staking

* chore: claim -> withdraw

Author: gbarkhatov

packages/babylon-proto-ts/README.md

@@ -31,8 +31,15 @@ const sdk = createBabylonSDK({ rpcUrl: "https://babylon-rpc.example.com" });
 // Connect the client
 await sdk.connect();
 
-// Query rewards for an address
-const rewards = await sdk.client.getRewards("bbn1...");
+// Query BTC staking rewards for an address
+const btcRewards = await sdk.client.getRewardsForBTCStaking("bbn1...");
+
+// Query BABY staking delegations and rewards
+const delegations = await sdk.client.getBabyStakingDelegations("bbn1...");
+const babyRewards = await sdk.client.getRewardsForBABYStaking("bbn1...");
+
+// Get all validators
+const validators = await sdk.client.getValidators();
 
 // Query balance
 const balance = await sdk.client.getBalance("bbn1...", "ubbn");
@@ -61,11 +68,22 @@ const client = await SigningStargateClient.connectWithSigner(
   },
 );
 
-// Create messages using the SDK
-const withdrawMsg = sdk.messages.createWithdrawRewardMsg("bbn1...");
+// BTC Staking: Withdraw BTC staking rewards message
+const btcWithdrawMsg =
+  sdk.messages.createWithdrawRewardForBTCStakingMsg("bbn1...");
+
+// BABY Staking: Stake BABY tokens message
+const stakeMsg = sdk.messages.createStakeBABYMsg(
+  "bbn1...", // delegator address
+  "bbnvalidator...", // validator address
+  { denom: "ubbn", amount: "1000000" }, // 1 BBN
+);
 
-// Sign and broadcast
-const result = await client.signAndBroadcast("bbn1...", [withdrawMsg], "auto");
+// BABY Staking: Withdraw BABY staking rewards message
+const withdrawMsg = sdk.messages.createWithdrawRewardForBABYStakingMsg(
+  "bbn1...",
+  "bbnvalidator...",
+);
 ```
 
 ### Direct Protobuf Access
@@ -100,21 +118,21 @@ Initializes the client connection to the Babylon network.
 
 ### Client Methods (via sdk.client)
 
-#### `sdk.client.getRewards(address: string): Promise<number>`
+#### `sdk.client.getRewardsForBTCStaking(address: string): Promise<number>`
 
-Retrieves the total rewards for a given address.
+Retrieves the total BTC staking rewards for a given address.
 
 - **Parameters:**
   - `address`: The Babylon address to query
-- **Returns:** Total rewards amount (number)
+- **Returns:** Total BTC staking rewards amount (number)
 
 #### `sdk.client.getBalance(address: string, denom?: string): Promise<number>`
 
 Gets the balance of a specific token for an address.
 
 - **Parameters:**
   - `address`: The Babylon address to query
-  - `denom`: Token denomination (defaults to "ubbn")
+  - `denom`: Token denomination (defaults to `"ubbn"`)
 - **Returns:** Balance amount (number)
 
 #### `sdk.client.getBTCTipHeight(): Promise<number>`
@@ -123,16 +141,67 @@ Retrieves the current Bitcoin blockchain tip height.
 
 - **Returns:** Bitcoin tip height (number)
 
+#### `sdk.client.getBabyStakingDelegations(address: string): Promise<DelegationResponse[]>`
+
+Gets all BABY token delegations for a given address.
+
+- **Parameters:**
+  - `address`: The Babylon address to query
+- **Returns:** Array of delegation responses
+
+#### `sdk.client.getRewardsForBABYStaking(address: string): Promise<DelegationDelegatorReward[]>`
+
+Retrieves all BABY staking delegation rewards for a given address.
+
+- **Parameters:**
+  - `address`: The Babylon address to query
+- **Returns:** Array of delegation rewards
+
+#### `sdk.client.getValidators(): Promise<Validator[]>`
+
+Gets all validators in the network.
+
+- **Returns:** Array of all validators
+
 ### SDK Messages
 
-#### `sdk.messages.createWithdrawRewardMsg(address: string)`
+#### `sdk.messages.createWithdrawRewardForBTCStakingMsg(address: string)`
 
 Creates a message for withdrawing rewards from Bitcoin staking.
 
 - **Parameters:**
   - `address`: The Babylon address to withdraw rewards for
 - **Returns:** Message object with proper typeUrl and value, ready for signing and broadcasting
 
+#### `sdk.messages.createStakeBABYMsg(delegatorAddress: string, validatorAddress: string, amount: Coin)`
+
+Creates a message for staking BABY tokens.
+
+- **Parameters:**
+  - `delegatorAddress`: The delegator's Babylon address
+  - `validatorAddress`: The validator's Babylon address
+  - `amount`: The amount to stake (Coin object with amount and denom)
+- **Returns:** Message object with proper typeUrl and value, ready for signing and broadcasting
+
+#### `sdk.messages.createUnstakeBABYMsg(delegatorAddress: string, validatorAddress: string, amount: Coin)`
+
+Creates a message for unstaking BABY tokens.
+
+- **Parameters:**
+  - `delegatorAddress`: The delegator's Babylon address
+  - `validatorAddress`: The validator's Babylon address
+  - `amount`: The amount to unstake (Coin object with amount and denom)
+- **Returns:** Message object with proper typeUrl and value, ready for signing and broadcasting
+
+#### `sdk.messages.createWithdrawRewardForBABYStakingMsg(delegatorAddress: string, validatorAddress: string)`
+
+Creates a message for withdrawing BABY staking rewards.
+
+- **Parameters:**
+  - `delegatorAddress`: The delegator's Babylon address
+  - `validatorAddress`: The validator's Babylon address
+- **Returns:** Message object with proper typeUrl and value, ready for signing and broadcasting
+
 ### SDK Utilities
 
 #### `sdk.utils.createRegistry(): Registry`
@@ -143,6 +212,22 @@ Creates a CosmJS registry with all Babylon message types registered.
 
 Creates amino types for Babylon messages, required for hardware wallet compatibility.
 
+#### `sdk.utils.normalizeCosmjsAmount(amount: string): string`
+
+Normalizes CosmJS amount from 18-decimal precision to standard `ubbn` format.
+
+- **Parameters:**
+  - `amount`: The amount string from CosmJS (with 18 decimal precision)
+- **Returns:** The normalized amount as a string in standard `ubbn` format
+
+#### `sdk.utils.normalizeRewardResponse(response: QueryDelegationTotalRewardsResponse)`
+
+Normalizes reward response from CosmJS to standard `ubbn` format.
+
+- **Parameters:**
+  - `response`: The QueryDelegationTotalRewardsResponse from CosmJS
+- **Returns:** The normalized reward array with amounts in standard `ubbn` format
+
 ### Protobuf Exports
 
 The library exports all generated protobuf types directly, allowing for advanced use cases

packages/babylon-proto-ts/src/constants.ts

@@ -2,5 +2,9 @@ export const REWARD_GAUGE_KEY_BTC_DELEGATION = "BTC_STAKER";
 export const BTC_STAKER = "btc_staker";
 export const REGISTRY_TYPE_URLS = {
   MsgCreateBTCDelegation: "/babylon.btcstaking.v1.MsgCreateBTCDelegation",
-  MsgWithdrawReward: "/babylon.incentive.MsgWithdrawReward",
+  MsgWithdrawRewardForBTCStaking: "/babylon.incentive.MsgWithdrawReward",
+  MsgStakeBABY: "/babylon.epoching.v1.MsgWrappedDelegate",
+  MsgUnstakeBABY: "/babylon.epoching.v1.MsgWrappedUndelegate",
+  MsgWithdrawRewardForBABYStaking:
+    "/cosmos.distribution.v1beta1.MsgWithdrawDelegatorReward",
 };

packages/babylon-proto-ts/src/lib/client.ts

@@ -1,20 +1,31 @@
 import {
   BankExtension,
   Coin,
+  DistributionExtension,
   QueryClient,
+  StakingExtension,
   createProtobufRpcClient,
   setupBankExtension,
+  setupDistributionExtension,
+  setupStakingExtension,
 } from "@cosmjs/stargate";
 import { Tendermint34Client } from "@cosmjs/tendermint-rpc";
+import { DelegationDelegatorReward } from "cosmjs-types/cosmos/distribution/v1beta1/distribution";
+import {
+  DelegationResponse,
+  Validator,
+} from "cosmjs-types/cosmos/staking/v1beta1/staking";
 
 import { REWARD_GAUGE_KEY_BTC_DELEGATION } from "../constants";
 import * as btclightclientquery from "../generated/babylon/btclightclient/v1/query";
 import * as incentivequery from "../generated/babylon/incentive/query";
 
 interface Clients {
-  incentive: incentivequery.QueryClientImpl;
-  btcLight: btclightclientquery.QueryClientImpl;
-  bank: BankExtension["bank"];
+  incentive: incentivequery.QueryClientImpl | null;
+  btcLight: btclightclientquery.QueryClientImpl | null;
+  bank: BankExtension["bank"] | null;
+  staking: StakingExtension["staking"] | null;
+  distribution: DistributionExtension["distribution"] | null;
 }
 
 type ClientNames = keyof Clients;
@@ -26,13 +37,17 @@ export class BabylonClient {
     incentive: null,
     btcLight: null,
     bank: null,
+    staking: null,
+    distribution: null,
   };
 
   constructor(rpcUrl: string) {
     this.rpcUrl = rpcUrl;
   }
 
-  private getClient<K extends ClientNames>(clientName: K): Clients[K] {
+  private getClient<K extends ClientNames>(
+    clientName: K,
+  ): NonNullable<Clients[K]> {
     if (!this.clients[clientName]) {
       throw Error("Babylon Client not initialized");
     }
@@ -45,20 +60,25 @@ export class BabylonClient {
     const queryClient = QueryClient.withExtensions(
       tmClient,
       setupBankExtension,
+      setupStakingExtension,
+      setupDistributionExtension,
     );
     const rpc = createProtobufRpcClient(queryClient);
 
     this.clients.incentive = new incentivequery.QueryClientImpl(rpc);
     this.clients.btcLight = new btclightclientquery.QueryClientImpl(rpc);
     this.clients.bank = setupBankExtension(queryClient).bank;
+    this.clients.staking = setupStakingExtension(queryClient).staking;
+    this.clients.distribution =
+      setupDistributionExtension(queryClient).distribution;
   }
 
   /**
-   * Gets the rewards of an address in the Babylon chain.
+   * [BTC Staking] Gets the rewards of the user's account.
    * @param {string} address - The address to get the rewards of.
    * @returns {Promise<number>} - The rewards of the address.
    */
-  async getRewards(address: string): Promise<number> {
+  async getRewardsForBTCStaking(address: string): Promise<number> {
     try {
       const req = incentivequery.QueryRewardGaugesRequest.fromPartial({
         address,
@@ -137,6 +157,62 @@ export class BabylonClient {
       });
     }
   }
+
+  /**
+   * [BABY Staking] Gets all delegations of the user's account.
+   * @param {string} address - The address to get the delegations of.
+   * @returns {Promise<DelegationResponse[]>} - The delegations of the address.
+   */
+  async getBabyStakingDelegations(
+    address: string,
+  ): Promise<DelegationResponse[]> {
+    try {
+      const response =
+        await this.getClient("staking").delegatorDelegations(address);
+      return response.delegationResponses || [];
+    } catch (error) {
+      throw new Error(`Failed to fetch delegations for ${address}`, {
+        cause: error,
+      });
+    }
+  }
+
+  /**
+   * [BABY Staking] Gets all delegation rewards of the user's account.
+   * @param {string} address - The address to get the delegation rewards of.
+   * @returns {Promise<DelegationDelegatorReward[]>} - The delegation rewards of the address.
+   */
+  async getRewardsForBABYStaking(
+    address: string,
+  ): Promise<DelegationDelegatorReward[]> {
+    try {
+      const response =
+        await this.getClient("distribution").delegationTotalRewards(address);
+      return response.rewards || [];
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("no delegation")) {
+        return [];
+      }
+      throw new Error(`Failed to fetch delegation rewards for ${address}`, {
+        cause: error,
+      });
+    }
+  }
+
+  /**
+   * [BABY Staking] Gets all the validators.
+   * @returns {Promise<Validator[]>} - All validators.
+   */
+  async getValidators(): Promise<Validator[]> {
+    try {
+      const response = await this.getClient("staking").validators("");
+      return response.validators || [];
+    } catch (error) {
+      throw new Error(`Failed to fetch validators`, {
+        cause: error,
+      });
+    }
+  }
 }
 
 export const createBabylonClient = (rpcUrl: string) =>