mx-bluesky 1.4.6__py3-none-any.whl → 1.4.7__py3-none-any.whl

mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -24,6 +24,9 @@ from dodal.devices.zocalo.zocalo_results import (
     XrcResult,
     get_full_processing_results,
 )
+from dodal.plans.preprocessors.verify_undulator_gap import (
+    verify_undulator_gap_before_run_decorator,
+)
 
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     ispyb_activation_wrapper,
@@ -34,6 +37,9 @@ from mx_bluesky.common.plans.read_hardware import (
     standard_read_hardware_during_collection,
     standard_read_hardware_pre_collection,
 )
+from mx_bluesky.common.preprocessors.preprocessors import (
+    transmission_and_xbpm_feedback_for_collection_decorator,
+)
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.exceptions import (
     CrystalNotFoundException,
@@ -52,9 +58,6 @@ from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
     setup_zebra_for_panda_flyscan,
     tidy_up_zebra_after_gridscan,
 )
-from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
-)
 from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
     change_aperture_then_move_to_xtal,
 )
@@ -102,13 +105,6 @@ def flyscan_xray_centre_no_move(
         }
     )
     @bpp.finalize_decorator(lambda: feature_controlled.tidy_plan(composite))
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.undulator,
-        composite.xbpm_feedback,
-        composite.attenuator,
-        composite.dcm,
-        parameters.transmission_frac,
-    )
     def run_gridscan_and_fetch_and_tidy(
         fgs_composite: HyperionFlyScanXRayCentreComposite,
         params: HyperionSpecifiedThreeDGridScan,
@@ -140,6 +136,11 @@ def flyscan_xray_centre(
     """
     xrc_event_handler = XRayCentreEventHandler()
 
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite,
+        parameters.transmission_frac,
+    )
+    @verify_undulator_gap_before_run_decorator(composite)
     @bpp.subs_decorator(xrc_event_handler)
     def flyscan_and_fetch_results() -> MsgGenerator:
         yield from ispyb_activation_wrapper(

mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py

@@ -9,6 +9,7 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.oav.oav_parameters import OAVParameters
 
 import mx_bluesky.common.xrc_result as flyscan_result
+from mx_bluesky.common.parameters.components import WithSnapshot
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.xrc_result import XRayCentreEventHandler
@@ -54,7 +55,10 @@ def load_centre_collect_full(
     @run_decorator(
         md={
             "metadata": {"sample_id": parameters.sample_id},
-            "activate_callbacks": ["SampleHandlingCallback"],
+            "activate_callbacks": ["BeamDrawingCallback", "SampleHandlingCallback"],
+            "with_snapshot": parameters.multi_rotation_scan.model_dump_json(
+                include=WithSnapshot.model_fields.keys()  # type: ignore
+            ),
         }
     )
     def plan_with_callback_subs():

mx_bluesky/hyperion/experiment_plans/oav_snapshot_plan.py

@@ -67,7 +67,7 @@ def _take_oav_snapshot(composite: OavSnapshotComposite, omega: float):
         composite.smargon.omega, omega, group=OAV_SNAPSHOT_SETUP_SHOT
     )
     time_now = datetime.now()
-    filename = f"{time_now.strftime('%H%M%S')}_oav_snapshot_{omega:.0f}"
+    filename = f"{time_now.strftime('%H%M%S%f')[:8]}_oav_snapshot_{omega:.0f}"
     yield from bps.abs_set(
         composite.oav.snapshot.filename,
         filename,
@@ -76,5 +76,5 @@ def _take_oav_snapshot(composite: OavSnapshotComposite, omega: float):
     yield from bps.wait(group=OAV_SNAPSHOT_SETUP_SHOT)
     yield from bps.trigger(composite.oav.snapshot, wait=True)
     yield from bps.create(DocDescriptorNames.OAV_ROTATION_SNAPSHOT_TRIGGERED)
-    yield from bps.read(composite.oav.snapshot)
+    yield from bps.read(composite.oav)
     yield from bps.save()

mx_bluesky/hyperion/experiment_plans/pin_tip_centring_plan.py

@@ -9,6 +9,7 @@ from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.oav_parameters import OAV_CONFIG_JSON, OAVParameters
 from dodal.devices.oav.pin_image_recognition import PinTipDetection, Tip
 from dodal.devices.oav.utils import (
+    PinNotFoundException,
     Pixel,
     get_move_required_so_that_beam_is_at_pixel,
     wait_for_tip_to_be_found,
@@ -16,7 +17,7 @@ from dodal.devices.oav.utils import (
 from dodal.devices.smargon import Smargon
 
 from mx_bluesky.common.utils.context import device_composite_from_context
-from mx_bluesky.common.utils.exceptions import SampleException
+from mx_bluesky.common.utils.exceptions import SampleException, catch_exception_and_warn
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.device_setup_plans.setup_oav import pre_centring_setup_oav
 from mx_bluesky.hyperion.device_setup_plans.smargon import (
@@ -159,6 +160,7 @@ def pin_tip_centre_plan(
     # need to wait for the OAV image to update
     # See #673 for improvements
     yield from bps.sleep(0.3)
-
-    tip = yield from wait_for_tip_to_be_found(pin_tip_detect)
+    tip = yield from catch_exception_and_warn(
+        PinNotFoundException, wait_for_tip_to_be_found, pin_tip_detect
+    )
     yield from offset_and_move(tip)

mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from math import isclose
 from typing import cast
 
-import bluesky.preprocessors as bpp
 import pydantic
 from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
@@ -37,14 +36,10 @@ from dodal.log import LOGGER
 from ophyd_async.fastcs.panda import HDFPanda
 
 from mx_bluesky.common.parameters.constants import OavConstants
-from mx_bluesky.common.xrc_result import XRayCentreEventHandler
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     fill_in_energy_if_not_supplied,
     start_preparing_data_collection_then_do_plan,
 )
-from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
-    change_aperture_then_move_to_xtal,
-)
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     GridDetectThenXRayCentreComposite,
 )
@@ -133,27 +128,6 @@ def _robot_load_then_flyscan_plan(
     yield from _flyscan_plan_from_robot_load_params(composite, params, oav_config_file)
 
 
-def robot_load_then_centre(
-    composite: RobotLoadThenCentreComposite,
-    parameters: RobotLoadThenCentre,
-) -> MsgGenerator:
-    """Perform pin-tip detection followed by a flyscan to determine centres of interest.
-    Performs a robot load if necessary. Centre on the best diffracting centre.
-    """
-
-    xray_centre_event_handler = XRayCentreEventHandler()
-
-    yield from bpp.subs_wrapper(
-        robot_load_then_xray_centre(composite, parameters), xray_centre_event_handler
-    )
-    flyscan_results = xray_centre_event_handler.xray_centre_results
-    if flyscan_results is not None:
-        yield from change_aperture_then_move_to_xtal(
-            flyscan_results[0], composite.smargon, composite.aperture_scatterguard
-        )
-    # else no chi change, no need to recentre.
-
-
 def robot_load_then_xray_centre(
     composite: RobotLoadThenCentreComposite,
     parameters: RobotLoadThenCentre,

mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -26,12 +26,19 @@ from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra.zebra import RotationDirection, Zebra
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
+from dodal.plans.preprocessors.verify_undulator_gap import (
+    verify_undulator_gap_before_run_decorator,
+)
 
+from mx_bluesky.common.parameters.components import WithSnapshot
 from mx_bluesky.common.plans.read_hardware import (
     read_hardware_for_zocalo,
     standard_read_hardware_during_collection,
     standard_read_hardware_pre_collection,
 )
+from mx_bluesky.common.preprocessors.preprocessors import (
+    transmission_and_xbpm_feedback_for_collection_decorator,
+)
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import (
@@ -48,9 +55,6 @@ from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
 )
-from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
-)
 from mx_bluesky.hyperion.experiment_plans.oav_snapshot_plan import (
     OavSnapshotComposite,
     oav_snapshot_plan,
@@ -369,24 +373,26 @@ def rotation_scan(
     if not oav_params:
         oav_params = OAVParameters(context="xrayCentring")
 
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite,
+        parameters.transmission_frac,
+    )
+    @verify_undulator_gap_before_run_decorator(composite)
     @bpp.set_run_key_decorator("rotation_scan")
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": CONST.PLAN.ROTATION_OUTER,
             "mx_bluesky_parameters": parameters.model_dump_json(),
+            "with_snapshot": parameters.model_dump_json(
+                include=WithSnapshot.model_fields.keys()  # type: ignore
+            ),
             "activate_callbacks": [
+                "BeamDrawingCallback",
                 "RotationISPyBCallback",
                 "RotationNexusFileCallback",
             ],
         }
     )
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.undulator,
-        composite.xbpm_feedback,
-        composite.attenuator,
-        composite.dcm,
-        parameters.transmission_frac,
-    )
     def rotation_scan_plan_with_stage_and_cleanup(
         params: RotationScan,
     ):
@@ -422,6 +428,10 @@ def multi_rotation_scan(
     eiger: EigerDetector = composite.eiger
     eiger.set_detector_parameters(parameters.detector_params)
 
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite,
+        parameters.transmission_frac,
+    )
     @bpp.set_run_key_decorator("multi_rotation_scan")
     @bpp.run_decorator(
         md={
@@ -434,17 +444,11 @@ def multi_rotation_scan(
             ],
         }
     )
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.undulator,
-        composite.xbpm_feedback,
-        composite.attenuator,
-        composite.dcm,
-        parameters.transmission_frac,
-    )
     @bpp.finalize_decorator(lambda: _cleanup_plan(composite))
     def _multi_rotation_scan():
         for single_scan in parameters.single_rotation_scans:
 
+            @verify_undulator_gap_before_run_decorator(composite)
             @bpp.set_run_key_decorator("rotation_scan")
             @bpp.run_decorator(  # attach experiment metadata to the start document
                 md={
@@ -459,6 +463,8 @@ def multi_rotation_scan(
 
             yield from rotation_scan_core(single_scan)
 
+        yield from bps.unstage(eiger)
+
     LOGGER.info("setting up and staging eiger...")
     yield from start_preparing_data_collection_then_do_plan(
         composite.beamstop,
@@ -468,4 +474,3 @@ def multi_rotation_scan(
         _multi_rotation_scan(),
         group=CONST.WAIT.ROTATION_READY_FOR_DC,
     )
-    yield from bps.unstage(eiger)

mx_bluesky/hyperion/experiment_plans/set_energy_plan.py

@@ -5,18 +5,21 @@
 * reenable feedback
 """
 
+import bluesky.preprocessors as bpp
 import pydantic
 from bluesky import plan_stubs as bps
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.dcm import DCM
 from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, MirrorVoltages
+from dodal.devices.undulator import Undulator
 from dodal.devices.undulator_dcm import UndulatorDCM
 from dodal.devices.xbpm_feedback import XBPMFeedback
 
-from mx_bluesky.hyperion.device_setup_plans import dcm_pitch_roll_mirror_adjuster
-from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
+from mx_bluesky.common.parameters.constants import PlanNameConstants
+from mx_bluesky.common.preprocessors.preprocessors import (
     transmission_and_xbpm_feedback_for_collection_wrapper,
 )
+from mx_bluesky.hyperion.device_setup_plans import dcm_pitch_roll_mirror_adjuster
 
 DESIRED_TRANSMISSION_FRACTION = 0.1
 
@@ -33,6 +36,17 @@ class SetEnergyComposite:
     attenuator: BinaryFilterAttenuator
 
 
+# Remove composite after https://github.com/DiamondLightSource/dodal/issues/1092
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class XBPMWrapperComposite:
+    undulator: Undulator
+    xbpm_feedback: XBPMFeedback
+    attenuator: BinaryFilterAttenuator
+    dcm: DCM
+
+
+@bpp.set_run_key_decorator(PlanNameConstants.SET_ENERGY)
+@bpp.run_decorator()
 def _set_energy_plan(
     energy_kev,
     composite: SetEnergyComposite,
@@ -51,12 +65,17 @@ def set_energy_plan(
     energy_ev: float | None,
     composite: SetEnergyComposite,
 ):
+    # Remove conversion after https://github.com/DiamondLightSource/dodal/issues/1092
+    composite_for_wrapper = XBPMWrapperComposite(
+        composite.undulator_dcm.undulator_ref._obj,  # noqa: SLF001
+        composite.xbpm_feedback,
+        composite.attenuator,
+        composite.dcm,
+    )
+
     if energy_ev:
         yield from transmission_and_xbpm_feedback_for_collection_wrapper(
             _set_energy_plan(energy_ev / 1000, composite),
-            composite.undulator_dcm.undulator_ref(),
-            composite.xbpm_feedback,
-            composite.attenuator,
-            composite.dcm,
+            composite_for_wrapper,
             DESIRED_TRANSMISSION_FRACTION,
         )

mx_bluesky/hyperion/external_interaction/agamemnon.py

@@ -1,21 +1,52 @@
 import dataclasses
 import json
 import re
-from typing import TypeVar
+from os import path
+from typing import Any, TypeVar
 
 import requests
+from deepdiff.diff import DeepDiff
 from dodal.utils import get_beamline_name
-
-from mx_bluesky.common.parameters.components import WithVisit
-from mx_bluesky.common.parameters.constants import GridscanParamConstants
+from jsonschema import ValidationError
+from pydantic_extra_types.semantic_version import SemanticVersion
+
+from mx_bluesky.common.parameters.components import (
+    PARAMETER_VERSION,
+    MxBlueskyParameters,
+    TopNByMaxCountSelection,
+    WithCentreSelection,
+    WithOptionalEnergyChange,
+    WithSample,
+    WithVisit,
+)
+from mx_bluesky.common.parameters.constants import (
+    GridscanParamConstants,
+)
 from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.common.utils.utils import convert_angstrom_to_eV
+from mx_bluesky.hyperion.parameters.components import WithHyperionUDCFeatures
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
+from mx_bluesky.hyperion.parameters.robot_load import RobotLoadThenCentre
 
 T = TypeVar("T", bound=WithVisit)
 AGAMEMNON_URL = "http://agamemnon.diamond.ac.uk/"
 MULTIPIN_PREFIX = "multipin"
 MULTIPIN_FORMAT_DESC = "Expected multipin format is multipin_{number_of_wells}x{well_size}+{distance_between_tip_and_first_well}"
 MULTIPIN_REGEX = rf"^{MULTIPIN_PREFIX}_(\d+)x(\d+(?:\.\d+)?)\+(\d+(?:\.\d+)?)$"
+MX_GENERAL_ROOT_REGEX = r"^/dls/(?P<beamline>[^/]+)/data/[^/]*/(?P<visit>[^/]+)(?:/|$)"
+
+
+class AgamemnonLoadCentreCollect(
+    MxBlueskyParameters,
+    WithVisit,
+    WithSample,
+    WithCentreSelection,
+    WithHyperionUDCFeatures,
+    WithOptionalEnergyChange,
+):
+    """Experiment parameters to compare against GDA populated LoadCentreCollect."""
+
+    robot_load_then_centre: RobotLoadThenCentre
 
 
 @dataclasses.dataclass
@@ -60,7 +91,7 @@ def _get_parameters_from_url(url: str) -> dict:
         raise KeyError(f"Unexpected json from agamemnon: {response_json}") from e
 
 
-def _get_pin_type_from_agamemnon_parameters(parameters: dict) -> PinType:
+def get_pin_type_from_agamemnon_parameters(parameters: dict) -> PinType:
     loop_type_name: str | None = parameters["sample"]["loopType"]
     if loop_type_name:
         regex_search = re.search(MULTIPIN_REGEX, loop_type_name)
@@ -81,15 +112,119 @@ def get_next_instruction(beamline: str) -> dict:
     return _get_parameters_from_url(AGAMEMNON_URL + f"getnextcollect/{beamline}")
 
 
-def get_pin_type_from_agamemnon(beamline: str) -> PinType:
-    params = get_next_instruction(beamline)
-    return _get_pin_type_from_agamemnon_parameters(params)
+def get_withvisit_parameters_from_agamemnon(parameters: dict) -> tuple:
+    try:
+        prefix = parameters["prefix"]
+        collection = parameters["collection"]
+        # Assuming distance is identical for multiple collections. Remove after https://github.com/DiamondLightSource/mx-bluesky/issues/773
+        detector_distance = collection[0]["distance"]
+    except KeyError as e:
+        raise KeyError("Unexpected json from agamemnon") from e
+
+    match = re.match(MX_GENERAL_ROOT_REGEX, prefix) if prefix else None
+
+    if match:
+        return (match.group("visit"), detector_distance)
+
+    raise ValueError(
+        f"Agamemnon prefix '{prefix}' does not match MX-General root structure"
+    )
+
+
+def get_withsample_parameters_from_agamemnon(parameters: dict) -> dict[str, Any]:
+    assert parameters.get("sample"), "instruction does not have a sample"
+    return {
+        "sample_id": parameters["sample"]["id"],
+        "sample_puck": parameters["sample"]["container"],
+        "sample_pin": parameters["sample"]["position"],
+    }
+
+
+def get_withenergy_parameters_from_agamemnon(parameters: dict) -> dict[str, Any]:
+    try:
+        first_collection: dict = parameters["collection"][0]
+        wavelength = first_collection.get("wavelength")
+        assert isinstance(wavelength, float)
+        demand_energy_ev = convert_angstrom_to_eV(wavelength)
+        return {"demand_energy_ev": demand_energy_ev}
+    except (KeyError, IndexError, AttributeError, TypeError):
+        return {"demand_energy_ev": None}
+
+
+def get_param_version() -> SemanticVersion:
+    return SemanticVersion.validate_from_str(str(PARAMETER_VERSION))
+
+
+def create_robot_load_then_centre_params_from_agamemnon(
+    parameters: dict,
+) -> RobotLoadThenCentre:
+    visit, detector_distance = get_withvisit_parameters_from_agamemnon(parameters)
+    with_sample_params = get_withsample_parameters_from_agamemnon(parameters)
+    with_energy_params = get_withenergy_parameters_from_agamemnon(parameters)
+    visit_directory, file_name = path.split(parameters["prefix"])
+    return RobotLoadThenCentre(
+        parameter_model_version=get_param_version(),
+        storage_directory=visit_directory + "/xraycentring",
+        visit=visit,
+        detector_distance_mm=detector_distance,
+        snapshot_directory=visit_directory + "/snapshots",
+        file_name=file_name,
+        **with_energy_params,
+        **with_sample_params,
+    )
+
+
+def populate_parameters_from_agamemnon(agamemnon_params):
+    visit, detector_distance = get_withvisit_parameters_from_agamemnon(agamemnon_params)
+    with_sample_params = get_withsample_parameters_from_agamemnon(agamemnon_params)
+    with_energy_params = get_withenergy_parameters_from_agamemnon(agamemnon_params)
+    pin_type = get_pin_type_from_agamemnon_parameters(agamemnon_params)
+    robot_load_params = create_robot_load_then_centre_params_from_agamemnon(
+        agamemnon_params
+    )
+    return AgamemnonLoadCentreCollect(
+        parameter_model_version=SemanticVersion.validate_from_str(
+            str(PARAMETER_VERSION)
+        ),
+        visit=visit,
+        detector_distance_mm=detector_distance,
+        select_centres=TopNByMaxCountSelection(n=pin_type.expected_number_of_crystals),
+        robot_load_then_centre=robot_load_params,
+        **with_sample_params,
+        **with_energy_params,
+    )
+
+
+def create_parameters_from_agamemnon() -> AgamemnonLoadCentreCollect:
+    beamline_name = get_beamline_name("i03")
+    agamemnon_params = get_next_instruction(beamline_name)
+
+    return populate_parameters_from_agamemnon(agamemnon_params)
+
+
+def compare_params(load_centre_collect_params):
+    try:
+        parameters = create_parameters_from_agamemnon()
+
+        # Log differences against GDA populated parameters
+        differences = DeepDiff(
+            parameters, load_centre_collect_params, math_epsilon=1e-5
+        )
+        if differences:
+            LOGGER.info(
+                f"Different parameters found when directly reading from Hyperion: {differences}"
+            )
+    except (ValueError, KeyError) as e:
+        LOGGER.warning(f"Failed to compare parameters: {e}")
+    except Exception as e:
+        LOGGER.warning(f"Unexpected error occurred. Failed to compare parameters: {e}")
 
 
 def update_params_from_agamemnon(parameters: T) -> T:
     try:
         beamline_name = get_beamline_name("i03")
-        pin_type = get_pin_type_from_agamemnon(beamline_name)
+        agamemnon_params = get_next_instruction(beamline_name)
+        pin_type = get_pin_type_from_agamemnon_parameters(agamemnon_params)
         if isinstance(parameters, LoadCentreCollect):
             parameters.robot_load_then_centre.tip_offset_um = pin_type.full_width / 2
             parameters.robot_load_then_centre.grid_width_um = pin_type.full_width
@@ -99,6 +234,9 @@ def update_params_from_agamemnon(parameters: T) -> T:
                 # Before we do https://github.com/DiamondLightSource/mx-bluesky/issues/226
                 # this will give no snapshots but that's preferable
                 parameters.multi_rotation_scan.snapshot_omegas_deg = []
+    except (ValueError, ValidationError) as e:
+        LOGGER.warning(f"Failed to update parameters: {e}")
     except Exception as e:
-        LOGGER.warning(f"Failed to get pin type from agamemnon, using single pin {e}")
+        LOGGER.warning(f"Unexpected error occurred. Failed to update parameters: {e}")
+
     return parameters

mx_bluesky/hyperion/external_interaction/callbacks/__main__.py

@@ -1,8 +1,8 @@
 import logging
 from collections.abc import Callable, Sequence
 from threading import Thread
-from time import sleep
 
+import bluesky.plan_stubs as bps
 from bluesky.callbacks import CallbackBase
 from bluesky.callbacks.zmq import Proxy, RemoteDispatcher
 from dodal.log import LOGGER as dodal_logger
@@ -14,6 +14,9 @@ from mx_bluesky.common.external_interaction.callbacks.common.log_uid_tag_callbac
 from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
     ZocaloCallback,
 )
+from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import (
+    SampleHandlingCallback,
+)
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     GridscanISPyBCallback,
 )
@@ -35,8 +38,8 @@ from mx_bluesky.hyperion.external_interaction.callbacks.rotation.ispyb_callback
 from mx_bluesky.hyperion.external_interaction.callbacks.rotation.nexus_callback import (
     RotationNexusFileCallback,
 )
-from mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback import (
-    SampleHandlingCallback,
+from mx_bluesky.hyperion.external_interaction.callbacks.snapshot_callback import (
+    BeamDrawingCallback,
 )
 from mx_bluesky.hyperion.parameters.cli import parse_callback_dev_mode_arg
 from mx_bluesky.hyperion.parameters.constants import CONST
@@ -67,15 +70,18 @@ def create_rotation_callbacks() -> tuple[
     return (
         RotationNexusFileCallback(),
         RotationISPyBCallback(
-            emit=ZocaloCallback(CONST.PLAN.ROTATION_MAIN, CONST.ZOCALO_ENV)
+            emit=ZocaloCallback(CONST.PLAN.ROTATION_MULTI, CONST.ZOCALO_ENV)
         ),
     )
 
 
 def setup_callbacks() -> list[CallbackBase]:
+    rot_nexus_cb, rot_ispyb_cb = create_rotation_callbacks()
+    snapshot_cb = BeamDrawingCallback(emit=rot_ispyb_cb)
     return [
         *create_gridscan_callbacks(),
-        *create_rotation_callbacks(),
+        rot_nexus_cb,
+        snapshot_cb,
         LogUidTaggingCallback(),
         RobotLoadISPyBCallback(),
         SampleHandlingCallback(),
@@ -134,7 +140,7 @@ def wait_for_threads_forever(threads: Sequence[Thread]):
     try:
         log_debug("Trying to wait forever on callback and dispatcher threads")
         while all(alive):
-            sleep(LIVENESS_POLL_SECONDS)
+            yield from bps.sleep(LIVENESS_POLL_SECONDS)
             alive = [t.is_alive() for t in threads]
     except KeyboardInterrupt:
         log_info("Main thread received interrupt - exiting.")