manual_probe: Add extra checks to catch manual probe not finishing correctly

Signed-off-by: Kevin O'Connor <kevin@koconnor.net>

1 files changed, 13 insertions, 8 deletions

diff --git a/klippy/extras/manual_probe.py b/klippy/extras/manual_probe.py
index 2fabd4f6..38160fef 100644
--- a/klippy/extras/manual_probe.py
+++ b/klippy/extras/manual_probe.py
@@ -41,6 +41,16 @@ class ManualProbe:
     def cmd_Z_ENDSTOP_CALIBRATE(self, params):
         ManualProbeHelper(self.printer, params, self.z_endstop_finalize)
 
+# Verify that a manual probe isn't already in progress
+def verify_no_manual_probe(printer):
+    gcode = printer.lookup_object('gcode')
+    try:
+        gcode.register_command('ACCEPT', 'dummy')
+    except printer.config_error as e:
+        raise gcode.error(
+            "Already in a manual Z probe. Use ABORT to abort it.")
+    gcode.register_command('ACCEPT', None)
+
 Z_BOB_MINIMUM = 0.500
 BISECT_MAX = 0.200
 
@@ -55,14 +65,9 @@ class ManualProbeHelper:
         self.past_positions = []
         self.last_toolhead_pos = self.last_kinematics_pos = None
         # Register commands
-        try:
-            self.gcode.register_command('ACCEPT', self.cmd_ACCEPT,
-                                        desc=self.cmd_ACCEPT_help)
-        except self.printer.config_error as e:
-            self.gcode.respond_error(
-                "Already in a manual Z probe. Use ABORT to abort it.")
-            self.finalize_callback(None)
-            return
+        verify_no_manual_probe(printer)
+        self.gcode.register_command('ACCEPT', self.cmd_ACCEPT,
+                                    desc=self.cmd_ACCEPT_help)
         self.gcode.register_command('NEXT', self.cmd_ACCEPT)
         self.gcode.register_command('ABORT', self.cmd_ABORT,
                                     desc=self.cmd_ABORT_help)