feat: granular states

internal/db/delegation.go

@@ -36,13 +36,22 @@ func (db *Database) SaveNewBTCDelegation(
 }
 
 func (db *Database) UpdateBTCDelegationState(
-	ctx context.Context, stakingTxHash string, newState types.DelegationState,
+	ctx context.Context,
+	stakingTxHash string,
+	newState types.DelegationState,
+	subState *types.DelegationSubState,
 ) error {
 	filter := bson.M{"_id": stakingTxHash}
+	updateFields := bson.M{
+		"state": newState.String(),
+	}
+
+	if subState != nil {
+		updateFields["sub_state"] = subState.String()
+	}
+
 	update := bson.M{
-		"$set": bson.M{
-			"state": newState.String(),
-		},
+		"$set": updateFields,
 	}
 
 	res := db.client.Database(db.dbName).

internal/db/interface.go

@@ -98,7 +98,10 @@ type DbInterface interface {
 	 * @return An error if the operation failed
 	 */
 	UpdateBTCDelegationState(
-		ctx context.Context, stakingTxHash string, newState types.DelegationState,
+		ctx context.Context,
+		stakingTxHash string,
+		newState types.DelegationState,
+		subState *types.DelegationSubState,
 	) error
 	/**
 	 * SaveBTCDelegationUnbondingCovenantSignature saves a BTC delegation
@@ -160,7 +163,10 @@ type DbInterface interface {
 	 * @return An error if the operation failed
 	 */
 	SaveNewTimeLockExpire(
-		ctx context.Context, stakingTxHashHex string, expireHeight uint32, txType string,
+		ctx context.Context,
+		stakingTxHashHex string,
+		expireHeight uint32,
+		subState types.DelegationSubState,
 	) error
 	/**
 	 * FindExpiredDelegations finds the expired delegations.

internal/db/model/delegation.go

@@ -32,6 +32,7 @@ type BTCDelegationDetails struct {
 	StartHeight                 uint32                       `bson:"start_height"`
 	EndHeight                   uint32                       `bson:"end_height"`
 	State                       types.DelegationState        `bson:"state"`
+	SubState                    types.DelegationSubState     `bson:"sub_state,omitempty"`
 	ParamsVersion               uint32                       `bson:"params_version"`
 	UnbondingTime               uint32                       `bson:"unbonding_time"`
 	UnbondingTx                 string                       `bson:"unbonding_tx"`

internal/db/model/timelock.go

@@ -1,15 +1,19 @@
 package model
 
+import "github.com/babylonlabs-io/babylon-staking-indexer/internal/types"
+
 type TimeLockDocument struct {
-	StakingTxHashHex string `bson:"_id"` // Primary key
-	ExpireHeight     uint32 `bson:"expire_height"`
-	TxType           string `bson:"tx_type"`
+	StakingTxHashHex   string                   `bson:"_id"` // Primary key
+	ExpireHeight       uint32                   `bson:"expire_height"`
+	DelegationSubState types.DelegationSubState `bson:"delegation_sub_state"`
 }
 
-func NewTimeLockDocument(stakingTxHashHex string, expireHeight uint32, txType string) *TimeLockDocument {
+func NewTimeLockDocument(
+	stakingTxHashHex string, expireHeight uint32, subState types.DelegationSubState,
+) *TimeLockDocument {
 	return &TimeLockDocument{
-		StakingTxHashHex: stakingTxHashHex,
-		ExpireHeight:     expireHeight,
-		TxType:           txType,
+		StakingTxHashHex:   stakingTxHashHex,
+		ExpireHeight:       expireHeight,
+		DelegationSubState: subState,
 	}
 }

internal/db/timelock.go

@@ -6,16 +6,19 @@ import (
 	"fmt"
 
 	"github.com/babylonlabs-io/babylon-staking-indexer/internal/db/model"
+	"github.com/babylonlabs-io/babylon-staking-indexer/internal/types"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 )
 
 func (db *Database) SaveNewTimeLockExpire(
-	ctx context.Context, stakingTxHashHex string,
-	expireHeight uint32, txType string,
+	ctx context.Context,
+	stakingTxHashHex string,
+	expireHeight uint32,
+	subState types.DelegationSubState,
 ) error {
-	tlDoc := model.NewTimeLockDocument(stakingTxHashHex, expireHeight, txType)
+	tlDoc := model.NewTimeLockDocument(stakingTxHashHex, expireHeight, subState)
 	_, err := db.client.Database(db.dbName).
 		Collection(model.TimeLockCollection).
 		InsertOne(ctx, tlDoc)

internal/services/delegation.go

@@ -158,7 +158,10 @@ func (s *Service) processCovenantQuorumReachedEvent(
 	}
 
 	if dbErr := s.db.UpdateBTCDelegationState(
-		ctx, covenantQuorumReachedEvent.StakingTxHash, newState,
+		ctx,
+		covenantQuorumReachedEvent.StakingTxHash,
+		newState,
+		nil,
 	); dbErr != nil {
 		return types.NewError(
 			http.StatusInternalServerError,
@@ -205,12 +208,14 @@ func (s *Service) processBTCDelegationInclusionProofReceivedEvent(
 	}
 
 	if dbErr := s.db.UpdateBTCDelegationDetails(
-		ctx, inclusionProofEvent.StakingTxHash, model.FromEventBTCDelegationInclusionProofReceived(inclusionProofEvent),
+		ctx,
+		inclusionProofEvent.StakingTxHash,
+		model.FromEventBTCDelegationInclusionProofReceived(inclusionProofEvent),
 	); dbErr != nil {
 		return types.NewError(
 			http.StatusInternalServerError,
 			types.InternalServiceError,
-			fmt.Errorf("failed to update BTC delegation state: %w", dbErr),
+			fmt.Errorf("failed to update BTC delegation details: %w", dbErr),
 		)
 	}
 
@@ -260,13 +265,15 @@ func (s *Service) processBTCDelegationUnbondedEarlyEvent(
 		)
 	}
 
+	subState := types.SubStateEarlyUnbonding
+
 	// Save timelock expire
 	unbondingExpireHeight := uint32(unbondingStartHeight) + delegation.UnbondingTime
 	if err := s.db.SaveNewTimeLockExpire(
 		ctx,
 		delegation.StakingTxHashHex,
 		unbondingExpireHeight,
-		types.EarlyUnbondingTxType.String(),
+		subState,
 	); err != nil {
 		return types.NewError(
 			http.StatusInternalServerError,
@@ -280,6 +287,7 @@ func (s *Service) processBTCDelegationUnbondedEarlyEvent(
 		ctx,
 		unbondedEarlyEvent.StakingTxHash,
 		types.StateUnbonding,
+		&subState,
 	); err != nil {
 		return types.NewError(
 			http.StatusInternalServerError,
@@ -330,12 +338,14 @@ func (s *Service) processBTCDelegationExpiredEvent(
 		return err
 	}
 
+	subState := types.SubStateTimelock
+
 	// Save timelock expire
 	if err := s.db.SaveNewTimeLockExpire(
 		ctx,
 		delegation.StakingTxHashHex,
 		delegation.EndHeight,
-		types.ExpiredTxType.String(),
+		subState,
 	); err != nil {
 		return types.NewError(
 			http.StatusInternalServerError,
@@ -349,6 +359,7 @@ func (s *Service) processBTCDelegationExpiredEvent(
 		ctx,
 		delegation.StakingTxHashHex,
 		types.StateUnbonding,
+		&subState,
 	); err != nil {
 		return types.NewError(
 			http.StatusInternalServerError,

internal/services/expiry_checker.go

@@ -58,7 +58,12 @@ func (s *Service) checkExpiry(ctx context.Context) *types.Error {
 			return consumerErr
 		}
 
-		if err := s.db.UpdateBTCDelegationState(ctx, delegation.StakingTxHashHex, types.StateWithdrawable); err != nil {
+		if err := s.db.UpdateBTCDelegationState(
+			ctx,
+			delegation.StakingTxHashHex,
+			types.StateWithdrawable,
+			&tlDoc.DelegationSubState,
+		); err != nil {
 			log.Error().Err(err).Msg("Error updating BTC delegation state to withdrawable")
 			return types.NewInternalServiceError(
 				fmt.Errorf("failed to update BTC delegation state to withdrawable: %w", err),

internal/services/watch_btc_events.go

@@ -34,6 +34,7 @@ func (s *Service) watchForSpendStakingTx(
 			quitCtx,
 			spendDetail.SpendingTx,
 			spendDetail.SpenderInputIndex,
+			uint32(spendDetail.SpendingHeight),
 			delegation,
 		); err != nil {
 			log.Error().Err(err).Msg("failed to handle spending staking transaction")
@@ -62,6 +63,7 @@ func (s *Service) watchForSpendUnbondingTx(
 		if err := s.handleSpendingUnbondingTransaction(
 			quitCtx,
 			spendDetail.SpendingTx,
+			uint32(spendDetail.SpendingHeight),
 			spendDetail.SpenderInputIndex,
 			delegation,
 		); err != nil {
@@ -79,6 +81,7 @@ func (s *Service) watchForSpendUnbondingTx(
 func (s *Service) watchForSpendSlashingChange(
 	spendEvent *notifier.SpendEvent,
 	delegation *model.BTCDelegationDetails,
+	subState types.DelegationSubState,
 ) {
 	defer s.wg.Done()
 	quitCtx, cancel := s.quitContext()
@@ -103,10 +106,12 @@ func (s *Service) watchForSpendSlashingChange(
 		}
 
 		// Update to withdrawn state
+		delegationSubState := subState
 		if err := s.db.UpdateBTCDelegationState(
 			quitCtx,
 			delegation.StakingTxHashHex,
-			types.StateSlashedWithdrawn,
+			types.StateWithdrawn,
+			&delegationSubState,
 		); err != nil {
 			log.Error().Err(err).Msg("failed to update delegation state")
 			return
@@ -121,8 +126,9 @@ func (s *Service) watchForSpendSlashingChange(
 
 func (s *Service) handleSpendingStakingTransaction(
 	ctx context.Context,
-	tx *wire.MsgTx,
+	spendingTx *wire.MsgTx,
 	spendingInputIdx uint32,
+	spendingHeight uint32,
 	delegation *model.BTCDelegationDetails,
 ) error {
 	params, err := s.db.GetStakingParams(ctx, delegation.ParamsVersion)
@@ -131,7 +137,7 @@ func (s *Service) handleSpendingStakingTransaction(
 	}
 
 	// First try to validate as unbonding tx
-	isUnbonding, err := s.IsValidUnbondingTx(tx, delegation, params)
+	isUnbonding, err := s.IsValidUnbondingTx(spendingTx, delegation, params)
 	if err != nil {
 		return fmt.Errorf("failed to validate unbonding tx: %w", err)
 	}
@@ -141,10 +147,10 @@ func (s *Service) handleSpendingStakingTransaction(
 	}
 
 	// Try to validate as withdrawal transaction
-	withdrawalErr := s.validateWithdrawalTxFromStaking(tx, spendingInputIdx, delegation, params)
+	withdrawalErr := s.validateWithdrawalTxFromStaking(spendingTx, spendingInputIdx, delegation, params)
 	if withdrawalErr == nil {
 		// It's a valid withdrawal, process it
-		return s.handleWithdrawal(ctx, delegation)
+		return s.handleWithdrawal(ctx, delegation, types.SubStateTimelock)
 	}
 
 	// If it's not a valid withdrawal, check if it's a valid slashing
@@ -153,7 +159,7 @@ func (s *Service) handleSpendingStakingTransaction(
 	}
 
 	// Try to validate as slashing transaction
-	if err := s.validateSlashingTxFromStaking(tx, spendingInputIdx, delegation, params); err != nil {
+	if err := s.validateSlashingTxFromStaking(spendingTx, spendingInputIdx, delegation, params); err != nil {
 		if errors.Is(err, types.ErrInvalidSlashingTx) {
 			// Neither withdrawal nor slashing - this is an invalid spend
 			return fmt.Errorf("transaction is neither valid unbonding, withdrawal, nor slashing: %w", err)
@@ -162,12 +168,19 @@ func (s *Service) handleSpendingStakingTransaction(
 	}
 
 	// It's a valid slashing tx, watch for spending change output
-	return s.startWatchingSlashingChange(ctx, tx, delegation)
+	return s.startWatchingSlashingChange(
+		ctx,
+		spendingTx,
+		spendingHeight,
+		delegation,
+		types.SubStateTimelockSlashing,
+	)
 }
 
 func (s *Service) handleSpendingUnbondingTransaction(
 	ctx context.Context,
-	tx *wire.MsgTx,
+	spendingTx *wire.MsgTx,
+	spendingHeight uint32,
 	spendingInputIdx uint32,
 	delegation *model.BTCDelegationDetails,
 ) error {
@@ -177,10 +190,10 @@ func (s *Service) handleSpendingUnbondingTransaction(
 	}
 
 	// First try to validate as withdrawal transaction
-	withdrawalErr := s.validateWithdrawalTxFromUnbonding(tx, delegation, spendingInputIdx, params)
+	withdrawalErr := s.validateWithdrawalTxFromUnbonding(spendingTx, delegation, spendingInputIdx, params)
 	if withdrawalErr == nil {
 		// It's a valid withdrawal, process it
-		return s.handleWithdrawal(ctx, delegation)
+		return s.handleWithdrawal(ctx, delegation, types.SubStateEarlyUnbonding)
 	}
 
 	// If it's not a valid withdrawal, check if it's a valid slashing
@@ -189,7 +202,7 @@ func (s *Service) handleSpendingUnbondingTransaction(
 	}
 
 	// Try to validate as slashing transaction
-	if err := s.validateSlashingTxFromUnbonding(tx, delegation, spendingInputIdx, params); err != nil {
+	if err := s.validateSlashingTxFromUnbonding(spendingTx, delegation, spendingInputIdx, params); err != nil {
 		if errors.Is(err, types.ErrInvalidSlashingTx) {
 			// Neither withdrawal nor slashing - this is an invalid spend
 			return fmt.Errorf("transaction is neither valid withdrawal nor slashing: %w", err)
@@ -198,10 +211,20 @@ func (s *Service) handleSpendingUnbondingTransaction(
 	}
 
 	// It's a valid slashing tx, watch for spending change output
-	return s.startWatchingSlashingChange(ctx, tx, delegation)
+	return s.startWatchingSlashingChange(
+		ctx,
+		spendingTx,
+		spendingHeight,
+		delegation,
+		types.SubStateEarlyUnbondingSlashing,
+	)
 }
 
-func (s *Service) handleWithdrawal(ctx context.Context, delegation *model.BTCDelegationDetails) error {
+func (s *Service) handleWithdrawal(
+	ctx context.Context,
+	delegation *model.BTCDelegationDetails,
+	subState types.DelegationSubState,
+) error {
 	delegationState, err := s.db.GetBTCDelegationState(ctx, delegation.StakingTxHashHex)
 	if err != nil {
 		return fmt.Errorf("failed to get delegation state: %w", err)
@@ -217,10 +240,17 @@ func (s *Service) handleWithdrawal(ctx context.Context, delegation *model.BTCDel
 		ctx,
 		delegation.StakingTxHashHex,
 		types.StateWithdrawn,
+		&subState,
 	)
 }
 
-func (s *Service) startWatchingSlashingChange(ctx context.Context, slashingTx *wire.MsgTx, delegation *model.BTCDelegationDetails) error {
+func (s *Service) startWatchingSlashingChange(
+	ctx context.Context,
+	slashingTx *wire.MsgTx,
+	spendingHeight uint32,
+	delegation *model.BTCDelegationDetails,
+	subState types.DelegationSubState,
+) error {
 	// Create outpoint for the change output (index 1)
 	changeOutpoint := wire.OutPoint{
 		Hash:  slashingTx.TxHash(),
@@ -237,8 +267,24 @@ func (s *Service) startWatchingSlashingChange(ctx context.Context, slashingTx *w
 		return fmt.Errorf("failed to register spend ntfn for slashing change output: %w", err)
 	}
 
+	stakingParams, err := s.db.GetStakingParams(ctx, delegation.ParamsVersion)
+	if err != nil {
+		return fmt.Errorf("failed to get staking params: %w", err)
+	}
+	slashingChangeTimelockExpireHeight := spendingHeight + stakingParams.MinUnbondingTimeBlocks
+
+	// Save timelock expire to mark it as Withdrawn (sub state - timelock_slashing/early_unbonding_slashing)
+	if err := s.db.SaveNewTimeLockExpire(
+		ctx,
+		delegation.StakingTxHashHex,
+		slashingChangeTimelockExpireHeight,
+		subState,
+	); err != nil {
+		return fmt.Errorf("failed to save timelock expire: %w", err)
+	}
+
 	s.wg.Add(1)
-	go s.watchForSpendSlashingChange(spendEv, delegation)
+	go s.watchForSpendSlashingChange(spendEv, delegation, subState)
 
 	return nil
 }