updated implementation to include Aspirate and Dispense results

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

@@ -54,19 +54,6 @@ async def execute(
         self._state_view.labware.validate_liquid_allowed_in_labware(
             labware_id=params.labwareId, wells=params.volumeByWell
         )
-        # do this in WellStore? Don't have access to GeometryView
-        labware_id = params.labwareId
-        well_name = next(iter(params.volumeByWell))
-        volume = next(iter(params.volumeByWell.values()))
-        height = self._state_view.geometry.get_well_height_at_volume(
-            labware_id=labware_id, well_name=well_name, volume=volume
-        )
-        self._state_view.wells.set_liquid_height_from_load(
-            labware_id=labware_id,
-            well_name=well_name,
-            height=height,
-            time=datetime.now(),
-        )
 
         return SuccessData(public=LoadLiquidResult(), private=None)
 

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

@@ -1386,9 +1386,9 @@ def get_well_offset_adjustment(
             volume = operation_volume or 0.0
 
         if volume:
-            well_geometry = self._labware.get_well_geometry(labware_id, well_name)
             return self.get_well_height_after_volume(
-                well_geometry=well_geometry,
+                labware_id=labware_id,
+                well_name=well_name,
                 initial_height=initial_handling_height,
                 volume=volume,
             )
@@ -1431,12 +1431,13 @@ def get_well_handling_height(
         return float(handling_height)
 
     def get_well_height_after_volume(
-        self, well_geometry: InnerWellGeometry, initial_height: float, volume: float
+        self, labware_id: str, well_name: str, initial_height: float, volume: float
     ) -> float:
         """Return the height of liquid in a labware well after a given volume has been handled.
 
         This is given an initial handling height, with reference to the well bottom.
         """
+        well_geometry = self._labware.get_well_geometry(labware_id=labware_id, well_name=well_name)
         initial_volume = find_volume_at_well_height(
             target_height=initial_height, well_geometry=well_geometry
         )

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

@@ -7,12 +7,16 @@
     SucceedCommandAction,
 )
 from opentrons.protocol_engine.commands.liquid_probe import LiquidProbeResult
+from opentrons.protocol_engine.commands.load_liquid import LoadLiquidResult
+from opentrons.protocol_engine.commands.aspirate import AspirateResult
+from opentrons.protocol_engine.commands.dispense import DispenseResult
 from opentrons.protocol_engine.commands.pipetting_common import LiquidNotFoundError
 from opentrons.protocol_engine.types import LiquidHeightInfo, LiquidHeightSummary
 
 from ._abstract_store import HasState, HandlesActions
 from ..actions import Action
 from ..commands import Command
+from .geometry import get_well_height_at_volume, get_well_height_after_volume
 
 
 @dataclass
@@ -40,11 +44,32 @@ def handle_action(self, action: Action) -> None:
 
     def _handle_succeeded_command(self, command: Command) -> None:
         if isinstance(command.result, LiquidProbeResult):
-            self._set_liquid_height_from_probe(
+            self._set_liquid_height_after_probe(
                 labware_id=command.params.labwareId,
                 well_name=command.params.wellName,
                 height=command.result.z_position,
-                time=command.createdAt,
+                time=command.completedAt,
+            )
+        if isinstance(command.result, LoadLiquidResult):
+            self._set_liquid_height_after_load(
+                labware_id=command.params.labwareId,
+                well_name=next(iter(command.params.volumeByWell)),
+                volume=next(iter(command.params.volumeByWell.values())),
+                time=command.completedAt,
+            )
+        if isinstance(command.result, AspirateResult):
+            self._update_liquid_height_after_operation(
+                labware_id=command.params.labwareId,
+                well_name=command.params.wellName,
+                volume=-command.result.volume,
+                time=command.completedAt,
+            )
+        if isinstance(command.result, DispenseResult):
+            self._update_liquid_height_after_operation(
+                labware_id=command.params.labwareId,
+                well_name=command.params.wellName,
+                volume=command.result.volume,
+                time=command.completedAt,
             )
 
     def _handle_failed_command(self, action: FailCommandAction) -> None:
@@ -56,7 +81,7 @@ def _handle_failed_command(self, action: FailCommandAction) -> None:
                 time=action.failed_at,
             )
 
-    def _set_liquid_height_from_probe(
+    def _set_liquid_height_after_probe(
         self, labware_id: str, well_name: str, height: float, time: datetime
     ) -> None:
         """Set the liquid height of the well from a LiquidProbe command."""
@@ -65,6 +90,27 @@ def _set_liquid_height_from_probe(
             self._state.measured_liquid_heights[labware_id] = {}
         self._state.measured_liquid_heights[labware_id][well_name] = lhi
 
+    def _set_liquid_height_after_load(
+        self, labware_id: str, well_name: str, volume: float, time: datetime
+    ) -> None:
+        """Set the liquid height of the well from a LoadLiquid command."""
+        height = get_well_height_at_volume(labware_id=labware_id, well_name=well_name, volume=volume)
+        lhi = LiquidHeightInfo(height=height, last_measured=time)
+        if labware_id not in self._state.measured_liquid_heights:
+            self._state.measured_liquid_heights[labware_id] = {}
+        self._state.measured_liquid_heights[labware_id][well_name] = lhi
+
+    def _update_liquid_height_after_operation(
+        self, labware_id: str, well_name: str, volume: float, time: datetime
+    ) -> None:
+        """Set the liquid height of the well from a LoadLiquid command."""
+        initial_height = self._state.measured_liquid_heights[labware_id][well_name]
+        height = get_well_height_after_volume(labware_id=labware_id, well_name=well_name, initial_height=initial_height, volume=volume)
+        lhi = LiquidHeightInfo(height=height, last_measured=time)
+        if labware_id not in self._state.measured_liquid_heights:
+            self._state.measured_liquid_heights[labware_id] = {}
+        self._state.measured_liquid_heights[labware_id][well_name] = lhi
+
 
 class WellView(HasState[WellState]):
     """Read-only well state view."""
@@ -127,12 +173,3 @@ def has_measured_liquid_height(self, labware_id: str, well_name: str) -> bool:
             )
         except KeyError:
             return False
-
-    def set_liquid_height_from_load(
-        self, labware_id: str, well_name: str, height: float, time: datetime
-    ) -> None:
-        """Set the liquid height of the well from a LoadLiquid command."""
-        lhi = LiquidHeightInfo(height=height, last_measured=time)
-        if labware_id not in self._state.measured_liquid_heights:
-            self._state.measured_liquid_heights[labware_id] = {}
-        self._state.measured_liquid_heights[labware_id][well_name] = lhi