refactor(api): Ensure we handle state updates even for failed commands

api/src/opentrons/protocol_engine/actions/__init__.py

@@ -30,6 +30,7 @@
     SetPipetteMovementSpeedAction,
     AddAbsorbanceReaderLidAction,
 )
+from .get_state_update import get_state_update
 
 __all__ = [
     # action pipeline interface
@@ -61,4 +62,6 @@
     # action payload values
     "PauseSource",
     "FinishErrorDetails",
+    # helper functions
+    "get_state_update",
 ]

api/src/opentrons/protocol_engine/actions/get_state_update.py

@@ -0,0 +1,18 @@
+# noqa: D100
+
+
+from .actions import Action, SucceedCommandAction, FailCommandAction
+from ..commands.command import DefinedErrorData
+from ..state.update_types import StateUpdate
+
+
+def get_state_update(action: Action) -> StateUpdate | None:
+    """Extract the StateUpdate from an action, if there is one."""
+    if isinstance(action, SucceedCommandAction):
+        return action.state_update
+    elif isinstance(action, FailCommandAction) and isinstance(
+        action.error, DefinedErrorData
+    ):
+        return action.error.state_update
+    else:
+        return None

api/src/opentrons/protocol_engine/state/labware.py

@@ -48,9 +48,9 @@
 )
 from ..actions import (
     Action,
-    SucceedCommandAction,
     AddLabwareOffsetAction,
     AddLabwareDefinitionAction,
+    get_state_update,
 )
 from ._abstract_store import HasState, HandlesActions
 from ._move_types import EdgePathType
@@ -64,8 +64,6 @@
     "opentrons/usascientific_96_wellplate_2.4ml_deep/1",
 }
 
