WIP

src/blueapi/core/context.py

@@ -109,12 +109,6 @@ def with_dodal_module(self, module: ModuleType, **kwargs) -> None:
         for device in devices.values():
             self.device(device)
 
-        directory_provider = get_directory_provider()
-        if directory_provider is not None:
-            self.run_engine.scan_id_source = (
-                lambda _: directory_provider.next_scan_number
-            )
-
         # If exceptions have occurred, we log them but we do not make blueapi
         # fall over
         if len(exceptions) > 0:
@@ -123,6 +117,14 @@ def with_dodal_module(self, module: ModuleType, **kwargs) -> None:
             )
             LOGGER.exception(NotConnected(exceptions))
 
+    def with_dodal_directory_provider(self) -> None:
+        directory_provider = get_directory_provider()
+        if directory_provider is not None:
+            directory_provider.scan_id_source = lambda _: self.run_engine.md["scan_id"]
+
+    def with_scan_id_source(self, source) -> None:
+        self.run_engine.scan_id_source = lambda _: source
+
     def plan(self, plan: PlanGenerator) -> PlanGenerator:
         """
         from dodal.beamlines.beamline_utils import get_directory_provider