Z-A coupling: auto-coordinate A on jogs and MDI

Match the file-preprocessor behaviour for live operator input. When a
Z-down jog or MDI line would push (A-Z) above the safe band, append
the matching A delta to the same line so the planner runs Z and A
together. Same direction-aware refusal: only error when the operator
explicitly asks A to move *up* (delta > 0) past the bound, or when
the required A would violate A's soft minimum.

Implementation:
  * ExternalAxis.coordinate_mdi rewrites a multi-line MDI burst,
    tracking G90/G91 modal across lines (jogs always emit
    M70/G91/G0/M72; standard MDI defaults to G90). Z and A targets
    are computed in machine coords using offset_z and offset_a so
    the work-coord A token we emit is consistent with the operator's
    frame.
  * The 'A0' the jog UI emits for axes that aren't moving is treated
    as 'no A intent' (G91 delta of zero) and freely overridden.
  * Hooked into Mach.mdi after the existing ATC rewrite. On
    ExternalAxisError the burst is dropped with a user message; the
    planner check downstream still fires as defense in depth.
  * Planner.__encode also catches ExternalAxisError now (vs
    bricking on uncaught) - logs to the operator messages list and
    halts the cycle cleanly so subsequent jogs work.
  * check_coupling itself is now improvement-aware: only refuses
    moves that worsen an existing violation. Pure XY jogs and
    Z-up/A-down recovery moves pass even when (A-Z) is currently
    above the bound.

Tested locally with synthetic MDI: small Z jog within band, Z jog
across the boundary (auto-injects A delta), G90 MDI G0 Z-50
(appends A106), explicit A-lift while Z deep (refuses), pure XY
jog (unchanged), G91 A-down (unchanged), G90 G0 A0 with
offset_a=134 (refuses as lift to home).

src/py/bbctrl/ExternalAxis.py

@@ -38,6 +38,7 @@
 #
 ################################################################################
 
+import re
 import threading
 
 try:
@@ -258,6 +259,212 @@ class ExternalAxis(object):
         except Exception:
             return None
 
