From 9585307f729247b559cddbd2636b8762a8f3e7bc Mon Sep 17 00:00:00 2001
From: Evan B <ebatsell@gmail.com>
Date: Wed, 11 Sep 2024 16:28:54 -0400
Subject: [PATCH] Force state transitions to Idle before doing work in new
 epoch (#82)

Currently a single lingering ComputeInstantUnstake or Rebalance step can
be executed if the Steward was left in this state after rolling over to
a new epoch. This can affect stake based on outdated information which
we don't want. Force the transition to Idle so everything is handled
normally.
---
 .../src/instructions/compute_instant_unstake.rs       | 11 +++++++++++
 programs/steward/src/instructions/rebalance.rs        | 11 +++++++++++
 2 files changed, 22 insertions(+)

diff --git a/programs/steward/src/instructions/compute_instant_unstake.rs b/programs/steward/src/instructions/compute_instant_unstake.rs
index 367f57f4..28f52135 100644
--- a/programs/steward/src/instructions/compute_instant_unstake.rs
+++ b/programs/steward/src/instructions/compute_instant_unstake.rs
@@ -42,6 +42,17 @@ pub fn handler(ctx: Context<ComputeInstantUnstake>, validator_list_index: usize)
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
+    // Transitions to Idle before doing compute_instant_unstake if RESET_TO_IDLE is set
+    if let Some(event) = maybe_transition(
+        &mut state_account.state,
+        &clock,
+        &config.parameters,
+        &epoch_schedule,
+    )? {
+        emit!(event);
+        return Ok(());
+    }
+
     state_checks(
         &clock,
         &config,
diff --git a/programs/steward/src/instructions/rebalance.rs b/programs/steward/src/instructions/rebalance.rs
index e706ffd4..b2157806 100644
--- a/programs/steward/src/instructions/rebalance.rs
+++ b/programs/steward/src/instructions/rebalance.rs
@@ -153,6 +153,17 @@ pub fn handler(ctx: Context<Rebalance>, validator_list_index: usize) -> Result<(
     {
         let mut state_account = ctx.accounts.state_account.load_mut()?;
 
+        // Transitions to Idle before doing rebalance if RESET_TO_IDLE is set
+        if let Some(event) = maybe_transition(
+            &mut state_account.state,
+            &clock,
+            &config.parameters,
+            &epoch_schedule,
+        )? {
+            emit!(event);
+            return Ok(());
+        }
+
         state_checks(
             &clock,
             &config,