@servicenow/sdk-build-plugins 4.6.1 → 4.7.1

package/src/flow/utils/flow-stage-processor.ts

@@ -0,0 +1,831 @@
+import {
+    CallExpressionShape,
+    DATA_HELPER_NAMES,
+    type Diagnostics,
+    DurationShape,
+    type Factory,
+    ObjectShape,
+    PropertyAccessShape,
+    Record,
+    Shape,
+    type Source,
+    VariableStatementShape,
+} from '@servicenow/sdk-build-core'
+import { ArrowFunctionShape } from '../../arrow-function-plugin'
+import { NowIdShape } from '../../now-id-plugin'
+import {
+    FLOW_STAGE_CALLEE,
+    FLOW_STAGE_TABLE,
+    STAGE_API_NAME,
+    STAGE_DEFAULT_DURATION,
+    SUBFLOW_INSTANCE_API_NAME,
+} from './flow-constants'
+import { FLOW_LOGIC } from '../flow-logic/flow-logic-constants'
+import { sysIdToUuid } from './utils'
+
+// ---------------------------------------------------------------------------
+// Duration conversion (Fluent → XML)
+// ---------------------------------------------------------------------------
+
+/**
+ * Converts a `Duration({...})` CallExpressionShape from the Fluent API into the
+ * `YYYY-MM-DD HH:MM:SS` string stored in `sys_hub_flow_stage.duration`.
+ *
+ * `Duration()` is an identity function — ObjectShape.get() resolves it with
+ * resolve=true by default, but CallExpressionShape.resolve() returns itself
+ * (can't evaluate the body). We therefore read with resolve=false to get the
+ * raw CallExpressionShape, then delegate to DurationShape for formatting.
+ * Falls back to zero duration when absent or not a Duration() call.
+ */
+export function stageDurationToString(stageConfig: ObjectShape): string {
+    const raw = stageConfig.get('duration', false)
+    const ce = raw.if(CallExpressionShape) ? raw.as(CallExpressionShape) : undefined
+    if (ce && ce.getCallee() === DATA_HELPER_NAMES.DURATION) {
+        return new DurationShape({
+            source: raw.getSource(),
+            value: ce.getArgument(0).asObject(),
+        })
+            .toString()
+            .getValue()
+    }
+    return STAGE_DEFAULT_DURATION
+}
+
+// camelCase API key → snake_case DB key for FlowStageStates fields
+const STATE_KEY_MAP: { [k: string]: string } = {
+    pending: 'pending',
+    inProgress: 'in_progress',
+    complete: 'complete',
+    error: 'error',
+    skipped: 'skipped',
+}
+
+const DEFAULT_STATES: { [k: string]: string } = {
+    in_progress: 'In progress',
+    complete: 'Completed',
+    error: 'Error',
+    pending: 'Pending - has not started',
+    skipped: 'Skipped',
+}
+
+// ---------------------------------------------------------------------------
+// Phase 1 — Header stage pre-creation
+// ---------------------------------------------------------------------------
+
+/**
+ * Reads the `stages: {}` config object from the Flow/Subflow first argument
+ * and creates `sys_hub_flow_stage` records for each declared stage.
+ *
+ * Records are created WITHOUT `component_indexes` — those are computed by
+ * `finalizeStageRecords` after the body loop has determined which actions
+ * each stage covers.
+ *
+ * `stage_id` is derived from the record's own coalesce-based sys_id (which the
+ * factory makes stable from `flow + value`), so it is deterministic across
+ * builds without requiring a separate hash.
+ *
+ * @returns Map keyed by the stages object property name (e.g. `'triage'`).
+ */
+export async function processHeaderStages(
+    flowConfiguration: ObjectShape,
+    flowDefinitionRecord: Record,
+    factory: Factory
+): Promise<Map<string, Record>> {
+    const headerStageRecords = new Map<string, Record>()
+    if (!flowConfiguration.has('stages')) {
+        return headerStageRecords
+    }
+
+    const stagesObj = flowConfiguration.get('stages').asObject()
+
+    for (const [key, stageVal] of stagesObj.entries()) {
+        const stageCE = stageVal.if(CallExpressionShape) ? stageVal.as(CallExpressionShape) : undefined
+        if (!stageCE) {
+            continue
+        }
+
+        const stageConfig = stageCE.getArgument(0).asObject()
+        const stageValue = stageConfig.get('value').ifString()?.getValue() ?? ''
+
+        const statesArg = stageConfig.get('states').ifObject()?.asObject()
+        let statesJson: { [k: string]: string }
+        if (statesArg) {
+            statesJson = {}
+            for (const [apiKey, dbKey] of Object.entries(STATE_KEY_MAP)) {
+                const val = statesArg.get(apiKey).ifString()?.getValue()
+                if (val) {
+                    statesJson[dbKey] = val
+                }
+            }
+        } else {
+            statesJson = { ...DEFAULT_STATES }
+        }
+
+        const stageRecord = await factory.createRecord({
+            source: stageCE,
+            table: FLOW_STAGE_TABLE,
+            properties: {
+                flow: flowDefinitionRecord.getId().getValue(),
+                label: stageConfig.get('label').ifString()?.getValue() ?? '',
+                value: stageValue,
+                duration: stageDurationToString(stageConfig),
+                always_show: stageConfig.get('alwaysShow').ifBoolean()?.getValue() ?? false,
+                states: JSON.stringify(statesJson),
+                type: 'standard',
+                ancestor_array_position: 0,
+                ancestor_component_id: '',
+                ancestor_stage_id: '',
+                ancestral_if_else_logic: '',
+            },
+        })
+
+        // stage_id reuses the record's coalesce-based sys_id — stable across builds
+        // because sys_hub_flow_stage coalesces on ['flow', 'value'].
+        headerStageRecords.set(
+            key,
+            stageRecord.merge({ stage_id: sysIdToUuid(String(stageRecord.getId().getValue())) })
+        )
+    }
+
+    return headerStageRecords
+}
+
+// ---------------------------------------------------------------------------
+// Phase 2 — Subflow stage creation
+// ---------------------------------------------------------------------------
+
+/**
+ * Scans `relatedRecords` for subflow instances with `show_stages=true` and
+ * creates a separate `sys_hub_flow_stage` record with `type:'subflow'` for each.
+ *
+ * These differ from developer-declared stages:
+ *  - `type`             = `'subflow'` (not `'standard'`)
+ *  - `value`            = the subflow instance's `ui_id` (UUID)
+ *  - `label`            = the subflow definition's name (resolved via `subflowDefNames` map)
+ *  - `component_indexes`= empty — the platform links via `value` (the UUID)
+ *
+ * @param pendingStagesLength  Number of developer-declared stages — subflow stages
+ *                             are ordered after them.
+ * @param subflowDefNames      Map of subflow definition sys_id → name, built by the
+ *                             caller from the original shapes.
+ */
+export async function createSubflowStageRecords(
+    relatedRecords: Record[],
+    pendingStagesLength: number,
+    flowDefinitionRecord: Record,
+    factory: Factory,
+    subflowDefNames: Map<string, string>
+): Promise<Record[]> {
+    const subflowStageRecords: Record[] = []
+    let subflowStageOrder = pendingStagesLength
+
+    for (const rec of relatedRecords) {
+        if (rec.getTable() !== 'sys_hub_sub_flow_instance_v2') {
+            continue
+        }
+        const showStages = rec.get('show_stages')?.getValue()
+        if (showStages !== true && showStages !== 'true') {
+            continue
+        }
+        const uiId = String(rec.get('ui_id')?.getValue() ?? '')
+        if (!uiId) {
+            continue
+        }
+        const defSysId = String(rec.get('subflow')?.getValue() ?? '')
+        const label = subflowDefNames.get(defSysId) ?? ''
+
+        const stageRecord = await factory.createRecord({
+            source: flowDefinitionRecord,
+            table: FLOW_STAGE_TABLE,
+            properties: {
+                flow: flowDefinitionRecord.getId().getValue(),
+                type: 'subflow',
+                value: uiId,
+                label,
+                always_show: false,
+                states: '{}',
+                component_indexes: '',
+                order: subflowStageOrder,
+                ancestor_array_position: 0,
+                ancestor_component_id: '',
+                ancestor_stage_id: '',
+                ancestral_if_else_logic: '',
+                duration: STAGE_DEFAULT_DURATION,
+            },
+        })
+
+        // stage_id reuses the record's coalesce-based sys_id — stable because
+        // sys_hub_flow_stage coalesces on ['flow', 'value'] and value = uiId.
+        subflowStageRecords.push(stageRecord.merge({ stage_id: sysIdToUuid(String(stageRecord.getId().getValue())) }))
+        subflowStageOrder++
+    }
+
+    return subflowStageRecords
+}
+
+// ---------------------------------------------------------------------------
+// Body scanning — collect stage markers & subflow definition names
+// ---------------------------------------------------------------------------
+
+/**
+ * Extracts the stage key from a `wfa.stage(params.stages.<key>)` call.
+ * Returns `undefined` when the argument is not a PropertyAccessShape.
+ */
+export function extractStageKey(callExpr: CallExpressionShape): string | undefined {
+    const arg = callExpr.getArgument(0, false)
+    if (!arg.if(PropertyAccessShape)) {
+        return undefined
+    }
+    return arg.as(PropertyAccessShape).getLastElement().getName()
+}
+
+/**
+ * Extracts the NowIdKey (`$id` value) from the config object of an instance call.
+ * Checks argument positions 0 and 1 (actions use position 1, flow logic uses 0).
+ */
+function extractNowIdKey(callExpr: CallExpressionShape): string | undefined {
+    for (const idx of [0, 1]) {
+        try {
+            const a = callExpr.getArgument(idx)
+            if (a?.if(ObjectShape)) {
+                const $id = a.as(ObjectShape).get('$id')
+                if ($id?.if(NowIdShape)) {
+                    return String($id.as(NowIdShape).getValue())
+                }
+            }
+        } catch {
+            /* skip */
+        }
+    }
+    return undefined
+}
+
+/** Unwrap a statement to the underlying CallExpressionShape if applicable. */
+export function unwrapCall(stmt: Shape): CallExpressionShape | undefined {
+    const inner = stmt instanceof VariableStatementShape ? (stmt as VariableStatementShape).getInitializer() : stmt
+    return inner instanceof CallExpressionShape ? inner : undefined
+}
+
+/**
+ * Collects subflow definition name from a `wfa.subflow(defRecord, ...)` call
+ * and stores it in the provided map.
+ */
+export function collectSubflowDefName(callExpr: CallExpressionShape, subflowDefNames: Map<string, string>): void {
+    const defArg = callExpr.getArgument(0)
+    if (defArg?.isRecord()) {
+        const defRecord = defArg.as(Record)
+        const defSysId = String(defRecord.getId().getValue())
+        if (!subflowDefNames.has(defSysId)) {
+            subflowDefNames.set(defSysId, String(defRecord.get('name')?.getValue() ?? ''))
+        }
+    }
+}
+
+/**
+ * Resolves nested stage indexes inline during the main instance loop.
+ *
+ * Called when a flow logic instance (if/elseIf/else) is encountered.
+ * Builds a local NowIdKey → order map from the flat records (whose orders
+ * are known: `baseOrder + index`), then recursively scans the body for
+ * `wfa.stage()` calls and pushes resolved entries directly into `pendingStages`.
+ *
+ * Also collects subflow definition names from nested `wfa.subflow()` calls.
+ */
+export function resolveNestedStagesInline(
+    callExpr: CallExpressionShape,
+    flatRecords: Record[],
+    baseOrder: number,
+    headerStageRecords: Map<string, Record>,
+    pendingStages: Array<{ startActionIndex: number; key: string }>,
+    subflowDefNames: Map<string, string>,
+    diagnostics: Diagnostics
+): void {
+    const bodyArg = callExpr.getArgument(1, false)
+    if (!bodyArg.if(ArrowFunctionShape)) {
+        return
+    }
+
+    // Build NowIdKey → order map for these flat records
+    const orderByNowIdKey = new Map<string, number>()
+    flatRecords.forEach((rec, idx) => {
+        const nk = String(rec.getId().getNowIdKey() ?? '')
+        if (nk) {
+            orderByNowIdKey.set(nk, baseOrder + idx)
+        }
+    })
+
+    resolveStagesInBody(
+        bodyArg.as(ArrowFunctionShape),
+        headerStageRecords,
+        orderByNowIdKey,
+        pendingStages,
+        subflowDefNames,
+        diagnostics
+    )
+}
+
+/**
+ * Resolves nested stages inside a `tryCatch(config, { try, catch })` call.
+ * Scans both the `try` and `catch` arrow function bodies for `wfa.stage()` calls.
+ */
+export function resolveNestedStagesForTryCatch(
+    callExpr: CallExpressionShape,
+    flatRecords: Record[],
+    baseOrder: number,
+    headerStageRecords: Map<string, Record>,
+    pendingStages: Array<{ startActionIndex: number; key: string }>,
+    subflowDefNames: Map<string, string>,
+    diagnostics: Diagnostics
+): void {
+    const handlersArg = callExpr.getArgument(1, false)
+    if (!handlersArg.if(ObjectShape)) {
+        return
+    }
+
+    // Build NowIdKey → order map for these flat records
+    const orderByNowIdKey = new Map<string, number>()
+    flatRecords.forEach((rec, idx) => {
+        const nk = String(rec.getId().getNowIdKey() ?? '')
+        if (nk) {
+            orderByNowIdKey.set(nk, baseOrder + idx)
+        }
+    })
+
+    const handlers = handlersArg.as(ObjectShape)
+    const tryHandler = handlers.get('try', false)
+    const catchHandler = handlers.get('catch', false)
+
+    if (tryHandler.if(ArrowFunctionShape)) {
+        resolveStagesInBody(
+            tryHandler.as(ArrowFunctionShape),
+            headerStageRecords,
+            orderByNowIdKey,
+            pendingStages,
+            subflowDefNames,
+            diagnostics
+        )
+    }
+    if (catchHandler.if(ArrowFunctionShape)) {
+        resolveStagesInBody(
+            catchHandler.as(ArrowFunctionShape),
+            headerStageRecords,
+            orderByNowIdKey,
+            pendingStages,
+            subflowDefNames,
+            diagnostics
+        )
+    }
+}
+
+/**
+ * Recursively scans an arrow-function body for:
+ *  - `wfa.stage()` calls: resolves the order of the *next* non-stage instance
+ *    from `orderByNowIdKey` and pushes directly to `pendingStages`.
+ *  - `wfa.subflow()` calls: collects the definition name for showSubflowStage labels.
+ */
+function resolveStagesInBody(
+    body: ArrowFunctionShape,
+    headerStageRecords: Map<string, Record>,
+    orderByNowIdKey: Map<string, number>,
+    pendingStages: Array<{ startActionIndex: number; key: string }>,
+    subflowDefNames: Map<string, string>,
+    diagnostics: Diagnostics
+): void {
+    const stmts = body.getStatements()
+    for (let i = 0; i < stmts.length; i++) {
+        const inner = unwrapCall(stmts[i]!)
+        if (!inner) {
+            continue
+        }
+        const callee = inner.getCallee()
+
+        // Recurse into flow logic bodies (if / elseIf / else)
+        if (callee === FLOW_LOGIC.IF || callee === FLOW_LOGIC.ELSEIF || callee === FLOW_LOGIC.ELSE) {
+            const bodyArg = inner.getArgument(1, false)
+            if (bodyArg.if(ArrowFunctionShape)) {
+                resolveStagesInBody(
+                    bodyArg.as(ArrowFunctionShape),
+                    headerStageRecords,
+                    orderByNowIdKey,
+                    pendingStages,
+                    subflowDefNames,
+                    diagnostics
+                )
+            }
+            continue
+        }
+
+        // Recurse into tryCatch try/catch handler bodies
+        if (callee === FLOW_LOGIC.TRY_CATCH) {
+            const handlersArg = inner.getArgument(1, false)
+            if (handlersArg.if(ObjectShape)) {
+                const handlers = handlersArg.as(ObjectShape)
+                const tryHandler = handlers.get('try', false)
+                const catchHandler = handlers.get('catch', false)
+                if (tryHandler.if(ArrowFunctionShape)) {
+                    resolveStagesInBody(
+                        tryHandler.as(ArrowFunctionShape),
+                        headerStageRecords,
+                        orderByNowIdKey,
+                        pendingStages,
+                        subflowDefNames,
+                        diagnostics
+                    )
+                }
+                if (catchHandler.if(ArrowFunctionShape)) {
+                    resolveStagesInBody(
+                        catchHandler.as(ArrowFunctionShape),
+                        headerStageRecords,
+                        orderByNowIdKey,
+                        pendingStages,
+                        subflowDefNames,
+                        diagnostics
+                    )
+                }
+            }
+            continue
+        }
+
+        // Collect subflow definition names from nested bodies
+        if (callee === SUBFLOW_INSTANCE_API_NAME) {
+            collectSubflowDefName(inner, subflowDefNames)
+        }
+
+        if (callee !== STAGE_API_NAME) {
+            continue
+        }
+
+        // Extract the stage key and resolve the order of the next non-stage instance
+        const key = extractStageKey(inner)
+        if (!key) {
+            diagnostics.error(inner, 'wfa.stage() argument must be a stage reference like params.stages.<key>.')
+            continue
+        }
+        if (!headerStageRecords.has(key)) {
+            diagnostics.error(
+                inner,
+                `wfa.stage() references stage '${key}' which is not declared in the stages config.`
+            )
+            continue
+        }
+
+        for (let j = i + 1; j < stmts.length; j++) {
+            const next = unwrapCall(stmts[j]!)
+            if (!next || next.getCallee() === STAGE_API_NAME) {
+                continue
+            }
+            const nowIdKey = extractNowIdKey(next)
+            if (nowIdKey) {
+                const recOrder = orderByNowIdKey.get(nowIdKey)
+                if (recOrder != null) {
+                    pendingStages.push({ startActionIndex: recOrder - 1, key })
+                }
+            }
+            break
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Header stage finalisation
+// ---------------------------------------------------------------------------
+
+/**
+ * Computes `component_indexes` for each developer-declared stage and merges it
+ * onto the pre-created header record.
+ *
+ * Matching is by **key name**: the `key` stored in each `pendingStage` entry is
+ * looked up in `headerStageRecords`. This is correct even when the order of
+ * `wfa.stage()` calls in the body differs from the declaration order in the
+ * stages config header.
+ *
+ * A stage may appear multiple times in `pendingStages` (e.g. inside both an
+ * `if` and an `elseIf` body).  All indexes for the same key are collected and
+ * joined with `,` into a single `component_indexes` value on one record.
+ *
+ * @returns Array of finalised stage records ready to be added to `relatedRecords`.
+ */
+export function finalizeStageRecords(
+    pendingStages: Array<{ startActionIndex: number; key: string }>,
+    headerStageRecords: Map<string, Record>
+): Record[] {
+    // Group indexes by stage key, preserving encounter order for `order`.
+    const indexesByKey = new Map<string, number[]>()
+    const keyOrder: string[] = []
+
+    for (const entry of pendingStages) {
+        if (!entry) {
+            continue
+        }
+        const existing = indexesByKey.get(entry.key)
+        if (existing) {
+            existing.push(entry.startActionIndex)
+        } else {
+            indexesByKey.set(entry.key, [entry.startActionIndex])
+            keyOrder.push(entry.key)
+        }
+    }
+
+    const finalised: Record[] = []
+    for (let i = 0; i < keyOrder.length; i++) {
+        const key = keyOrder[i]!
+        const indexes = indexesByKey.get(key)!
+        const headerRecord = headerStageRecords.get(key)
+        if (headerRecord) {
+            finalised.push(headerRecord.merge({ component_indexes: indexes.join(','), order: i }))
+        }
+    }
+
+    // Emit declared-but-never-activated stages with empty component_indexes.
+    // These are stages declared in the header config that have no wfa.stage()
+    // call anywhere in the flow body (e.g. a stage only activated inside a
+    // subflow).  The platform still needs the sys_hub_flow_stage record so the
+    // stage tracker can display it.
+    const activatedKeys = new Set(keyOrder)
+    let nextOrder = keyOrder.length
+    for (const [key, record] of headerStageRecords) {
+        if (!activatedKeys.has(key)) {
+            finalised.push(record.merge({ component_indexes: '', order: nextOrder++ }))
+        }
+    }
+
+    return finalised
+}
+
+// ---------------------------------------------------------------------------
+// toShape helpers — reconstruct stages from DB records back to Fluent shapes
+// ---------------------------------------------------------------------------
+
+// snake_case DB key → camelCase API key (reverse of STATE_KEY_MAP)
+const REVERSE_STATE_KEY_MAP: { [k: string]: string } = {
+    pending: 'pending',
+    in_progress: 'inProgress',
+    complete: 'complete',
+    error: 'error',
+    skipped: 'skipped',
+}
+
+/**
+ * Converts a stage `value` string (which may contain spaces, dots, underscores
+ * etc.) into a valid camelCase JavaScript identifier suitable for use as an
+ * object property key and in `params.stages.<key>` property-access expressions.
+ *
+ * Examples:
+ *   "Request Cancelled" → "requestCancelled"
+ *   "Dept. Head Approval" → "deptHeadApproval"
+ *   "CIO Approval"        → "cioApproval"
+ *   "manager_approval"    → "managerApproval"
+ *   "triage"              → "triage"
+ */
+export function toSafeIdentifier(value: string): string {
+    const parts = value.split(/[^a-zA-Z0-9]+/).filter(Boolean)
+    if (parts.length === 0) {
+        return value
+    }
+    const result = parts
+        .map((part, index) => {
+            const lower = part.toLowerCase()
+            if (index === 0) {
+                return lower
+            }
+            return lower.charAt(0).toUpperCase() + lower.slice(1)
+        })
+        .join('')
+    // Prefix with underscore if starts with a digit
+    return /^\d/.test(result) ? `_${result}` : result
+}
+
+/**
+ * Builds `FlowStage()` CallExpressionShapes from `sys_hub_flow_stage` records
+ * for the header `stages: {}` config property.
+ *
+ * Only standard stages are included (type='subflow' stages are skipped).
+ *
+ * @returns Plain object keyed by stage `value` — e.g. `{ triage: FlowStage({...}) }`.
+ */
+export function buildStageShapes(stageRecords: Record[]): globalThis.Record<string, CallExpressionShape> {
+    const shapes: globalThis.Record<string, CallExpressionShape> = {}
+
+    const sorted = [...stageRecords].sort(
+        (a, b) => Number(a.get('order')?.getValue() ?? 0) - Number(b.get('order')?.getValue() ?? 0)
+    )
+
+    for (const rec of sorted) {
+        if (rec.get('type')?.getValue() !== 'standard') {
+            continue
+        }
+        const value = rec.get('value')?.ifString()?.getValue() ?? ''
+        if (!value) {
+            continue
+        }
+
+        // Build config properties
+        const props: globalThis.Record<string, unknown> = {
+            label: rec.get('label')?.ifString()?.getValue() ?? '',
+            value,
+        }
+
+        // Duration — only emit when not the default (zero duration)
+        const durationStr = rec.get('duration')?.ifString()?.getValue() ?? STAGE_DEFAULT_DURATION
+        if (durationStr && durationStr !== STAGE_DEFAULT_DURATION) {
+            try {
+                props['duration'] = DurationShape.from(rec, Shape.from(rec, durationStr).asString())
+            } catch {
+                /* keep default */
+            }
+        }
+
+        // alwaysShow
+        const alwaysShow = rec.get('always_show')?.getValue()
+        if (alwaysShow === true || alwaysShow === 'true') {
+            props['alwaysShow'] = true
+        }
+
+        // states — only emit when they differ from platform defaults
+        const statesStr = rec.get('states')?.ifString()?.getValue()
+        if (statesStr) {
+            try {
+                const parsed: { [k: string]: string } = JSON.parse(statesStr)
+                const isDefault =
+                    Object.keys(parsed).length === Object.keys(DEFAULT_STATES).length &&
+                    Object.entries(parsed).every(([k, v]) => DEFAULT_STATES[k] === v)
+                if (!isDefault) {
+                    const stateProps: globalThis.Record<string, string> = {}
+                    for (const [dbKey, label] of Object.entries(parsed)) {
+                        const apiKey = REVERSE_STATE_KEY_MAP[dbKey]
+                        if (apiKey && label) {
+                            stateProps[apiKey] = label
+                        }
+                    }
+                    if (Object.keys(stateProps).length > 0) {
+                        props['states'] = Shape.from(rec, stateProps).asObject()
+                    }
+                }
+            } catch {
+                /* skip malformed JSON */
+            }
+        }
+
+        let key = toSafeIdentifier(value)
+        if (key in shapes) {
+            let suffix = 2
+            while (`${key}_${suffix}` in shapes) {
+                suffix++
+            }
+            key = `${key}_${suffix}`
+        }
+        shapes[key] = new CallExpressionShape({
+            source: rec,
+            callee: FLOW_STAGE_CALLEE,
+            args: [new ObjectShape({ source: rec, properties: props })],
+        })
+    }
+
+    return shapes
+}
+
+/**
+ * Parses `component_indexes` from stage records and builds a map from
+ * action record order to the stage key(s) that should be inserted before it.
+ *
+ * component_index N means the stage belongs before the record with order = N + 1.
+ */
+export function buildStageInsertionMap(stageRecords: Record[]): Map<number, string[]> {
+    const map = new Map<number, string[]>()
+    for (const rec of stageRecords) {
+        if (rec.get('type')?.getValue() !== 'standard') {
+            continue
+        }
+        const value = rec.get('value')?.ifString()?.getValue() ?? ''
+        if (!value) {
+            continue
+        }
+        const indexes = rec.get('component_indexes')?.ifString()?.getValue() ?? ''
+        if (!indexes) {
+            continue
+        }
+        for (const part of indexes.split(',')) {
+            const actionOrder = parseInt(part.trim(), 10) + 1
+            if (isNaN(actionOrder)) {
+                continue
+            }
+            const existing = map.get(actionOrder) ?? []
+            existing.push(toSafeIdentifier(value))
+            map.set(actionOrder, existing)
+        }
+    }
+    return map
+}
+
+/**
+ * Creates a `wfa.stage(params.stages.<key>)` CallExpressionShape.
+ */
+export function createStageCallShape(stageKey: string, paramName: string, source: Source): CallExpressionShape {
+    return new CallExpressionShape({
+        source,
+        callee: STAGE_API_NAME,
+        args: [
+            new PropertyAccessShape({
+                source,
+                elements: [paramName, 'stages', stageKey],
+            }),
+        ],
+    })
+}
+
+/**
+ * Scans `stmts` for records whose order matches an entry in `stageInsertionMap`
+ * and splices `wfa.stage()` calls before them. Matched entries are removed from
+ * the map so they aren't inserted twice.
+ */
+function insertStagesIntoStatements(
+    stmts: Shape[],
+    stageInsertionMap: Map<number, string[]>,
+    paramName: string,
+    source: Source,
+    orderFn: (rec: Record) => number
+): void {
+    const insertions: Array<{ before: number; keys: string[] }> = []
+    for (let i = 0; i < stmts.length; i++) {
+        const stmtSource = stmts[i]!.getSource()
+        if (!(stmtSource instanceof Record)) {
+            continue
+        }
+        const recOrder = orderFn(stmtSource)
+        const keys = stageInsertionMap.get(recOrder)
+        if (keys && keys.length > 0) {
+            insertions.push({ before: i, keys })
+            stageInsertionMap.delete(recOrder)
+        }
+    }
+    for (let j = insertions.length - 1; j >= 0; j--) {
+        const { before, keys } = insertions[j]!
+        const stageCalls = keys.map((k) => createStageCallShape(k, paramName, source))
+        stmts.splice(before, 0, ...stageCalls)
+    }
+}
+
+/**
+ * Walks through an array of body shapes and mutates flow logic bodies
+ * (if/elseIf/else/tryCatch) to insert `wfa.stage()` calls at positions identified
+ * by `stageInsertionMap`.
+ *
+ * Matched entries are removed from the map so they aren't inserted twice.
+ * Call this AFTER top-level stages have been handled.
+ */
+export function insertNestedStageShapes(
+    shapes: Shape[],
+    stageInsertionMap: Map<number, string[]>,
+    paramName: string,
+    orderFn: (rec: Record) => number
+): void {
+    for (const shape of shapes) {
+        const callExpr =
+            shape instanceof VariableStatementShape ? (shape as VariableStatementShape).getInitializer() : shape
+        if (!(callExpr instanceof CallExpressionShape)) {
+            continue
+        }
+        const callee = callExpr.getCallee()
+
+        // Recurse into tryCatch try/catch handler bodies
+        if (callee === FLOW_LOGIC.TRY_CATCH) {
+            const handlersArg = callExpr.getArgument(1, false)
+            if (handlersArg.if(ObjectShape)) {
+                const handlers = handlersArg.as(ObjectShape)
+                for (const handlerKey of ['try', 'catch'] as const) {
+                    const handler = handlers.get(handlerKey, false)
+                    if (handler.if(ArrowFunctionShape)) {
+                        const handlerBody = handler.as(ArrowFunctionShape)
+                        const handlerStmts = handlerBody.getStatements()
+                        insertStagesIntoStatements(
+                            handlerStmts,
+                            stageInsertionMap,
+                            paramName,
+                            handlerBody.getSource(),
+                            orderFn
+                        )
+                        insertNestedStageShapes(handlerStmts, stageInsertionMap, paramName, orderFn)
+                    }
+                }
+            }
+            continue
+        }
+
+        // Only recurse into conditional bodies
+        if (callee !== FLOW_LOGIC.IF && callee !== FLOW_LOGIC.ELSEIF && callee !== FLOW_LOGIC.ELSE) {
+            continue
+        }
+
+        const bodyArg = callExpr.getArgument(1, false)
+        if (!bodyArg.if(ArrowFunctionShape)) {
+            continue
+        }
+        const body = bodyArg.as(ArrowFunctionShape)
+        const stmts = body.getStatements() // mutable reference
+
+        insertStagesIntoStatements(stmts, stageInsertionMap, paramName, body.getSource(), orderFn)
+
+        // Recurse into children for deeper nesting
+        insertNestedStageShapes(stmts, stageInsertionMap, paramName, orderFn)
+    }
+}