stageflow 0.3.0 → 0.9.0

package/dist/runtime/pipelineScheduler.js

@@ -1,9 +1,25 @@
 import { normalizeForkChoice } from "../envelope/forkChoice.js";
-import { collectDownstreamClosure, collectDownstreamStageIds, } from "./dagTraversal.js";
+import { buildPipelineDagSnapshotFromLoaded } from "../runstore/pipelineDagSnapshot.js";
+import { definitionIdForInstance } from "../runstore/stageInstanceId.js";
+import { applyCloneForksFromEnvelopes, applyCloneForksToSchedule, cloneFailureContinuesSchedule, cloneFanoutConflict, cloneRetryDownstream, cloneFailFastSkipIds, cloneScheduleAllowsRun, isCloneInstance, protectedClonableChildIds, sequentialLaterCloneIds, } from "./cloneSchedule.js";
+import { collectDownstreamStageIds } from "./dagTraversal.js";
 import { buildCompletedEnvelopesFromRun, buildStageConfigById, } from "./envelopeRouting.js";
 import { WAIT_WITHOUT_WORKER_DISPATCH } from "./answerResume.js";
 import { runStage, isRunStageWaiting } from "./stageRunner.js";
 import { attemptContext } from "./stageAttemptContext.js";
+function asDagSnapshot(dag) {
+    const snapshot = dag;
+    if (Array.isArray(snapshot.stage_ids))
+        return snapshot;
+    return {
+        ...dag,
+        stage_ids: dag.nodes.map((n) => n.id),
+    };
+}
+export { applyCloneForksFromEnvelopes, applyCloneForksToSchedule, };
+function synthesizedFailureEnvelope(reason) {
+    return { status: "failure", summary: reason, artifacts: [] };
+}
 export function snapshotToScheduleState(snap, overrides, stageId, executionMode) {
     const override = overrides.get(stageId);
     if (override !== undefined)
@@ -17,23 +33,28 @@ export function snapshotToScheduleState(snap, overrides, stageId, executionMode)
             return "active";
         case "waiting_for_input":
             return executionMode === "process" ? "waiting" : "active";
+        case "skipped":
+            return "skipped";
         default:
             return "pending";
     }
 }
 export async function hydrateScheduleFromStore(store, runId, dag, executionMode = "process") {
+    const runDetail = await store.readRun(runId);
+    const meta = await store.readRunMeta(runId);
+    const resolvedDag = meta.pipeline_dag ?? asDagSnapshot(dag);
     const states = new Map();
-    for (const node of dag.nodes) {
+    for (const node of resolvedDag.nodes) {
         states.set(node.id, "pending");
     }
-    const runDetail = await store.readRun(runId);
-    const completedEnvelopes = await buildCompletedEnvelopesFromRun(store, runId, runDetail.stages);
+    const completedEnvelopes = await buildCompletedEnvelopesFromRun(store, runId, runDetail.stages, resolvedDag);
     let schedulingHalted = false;
     const overrides = new Map();
     for (const snap of runDetail.stages) {
         const state = snapshotToScheduleState(snap, overrides, snap.stage_id, executionMode);
         states.set(snap.stage_id, state);
-        if (state === "failed") {
+        if (state === "failed" &&
+            !cloneFailureContinuesSchedule(resolvedDag, snap.stage_id, completedEnvelopes)) {
             schedulingHalted = true;
         }
     }
@@ -42,15 +63,23 @@ export async function hydrateScheduleFromStore(store, runId, dag, executionMode
         completedEnvelopes,
         schedulingHalted,
         firstFailureReason: undefined,
+        dag: resolvedDag,
     };
 }
 export async function hydrateScheduleForRetryRoots(store, runId, dag, retryStageIds) {
     const runDetail = await store.readRun(runId);
-    const downstream = collectDownstreamClosure(dag, retryStageIds);
+    const meta = await store.readRunMeta(runId);
+    const resolvedDag = meta.pipeline_dag ?? asDagSnapshot(dag);
+    const completedEnvelopes = await buildCompletedEnvelopesFromRun(store, runId, runDetail.stages, resolvedDag);
+    const downstream = cloneRetryDownstream(resolvedDag, retryStageIds);
+    for (const rootId of retryStageIds) {
+        for (const id of sequentialLaterCloneIds(resolvedDag, rootId, completedEnvelopes)) {
+            downstream.add(id);
+        }
+    }
     const retryRootSet = new Set(retryStageIds);
-    const completedEnvelopes = await buildCompletedEnvelopesFromRun(store, runId, runDetail.stages);
     const states = new Map();
-    for (const node of dag.nodes) {
+    for (const node of resolvedDag.nodes) {
         states.set(node.id, "pending");
     }
     const overrides = new Map();
@@ -73,6 +102,7 @@ export async function hydrateScheduleForRetryRoots(store, runId, dag, retryStage
         completedEnvelopes,
         schedulingHalted: false,
         firstFailureReason: undefined,
+        dag: resolvedDag,
     };
 }
 export async function hydrateScheduleForRetry(store, runId, dag, retryStageId) {
@@ -80,7 +110,8 @@ export async function hydrateScheduleForRetry(store, runId, dag, retryStageId) {
 }
 export async function resumeRun(options) {
     const { prepared, resumeFromStageId, initialPrior } = options;
-    const dag = prepared.loaded.dag;
+    const meta = await prepared.store.readRunMeta(prepared.run.runId);
+    const dag = meta.pipeline_dag ?? buildPipelineDagSnapshotFromLoaded(prepared.loaded);
     const downstream = dag.childrenOf[resumeFromStageId];
     if (downstream?.length) {
         return runPipelineDag(options);
@@ -124,7 +155,7 @@ export async function resumeRun(options) {
 }
 export async function retryRun(options) {
     const { prepared, retryRoots } = options;
-    const hydrated = await hydrateScheduleForRetryRoots(prepared.store, prepared.run.runId, prepared.loaded.dag, [...retryRoots.keys()]);
+    const hydrated = await hydrateScheduleForRetryRoots(prepared.store, prepared.run.runId, buildPipelineDagSnapshotFromLoaded(prepared.loaded), [...retryRoots.keys()]);
     return runPipelineDag({
         prepared: options.prepared,
         maxActiveStagesPerRun: options.maxActiveStagesPerRun,
@@ -145,6 +176,8 @@ function applyForkChoiceToSchedule(dag, forkStageId, chosen, states, options) {
         }
     };
     for (const childId of dag.childrenOf[forkStageId] ?? []) {
+        if (options?.protectedIds?.has(childId))
+            continue;
         if (!chosen.has(childId)) {
             skipPending(childId);
             for (const desc of collectDownstreamStageIds(dag, childId)) {
@@ -166,8 +199,11 @@ export function applyForkSkipsFromEnvelopes(dag, states, completedEnvelopes) {
 }
 export function applyRetryRootDelta(dag, states, completedEnvelopes, retryRoots, stageId, attempt, clearSchedulingHalt) {
     retryRoots.set(stageId, attempt);
-    const downstream = collectDownstreamClosure(dag, [stageId]);
+    const downstream = cloneRetryDownstream(dag, [stageId]);
     const resetIds = new Set([stageId, ...downstream]);
+    for (const id of sequentialLaterCloneIds(dag, stageId, completedEnvelopes)) {
+        resetIds.add(id);
+    }
     for (const id of resetIds) {
         const state = states.get(id);
         if (state === "active" || state === "succeeded") {
@@ -221,28 +257,28 @@ function sortRunnableIds(dag, ids) {
     const indexById = new Map(dag.nodes.map((node) => [node.id, node.stageIndex]));
     return ids.slice().sort((a, b) => (indexById.get(a) ?? 0) - (indexById.get(b) ?? 0));
 }
-function isRunnable(dag, stageId, states, schedulingHalted) {
+function isRunnable(dag, stageId, states, schedulingHalted, completedEnvelopes) {
     if (schedulingHalted)
         return false;
     const state = states.get(stageId);
     if (state !== "pending")
         return false;
-    const node = dag.nodes.find((n) => n.id === stageId);
-    if (!node)
-        return false;
-    if (node.needs === null)
-        return true;
-    return states.get(node.needs) === "succeeded";
+    return cloneScheduleAllowsRun(dag, stageId, states, completedEnvelopes);
 }
 export async function runPipelineDag(options) {
     const { prepared, maxActiveStagesPerRun } = options;
     const executionMode = options.executionMode ?? "process";
     const { agent, task, loaded, run, store, checkoutRoot, hitl, cwd, projectRoot, } = prepared;
     const factoryCwd = projectRoot ?? cwd ?? process.cwd();
-    const dag = loaded.dag;
+    const fallbackDag = buildPipelineDagSnapshotFromLoaded(loaded);
+    let { states, completedEnvelopes, schedulingHalted, firstFailureReason, dag: hydratedDag, } = options.initialSchedule ??
+        (await hydrateScheduleFromStore(store, run.runId, fallbackDag, executionMode));
+    let dag = hydratedDag ?? fallbackDag;
+    if (hydratedDag === undefined) {
+        const meta = await store.readRunMeta(run.runId);
+        dag = meta.pipeline_dag ?? fallbackDag;
+    }
     const stageById = buildStageConfigById(loaded);
-    let { states, completedEnvelopes, schedulingHalted, firstFailureReason } = options.initialSchedule ??
-        (await hydrateScheduleFromStore(store, run.runId, dag, executionMode));
     const retryContext = options.retryContext;
     const resolveLaunchAttempt = async (stageId) => {
         if (retryContext) {
@@ -293,6 +329,13 @@ export async function runPipelineDag(options) {
         }
     }
     applyForkSkipsFromEnvelopes(dag, states, completedEnvelopes);
+    {
+        const applied = applyCloneForksFromEnvelopes(dag, states, completedEnvelopes);
+        if (applied.dag !== dag) {
+            await store.updatePipelineDag(run.runId, applied.dag);
+        }
+        dag = applied.dag;
+    }
     let activeCount = 0;
     const inFlight = new Set();
     const notifyRetryRootTerminal = (stageId, outcome) => {
@@ -303,6 +346,20 @@ export async function runPipelineDag(options) {
             return;
         options.onRetryRootTerminal?.(stageId, attempt, outcome);
     };
+    const persistSkipPending = async (id) => {
+        if (states.get(id) !== "pending")
+            return;
+        states.set(id, "skipped");
+        notifyRetryRootTerminal(id, "skipped");
+        await store.appendStageEvent(run.runId, id, { event: "skipped" });
+    };
+    for (const [stageId, state] of states) {
+        if (state !== "failed")
+            continue;
+        for (const id of cloneFailFastSkipIds(dag, stageId, completedEnvelopes)) {
+            await persistSkipPending(id);
+        }
+    }
     const markSkippedPending = () => {
         for (const node of dag.nodes) {
             if (states.get(node.id) === "pending") {
@@ -333,9 +390,41 @@ export async function runPipelineDag(options) {
             });
         });
     };
-    const onStageFailure = async (stageId, reason) => {
+    const recordCloneFailureEnvelope = async (stageId, reason, envelope) => {
+        if (!isCloneInstance(dag, stageId)) {
+            if (envelope !== undefined) {
+                completedEnvelopes.set(stageId, envelope);
+            }
+            return;
+        }
+        const failureEnvelope = envelope ??
+            completedEnvelopes.get(stageId) ??
+            synthesizedFailureEnvelope(reason);
+        completedEnvelopes.set(stageId, failureEnvelope);
+        if (envelope === undefined) {
+            try {
+                await store.writeEnvelope(run.runId, stageId, failureEnvelope);
+            }
+            catch {
+            }
+        }
+    };
+    const onStageFailure = async (stageId, reason, envelope) => {
         states.set(stageId, "failed");
         notifyRetryRootTerminal(stageId, "failed");
+        await recordCloneFailureEnvelope(stageId, reason, envelope);
+        if (cloneFailureContinuesSchedule(dag, stageId, completedEnvelopes)) {
+            if (firstFailureReason === undefined) {
+                firstFailureReason = reason;
+            }
+            for (const id of cloneFailFastSkipIds(dag, stageId, completedEnvelopes)) {
+                await persistSkipPending(id);
+            }
+            if (retryContext) {
+                drainRetryMutations();
+            }
+            return;
+        }
         if (retryContext) {
             markBranchDownstreamSkipped(stageId);
             drainRetryMutations();
@@ -352,20 +441,59 @@ export async function runPipelineDag(options) {
             firstFailureReason = reason;
         }
     };
-    const onStageSuccess = (stageId, envelope) => {
+    const onStageSuccess = async (stageId, envelope) => {
+        const conflict = cloneFanoutConflict(dag, envelope);
+        if (conflict !== undefined) {
+            await store.appendStageEvent(run.runId, stageId, {
+                event: "failed",
+                reason: conflict,
+            });
+            await onStageFailure(stageId, conflict);
+            return;
+        }
         states.set(stageId, "succeeded");
         completedEnvelopes.set(stageId, envelope);
         notifyRetryRootTerminal(stageId, "succeeded");
         const node = dag.nodes.find((n) => n.id === stageId);
+        const protectedIds = protectedClonableChildIds(dag, stageId, envelope);
         if (node?.fork) {
             const chosen = normalizeForkChoice(envelope.fork_choice, "stored");
+            const skippedIds = [];
             applyForkChoiceToSchedule(dag, stageId, chosen, states, {
-                onSkip: (id) => notifyRetryRootTerminal(id, "skipped"),
+                protectedIds,
+                onSkip: (id) => {
+                    notifyRetryRootTerminal(id, "skipped");
+                    skippedIds.push(id);
+                },
+            });
+            for (const id of skippedIds) {
+                await store.appendStageEvent(run.runId, id, { event: "skipped" });
+            }
+        }
+        const applied = applyCloneForksToSchedule(dag, stageId, envelope, states);
+        const dagChanged = applied.dag !== dag;
+        dag = applied.dag;
+        if (dagChanged) {
+            await store.updatePipelineDag(run.runId, dag);
+        }
+        for (const id of applied.skippedIds) {
+            notifyRetryRootTerminal(id, "skipped");
+            await store.appendStageEvent(run.runId, id, { event: "skipped" });
+        }
+    };
+    const handlePostSuccessError = async (stageId, reason) => {
+        if (states.get(stageId) === "succeeded") {
+            completedEnvelopes.delete(stageId);
+            await store.appendStageEvent(run.runId, stageId, {
+                event: "failed",
+                reason,
             });
         }
+        await onStageFailure(stageId, reason);
     };
     const launchStage = async (stageId) => {
-        const stage = stageById.get(stageId);
+        const definitionId = definitionIdForInstance(dag, stageId);
+        const stage = stageById.get(definitionId);
         if (!stage) {
             await onStageFailure(stageId, `stage config missing for "${stageId}"`);
             return;
@@ -391,10 +519,11 @@ export async function runPipelineDag(options) {
             if (launchResult.type === "succeeded") {
                 try {
                     const envelope = await store.readEnvelope(run.runId, stageId);
-                    onStageSuccess(stageId, envelope);
+                    await onStageSuccess(stageId, envelope);
                 }
                 catch (err) {
-                    await onStageFailure(stageId, err instanceof Error ? err.message : String(err));
+                    const reason = err instanceof Error ? err.message : String(err);
+                    await handlePostSuccessError(stageId, reason);
                 }
                 return;
             }
@@ -410,6 +539,7 @@ export async function runPipelineDag(options) {
             store,
             runId: run.runId,
             stage,
+            stageId,
             task,
             dag,
             checkoutRoot,
@@ -426,11 +556,17 @@ export async function runPipelineDag(options) {
             return;
         }
         if (!result.ok) {
-            await onStageFailure(stageId, result.reason ?? "stage failed");
+            await onStageFailure(stageId, result.reason ?? "stage failed", result.envelope);
             return;
         }
         if (result.envelope) {
-            onStageSuccess(stageId, result.envelope);
+            try {
+                await onStageSuccess(stageId, result.envelope);
+            }
+            catch (err) {
+                const reason = err instanceof Error ? err.message : String(err);
+                await handlePostSuccessError(stageId, reason);
+            }
         }
         else {
             states.set(stageId, "succeeded");
@@ -451,7 +587,7 @@ export async function runPipelineDag(options) {
     const pickRunnable = () => {
         const ids = [];
         for (const node of dag.nodes) {
-            if (isRunnable(dag, node.id, states, schedulingHalted)) {
+            if (isRunnable(dag, node.id, states, schedulingHalted, completedEnvelopes)) {
                 ids.push(node.id);
             }
         }

package/dist/runtime/resumeReconstruct.js

@@ -4,6 +4,7 @@ import { fakeHitlResumePath } from "../agent/fakeAgent.js";
 import { loadTaskFromYaml } from "../config/loadTask.js";
 import { normalizeCatalogPath } from "../runstore/normalizeCatalogPath.js";
 import { reloadPipelineForRun, reloadTaskForRun, } from "./reloadRunCatalog.js";
+import { definitionIdForInstance } from "../runstore/stageInstanceId.js";
 import { WAIT_WITHOUT_WORKER_DISPATCH } from "./answerResume.js";
 import { resumeRun } from "./pipelineScheduler.js";
 import { buildStageRoots, withResolvedAuthPath } from "./stageRoots.js";
@@ -38,7 +39,8 @@ export async function reconstructAndContinue(ctx) {
         const attemptCtx = attemptContext(latestExecution?.attempt ?? 1);
         const attemptOpt = attemptCtx.eventOptions();
         const { meta, task, loaded } = await loadRunContext(store, runId, cwd);
-        const stageIndex = loaded.stages.findIndex((s) => s.id === stageId);
+        const definitionId = definitionIdForInstance(meta.pipeline_dag, stageId);
+        const stageIndex = loaded.stages.findIndex((s) => s.id === definitionId);
         if (stageIndex < 0) {
             const reason = `Stage ${stageId} not in pipeline ${meta.pipeline_id}`;
             await store.appendStageEvent(runId, stageId, {
@@ -49,6 +51,7 @@ export async function reconstructAndContinue(ctx) {
             return { ok: false, reason };
         }
         const stage = loaded.stages[stageIndex];
+        const dag = meta.pipeline_dag ?? loaded.dag;
         const workspaceDir = store.getWorkspaceDir(runId);
         const checkoutRoot = meta.checkout_root;
         const roots = withResolvedAuthPath(buildStageRoots(workspaceDir, stageId, checkoutRoot, attemptCtx), factoryCwd);
@@ -68,8 +71,9 @@ export async function reconstructAndContinue(ctx) {
             store,
             runId,
             stage,
+            stageId,
             task,
-            dag: loaded.dag,
+            dag,
             checkoutRoot,
             workspaceDir,
             factoryCwd,
@@ -93,8 +97,9 @@ export async function reconstructAndContinue(ctx) {
             store,
             runId,
             stage,
+            stageId,
             task,
-            dag: loaded.dag,
+            dag,
             checkoutRoot,
             workspaceDir,
             hitl: ctx.hitl,

package/dist/runtime/runChangeBus.d.ts

@@ -0,0 +1,18 @@
+import type { RunStore } from "../runstore/port.js";
+export type RunChangeKind = "created" | "status" | "waiting";
+export type RunChangeEvent = {
+    runId: string;
+    kind: RunChangeKind;
+};
+type RunChangeListener = (event: RunChangeEvent) => void;
+export declare class RunChangeBus {
+    private readonly listeners;
+    on(listener: RunChangeListener): () => void;
+    off(listener: RunChangeListener): void;
+    emit(event: RunChangeEvent): void;
+}
+export declare function createRunChangeBus(): RunChangeBus;
+export declare function isRunStoreWrapped(store: RunStore): boolean;
+export declare function getRunChangeBusFromWrappedStore(store: RunStore): RunChangeBus | undefined;
+export declare function wrapRunStoreWithChangeBus(store: RunStore, bus: RunChangeBus): RunStore;
+export {};

package/dist/runtime/runChangeBus.js

@@ -0,0 +1,83 @@
+export class RunChangeBus {
+    listeners = new Set();
+    on(listener) {
+        this.listeners.add(listener);
+        return () => {
+            this.listeners.delete(listener);
+        };
+    }
+    off(listener) {
+        this.listeners.delete(listener);
+    }
+    emit(event) {
+        for (const listener of [...this.listeners]) {
+            try {
+                listener(event);
+            }
+            catch {
+                // Listener isolation: one throw must not block others.
+            }
+        }
+    }
+}
+export function createRunChangeBus() {
+    return new RunChangeBus();
+}
+const WRAPPED_MARKER = Symbol.for("stageflow.wrapRunStoreWithChangeBus");
+const BUS_MARKER = Symbol.for("stageflow.runChangeBus");
+export function isRunStoreWrapped(store) {
+    return store[WRAPPED_MARKER] === true;
+}
+export function getRunChangeBusFromWrappedStore(store) {
+    const bus = store[BUS_MARKER];
+    return bus instanceof RunChangeBus ? bus : undefined;
+}
+export function wrapRunStoreWithChangeBus(store, bus) {
+    if (isRunStoreWrapped(store)) {
+        const bound = getRunChangeBusFromWrappedStore(store);
+        if (bound !== bus) {
+            throw new Error("runChangeBus does not match the bus already bound to the provided store");
+        }
+        return store;
+    }
+    const createRun = async (input) => {
+        const created = await store.createRun(input);
+        bus.emit({ runId: created.runId, kind: "created" });
+        return created;
+    };
+    const updateRunStatus = async (runId, status) => {
+        await store.updateRunStatus(runId, status);
+        bus.emit({ runId, kind: "status" });
+    };
+    const appendStageEvent = async (runId, stageId, event, options) => {
+        await store.appendStageEvent(runId, stageId, event, options);
+        if (event.event === "waiting_for_input" || event.event === "resumed") {
+            bus.emit({ runId, kind: "waiting" });
+        }
+    };
+    const boundCache = new Map();
+    return new Proxy(store, {
+        get(target, prop, receiver) {
+            if (prop === WRAPPED_MARKER)
+                return true;
+            if (prop === BUS_MARKER)
+                return bus;
+            if (prop === "createRun")
+                return createRun;
+            if (prop === "updateRunStatus")
+                return updateRunStatus;
+            if (prop === "appendStageEvent")
+                return appendStageEvent;
+            const cached = boundCache.get(prop);
+            if (cached !== undefined)
+                return cached;
+            const value = Reflect.get(target, prop, receiver);
+            if (typeof value === "function") {
+                const bound = value.bind(target);
+                boundCache.set(prop, bound);
+                return bound;
+            }
+            return value;
+        },
+    });
+}

package/dist/runtime/runManager.d.ts

@@ -111,6 +111,13 @@ export declare class RunManager {
     }): Promise<StartRunResult>;
     rerun(runId: string): Promise<StartRunResult>;
     retryStage(runId: string, stageId: string): Promise<RetryStageResult>;
+    retryStageUntilStop(runId: string, stageId: string): Promise<{
+        ok: true;
+        pipeline: PipelineRunResult;
+    } | Extract<RetryStageResult, {
+        ok: false;
+    }>>;
+    private retryStageInternal;
     /**
      * Deliver an operator answer for `(runId, stageId)` (KTD2 / KTD6).
      * Same-process: unparks the live yield loop. After restart: reconstructs

package/dist/runtime/runManager.js

@@ -389,6 +389,29 @@ export class RunManager {
         return this.reserveAndStartPipeline(taskYaml, pipeline, `run ${runId} task`, rerunCwd, undefined, undefined, undefined, rerunProjectRoot);
     }
     async retryStage(runId, stageId) {
+        return this.retryStageInternal(runId, stageId, true);
+    }
+    async retryStageUntilStop(runId, stageId) {
+        const retryKey = waitKey(runId, stageId);
+        const retried = await this.retryStageInternal(runId, stageId, false);
+        if (!retried.ok)
+            return retried;
+        try {
+            if (retried.done === undefined) {
+                return {
+                    ok: false,
+                    reason: `Retry orchestration did not start for run ${runId} stage ${stageId}`,
+                    status: 500,
+                };
+            }
+            const pipeline = await retried.done;
+            return { ok: true, pipeline };
+        }
+        finally {
+            this.retryInFlight.delete(retryKey);
+        }
+    }
+    async retryStageInternal(runId, stageId, awaitRoot) {
         const retryKey = waitKey(runId, stageId);
         if (this.retryInFlight.has(retryKey)) {
             return {
@@ -401,7 +424,7 @@ export class RunManager {
         try {
             const orchestrationConflict = this.active.has(runId) &&
                 !this.attachedWaiting.has(waitKey(runId, stageId));
-            return await this.retryCoordinator.retryStage({
+            const result = await this.retryCoordinator.retryStage({
                 runId,
                 stageId,
                 store: this.options.store,
@@ -414,10 +437,16 @@ export class RunManager {
                 hitl: this.hitl,
                 orchestrationConflict,
                 tracking: this.retryTracking,
+                awaitRoot,
             });
+            if (!result.ok || awaitRoot !== false) {
+                this.retryInFlight.delete(retryKey);
+            }
+            return result;
         }
-        finally {
+        catch (err) {
             this.retryInFlight.delete(retryKey);
+            throw err;
         }
     }
     /**

package/dist/runtime/runRetryCoordinator.d.ts

@@ -29,6 +29,7 @@ export type RetryStageResult = {
     runId: string;
     stageId: string;
     attemptIndex: number;
+    done?: Promise<PipelineRunResult>;
 } | {
     ok: false;
     reason: string;
@@ -62,6 +63,7 @@ export type RetryStageRequest = {
     hitl: StageHitlController;
     orchestrationConflict: boolean;
     tracking: RetryTrackingPort;
+    awaitRoot?: boolean;
 };
 export declare function assertStageRetryEligible(detail: RunDetail, stageId: string, opts?: {
     recoveryActive?: boolean;

package/dist/runtime/runRetryCoordinator.js

@@ -315,12 +315,15 @@ export class RunRetryCoordinator {
                     tracking.onOrchestrationStarted(runId, orchestrationPromise);
                 }
             }
-            await this.waitForRoot(runId, stageId, store);
+            if (req.awaitRoot !== false) {
+                await this.waitForRoot(runId, stageId, store);
+            }
             return {
                 ok: true,
                 runId,
                 stageId,
                 attemptIndex: execution.attempt,
+                done: this.getOrchestrationPromise(runId),
             };
         }
         catch (err) {

package/dist/runtime/stageAttemptBootstrap.d.ts

@@ -27,6 +27,7 @@ export type StageAttemptOpenInput = {
     resumeToken?: string;
     workerRoots?: boolean;
     completedEnvelopes?: Map<string, StageEnvelope>;
+    stageId?: string;
 };
 export type StageAttemptOpenResult = {
     ok: true;