+    # ---- MDI / jog auto-coordination -----------------------------------
+    #
+    # The file preprocessor injects pre-position A moves at upload time,
+    # but jogs and ad-hoc MDI lines bypass that path. Without help, a
+    # jog past the safe band would be refused and the operator has to
+    # manually drop A before retrying. That's no fun and no longer
+    # matches the documented behaviour, so we apply the same rule here:
+    # if a Z-down move would worsen the gap, attach an A delta on the
+    # *same* line so the planner runs Z and A together.
+    #
+    # Jogs always come in as a multi-line MDI sequence:
+    #     M70           ; push modal state
+    #     G91           ; relative
+    #     G0 X.. Y.. Z.. A..
+    #     M72           ; pop modal state
+    # so we have to be careful: G90/G91 modal must be tracked across
+    # lines (the G91 set up by the jog applies to the G0 below it),
+    # and we must not re-set absolute mode for free-form MDI.
+    #
+    # Coordinate frame: gcode is in *work coords*. To compare against
+    # the machine-coord constraint we use offset_z and offset_a. We
+    # work in machine coords throughout; emitted A token is in work
+    # coords (= A_machine - offset_a) so the resulting gcode is
+    # consistent with whatever modal frame the operator is in.
+
+    _MDI_AXIS_RES = {
+        'x': re.compile(r'(?<![A-Za-z_0-9])[Xx]\s*([-+]?\d*\.?\d+)'),
+        'y': re.compile(r'(?<![A-Za-z_0-9])[Yy]\s*([-+]?\d*\.?\d+)'),
+        'z': re.compile(r'(?<![A-Za-z_0-9])[Zz]\s*([-+]?\d*\.?\d+)'),
+        'a': re.compile(r'(?<![A-Za-z_0-9])[Aa]\s*([-+]?\d*\.?\d+)'),
+    }
+    _MDI_G_CODE_RE = re.compile(
+        r'(?<![A-Za-z_0-9])[Gg]\s*0*(\d+(?:\.\d+)?)')
+    _MDI_PAREN_RE = re.compile(r'\([^)]*\)')
+
+    def _state_machine_z(self):
+        try:
+            zp = self.ctrl.state.get('zp', None)
+            if zp is None: return None
+            return float(zp)
+        except Exception:
+            return None
+
+    def _state_offset(self, axis):
+        try:
+            return float(self.ctrl.state.get('offset_' + axis, 0.0) or 0.0)
+        except Exception:
+            return 0.0
+
+    def _soft_min_machine(self):
+        """A's soft minimum in machine coords (the deepest A may go).
+        Falls back to None when soft limits are disabled."""
+        lo, _hi = self._soft_limits()
+        return lo
+
+    def coordinate_mdi(self, cmd):
+        """Rewrite an MDI/jog command sequence to keep A coupled with
+        Z. Returns the (possibly modified) command string.
+
+        Behaviour:
+          * If coupling is disabled or A isn't homed, returns cmd
+            unchanged.
+          * Tracks G90/G91 modal across lines; the modal state at
+            entry is taken from gplan-side global modal (we assume
+            G90 if not seen on a previous line - jogs always set
+            G91 as their first line, MDI typed by the operator is
+            usually G90 unless they say otherwise).
+          * For each motion line that includes Z, computes the
+            target Z in machine coords and the resulting (A-Z)
+            gap. If the gap would exceed K, augments / overrides
+            the line's A token so the resulting A drops Z + K.
+          * Refuses (raises ExternalAxisError) only when the required
+            A would violate A's soft minimum.
+        """
+        if not self.couple_z_enabled:
+            return cmd
+        try:
+            homed = bool(self.aux._homed)
+        except Exception:
+            return cmd
+        if not homed:
+            return cmd
+        K = self.couple_K
+        if K is None:
+            return cmd
+
+        # Live state. Use machine-coord values throughout.
+        a_mach = self._a_machine_now()
+        z_mach = self._state_machine_z()
+        if a_mach is None or z_mach is None:
+            return cmd
+        a_off = self._state_offset(self.axis_letter)
+        z_off = self._state_offset('z')
+        a_soft_min = self._soft_min_machine()
+
+        # Modal across the lines we rewrite. We default to absolute
+        # because jogs explicitly set G91 first, and standard MDI is
+        # G90. If the user's prior session ended in G91 we'll be
+        # wrong on the very first line of free-form MDI; that's a
+        # known shortcoming - file uploads use the file preprocessor
+        # which has the same default.
+        absolute = True
+
+        out_lines = []
+        for raw in cmd.splitlines():
+            line = raw.rstrip('\n')
+            code = self._MDI_PAREN_RE.sub('', line)
+            code = code.split(';', 1)[0]
+            if not code.strip():
+                out_lines.append(line)
+                continue
+
+            # Track modal G90/G91.
+            for m in self._MDI_G_CODE_RE.finditer(code):
+                try: g = int(float(m.group(1)))
+                except Exception: continue
+                if g == 90: absolute = True
+                elif g == 91: absolute = False
+
+            # Find Z and A tokens (last one wins).
+            z_tok = None
+            a_tok = None
+            for m in self._MDI_AXIS_RES['z'].finditer(code):
+                try: z_tok = float(m.group(1))
+                except Exception: pass
+            for m in self._MDI_AXIS_RES['a'].finditer(code):
+                try: a_tok = float(m.group(1))
+                except Exception: pass
+
+            if z_tok is None and a_tok is None:
+                out_lines.append(line)
+                continue
+
+            # Compute machine-coord targets for this line.
+            if z_tok is not None:
+                z_target_mach = (z_off + z_tok) if absolute \
+                                else (z_mach + z_tok)
+            else:
+                z_target_mach = z_mach
+            if a_tok is not None:
+                a_target_mach = (a_off + a_tok) if absolute \
+                                else (a_mach + a_tok)
+            else:
+                a_target_mach = a_mach
+
+            eps = 1e-4
+            gap_after  = a_target_mach - z_target_mach
+            gap_before = a_mach - z_mach
+
+            if gap_after > K + eps and gap_after > gap_before - eps:
+                # Need to drop A so the gap is exactly K (the deepest
+                # we may take A is whatever Z's target requires).
+                a_required_mach = z_target_mach + K
+                # Refuse when A would have to go below its soft
+                # minimum to satisfy the rule. The operator should
+                # raise Z first.
+                if a_soft_min is not None and a_required_mach < a_soft_min - eps:
+                    raise ExternalAxisError(
+                        'Z-A coupling: jog would require A=%.3f mm '
+                        '(machine), below soft minimum %.3f mm. '
+                        'Raise Z first.' % (
+                            a_required_mach, a_soft_min))
+                # Refuse only when the operator explicitly asked A
+                # to *move upward* (delta > 0) and the unsafe target
+                # is above the bound. A0 from the jog UI (G91 delta
+                # of zero) is treated as 'no A intent' and is freely
+                # overridden; an A token that already drops A below
+                # what we need is also overridden to the required
+                # depth (deeper-than-needed is fine in the other
+                # direction handled by the no-rewrite branch).
+                if a_tok is not None:
+                    delta_a = a_target_mach - a_mach
+                    if delta_a > eps and a_target_mach > a_required_mach + eps:
+                        raise ExternalAxisError(
+                            'Z-A coupling: line requests A=%.3f mm '
+                            'while Z=%.3f mm needs A<=%.3f mm.' % (
+                                a_target_mach, z_target_mach,
+                                a_required_mach))
+                # Convert required machine A back to a token in the
+                # operator's frame.
+                if absolute:
+                    a_token_value = a_required_mach - a_off
+                else:
+                    a_token_value = a_required_mach - a_mach
+                new_token = 'A%.4f' % a_token_value
+                if a_tok is None:
+                    # Append A to the line.
+                    line = line.rstrip() + ' ' + new_token
+                else:
+                    # Replace existing A token.
+                    line = self._MDI_AXIS_RES['a'].sub(
+                        new_token, line, count=1)
+                # Update the running mirror so subsequent lines in
+                # this same MDI burst compute correctly.
+                a_mach = a_required_mach
+                if z_tok is not None:
+                    z_mach = z_target_mach
+            else:
+                # Move was already safe; just update mirrors for
+                # subsequent lines.
+                if z_tok is not None: z_mach = z_target_mach
+                if a_tok is not None: a_mach = a_target_mach
+
+            out_lines.append(line)
+        return '\n'.join(out_lines)
+
     def coupling_for_preprocessor(self):
         """Return the dict the AuxPreprocessor wants for in-file
         injection, or None when coupling is off. We assume the
@@ -278,17 +485,19 @@ class ExternalAxis(object):
         """Validate that a proposed motion respects the Z-A coupling.
 
         Each argument is a target *machine* mm position; pass None to
