From 4504446e2eda67d96f41ad007bf1415d33279d32 Mon Sep 17 00:00:00 2001
From: Scott Lahteine <thinkyhead@users.noreply.github.com>
Date: Wed, 31 Jan 2024 17:57:21 -0600
Subject: [PATCH] onLevelingDone for G29_RETRY_AND_RECOVER

---
 Marlin/src/gcode/bedlevel/abl/G29.cpp |  4 +++-
 Marlin/src/gcode/gcode.cpp            | 31 ++++++++++++++++-----------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/Marlin/src/gcode/bedlevel/abl/G29.cpp b/Marlin/src/gcode/bedlevel/abl/G29.cpp
index 1ca3826c81..75e0f9dcaf 100644
--- a/Marlin/src/gcode/bedlevel/abl/G29.cpp
+++ b/Marlin/src/gcode/bedlevel/abl/G29.cpp
@@ -78,7 +78,9 @@ static void pre_g29_return(const bool retry, const bool did) {
   }
   if (did) {
     TERN_(HAS_DWIN_E3V2_BASIC, dwinLevelingDone());
-    TERN_(EXTENSIBLE_UI, ExtUI::onLevelingDone());
+    #if DISABLED(G29_RETRY_AND_RECOVER) && ENABLED(EXTENSIBLE_UI)
+      ExtUI::onLevelingDone();
+    #endif
   }
 }
 
diff --git a/Marlin/src/gcode/gcode.cpp b/Marlin/src/gcode/gcode.cpp
index 1ff643dc21..8502e28fae 100644
--- a/Marlin/src/gcode/gcode.cpp
+++ b/Marlin/src/gcode/gcode.cpp
@@ -289,22 +289,27 @@ void GcodeSuite::dwell(millis_t time) {
 
   void GcodeSuite::G29_with_retry() {
     uint8_t retries = G29_MAX_RETRIES;
-    while (G29()) { // G29 should return true for failed probes ONLY
-      if (retries) {
-        event_probe_recover();
-        --retries;
-      }
-      else {
-        event_probe_failure();
-        return;
-      }
+    bool fail = false;
+    for (;;) {
+      fail = G29(); // G29 should return true for failed probes ONLY
+      if (!fail || !retries) break;
+      event_probe_recover();
+      --retries;
     }
 
-    TERN_(HOST_PROMPT_SUPPORT, hostui.prompt_end());
+    if (fail) {
+      event_probe_failure();
+      TERN_(G29_HALT_ON_FAILURE, return);
+    }
+    else {
+      TERN_(HOST_PROMPT_SUPPORT, hostui.prompt_end());
+      #ifdef G29_SUCCESS_COMMANDS
+        process_subcommands_now(F(G29_SUCCESS_COMMANDS));
+      #endif
+    }
 
-    #ifdef G29_SUCCESS_COMMANDS
-      process_subcommands_now(F(G29_SUCCESS_COMMANDS));
-    #endif
+    TERN_(HAS_DWIN_E3V2_BASIC, dwinLevelingDone());
+    TERN_(EXTENSIBLE_UI, ExtUI::onLevelingDone());
   }
 
 #endif // G29_RETRY_AND_RECOVER