From aae2bfab3b783e34f7ae717d1e189b3de94c6720 Mon Sep 17 00:00:00 2001
From: Nick Diehl <47604184+ncdiehl11@users.noreply.github.com>
Date: Thu, 6 Feb 2025 10:56:08 -0500
Subject: [PATCH] fix(protocol-designer, step-generation): update prewet
 behavior in PD (#17433)

This PR adds prewet to distribute, and updates the logic for when to use
prewet for all compound liquid handling command creators. The prewet
volume should always be set to the volume of the aspiration that is
about to be called (including disposal volume for `distribute`). Prewet
commands should be emmitted anytime a new tip is used for an aspiration.
Also, I update the UI for the prewet form field in MoveLiquidTools

Closes AUTH-908
---
 .../src/assets/localization/en/tooltip.json   |   2 +-
 .../molecules/ToggleStepFormField/index.tsx   |  36 ++--
 .../StepForm/PipetteFields/LabwareField.tsx   |   1 +
 .../StepForm/StepFormToolbox.tsx              |   1 +
 .../StepTools/MoveLiquidTools/index.tsx       |  12 +-
 .../src/__tests__/distribute.test.ts          | 198 ++++++++++++++++++
 .../commandCreators/compound/distribute.ts    |  27 +++
 .../src/commandCreators/compound/transfer.ts  |   2 +-
 .../src/fixtures/commandFixtures.ts           |  23 +-
 9 files changed, 277 insertions(+), 25 deletions(-)

diff --git a/protocol-designer/src/assets/localization/en/tooltip.json b/protocol-designer/src/assets/localization/en/tooltip.json
index 849844a0c34..1445cedf0bd 100644
--- a/protocol-designer/src/assets/localization/en/tooltip.json
+++ b/protocol-designer/src/assets/localization/en/tooltip.json
@@ -65,7 +65,7 @@
       "newLocation": "New location to move the selected labware",
       "nozzles": "Partial pickup requires a tip rack directly on the deck. Full rack pickup requires the Flex 96 Tip Rack Adapter.",
       "pipette": "Select the pipette you want to use",
-      "preWetTip": "Pre-wet by aspirating and dispensing 2/3 of the tip’s max volume",
+      "preWetTip": "Pre-wet by aspirating and dispensing the total aspiration volume",
       "setTemperature": "Select the temperature to set your module to",
       "wells": "Select wells",
       "volume": "Volume to dispense in each well"
diff --git a/protocol-designer/src/molecules/ToggleStepFormField/index.tsx b/protocol-designer/src/molecules/ToggleStepFormField/index.tsx
index ad609acf0b9..3068f4e0e69 100644
--- a/protocol-designer/src/molecules/ToggleStepFormField/index.tsx
+++ b/protocol-designer/src/molecules/ToggleStepFormField/index.tsx
@@ -1,5 +1,6 @@
 import {
   ALIGN_CENTER,
+  Check,
   COLORS,
   DIRECTION_COLUMN,
   Flex,
@@ -16,12 +17,13 @@ import { ToggleButton } from '../../atoms/ToggleButton'
 interface ToggleStepFormFieldProps {
   title: string
   isSelected: boolean
-  onLabel: string
-  offLabel: string
   toggleUpdateValue: (value: unknown) => void
   toggleValue: unknown
   tooltipContent: string | null
-  isDisabled: boolean
+  isDisabled?: boolean
+  onLabel?: string
+  offLabel?: string
+  toggleElement?: 'toggle' | 'checkbox'
 }
 export function ToggleStepFormField(
   props: ToggleStepFormFieldProps
@@ -34,7 +36,8 @@ export function ToggleStepFormField(
     toggleUpdateValue,
     toggleValue,
     tooltipContent,
-    isDisabled,
+    isDisabled = false,
+    toggleElement = 'toggle',
   } = props
   const [targetProps, tooltipProps] = useHoverTooltip()
 
@@ -42,7 +45,7 @@ export function ToggleStepFormField(
     <>
       <ListButton
         type="noActive"
-        padding={SPACING.spacing12}
+        padding={SPACING.spacing16}
         onClick={() => {
           if (!isDisabled) {
             toggleUpdateValue(!toggleValue)
@@ -50,11 +53,7 @@ export function ToggleStepFormField(
         }}
         disabled={isDisabled}
       >
-        {tooltipContent != null ? (
-          <Tooltip tooltipProps={tooltipProps}>{tooltipContent}</Tooltip>
-        ) : null}
-
-        <Flex width="100%" flexDirection={DIRECTION_COLUMN} {...targetProps}>
+        <Flex width="100%" flexDirection={DIRECTION_COLUMN}>
           <Flex
             justifyContent={JUSTIFY_SPACE_BETWEEN}
             alignItems={ALIGN_CENTER}
@@ -62,6 +61,7 @@ export function ToggleStepFormField(
             <StyledText
               desktopStyle="bodyDefaultRegular"
               color={isDisabled ? COLORS.grey40 : COLORS.black90}
+              {...targetProps}
             >
               {title}
             </StyledText>
@@ -72,15 +72,21 @@ export function ToggleStepFormField(
               >
                 {isSelected ? onLabel : offLabel}
               </StyledText>
-              <ToggleButton
-                disabled={isDisabled}
-                label={isSelected ? onLabel : offLabel}
-                toggledOn={isSelected}
-              />
+              {toggleElement === 'toggle' ? (
+                <ToggleButton
+                  label={isSelected ? onLabel : offLabel}
+                  toggledOn={isSelected}
+                />
+              ) : (
+                <Check color={COLORS.blue50} isChecked={isSelected} />
+              )}
             </Flex>
           </Flex>
         </Flex>
       </ListButton>
+      {tooltipContent != null ? (
+        <Tooltip tooltipProps={tooltipProps}>{tooltipContent}</Tooltip>
+      ) : null}
     </>
   )
 }
diff --git a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/PipetteFields/LabwareField.tsx b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/PipetteFields/LabwareField.tsx
index 07e08b8a299..00f68dcd46c 100644
--- a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/PipetteFields/LabwareField.tsx
+++ b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/PipetteFields/LabwareField.tsx
@@ -31,6 +31,7 @@ export function LabwareField(props: FieldProps): JSX.Element {
       onExit={() => {
         dispatch(hoverSelection({ id: null, text: null }))
       }}
+      width="100%"
     />
   )
 }
diff --git a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepFormToolbox.tsx b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepFormToolbox.tsx
index 253889423ca..45eee9219f8 100644
--- a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepFormToolbox.tsx
+++ b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepFormToolbox.tsx
@@ -360,6 +360,7 @@ export function StepFormToolbox(props: StepFormToolboxProps): JSX.Element {
             </StyledText>
           </Flex>
         }
+        width="21.875rem"
       >
         <div
           ref={toolsComponentRef}
diff --git a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepTools/MoveLiquidTools/index.tsx b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepTools/MoveLiquidTools/index.tsx
index 652baebd843..da137eccf9c 100644
--- a/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepTools/MoveLiquidTools/index.tsx
+++ b/protocol-designer/src/pages/Designer/ProtocolSteps/StepForm/StepTools/MoveLiquidTools/index.tsx
@@ -27,6 +27,7 @@ import {
 import {
   CheckboxExpandStepFormField,
   InputStepFormField,
+  ToggleStepFormField,
 } from '../../../../../../molecules'
 import {
   BlowoutLocationField,
@@ -398,15 +399,16 @@ export function MoveLiquidTools(props: StepFormProps): JSX.Element {
           {t('protocol_steps:advanced_settings')}
         </StyledText>
         {tab === 'aspirate' ? (
-          <CheckboxExpandStepFormField
+          <ToggleStepFormField
             title={i18n.format(
               t('form:step_edit_form.field.preWetTip.label'),
               'capitalize'
             )}
-            checkboxValue={propsForFields.preWetTip.value}
-            isChecked={propsForFields.preWetTip.value === true}
-            checkboxUpdateValue={propsForFields.preWetTip.updateValue}
-            tooltipText={propsForFields.preWetTip.tooltipContent}
+            toggleValue={propsForFields.preWetTip.value}
+            isSelected={propsForFields.preWetTip.value === true}
+            toggleUpdateValue={propsForFields.preWetTip.updateValue}
+            toggleElement="checkbox"
+            tooltipContent={propsForFields.preWetTip.tooltipContent ?? null}
           />
         ) : null}
         <CheckboxExpandStepFormField
diff --git a/step-generation/src/__tests__/distribute.test.ts b/step-generation/src/__tests__/distribute.test.ts
index 2264f4b44b3..638cc95e7a2 100644
--- a/step-generation/src/__tests__/distribute.test.ts
+++ b/step-generation/src/__tests__/distribute.test.ts
@@ -896,6 +896,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
       const result = distribute(args, invariantContext, robotStateWithTip)
       const res = getSuccessResult(result)
       expect(res.commands).toEqual([
+        // prewet
+        aspirateHelper('A1', 150),
+        delayCommand(11),
+        dispenseHelper('A1', 150, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         aspirateHelper('A1', 35),
         delayCommand(11),
@@ -957,6 +972,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
       const res = getSuccessResult(result)
       expect(res.commands).toEqual([
         // no need to pickup tip/drop tip since change tip is never
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1040,6 +1070,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         // replace tip since change tip is always
         ...dropTipHelper(),
         pickUpTipHelper('A1'),
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1079,6 +1124,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         ...dropTipHelper(),
         // next chunk from A1: remaining volume
         pickUpTipHelper('B1'),
+        // prewet
+        aspirateHelper('A1', 160),
+        delayCommand(11),
+        dispenseHelper('A1', 160, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate 100 liquid + 60 for disposal vol
@@ -1130,6 +1190,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         // replace tip at the beginning of the step
         ...dropTipHelper(),
         pickUpTipHelper('A1'),
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1210,6 +1285,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
       const res = getSuccessResult(result)
       expect(res.commands).toEqual([
         // no need to replace tip since change tip is never
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1295,6 +1385,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         // replace tip
         ...dropTipHelper(),
         pickUpTipHelper('A1'),
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1336,6 +1441,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         ...dropTipHelper(),
         // next chunk from A1: remaining volume
         pickUpTipHelper('B1'),
+        // prewet
+        aspirateHelper('A1', 160),
+        delayCommand(11),
+        dispenseHelper('A1', 160, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate 100 liquid + 60 for disposal vol
@@ -1387,6 +1507,22 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         // replace tip
         ...dropTipHelper(),
         pickUpTipHelper('A1'),
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
+
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1469,6 +1605,22 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
       const res = getSuccessResult(result)
       expect(res.commands).toEqual([
         // no need to replace tip since changeTip is never
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
+
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1554,6 +1706,21 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         // replace tip
         ...dropTipHelper(),
         pickUpTipHelper('A1'),
+        // prewet
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
@@ -1595,6 +1762,20 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         ...dropTipHelper(),
         // next chunk from A1: remaining volume
         pickUpTipHelper('B1'),
+        aspirateHelper('A1', 160),
+        delayCommand(11),
+        dispenseHelper('A1', 160, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate 100 liquid + 60 for disposal vol
@@ -1639,13 +1820,30 @@ describe('advanced settings: volume, mix, pre-wet tip, tip touch, tip position',
         changeTip: 'once',
         blowoutLocation: DEST_WELL_BLOWOUT_DESTINATION,
       } as DistributeArgs
+      console.log(args.aspirateDelay, args.dispenseDelay)
 
       const result = distribute(args, invariantContext, robotStateWithTip)
       const res = getSuccessResult(result)
       expect(res.commands).toEqual([
         // replace tip
         ...dropTipHelper(),
+
+        // prewet
         pickUpTipHelper('A1'),
+        aspirateHelper('A1', 260),
+        delayCommand(11),
+        dispenseHelper('A1', 260, {
+          labwareId: SOURCE_LABWARE,
+          wellLocation: {
+            origin: 'bottom',
+            offset: {
+              x: 0,
+              y: 0,
+              z: ASPIRATE_OFFSET_FROM_BOTTOM_MM,
+            },
+          },
+        }),
+        delayCommand(12),
         // mix (asp)
         ...mixCommandsWithDelay,
         // aspirate
diff --git a/step-generation/src/commandCreators/compound/distribute.ts b/step-generation/src/commandCreators/compound/distribute.ts
index 111626c79ce..8269250b79d 100644
--- a/step-generation/src/commandCreators/compound/distribute.ts
+++ b/step-generation/src/commandCreators/compound/distribute.ts
@@ -496,8 +496,35 @@ export const distribute: CommandCreator<DistributeArgs> = (
           ]
         : []
       const aspirateDisposalVolumeOnce = chunkIndex === 0 ? disposalVolume : 0
+      // if changeTip is 'once' or 'never', only prewet for first aspiration
+      // if changeTip is 'always', pre-wet for each chunk
+      const preWetTipCommands =
+        args.preWetTip &&
+        (args.changeTip === 'always' ? true : chunkIndex === 0)
+          ? mixUtil({
+              pipette: args.pipette,
+              labware: args.sourceLabware,
+              well: args.sourceWell,
+              volume:
+                args.volume * destWellChunk.length +
+                (args.changeTip === 'always'
+                  ? disposalVolume
+                  : aspirateDisposalVolumeOnce),
+              times: 1,
+              offsetFromBottomMm: aspirateOffsetFromBottomMm,
+              aspirateFlowRateUlSec,
+              dispenseFlowRateUlSec,
+              aspirateDelaySeconds: aspirateDelay?.seconds,
+              dispenseDelaySeconds: dispenseDelay?.seconds,
+              tipRack: args.tipRack,
+              xOffset: aspirateXOffset,
+              yOffset: aspirateYOffset,
+              nozzles,
+            })
+          : []
       return [
         ...tipCommands,
+        ...preWetTipCommands,
         ...configureForVolumeCommand,
         ...mixBeforeAspirateCommands,
         curryCommandCreator(aspirate, {
diff --git a/step-generation/src/commandCreators/compound/transfer.ts b/step-generation/src/commandCreators/compound/transfer.ts
index c4b3d183ea2..282b92ad3e1 100644
--- a/step-generation/src/commandCreators/compound/transfer.ts
+++ b/step-generation/src/commandCreators/compound/transfer.ts
@@ -286,7 +286,7 @@ export const transfer: CommandCreator<TransferArgs> = (
                   pipette: args.pipette,
                   labware: args.sourceLabware,
                   well: sourceWell,
-                  volume: Math.max(subTransferVol),
+                  volume: subTransferVol,
                   times: 1,
                   offsetFromBottomMm: aspirateOffsetFromBottomMm,
                   aspirateFlowRateUlSec,
diff --git a/step-generation/src/fixtures/commandFixtures.ts b/step-generation/src/fixtures/commandFixtures.ts
index 37e4ed07a44..056ad5114cd 100644
--- a/step-generation/src/fixtures/commandFixtures.ts
+++ b/step-generation/src/fixtures/commandFixtures.ts
@@ -138,7 +138,12 @@ export const makeAspirateHelper: MakeAspDispHelper<AspDispAirgapParams> = bakedP
     ...params,
   },
 })
-export const makeMoveToWellHelper = (wellName: string, labwareId?: string) => ({
+export const makeMoveToWellHelper = (
+  wellName: string,
+  labwareId?: string,
+  forceDirect?: boolean,
+  minimumZHeight?: number
+) => ({
   commandType: 'moveToWell',
   key: expect.any(String),
   params: {
@@ -153,6 +158,8 @@ export const makeMoveToWellHelper = (wellName: string, labwareId?: string) => ({
         z: 11.54,
       },
     },
+    forceDirect,
+    minimumZHeight,
   },
 })
 export const makeAirGapHelper = (volume: number) => ({
@@ -268,18 +275,25 @@ export const makeTouchTipHelper: MakeTouchTipHelper = bakedParams => (
   key: expect.any(String),
   params: { ..._defaultTouchTipParams, ...bakedParams, wellName, ...params },
 })
-export const delayCommand = (seconds: number): CreateCommand => ({
+export const delayCommand = (
+  seconds: number,
+  message?: string
+): CreateCommand => ({
   commandType: 'waitForDuration',
   key: expect.any(String),
   params: {
     seconds: seconds,
+    message,
   },
 })
 export const delayWithOffset = (
   wellName: string,
   labwareId: string,
   seconds?: number,
-  zOffset?: number
+  zOffset?: number,
+  forceDirect?: boolean,
+  minimumZHeight?: number,
+  message?: string
 ): CreateCommand[] => [
   {
     commandType: 'moveToWell',
@@ -296,6 +310,8 @@ export const delayWithOffset = (
           z: zOffset || 14,
         },
       },
+      forceDirect,
+      minimumZHeight,
     },
   },
   {
@@ -303,6 +319,7 @@ export const delayWithOffset = (
     key: expect.any(String),
     params: {
       seconds: seconds ?? 12,
+      message,
     },
   },
 ]