From 05c9986ffa114e1ca5c193630b1adb434db0e4d3 Mon Sep 17 00:00:00 2001
From: Max Marrone <max@opentrons.com>
Date: Tue, 19 Mar 2024 13:32:15 -0400
Subject: [PATCH] Allow `resume-from-recovery` and `stop` (but nothing else)
 while `awaiting-recovery`.

---
 .../protocol_engine/state/commands.py         | 43 +++++++++++--------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/api/src/opentrons/protocol_engine/state/commands.py b/api/src/opentrons/protocol_engine/state/commands.py
index b27154348bd..fedbb126aee 100644
--- a/api/src/opentrons/protocol_engine/state/commands.py
+++ b/api/src/opentrons/protocol_engine/state/commands.py
@@ -725,38 +725,47 @@ def validate_action_allowed(  # noqa: C901
             SetupCommandNotAllowedError: The engine is running, so a setup command
                 may not be added.
         """
-        if self.get_status() == EngineStatus.AWAITING_RECOVERY:
-            # While we're developing error recovery, we'll conservatively disallow
-            # all actions, to avoid putting the engine in weird undefined states.
-            # We'll allow specific actions here as we flesh things out and add support
-            # for them.
-            raise NotImplementedError()
-
-        if isinstance(action, ResumeFromRecoveryAction):
-            # https://opentrons.atlassian.net/browse/EXEC-301
-            raise NotImplementedError()
-
         if self._state.run_result is not None:
             raise RunStoppedError("The run has already stopped.")
 
         elif isinstance(action, PlayAction):
             if self.get_status() == EngineStatus.BLOCKED_BY_OPEN_DOOR:
                 raise RobotDoorOpenError("Front door or top window is currently open.")
+            elif self.get_status() == EngineStatus.AWAITING_RECOVERY:
+                raise NotImplementedError()
+            else:
+                return action
 
         elif isinstance(action, PauseAction):
             if not self.get_is_running():
                 raise PauseNotAllowedError("Cannot pause a run that is not running.")
+            elif self.get_status() == EngineStatus.AWAITING_RECOVERY:
+                raise NotImplementedError()
+            else:
+                return action
 
-        elif (
-            isinstance(action, QueueCommandAction)
-            and action.request.intent == CommandIntent.SETUP
-        ):
-            if self._state.queue_status != QueueStatus.SETUP:
+        elif isinstance(action, QueueCommandAction):
+            if (
+                action.request.intent == CommandIntent.SETUP
+                and self._state.queue_status != QueueStatus.SETUP
+            ):
                 raise SetupCommandNotAllowedError(
                     "Setup commands are not allowed after run has started."
                 )
+            else:
+                return action
+
+        elif isinstance(action, ResumeFromRecoveryAction):
+            if self.get_status() != EngineStatus.AWAITING_RECOVERY:
+                raise NotImplementedError()
+            else:
+                return action
 
-        return action
+        elif isinstance(action, StopAction):
+            return action
+
+        else:
+            assert_never(action)
 
     def get_status(self) -> EngineStatus:
         """Get the current execution status of the engine."""