chore: migrate to height upgrade style

app/app.go

@@ -781,6 +781,8 @@ func (app *NillionApp) Name() string { return app.BaseApp.Name() }
 
 // PreBlocker application updates every pre block
 func (app *NillionApp) PreBlocker(ctx sdk.Context, _ *abci.RequestFinalizeBlock) (*sdk.ResponsePreBlock, error) {
+	// Schedule height-based upgrades before the upgrade module's PreBlock runs
+	app.scheduleHeightBasedUpgrades(ctx)
 	return app.ModuleManager.PreBlock(ctx)
 }
 

app/upgrades.go

@@ -1,8 +1,11 @@
 package nillionapp
 
 import (
-	upgradetypes "cosmossdk.io/x/upgrade/types"
 	"fmt"
+
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
 	"github.com/NillionNetwork/nilchain/app/upgrades"
 )
 
@@ -13,12 +16,16 @@ var Upgrades = []upgrades.Upgrade{
 }
 
 func (app NillionApp) setupUpgradeHandlers() {
+	keepers := upgrades.Keepers{
+		MigrationMngrKeeper: &app.MigrationMngrKeeper,
+	}
 	for _, upgrade := range Upgrades {
 		app.UpgradeKeeper.SetUpgradeHandler(
 			upgrade.UpgradeName,
 			upgrade.CreateUpgradeHandler(
 				*app.ModuleManager,
 				app.configurator,
+				keepers,
 			),
 		)
 	}
@@ -35,8 +42,65 @@ func (app NillionApp) setupUpgradeStoreLoaders() {
 	}
 
 	for _, upgrade := range Upgrades {
+		// Handle upgrade info from disk (standard upgrade module halt)
 		if upgradeInfo.Name == upgrade.UpgradeName {
 			app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &upgrade.StoreUpgrades))
+			return
+		}
+	}
+
+	// Handle height-based upgrades with manual --halt-height (no disk info).
+	// UpgradeStoreLoader only applies when latestVersion == upgradeHeight - 1,
+	// so setting this unconditionally is safe - it's a no-op at other heights.
+	for _, upgrade := range Upgrades {
+		if upgrade.UpgradeHeight > 0 && len(upgrade.StoreUpgrades.Added) > 0 {
+			app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgrade.UpgradeHeight, &upgrade.StoreUpgrades))
+			return
+		}
+	}
+}
+
+// scheduleHeightBasedUpgrades schedules upgrades that have a non-zero UpgradeHeight.
+// This should be called in PreBlocker before the upgrade module's PreBlock runs.
+// Supports manual halt-height flow: validators set --halt-height, restart with new binary,
+// and this function schedules the upgrade to run on the first block.
+func (app NillionApp) scheduleHeightBasedUpgrades(ctx sdk.Context) {
+	for _, upgrade := range Upgrades {
+		if upgrade.UpgradeHeight == 0 {
+			continue
+		}
+
+		// Skip if we've already passed the upgrade height (upgrade already applied)
+		if ctx.BlockHeight() > upgrade.UpgradeHeight {
+			continue
+		}
+
+		// Check if this upgrade is already scheduled
+		plan, err := app.UpgradeKeeper.GetUpgradePlan(ctx)
+		if err == nil && plan.Name == upgrade.UpgradeName {
+			continue
+		}
+
+		// Check if upgrade was already applied
+		applied, err := app.UpgradeKeeper.GetDoneHeight(ctx, upgrade.UpgradeName)
+		if err == nil && applied > 0 {
+			continue
+		}
+
+		// Schedule the upgrade at the target height, or current height if we're at/past it
+		scheduleHeight := upgrade.UpgradeHeight
+		if ctx.BlockHeight() == upgrade.UpgradeHeight {
+			// We're at the upgrade height (post manual halt-height restart)
+			// Schedule for current height so upgrade module applies it
+			scheduleHeight = ctx.BlockHeight()
+		}
+
+		err = app.UpgradeKeeper.ScheduleUpgrade(ctx, upgradetypes.Plan{
+			Name:   upgrade.UpgradeName,
+			Height: scheduleHeight,
+		})
+		if err != nil {
+			app.Logger().Error("failed to schedule upgrade", "name", upgrade.UpgradeName, "height", scheduleHeight, "error", err)
 		}
 	}
 }

