feat: Ability to restart flow on loop/branchall iteration (#2526)

* feat: Ability to restart flow on loop/branchall iteration

* Add popover to frontend

* refresh restart button onSelectedJobStepChange

* Display branch names for branchall

* Fix flow preview

* Fix compile break

---------

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>

backend/tests/worker.rs

@@ -2791,6 +2791,7 @@ async fn test_complex_flow_restart(db: Pool<Postgres>) {
         restarted_from: Some(RestartedFrom {
             flow_job_id: first_run_result.id,
             step_id: "h".to_owned(),
+            branch_or_iteration_n: None,
         }),
     }).run_until_complete(&db, port).await;
 

backend/windmill-api/openapi-deref.yaml

@@ -6134,7 +6134,7 @@ paths:
               schema:
                 type: string
                 format: uuid
-  /w/{workspace}/jobs/restart/f/{id}/from/{step_id}:
+  /w/{workspace}/jobs/restart/f/{id}/from/{step_id}/{branch_or_iteration_n}:
     post:
       summary: restart a completed flow at a given step
       operationId: restartFlowAtStep
@@ -6157,6 +6157,14 @@ paths:
           in: path
           schema:
             type: string
+        - name: branch_or_iteration_n
+          description: >-
+            for branchall or loop, the iteration at which the flow should
+            restart
+          required: true
+          in: path
+          schema:
+            type: integer
         - name: scheduled_for
           description: when to schedule this job (leave empty for immediate run)
           in: query
@@ -6422,6 +6430,8 @@ paths:
                       format: uuid
                     step_id:
                       type: string
+                    branch_or_iteration_n:
+                      type: integer
               required: &ref_145
                 - value
                 - content

backend/windmill-api/openapi.yaml

@@ -4230,7 +4230,7 @@ paths:
                 type: string
                 format: uuid
 
-  /w/{workspace}/jobs/restart/f/{id}/from/{step_id}:
+  /w/{workspace}/jobs/restart/f/{id}/from/{step_id}/{branch_or_iteration_n}:
     post:
       summary: restart a completed flow at a given step
       operationId: restartFlowAtStep
@@ -4245,6 +4245,12 @@ paths:
           in: path
           schema:
             type: string
+        - name: branch_or_iteration_n
+          description: for branchall or loop, the iteration at which the flow should restart
+          required: true
+          in: path
+          schema:
+            type: integer
         - name: scheduled_for
           description: when to schedule this job (leave empty for immediate run)
           in: query
@@ -7934,6 +7940,8 @@ components:
           format: uuid
         step_id:
           type: string
+        branch_or_iteration_n:
+          type: integer
 
     Policy:
       type: object

backend/windmill-api/src/jobs.rs

