feat(server): pre-flight disk check for the TCP run commands

software/control/core/disk_space.py

@@ -199,3 +199,49 @@ def _directory_usage(self) -> int:
         if not self._directory.is_dir():
             return 0
         return control.utils.get_directory_disk_usage(self._directory)
+
+
+PREFLIGHT_SAFETY_FACTOR = 1.03
+
+
+@dataclass(frozen=True)
+class PreflightDiskCheck:
+    """Whether the configured acquisition fits on the save disk, as checked before it starts."""
+
+    required_bytes: float
+    available_bytes: int
+    image_count: int
+    save_directory: str
+
+    @property
+    def fits(self) -> bool:
+        return self.required_bytes <= self.available_bytes
+
+    def describe(self) -> str:
+        megabytes_required = int(self.required_bytes / 1024 / 1024)
+        megabytes_available = int(self.available_bytes / 1024 / 1024)
+        return (
+            f"This acquisition will capture {self.image_count:,} images, which will"
+            f" require {megabytes_required:,} [MB], but '{self.save_directory}' only has"
+            f" {megabytes_available:,} [MB] available."
+        )
+
+
+def preflight_disk_check(
+    multi_point_controller, save_directory: str, safety_factor: float = PREFLIGHT_SAFETY_FACTOR
+) -> PreflightDiskCheck:
+    """Compare the fully configured controller's storage estimate with the free space where it will save.
+
+    Shared by the GUI's "Not Enough Disk Space" dialog and the TCP run commands so both decide alike.
+    ``save_directory`` may not exist yet (the TCP commands check before creating anything); free space is
+    then measured on its nearest existing parent, which is the disk it will be created on.
+    """
+    existing = pathlib.Path(save_directory)
+    while not existing.exists() and existing != existing.parent:
+        existing = existing.parent
+    return PreflightDiskCheck(
+        required_bytes=safety_factor * multi_point_controller.get_estimated_acquisition_disk_storage(),
+        available_bytes=control.utils.get_available_disk_space(existing),
+        image_count=multi_point_controller.get_acquisition_image_count(),
+        save_directory=str(save_directory),
+    )

software/control/microscope_control_server.py

@@ -21,6 +21,7 @@
 from control.core.acquisition_settings import apply_acquisition_settings, parse_wells
 
 import control._def  # Module import for runtime access to MCP-modifiable settings
+from control.core.disk_space import preflight_disk_check
 
 # Qt imports for thread-safe GUI operations
 try:
@@ -693,7 +694,11 @@ def _cmd_run_acquisition(
             le=99,
         ),
     ) -> Dict[str, Any]:
-        """Run a multi-point acquisition across wells using the MultiPointController."""
+        """Run a multi-point acquisition across wells using the MultiPointController.
+
+        An acquisition the save disk cannot hold is refused with an error before anything is created,
+        unless the Large Acquisitions setting is on.
+        """
         import os
 
         import control._def
