From ed8da3553807ba06fd12367be13b0a8da264d07d Mon Sep 17 00:00:00 2001
From: Adamantios <adamantios.zaras@valory.xyz>
Date: Wed, 15 Nov 2023 12:40:46 +0200
Subject: [PATCH] fix: add a separate exit round for the checkpoint call

---
 packages/valory/skills/staking_abci/rounds.py     | 14 +++++++++++---
 packages/valory/skills/trader_abci/composition.py |  4 +++-
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/packages/valory/skills/staking_abci/rounds.py b/packages/valory/skills/staking_abci/rounds.py
index 36d1127e7..612f26997 100644
--- a/packages/valory/skills/staking_abci/rounds.py
+++ b/packages/valory/skills/staking_abci/rounds.py
@@ -112,6 +112,10 @@ def end_block(self) -> Optional[Tuple[BaseSynchronizedData, Enum]]:
         return res
 
 
+class CheckpointCallPreparedRound(DegenerateRound, ABC):
+    """A round that represents staking has finished with a checkpoint call safe tx prepared."""
+
+
 class FinishedStakingRound(DegenerateRound, ABC):
     """A round that represents staking has finished."""
 
@@ -137,26 +141,30 @@ class StakingAbciApp(AbciApp[Event]):  # pylint: disable=too-few-public-methods
     initial_round_cls: Type[AbstractRound] = CallCheckpointRound
     transition_function: AbciAppTransitionFunction = {
         CallCheckpointRound: {
-            Event.DONE: FinishedStakingRound,
+            Event.DONE: CheckpointCallPreparedRound,
             Event.SERVICE_NOT_STAKED: FinishedStakingRound,
             Event.NEXT_CHECKPOINT_NOT_REACHED_YET: FinishedStakingRound,
             Event.ROUND_TIMEOUT: CallCheckpointRound,
             Event.NO_MAJORITY: CallCheckpointRound,
         },
+        CheckpointCallPreparedRound: {},
         FinishedStakingRound: {},
     }
     cross_period_persisted_keys = frozenset(
         {get_name(SynchronizedData.is_service_staked)}
     )
-    final_states: Set[AppState] = {FinishedStakingRound}
+    final_states: Set[AppState] = {CheckpointCallPreparedRound, FinishedStakingRound}
     event_to_timeout: Dict[Event, float] = {
         Event.ROUND_TIMEOUT: 30.0,
     }
     db_pre_conditions: Dict[AppState, Set[str]] = {CallCheckpointRound: set()}
     db_post_conditions: Dict[AppState, Set[str]] = {
-        FinishedStakingRound: {
+        CheckpointCallPreparedRound: {
             get_name(SynchronizedData.tx_submitter),
             get_name(SynchronizedData.most_voted_tx_hash),
             get_name(SynchronizedData.is_service_staked),
         },
+        FinishedStakingRound: {
+            get_name(SynchronizedData.is_service_staked),
+        },
     }
diff --git a/packages/valory/skills/trader_abci/composition.py b/packages/valory/skills/trader_abci/composition.py
index 7df594810..dbc896105 100644
--- a/packages/valory/skills/trader_abci/composition.py
+++ b/packages/valory/skills/trader_abci/composition.py
@@ -58,6 +58,7 @@
 )
 from packages.valory.skills.staking_abci.rounds import (
     CallCheckpointRound,
+    CheckpointCallPreparedRound,
     FinishedStakingRound,
     StakingAbciApp,
 )
@@ -97,7 +98,8 @@
     FinishedRedeemingTxRound: CallCheckpointRound,
     FinishedWithoutDecisionRound: RedeemRound,
     FinishedWithoutRedeemingRound: CallCheckpointRound,
-    FinishedStakingRound: RandomnessTransactionSubmissionRound,
+    FinishedStakingRound: ResetAndPauseRound,
+    CheckpointCallPreparedRound: RandomnessTransactionSubmissionRound,
     FinishedStakingTxRound: ResetAndPauseRound,
     FinishedResetAndPauseRound: UpdateBetsRound,
     FinishedResetAndPauseErrorRound: RegistrationRound,