From 33225f7218db215329abea804ab7bdb7505fa4f4 Mon Sep 17 00:00:00 2001
From: Max Marrone <max@opentrons.com>
Date: Thu, 28 Mar 2024 16:22:48 -0400
Subject: [PATCH] Internal support for waiting for a specific command's
 recovery.

---
 .../protocol_engine/protocol_engine.py        | 21 +++++++++++++++++++
 .../protocol_engine/state/commands.py         | 20 ++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/api/src/opentrons/protocol_engine/protocol_engine.py b/api/src/opentrons/protocol_engine/protocol_engine.py
index 5056385dea5..e7f9e4eb097 100644
--- a/api/src/opentrons/protocol_engine/protocol_engine.py
+++ b/api/src/opentrons/protocol_engine/protocol_engine.py
@@ -242,6 +242,27 @@ async def add_and_execute_command(
         await self.wait_for_command(command.id)
         return self._state_store.commands.get(command.id)
 
+    async def add_and_execute_command_wait_for_recovery(
+        self, request: commands.CommandCreate
+    ) -> None:
+        """Like `add_and_execute_command()`, except wait for error recovery.
+
+        Unlike `add_and_execute_command()`, if the command fails, this will not
+        immediately return the failed command. Instead, if the error is recoverable,
+        it will wait until error recovery has completed (e.g. when some other task
+        calls `self.resume_from_recovery()`).
+        """
+        command = self.add_command(request)
+        await self.wait_for_command(command_id=command.id)
+        await self._state_store.wait_for(
+            # Per the warnings on `wait_for()`, we want our condition function to
+            # specifically check that *this* command's recovery has completed,
+            # rather than just checking that the overall run state
+            # != "awaiting-recovery."
+            self.state_view.commands.get_command_recovery_complete,
+            command_id=command.id,
+        )
+
     def estop(
         self,
         # TODO(mm, 2024-03-26): Maintenance runs are a robot-server concept that
diff --git a/api/src/opentrons/protocol_engine/state/commands.py b/api/src/opentrons/protocol_engine/state/commands.py
index 7a4b88c7cd2..ca8536367b4 100644
--- a/api/src/opentrons/protocol_engine/state/commands.py
+++ b/api/src/opentrons/protocol_engine/state/commands.py
@@ -715,6 +715,26 @@ def get_all_commands_final(self) -> bool:
 
         return no_command_running and no_command_to_execute
 
+    def get_command_recovery_complete(self, command_id: str) -> bool:
+        """Return whether we're finished with error recovery for the given command.
+
+        If the given command didn't have a recovery phase (because it hasn't completed
+        yet, or because it succeeded without an error, or because it failed with an
+        error that wasn't recoverable), that counts as `True`.
+
+        If the given command did have a recovery phase, but it was interrupted by a
+        stop, that also counts as `True`.
+        """
+        command_is_most_recent_to_fail = (
+            self._state.failed_command is not None
+            and command_id == self._state.failed_command.command.id
+        )
+        if command_is_most_recent_to_fail:
+            recovery_is_ongoing = self.get_status() == EngineStatus.AWAITING_RECOVERY
+            return not recovery_is_ongoing
+        else:
+            return True
+
     def raise_fatal_command_error(self) -> None:
         """Raise the run's fatal command error, if there was one, as an exception.