app/upgrades/types.go

@@ -3,11 +3,17 @@ package upgrades
 import (
 	"cosmossdk.io/store/types"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	migrationmngr_keeper "github.com/evstack/ev-abci/modules/migrationmngr/keeper"
 	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
+type Keepers struct {
+	MigrationMngrKeeper *migrationmngr_keeper.Keeper
+}
+
 type Upgrade struct {
 	UpgradeName          string
-	CreateUpgradeHandler func(mm module.Manager, configurator module.Configurator) upgradetypes.UpgradeHandler
+	UpgradeHeight        int64 // 0 = governance-triggered, >0 = height-triggered
+	CreateUpgradeHandler func(mm module.Manager, configurator module.Configurator, keepers Keepers) upgradetypes.UpgradeHandler
 	StoreUpgrades        types.StoreUpgrades
 }

app/upgrades/upgrade_0_2_1.go

@@ -10,7 +10,7 @@ import (
 
 var Upgrade_0_2_1 = Upgrade{
 	UpgradeName: "v0.2.1",
-	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator) upgradetypes.UpgradeHandler {
+	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator, _ Keepers) upgradetypes.UpgradeHandler {
 		return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 			return mm.RunMigrations(ctx, configurator, fromVM)
 		}

app/upgrades/upgrade_0_2_4.go

@@ -10,7 +10,7 @@ import (
 
 var Upgrade_0_2_4 = Upgrade{
 	UpgradeName: "v0.2.4-rc11",
-	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator) upgradetypes.UpgradeHandler {
+	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator, _ Keepers) upgradetypes.UpgradeHandler {
 		return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 			return mm.RunMigrations(ctx, configurator, fromVM)
 		}

app/upgrades/upgrade_0_3_0.go

@@ -2,16 +2,51 @@ package upgrades
 
 import (
 	"context"
+	"encoding/base64"
+
+	"cosmossdk.io/api/cosmos/crypto/ed25519"
 	"cosmossdk.io/store/types"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
-    migrationmngrtypes "github.com/evstack/ev-abci/modules/migrationmngr/types"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	migrationmngrtypes "github.com/evstack/ev-abci/modules/migrationmngr/types"
 )
 
 var Upgrade_0_3_0 = Upgrade{
-	UpgradeName: "v0.3.0",
-	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator) upgradetypes.UpgradeHandler {
+	UpgradeName:   "v0.3.0",
+	UpgradeHeight: 5335502,
+	CreateUpgradeHandler: func(mm module.Manager, configurator module.Configurator, keepers Keepers) upgradetypes.UpgradeHandler {
 		return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+			sdkCtx := sdk.UnwrapSDKContext(ctx)
+			migrationHeight := uint64(sdkCtx.BlockHeight() + 2)
+
+			// Decode the base64 key from `tendermint show-validator`
+			keyBytes, err := base64.StdEncoding.DecodeString("z4OQC3UzB8xUlJKigQu/mI7mclgofa8qGQtp3RUPcqc=")
+			if err != nil {
+				return nil, err
+			}
+
+			// Create the ed25519 public key
+			pubKey := &ed25519.PubKey{Key: keyBytes}
+
+			// Wrap it in a codectypes.Any
+			pubKeyAny, err := codectypes.NewAnyWithValue(pubKey)
+			if err != nil {
+				return nil, err
+			}
+
+			if err := keepers.MigrationMngrKeeper.Migration.Set(ctx, migrationmngrtypes.EvolveMigration{
+				BlockHeight: migrationHeight,
+				Sequencer: migrationmngrtypes.Sequencer{
+					Name:            "polkachu",
+					ConsensusPubkey: pubKeyAny,
+				},
+				StayOnComet: true,
+			}); err != nil {
+				return nil, err
+			}
+
 			return mm.RunMigrations(ctx, configurator, fromVM)
 		}
 	},