imp(costaking): track baby staked to active validators

.github/workflows/ci.yml

@@ -313,6 +313,28 @@ jobs:
         run: |
           make test-e2e-cache-btc-stake-expansion
 
+  e2e-run-validator-jailing:
+    needs: [e2e-docker-build-babylon]
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Download babylon artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: babylond-${{ github.sha }}
+          path: /tmp
+      - name: Docker load babylond
+        run: |
+          docker load < /tmp/docker-babylond.tar.gz
+      - name: Cache Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: 1.23
+      - name: Run e2e TestValidatorJailingTestSuite
+        run: |
+          make test-e2e-cache-validator-jailing
+
   e2e-V2:
     needs: [e2e-docker-build-babylon]
     runs-on: ubuntu-22.04

CHANGELOG.md

@@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [#1762](https://github.com/babylonlabs-io/babylon/pull/1762) Add new test case on stake expansion
 - [#1764](https://github.com/babylonlabs-io/babylon/pull/1764) Add mergify yaml file for automatic backporting
 - [#1785](https://github.com/babylonlabs-io/babylon/pull/1785) CI reusable to v0.13.5 and golang lint version 2
+- [#1776](https://github.com/babylonlabs-io/babylon/pull/1776) Track baby staked to active validators only
 
 ### Bug fixes
 

Makefile

@@ -248,6 +248,7 @@ test-e2e-cache:
 	$(MAKE) test-e2e-cache-upgrade-v2
 	$(MAKE) test-e2e-cache-epoching-spam-prevention
 	$(MAKE) test-e2e-cache-btc-stake-expansion
+	$(MAKE) test-e2e-cache-validator-jailing
 
 clean-e2e:
 	docker container rm -f $(shell docker container ls -a -q) || true
@@ -295,6 +296,9 @@ test-e2e-cache-upgrade-v4:
 test-e2e-cache-btc-stake-expansion:
 	go test -run TestBTCStakeExpansionTestSuite -mod=readonly -timeout=60m -v $(PACKAGES_E2E) --tags=e2e
 
+test-e2e-cache-validator-jailing:
+	go test -run TestValidatorJailingTestSuite -mod=readonly -timeout=60m -v $(PACKAGES_E2E) --tags=e2e
+
 test-sim-nondeterminism:
 	@echo "Running non-determinism test..."
 	@go test -mod=readonly $(SIMAPP) -run TestAppStateDeterminism -Enabled=true \

app/keepers/keepers.go

@@ -544,7 +544,10 @@ func (ak *AppKeepers) InitKeepers(
 
 	// make ZoneConcierge and Monitor to subscribe to the epoching's hooks
 	ak.EpochingKeeper = *epochingKeeper.SetHooks(
-		epochingtypes.NewMultiEpochingHooks(ak.MonitorKeeper.Hooks()),
+		epochingtypes.NewMultiEpochingHooks(
+			ak.MonitorKeeper.Hooks(),
+			ak.CostakingKeeper.HookEpoching(),
+		),
 	)
 
 	// set up Checkpointing, BTCCheckpoint, and BTCLightclient keepers

app/upgrades/v4/upgrades.go

@@ -134,7 +134,7 @@ func InitializeCoStakerRwdsTracker(
 // Returns the total score of the co-staker rewards tracker
 func saveBABYStakersRwdTracker(ctx context.Context, cdc codec.BinaryCodec, costkStoreService corestoretypes.KVStoreService, stkKeeper *stkkeeper.Keeper, params costktypes.Params) (math.Int, error) {
 	totalScore := math.ZeroInt()
-	// Get all BABY stakers
+	// Get all BABY stakers that are staking to an active validator
 	babyStakers, err := getAllBABYStakers(ctx, stkKeeper)
 	if err != nil {
 		return totalScore, fmt.Errorf("failed to get all BABY stakers: %w", err)
@@ -153,36 +153,23 @@ func saveBABYStakersRwdTracker(ctx context.Context, cdc codec.BinaryCodec, costk
 	return totalScore, nil
 }
 
-// getAllBABYStakers retrieves all BABY stakers with pagination
+// getAllBABYStakers retrieves all BABY stakers by iterating only over active validators
 func getAllBABYStakers(ctx context.Context, stkKeeper *stkkeeper.Keeper) (map[string]math.Int, error) {
 	stkQuerier := stkkeeper.NewQuerier(stkKeeper)
 	babyStakers := make(map[string]math.Int)
 
-	// First get all validators
-	var nextKey []byte
-	for {
-		req := &stktypes.QueryValidatorsRequest{
-			Pagination: &query.PageRequest{
-				Key: nextKey,
-			},
-		}
-
-		res, err := stkQuerier.Validators(ctx, req)
-		if err != nil {
-			return nil, err
+	// Iterate directly over active validators (last validator powers)
+	err := stkKeeper.IterateLastValidatorPowers(ctx, func(valAddr sdk.ValAddress, power int64) bool {
+		// Get all delegations for this active validator
+		if err := getValidatorDelegations(ctx, stkQuerier, valAddr.String(), babyStakers); err != nil {
+			// Return true to stop iteration on error
+			return true
 		}
+		return false // continue iteration
+	})
 
-		// For each validator, get all delegations
-		for _, validator := range res.Validators {
-			if err := getValidatorDelegations(ctx, stkQuerier, validator.OperatorAddress, babyStakers); err != nil {
-				return nil, err
-			}
-		}
-
-		if res.Pagination == nil || len(res.Pagination.NextKey) == 0 {
-			break
-		}
-		nextKey = res.Pagination.NextKey
+	if err != nil {
+		return nil, fmt.Errorf("failed to iterate active validators: %w", err)
 	}
 
 	return babyStakers, nil

app/upgrades/v4/upgrades_costaking_test.go

@@ -199,6 +199,54 @@ func TestInitializeCoStakerRwdsTracker_FpNotActive(t *testing.T) {
 	verifyCoStakerCreated(t, ctx, cdc, storeService, stakerAddr, math.ZeroInt(), babyAmount)
 }
 
+func TestInitializeCoStakerRwdsTracker_ValidatorNotActive(t *testing.T) {
+	ctx, cdc, storeService, stkKeeper, btcStkKeeper, _, costkKeeper, fKeeper, ctrl := setupTestKeepers(t, 10)
+	defer ctrl.Finish()
+
+	require.NoError(t, btcStkKeeper.SetParams(ctx, btcstktypes.DefaultParams()))
+	require.NoError(t, stkKeeper.SetParams(ctx, stktypes.DefaultParams()))
+
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	// Create a test staker address
+	stakerAddr := datagen.GenRandomAccount().GetAddress()
+
+	// Create BTC delegation
+	btcDel := createTestBTCDelegation(t, r, ctx, btcStkKeeper, stakerAddr, 50000)
+
+	// Create baby staking delegation to an INACTIVE validator (not in LastValidatorPowers)
+	validatorAddr := datagen.GenRandomValidatorAddress()
+	babyAmount := math.NewInt(25000)
+	delegation := stktypes.Delegation{
+		DelegatorAddress: stakerAddr.String(),
+		ValidatorAddress: validatorAddr.String(),
+		Shares:           math.LegacyNewDecFromInt(babyAmount),
+	}
+
+	// Create validator but DON'T add to LastValidatorPowers (making it inactive)
+	validator := stktypes.Validator{
+		OperatorAddress: validatorAddr.String(),
+		Tokens:          babyAmount,
+		DelegatorShares: math.LegacyNewDecFromInt(babyAmount),
+		Status:          stktypes.Unbonded, // Inactive validator
+	}
+	require.NoError(t, stkKeeper.SetValidator(ctx, validator))
+	require.NoError(t, stkKeeper.SetDelegation(ctx, delegation))
+	// NOTE: Not calling SetLastValidatorPower - validator is NOT in active set
+
+	// seed voting power dist cache with FP as active
+	setupVotingPowerDistCacheWithActiveFPs(t, r, ctx, fKeeper, btcDel.FpBtcPkList)
+
+	// Execute upgrade function
+	err := v4.InitializeCoStakerRwdsTracker(
+		ctx, cdc, storeService, stkKeeper, btcStkKeeper, *costkKeeper, *fKeeper,
+	)
+	require.NoError(t, err)
+
+	// Verify co-staker was created with zero active baby (validator not active, BTC only)
+	verifyCoStakerCreated(t, ctx, cdc, storeService, stakerAddr, math.NewIntFromUint64(btcDel.TotalSat), math.ZeroInt())
+}
+
 func TestInitializeCoStakerRwdsTracker_WithRealDelegations(t *testing.T) {
 	ctx, cdc, storeService, stkKeeper, btcStkKeeper, _, costkKeeper, fKeeper, ctrl := setupTestKeepers(t, 10)
 	defer ctrl.Finish()
@@ -259,6 +307,73 @@ func TestInitializeCoStakerRwdsTracker_OnlyBTCStaking(t *testing.T) {
 	verifyCoStakerCreated(t, ctx, cdc, storeService, stakerAddr, math.NewIntFromUint64(btcDel.TotalSat), math.ZeroInt())
 }
 
+func TestInitializeCoStakerRwdsTracker_MixedActiveInactiveValidators(t *testing.T) {
+	ctx, cdc, storeService, stkKeeper, btcStkKeeper, _, costkKeeper, fKeeper, ctrl := setupTestKeepers(t, 10)
+	defer ctrl.Finish()
+
+	require.NoError(t, btcStkKeeper.SetParams(ctx, btcstktypes.DefaultParams()))
+	require.NoError(t, stkKeeper.SetParams(ctx, stktypes.DefaultParams()))
+
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	// Create a test staker address
+	stakerAddr := datagen.GenRandomAccount().GetAddress()
+
+	// Create BTC delegation
+	btcDel := createTestBTCDelegation(t, r, ctx, btcStkKeeper, stakerAddr, 50000)
+
+	// Create delegation to ACTIVE validator
+	activeValAddr := datagen.GenRandomValidatorAddress()
+	activeBabyAmount := math.NewInt(15000)
+	activeDelegation := stktypes.Delegation{
+		DelegatorAddress: stakerAddr.String(),
+		ValidatorAddress: activeValAddr.String(),
+		Shares:           math.LegacyNewDecFromInt(activeBabyAmount),
+	}
+	activeValidator := stktypes.Validator{
+		OperatorAddress: activeValAddr.String(),
+		Tokens:          activeBabyAmount,
+		DelegatorShares: math.LegacyNewDecFromInt(activeBabyAmount),
+		Status:          stktypes.Bonded,
+	}
+	require.NoError(t, stkKeeper.SetValidator(ctx, activeValidator))
+	require.NoError(t, stkKeeper.SetDelegation(ctx, activeDelegation))
+	// Mark as active validator
+	power := stkKeeper.TokensToConsensusPower(ctx, activeBabyAmount)
+	require.NoError(t, stkKeeper.SetLastValidatorPower(ctx, activeValAddr, power))
+
+	// Create delegation to INACTIVE validator
+	inactiveValAddr := datagen.GenRandomValidatorAddress()
+	inactiveBabyAmount := math.NewInt(10000)
+	inactiveDelegation := stktypes.Delegation{
+		DelegatorAddress: stakerAddr.String(),
+		ValidatorAddress: inactiveValAddr.String(),
+		Shares:           math.LegacyNewDecFromInt(inactiveBabyAmount),
+	}
+	inactiveValidator := stktypes.Validator{
+		OperatorAddress: inactiveValAddr.String(),
+		Tokens:          inactiveBabyAmount,
+		DelegatorShares: math.LegacyNewDecFromInt(inactiveBabyAmount),
+		Status:          stktypes.Unbonded, // Not bonded
+	}
+	require.NoError(t, stkKeeper.SetValidator(ctx, inactiveValidator))
+	require.NoError(t, stkKeeper.SetDelegation(ctx, inactiveDelegation))
+	// NOTE: Not calling SetLastValidatorPower - validator is NOT active
+
+	// seed voting power dist cache with FP as active
+	setupVotingPowerDistCacheWithActiveFPs(t, r, ctx, fKeeper, btcDel.FpBtcPkList)
+
+	// Execute upgrade function
+	err := v4.InitializeCoStakerRwdsTracker(
+		ctx, cdc, storeService, stkKeeper, btcStkKeeper, *costkKeeper, *fKeeper,
+	)
+	require.NoError(t, err)
+
+	// Verify co-staker was created with only the active validator's baby amount
+	// Total baby should be 15000 (from active validator), not 25000 (15000 + 10000)
+	verifyCoStakerCreated(t, ctx, cdc, storeService, stakerAddr, math.NewIntFromUint64(btcDel.TotalSat), activeBabyAmount)
+}
+
 func TestInitializeCoStakerRwdsTracker_MultipleCombinations(t *testing.T) {
 	ctx, cdc, storeService, stkKeeper, btcStkKeeper, _, costkKeeper, fKeeper, ctrl := setupTestKeepers(t, 10)
 	defer ctrl.Finish()
@@ -646,12 +761,19 @@ func createBabyDelegation(t *testing.T, ctx context.Context, stkKeeper *stkkeepe
 		Shares:           math.LegacyNewDecFromInt(delAmount),
 	}
 
-	require.NoError(t, stkKeeper.SetValidator(ctx, stktypes.Validator{
+	// Create validator and mark it as active (bonded with power)
+	validator := stktypes.Validator{
 		OperatorAddress: validatorAddr.String(),
 		Tokens:          delAmount,
 		DelegatorShares: math.LegacyNewDecFromInt(delAmount),
-	}))
+		Status:          stktypes.Bonded,
+	}
+	require.NoError(t, stkKeeper.SetValidator(ctx, validator))
 	require.NoError(t, stkKeeper.SetDelegation(ctx, delegation))
+
+	// Add to LastValidatorPowers to mark as active validator
+	power := stkKeeper.TokensToConsensusPower(ctx, delAmount)
+	require.NoError(t, stkKeeper.SetLastValidatorPower(ctx, validatorAddr, power))
 }
 
 func verifyCoStakerCreated(t *testing.T, ctx sdk.Context, cdc codec.BinaryCodec, storeService corestore.KVStoreService, stakerAddr sdk.AccAddress, expectedBTCAmount, expectedBabyAmount math.Int) {
@@ -949,6 +1071,7 @@ func loadAndSeedCosmosDelegations(t *testing.T, ctx sdk.Context, env string, stk
 				OperatorAddress: rawDel.Delegation.ValidatorAddress,
 				Tokens:          math.ZeroInt(),
 				DelegatorShares: math.LegacyZeroDec(),
+				Status:          stktypes.Bonded, // Mark as bonded so it's considered active
 			}
 			if err := stkKeeper.SetValidator(ctx, validator); err != nil {
 				return 0, fmt.Errorf("failed to set validator %s: %w", rawDel.Delegation.ValidatorAddress, err)
@@ -968,7 +1091,8 @@ func loadAndSeedCosmosDelegations(t *testing.T, ctx sdk.Context, env string, stk
 		}
 
 		// Update validator shares and tokens
-		validator, err := stkKeeper.GetValidator(ctx, sdk.MustValAddressFromBech32(rawDel.Delegation.ValidatorAddress))
+		valAddr := sdk.MustValAddressFromBech32(rawDel.Delegation.ValidatorAddress)
+		validator, err := stkKeeper.GetValidator(ctx, valAddr)
 		if err != nil {
 			return 0, fmt.Errorf("validator %s not found after creation", rawDel.Delegation.ValidatorAddress)
 		}
@@ -978,6 +1102,12 @@ func loadAndSeedCosmosDelegations(t *testing.T, ctx sdk.Context, env string, stk
 			return 0, fmt.Errorf("failed to update validator %s: %w", rawDel.Delegation.ValidatorAddress, err)
 		}
 
+		// Mark validator as active by adding to LastValidatorPowers
+		power := stkKeeper.TokensToConsensusPower(ctx, validator.Tokens)
+		if err := stkKeeper.SetLastValidatorPower(ctx, valAddr, power); err != nil {
+			return 0, fmt.Errorf("failed to set last validator power for %s: %w", rawDel.Delegation.ValidatorAddress, err)
+		}
+
 		count++
 		if count%1000 == 0 {
 			t.Logf("Processed %d cosmos delegations...", count)

proto/babylon/costaking/v1/costaking.proto

@@ -35,3 +35,27 @@ message Params {
     (gogoproto.nullable) = false
   ];
 }
+
+// Validator is a message that denotes a validator
+message Validator {
+  // addr is the validator's address (in sdk.ValAddress)
+  bytes addr = 1;
+  // tokens define the delegated tokens (incl. self-delegation).
+  bytes tokens = 2 [
+    (cosmos_proto.scalar) = "cosmos.Int",
+    (gogoproto.customtype) = "cosmossdk.io/math.Int",
+    (gogoproto.nullable) = false
+  ];
+  // shares defines total shares issued to a validator's delegators.
+  string shares = 3 [
+    (cosmos_proto.scalar) = "cosmos.Dec",
+    (gogoproto.customtype) = "cosmossdk.io/math.LegacyDec",
+    (gogoproto.nullable) = false
+  ];
+}
+
+// ValidatorSet is a message that denotes a set of validators
+message ValidatorSet {
+  // validators is the list of all validators and their delegated tokens.
+  repeated Validator validators = 1;
+}

proto/babylon/costaking/v1/genesis.proto

@@ -22,6 +22,9 @@ message GenesisState {
   // costaker addresses.
   repeated CostakerRewardsTrackerEntry costakers_rewards_tracker = 4
       [ (gogoproto.nullable) = false ];
+  // validator_set contains all validators and their delegated tokens.
+  ValidatorSet validator_set = 5
+      [ (gogoproto.nullable) = false ];
 }
 
 // CurrentRewardsEntry represents the chain current rewards.

test/e2e/configurer/chain/commands.go

@@ -22,6 +22,9 @@ import (
 	btccheckpointtypes "github.com/babylonlabs-io/babylon/v4/x/btccheckpoint/types"
 	blc "github.com/babylonlabs-io/babylon/v4/x/btclightclient/types"
 	cttypes "github.com/babylonlabs-io/babylon/v4/x/checkpointing/types"
+	cmtjson "github.com/cometbft/cometbft/libs/json"
+	"github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/bech32"
 	sdkquerytypes "github.com/cosmos/cosmos-sdk/types/query"
@@ -573,11 +576,29 @@ func (n *NodeConfig) FailICASendTx(from, connectionID, packetMsgPath string) {
 	n.LogActionF("Failed to perform ICA send (as expected)")
 }
 
-func (n *NodeConfig) Delegate(fromWallet, validator string, amount string, overallFlags ...string) {
+func (n *NodeConfig) Delegate(fromWallet, validator string, amount string, overallFlags ...string) (txHash string) {
 	n.LogActionF("delegating from %s to validator %s", fromWallet, validator)
 	cmd := []string{"babylond", "tx", "epoching", "delegate", validator, amount, fmt.Sprintf("--from=%s", fromWallet)}
-	_, _, err := n.containerManager.ExecTxCmd(n.t, n.chainId, n.Name, append(cmd, overallFlags...))
+	outBuf, _, err := n.containerManager.ExecTxCmd(n.t, n.chainId, n.Name, append(cmd, overallFlags...))
 
 	require.NoError(n.t, err)
 	n.LogActionF("successfully delegated %s to validator %s", fromWallet, validator)
+	return GetTxHashFromOutput(outBuf.String())
+}
+
+// ValidatorConsPubKey gets the consensus pubkey base64 key from a validator node
+func (n *NodeConfig) ValidatorConsPubKey() cryptotypes.PubKey {
+	cmd := []string{"babylond", "comet", "show-validator", "--home=/home/babylon/babylondata"}
+	outBuf, _, err := n.containerManager.ExecCmd(n.t, n.Name, cmd, "")
+	require.NoError(n.t, err)
+
+	// Parse the JSON output to extract the base64 key
+	// Format: {"@type":"/cosmos.crypto.ed25519.PubKey","key":"<base64>"}
+	var pubKey ed25519.PubKey
+	output := strings.TrimSpace(outBuf.String())
+
+	err = cmtjson.Unmarshal([]byte(output), &pubKey)
+	require.NoError(n.t, err)
+
+	return &pubKey
 }