add store command capability as well

api/src/opentrons/protocol_api/core/module.py

@@ -397,4 +397,3 @@ def retrieve(self) -> None:
     @abstractmethod
     def store(self) -> None:
         """Store a labware at the bottom of the labware stack."""
-        pass

api/src/opentrons/protocol_api/module_contexts.py

@@ -1163,4 +1163,5 @@ def store(self, labware: Labware) -> None:
 
         :param labware: The labware object to store.
         """
+        assert labware._core is not None
         self._core.store()

api/src/opentrons/protocol_engine/commands/flex_stacker/retrieve.py

@@ -53,7 +53,6 @@ def __init__(
 
     async def execute(self, params: RetrieveParams) -> SuccessData[RetrieveResult]:
         """Execute the labware retrieval command."""
-
         stacker_state = self._state_view.modules.get_flex_stacker_substate(
             params.moduleId
         )
@@ -90,7 +89,7 @@ async def execute(self, params: RetrieveParams) -> SuccessData[RetrieveResult]:
             new_location=ModuleLocation(moduleId=params.moduleId),
             new_offset_id=None,
         )
-        state_update.remove_flex_stacker_hopper_labware(
+        state_update.retrieve_flex_stacker_labware(
             module_id=params.moduleId, labware_id=lw_id
         )
         return SuccessData(

api/src/opentrons/protocol_engine/commands/flex_stacker/store.py

@@ -6,8 +6,15 @@
 from pydantic import BaseModel, Field
 
 from ..command import AbstractCommandImpl, BaseCommand, BaseCommandCreate, SuccessData
-from ...errors.error_occurrence import ErrorOccurrence
+from ...errors import ErrorOccurrence, CannotPerformModuleAction
 from ...state import update_types
+from ...types import (
+    ModuleModel,
+    AddressableAreaLocation,
+    ModuleLocation,
+    OFF_DECK_LOCATION,
+)
+
 
 if TYPE_CHECKING:
     from opentrons.protocol_engine.state.state import StateView
@@ -44,17 +51,36 @@ def __init__(
 
     async def execute(self, params: StoreParams) -> SuccessData[StoreResult]:
         """Execute the labware storage command."""
-        state_update = update_types.StateUpdate()
-        stacker_substate = self._state_view.modules.get_flex_stacker_substate(
-            module_id=params.moduleId
+        stacker_state = self._state_view.modules.get_flex_stacker_substate(
+            params.moduleId
         )
-
         # Allow propagation of ModuleNotAttachedError.
-        stacker = self._equipment.get_module_hardware_api(stacker_substate.module_id)
+        stacker_hw = self._equipment.get_module_hardware_api(stacker_state.module_id)
+
+        try:
+            lw_id = self._state_view.labware.get_id_by_module(params.moduleId)
+        except Exception as e:
+            raise CannotPerformModuleAction(
+                f"Cannot store labware if Flex Stacker carriage is empty"
+            )
+        lw_dim = self._state_view.labware.get_dimensions(labware_id=lw_id)
+
+        # TODO: check the type of the labware should match that already in the stack
 
-        if stacker is not None:
-            # TODO: get labware height from labware state view
-            await stacker.store_labware(labware_height=50.0)
+        state_update = update_types.StateUpdate()
+
+        if stacker_hw is not None:
+            await stacker_hw.store_labware(labware_height=lw_dim.z)
+
+        # update the state to reflect the labware is store in the stack
+        state_update.set_labware_location(
+            labware_id=lw_id,
+            new_location=OFF_DECK_LOCATION,
+            new_offset_id=None,
+        )
+        state_update.store_flex_stacker_labware(
+            module_id=params.moduleId, labware_id=lw_id
+        )
 
         return SuccessData(public=StoreResult(), state_update=state_update)
 

api/src/opentrons/protocol_engine/commands/load_labware.py

@@ -217,7 +217,7 @@ async def execute(  # noqa: C901
         if self._is_loading_to_module(
             params.location, ModuleModel.FLEX_STACKER_MODULE_V1
         ):
-            state_update.add_flex_stacker_hopper_labware(
+            state_update.load_flex_stacker_hopper_labware(
                 module_id=params.location.moduleId,
                 labware_id=loaded_labware.labware_id,
             )

api/src/opentrons/protocol_engine/state/module_substates/flex_stacker_substate.py

@@ -3,8 +3,9 @@
 from typing import NewType, List
 from opentrons.protocol_engine.state.update_types import (
     FlexStackerStateUpdate,
-    FlexStackerAddHopperLabware,
-    FlexStackerRemoveHopperLabware,
+    FlexStackerLoadHopperLabware,
+    FlexStackerRetrieveLabware,
+    FlexStackerStoreLabware,
 )
 
 
@@ -28,10 +29,17 @@ def new_from_state_change(
         """Return a new state with the given update applied."""
         lw_change = update.hopper_labware_update
         new_labware_ids = self.hopper_labware_ids.copy()
-        if isinstance(lw_change, FlexStackerAddHopperLabware):
+
+        # TODO the labware stack needs to be handled more elegantly
+        # this is a temporary solution to enable evt testing
+        if isinstance(lw_change, FlexStackerLoadHopperLabware):
+            # for manually loading labware in the stacker
             new_labware_ids.append(lw_change.labware_id)
-        elif isinstance(lw_change, FlexStackerRemoveHopperLabware):
+        elif isinstance(lw_change, FlexStackerRetrieveLabware):
             new_labware_ids.remove(lw_change.labware_id)
+        elif isinstance(lw_change, FlexStackerStoreLabware):
+            # automatically store labware at the bottom of the stack
+            new_labware_ids.insert(0, lw_change.labware_id)
 
         return FlexStackerSubState(
             module_id=self.module_id,

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

@@ -317,14 +317,21 @@ class AbsorbanceReaderStateUpdate:
 
 
 @dataclasses.dataclass
-class FlexStackerAddHopperLabware:
+class FlexStackerLoadHopperLabware:
     """An update to the Flex Stacker module static state."""
 
     labware_id: str
 
 
 @dataclasses.dataclass
-class FlexStackerRemoveHopperLabware:
+class FlexStackerRetrieveLabware:
+    """An update to the Flex Stacker module static state."""
+
+    labware_id: str
+
+
+@dataclasses.dataclass
+class FlexStackerStoreLabware:
     """An update to the Flex Stacker module static state."""
 
     labware_id: str
@@ -335,7 +342,7 @@ class FlexStackerStateUpdate:
     """An update to the Flex Stacker module state."""
 
     module_id: str
-    hopper_labware_update: FlexStackerAddHopperLabware | FlexStackerRemoveHopperLabware | NoChangeType = (
+    hopper_labware_update: FlexStackerLoadHopperLabware | FlexStackerRetrieveLabware | FlexStackerStoreLabware | NoChangeType = (
         NO_CHANGE
     )
 
@@ -744,26 +751,38 @@ def set_addressable_area_used(self: Self, addressable_area_name: str) -> Self:
         )
         return self
 
-    def add_flex_stacker_hopper_labware(
+    def load_flex_stacker_hopper_labware(
+        self,
+        module_id: str,
+        labware_id: str,
+    ) -> Self:
+        """Add a labware definition to the engine."""
+        self.flex_stacker_state_update = FlexStackerStateUpdate(
+            module_id=module_id,
+            hopper_labware_update=FlexStackerLoadHopperLabware(labware_id=labware_id),
+        )
+        return self
+
+    def retrieve_flex_stacker_labware(
         self,
         module_id: str,
         labware_id: str,
     ) -> Self:
         """Add a labware definition to the engine."""
         self.flex_stacker_state_update = FlexStackerStateUpdate(
             module_id=module_id,
-            hopper_labware_update=FlexStackerAddHopperLabware(labware_id=labware_id),
+            hopper_labware_update=FlexStackerRetrieveLabware(labware_id=labware_id),
         )
         return self
 
-    def remove_flex_stacker_hopper_labware(
+    def store_flex_stacker_labware(
         self,
         module_id: str,
         labware_id: str,
     ) -> Self:
         """Add a labware definition to the engine."""
         self.flex_stacker_state_update = FlexStackerStateUpdate(
             module_id=module_id,
-            hopper_labware_update=FlexStackerRemoveHopperLabware(labware_id=labware_id),
+            hopper_labware_update=FlexStackerStoreLabware(labware_id=labware_id),
         )
         return self

api/tests/opentrons/protocol_engine/commands/flex_stacker/test_retrieve.py

@@ -12,6 +12,7 @@
 from opentrons.protocol_engine.commands import flex_stacker
 from opentrons.protocol_engine.commands.command import SuccessData
 from opentrons.protocol_engine.commands.flex_stacker.retrieve import RetrieveImpl
+from opentrons.protocol_engine.types import Dimensions
 
 
 async def test_retrieve(
@@ -33,6 +34,10 @@ async def test_retrieve(
     decoy.when(fs_module_substate.module_id).then_return(
         FlexStackerId("flex-stacker-id")
     )
+    decoy.when(fs_module_substate.hopper_labware_ids).then_return(["labware-id"])
+    decoy.when(state_view.labware.get_dimensions(labware_id="labware-id")).then_return(
+        Dimensions(x=1, y=1, z=1)
+    )
 
     decoy.when(
         equipment.get_module_hardware_api(FlexStackerId("flex-stacker-id"))
@@ -41,5 +46,5 @@ async def test_retrieve(
     result = await subject.execute(data)
     decoy.verify(await fs_hardware.dispense_labware(labware_height=50.0), times=1)
     assert result == SuccessData(
-        public=flex_stacker.RetrieveResult(),
+        public=flex_stacker.RetrieveResult(labware_id="labware-id"),
     )

api/tests/opentrons/protocol_engine/commands/test_load_labware.py

@@ -32,7 +32,7 @@
     AddressableAreaUsedUpdate,
     LoadedLabwareUpdate,
     FlexStackerStateUpdate,
-    FlexStackerAddHopperLabware,
+    FlexStackerLoadHopperLabware,
     StateUpdate,
 )
 
@@ -303,7 +303,7 @@ async def test_load_labware_in_flex_stacker(
             ),
             flex_stacker_state_update=FlexStackerStateUpdate(
                 module_id="some-module-id",
-                hopper_labware_update=FlexStackerAddHopperLabware(
+                hopper_labware_update=FlexStackerLoadHopperLabware(
                     labware_id="labware-id"
                 ),
             ),