-        keep the current value of that axis. Raises ExternalAxisError
+        keep the current value of that axis.
-        if the resulting (A,Z) pair would violate the constraint.
 
-        Skipped when:
+        Improvement-aware: a move is rejected only when it *worsens*
-          * coupling is disabled in aux.json,
+        an already-violating state (or moves a healthy state into
-          * the auxiliary axis isn't homed (we haven't established a
+        violation). Pure XY jogs that touch neither Z nor A are not
-            reliable A position yet - the soft-limit gate uses the
+        passed through here; jogs that hold Z or A at their current
-            same convention),
+        value (gplan emits the unchanged value in `target`) pass
-          * Z's current position is unknown (no zp reported yet -
+        because (a-z) doesn't change. Z-up moves while in violation
-            avoid spurious failures during boot before the AVR has
+        also pass because they reduce (a-z) toward the bound.
-            sent its first status frame).
+
+        Raises ExternalAxisError on violation. Skipped when coupling
+        is disabled, the aux axis isn't homed, or current positions
+        aren't yet known.
         """
         if not self.couple_z_enabled:
             return
@@ -303,17 +512,26 @@ class ExternalAxis(object):
             return
         a_now = self._a_machine_now()
         z_now = self._z_machine_now()
-        a = float(target_a_machine) if target_a_machine is not None else a_now
+        if a_now is None or z_now is None:
-        z = float(target_z_machine) if target_z_machine is not None else z_now
-        if a is None or z is None:
             return
+        a_after = (float(target_a_machine)
+                   if target_a_machine is not None else a_now)
+        z_after = (float(target_z_machine)
+                   if target_z_machine is not None else z_now)
         eps = 1e-4
-        if a - z > K + eps:
+        gap_after  = a_after - z_after
+        gap_before = a_now   - z_now
+        # Only refuse when (a) the resulting state would violate the
+        # constraint AND (b) the move makes things at least as bad
+        # as the current state. This lets the operator escape an
+        # already-violating state by moving in the right direction
+        # (Z up, A down).
+        if gap_after > K + eps and gap_after > gap_before - eps:
             raise ExternalAxisError(
                 'Z-A coupling violation: A=%.3f mm and Z=%.3f mm '
                 '(machine) would put A above Z by %.3f mm; max '
                 'allowed is %.3f mm. Drop A or raise Z first.' % (
-                    a, z, a - z, K))
+                    a_after, z_after, gap_after, K))
 
     # ----------------------------------------------------------- conversion
 

src/py/bbctrl/Mach.py

@@ -266,6 +266,14 @@ class Mach(Comm):
                 # now that the auxcnc stepper is exposed as a virtual
                 # A axis (see ExternalAxis).
                 cmd = self._rewrite_aux_mdi(cmd)
+                # Z-A coupling: if a line in this MDI/jog burst would
+                # exceed the safe band, attach an A delta on the same
+                # line so the planner runs Z and A together. Refuses
+                # only when the operator explicitly asks for an unsafe
+                # A lift or A would have to violate its soft min.
+                cmd = self._coordinate_mdi_for_coupling(cmd)
+                if not cmd or not cmd.strip():
+                    return
                 self._begin_cycle('mdi')
                 self.planner.mdi(cmd, with_limits)
                 super().resume()
@@ -273,6 +281,26 @@ class Mach(Comm):
             self.mlog.info("Exception during MDI: %s" % err)
             pass
 
+    def _coordinate_mdi_for_coupling(self, cmd):
+        """Apply ExternalAxis.coordinate_mdi to the MDI string. On
+        ExternalAxisError surface a user message and re-raise so the
+        outer try/except in mdi() halts cleanly."""
+        ext = getattr(self.ctrl, 'ext_axis', None)
+        if ext is None:
+            return cmd
+        try:
+            return ext.coordinate_mdi(cmd)
+        except Exception as e:
+            try:
+                self.ctrl.state.add_message(
+                    'Z-A coupling refused move: ' + str(e))
+            except Exception: pass
+            self.mlog.warning('Z-A coupling rewrite refused: %s' % e)
+            # Returning the original cmd would let the unsafe move
+            # through to the planner check, which would also refuse.
+            # Better to bail here and skip the MDI burst entirely.
+            return ''
+
     def _rewrite_aux_mdi(self, cmd):
         """Apply the ATC M-code preprocessor to a single MDI line.
         Returns possibly-multi-line G-code with HOOK: comments inserted."""

src/py/bbctrl/Planner.py

@@ -272,13 +272,14 @@ class Planner():
         if type == 'line':
             # Z-A coupling check: every line block that touches Z (or
             # A) is validated against the projected (A,Z) machine
-            # pair. The ExternalAxis check raises on violation so
+            # pair. The ExternalAxis check is improvement-aware: it
-            # gplan aborts the program with a useful error rather
+            # only refuses moves that worsen an existing violation
-            # than driving the gantry into the auxiliary tool. The
+            # or push a healthy state into one. So pure-XY jogs and
-            # check is skipped when coupling is disabled or A isn't
+            # recovery moves (Z up, A down) are not rejected even
-            # homed (see ExternalAxis.check_coupling).
+            # when (A-Z) is currently above the bound.
             ext_check = getattr(self.ctrl, 'ext_axis', None)
             if ext_check is not None:
+                from bbctrl.ExternalAxis import ExternalAxisError
                 target = block.get('target') or {}
                 z_target = target.get('z')
                 if z_target is None: z_target = target.get('Z')
@@ -287,9 +288,28 @@ class Planner():
                 if a_target is None:
                     a_target = target.get(a_letter.upper())
                 if z_target is not None or a_target is not None:
-                    ext_check.check_coupling(
+                    try:
-                        target_a_machine=a_target,
+                        ext_check.check_coupling(
-                        target_z_machine=z_target)
+                            target_a_machine=a_target,
+                            target_z_machine=z_target)
+                    except ExternalAxisError as e:
+                        # Convert the raw error into a clean abort:
+                        # surface the message to the operator, stop
+                        # the cycle, and skip this block. Returning
+                        # None drops the block from the AVR queue;
+                        # mach.stop() halts further planner output
+                        # so the rest of an offending program can't
+                        # leak through. The planner stays usable
+                        # for new MDI / jog commands.
+                        self.log.warning('Z-A coupling refused: %s' % e)
+                        try:
+                            self.ctrl.state.add_message(
+                                'Z-A coupling refused move: ' + str(e))
+                        except Exception: pass
+                        try:
+                            self.ctrl.mach.stop()
+                        except Exception: pass
+                        return None
 
             ext = self._external_axis_for_line(block)
             if ext is not None: