diff --git a/execution/executor-types/src/state_compute_result.rs b/execution/executor-types/src/state_compute_result.rs
index c47ceeb97c4205..390818228b8da2 100644
--- a/execution/executor-types/src/state_compute_result.rs
+++ b/execution/executor-types/src/state_compute_result.rs
@@ -165,7 +165,7 @@ impl StateComputeResult {
                 .state_updates_before_last_checkpoint
                 .as_ref(),
             sharded_state_cache: Some(&self.execution_output.state_cache.sharded_state_cache),
-            is_reconfig: self.execution_output.block_end_info.is_some(),
+            is_reconfig: self.execution_output.next_epoch_state.is_some(),
         }
     }
 }
diff --git a/storage/aptosdb/src/db/include/aptosdb_writer.rs b/storage/aptosdb/src/db/include/aptosdb_writer.rs
index e2e9416933bcc7..a67f8f19ac9f04 100644
--- a/storage/aptosdb/src/db/include/aptosdb_writer.rs
+++ b/storage/aptosdb/src/db/include/aptosdb_writer.rs
@@ -615,6 +615,18 @@ impl AptosDB {
             current_epoch,
         );
 
+        // Ensure that state tree at the end of the epoch is persisted
+        if ledger_info_with_sig.ledger_info().ends_epoch() {
+            let state_snapshot = self.state_store.get_state_snapshot_before(version + 1)?;
+            ensure!(
+                state_snapshot.is_some()
+                && state_snapshot.unwrap().0 == version,
+                "State checkpoint not persisted at the end of the epoch, version {}, epoch {}",
+                version,
+                ledger_info_with_sig.ledger_info().epoch(),
+            );
+        }
+
         // Put write to batch.
         self.ledger_db
             .metadata_db()