@@ -781,6 +786,10 @@ def _cmd_run_acquisition(
 
             # Set the base path and start new experiment
             self.multipoint_controller.set_base_path(base_path)
+            self._refuse_if_acquisition_will_not_fit(
+                base_path,
+                opt_in_hint="enable Settings > Acquisition > Large Acquisitions (large_acquisition_mode = True in the INI)",
+            )
             self.multipoint_controller.start_new_experiment(experiment_id)
 
             # Calculate total FOVs for status reporting
@@ -1066,6 +1075,11 @@ def _cmd_run_acquisition_from_yaml(
         and region coordinates - wellplate or flexible), updates the GUI to reflect these
         settings, and starts the acquisition. Regions and every controller setting come from the
         file through control.core.acquisition_settings (shared with the fluidics protocol runner).
+
+        An acquisition the save disk cannot hold is refused with an error before anything is created,
+        unless large acquisition mode is on for it ('large_acquisition_mode: true' under 'acquisition:'
+        in the YAML, or the Large Acquisitions setting); it then pauses when space runs low and resumes
+        as finished data is offloaded.
         """
         import os
 
@@ -1150,6 +1164,13 @@ def _cmd_run_acquisition_from_yaml(
 
             # Set the base path and start new experiment
             self.multipoint_controller.set_base_path(base_path)
+            self._refuse_if_acquisition_will_not_fit(
+                base_path,
+                opt_in_hint=(
+                    "add 'large_acquisition_mode: true' under 'acquisition:' in the YAML, or enable "
+                    "Settings > Acquisition > Large Acquisitions"
+                ),
+            )
             self.multipoint_controller.start_new_experiment(experiment_id)
 
             # Calculate total FOVs for status reporting
@@ -1187,6 +1208,30 @@ def _cmd_run_acquisition_from_yaml(
             self._log.error(traceback.format_exc())
             raise RuntimeError(f"Failed to start acquisition: {str(e)}") from e
 
+    def _refuse_if_acquisition_will_not_fit(self, base_path: str, opt_in_hint: str) -> None:
+        """The GUI's "Not Enough Disk Space" check for TCP clients, who never see that dialog.
+
+        Call with the controller fully configured and before anything is created on disk. A run the save
+        disk cannot hold is refused unless large acquisition mode is on for it (then the disk-space guard
+        pauses it while data is offloaded). Runs that save nothing are not checked.
+        """
+        controller = self.multipoint_controller
+        if controller.skip_saving:
+            return
+        check = preflight_disk_check(controller, base_path)
+        if check.fits:
+            return
+        if controller.large_acquisition_mode or control._def.LARGE_ACQUISITION_MODE:
+            self._log.warning(
+                f"{check.describe()} Large acquisition mode is on: the run will pause when the disk runs low "
+                "and resume as data is offloaded."
+            )
+            return
+        raise RuntimeError(
+            f"{check.describe()} Free up disk space or reduce the acquisition. To run an acquisition larger "
+            f"than the disk (it pauses when space runs low while you offload finished data), {opt_in_hint}."
+        )
+
     @schema_method
     def _cmd_set_performance_mode(
         self,

software/control/widgets.py

@@ -21,6 +21,7 @@
 from control.core.multi_point_controller import MultiPointController
 from control.core.mosaic_utils import format_well_id
 from control.core.geometry_utils import get_effective_well_size, calculate_well_coverage
+from control.core.disk_space import preflight_disk_check
 from control.microcontroller import Microcontroller
 from control.piezo import PiezoStage
 from control.channel_sequence import enable_channel_sequence
@@ -126,9 +127,10 @@ def check_space_available_with_error_dialog(
     # To check how much disk space is required, we need to have the MultiPointController all configured.  That is
     # a precondition of this function.
     save_directory = multi_point_controller.base_path
-    available_disk_space = utils.get_available_disk_space(save_directory)
-    space_required = factor_of_safecty * multi_point_controller.get_estimated_acquisition_disk_storage()
-    image_count = multi_point_controller.get_acquisition_image_count()
+    check = preflight_disk_check(multi_point_controller, save_directory, safety_factor=factor_of_safecty)
+    available_disk_space = check.available_bytes
+    space_required = check.required_bytes
+    image_count = check.image_count
 
     logger.info(
         f"Checking space available: {space_required=}, {available_disk_space=}, {image_count=}, {save_directory=}"
@@ -137,12 +139,7 @@ def check_space_available_with_error_dialog(
     if not options.needs_dialog:
         return True
 
-    megabytes_required = int(space_required / 1024 / 1024)
-    megabytes_available = int(available_disk_space / 1024 / 1024)
-    error_message = (
-        f"This acquisition will capture {image_count:,} images, which will"
-        f" require {megabytes_required:,} [MB], but '{save_directory}' only has {megabytes_available:,} [MB] available."
-    )
+    error_message = check.describe()
     # Module-level lookup (not a default argument) so tests can monkeypatch the dialog.
     action = preflight_disk_dialog(f"{error_message}\n\n{LARGE_ACQUISITION_MODE_EXPLANATION}", options)
     if action == "enable":

software/docs/mcp_integration.md

@@ -152,6 +152,11 @@ The `python_exec` command is disabled by default for security. To enable it:
 | `get_acquisition_status` | - | Check acquisition progress |
 | `abort_acquisition` | - | Stop running acquisition |
 
+> Both run commands check the save disk first, like the GUI's "Not Enough Disk Space" dialog: an acquisition that
+> will not fit is refused with an error unless large acquisition mode is on for it (`large_acquisition_mode: true`
+> under `acquisition:` in the YAML, or Settings > Acquisition > Large Acquisitions). With the mode on, the run
+> pauses when space runs low and resumes as finished data is offloaded.
+>
 > For scripted automation, see [Automation](automation.md).
 
 ### Performance

software/tests/control/core/test_disk_space.py

@@ -6,6 +6,7 @@
 import pytest
 
 import control._def
+import control.utils
 from control.core.disk_space import DiskSpaceGuard, DiskStatus
 from control.core.pause_gate import PauseGate
 
@@ -373,3 +374,49 @@ def test_simulated_capacity_poll_interval_defaults_to_setting(tmp_path, monkeypa
     assert guard.free_bytes() == 100_000
     _write_file(tmp_path / "img.bin", 40_000)
     assert guard.free_bytes() == 60_000
+
+
+# --- pre-flight check shared by the GUI dialog and the TCP run commands ----------------------------
+
+
+class _EstimateOnlyController:
+    def __init__(self, estimate_bytes, image_count=12):
+        self._estimate = estimate_bytes
+        self._count = image_count
+
+    def get_estimated_acquisition_disk_storage(self):
+        return self._estimate
+
+    def get_acquisition_image_count(self):
+        return self._count
+
+
+def test_preflight_check_applies_the_safety_factor_and_reports_fit(tmp_path, monkeypatch):
+    from control.core.disk_space import preflight_disk_check
+
+    monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 1030)
+    fits = preflight_disk_check(_EstimateOnlyController(1000), str(tmp_path))
+    assert fits.fits and fits.required_bytes == 1030 and fits.available_bytes == 1030 and fits.image_count == 12
+
+    too_big = preflight_disk_check(_EstimateOnlyController(1001), str(tmp_path))
+    assert not too_big.fits
+
+
+def test_preflight_check_measures_the_nearest_existing_folder(tmp_path, monkeypatch):
+    from control.core.disk_space import preflight_disk_check
+
+    asked = []
+    monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: asked.append(str(d)) or 10**9)
+    # The TCP commands check before the experiment folder (or even the base path) has been created.
+    check = preflight_disk_check(_EstimateOnlyController(1), str(tmp_path / "not" / "created" / "yet"))
+    assert asked == [str(tmp_path)]
+    assert check.save_directory == str(tmp_path / "not" / "created" / "yet")
+
+
+def test_preflight_check_describes_the_shortfall(tmp_path, monkeypatch):
+    from control.core.disk_space import preflight_disk_check
+
+    monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 5 * 1024 * 1024)
+    check = preflight_disk_check(_EstimateOnlyController(2000 * 1024 * 1024, image_count=1234), str(tmp_path))
+    text = check.describe()
+    assert "1,234 images" in text and "2,060 [MB]" in text and "5 [MB] available" in text and str(tmp_path) in text

software/tests/control/test_microscope_control_server.py

@@ -35,6 +35,15 @@ def create_mock_server(objective: str = "20x", channels: list = None):
     mock_multipoint = MagicMock()
     mock_multipoint.acquisition_in_progress.return_value = False
     mock_multipoint.experiment_ID = "test_experiment"
+    # Pre-flight disk check inputs: a tiny acquisition that saves, with large acquisition mode off.
+    mock_multipoint.get_estimated_acquisition_disk_storage.return_value = 1000
+    mock_multipoint.get_acquisition_image_count.return_value = 12
+    mock_multipoint.skip_saving = False
+    mock_multipoint.large_acquisition_mode = False
+    mock_multipoint.set_large_acquisition_mode.side_effect = lambda on: setattr(
+        mock_multipoint, "large_acquisition_mode", on
+    )
+    mock_multipoint.set_skip_saving.side_effect = lambda on: setattr(mock_multipoint, "skip_saving", on)
     # What set_selected_configurations would have resolved for these channels
     mock_multipoint.selected_configurations = mock_channels
 
@@ -641,3 +650,76 @@ def test_schemas_document_the_commands(self, real_server):
             assert schemas[name]["parameters"] == {}
             assert schemas[name]["required"] == []
             assert len(schemas[name]["description"]) > 20
+
+
+class TestTcpDiskPreflight:
+    """Both TCP run commands refuse an acquisition the save disk cannot hold unless large acquisition mode is
+    on, mirroring the GUI's "Not Enough Disk Space" check (which a TCP client never sees)."""
+
+    @pytest.fixture
+    def yaml_file(self, tmp_path):
+        path = tmp_path / "acq.yaml"
+        path.write_text(SAMPLE_WELLPLATE_YAML)
+        return str(path)
+
+    @pytest.fixture
+    def opt_in_yaml_file(self, tmp_path):
+        path = tmp_path / "acq_large.yaml"
+        path.write_text(
+            SAMPLE_WELLPLATE_YAML.replace("acquisition:\n", "acquisition:\n  large_acquisition_mode: true\n", 1)
+        )
+        assert "large_acquisition_mode: true" in path.read_text()
+        return str(path)
+
+    @pytest.fixture
+    def tiny_disk(self, monkeypatch):
+        import control._def
+        import control.utils
+
+        monkeypatch.setattr(control._def, "LARGE_ACQUISITION_MODE", False)
+        monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 500)  # the mock run needs 1030
+
+    def test_yaml_run_that_does_not_fit_is_refused_before_anything_is_created(self, tiny_disk, yaml_file, tmp_path):
+        server = create_mock_server()
+        with pytest.raises(RuntimeError) as err:
+            server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path))
+        message = str(err.value)
+        assert "large_acquisition_mode: true" in message, "the error must name the opt-in"
+        assert "12 images" in message and str(tmp_path) in message
+        server.multipoint_controller.start_new_experiment.assert_not_called()
+        server.multipoint_controller.run_acquisition.assert_not_called()
+
+    def test_yaml_opt_in_lets_the_run_start(self, tiny_disk, opt_in_yaml_file, tmp_path):
+        server = create_mock_server()
+        result = server._cmd_run_acquisition_from_yaml(yaml_path=opt_in_yaml_file, base_path=str(tmp_path))
+        assert result["started"] is True
+        server.multipoint_controller.run_acquisition.assert_called_once()
+
+    def test_global_setting_lets_the_run_start(self, tiny_disk, yaml_file, tmp_path, monkeypatch):
+        import control._def
+
+        monkeypatch.setattr(control._def, "LARGE_ACQUISITION_MODE", True)
+        server = create_mock_server()
+        assert server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path))["started"] is True
+
+    def test_runs_that_save_nothing_are_not_checked(self, tiny_disk, tmp_path):
+        path = tmp_path / "acq_skip.yaml"
+        path.write_text(SAMPLE_WELLPLATE_YAML.replace("acquisition:\n", "acquisition:\n  skip_saving: true\n", 1))
+        server = create_mock_server()
+        assert server._cmd_run_acquisition_from_yaml(yaml_path=str(path), base_path=str(tmp_path))["started"] is True
+
+    def test_a_run_that_fits_starts_as_before(self, yaml_file, tmp_path, monkeypatch):
+        import control.utils
+
+        monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 10**9)
+        server = create_mock_server()
+        assert server._cmd_run_acquisition_from_yaml(yaml_path=yaml_file, base_path=str(tmp_path))["started"] is True
+
+    def test_the_wells_command_is_refused_too_and_names_the_setting(self, tiny_disk, tmp_path):
+        server = create_mock_server()
+        # This command validates channels against the live controller rather than the config repo.
+        server.microscope.live_controller.get_channels.return_value = [create_mock_channel("BF LED matrix full")]
+        with pytest.raises(RuntimeError) as err:
+            server._cmd_run_acquisition(wells="B6", channels=["BF LED matrix full"], base_path=str(tmp_path))
+        assert "Large Acquisitions" in str(err.value)
+        server.multipoint_controller.run_acquisition.assert_not_called()

