Allow restarting the step chain at a specific position

Closes #44789

Signed-off-by: Pedro Igor <pigor.craveiro@gmail.com>

model/jpa/src/main/java/org/keycloak/models/workflow/DefaultWorkflowExecutionContext.java

@@ -16,8 +16,8 @@ final class DefaultWorkflowExecutionContext implements WorkflowExecutionContext
     private final Workflow workflow;
     private final WorkflowEvent event;
     private final KeycloakSession session;
-    private WorkflowStep currentStep;
+    private final WorkflowStep step;
-    private boolean restarted;
+    private boolean completed;
 
     /**
      * A new execution context for a workflow event. The execution ID is randomly generated.
@@ -52,15 +52,38 @@ final class DefaultWorkflowExecutionContext implements WorkflowExecutionContext
         this(session, workflow, null, step.stepId(), step.executionId(), step.resourceId());
     }
 
-    DefaultWorkflowExecutionContext(KeycloakSession session, Workflow workflow, WorkflowEvent event, String stepId, String executionId, String resourceId) {
+    /**
+     * A copy constructor that creates a new execution context based on an existing one but bound to a different {@link KeycloakSession}.
+
+     * @param session the session
+     * @param context the existing context
+     */
+    DefaultWorkflowExecutionContext(KeycloakSession session, DefaultWorkflowExecutionContext context) {
+        this(session, context.getWorkflow(), context.getEvent(), context.getCurrentStepId(), context.getExecutionId(), context.getResourceId());
+        completed = context.isCompleted();
+    }
+
+    /**
+     * A copy constructor that creates a new execution context based on an existing one but bound to a different {@link KeycloakSession} and the given {@link WorkflowStep}.
+
+     * @param session the session
+     * @param context the existing context
+     * @param step the current step
+     */
+    DefaultWorkflowExecutionContext(KeycloakSession session, DefaultWorkflowExecutionContext context, WorkflowStep step) {
+        this(session, context.getWorkflow(), context.getEvent(), step.getId(), context.getExecutionId(), context.getResourceId());
+        completed = context.isCompleted();
+    }
+
+    private DefaultWorkflowExecutionContext(KeycloakSession session, Workflow workflow, WorkflowEvent event, String stepId, String executionId, String resourceId) {
         this.session = session;
         this.workflow = workflow;
         this.event = event;
 
         if (stepId != null) {
-            this.currentStep = workflow.getStepById(stepId);
+            this.step = workflow.getStepById(stepId);
         } else {
-            this.currentStep = null;
+            this.step = null;
         }
 
         this.executionId = executionId;
@@ -79,7 +102,7 @@ final class DefaultWorkflowExecutionContext implements WorkflowExecutionContext
 
     @Override
     public WorkflowStep getNextStep() {
-        return workflow.getSteps(currentStep.getId()).skip(1).findFirst().orElse(null);
+        return workflow.getSteps(step.getId()).skip(1).findFirst().orElse(null);
     }
 
     String getExecutionId() {
@@ -90,27 +113,27 @@ final class DefaultWorkflowExecutionContext implements WorkflowExecutionContext
         return workflow;
     }
 
-    WorkflowStep getCurrentStep() {
+    WorkflowStep getStep() {
-        return currentStep;
+        return step;
     }
 
-    void setCurrentStep(WorkflowStep step) {
+    boolean isCompleted() {
-        this.currentStep = step;
+        return this.completed;
     }
 
-    boolean isRestarted() {
+    void complete() {
-        return this.restarted;
+        completed = true;
     }
 
-    void restart() {
+    void restart(int position) {
-        log.debugf("Restarting workflow '%s' for resource %s (execution id: %s)", workflow.getName(), resourceId, executionId);
+        new RestartWorkflowTask(this, position).run(session);
-        this.restarted = false;
-        this.currentStep = null;
-        new RunWorkflowTask(this).run(session);
-        this.restarted = true;
     }
 
     KeycloakSession getSession() {
         return session;
     }
+
+    private String getCurrentStepId() {
+        return step != null ? step.getId() : null;
+    }
 }

model/jpa/src/main/java/org/keycloak/models/workflow/DefaultWorkflowProvider.java

@@ -156,7 +156,7 @@ public class DefaultWorkflowProvider implements WorkflowProvider {
                             scheduled.resourceId(), scheduled.workflowId());
                     stateProvider.remove(scheduled.executionId());
                 } else {
-                    WorkflowStep step = context.getCurrentStep();
+                    WorkflowStep step = context.getStep();
                     if (step == null) {
                         log.warnf("Could not find step %s in workflow %s for resource %s. Cancelling execution of the workflow.",
                                 scheduled.stepId(), scheduled.workflowId(), scheduled.resourceId());
@@ -279,7 +279,7 @@ public class DefaultWorkflowProvider implements WorkflowProvider {
                     String executionId = scheduledStep.executionId();
                     String resourceId = scheduledStep.resourceId();
                     if (provider.restart(context)) {
-                        new DefaultWorkflowExecutionContext(session, workflow, event, scheduledStep).restart();
+                        new DefaultWorkflowExecutionContext(session, workflow, event, scheduledStep).restart(0);
                     } else if (provider.deactivate(context)) {
                         log.debugf("Workflow '%s' cancelled for resource %s (execution id: %s)", workflow.getName(), resourceId, executionId);
                         stateProvider.remove(executionId);

model/jpa/src/main/java/org/keycloak/models/workflow/RestartWorkflowStepProvider.java

@@ -1,11 +1,11 @@
 package org.keycloak.models.workflow;
 
-public class RestartWorkflowStepProvider implements WorkflowStepProvider {
+public record RestartWorkflowStepProvider(int position) implements WorkflowStepProvider {
 
     @Override
     public void run(WorkflowExecutionContext context) {
         if (context instanceof DefaultWorkflowExecutionContext) {
-            ((DefaultWorkflowExecutionContext) context).restart();
+            ((DefaultWorkflowExecutionContext) context).restart(position);
         } else {
             throw new IllegalArgumentException("Context must be DefaultWorkflowExecutionContext");
         }

model/jpa/src/main/java/org/keycloak/models/workflow/RestartWorkflowStepProviderFactory.java

@@ -3,18 +3,26 @@ package org.keycloak.models.workflow;
 import java.util.List;
 
 import org.keycloak.component.ComponentModel;
+import org.keycloak.component.ComponentValidationException;
 import org.keycloak.models.KeycloakSession;
+import org.keycloak.models.RealmModel;
 import org.keycloak.provider.ProviderConfigProperty;
 
-public class RestartWorkflowStepProviderFactory implements WorkflowStepProviderFactory<RestartWorkflowStepProvider> {
+public final class RestartWorkflowStepProviderFactory implements WorkflowStepProviderFactory<RestartWorkflowStepProvider> {
 
     public static final String ID = "restart";
-
+    public static final String CONFIG_POSITION = "position";
-    private final RestartWorkflowStepProvider provider = new RestartWorkflowStepProvider();
 
     @Override
     public RestartWorkflowStepProvider create(KeycloakSession session, ComponentModel model) {
-        return provider;
+        return new RestartWorkflowStepProvider(getPosition(model));
+    }
+
+    @Override
+    public void validateConfiguration(KeycloakSession session, RealmModel realm, ComponentModel model) throws ComponentValidationException {
+        if (getPosition(model) < 0) {
+            throw new ComponentValidationException("Position must be a non-negative integer");
+        }
     }
 
     @Override
@@ -52,4 +60,8 @@ public class RestartWorkflowStepProviderFactory implements WorkflowStepProviderF
     public String getHelpText() {
         return "Restarts the current workflow";
     }
+
+    private int getPosition(ComponentModel model) {
+        return model.get(CONFIG_POSITION, 0);
+    }
 }

model/jpa/src/main/java/org/keycloak/models/workflow/RestartWorkflowTask.java

@@ -0,0 +1,52 @@
+package org.keycloak.models.workflow;
+
+import java.util.List;
+
+import org.keycloak.models.KeycloakSession;
+
+import org.jboss.logging.Logger;
+
+class RestartWorkflowTask extends RunWorkflowTask {
+
+    private static final Logger log = Logger.getLogger(RestartWorkflowTask.class);
+
+    private final int position;
+
+    RestartWorkflowTask(DefaultWorkflowExecutionContext context, int position) {
+        super(context);
+        this.position = position;
+    }
+
+    @Override
+    protected WorkflowStep runCurrentStep(DefaultWorkflowExecutionContext context) {
+        Workflow workflow = context.getWorkflow();
+        List<WorkflowStep> steps = workflow.getSteps().toList();
+
+        if (position < 0 || position >= steps.size()) {
+            throw new IllegalArgumentException("Invalid position to restart workflow: " + position);
+        }
+
+        return steps.get(position);
+    }
+
+    @Override
+    public void run(KeycloakSession session) {
+        if (log.isDebugEnabled()) {
+            Workflow workflow = context.getWorkflow();
+            String resourceId = context.getResourceId();
+            String executionId = context.getExecutionId();
+            WorkflowStep currentStep = context.getStep();
+
+            if (currentStep == null) {
+                currentStep = workflow.getSteps().findFirst().orElse(null);
+            }
+
+            log.debugf("Restarting workflow '%s' for resource %s (execution id: %s) at step %s", workflow.getName(), resourceId, executionId, currentStep.getProviderId());
+        }
+        try {
+            super.run(session);
+        } finally {
+            context.complete();
+        }
+    }
+}

model/jpa/src/main/java/org/keycloak/models/workflow/RunWorkflowTask.java

@@ -1,6 +1,5 @@
 package org.keycloak.models.workflow;
 
-import java.util.List;
 
 import org.keycloak.common.util.DurationConverter;
 import org.keycloak.models.KeycloakSession;
@@ -13,54 +12,37 @@ class RunWorkflowTask extends WorkflowTransactionalTask {
 
     private static final Logger log = Logger.getLogger(RunWorkflowTask.class);
 
-    private final String executionId;
+    protected final DefaultWorkflowExecutionContext context;
-    private final String resourceId;
-    private final Workflow workflow;
-    private final WorkflowStep currentStep;
-    private final WorkflowEvent event;
 
     RunWorkflowTask(DefaultWorkflowExecutionContext context) {
         super(context.getSession());
-        this.executionId = context.getExecutionId();
+        this.context = context;
-        this.resourceId = context.getResourceId();
-        this.workflow = context.getWorkflow();
-        this.currentStep = context.getCurrentStep();
-        this.event = context.getEvent();
     }
 
     @Override
     public void run(KeycloakSession session) {
-        DefaultWorkflowExecutionContext context = new DefaultWorkflowExecutionContext(session, workflow, event, currentStep == null ? null : currentStep.getId(), executionId, resourceId);
+        DefaultWorkflowExecutionContext context = new DefaultWorkflowExecutionContext(session, this.context);
-        String executionId = context.getExecutionId();
-        String resourceId = context.getResourceId();
-        Workflow workflow = context.getWorkflow();
-        WorkflowStep currentStep = context.getCurrentStep();
-
-        if (currentStep != null) {
-            // we are resuming from a scheduled step - run it and then continue with the rest of the workflow
-            runWorkflowStep(context);
-        }
-
-        List<WorkflowStep> stepsToRun = workflow.getSteps()
-                .skip(currentStep != null ? currentStep.getPriority() : 0).toList();
         WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+        Workflow workflow = context.getWorkflow();
+        String resourceId = context.getResourceId();
+        String executionId = context.getExecutionId();
+        WorkflowStep nextStep = runCurrentStep(context);
 
-        for (WorkflowStep step : stepsToRun) {
+        while (nextStep != null) {
-            if (DurationConverter.isPositiveDuration(step.getAfter())) {
+            if (DurationConverter.isPositiveDuration(nextStep.getAfter())) {
+                log.debugf("Scheduling step %s to run in %s for resource %s (execution id: %s)",
+                        nextStep.getProviderId(), nextStep.getAfter(), resourceId, executionId);
                 // If a step has a time defined, schedule it and stop processing the other steps of workflow
-                log.debugf("Scheduling step %s to run in %s ms for resource %s (execution id: %s)",
+                stateProvider.scheduleStep(workflow, nextStep, resourceId, executionId);
-                        step.getProviderId(), step.getAfter(), resourceId, executionId);
-                stateProvider.scheduleStep(workflow, step, resourceId, executionId);
                 return;
-            } else {
+            }
-                // Otherwise, run the step right away
-                context.setCurrentStep(step);
 
-                runWorkflowStep(context);
+            DefaultWorkflowExecutionContext stepContext = new DefaultWorkflowExecutionContext(session, this.context, nextStep);
 
-                if (context.isRestarted()) {
+            nextStep = runWorkflowStep(stepContext);
-                    return;
+
-                }
+            if (stepContext.isCompleted()) {
+                return;
             }
         }
 
@@ -69,15 +51,22 @@ class RunWorkflowTask extends WorkflowTransactionalTask {
         stateProvider.remove(executionId);
     }
 
-    private void runWorkflowStep(DefaultWorkflowExecutionContext context) {
+    protected WorkflowStep runCurrentStep(DefaultWorkflowExecutionContext context) {
+        if (context.getStep() != null) {
+            return runWorkflowStep(context);
+        }
+        return context.getWorkflow().getSteps().findFirst().orElse(null);
+    }
+
+    private WorkflowStep runWorkflowStep(DefaultWorkflowExecutionContext context) {
+        WorkflowStep step = context.getStep();
         String executionId = context.getExecutionId();
-        WorkflowStep step = context.getCurrentStep();
         String resourceId = context.getResourceId();
         log.debugf("Running step %s on resource %s (execution id: %s)", step.getProviderId(), resourceId, executionId);
         try {
             getStepProvider(context.getSession(), step).run(context);
             log.debugf("Step %s completed successfully (execution id: %s)", step.getProviderId(), executionId);
-        } catch(WorkflowExecutionException e) {
+        } catch (WorkflowExecutionException e) {
             StringBuilder sb = new StringBuilder();
             sb.append("Step %s failed (execution id: %s)");
             String errorMessage = e.getMessage();
@@ -90,5 +79,7 @@ class RunWorkflowTask extends WorkflowTransactionalTask {
             }
             throw e;
         }
+
+        return context.getNextStep();
     }
 }

model/jpa/src/main/java/org/keycloak/models/workflow/ScheduleWorkflowTask.java

@@ -9,18 +9,11 @@ import org.jboss.logging.Logger;
 final class ScheduleWorkflowTask extends WorkflowTransactionalTask {
 
     private static final Logger log = Logger.getLogger(ScheduleWorkflowTask.class);
-
+    private final DefaultWorkflowExecutionContext context;
-    private final String executionId;
-    private final String resourceId;
-    private final Workflow workflow;
-    private final WorkflowEvent event;
 
     ScheduleWorkflowTask(DefaultWorkflowExecutionContext context) {
         super(context.getSession());
-        this.executionId = context.getExecutionId();
+        this.context = context;
-        this.resourceId = context.getResourceId();
-        this.workflow = context.getWorkflow();
-        this.event = context.getEvent();
     }
 
     @Override
@@ -33,7 +26,7 @@ final class ScheduleWorkflowTask extends WorkflowTransactionalTask {
             return;
         }
 
-        DefaultWorkflowExecutionContext workflowContext = new DefaultWorkflowExecutionContext(session, workflow, event, null, executionId, resourceId);
+        DefaultWorkflowExecutionContext workflowContext = new DefaultWorkflowExecutionContext(session, this.context);
         Workflow workflow = workflowContext.getWorkflow();
         WorkflowEvent event = workflowContext.getEvent();
         WorkflowStep firstStep = workflow.getSteps().findFirst().orElseThrow(() -> new WorkflowInvalidStateException("No steps found for workflow " + workflow.getName()));
@@ -52,6 +45,7 @@ final class ScheduleWorkflowTask extends WorkflowTransactionalTask {
 
     @Override
     public String toString() {
+        WorkflowEvent event = context.getEvent();
         return "eventType=" + event.getOperation() +
                 ",resourceType=" + event.getResourceType() +
                 ",resourceId=" + event.getResourceId();

model/jpa/src/main/java/org/keycloak/models/workflow/WorkflowValidator.java

@@ -5,6 +5,7 @@ import java.util.List;
 import java.util.Objects;
 
 import org.keycloak.common.util.DurationConverter;
+import org.keycloak.common.util.MultivaluedHashMap;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.workflow.conditions.expression.BooleanConditionParser;
 import org.keycloak.models.workflow.conditions.expression.ConditionNameCollector;
@@ -53,10 +54,16 @@ public class WorkflowValidator {
             if (steps.indexOf(restartStep) != steps.size() - 1) {
                 throw new WorkflowInvalidStateException("Workflow restart step must be the last step.");
             }
+            MultivaluedHashMap<String, String> config = restartStep.getConfig();
+            int position = config == null ? 0 : Integer.parseInt(config.getFirstOrDefault("position", "0"));
+            if (position < 0 || position >= steps.size()) {
+                throw new WorkflowInvalidStateException("Workflow restart step has invalid position: " + position);
+            }
             boolean hasScheduledStep = steps.stream()
+                    .skip(position)
                     .anyMatch(step -> DurationConverter.isPositiveDuration(step.getAfter()));
             if (!hasScheduledStep) {
-                throw new WorkflowInvalidStateException("A workflow with a restart step must have at least one step with a time delay.");
+                throw new WorkflowInvalidStateException("No scheduled step found if restarting at position " + position);
             }
         }
     }

tests/base/src/test/java/org/keycloak/tests/admin/model/workflow/WorkflowManagementTest.java

@@ -66,6 +66,7 @@ import org.keycloak.testframework.mail.annotations.InjectMailServer;
 import org.keycloak.testframework.realm.ManagedUser;
 import org.keycloak.testframework.realm.UserConfig;
 import org.keycloak.testframework.realm.UserConfigBuilder;
+import org.keycloak.testframework.remote.providers.runonserver.FetchOnServer;
 import org.keycloak.testframework.remote.providers.runonserver.RunOnServer;
 import org.keycloak.testframework.server.KeycloakUrls;
 import org.keycloak.testframework.util.ApiUtil;
@@ -83,6 +84,7 @@ import org.junit.jupiter.api.Test;
 import static org.keycloak.models.workflow.ResourceOperationType.USER_ADDED;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
@@ -768,7 +770,7 @@ public class WorkflowManagementTest extends AbstractWorkflowTest {
     }
 
     @Test
-    public void testRecurringWorkflow() {
+    public void testRestartWorkflow() {
         managedRealm.admin().workflows().create(WorkflowRepresentation.withName("myworkflow")
                 .onEvent(ResourceOperationType.USER_ADDED.toString())
                 .withSteps(
@@ -782,8 +784,38 @@ public class WorkflowManagementTest extends AbstractWorkflowTest {
         // create a new user - should bind the user to the workflow and setup the only step in the workflow
         managedRealm.admin().users().create(UserConfigBuilder.create().username("testuser").email("testuser@example.com").build()).close();
 
+        Long scheduledAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            return scheduledStep.scheduledAt();
+        }, Long.class);
+
         runScheduledSteps(Duration.ofDays(6));
 
+        Long reScheduleAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep step = workflow.getSteps().toList().get(0);
+
+            // Verify that the step was scheduled again for the user
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(scheduledAt));
+            return  scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(12));
+
         runOnServer.run((RunOnServer) session -> {
             RealmModel realm = session.getContext().getRealm();
             WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
@@ -797,15 +829,249 @@ public class WorkflowManagementTest extends AbstractWorkflowTest {
             ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
             assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
             assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(reScheduleAt));
         });
 
-        runScheduledSteps(Duration.ofDays(12));
 
         // Verify that there should be two emails sent
         assertEquals(2, findEmailsByRecipient(mailServer, "testuser@example.com").size());
         mailServer.runCleanup();
     }
 
+    @Test
+    public void testRestartFromPosition() {
+        try (Response response = managedRealm.admin().workflows().create(WorkflowRepresentation.withName("myworkflow")
+                .onEvent(ResourceOperationType.USER_ADDED.toString())
+                .withSteps(
+                        WorkflowStepRepresentation.create().of(NotifyUserStepProviderFactory.ID)
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("test", "value")
+                                .build(),
+                        WorkflowStepRepresentation.create().of(RestartWorkflowStepProviderFactory.ID)
+                                .withConfig(RestartWorkflowStepProviderFactory.CONFIG_POSITION, "1")
+                                .build()
+                ).build())) {
+            assertThat(response.getStatus(), is(Status.BAD_REQUEST.getStatusCode()));
+            assertThat(response.readEntity(ErrorRepresentation.class).getErrorMessage(),
+                    is("No scheduled step found if restarting at position 1"));
+        }
+        try (Response response = managedRealm.admin().workflows().create(WorkflowRepresentation.withName("myworkflow")
+                .onEvent(ResourceOperationType.USER_ADDED.toString())
+                .withSteps(
+                        WorkflowStepRepresentation.create().of(NotifyUserStepProviderFactory.ID)
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("test", "value")
+                                .build(),
+                        WorkflowStepRepresentation.create().of(RestartWorkflowStepProviderFactory.ID)
+                                .withConfig(RestartWorkflowStepProviderFactory.CONFIG_POSITION, "2")
+                                .build()
+                ).build())) {
+            assertThat(response.getStatus(), is(Status.BAD_REQUEST.getStatusCode()));
+            assertThat(response.readEntity(ErrorRepresentation.class).getErrorMessage(),
+                    is("No scheduled step found if restarting at position 2"));
+        }
+        managedRealm.admin().workflows().create(WorkflowRepresentation.withName("myworkflow")
+                .onEvent(ResourceOperationType.USER_ADDED.toString())
+                .withSteps(
+                        WorkflowStepRepresentation.create().of(NotifyUserStepProviderFactory.ID)
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("first", "first")
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("second", "second")
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(RestartWorkflowStepProviderFactory.ID)
+                                .withConfig(RestartWorkflowStepProviderFactory.CONFIG_POSITION, "1")
+                                .build()
+                ).build()).close();
+
+        // create a new user - should bind the user to the workflow and setup the only step in the workflow
+        managedRealm.admin().users().create(UserConfigBuilder.create().username("testuser").email("testuser@example.com").build()).close();
+
+        Long scheduleAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep step = workflow.getSteps().toList().get(0);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            return  scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(6));
+
+        Long reScheduledAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            // Verify that the first attribute was set, and the second is not yet set
+            assertThat(user.getAttributes().get("first"), containsInAnyOrder("first"));
+            assertThat(user.getAttributes().get("second"), nullValue());
+
+            // remove the first attribute to verify it gets set again after restart
+            user.removeAttribute("first");
+
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep step = workflow.getSteps().toList().get(2);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(scheduleAt));
+            return scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(6));
+
+        Long reScheduledAtLast = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            // Verify that both attributes are set
+            assertThat(user.getAttributes().get("first"), containsInAnyOrder("first"));
+            assertThat(user.getAttributes().get("second"), containsInAnyOrder("second"));
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep expectedStep = workflow.getSteps().toList().get(2);
+
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(expectedStep.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(reScheduledAt));
+            return scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(12));
+
+        // Verify that there should be one email sent, the first step should not have run again
+        assertEquals(1, findEmailsByRecipient(mailServer, "testuser@example.com").size());
+
+        runOnServer.run((RunOnServer) session -> {
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep expectedStep = workflow.getSteps().toList().get(2);
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(expectedStep.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(reScheduledAtLast));
+        });
+
+        mailServer.runCleanup();
+    }
+
+    @Test
+    public void testRestartFromLastStep() {
+        managedRealm.admin().workflows().create(WorkflowRepresentation.withName("myworkflow")
+                .onEvent(ResourceOperationType.USER_ADDED.toString())
+                .withSteps(
+                        WorkflowStepRepresentation.create().of(NotifyUserStepProviderFactory.ID)
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("first", "first")
+                                .build(),
+                        WorkflowStepRepresentation.create().of(SetUserAttributeStepProviderFactory.ID)
+                                .withConfig("second", "second")
+                                .after(Duration.ofDays(5))
+                                .build(),
+                        WorkflowStepRepresentation.create().of(RestartWorkflowStepProviderFactory.ID)
+                                .withConfig(RestartWorkflowStepProviderFactory.CONFIG_POSITION, "2")
+                                .build()
+                ).build()).close();
+
+        // create a new user - should bind the user to the workflow and setup the only step in the workflow
+        managedRealm.admin().users().create(UserConfigBuilder.create().username("testuser").email("testuser@example.com").build()).close();
+
+        Long scheduleAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep step = workflow.getSteps().toList().get(0);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            return  scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(6));
+
+        Long reScheduledAt = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            // Verify that the first attribute was set, and the second is not yet set
+            assertThat(user.getAttributes().get("first"), containsInAnyOrder("first"));
+            assertThat(user.getAttributes().get("second"), nullValue());
+
+            // remove the first attribute to verify it gets set again after restart
+            user.removeAttribute("first");
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep step = workflow.getSteps().toList().get(2);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(step.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(scheduleAt));
+            return scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(6));
+
+        Long reScheduledAtLast = runOnServer.fetch((FetchOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            // Verify that first attribute is not set, and the second is set
+            assertThat(user.getAttributes().get("first"), nullValue());
+            assertThat(user.getAttributes().get("second"), containsInAnyOrder("second"));
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep expectedStep = workflow.getSteps().toList().get(2);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(expectedStep.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(reScheduledAt));
+            user.removeAttribute("second");
+            return scheduledStep.scheduledAt();
+        }, Long.class);
+
+        runScheduledSteps(Duration.ofDays(12));
+
+        // Verify that there should be one email sent, the first step should not have run again
+        assertEquals(1, findEmailsByRecipient(mailServer, "testuser@example.com").size());
+        runOnServer.run((RunOnServer) session -> {
+            RealmModel realm = session.getContext().getRealm();
+            UserModel user = session.users().getUserByUsername(realm, "testuser");
+            // Verify that the first attribute is not set, and the second is set again
+            assertThat(user.getAttributes().get("first"), nullValue());
+            assertThat(user.getAttributes().get("second"), containsInAnyOrder(("second")));
+            WorkflowProvider provider = session.getProvider(WorkflowProvider.class);
+            Workflow workflow = provider.getWorkflows().toList().get(0);
+            WorkflowStep expectedStep = workflow.getSteps().toList().get(2);
+            WorkflowStateProvider stateProvider = session.getProvider(WorkflowStateProvider.class);
+            ScheduledStep scheduledStep = stateProvider.getScheduledStep(workflow.getId(), user.getId());
+            assertNotNull(scheduledStep, "A step should have been scheduled for the user " + user.getUsername());
+            assertEquals(expectedStep.getId(), scheduledStep.stepId(), "The step should have been scheduled again");
+            assertThat("The step should have been scheduled again at a later time", scheduledStep.scheduledAt(), greaterThan(reScheduledAtLast));
+        });
+        mailServer.runCleanup();
+    }
+
     @Test
     public void testRunImmediateWorkflow() {
         // create a test workflow with no time conditions - should run immediately when scheduled