diff --git a/backend/tests/job_payload.rs b/backend/tests/job_payload.rs
index 545665c39f..ed1413439c 100644
--- a/backend/tests/job_payload.rs
+++ b/backend/tests/job_payload.rs
@@ -52,8 +52,10 @@ mod job_payload {
             let result = RunJob::from(JobPayload::ScriptHash {
                 hash: ScriptHash(123412),
                 path: "f/system/hello".to_string(),
-                concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
-                debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
+                concurrency_settings:
+                    windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+                debouncing_settings:
+                    windmill_common::runnable_settings::DebouncingSettings::default(),
                 cache_ttl: None,
                 cache_ignore_s3_path: None,
                 dedicated_worker: None,
@@ -90,8 +92,10 @@ mod job_payload {
                 language: ScriptLang::Deno,
                 priority: None,
                 apply_preprocessor: true,
-                concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default(),
-                debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
+                concurrency_settings:
+                    windmill_common::runnable_settings::ConcurrencySettings::default(),
+                debouncing_settings:
+                    windmill_common::runnable_settings::DebouncingSettings::default(),
             })
             .run_until_complete_with(db, false, port, |id| async move {
                 let job = sqlx::query!("SELECT preprocessed FROM v2_job WHERE id = $1", id)
@@ -163,7 +167,8 @@ mod job_payload {
             let result = RunJob::from(JobPayload::FlowScript {
                 id: flow_scripts[0],
                 language: ScriptLang::Deno,
-                concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default(),
+                concurrency_settings:
+                    windmill_common::runnable_settings::ConcurrencySettings::default(),
                 cache_ttl: None,
                 cache_ignore_s3_path: None,
                 dedicated_worker: None,
@@ -182,7 +187,8 @@ mod job_payload {
             let result = RunJob::from(JobPayload::FlowScript {
                 id: flow_scripts[1],
                 language: ScriptLang::Deno,
-                concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default(),
+                concurrency_settings:
+                    windmill_common::runnable_settings::ConcurrencySettings::default(),
                 cache_ttl: None,
                 cache_ignore_s3_path: None,
                 dedicated_worker: None,
@@ -547,6 +553,7 @@ mod job_payload {
                 completed_job_id,
                 step_id: "a".into(),
                 branch_or_iteration_n: None,
+                flow_version: None,
             })
             .arg("iter", json!({ "value": "tests", "index": 0 }))
             .run_until_complete(&db, false, port)
@@ -714,7 +721,12 @@ mod job_payload {
         )
         .await;
         let flow_job_id = test(
-            Some(RestartedFrom { flow_job_id, step_id: "a".into(), branch_or_iteration_n: None }),
+            Some(RestartedFrom {
+                flow_job_id,
+                step_id: "a".into(),
+                branch_or_iteration_n: None,
+                flow_version: None,
+            }),
             json!("foo"),
             json!([
                 "a: Hello foo! foo! foo!",
@@ -724,7 +736,12 @@ mod job_payload {
         )
         .await;
         let flow_job_id = test(
-            Some(RestartedFrom { flow_job_id, step_id: "b".into(), branch_or_iteration_n: None }),
+            Some(RestartedFrom {
+                flow_job_id,
+                step_id: "b".into(),
+                branch_or_iteration_n: None,
+                flow_version: None,
+            }),
             json!("bar"),
             json!([
                 "a: Hello foo! bar! bar!",
@@ -738,6 +755,7 @@ mod job_payload {
                 flow_job_id,
                 step_id: "c".into(),
                 branch_or_iteration_n: Some(1),
+                flow_version: None,
             }),
             json!("yolo"),
             json!([
diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs
index 5cef946315..3328296faf 100644
--- a/backend/tests/worker.rs
+++ b/backend/tests/worker.rs
@@ -865,7 +865,8 @@ func main(derp string) (string, error) {
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("derp", json!("world"))
@@ -900,7 +901,8 @@ fn main(world: String) -> Result<String, String> {
         lock: None,
         language: ScriptLang::Rust,
         cache_ignore_s3_path: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
         cache_ttl: None,
         dedicated_worker: None,
@@ -978,7 +980,8 @@ echo "hello $msg"
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("msg", json!("world"))
@@ -1011,7 +1014,8 @@ def main [ msg: string ] {
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("msg", json!("world"))
@@ -1064,7 +1068,8 @@ def main [
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("a", json!("3"))
@@ -1126,7 +1131,8 @@ public class Main {
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("a", json!(3))
@@ -1161,7 +1167,8 @@ export async function main(a: Date) {
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("a", json!("2024-09-24T10:00:00.000Z"))
@@ -1196,7 +1203,8 @@ export async function main(a: Date) {
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("a", json!("2024-09-24T10:00:00.000Z"))
@@ -1232,7 +1240,8 @@ def main(a: datetime, b: bytes):
         cache_ttl: None,
         cache_ignore_s3_path: None,
         dedicated_worker: None,
-        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default().into(),
+        concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default()
+            .into(),
         debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
     }))
     .arg("a", json!("2024-09-24T10:00:00.000Z"))
@@ -2279,6 +2288,7 @@ async fn test_complex_flow_restart(db: Pool<Postgres>) -> anyhow::Result<()> {
             flow_job_id: first_run_result.id,
             step_id: "h".to_owned(),
             branch_or_iteration_n: None,
+            flow_version: None,
         }),
     })
     .run_until_complete(&db, false, port)
diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml
index 2aa77c1d1c..f277a2cd6a 100644
--- a/backend/windmill-api/openapi.yaml
+++ b/backend/windmill-api/openapi.yaml
@@ -8317,7 +8317,7 @@ paths:
               schema:
                 type: string
 
-  /w/{workspace}/jobs/restart/f/{id}/from/{step_id}/{branch_or_iteration_n}:
+  /w/{workspace}/jobs/restart/f/{id}:
     post:
       summary: restart a completed flow at a given step
       operationId: restartFlowAtStep
@@ -8326,20 +8326,6 @@ paths:
       parameters:
         - $ref: "#/components/parameters/WorkspaceId"
         - $ref: "#/components/parameters/JobId"
-        - name: step_id
-          description: step id to restart the flow from
-          required: true
-          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
@@ -8362,12 +8348,24 @@ paths:
             type: boolean
 
       requestBody:
-        description: flow args
+        description: restart flow parameters
         required: true
         content:
           application/json:
             schema:
-              $ref: "#/components/schemas/ScriptArgs"
+              type: object
+              required:
+                - step_id
+              properties:
+                step_id:
+                  type: string
+                  description: step id to restart the flow from
+                branch_or_iteration_n:
+                  type: integer
+                  description: for branchall or loop, the iteration at which the flow should restart (optional)
+                flow_version:
+                  type: integer
+                  description: specific flow version to use for restart (optional, uses current version if not specified)
 
       responses:
         "201":
@@ -20182,6 +20180,8 @@ components:
           type: string
         branch_or_iteration_n:
           type: integer
+        flow_version:
+          type: integer
 
     Policy:
       type: object
diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs
index e84fa25ad7..d899047942 100644
--- a/backend/windmill-api/src/jobs.rs
+++ b/backend/windmill-api/src/jobs.rs
@@ -158,11 +158,7 @@ pub fn workspaced_service() -> Router {
                 .layer(ce_headers.clone()),
         )
         .route(
-            "/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",
+            "/restart/f/:job_id",
             post(restart_flow).head(|| async { "" }).layer(cors.clone()),
         )
         .route(
@@ -4395,18 +4391,24 @@ pub async fn restart_flow(
     ));
 }
 
+#[cfg(feature = "enterprise")]
+#[derive(Deserialize)]
+pub struct RestartFlowRequestBody {
+    step_id: String,
+    branch_or_iteration_n: Option<usize>,
+    flow_version: Option<i64>,
+}
+
 #[cfg(feature = "enterprise")]
 pub async fn restart_flow(
     authed: ApiAuthed,
     Extension(db): Extension<DB>,
     Extension(user_db): Extension<UserDB>,
-    Path((w_id, job_id, step_id, branch_or_iteration_n)): Path<(
-        String,
-        Uuid,
-        String,
-        Option<usize>,
-    )>,
+    Path((w_id, job_id)): Path<(String, Uuid)>,
     Query(run_query): Query<RunJobQuery>,
+    Json(RestartFlowRequestBody { step_id, branch_or_iteration_n, flow_version }): Json<
+        RestartFlowRequestBody,
+    >,
 ) -> error::Result<(StatusCode, String)> {
     check_license_key_valid().await?;
 
@@ -4445,7 +4447,12 @@ pub async fn restart_flow(
         &db,
         tx,
         &w_id,
-        JobPayload::RestartedFlow { completed_job_id: job_id, step_id, branch_or_iteration_n },
+        JobPayload::RestartedFlow {
+            completed_job_id: job_id,
+            step_id,
+            branch_or_iteration_n,
+            flow_version,
+        },
         push_args,
         &authed.username,
         &authed.email,
diff --git a/backend/windmill-common/src/flow_status.rs b/backend/windmill-common/src/flow_status.rs
index 44da5f6639..307d28313e 100644
--- a/backend/windmill-common/src/flow_status.rs
+++ b/backend/windmill-common/src/flow_status.rs
@@ -72,6 +72,7 @@ pub struct RestartedFrom {
     pub flow_job_id: Uuid,
     pub step_id: String,
     pub branch_or_iteration_n: Option<usize>,
+    pub flow_version: Option<i64>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
diff --git a/backend/windmill-common/src/jobs.rs b/backend/windmill-common/src/jobs.rs
index 0fb389ecd5..534db20bf0 100644
--- a/backend/windmill-common/src/jobs.rs
+++ b/backend/windmill-common/src/jobs.rs
@@ -448,6 +448,7 @@ pub enum JobPayload {
         completed_job_id: Uuid,
         step_id: String,
         branch_or_iteration_n: Option<usize>,
+        flow_version: Option<i64>,
     },
 
     /// Flow Preview
diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs
index 9566c9869b..e9d9010ae7 100644
--- a/backend/windmill-queue/src/jobs.rs
+++ b/backend/windmill-queue/src/jobs.rs
@@ -4379,6 +4379,7 @@ pub async fn push<'c, 'd>(
                             restarted_from_val.flow_job_id,
                             restarted_from_val.step_id.as_str(),
                             restarted_from_val.branch_or_iteration_n,
+                            restarted_from_val.flow_version,
                         )
                         .await?;
                     FlowStatus {
@@ -4400,6 +4401,7 @@ pub async fn push<'c, 'd>(
                             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,
+                            flow_version: restarted_from_val.flow_version,
                         }),
                         user_states,
                         preprocessor_module: None,
@@ -4674,7 +4676,12 @@ pub async fn push<'c, 'd>(
                 ..Default::default()
             }
         }
-        JobPayload::RestartedFlow { completed_job_id, step_id, branch_or_iteration_n } => {
+        JobPayload::RestartedFlow {
+            completed_job_id,
+            step_id,
+            branch_or_iteration_n,
+            flow_version,
+        } => {
             let (
                 version,
                 flow_path,
@@ -4689,6 +4696,7 @@ pub async fn push<'c, 'd>(
                 completed_job_id,
                 step_id.as_str(),
                 branch_or_iteration_n,
+                flow_version,
             )
             .await?;
 
@@ -4711,6 +4719,7 @@ pub async fn push<'c, 'd>(
                     flow_job_id: completed_job_id,
                     step_id,
                     branch_or_iteration_n,
+                    flow_version,
                 }),
                 user_states,
                 preprocessor_module: None,
@@ -5523,12 +5532,114 @@ pub fn canceled_job_to_result(job: &MiniPulledJob) -> serde_json::Value {
     serde_json::json!({"message": format!("Job canceled: {reason} by {canceler}"), "name": "Canceled", "reason": reason, "canceler": canceler})
 }
 
+/// Helper function to create a restarted module for branch/iteration restart
+fn create_restarted_module(
+    module: &FlowStatusModule,
+    module_definition: &FlowModule,
+    branch_or_iteration_n: usize,
+    restart_step_id: &str,
+) -> Result<FlowStatusModule, Error> {
+    match module_definition.get_value() {
+        Ok(FlowModuleValue::BranchAll { branches, parallel, .. }) => {
+            if parallel {
+                return Err(Error::internal_err(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::internal_err(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);
+            let mut new_flow_jobs_success = module.flow_jobs_success();
+            if let Some(new_flow_jobs_success) = new_flow_jobs_success.as_mut() {
+                new_flow_jobs_success.truncate(branch_or_iteration_n);
+            }
+            let mut new_flow_jobs_timeline = module.flow_jobs_duration();
+            if let Some(new_flow_jobs_timeline) = new_flow_jobs_timeline.as_mut() {
+                new_flow_jobs_timeline.truncate(branch_or_iteration_n);
+            }
+            Ok(FlowStatusModule::InProgress {
+                id: module.id(),
+                job: new_flow_jobs[new_flow_jobs.len() - 1],
+                iterator: None,
+                flow_jobs: Some(new_flow_jobs),
+                flow_jobs_success: new_flow_jobs_success,
+                flow_jobs_duration: new_flow_jobs_timeline,
+                branch_chosen: None,
+                branchall: Some(BranchAllStatus {
+                    branch: branch_or_iteration_n - 1,
+                    len: branches.len(),
+                }),
+                parallel,
+                while_loop: false,
+                progress: None,
+                agent_actions: None,
+                agent_actions_success: None,
+            })
+        }
+        Ok(FlowModuleValue::ForloopFlow { parallel, .. }) => {
+            if parallel {
+                return Err(Error::internal_err(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::internal_err(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);
+            let mut new_flow_jobs_success = module.flow_jobs_success();
+            if let Some(new_flow_jobs_success) = new_flow_jobs_success.as_mut() {
+                new_flow_jobs_success.truncate(branch_or_iteration_n);
+            }
+            let mut new_flow_jobs_timeline = module.flow_jobs_duration();
+            if let Some(new_flow_jobs_timeline) = new_flow_jobs_timeline.as_mut() {
+                new_flow_jobs_timeline.truncate(branch_or_iteration_n);
+            }
+            Ok(FlowStatusModule::InProgress {
+                id: module.id(),
+                job: new_flow_jobs[new_flow_jobs.len() - 1],
+                iterator: Some(FlowIterator { index: branch_or_iteration_n - 1, itered: vec![] }),
+                flow_jobs: Some(new_flow_jobs),
+                flow_jobs_success: new_flow_jobs_success,
+                flow_jobs_duration: new_flow_jobs_timeline,
+                branch_chosen: None,
+                branchall: None,
+                parallel,
+                while_loop: false,
+                progress: None,
+                agent_actions: None,
+                agent_actions_success: None,
+            })
+        }
+        _ => Err(Error::internal_err(format!(
+            "Module {} is not a branchall or forloop, unable to restart it at step {:?}",
+            restart_step_id, branch_or_iteration_n
+        ))),
+    }
+}
+
 async fn restarted_flows_resolution(
     db: &Pool<Postgres>,
     workspace_id: &str,
     completed_flow_id: Uuid,
     restart_step_id: &str,
     branch_or_iteration_n: Option<usize>,
+    flow_version: Option<i64>,
+    // parents: Vec<RestartedParent>,
 ) -> Result<
     (
         Option<i64>,
@@ -5560,9 +5671,24 @@ async fn restarted_flows_resolution(
         ))
     })?;
 
-    let flow_data = cache::job::fetch_flow(db, &row.job_kind, row.script_hash)
-        .or_else(|_| cache::job::fetch_preview_flow(db.into(), &completed_flow_id, row.raw_flow))
-        .await?;
+    let current_flow_version = row.script_hash.map(|x| x.0);
+    let is_version_change = flow_version.is_some()
+        && current_flow_version.is_some()
+        && flow_version != current_flow_version
+        && row.job_kind == JobKind::Flow;
+
+    let flow_data = if is_version_change {
+        // Fetch the new flow version
+        let new_version = flow_version.unwrap();
+        cache::flow::fetch_version(db, new_version).await?
+    } else {
+        cache::job::fetch_flow(db, &row.job_kind, row.script_hash)
+            .or_else(|_| {
+                cache::job::fetch_preview_flow(db.into(), &completed_flow_id, row.raw_flow)
+            })
+            .await?
+    };
+
     let flow_value = flow_data.value();
     let flow_status = row
         .flow_status
@@ -5576,137 +5702,93 @@ async fn restarted_flows_resolution(
     let mut step_n = 0;
     let mut dependent_module = false;
     let mut truncated_modules: Vec<FlowStatusModule> = vec![];
-    for module in flow_status.modules {
-        let Some(module_definition) = flow_value
-            .modules
-            .iter()
-            .find(|flow_value_module| flow_value_module.id == module.id())
-        else {
-            // skip module as it doesn't appear in the flow_value anymore
-            continue;
-        };
-        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
-            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();
 
-                match module_definition.get_value() {
-                    Ok(FlowModuleValue::BranchAll { branches, parallel, .. }) => {
-                        if parallel {
-                            return Err(Error::internal_err(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::internal_err(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);
-                        let mut new_flow_jobs_success = module.flow_jobs_success();
-                        if let Some(new_flow_jobs_success) = new_flow_jobs_success.as_mut() {
-                            new_flow_jobs_success.truncate(branch_or_iteration_n);
-                        }
-                        let mut new_flow_jobs_timeline = module.flow_jobs_duration();
-                        if let Some(new_flow_jobs_timeline) = new_flow_jobs_timeline.as_mut() {
-                            new_flow_jobs_timeline.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),
-                            flow_jobs_success: new_flow_jobs_success,
-                            flow_jobs_duration: new_flow_jobs_timeline,
-                            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,
-                            while_loop: false,
-                            progress: None,
-                            agent_actions: None,
-                            agent_actions_success: None,
-                        });
-                    }
-                    Ok(FlowModuleValue::ForloopFlow { parallel, .. }) => {
-                        if parallel {
-                            return Err(Error::internal_err(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::internal_err(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);
-                        let mut new_flow_jobs_success = module.flow_jobs_success();
-                        if let Some(new_flow_jobs_success) = new_flow_jobs_success.as_mut() {
-                            new_flow_jobs_success.truncate(branch_or_iteration_n);
-                        }
-                        let mut new_flow_jobs_timeline = module.flow_jobs_duration();
-                        if let Some(new_flow_jobs_timeline) = new_flow_jobs_timeline.as_mut() {
-                            new_flow_jobs_timeline.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(FlowIterator {
-                                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),
-                            flow_jobs_success: new_flow_jobs_success,
-                            flow_jobs_duration: new_flow_jobs_timeline,
-                            branch_chosen: None,
-                            branchall: None,
-                            parallel,
-                            while_loop: false,
-                            progress: None,
-                            agent_actions: None,
-                            agent_actions_success: None,
-                        });
-                    }
-                    _ => {
-                        return Err(Error::internal_err(format!(
-                            "Module {} is not a branchall or forloop, unable to restart it at step {:?}",
+    if is_version_change {
+        // When the flow version has changed, create flow status from scratch
+        // based on the new flow, but match modules from the old flow status where possible
+        for module_definition in &flow_value.modules {
+            let module_id = &module_definition.id;
+
+            if module_id == restart_step_id {
+                // Mark this and all following modules as WaitingForPriorSteps
+                if branch_or_iteration_n.is_none() || branch_or_iteration_n.unwrap() == 0 {
+                    truncated_modules
+                        .push(FlowStatusModule::WaitingForPriorSteps { id: module_id.clone() });
+                } else {
+                    // Handle branch/iteration restart for version changes
+                    let branch_n = branch_or_iteration_n.unwrap();
+                    // Try to find matching module in old flow status
+                    if let Some(old_module) =
+                        flow_status.modules.iter().find(|m| &m.id() == module_id)
+                    {
+                        truncated_modules.push(create_restarted_module(
+                            old_module,
+                            module_definition,
+                            branch_n,
                             restart_step_id,
-                            branch_or_iteration_n
-                        )));
+                        )?);
+                    } else {
+                        // Module not found in old flow, mark as waiting
+                        truncated_modules
+                            .push(FlowStatusModule::WaitingForPriorSteps { id: module_id.clone() });
                     }
                 }
+                dependent_module = true;
+            } else if dependent_module {
+                truncated_modules
+                    .push(FlowStatusModule::WaitingForPriorSteps { id: module_id.clone() });
+            } else {
+                // Before the restart step, try to match with old flow status
+                if let Some(old_module) = flow_status.modules.iter().find(|m| &m.id() == module_id)
+                {
+                    truncated_modules.push(old_module.clone());
+                } else {
+                    truncated_modules
+                        .push(FlowStatusModule::WaitingForPriorSteps { id: module_id.clone() });
+                }
+                step_n += 1;
+            }
+        }
+    } else {
+        // Original logic for same version
+        for module in flow_status.modules {
+            let Some(module_definition) = flow_value
+                .modules
+                .iter()
+                .find(|flow_value_module| flow_value_module.id == module.id())
+            else {
+                // skip module as it doesn't appear in the flow_value anymore
+                continue;
+            };
+            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
+                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)
+                    truncated_modules.push(create_restarted_module(
+                        &module,
+                        module_definition,
+                        branch_or_iteration_n.unwrap(),
+                        restart_step_id,
+                    )?);
+                }
+                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;
+                match module.clone() {
+                    FlowStatusModule::Success { .. } => Ok(truncated_modules.push(module)),
+                    _ => Err(Error::internal_err(format!(
+                        "Flow cannot be restarted from a non successful module",
+                    ))),
+                }?;
             }
-            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;
-            match module.clone() {
-                FlowStatusModule::Success { .. } => Ok(truncated_modules.push(module)),
-                _ => Err(Error::internal_err(format!(
-                    "Flow cannot be restarted from a non successful module",
-                ))),
-            }?;
         }
     }
 
@@ -5719,7 +5801,11 @@ async fn restarted_flows_resolution(
     }
 
     Ok((
-        row.script_hash.map(|x| x.0),
+        if is_version_change {
+            flow_version
+        } else {
+            row.script_hash.map(|x| x.0)
+        },
         row.script_path,
         flow_data,
         step_n,
diff --git a/frontend/src/lib/components/FlowPreviewContent.svelte b/frontend/src/lib/components/FlowPreviewContent.svelte
index b6f2fc4e18..af002a21ec 100644
--- a/frontend/src/lib/components/FlowPreviewContent.svelte
+++ b/frontend/src/lib/components/FlowPreviewContent.svelte
@@ -1,6 +1,4 @@
 <script lang="ts">
-	import { stopPropagation } from 'svelte/legacy'
-
 	import {
 		type Job,
 		JobService,
@@ -10,7 +8,6 @@
 	} from '$lib/gen'
 	import { workspaceStore } from '$lib/stores'
 	import { Badge, Button } from './common'
-	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import { createEventDispatcher, getContext, untrack } from 'svelte'
 	import type { FlowEditorContext } from './flows/types'
 	import { runFlowPreview } from './flows/utils.svelte'
@@ -18,15 +15,7 @@
 	import SchemaFormWithArgPicker from './SchemaFormWithArgPicker.svelte'
 	import FlowStatusViewer from '../components/FlowStatusViewer.svelte'
 	import FlowProgressBar from './flows/FlowProgressBar.svelte'
-	import {
-		AlertTriangle,
-		ArrowRight,
-		CornerDownLeft,
-		Loader2,
-		Play,
-		RefreshCw,
-		X
-	} from 'lucide-svelte'
+	import { AlertTriangle, CornerDownLeft, Loader2, Play, RefreshCw, X } from 'lucide-svelte'
 	import { emptyString, sendUserToast, type StateStore } from '$lib/utils'
 	import { dfs } from './flows/dfs'
 	import { sliceModules } from './flows/flowStateUtils.svelte'
@@ -38,6 +27,7 @@
 	import { getStepHistoryLoaderContext } from './stepHistoryLoader.svelte'
 	import FlowChat from './flows/conversations/FlowChat.svelte'
 	import { stateSnapshot } from '$lib/svelte5Utils.svelte'
+	import FlowRestartButton from './FlowRestartButton.svelte'
 
 	interface Props {
 		previewMode: 'upTo' | 'whole'
@@ -333,98 +323,19 @@
 			{:else}
 				<div class="grow justify-center flex flex-row gap-2">
 					{#if jobId !== undefined && selectedJobStep !== undefined && selectedJobStepIsTopLevel}
-						{#if selectedJobStepType == 'single'}
-							<Button
-								unifiedSize="md"
-								variant="default"
-								title={`Re-start this flow from step ${selectedJobStep} (included).`}
-								on:click={() => {
-									runPreview(previewArgs.val, {
-										flow_job_id: jobId,
-										step_id: selectedJobStep,
-										branch_or_iteration_n: 0
-									})
-								}}
-								startIcon={{ icon: Play }}
-							>
-								Re-start from
-								<Badge baseClass="ml-1" color="indigo">
-									{selectedJobStep}
-								</Badge>
-							</Button>
-						{:else}
-							<Popover
-								floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}
-								contentClasses="p-4"
-							>
-								{#snippet button()}
-									<Button
-										title={`Re-start this flow from step ${selectedJobStep} (included).`}
-										variant="default"
-										startIcon={{ icon: RefreshCw }}
-										on:click={() => {
-											runPreview(previewArgs.val, {
-												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>
-								{/snippet}
-								{#snippet content()}
-									<label class="block text-primary p-4">
-										<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"
-													onclick={stopPropagation(() => {})}
-												/>
-											{:else}
-												<select
-													bind:value={branchOrIterationN}
-													class="!w-32 grow"
-													onclick={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.val, {
-														flow_job_id: jobId,
-														step_id: selectedJobStep,
-														branch_or_iteration_n: branchOrIterationN
-													})
-												}}
-											>
-												<ArrowRight size={18} />
-											</Button>
-										</div>
-									</label>
-								{/snippet}
-							</Popover>
-						{/if}
+						<FlowRestartButton
+							{jobId}
+							{selectedJobStep}
+							{selectedJobStepType}
+							{restartBranchNames}
+							onRestart={(stepId, branchOrIterationN) => {
+								runPreview(previewArgs.val, {
+									flow_job_id: jobId,
+									step_id: stepId,
+									branch_or_iteration_n: branchOrIterationN
+								})
+							}}
+						/>
 					{/if}
 					{#if !flowStore.val.value?.chat_input_enabled}
 						<Button
diff --git a/frontend/src/lib/components/FlowRestartButton.svelte b/frontend/src/lib/components/FlowRestartButton.svelte
new file mode 100644
index 0000000000..8b387b3c11
--- /dev/null
+++ b/frontend/src/lib/components/FlowRestartButton.svelte
@@ -0,0 +1,218 @@
+<script lang="ts">
+	import { Badge, Button } from './common'
+	import Popover from './meltComponents/Popover.svelte'
+	import { Play, RefreshCw } from 'lucide-svelte'
+	import { FlowService, JobService, type FlowVersion } from '$lib/gen'
+	import { workspaceStore } from '$lib/stores'
+	import { emptyString, sendUserToast } from '$lib/utils'
+	import { goto } from '$lib/navigation'
+
+	interface Props {
+		jobId: string
+		selectedJobStep: string
+		selectedJobStepType: 'single' | 'forloop' | 'branchall'
+		restartBranchNames?: [number, string][]
+		flowPath?: string
+		disabled?: boolean
+		enterpriseOnly?: boolean
+		variant?: 'default' | 'accent'
+		unifiedSize?: 'xs' | 'sm' | 'md' | 'lg'
+		onRestart?: (stepId: string, branchOrIterationN: number, flowVersion?: number) => void
+	}
+
+	let {
+		jobId,
+		selectedJobStep,
+		selectedJobStepType,
+		restartBranchNames = [],
+		flowPath = undefined,
+		disabled = false,
+		enterpriseOnly = false,
+		variant = 'default',
+		unifiedSize = 'md',
+		onRestart
+	}: Props = $props()
+
+	let branchOrIterationN = $state(0)
+	let selectedVersionMode: 'run' | 'custom' = $state('run')
+	let customFlowVersion: number | undefined = $state(undefined)
+	let flowVersions: Array<FlowVersion> = $state([])
+	let loadingVersions: boolean = $state(false)
+
+	async function restartFlow(stepId: string, branchOrIterationN: number, flowVersion?: number) {
+		let run = await JobService.restartFlowAtStep({
+			workspace: $workspaceStore!,
+			id: jobId,
+			requestBody: {
+				step_id: stepId,
+				branch_or_iteration_n: branchOrIterationN,
+				flow_version: flowVersion
+			}
+		})
+		await goto('/run/' + run + '?workspace=' + $workspaceStore)
+	}
+
+	async function loadFlowVersions() {
+		if (!flowPath || loadingVersions) return
+		loadingVersions = true
+		try {
+			flowVersions = await FlowService.getFlowHistory({
+				workspace: $workspaceStore!,
+				path: flowPath
+			})
+			if (flowVersions.length > 0 && customFlowVersion === undefined) {
+				customFlowVersion = flowVersions[0].id
+			}
+		} catch (e) {
+			sendUserToast('Failed to load flow versions', true)
+		} finally {
+			loadingVersions = false
+		}
+	}
+
+	function getFlowVersionForRestart(): number | undefined {
+		if (selectedVersionMode === 'run') {
+			return undefined // use run version
+		} else if (selectedVersionMode === 'custom') {
+			return customFlowVersion
+		}
+		return undefined
+	}
+
+	function handleRestart() {
+		if (onRestart) {
+			onRestart(selectedJobStep, branchOrIterationN, getFlowVersionForRestart())
+		} else {
+			const flowVersion = getFlowVersionForRestart()
+			restartFlow(selectedJobStep, branchOrIterationN, flowVersion)
+		}
+	}
+</script>
+
+{#snippet flowVersionSelector()}
+	<label>
+		<div class="pb-1 text-xs font-semibold text-emphasis">Flow version</div>
+		<div class="flex flex-col gap-2">
+			<select
+				bind:value={selectedVersionMode}
+				class="grow"
+				onchange={() => {
+					if (selectedVersionMode === 'custom' && flowVersions.length === 0) {
+						loadFlowVersions()
+					}
+				}}
+			>
+				<option value="run">Run version</option>
+				<option value="custom">Specific version</option>
+			</select>
+
+			{#if selectedVersionMode === 'custom'}
+				{#if loadingVersions}
+					<div class="text-xs text-secondary">Loading versions...</div>
+				{:else if flowVersions.length > 0}
+					<select bind:value={customFlowVersion} class="grow text-xs">
+						{#each flowVersions as version}
+							<option value={version.id}>
+								{#if emptyString(version.deployment_msg)}Version {version.id}{:else}{version.deployment_msg}{/if}
+								- {new Date(version.created_at).toLocaleString()}
+							</option>
+						{/each}
+					</select>
+				{:else}
+					<div class="text-xs text-tertiary">No versions available</div>
+				{/if}
+			{/if}
+		</div>
+	</label>
+{/snippet}
+{#snippet singleRestartButton()}
+	<Button
+		title={`Re-start this flow from step ${selectedJobStep} (included).${enterpriseOnly ? ' This is a feature only available in enterprise edition.' : ''}`}
+		{variant}
+		{unifiedSize}
+		{disabled}
+		startIcon={{ icon: Play }}
+		nonCaptureEvent={!!flowPath}
+		onClick={() => {
+			if (!flowPath) {
+				handleRestart()
+			}
+		}}
+	>
+		Re-start from
+		<Badge baseClass="ml-1" color="indigo">
+			{selectedJobStep}
+		</Badge>
+		{#if enterpriseOnly && disabled}
+			(EE)
+		{/if}
+	</Button>
+{/snippet}
+
+{#if selectedJobStepType === 'single'}
+	{#if !flowPath}
+		{@render singleRestartButton()}
+	{:else}
+		<Popover
+			floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}
+			disablePopup={!flowPath}
+		>
+			{#snippet trigger()}
+				{@render singleRestartButton()}
+			{/snippet}
+			{#snippet content()}
+				<div class="flex flex-col gap-4 text-primary p-4 min-w-64">
+					{@render flowVersionSelector()}
+
+					<Button variant="accent" onClick={handleRestart}>Restart</Button>
+				</div>
+			{/snippet}
+		</Popover>
+	{/if}
+{:else}
+	<Popover floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}>
+		{#snippet trigger()}
+			<Button
+				title={`Re-start this flow from step ${selectedJobStep} (included).${enterpriseOnly ? ' This is a feature only available in enterprise edition.' : ''}`}
+				{variant}
+				{unifiedSize}
+				{disabled}
+				startIcon={{ icon: RefreshCw }}
+				nonCaptureEvent={true}
+			>
+				Re-start from
+				<Badge baseClass="ml-1" color="indigo">
+					{selectedJobStep}
+				</Badge>
+				{#if enterpriseOnly && disabled}
+					(EE)
+				{/if}
+			</Button>
+		{/snippet}
+		{#snippet content()}
+			<div class="flex flex-col gap-4 text-primary p-4 min-w-64">
+				<label>
+					<div class="pb-1 text-xs font-semibold text-emphasis"
+						>{selectedJobStepType == 'forloop' ? 'From iteration #' : 'From branch'}</div
+					>
+					<div class="flex w-full gap-2">
+						{#if selectedJobStepType === 'forloop'}
+							<input type="number" min="0" bind:value={branchOrIterationN} class="!w-32 grow" />
+						{:else}
+							<select bind:value={branchOrIterationN} class="!w-32 grow">
+								{#each restartBranchNames as [branchIdx, branchName]}
+									<option value={branchIdx}>{branchName}</option>
+								{/each}
+							</select>
+						{/if}
+					</div>
+				</label>
+
+				{#if flowPath}
+					{@render flowVersionSelector()}
+				{/if}
+				<Button variant="accent" onClick={handleRestart}>Restart</Button>
+			</div>
+		{/snippet}
+	</Popover>
+{/if}
diff --git a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
index 0418cb9a08..be8639db74 100644
--- a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
+++ b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
@@ -1,6 +1,4 @@
 <script lang="ts">
-	import { stopPropagation } from 'svelte/legacy'
-
 	import { base } from '$lib/base'
 	import {
 		JobService,
@@ -28,7 +26,6 @@
 
 	import {
 		Activity,
-		ArrowRight,
 		Calendar,
 		CheckCircle2,
 		Circle,
@@ -60,13 +57,11 @@
 	import JobLoader from '$lib/components/JobLoader.svelte'
 	import LogViewer from '$lib/components/LogViewer.svelte'
 	import { ActionRow, Button, Skeleton, Tab, Alert, DrawerContent } from '$lib/components/common'
-	import Popover from '$lib/components/meltComponents/Popover.svelte'
 	import FlowMetadata from '$lib/components/FlowMetadata.svelte'
 	import JobArgs from '$lib/components/JobArgs.svelte'
 	import FlowProgressBar from '$lib/components/flows/FlowProgressBar.svelte'
 	import JobProgressBar from '$lib/components/jobs/JobProgressBar.svelte'
 	import Tabs from '$lib/components/common/tabs/TabsV2.svelte'
-	import Badge from '$lib/components/common/badge/Badge.svelte'
 	import { goto } from '$lib/navigation'
 	import { sendUserToast } from '$lib/toast'
 	import { forLater } from '$lib/forLater'
@@ -95,6 +90,7 @@
 	import { page } from '$app/state'
 	import RunBadges from '$lib/components/runs/RunBadges.svelte'
 	import { twMerge } from 'tailwind-merge'
+	import FlowRestartButton from '$lib/components/FlowRestartButton.svelte'
 	let job: (Job & { result?: any; result_stream?: string }) | undefined = $state()
 	let jobUpdateLastFetch: Date | undefined = $state()
 
@@ -103,7 +99,6 @@
 
 	let viewTab: 'result' | 'logs' | 'code' | 'stats' | 'assets' = $state('result')
 	let selectedJobStep: string | undefined = $state(undefined)
-	let branchOrIterationN: number = $state(0)
 
 	let selectedJobStepIsTopLevel: boolean | undefined = $state(undefined)
 	let selectedJobStepType: 'single' | 'forloop' | 'branchall' = $state('single')
@@ -153,24 +148,6 @@
 		}
 	}
 
-	async function restartFlow(
-		id: string | undefined,
-		stepId: string | undefined,
-		branchOrIterationN: number
-	) {
-		if (id === undefined || stepId === undefined) {
-			return
-		}
-		let run = await JobService.restartFlowAtStep({
-			workspace: $workspaceStore!,
-			id,
-			stepId,
-			branchOrIterationN,
-			requestBody: {}
-		})
-		await goto('/run/' + run + '?workspace=' + $workspaceStore)
-	}
-
 	// If we get results, focus on that tab. Else, focus on logs
 	function initView(): void {
 		if (job && (job.result || job.result_stream)) {
@@ -607,95 +584,16 @@
 					startIcon={{ icon: Calendar }}>Edit schedule</Button
 				>
 			{/if}
-			{#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="default"
-						unifiedSize="md"
-						disabled={!$enterpriseLicense}
-						on:click|once={() => {
-							restartFlow(job?.id, selectedJobStep, 0)
-						}}
-						startIcon={{ icon: RefreshCw }}
-					>
-						Re-start from
-						<Badge baseClass="ml-1" color="indigo">
-							{selectedJobStep}
-						</Badge>
-						{#if !$enterpriseLicense}
-							(EE)
-						{/if}
-					</Button>
-				{:else}
-					<Popover
-						floatingConfig={{ strategy: 'absolute', placement: 'bottom-start' }}
-						contentClasses="p-4"
-					>
-						{#snippet trigger()}
-							<Button
-								title={`Re-start this flow from step ${selectedJobStep} (included). ${
-									!$enterpriseLicense
-										? ' This is a feature only available in enterprise edition.'
-										: ''
-								}`}
-								variant="default"
-								unifiedSize="md"
-								disabled={!$enterpriseLicense}
-								startIcon={{ icon: RefreshCw }}
-								nonCaptureEvent={true}
-							>
-								Re-start from
-								<Badge baseClass="ml-1" color="indigo">
-									{selectedJobStep}
-								</Badge>
-							</Button>
-						{/snippet}
-						{#snippet content()}
-							<label class="block text-primary">
-								<div class="pb-1 text-xs font-semibold text-emphasis"
-									>{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"
-											onclick={stopPropagation(() => {})}
-										/>
-									{:else}
-										<select
-											bind:value={branchOrIterationN}
-											class="!w-32 grow"
-											onclick={stopPropagation(() => {})}
-										>
-											{#each restartBranchNames as [branchIdx, branchName]}
-												<option value={branchIdx}>{branchName}</option>
-											{/each}
-										</select>
-									{/if}
-
-									<Button
-										unifiedSize="md"
-										variant="accent"
-										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>
-						{/snippet}
-					</Popover>
-				{/if}
+			{#if job?.type === 'CompletedJob' && job?.job_kind === 'flow' && selectedJobStep !== undefined && selectedJobStepIsTopLevel && job.id}
+				<FlowRestartButton
+					jobId={job.id}
+					{selectedJobStep}
+					{selectedJobStepType}
+					{restartBranchNames}
+					flowPath={job.script_path}
+					disabled={!$enterpriseLicense}
+					enterpriseOnly={!$enterpriseLicense}
+				/>
 			{/if}
 			{#if job?.job_kind === 'script' || job?.job_kind === 'flow'}
 				<Button