diff --git a/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java b/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
index 63d0d7587..90734bf02 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/AbstractRetryableWorkflowStep.java
@@ -75,7 +75,7 @@ protected void retryableGetMlTask(
         ActionListener<MLTask> mlTaskListener
     ) {
         CompletableFuture.runAsync(() -> {
-            while (!future.isDone()) {
+            do {
                 mlClient.getTask(taskId, ActionListener.wrap(response -> {
                     switch (response.getState()) {
                         case COMPLETED:
@@ -128,7 +128,7 @@ protected void retryableGetMlTask(
                     FutureUtils.cancel(future);
                     Thread.currentThread().interrupt();
                 }
-            }
+            } while (!future.isDone());
         }, threadPool.executor(WORKFLOW_THREAD_POOL));
     }