@@ -65,7 +65,11 @@ pub fn workspaced_service() -> Router {
                 .layer(cors.clone()),
         )
         .route(
-            "/restart/f/:job_id/from/*step_id",
+            "/restart/f/:job_id/from/:step_id",
+            post(restart_flow).head(|| async { "" }).layer(cors.clone()),
+        )
+        .route(
+            "/restart/f/:job_id/from/:step_id/:branch_of_iteration_n",
             post(restart_flow).head(|| async { "" }).layer(cors.clone()),
         )
         .route(
@@ -1633,7 +1637,12 @@ pub async fn restart_flow(
     Extension(db): Extension<DB>,
     Extension(user_db): Extension<UserDB>,
     Extension(rsmq): Extension<Option<rsmq_async::MultiplexedRsmq>>,
-    Path((w_id, job_id, step_id)): Path<(String, Uuid, String)>,
+    Path((w_id, job_id, step_id, branch_or_iteration_n)): Path<(
+        String,
+        Uuid,
+        String,
+        Option<usize>,
+    )>,
     Query(run_query): Query<RunJobQuery>,
 ) -> error::Result<(StatusCode, String)> {
     #[cfg(not(feature = "enterprise"))]
@@ -1671,7 +1680,11 @@ pub async fn restart_flow(
         &db,
         tx,
         &w_id,
-        JobPayload::RestartedFlow { completed_job_id: job_id, step_id: step_id },
+        JobPayload::RestartedFlow {
+            completed_job_id: job_id,
+            step_id: step_id,
+            branch_or_iteration_n: branch_or_iteration_n,
+        },
         push_args,
         &authed.username,
         &authed.email,

backend/windmill-common/src/flow_status.rs

@@ -57,6 +57,7 @@ pub struct ApprovalConditions {
 pub struct RestartedFrom {
     pub flow_job_id: Uuid,
     pub step_id: String,
+    pub branch_or_iteration_n: Option<usize>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

backend/windmill-common/src/jobs.rs

@@ -295,6 +295,7 @@ pub enum JobPayload {
     RestartedFlow {
         completed_job_id: Uuid,
         step_id: String,
+        branch_or_iteration_n: Option<usize>,
     },
     RawFlow {
         value: FlowValue,

backend/windmill-queue/src/jobs.rs

@@ -45,8 +45,8 @@ use windmill_common::{
     db::{Authed, UserDB},
     error::{self, Error},
     flow_status::{
-        FlowStatus, FlowStatusModule, FlowStatusModuleWParent, JobResult, RestartedFrom,
-        RetryStatus, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL,
+        BranchAllStatus, FlowStatus, FlowStatusModule, FlowStatusModuleWParent, Iterator,
+        JobResult, RestartedFrom, RetryStatus, MAX_RETRY_ATTEMPTS, MAX_RETRY_INTERVAL,
     },
     flows::{FlowModule, FlowModuleValue, FlowValue},
     jobs::{
@@ -2242,6 +2242,7 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection
                         Some(value.clone()),
                         restarted_from_val.flow_job_id,
                         restarted_from_val.step_id.as_str(),
+                        restarted_from_val.branch_or_iteration_n,
                     )
                     .await?;
                     FlowStatus {
@@ -2261,6 +2262,7 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection
                         restarted_from: Some(RestartedFrom {
                             flow_job_id: restarted_from_val.flow_job_id,
                             step_id: restarted_from_val.step_id,
+                            branch_or_iteration_n: restarted_from_val.branch_or_iteration_n,
                         }),
                     }
                 }
@@ -2311,7 +2313,7 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection
                 value.priority,
             )
         }
-        JobPayload::RestartedFlow { completed_job_id, step_id } => {
+        JobPayload::RestartedFlow { completed_job_id, step_id, branch_or_iteration_n } => {
             let (flow_path, raw_flow, step_n, truncated_modules, priority) =
                 restarted_flows_resolution(
                     _db,
@@ -2319,6 +2321,7 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection
                     None,
                     completed_job_id,
                     step_id.as_str(),
+                    branch_or_iteration_n,
                 )
                 .await?;
             let restarted_flow_status = FlowStatus {
@@ -2335,7 +2338,11 @@ pub async fn push<'c, T: Serialize + Send + Sync, R: rsmq_async::RsmqConnection
                 retry: RetryStatus { fail_count: 0, failed_jobs: vec![] },
                 // TODO: for now, flows with approval conditions aren't supported for restart
                 approval_conditions: None,
-                restarted_from: Some(RestartedFrom { flow_job_id: completed_job_id, step_id }),
+                restarted_from: Some(RestartedFrom {
+                    flow_job_id: completed_job_id,
+                    step_id,
+                    branch_or_iteration_n,
+                }),
             };
             (
                 None,
@@ -2617,6 +2624,7 @@ async fn restarted_flows_resolution(
     flow_value_if_any: Option<FlowValue>,
     completed_flow_id: Uuid,
     restart_step_id: &str,
+    branch_or_iteration_n: Option<usize>,
 ) -> Result<
     (
         Option<String>,
@@ -2671,11 +2679,100 @@ async fn restarted_flows_resolution(
             // skip module as it doesn't appear in the flow_value anymore
             continue;
         }
-        if module.id() == restart_step_id || dependent_module {
+        if module.id() == restart_step_id {
             // if the module ID is the one we want to restart the flow at, or if it's past it in the flow,
             // set the module as WaitingForPriorSteps as it needs to be re-run
-            truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id: module.id() });
+            if branch_or_iteration_n.is_none() || branch_or_iteration_n.unwrap() == 0 {
+                // The module as WaitingForPriorSteps as the entire module (i.e. all the branches) need to be re-run
+                truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id: module.id() });
+            } else {
+                // expect a module to be either a branchall (resp. loop), and resume the flow from this branch (resp. iteration)
+                let branch_or_iteration_n = branch_or_iteration_n.unwrap();
+                let module_definition = raw_flow
+                    .modules
+                    .iter()
+                    .find(|flow_value_module| flow_value_module.id == restart_step_id)
+                    .ok_or(Error::InternalErr(format!(
+                        "Module {} not found in flow definition",
+                        module.id()
+                    )))?;
+
+                match module_definition.value.clone() {
+                    FlowModuleValue::BranchAll { branches, parallel, .. } => {
+                        if parallel {
+                            return Err(Error::InternalErr(format!(
+                                "Module {} is a parallel branchall. It can only be restarted at a given branch if it's sequential",
+                                restart_step_id,
+                            )));
+                        }
+                        let total_branch_number = module.flow_jobs().map(|v| v.len()).unwrap_or(0);
+                        if total_branch_number <= branch_or_iteration_n {
+                            return Err(Error::InternalErr(format!(
+                                "Branch-all module {} has only {} branches. It can't be restarted on branch {}",
+                                restart_step_id,
+                                total_branch_number,
+                                branch_or_iteration_n,
+                            )));
+                        }
+                        let mut new_flow_jobs = module.flow_jobs().unwrap_or_default();
+                        new_flow_jobs.truncate(branch_or_iteration_n);
+                        truncated_modules.push(FlowStatusModule::InProgress {
+                            id: module.id(),
+                            job: new_flow_jobs[new_flow_jobs.len() - 1], // set to last finished job from completed flow
+                            iterator: None,
+                            flow_jobs: Some(new_flow_jobs),
+                            branch_chosen: None,
+                            branchall: Some(BranchAllStatus {
+                                branch: branch_or_iteration_n - 1, // Doing minus one here as this variable reflects the latest finished job in the iteration
+                                len: branches.len(),
+                            }),
+                            parallel: parallel,
+                        });
+                    }
+                    FlowModuleValue::ForloopFlow { parallel, .. } => {
+                        if parallel {
+                            return Err(Error::InternalErr(format!(
+                                "Module {} is not parallel loop. It can only be restarted at a given iteration if it's sequential",
+                                restart_step_id,
+                            )));
+                        }
+                        let total_iterations = module.flow_jobs().map(|v| v.len()).unwrap_or(0);
+                        if total_iterations <= branch_or_iteration_n {
+                            return Err(Error::InternalErr(format!(
+                                "For-loop module {} doesn't cannot be restarted on iteration number {} as it has only {} iterations",
+                                restart_step_id,
+                                branch_or_iteration_n,
+                                total_iterations,
+                            )));
+                        }
+                        let mut new_flow_jobs = module.flow_jobs().unwrap_or_default();
+                        new_flow_jobs.truncate(branch_or_iteration_n);
+
+                        truncated_modules.push(FlowStatusModule::InProgress {
+                            id: module.id(),
+                            job: new_flow_jobs[new_flow_jobs.len() - 1], // set to last finished job from completed flow
+                            iterator: Some(Iterator {
+                                index: branch_or_iteration_n - 1, // same deal as above, this refers to the last finished job
+                                itered: vec![], // Setting itered to empty array here, such that input transforms will be re-computed by worker_flows
+                            }),
+                            flow_jobs: Some(new_flow_jobs),
+                            branch_chosen: None,
+                            branchall: None,
+                            parallel: parallel,
+                        });
+                    }
+                    _ => {
+                        return Err(Error::InternalErr(format!(
+                            "Module {} is not a branchall or forloop, unable to restart it at step {:?}",
+                            restart_step_id,
+                            branch_or_iteration_n
+                        )));
+                    }
+                }
+            }
             dependent_module = true;
+        } else if dependent_module {
+            truncated_modules.push(FlowStatusModule::WaitingForPriorSteps { id: module.id() });
         } else {
             // else we simply "transfer" the module from the completed flow to the new one if it's a success
             step_n = step_n + 1;

backend/windmill-worker/src/worker_flow.rs

@@ -2123,7 +2123,7 @@ async fn compute_next_flow_transform(
                     let itered_raw = match iterator {
                         InputTransform::Static { value } => to_raw_value(value),
                         InputTransform::Javascript { expr } => {
-                            let mut context = HashMap::with_capacity(3);
+                            let mut context = HashMap::with_capacity(5);
                             context.insert("result".to_string(), arc_last_job_result.clone());
                             context.insert("previous_result".to_string(), arc_last_job_result);
                             context.insert("resumes".to_string(), resumes);
@@ -2167,20 +2167,53 @@ async fn compute_next_flow_transform(
                     flow_jobs: Some(flow_jobs),
                     ..
                 } if !*parallel => {
+                    let itered_new = if itered.is_empty() {
+                        // it's possible we need to re-compute the iterator Input Transforms here, in particular if the flow is being restarted inside the loop
+                        let by_id = if let Some(x) = by_id {
+                            x
+                        } else {
+                            get_transform_context(&flow_job, previous_id, &status).await?
+                        };
+                        let itered_raw = match iterator {
+                            InputTransform::Static { value } => to_raw_value(value),
+                            InputTransform::Javascript { expr } => {
+                                let mut context = HashMap::with_capacity(5);
+                                context.insert("result".to_string(), arc_last_job_result.clone());
+                                context.insert("previous_result".to_string(), arc_last_job_result);
+                                context.insert("resumes".to_string(), resumes);
+                                context.insert("resume".to_string(), resume);
+                                context.insert("approvers".to_string(), approvers);
+
+                                eval_timeout(
+                                    expr.to_string(),
+                                    context,
+                                    Some(arc_flow_job_args),
+                                    Some(client),
+                                    Some(by_id),
+                                )
+                                .await?
+                            }
+                        };
+                        serde_json::from_str::<Vec<serde_json::Value>>(itered_raw.get()).map_err(
+                            |not_array| {
+                                Error::ExecutionErr(format!(
+                                    "Expected an array value, found: {not_array}"
+                                ))
+                            },
+                        )?
+                    } else {
+                        itered.clone()
+                    };
                     let (index, next) = index
                         .checked_add(1)
-                        .and_then(|i| itered.get(i).map(|next| (i, next)))
-                        /* we shouldn't get here because update_flow_status_after_job_completion
-                         * should leave this state if there iteration is complete, but also it should
-                         * be reasonable to just enter a completed state instead of failing, similar to
-                         * iterating an empty list above */
+                        .and_then(|i| itered_new.get(i).map(|next| (i, next)))
                         .with_context(|| {
-                            format!("could not iterate index {index} of {itered:?}")
+                            format!("Could not find iteration number {index} restarting inside the for-loop flow. It's possible the itered-array has changed and this value isn't available anymore.")
                         })?;
 
                     LoopStatus::NextIteration(NextIteration {
                         index,
-                        itered: itered.clone(),
+                        itered: itered_new.clone(),
                         flow_jobs: flow_jobs.clone(),
                         new_args: Iter { index: index as i32, value: next.to_owned() },
                     })

frontend/src/lib/components/FlowPreviewContent.svelte

@@ -2,7 +2,7 @@
 	import { Job, JobService, type Flow, type FlowModule, type RestartedFrom } from '$lib/gen'
 	import { workspaceStore } from '$lib/stores'
 	import { faClose, faPlay, faRefresh } from '@fortawesome/free-solid-svg-icons'
-	import { Badge, Button, Drawer, Kbd } from './common'
+	import { Badge, Button, Drawer, Kbd, Popup } from './common'
 	import { createEventDispatcher, getContext } from 'svelte'
 	import Icon from 'svelte-awesome'
 	import type { FlowEditorContext } from './flows/types'
@@ -11,8 +11,8 @@
 	import FlowStatusViewer from '../components/FlowStatusViewer.svelte'
 	import FlowProgressBar from './flows/FlowProgressBar.svelte'
 	import CapturePayload from './flows/content/CapturePayload.svelte'
-	import { Loader2 } from 'lucide-svelte'
-	import { getModifierKey } from '$lib/utils'
+	import { ArrowRight, Loader2 } from 'lucide-svelte'
+	import { emptyString, getModifierKey } from '$lib/utils'
 	import DrawerContent from './common/drawer/DrawerContent.svelte'
 	import SavedInputs from './SavedInputs.svelte'
 	import { dfs } from './flows/dfs'
@@ -24,6 +24,12 @@
 	export let jobId: string | undefined = undefined
 	export let job: Job | undefined = undefined
 	let selectedJobStep: string | undefined = undefined
+	let branchOrIterationN: number = 0
+	let restartBranchNames: [number, string][] = []
+
+	let selectedJobStepIsTopLevel: boolean | undefined = undefined
+	let selectedJobStepType: 'single' | 'forloop' | 'branchall' = 'single'
+
 	let isRunning: boolean = false
 	let jobProgressReset: () => void
 
@@ -94,10 +100,33 @@
 		}
 	}
 
+	function onSelectedJobStepChange() {
+		if (selectedJobStep !== undefined && job?.flow_status?.modules !== undefined) {
+			selectedJobStepIsTopLevel =
+				job?.flow_status?.modules.map((m) => m.id).indexOf(selectedJobStep) >= 0
+			let moduleDefinition = job?.raw_flow?.modules.find((m) => m.id == selectedJobStep)
+			if (moduleDefinition?.value.type == 'forloopflow') {
+				selectedJobStepType = 'forloop'
+			} else if (moduleDefinition?.value.type == 'branchall') {
+				selectedJobStepType = 'branchall'
+				moduleDefinition?.value.branches.forEach((branch, idx) => {
+					restartBranchNames.push([
+						idx,
+						emptyString(branch.summary) ? `Branch #${idx}` : branch.summary!
+					])
+				})
+			} else {
+				selectedJobStepType = 'single'
+			}
+		}
+	}
+
 	$: if (job?.type === 'CompletedJob') {
 		isRunning = false
 	}
 
+	$: selectedJobStep !== undefined && onSelectedJobStepChange()
+
 	let inputLibraryDrawer: Drawer
 </script>
 
@@ -155,27 +184,94 @@
 			</Button>
 		{:else}
 			<div class="flex flex-row gap-4">
-				{#if jobId !== undefined && selectedJobStep !== undefined && job?.flow_status?.modules !== undefined && job?.flow_status?.modules
-						.map((m) => m.id)
-						.indexOf(selectedJobStep) >= 0}
-					<Button
-						size="xs"
-						color="light"
-						variant="border"
-						title={`Re-start this flow from step ${selectedJobStep} (included).`}
-						on:click={() => {
-							runPreview($previewArgs, {
-								flow_job_id: jobId,
-								step_id: selectedJobStep
-							})
-						}}
-						startIcon={{ icon: faPlay }}
-					>
-						Re-start from
-						<Badge baseClass="ml-1" color="indigo">
-							{selectedJobStep}
-						</Badge>
-					</Button>
+				{#if jobId !== undefined && selectedJobStep !== undefined && selectedJobStepIsTopLevel}
+					{#if selectedJobStepType == 'single'}
+						<Button
+							size="xs"
+							color="light"
+							variant="border"
+							title={`Re-start this flow from step ${selectedJobStep} (included).`}
+							on:click={() => {
+								runPreview($previewArgs, {
+									flow_job_id: jobId,
+									step_id: selectedJobStep,
+									branch_or_iteration_n: 0
+								})
+							}}
+							startIcon={{ icon: faPlay }}
+						>
+							Re-start from
+							<Badge baseClass="ml-1" color="indigo">
+								{selectedJobStep}
+							</Badge>
+						</Button>
+					{:else}
+						<Popup floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}>
+							<svelte:fragment slot="button">
+								<Button
+									title={`Re-start this flow from step ${selectedJobStep} (included).`}
+									variant="border"
+									color="blue"
+									startIcon={{ icon: faRefresh }}
+									on:click={() => {
+										runPreview($previewArgs, {
+											flow_job_id: jobId,
+											step_id: selectedJobStep,
+											branch_or_iteration_n: 0
+										})
+									}}
+									nonCaptureEvent={true}
+								>
+									Re-start from
+									<Badge baseClass="ml-1" color="indigo">
+										{selectedJobStep}
+									</Badge>
+								</Button>
+							</svelte:fragment>
+							<label class="block text-primary">
+								<div class="pb-1 text-sm text-secondary"
+									>{selectedJobStepType == 'forloop' ? 'From iteration #:' : 'From branch:'}</div
+								>
+								<div class="flex w-full">
+									{#if selectedJobStepType === 'forloop'}
+										<input
+											type="number"
+											min="0"
+											bind:value={branchOrIterationN}
+											class="!w-32 grow"
+											on:click|stopPropagation={() => {}}
+										/>
+									{:else}
+										<select
+											bind:value={branchOrIterationN}
+											class="!w-32 grow"
+											on:click|stopPropagation={() => {}}
+										>
+											{#each restartBranchNames as [branchIdx, branchName]}
+												<option value={branchIdx}>{branchName}</option>
+											{/each}
+										</select>
+									{/if}
+									<Button
+										size="xs"
+										color="blue"
+										buttonType="button"
+										btnClasses="!p-1 !w-[34px] !ml-1"
+										aria-label="Restart flow"
+										on:click|once={() => {
+											runPreview($previewArgs, {
+												flow_job_id: jobId,
+												step_id: selectedJobStep,
+												branch_or_iteration_n: branchOrIterationN
+											})
+										}}
+									>
+										<ArrowRight size={18} />
+									</Button>
+								</div>
+							</label>
+						</Popup>
+					{/if}
 				{/if}
 				<Button
 					variant="contained"

frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte

@@ -1,9 +1,10 @@
 <script lang="ts">
 	import { page } from '$app/stores'
 	import { JobService, Job } from '$lib/gen'
-	import { canWrite, displayDate, truncateHash } from '$lib/utils'
+	import { canWrite, displayDate, emptyString, truncateHash } from '$lib/utils'
 	import Icon from 'svelte-awesome'
 	import { check } from 'svelte-awesome/icons'
+	import { ArrowRight } from 'lucide-svelte'
 	import {
 		faRefresh,
 		faCircle,
@@ -31,7 +32,7 @@
 	import HighlightCode from '$lib/components/HighlightCode.svelte'
 	import TestJobLoader from '$lib/components/TestJobLoader.svelte'
 	import LogViewer from '$lib/components/LogViewer.svelte'
-	import { Button, ActionRow, Skeleton, Tab, Alert } from '$lib/components/common'
+	import { ActionRow, Button, Popup, Skeleton, Tab, Alert } from '$lib/components/common'
 	import FlowMetadata from '$lib/components/FlowMetadata.svelte'
 	import JobArgs from '$lib/components/JobArgs.svelte'
 	import FlowProgressBar from '$lib/components/flows/FlowProgressBar.svelte'
@@ -48,6 +49,11 @@
 
 	let viewTab: 'result' | 'logs' | 'code' = 'result'
 	let selectedJobStep: string | undefined = undefined
+	let branchOrIterationN: number = 0
+
+	let selectedJobStepIsTopLevel: boolean | undefined = undefined
+	let selectedJobStepType: 'single' | 'forloop' | 'branchall' = 'single'
+	let restartBranchNames: [number, string][] = []
 
 	// Test
 	let testIsLoading = false
@@ -75,7 +81,11 @@
 		}
 	}
 
-	async function restartFlow(id: string | undefined, stepId: string | undefined) {
+	async function restartFlow(
+		id: string | undefined,
+		stepId: string | undefined,
+		branchOrIterationN: number
+	) {
 		if (id === undefined || stepId === undefined) {
 			return
 		}
@@ -83,6 +93,7 @@
 			workspace: $workspaceStore!,
 			id,
 			stepId,
+			branchOrIterationN,
 			requestBody: {}
 		})
 		await goto('/run/' + run + '?workspace=' + $workspaceStore)
@@ -102,12 +113,37 @@
 		initView()
 	}
 
+	function onSelectedJobStepChange() {
+		console.log('yo')
+		if (selectedJobStep !== undefined && job?.flow_status?.modules !== undefined) {
+			selectedJobStepIsTopLevel =
+				job?.flow_status?.modules.map((m) => m.id).indexOf(selectedJobStep) >= 0
+			let moduleDefinition = job?.raw_flow?.modules.find((m) => m.id == selectedJobStep)
+			if (moduleDefinition?.value.type == 'forloopflow') {
+				selectedJobStepType = 'forloop'
+			} else if (moduleDefinition?.value.type == 'branchall') {
+				selectedJobStepType = 'branchall'
+				moduleDefinition?.value.branches.forEach((branch, idx) => {
+					restartBranchNames.push([
+						idx,
+						emptyString(branch.summary) ? `Branch #${idx}` : branch.summary!
+					])
+				})
+			} else {
+				selectedJobStepType = 'single'
+			}
+		}
+	}
+
 	$: {
 		if ($workspaceStore && $page.params.run && testJobLoader) {
 			forceCancel = false
 			getLogs()
 		}
 	}
+
+	$: selectedJobStep !== undefined && onSelectedJobStepChange()
+
 	let notfound = false
 	let forceCancel = false
 </script>
@@ -220,26 +256,87 @@
 					</Button>
 				{/if}
 			{/if}
-			{#if job?.job_kind === 'flow' && selectedJobStep !== undefined && job?.flow_status?.modules !== undefined && job?.flow_status?.modules
-					.map((m) => m.id)
-					.indexOf(selectedJobStep) >= 0}
-				<Button
-					title={`Re-start this flow from step ${selectedJobStep} (included). ${
-						!$enterpriseLicense ? ' This is a feature only available in enterprise edition.' : ''
-					}`}
-					variant="border"
-					color="blue"
-					disabled={!$enterpriseLicense}
-					on:click|once={() => {
-						restartFlow(job?.id, selectedJobStep)
-					}}
-					startIcon={{ icon: faRefresh }}
-				>
-					Re-start from
-					<Badge baseClass="ml-1" color="indigo">
-						{selectedJobStep}
-					</Badge>
-				</Button>
+			{#if job?.type === 'CompletedJob' && job?.job_kind === 'flow' && selectedJobStep !== undefined && selectedJobStepIsTopLevel}
+				{#if selectedJobStepType == 'single'}
+					<Button
+						title={`Re-start this flow from step ${selectedJobStep} (included). ${
+							!$enterpriseLicense ? ' This is a feature only available in enterprise edition.' : ''
+						}`}
+						variant="border"
+						color="blue"
+						disabled={!$enterpriseLicense}
+						on:click|once={() => {
+							restartFlow(job?.id, selectedJobStep, 0)
+						}}
+						startIcon={{ icon: faRefresh }}
+					>
+						Re-start from
+						<Badge baseClass="ml-1" color="indigo">
+							{selectedJobStep}
+						</Badge>
+					</Button>
+				{:else}
+					<Popup floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}>
+						<svelte:fragment slot="button">
+							<Button
+								title={`Re-start this flow from step ${selectedJobStep} (included). ${
+									!$enterpriseLicense
+										? ' This is a feature only available in enterprise edition.'
+										: ''
+								}`}
+								variant="border"
+								color="blue"
+								disabled={!$enterpriseLicense}
+								startIcon={{ icon: faRefresh }}
+								nonCaptureEvent={true}
+							>
+								Re-start from
+								<Badge baseClass="ml-1" color="indigo">
+									{selectedJobStep}
+								</Badge>
+							</Button>
+						</svelte:fragment>
+						<label class="block text-primary">
+							<div class="pb-1 text-sm text-secondary"
+								>{selectedJobStepType == 'forloop' ? 'From iteration #:' : 'From branch:'}</div
+							>
+							<div class="flex w-full">
+								{#if selectedJobStepType === 'forloop'}
+									<input
+										type="number"
+										min="0"
+										bind:value={branchOrIterationN}
+										class="!w-32 grow"
+										on:click|stopPropagation={() => {}}
+									/>
+								{:else}
+									<select
+										bind:value={branchOrIterationN}
+										class="!w-32 grow"
+										on:click|stopPropagation={() => {}}
+									>
+										{#each restartBranchNames as [branchIdx, branchName]}
+											<option value={branchIdx}>{branchName}</option>
+										{/each}
+									</select>
+								{/if}
+
+								<Button
+									size="xs"
+									color="blue"
+									buttonType="button"
+									btnClasses="!p-1 !w-[34px] !ml-1"
+									aria-label="Restart flow"
+									on:click|once={() => {
+										restartFlow(job?.id, selectedJobStep, branchOrIterationN)
+									}}
+								>
+									<ArrowRight size={18} />
+								</Button>
+							</div>
+						</label>
+					</Popup>
+				{/if}
 			{/if}
 			{#if job?.job_kind === 'script' || job?.job_kind === 'flow'}
 				<Button