feat: fix bugs with status command

cli/core/status.go

@@ -64,7 +64,7 @@ func getRegularBalancesGwei(allValidators []ValidatorWithIndex, state *spec.Vers
 	return validatorBalances
 }
 
-func sumActiveValidatorBalancesGwei(allValidators []ValidatorWithIndex, allBalances []phase0.Gwei, state *spec.VersionedBeaconState) phase0.Gwei {
+func sumActiveValidatorBeaconBalancesGwei(allValidators []ValidatorWithIndex, allBalances []phase0.Gwei, state *spec.VersionedBeaconState) phase0.Gwei {
 	var sumGwei phase0.Gwei = 0
 
 	for i := 0; i < len(allValidators); i++ {
@@ -75,31 +75,52 @@ func sumActiveValidatorBalancesGwei(allValidators []ValidatorWithIndex, allBalan
 	return sumGwei
 }
 
+func sumRestakedBalancesGwei(eth *ethclient.Client, eigenpodAddress string, activeValidators []ValidatorWithIndex) (phase0.Gwei, error) {
+	var sumGwei phase0.Gwei = 0
+
+	validatorInfos, err := GetOnchainValidatorInfo(eth, eigenpodAddress, activeValidators)
+	if err != nil {
+		return 0, err
+	}
+
+	for i := 0; i < len(activeValidators); i++ {
+		validatorInfo := validatorInfos[i]
+
+		sumGwei += phase0.Gwei(validatorInfo.RestakedBalanceGwei)
+	}
+
+	return sumGwei, nil
+}
+
 func GetStatus(ctx context.Context, eigenpodAddress string, eth *ethclient.Client, beaconClient BeaconClient) EigenpodStatus {
 	validators := map[string]Validator{}
 	var activeCheckpoint *Checkpoint = nil
 
 	eigenPod, err := onchain.NewEigenPod(common.HexToAddress(eigenpodAddress), eth)
 	PanicOnError("failed to reach eigenpod", err)
 
-	timestamp, err := eigenPod.CurrentCheckpointTimestamp(nil)
-	PanicOnError("failed to fetch current checkpoint timestamp", err)
+	checkpoint, err := eigenPod.CurrentCheckpoint(nil)
+	PanicOnError("failed to fetch checkpoint information", err)
 
-	state, err := beaconClient.GetBeaconState(ctx, "head")
-	PanicOnError("failed to fetch state", err)
+	// Fetch the beacon state associated with the checkpoint (or "head" if there is no checkpoint)
+	checkpointTimestamp, state, err := GetCheckpointTimestampAndBeaconState(ctx, eigenpodAddress, eth, beaconClient)
+	PanicOnError("failed to fetch checkpoint and beacon state", err)
 
 	allValidators, err := FindAllValidatorsForEigenpod(eigenpodAddress, state)
 	PanicOnError("failed to find validators", err)
 
-	allBalances := getRegularBalancesGwei(allValidators, state)
+	allBeaconBalances := getRegularBalancesGwei(allValidators, state)
 
 	activeValidators, err := SelectActiveValidators(eth, eigenpodAddress, allValidators)
 	PanicOnError("failed to find active validators", err)
 
-	checkpointableValidators, err := SelectCheckpointableValidators(eth, eigenpodAddress, allValidators, timestamp)
+	checkpointableValidators, err := SelectCheckpointableValidators(eth, eigenpodAddress, allValidators, checkpointTimestamp)
 	PanicOnError("failed to find checkpointable validators", err)
 
-	sumRegularBalancesGwei := sumActiveValidatorBalancesGwei(activeValidators, allBalances, state)
+	sumBeaconBalancesGwei := sumActiveValidatorBeaconBalancesGwei(activeValidators, allBeaconBalances, state)
+
+	sumRestakedBalancesGwei, err := sumRestakedBalancesGwei(eth, eigenpodAddress, activeValidators)
+	PanicOnError("failed to calculate sum of onchain validator balances", err)
 
 	for i := 0; i < len(allValidators); i++ {
 		validator := allValidators[i].Validator
@@ -116,13 +137,10 @@ func GetStatus(ctx context.Context, eigenpodAddress string, eth *ethclient.Clien
 			IsAwaitingActivationQueue:           validator.ActivationEpoch == FAR_FUTURE_EPOCH,
 			IsAwaitingWithdrawalCredentialProof: IsAwaitingWithdrawalCredentialProof(validatorInfo, validator),
 			EffectiveBalance:                    uint64(validator.EffectiveBalance),
-			CurrentBalance:                      uint64(allBalances[validatorIndex]),
+			CurrentBalance:                      uint64(allBeaconBalances[validatorIndex]),
 		}
 	}
 
-	checkpoint, err := eigenPod.CurrentCheckpoint(nil)
-	PanicOnError("failed to fetch checkpoint information", err)
-
 	eigenpodManagerContractAddress, err := eigenPod.EigenPodManager(nil)
 	PanicOnError("failed to get manager address", err)
 
@@ -136,56 +154,68 @@ func GetStatus(ctx context.Context, eigenpodAddress string, eth *ethclient.Clien
 	PanicOnError("failed to get eigenpod proof submitter", err)
 
 	currentOwnerShares, err := eigenPodManager.PodOwnerShares(nil, eigenPodOwner)
+	// currentOwnerShares = big.NewInt(0)
 	PanicOnError("failed to load pod owner shares", err)
 	currentOwnerSharesETH := IweiToEther(currentOwnerShares)
+	currentOwnerSharesGwei := WeiToGwei(currentOwnerShares)
 
 	withdrawableRestakedExecutionLayerGwei, err := eigenPod.WithdrawableRestakedExecutionLayerGwei(nil)
 	PanicOnError("failed to fetch withdrawableRestakedExecutionLayerGwei", err)
 
-	var pendingSharesGwei *big.Float
-	mustForceCheckpoint := false
-	// If we currently have an active checkpoint, estimate the total shares
-	// we'll have when we complete it:
+	// Estimate the total shares we'll have if we complete an existing checkpoint
+	// (or start a new one and complete that).
 	//
-	// pendingSharesGwei = withdrawableRestakedExecutionLayerGwei + checkpoint.PodBalanceGwei + sumRegularBalancesGwei
-	if timestamp != 0 {
-		pendingSharesGwei = new(big.Float).Add(
-			new(big.Float).Add(
-				new(big.Float).SetUint64(withdrawableRestakedExecutionLayerGwei),
-				new(big.Float).SetUint64(checkpoint.PodBalanceGwei),
-			),
-			new(big.Float).SetUint64(uint64(sumRegularBalancesGwei)),
-		)
+	// First, we need the change in the pod's native ETH balance since the last checkpoint:
+	var nativeETHDeltaGwei *big.Float
+	mustForceCheckpoint := false
+
+	if checkpointTimestamp != 0 {
+		// Change in the pod's native ETH balance (already calculated for us when the checkpoint was started)
+		nativeETHDeltaGwei = new(big.Float).SetUint64(checkpoint.PodBalanceGwei)
 
 		activeCheckpoint = &Checkpoint{
 			ProofsRemaining: checkpoint.ProofsRemaining.Uint64(),
-			StartedAt:       timestamp,
+			StartedAt:       checkpointTimestamp,
 		}
 	} else {
-		// If we don't have an active checkpoint, estimate the shares we'd have if
-		// we created one and then completed it:
-		//
-		// pendingSharesGwei = sumRegularBalancesGwei + latestPodBalanceGwei
 		latestPodBalanceWei, err := eth.BalanceAt(ctx, common.HexToAddress(eigenpodAddress), nil)
 		PanicOnError("failed to fetch pod balance", err)
 		latestPodBalanceGwei := WeiToGwei(latestPodBalanceWei)
 
-		pendingSharesGwei = new(big.Float).Add(
-			new(big.Float).SetUint64(uint64(sumRegularBalancesGwei)),
-			latestPodBalanceGwei,
-		)
-
-		// Determine whether the checkpoint needs to be run with `--force`
-		checkpointableBalance := new(big.Float).Sub(
+		// We don't have a checkpoint currently, so we need to calculate what
+		// checkpoint.PodBalanceGwei would be if we started one now:
+		nativeETHDeltaGwei = new(big.Float).Sub(
 			latestPodBalanceGwei,
 			new(big.Float).SetUint64(withdrawableRestakedExecutionLayerGwei),
 		)
 
-		if checkpointableBalance.Sign() == 0 {
+		// Determine whether the checkpoint needs to be started with `--force`
+		if nativeETHDeltaGwei.Sign() == 0 {
 			mustForceCheckpoint = true
 		}
 	}
 
+	// Next, we need the change in the pod's beacon chain balances since the last
+	// checkpoint:
+	//
+	// beaconETHDeltaGwei = sumBeaconBalancesGwei - sumRestakedBalancesGwei
+	beaconETHDeltaGwei := new(big.Float).Sub(
+		new(big.Float).SetUint64(uint64(sumBeaconBalancesGwei)),
+		new(big.Float).SetUint64(uint64(sumRestakedBalancesGwei)),
+	)
+
+	// Sum of these two deltas represents the change in shares after this checkpoint
+	totalShareDeltaGwei := new(big.Float).Add(
+		nativeETHDeltaGwei,
+		beaconETHDeltaGwei,
+	)
+
+	// Calculate new total shares by applying delta to current shares
+	pendingSharesGwei := new(big.Float).Add(
+		currentOwnerSharesGwei,
+		totalShareDeltaGwei,
+	)
+
 	pendingEth := GweiToEther(pendingSharesGwei)
 
 	return EigenpodStatus{

cli/core/utils.go

@@ -5,13 +5,15 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"crypto/sha256"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"log"
 	"math"
 	"math/big"
 	"os"
 	"sort"
+	"strconv"
 
 	eigenpodproofs "github.com/Layr-Labs/eigenpod-proofs-generation"
 	"github.com/Layr-Labs/eigenpod-proofs-generation/cli/core/onchain"
@@ -161,6 +163,68 @@ func GetCurrentCheckpoint(eigenpodAddress string, client *ethclient.Client) (uin
 	return timestamp, nil
 }
 
+// Fetch and return the current checkpoint timestamp for the pod
+// If the checkpoint exists (timestamp != 0), also return the beacon state for the checkpoint
+// If the checkpoint does not exist (timestamp == 0), return the head beacon state (i.e. the state we would use "if we start a checkpoint now")
+func GetCheckpointTimestampAndBeaconState(
+	ctx context.Context,
+	eigenpodAddress string,
+	eth *ethclient.Client,
+	beaconClient BeaconClient,
+) (uint64, *spec.VersionedBeaconState, error) {
+	tracing := GetContextTracingCallbacks(ctx)
+
+	tracing.OnStartSection("GetCurrentCheckpoint", map[string]string{})
+	checkpointTimestamp, err := GetCurrentCheckpoint(eigenpodAddress, eth)
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed to fetch current checkpoint: %w", err)
+	}
+	tracing.OnEndSection()
+
+	// stateId to look up beacon state. "head" by default (if we do not have a checkpoint)
+	beaconStateId := "head"
+
+	// If we have a checkpoint, get the state id for the checkpoint's block root
+	if checkpointTimestamp != 0 {
+		// Fetch the checkpoint's block root
+		tracing.OnStartSection("GetCurrentCheckpointBlockRoot", map[string]string{})
+		blockRoot, err := GetCurrentCheckpointBlockRoot(eigenpodAddress, eth)
+		if err != nil {
+			return 0, nil, fmt.Errorf("failed to fetch last checkpoint: %w", err)
+		}
+		if blockRoot == nil {
+			return 0, nil, fmt.Errorf("failed to fetch last checkpoint - nil blockRoot")
+		}
+		// Block root should be nonzero because we have an active checkpoint
+		rootBytes := *blockRoot
+		if AllZero(rootBytes[:]) {
+			return 0, nil, fmt.Errorf("failed to fetch last checkpoint - empty blockRoot")
+		}
+		tracing.OnEndSection()
+
+		headerBlock := "0x" + hex.EncodeToString((*blockRoot)[:])
+		tracing.OnStartSection("GetBeaconHeader", map[string]string{})
+		header, err := beaconClient.GetBeaconHeader(ctx, headerBlock)
+		if err != nil {
+			return 0, nil, fmt.Errorf("failed to fetch beacon header (%s): %w", headerBlock, err)
+		}
+		tracing.OnEndSection()
+
+		beaconStateId = strconv.FormatUint(uint64(header.Header.Message.Slot), 10)
+	} else {
+		beaconStateId = "head"
+	}
+
+	tracing.OnStartSection("GetBeaconState", map[string]string{})
+	beaconState, err := beaconClient.GetBeaconState(ctx, beaconStateId)
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed to fetch beacon state: %w", err)
+	}
+	tracing.OnEndSection()
+
+	return checkpointTimestamp, beaconState, nil
+}
+
 func SortByStatus(validators map[string]Validator) ([]Validator, []Validator, []Validator, []Validator) {
 	var awaitingActivationQueueValidators, inactiveValidators, activeValidators, withdrawnValidators []Validator