From ad9f1fa4541f2eefb0b013bac42cd626af424852 Mon Sep 17 00:00:00 2001
From: hugocasa <hugo@casademont.ch>
Date: Fri, 1 May 2026 20:03:40 +0200
Subject: [PATCH] fix: nested-restart iteration count for step-id collisions
 across subflow boundaries (#9003)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: support restart from steps inside BranchOne, ForLoop, Subflow

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: preserve original job kind in nested restart, support expanded subflow steps

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: read selected iteration from graph state for nested ForLoop restart

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* feat: iteration selectors per ForLoop in restart popup, more nested restart tests

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* refactor: extract useNestedRestartState composable

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test: cover deployed-subflow + FlowDependencies path in nested restart

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore: update sqlx prepare cache

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: detect BranchOne/ForLoop ancestors inside expanded subflows for nested restart

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: hide restart button for non-restartable steps (parallel containers, untaken branches)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: address review feedback on nested restart PR

- preview FlowRestartButton: hide nested case (chain UUIDs aren't resolvable in
  preview path; users can use the run page for nested restart instead)
- branchOneAncestorMatchesOriginal: be permissive when status isn't reachable
  (don't hide the button for BranchOnes nested deeper than top-level)
- worker_flow.rs: apply nested_restart_payload swap on the is_simple ForLoop
  fast path too, so simple iterations don't bypass restart spawn interception
- FlowStatusViewer: reset expandedSubflows cache on jobId change; drop
  $bindable({}) banned pattern for the new prop
- API resolver: validate the leaf step exists before returning (fail-fast)
- doc fix: branch_or_iteration_n is 0-based, not 1-based
- selectedJobStepIsTopLevel reset on early-return in composable
- comment iterationCounts collision caveat
- new HTTP-level integration tests covering the API endpoint contract:
  happy path (top-level + nested), unknown step, out-of-range iteration,
  parallel-loop rejection

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* revert: remove unreachable nested-restart swap on is_simple ForLoop fast path

The swap is unreachable in valid flows: `is_simple_modules` requires the body
to be a single `script` / `rawscript` / `flowscript` (per `FlowModule::is_simple`),
none of which spawn flow-kind children. Any nested-restart chain targeting a
leaf inside such an iteration is rejected by the API at leaf validation. Even
if a chain reached the worker via `JobPayload::RawFlow.restarted_from`, the
resulting `RestartedFlow` would fail to push (script kind isn't a flow kind).

Replaced the swap with an explanatory comment so the next reader knows why
the symmetry with the non-simple path was deliberately not added.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: handle undefined expandedSubflows + tighten branchOne match check

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix: iteration count for restart popup picks up wrong loop on step-id collision

When a top-level ForLoop step shares an id with a loop nested inside an expanded
subflow (e.g. parent has step `e` with 4 iterations and the subflow at step `h`
also has step `e` with 1 iteration), the popup's iteration `<select>` rendered
the subflow's count instead of the parent's because `iterationCounts` was
double-indexed by both the prefixed graph key and the bare leaf segment after
the last colon — last writer wins, and the bare key collided.

- `iterationCounts` now indexed only by the full graph key. The selected-step
  iteration picker still works because at the top level the step id IS the
  graph key.
- New `nestedPathIterationCounts` keyed by the popup's field-key
  ('top' / 'inner-N'). Populated by the composable during path construction so
  each entry uses the correct prefix-aware graph key, regardless of whether the
  ancestor lives in the parent flow or inside an expanded subflow.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .../lib/components/FlowRestartButton.svelte   | 14 +++-
 .../useNestedRestartState.svelte.ts           | 66 +++++++++++++------
 .../(root)/(logged)/run/[...run]/+page.svelte |  1 +
 3 files changed, 59 insertions(+), 22 deletions(-)

diff --git a/frontend/src/lib/components/FlowRestartButton.svelte b/frontend/src/lib/components/FlowRestartButton.svelte
index 7dc8f7b204..d470c51b54 100644
--- a/frontend/src/lib/components/FlowRestartButton.svelte
+++ b/frontend/src/lib/components/FlowRestartButton.svelte
@@ -40,9 +40,18 @@
 		 * Map from ForLoop step id to the number of iterations that ran in the
 		 * original execution (i.e. `flow_jobs.length`). When provided for a step,
 		 * the popup renders a `<select>` of `0..count-1` instead of a free-form
-		 * number input — same surface as the graph's iteration tabs.
+		 * number input — same surface as the graph's iteration tabs. Used for the
+		 * SELECTED step's iteration picker when it is itself a top-level ForLoop.
 		 */
 		iterationCounts?: Record<string, number>
+		/**
+		 * Iteration counts for nested-path entries, keyed by the popup's field-key
+		 * (`'top'` for the outer container, `'inner-N'` for nested ancestors).
+		 * Populated by the composable so each entry uses the *correct* graph-state
+		 * key (prefixed for in-subflow ancestors), avoiding collisions when the
+		 * same step id appears at multiple nesting levels.
+		 */
+		nestedPathIterationCounts?: Record<string, number>
 		/** Called when flow is restarted. If not provided, will navigate to the new run using goto (requires SvelteKit) */
 		onRestart?: (stepId: string, branchOrIterationN: number, flowVersion?: number) => void
 		/** Called when flow restart completes with the new job ID. Used for navigation in non-SvelteKit contexts */
@@ -65,6 +74,7 @@
 		nestedTopBranchOrIterationN = undefined,
 		presetIterationN = undefined,
 		iterationCounts = undefined,
+		nestedPathIterationCounts = undefined,
 		onRestart,
 		onRestartComplete
 	}: Props = $props()
@@ -303,7 +313,7 @@
 				{/if}
 
 				{#each iterationFields as field (field.key)}
-					{@const count = iterationCounts?.[field.label] ?? 0}
+					{@const count = nestedPathIterationCounts?.[field.key] ?? 0}
 					<label>
 						<div class="pb-1 text-xs font-semibold text-emphasis">
 							From iteration # of <code class="text-xs">{field.label}</code>
diff --git a/frontend/src/lib/components/useNestedRestartState.svelte.ts b/frontend/src/lib/components/useNestedRestartState.svelte.ts
index 67e3a9174a..38ea467cf0 100644
--- a/frontend/src/lib/components/useNestedRestartState.svelte.ts
+++ b/frontend/src/lib/components/useNestedRestartState.svelte.ts
@@ -73,6 +73,11 @@ export function useNestedRestartState(opts: {
 	let nestedRestartTopBranchOrIterationN: number | undefined = $state(undefined)
 	let nestedRestartPath: NestedRestartStep[] | undefined = $state(undefined)
 	let nestedRestartSupported = $state(false)
+	// Iteration counts keyed by the popup's field-key ('top' or 'inner-N').
+	// Built during path construction so each entry uses the *correct* graph key
+	// for its level (prefixed for in-subflow ancestors). Avoids the collision
+	// that would happen if we double-indexed `iterationCounts` by bare step_id.
+	let nestedPathIterationCounts: Record<string, number> = $state({})
 
 	$effect(() => {
 		const selectedJobStep = opts.selectedJobStep()
@@ -86,6 +91,7 @@ export function useNestedRestartState(opts: {
 		nestedRestartSupported = false
 		restartBranchNames = []
 		selectedJobStepIsTopLevel = undefined
+		nestedPathIterationCounts = {}
 
 		if (selectedJobStep === undefined || job?.flow_status?.modules === undefined) {
 			return
@@ -151,24 +157,40 @@ export function useNestedRestartState(opts: {
 						if (!blocked) {
 							// Inside the subflow, the graph state for ForLoop ancestors
 							// is keyed by the prefixed graph id. Build that key here so
-							// `selectedForloopIndex` lookups land on the right entry.
+							// `selectedForloopIndex` and `flow_jobs.length` lookups land
+							// on the right entry (avoiding collisions with same-step-id
+							// loops in the parent flow — e.g. parent has `e` with 4 iters
+							// and the subflow also has `e` with 1 iter).
 							const subflowPrefix = 'subflow:' + subflowParse.subflowSteps.join(':') + ':'
 							const iterationFor = (stepId: string): number =>
 								graphModuleStates[subflowPrefix + stepId]?.selectedForloopIndex ?? 0
-							for (const a of path.ancestors) {
+							const iterCountFor = (stepId: string): number =>
+								graphModuleStates[subflowPrefix + stepId]?.flow_jobs?.length ?? 0
+							const counts: Record<string, number> = {}
+							// Top-level (parent) container is the subflow step `h` —
+							// not a ForLoop, no count.
+							for (let i = 0; i < path.ancestors.length; i++) {
+								const a = path.ancestors[i]
 								const entry: NestedRestartStep = { step_id: a.stepId }
 								if (a.type === 'forloopflow') {
 									entry.branch_or_iteration_n = iterationFor(a.stepId)
+									// Path key: each subflow boundary already contributes an
+									// inner-N entry above (subflow steps without iterations).
+									// Match the index that FlowRestartButton will use when
+									// rendering iteration fields.
+									counts[`inner-${innerPath.length}`] = iterCountFor(a.stepId)
 								}
 								innerPath.push(entry)
 							}
 							const leafEntry: NestedRestartStep = { step_id: subflowParse.leaf }
 							if (path.target.value.type === 'forloopflow') {
 								leafEntry.branch_or_iteration_n = iterationFor(subflowParse.leaf)
+								counts[`inner-${innerPath.length}`] = iterCountFor(subflowParse.leaf)
 							}
 							innerPath.push(leafEntry)
 							nestedRestartTopStepId = top.id
 							nestedRestartPath = innerPath
+							nestedPathIterationCounts = counts
 							nestedRestartSupported = true
 							return
 						}
@@ -209,13 +231,17 @@ export function useNestedRestartState(opts: {
 		// for confirmation/editing before submit, so we never silently send a guess.
 		const iterationFor = (stepId: string): number =>
 			graphModuleStates[stepId]?.selectedForloopIndex ?? 0
+		const iterCountFor = (stepId: string): number =>
+			graphModuleStates[stepId]?.flow_jobs?.length ?? 0
 		const top = path.ancestors[0]
 		const inner = path.ancestors.slice(1)
 		const innerPath: NestedRestartStep[] = []
+		const counts: Record<string, number> = {}
 		for (const a of inner) {
 			const entry: NestedRestartStep = { step_id: a.stepId }
 			if (a.type === 'forloopflow') {
 				entry.branch_or_iteration_n = iterationFor(a.stepId)
+				counts[`inner-${innerPath.length}`] = iterCountFor(a.stepId)
 			}
 			innerPath.push(entry)
 		}
@@ -224,13 +250,17 @@ export function useNestedRestartState(opts: {
 		const leafEntry: NestedRestartStep = { step_id: selectedJobStep }
 		if (path.target.value.type === 'forloopflow') {
 			leafEntry.branch_or_iteration_n = iterationFor(selectedJobStep)
+			counts[`inner-${innerPath.length}`] = iterCountFor(selectedJobStep)
 		}
 		innerPath.push(leafEntry)
 
 		nestedRestartTopStepId = top.stepId
-		nestedRestartTopBranchOrIterationN =
-			top.type === 'forloopflow' ? iterationFor(top.stepId) : undefined
+		if (top.type === 'forloopflow') {
+			nestedRestartTopBranchOrIterationN = iterationFor(top.stepId)
+			counts['top'] = iterCountFor(top.stepId)
+		}
 		nestedRestartPath = innerPath
+		nestedPathIterationCounts = counts
 		nestedRestartSupported = true
 	})
 
@@ -262,28 +292,21 @@ export function useNestedRestartState(opts: {
 		return true
 	})
 
-	// Iteration counts indexed by the step id used in `nestedRestartPath`. For
-	// regular nesting that's the bare `step_id`; for steps inside expanded
-	// subflows we ALSO index by the leaf segment of the prefixed graph id (e.g.
-	// `subflow:h:loop_1` → also indexed at `loop_1`) so the popup can find the
-	// count regardless of which path produced the field.
-	//
-	// Known caveat: if the same step id appears at multiple levels (e.g. a
-	// top-level `loop_1` AND a `loop_1` inside a subflow), the last-writer-wins
-	// here may show the wrong option count. The submitted value is still 0-based
-	// and the backend validates it, so a wrong-count display only affects how
-	// many options the `<select>` shows. Step ids are globally unique within a
-	// single flow value, so the collision only happens across subflow boundaries.
+	// Iteration counts indexed by the graph module-state key (i.e. the prefixed
+	// `subflow:...:<step_id>` for in-subflow loops, or the bare `step_id` for
+	// top-level loops). Used by the popup for the SELECTED step's iteration
+	// picker — that step is always at the unprefixed top level (the run page's
+	// graph-state key matches the bare step_id), so the lookup is unambiguous.
+	// For nested-path iteration fields, see `nestedPathIterationCounts` instead,
+	// which is keyed by the popup's field-key ('top' / 'inner-N') and pulled
+	// from the path-aware graph key — avoiding collisions when the same step
+	// id appears at both the parent flow and inside a subflow.
 	const iterationCounts = $derived.by((): Record<string, number> => {
 		const out: Record<string, number> = {}
 		for (const [id, state] of Object.entries(opts.graphModuleStates())) {
 			const n = state.flow_jobs?.length
 			if (typeof n !== 'number' || n <= 0) continue
 			out[id] = n
-			const lastColon = id.lastIndexOf(':')
-			if (lastColon >= 0) {
-				out[id.slice(lastColon + 1)] = n
-			}
 		}
 		return out
 	})
@@ -318,6 +341,9 @@ export function useNestedRestartState(opts: {
 		},
 		get iterationCounts() {
 			return iterationCounts
+		},
+		get nestedPathIterationCounts() {
+			return nestedPathIterationCounts
 		}
 	}
 }
diff --git a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
index e8cd9b7287..69e9ff0b96 100644
--- a/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
+++ b/frontend/src/routes/(root)/(logged)/run/[...run]/+page.svelte
@@ -633,6 +633,7 @@
 					nestedTopBranchOrIterationN={restart.nestedRestartTopBranchOrIterationN}
 					presetIterationN={restart.topLevelLoopIteration}
 					iterationCounts={restart.iterationCounts}
+					nestedPathIterationCounts={restart.nestedPathIterationCounts}
 					onRestartComplete={(newJobId) => {
 						goto('/run/' + newJobId + '?workspace=' + $workspaceStore)
 					}}