imp(costaking): track baby staked to active validators

CHANGELOG.md

@@ -41,6 +41,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
+- [#1776](https://github.com/babylonlabs-io/babylon/pull/1776) Track baby staked to active validators only
 
 ### Bug fixes
 

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)

test/e2ev2/tmanager/genesis.go

@@ -309,4 +309,3 @@ func UpdateGenesisTokenFactory(tokenfactoryGenState *tokenfactorytypes.GenesisSt
 		DenomCreationFee: sdk.NewCoins(sdk.NewCoin(appparams.DefaultBondDenom, sdkmath.NewInt(10000))),
 	}
 }
-

types/sdk_math.go

@@ -48,4 +48,4 @@ func SafeNewCoin(denom string, amount sdkmath.Int) (sdk.Coin, error) {
 	}
 
 	return coin, nil
-}
+}