-_OT3_INSTRUMENT_ATTACH_SLOT = DeckSlotName.SLOT_D1
-
 _RIGHT_SIDE_SLOTS = {
     # OT-2:
     DeckSlotName.FIXED_TRASH,
@@ -148,10 +146,12 @@ def __init__(
 
     def handle_action(self, action: Action) -> None:
         """Modify state in reaction to an action."""
-        if isinstance(action, SucceedCommandAction):
-            self._handle_command(action)
+        state_update = get_state_update(action)
+        if state_update is not None:
+            self._add_loaded_labware(state_update)
+            self._set_labware_location(state_update)
 
-        elif isinstance(action, AddLabwareOffsetAction):
+        if isinstance(action, AddLabwareOffsetAction):
             labware_offset = LabwareOffset.construct(
                 id=action.labware_offset_id,
                 createdAt=action.created_at,
@@ -169,11 +169,6 @@ def handle_action(self, action: Action) -> None:
             )
             self._state.definitions_by_uri[uri] = action.definition
 
-    def _handle_command(self, action: Action) -> None:
-        """Modify state in reaction to a command."""
-        self._add_loaded_labware(action)
-        self._set_labware_location(action)
-
     def _add_labware_offset(self, labware_offset: LabwareOffset) -> None:
         """Add a new labware offset to state.
 
@@ -185,56 +180,50 @@ def _add_labware_offset(self, labware_offset: LabwareOffset) -> None:
 
         self._state.labware_offsets_by_id[labware_offset.id] = labware_offset
 
-    def _add_loaded_labware(self, action: Action) -> None:
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.loaded_labware != update_types.NO_CHANGE
-        ):
+    def _add_loaded_labware(self, state_update: update_types.StateUpdate) -> None:
+        loaded_labware_update = state_update.loaded_labware
+        if loaded_labware_update != update_types.NO_CHANGE:
             # If the labware load refers to an offset, that offset must actually exist.
-            if action.state_update.loaded_labware.offset_id is not None:
+            if loaded_labware_update.offset_id is not None:
                 assert (
-                    action.state_update.loaded_labware.offset_id
-                    in self._state.labware_offsets_by_id
+                    loaded_labware_update.offset_id in self._state.labware_offsets_by_id
                 )
 
             definition_uri = uri_from_details(
-                namespace=action.state_update.loaded_labware.definition.namespace,
-                load_name=action.state_update.loaded_labware.definition.parameters.loadName,
-                version=action.state_update.loaded_labware.definition.version,
+                namespace=loaded_labware_update.definition.namespace,
+                load_name=loaded_labware_update.definition.parameters.loadName,
+                version=loaded_labware_update.definition.version,
             )
 
             self._state.definitions_by_uri[
                 definition_uri
-            ] = action.state_update.loaded_labware.definition
+            ] = loaded_labware_update.definition
 
-            location = action.state_update.loaded_labware.new_location
+            location = loaded_labware_update.new_location
 
-            display_name = action.state_update.loaded_labware.display_name
+            display_name = loaded_labware_update.display_name
 
             self._state.labware_by_id[
-                action.state_update.loaded_labware.labware_id
+                loaded_labware_update.labware_id
             ] = LoadedLabware.construct(
-                id=action.state_update.loaded_labware.labware_id,
+                id=loaded_labware_update.labware_id,
                 location=location,
-                loadName=action.state_update.loaded_labware.definition.parameters.loadName,
+                loadName=loaded_labware_update.definition.parameters.loadName,
                 definitionUri=definition_uri,
-                offsetId=action.state_update.loaded_labware.offset_id,
+                offsetId=loaded_labware_update.offset_id,
                 displayName=display_name,
             )
 
-    def _set_labware_location(self, action: Action) -> None:
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.labware_location != update_types.NO_CHANGE
-        ):
-
-            labware_id = action.state_update.labware_location.labware_id
-            new_offset_id = action.state_update.labware_location.offset_id
+    def _set_labware_location(self, state_update: update_types.StateUpdate) -> None:
+        labware_location_update = state_update.labware_location
+        if labware_location_update != update_types.NO_CHANGE:
+            labware_id = labware_location_update.labware_id
+            new_offset_id = labware_location_update.offset_id
 
             self._state.labware_by_id[labware_id].offsetId = new_offset_id
 
-            if action.state_update.labware_location.new_location:
-                new_location = action.state_update.labware_location.new_location
+            if labware_location_update.new_location:
+                new_location = labware_location_update.new_location
 
                 if isinstance(
                     new_location, AddressableAreaLocation

api/src/opentrons/protocol_engine/state/pipettes.py

@@ -10,9 +10,7 @@
     Tuple,
     Union,
 )
-from typing_extensions import assert_type
 
-from opentrons_shared_data.errors import EnumeratedError
 from opentrons_shared_data.pipette import pipette_definition
 from opentrons.config.defaults_ot2 import Z_RETRACT_DISTANCE
 from opentrons.hardware_control.dev_types import PipetteDict
@@ -21,8 +19,6 @@
     NozzleConfigurationType,
     NozzleMap,
 )
-from opentrons.protocol_engine.actions.actions import FailCommandAction
-from opentrons.protocol_engine.commands.command import DefinedErrorData
 from opentrons.types import MountType, Mount as HwMount, Point
 
 from . import update_types
@@ -40,8 +36,10 @@
 )
 from ..actions import (
     Action,
+    FailCommandAction,
     SetPipetteMovementSpeedAction,
     SucceedCommandAction,
+    get_state_update,
 )
 from ._abstract_store import HasState, HandlesActions
 
@@ -140,37 +138,31 @@ def __init__(self) -> None:
 
     def handle_action(self, action: Action) -> None:
         """Modify state in reaction to an action."""
+        state_update = get_state_update(action)
+        if state_update is not None:
+            self._set_load_pipette(state_update)
+            self._update_current_location(state_update)
+            self._update_pipette_config(state_update)
+            self._update_pipette_nozzle_map(state_update)
+            self._update_tip_state(state_update)
+
         if isinstance(action, (SucceedCommandAction, FailCommandAction)):
-            self._handle_command(action)
+            self._update_volumes(action)
+
         elif isinstance(action, SetPipetteMovementSpeedAction):
             self._state.movement_speed_by_id[action.pipette_id] = action.speed
 
-    def _handle_command(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-        self._set_load_pipette(action)
-        self._update_current_location(action)
-        self._update_pipette_config(action)
-        self._update_pipette_nozzle_map(action)
-        self._update_tip_state(action)
-        self._update_volumes(action)
-
-    def _set_load_pipette(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.loaded_pipette != update_types.NO_CHANGE
-        ):
-            pipette_id = action.state_update.loaded_pipette.pipette_id
+    def _set_load_pipette(self, state_update: update_types.StateUpdate) -> None:
+        if state_update.loaded_pipette != update_types.NO_CHANGE:
+            pipette_id = state_update.loaded_pipette.pipette_id
 
             self._state.pipettes_by_id[pipette_id] = LoadedPipette(
                 id=pipette_id,
-                pipetteName=action.state_update.loaded_pipette.pipette_name,
-                mount=action.state_update.loaded_pipette.mount,
+                pipetteName=state_update.loaded_pipette.pipette_name,
+                mount=state_update.loaded_pipette.mount,
             )
             self._state.liquid_presence_detection_by_id[pipette_id] = (
-                action.state_update.loaded_pipette.liquid_presence_detection or False
+                state_update.loaded_pipette.liquid_presence_detection or False
             )
             self._state.aspirated_volume_by_id[pipette_id] = None
             self._state.movement_speed_by_id[pipette_id] = None
@@ -181,17 +173,11 @@ def _set_load_pipette(
                     pipette_id
                 ] = static_config.default_nozzle_map
 
-    def _update_tip_state(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.pipette_tip_state != update_types.NO_CHANGE
-        ):
-            pipette_id = action.state_update.pipette_tip_state.pipette_id
-            if action.state_update.pipette_tip_state.tip_geometry:
-                attached_tip = action.state_update.pipette_tip_state.tip_geometry
+    def _update_tip_state(self, state_update: update_types.StateUpdate) -> None:
+        if state_update.pipette_tip_state != update_types.NO_CHANGE:
+            pipette_id = state_update.pipette_tip_state.pipette_id
+            if state_update.pipette_tip_state.tip_geometry:
+                attached_tip = state_update.pipette_tip_state.tip_geometry
 
                 self._state.attached_tip_by_id[pipette_id] = attached_tip
                 self._state.aspirated_volume_by_id[pipette_id] = 0
@@ -220,7 +206,7 @@ def _update_tip_state(
                     )
 
             else:
-                pipette_id = action.state_update.pipette_tip_state.pipette_id
+                pipette_id = state_update.pipette_tip_state.pipette_id
                 self._state.aspirated_volume_by_id[pipette_id] = None
                 self._state.attached_tip_by_id[pipette_id] = None
 
@@ -236,17 +222,8 @@ def _update_tip_state(
                         default_dispense=tip_configuration.default_dispense_flowrate.values_by_api_level,
                     )
 
-    def _update_current_location(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-        if isinstance(action, SucceedCommandAction):
-            location_update = action.state_update.pipette_location
-        elif isinstance(action.error, DefinedErrorData):
-            location_update = action.error.state_update.pipette_location
-        else:
-            # The command failed with some undefined error. We have nothing to do.
-            assert_type(action.error, EnumeratedError)
-            return
+    def _update_current_location(self, state_update: update_types.StateUpdate) -> None:
+        location_update = state_update.pipette_location
 
         if location_update is update_types.NO_CHANGE:
             pass
@@ -282,18 +259,13 @@ def _update_current_location(
                     mount=loaded_pipette.mount, deck_point=new_deck_point
                 )
 
-    def _update_pipette_config(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
-    ) -> None:
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.pipette_config != update_types.NO_CHANGE
-        ):
-            config = action.state_update.pipette_config.config
+    def _update_pipette_config(self, state_update: update_types.StateUpdate) -> None:
+        if state_update.pipette_config != update_types.NO_CHANGE:
+            config = state_update.pipette_config.config
             self._state.static_config_by_id[
-                action.state_update.pipette_config.pipette_id
+                state_update.pipette_config.pipette_id
             ] = StaticPipetteConfig(
-                serial_number=action.state_update.pipette_config.serial_number,
+                serial_number=state_update.pipette_config.serial_number,
                 model=config.model,
                 display_name=config.display_name,
                 min_volume=config.min_volume,
@@ -325,26 +297,26 @@ def _update_pipette_config(
                 lld_settings=config.pipette_lld_settings,
             )
             self._state.flow_rates_by_id[
-                action.state_update.pipette_config.pipette_id
+                state_update.pipette_config.pipette_id
             ] = config.flow_rates
             self._state.nozzle_configuration_by_id[
-                action.state_update.pipette_config.pipette_id
+                state_update.pipette_config.pipette_id
             ] = config.nozzle_map
 
     def _update_pipette_nozzle_map(
-        self, action: Union[SucceedCommandAction, FailCommandAction]
+        self, state_update: update_types.StateUpdate
     ) -> None:
-        if (
-            isinstance(action, SucceedCommandAction)
-            and action.state_update.pipette_nozzle_map != update_types.NO_CHANGE
-        ):
+        if state_update.pipette_nozzle_map != update_types.NO_CHANGE:
             self._state.nozzle_configuration_by_id[
-                action.state_update.pipette_nozzle_map.pipette_id
-            ] = action.state_update.pipette_nozzle_map.nozzle_map
+                state_update.pipette_nozzle_map.pipette_id
+            ] = state_update.pipette_nozzle_map.nozzle_map
 
     def _update_volumes(
         self, action: Union[SucceedCommandAction, FailCommandAction]
     ) -> None:
+        # todo(mm, 2024-10-10): Port these isinstance checks to StateUpdate.
+        # https://opentrons.atlassian.net/browse/EXEC-754
+
         if isinstance(action, SucceedCommandAction) and isinstance(
             action.command.result,
             (commands.AspirateResult, commands.AspirateInPlaceResult),