Fix Z-Offset Calibration (#206)

# Description

- [ ] Feature
- [x] Bug fix
- [x] Code refactor
- [ ] Documentation

Summary of changes to the probe calibration flow (`probeHelperPage.py`,
`babystepPage.py`, `printTab.py`, `controlTab.py`, `mainWindow.py`,
`moonrakerComm.py`).

- Replace 3 boolean calibration flags with a `_CalibPhase` enum with 6
states
- Extract `_cancel_calibration()`, `_reset_calibration_state()`,
`_restore_ui()`, `_show_option_cards()` helpers
- Fix double `CLEAN_NOZZLE` execution: `Z_ENDSTOP_CALIBRATE` already
calls it internally
- Fix `EDDY_PHASE1_RESTART` phase: prevents post-restart subscription
updates from overwriting the paper-test loading message
- Fix `on_extruder_update` / `on_gcode_move_update` guards: use
`_calib_phase` instead of `helper_initialize` (was cleared prematurely
by 300ms manual_probe query)
- Fix heating message stuck after emergency stop:
`notify_klippy_shutdown` now triggers `evaluate_klippy_status()` in
`moonrakerComm.py`; `on_klippy_status("shutdown")` calls
`_cancel_calibration()`
- Add `lock_ui` signal: disables all tabs except `controlTab` and header
during calibration; re-asserted when probe session becomes active
- Fix `-0.000` display in `probeHelperPage` and `babystepPage` using
`round(x, 3) or 0.0`
- Fix save z-offset button showing for `-0.000` in `printTab`
- `babystepPage`: move buttons start disabled, enabled only when
`print_stats.state == "printing"`
- Inline `_DISCONNECT_MSG`/`_READY_MSG` module-level dicts into
`match/case` in `on_klippy_status`

# Motivation

The calibration flow had accumulated several interacting boolean flags
that made state transitions brittle and hard to reason about. The
`_CalibPhase` enum makes the state machine explicit and eliminates
entire categories of bugs (premature flag clears, wrong phase handling
on reconnect, no reset on emergency stop). The UI lock prevents
navigating away mid-calibration. The `-0.000` fixes prevent confusing
display values at the zero crossing.

# Future work

- Add unit tests covering the `_CalibPhase` state machine transitions

Author: gmmcosta15

BlocksScreen/lib/moonrakerComm.py

@@ -280,9 +280,10 @@ def on_message(self, *args) -> None:
                         metadata=_entry,
                     )
         elif "method" in response:
-            if (
-                str(response["method"]).lower() == "notify_klippy_disconnected"
-            ):  # Checkout for notify_klippy_disconnect
+            if str(response["method"]).lower() in (
+                "notify_klippy_disconnected",
+                "notify_klippy_shutdown",
+            ):
                 self.evaluate_klippy_status()
 
             message_event = (

BlocksScreen/lib/panels/controlTab.py

@@ -40,6 +40,9 @@ class ControlTab(QtWidgets.QStackedWidget):
     disable_popups: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
         bool, name="disable-popups"
     )
+    lock_ui: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
+        bool, name="lock-ui"
+    )
     request_numpad: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
         [str, int, "PyQt_PyObject"],
         [str, int, "PyQt_PyObject", int, int],
@@ -82,6 +85,7 @@ def __init__(
         self.probe_helper_page = ProbeHelper(self)
         self.probe_helper_page.toggle_conn_page.connect(self.toggle_conn_page)
         self.probe_helper_page.disable_popups.connect(self.disable_popups)
+        self.probe_helper_page.lock_ui.connect(self.lock_ui)
         self.addWidget(self.probe_helper_page)
         self.probe_helper_page.call_load_panel.connect(self.call_load_panel)
         self.printcores_page = SwapPrintcorePage(self)

BlocksScreen/lib/panels/mainWindow.py

@@ -34,6 +34,12 @@
 
 _logger = logging.getLogger(__name__)
 
+_GCODE_POPUP_MESSAGES: tuple[tuple[str, str], ...] = (
+    ("filament runout", "Filament Runout"),
+    ("no filament", "No Filament Detected"),
+    ("sensor not in valid range", "Eddy Current Sensor:\nnot in valid range"),
+)
+
 
 def api_handler(func):
     """Decorator for methods that handle api responses"""
@@ -112,6 +118,7 @@ def __init__(self):
         self.ui.setupUi(self)
         self.screensaver = ScreenSaver(self)
         self._popup_toggle: bool = False
+        self._klippy_ready: bool = False
         self.ui.main_content_widget.setCurrentIndex(0)
 
         usb_config = self.config.get_section("usb_manager", fallback=None)
@@ -145,6 +152,7 @@ def __init__(self):
             self.conn_window.on_websocket_connection_achieved
         )
         self.ws.connection_lost.connect(self.conn_window.on_websocket_connection_lost)