software/tests/control/test_preflight_disk_check_controller.py

@@ -0,0 +1,33 @@
+"""The Qt-free pre-flight disk check (control.core.disk_space.preflight_disk_check) against a real simulated
+controller, as the TCP run commands use it: before the experiment folder exists."""
+
+
+def test_shared_check_works_on_a_real_controller_before_the_experiment_folder_exists(tmp_path, monkeypatch):
+    """The TCP run commands check before start_new_experiment(), so the estimate must not need the folder."""
+    import control.microscope
+    import control.utils
+    import tests.control.test_stubs as ts
+    from control.core.disk_space import PREFLIGHT_SAFETY_FACTOR, preflight_disk_check
+    from tests.control.test_MultiPointController import select_some_configs
+
+    scope = control.microscope.Microscope.build_from_global_config(True)
+    mpc = ts.get_test_multi_point_controller(microscope=scope)
+    cfg = mpc.stage.get_config()
+    mpc.scanCoordinates.add_single_fov_region(
+        "region_1",
+        center_x=cfg.X_AXIS.MIN_POSITION + 1.0,
+        center_y=cfg.Y_AXIS.MIN_POSITION + 1.0,
+        center_z=cfg.Z_AXIS.MIN_POSITION + 1.0,
+    )
+    select_some_configs(mpc, scope.objective_store.current_objective)
+    mpc.set_Nt(3)
+    save_dir = tmp_path / "base" / "not_created_yet"
+    mpc.set_base_path(str(save_dir))
+
+    monkeypatch.setattr(control.utils, "get_available_disk_space", lambda d: 1)
+    check = preflight_disk_check(mpc, str(save_dir))
+
+    assert check.image_count == mpc.get_acquisition_image_count() > 0
+    assert check.required_bytes == PREFLIGHT_SAFETY_FACTOR * mpc.get_estimated_acquisition_disk_storage() > 1
+    assert not check.fits
+    assert not save_dir.exists(), "checking must not create anything"