+        self.ws.klippy_state_signal.connect(self._on_klippy_state)
         self.printer.webhooks_update.connect(self.conn_window.webhook_update)
         self.printPanel.request_back.connect(slot=self.global_back)
         self.printPanel.on_cancel_print.connect(slot=self.on_cancel_print)
@@ -214,7 +222,11 @@ def __init__(self):
         self.handle_error_response.connect(
             self.controlPanel.probe_helper_page.handle_error_response
         )
+        self.controlPanel.probe_helper_page.show_notifications.connect(
+            self.notiPage.new_notication
+        )
         self.controlPanel.disable_popups.connect(self.popup_toggle)
+        self.controlPanel.lock_ui.connect(self.set_ui_lock)
         self.on_update_message.connect(self.update_page.handle_update_message)
         self.update_page.request_full_update.connect(self.ws.api.full_update)
         self.update_page.request_recover_repo[str].connect(
@@ -432,6 +444,24 @@ def popup_toggle(self, toggle: bool) -> None:
         """Toggles app popups"""
         self._popup_toggle = toggle
 
+    @QtCore.pyqtSlot(bool, name="set-ui-lock")
+    def set_ui_lock(self, locked: bool) -> None:
+        """Lock or unlock navigation during calibration.
+
+        Disables all tabs except controlTab (where calibration lives) and
+        the header, so the user cannot navigate away mid-calibration.
+        """
+        for tab in (self.ui.printTab, self.ui.filamentTab, self.ui.utilitiesTab):
+            self.ui.main_content_widget.setTabEnabled(
+                self.ui.main_content_widget.indexOf(tab), not locked
+            )
+        self.ui.header_main_layout.setEnabled(not locked)
+
+    @QtCore.pyqtSlot(str, name="on-klippy-state")
+    def _on_klippy_state(self, state: str) -> None:
+        """Track Klippy readiness to suppress spurious error popups during disconnect."""
+        self._klippy_ready = state == "ready"
+
     def reset_tab_indexes(self):
         """
         Used to grantee all tabs reset to their
@@ -709,10 +739,18 @@ def _handle_notify_gcode_response_message(self, method, data, metadata) -> None:
             if self._popup_toggle:
                 return
             _gcode_msg_type, _message = str(_gcode_response[0]).split(" ", maxsplit=1)
-            popupWhitelist = ["filament runout", "no filament"]
-            if _message.lower() not in popupWhitelist or _gcode_msg_type != "!!":
+            _msg_lower = _message.lower()
+            _display = next(
+                (
+                    fmt
+                    for pattern, fmt in _GCODE_POPUP_MESSAGES
+                    if pattern in _msg_lower
+                ),
+                None,
+            )
+            if _gcode_msg_type != "!!" or _display is None:
                 return
-            self.show_notifications.emit("mainwindow", _message, 3, True)
+            self.show_notifications.emit("mainwindow", _display, 3, True)
 
     @api_handler
     def _handle_error_message(self, method, data, metadata) -> None:
@@ -721,6 +759,11 @@ def _handle_error_message(self, method, data, metadata) -> None:
         if self._popup_toggle:
             return
 
+        # Suppress error popups while Klippy is disconnected/shutting down.
+        # Those errors are side-effects of the disconnect, not actionable by the user.
+        if not self._klippy_ready:
+            return
+
         text = data.get("message", str(data)) if isinstance(data, dict) else str(data)
         lower_text = text.lower()
 
@@ -874,4 +917,4 @@ def event(self, event: QtCore.QEvent) -> bool:
     def sizeHint(self) -> QtCore.QSize:
         """Sets default size for the widget"""
         self.adjustSize()
-        return super().sizeHint(QtCore.QSize(800, 480))
+        return QtCore.QSize(800, 480)

BlocksScreen/lib/panels/printTab.py

@@ -62,9 +62,12 @@ class PrintTab(QtWidgets.QStackedWidget):
     on_cancel_print: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
         name="on_cancel_print"
     )
-    call_load_panel = QtCore.pyqtSignal(bool, str, name="call-load-panel")
-
-    call_cancel_panel = QtCore.pyqtSignal(bool, name="call-load-panel")
+    call_load_panel: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
+        bool, str, name="call-load-panel"
+    )
+    call_cancel_panel: typing.ClassVar[QtCore.pyqtSignal] = QtCore.pyqtSignal(
+        bool, name="call-load-panel"
+    )
 
     def __init__(
         self,
@@ -75,7 +78,9 @@ def __init__(
     ) -> None:
         super().__init__(parent)
         self._active_z_offset: float = 0.0
+        self._pending_save_offset: float = 0.0
         self._finish_print_handled: bool = False
+        self._cancel_z_snapshot: float = 0.0
         self._z_apply_command: str = "Z_OFFSET_APPLY_ENDSTOP"
 
         self.setupMainPrintPage()
@@ -224,6 +229,9 @@ def __init__(
         self.printer.gcode_move_update[str, list].connect(
             self.babystepPage.on_gcode_move_update
         )
+        self.printer.print_stats_update[str, str].connect(
+            self.babystepPage.on_print_state_update
+        )
         self.printer.gcode_move_update[str, list].connect(self.activate_save_button)
         self.tune_page.run_gcode.connect(self.ws.api.run_gcode)
         self.tune_page.request_sliderPage[str, int, "PyQt_PyObject"].connect(
@@ -268,6 +276,7 @@ def __init__(
         self.confirmPage_widget.on_delete.connect(self.delete_file)
         self.change_page(self.indexOf(self.print_page))  # force set the initial page
         self.save_config_btn.clicked.connect(self.save_config)
+        self.ws.klippy_state_signal.connect(self.on_klippy_state)
 
     @QtCore.pyqtSlot(str, dict, name="on_print_stats_update")
     @QtCore.pyqtSlot(str, float, name="on_print_stats_update")
@@ -276,10 +285,14 @@ def on_print_stats_update(self, field: str, value: dict | float | str) -> None:
         """
         unblocks tabs if on standby
         """
-        if isinstance(value, str):
-            if "state" in field:
-                if value in ("standby"):
-                    self.on_cancel_print.emit()
+        if isinstance(value, str) and "state" in field and value == "standby":
+            self.on_cancel_print.emit()
+            if not self._finish_print_handled and self._cancel_z_snapshot != 0:
+                self._active_z_offset = self._cancel_z_snapshot
+                self.save_config()
+                self._finish_print_handled = True
+                self.save_config_btn.setVisible(True)
+            self._cancel_z_snapshot = 0.0
 
     @QtCore.pyqtSlot(str, int, "PyQt_PyObject", name="on_numpad_request")
     @QtCore.pyqtSlot(str, int, "PyQt_PyObject", int, int, name="on_numpad_request")
@@ -292,6 +305,10 @@ def on_numpad_request(
         max_value: int = 100,
     ) -> None:
         """Handle numpad request"""
+        try:
+            self.numpadPage.value_selected.disconnect()
+        except (RuntimeError, TypeError):
+            pass
         self.numpadPage.value_selected.connect(callback)
         self.numpadPage.set_name(name)
         self.numpadPage.set_value(current_value)
@@ -311,6 +328,10 @@ def on_slidePage_request(
         max_value: int = 100,
     ) -> None:
         """Handle slider page request"""
+        try:
+            self.sliderPage.value_selected.disconnect()
+        except (RuntimeError, TypeError):
+            pass
         self.sliderPage.value_selected.connect(callback)
         self.sliderPage.set_name(name)
         self.sliderPage.set_slider_position(int(current_value))
@@ -334,12 +355,9 @@ def delete_file(self, filename: str, directory: str = "gcodes") -> None:
 
     def save_config(self) -> None:
         """Handle Save configuration behaviour, shows confirmation dialog"""
-
-        self.babystepPage.bbp_z_offset_title_label.setText(
-            f"Z: {self._active_z_offset:.3f}mm"
-        )
+        self._pending_save_offset = self._active_z_offset
         self.BasePopup_z_offset.set_message(
-            f"The Z-Offset is now {self._active_z_offset:.3f} mm.\n"
+            f"The Z-Offset is now {self._pending_save_offset + 0.0:.3f} mm.\n"
             "Would you like to save this change permanently?\n"
             "The machine will restart."
         )
@@ -352,28 +370,37 @@ def save_config(self) -> None:
         self.BasePopup_z_offset.open()
 
     def update_configuration_file(self) -> None:
-        """Restore the captured offset, apply it to the probe config, then save."""
+        """Runs the `SAVE_CONFIG` gcode"""
         try:
             self.BasePopup_z_offset.accepted.disconnect(self.update_configuration_file)
         except (RuntimeError, TypeError):
             pass
         self.run_gcode_signal.emit(
-            f"SET_GCODE_OFFSET Z={self._active_z_offset:.3f} MOVE=0"
+            f"SET_GCODE_OFFSET Z={self._pending_save_offset:.3f} MOVE=0"
         )
         self.run_gcode_signal.emit(self._z_apply_command)
         self.run_gcode_signal.emit("SAVE_CONFIG")
+        self.babystepPage.bbp_z_offset_title_label.setText(
+            f"Z: {self._pending_save_offset + 0.0:.3f}mm"
+        )
         self.save_config_btn.setVisible(False)
 
+    @QtCore.pyqtSlot(str, name="on_klippy_state")
+    def on_klippy_state(self, state: str) -> None:
+        """Dismiss the Z-offset save popup and reset save state on unexpected shutdown."""
+        if state in ("ready", "startup"):
+            return
+        self.BasePopup_z_offset.reject()
+        self.save_config_btn.setVisible(False)
+        self.babystepPage.baby_stepchange = False
+
     @QtCore.pyqtSlot(str, list, name="activate_save_button")
     def activate_save_button(self, name: str, value: list) -> None:
         """Sync the `Save config` popup with the save_config_pending state"""
-        if not value:
+        if not value or name != "homing_origin" or len(value) <= 2:
             return
-
-        if name == "homing_origin":
-            if len(value) > 2:
-                self._active_z_offset = value[2]
-                self.save_config_btn.setVisible(value[2] != 0)
+        self._active_z_offset = value[2]
+        self.save_config_btn.setVisible(round(value[2], 3) != 0)
 
     def _on_delete_file_confirmed(self, filename: str, directory: str) -> None:
         """Handle confirmed file deletion after user accepted the dialog."""
@@ -401,6 +428,12 @@ def setProperty(self, name: str, value: typing.Any) -> bool:
 
     def handle_cancel_print(self) -> None:
         """Handles the print cancel action"""
+        if (
+            not self._finish_print_handled
+            and self._active_z_offset != 0
+            and self.babystepPage.baby_stepchange
+        ):
+            self._cancel_z_snapshot = self._active_z_offset
         self.ws.api.cancel_print()
         self.call_load_panel.emit(True, "Cancelling print...\nPlease wait")
 
@@ -422,6 +455,7 @@ def klipper_ready_signal(self) -> None:
         """React to klipper ready signal"""
         self.babystepPage.baby_stepchange = False
         self._finish_print_handled = False
+        self._cancel_z_snapshot = 0.0
         self.printer.on_subscribe_config("stepper_z", self._on_stepper_z_config)
 
     def _on_stepper_z_config(self, config: dict | list) -> None:
@@ -442,6 +476,7 @@ def finish_print_signal(self) -> None:
         if self._active_z_offset != 0 and self.babystepPage.baby_stepchange:
             self.save_config()
             self._finish_print_handled = True
+        self.save_config_btn.setVisible(round(self._active_z_offset, 3) != 0)
 
     def setupMainPrintPage(self) -> None:
         """Setup UI for print page"""