@bastani/atomic 0.8.13-0 → 0.8.14-0

package/dist/builtin/workflows/src/runs/foreground/executor.ts

@@ -33,7 +33,14 @@ import type {
   WorkflowModelCatalogPort,
 } from "../../shared/types.js";
 import type { InternalStageContext, StageAdapters } from "./stage-runner.js";
-import type { RunStatus, StageNotice, StageSnapshot, RunSnapshot, WorkflowOverlayAdapter } from "../../shared/store-types.js";
+import type {
+  RunStatus,
+  StageNotice,
+  StageSnapshot,
+  RunSnapshot,
+  WorkflowOverlayAdapter,
+  WorkflowFailureKind,
+} from "../../shared/store-types.js";
 import type { StageControlHandle, StageControlRegistry, AgentSessionEventListener } from "./stage-control-registry.js";
 import type { Store } from "../../shared/store.js";
 import type { CancellationRegistry } from "../background/cancellation-registry.js";
@@ -59,9 +66,16 @@ import {
   appendRunEnd,
 } from "../../shared/persistence-session-entries.js";
 import { validateWorkflowModels } from "../shared/model-fallback.js";
+import type { WorkflowFailure } from "../../shared/workflow-failures.js";
+import { classifyWorkflowFailure } from "../../shared/workflow-failures.js";
 
 export interface ResolvedInputs extends Record<string, unknown> {}
 
+export interface RunContinuationOpts {
+  readonly source: RunSnapshot;
+  readonly resumeFromStageId: string;
+}
+
 export interface RunOpts {
   adapters?: StageAdapters;
   /** HIL adapter injected by the pi runtime or test harness. */
@@ -113,6 +127,8 @@ export interface RunOpts {
    * the runId before starting the background promise.
    */
   runId?: string;
+  /** Replay completed stages from a failed source run, then resume at this stage. */
+  continuation?: RunContinuationOpts;
   onRunStart?: (snapshot: RunSnapshot) => void;
   onStageStart?: (runId: string, snapshot: StageSnapshot) => void;
   onStageEnd?: (runId: string, snapshot: StageSnapshot) => void;
@@ -519,16 +535,22 @@ async function mapParallelSteps<T>(
   concurrency: number | undefined,
   failFast: boolean | undefined,
   mapper: (step: WorkflowTaskStep) => Promise<T>,
+  onFirstFailure?: (error: unknown) => void,
 ): Promise<T[]> {
   const limit = positiveConcurrency(concurrency) ?? steps.length;
+  const failFastEnabled = failFast !== false;
   const results = new Array<T>(steps.length);
   const failures: Array<{ readonly index: number; readonly error: unknown }> = [];
   let nextIndex = 0;
   let firstFailure: unknown;
+  let rejectFirstFailure: (reason: unknown) => void = () => {};
+  const firstFailurePromise = new Promise<never>((_, reject) => {
+    rejectFirstFailure = reject;
+  });
 
   async function worker(): Promise<void> {
     while (true) {
-      if (failFast !== false && firstFailure !== undefined) return;
+      if (failFastEnabled && firstFailure !== undefined) return;
       const index = nextIndex;
       nextIndex += 1;
       const step = steps[index];
@@ -537,15 +559,29 @@ async function mapParallelSteps<T>(
         results[index] = await mapper(step);
       } catch (err) {
         failures.push({ index, error: err });
-        firstFailure ??= err;
-        if (failFast !== false) throw err;
+        if (firstFailure === undefined) {
+          firstFailure = err;
+          onFirstFailure?.(err);
+          if (failFastEnabled) rejectFirstFailure(err);
+        }
+        if (failFastEnabled) return;
       }
     }
   }
 
-  await Promise.all(
-    Array.from({ length: Math.min(limit, steps.length) }, () => worker()),
-  );
+  const workers = Array.from({ length: Math.min(limit, steps.length) }, () => worker());
+  const allWorkers = Promise.all(workers);
+
+  if (!failFastEnabled) {
+    await allWorkers;
+  } else {
+    try {
+      await Promise.race([allWorkers, firstFailurePromise]);
+    } catch (err) {
+      void allWorkers.catch(() => {});
+      throw err;
+    }
+  }
 
   if (failures.length > 0) {
     throw new AggregateError(
@@ -723,6 +759,10 @@ async function writeDirectOutput(
   };
 }
 
+function directFailureMessage(error: unknown): string {
+  return classifyWorkflowFailure(error).userMessage;
+}
+
 function failedDirectDetails(
   mode: WorkflowDetails["mode"],
   runId: string,
@@ -738,7 +778,7 @@ function failedDirectDetails(
     ...(options.context !== undefined ? { context: options.context } : {}),
     results: [],
     progress: { completed: 0, total },
-    error: error instanceof Error ? error.message : String(error),
+    error: directFailureMessage(error),
   };
 }
 
@@ -1013,6 +1053,11 @@ function appendRunEndWhenRecorded(
     readonly runId: string;
     readonly status: RunStatus;
     readonly result?: Record<string, unknown>;
+    readonly error?: string;
+    readonly failureKind?: WorkflowFailureKind;
+    readonly failureMessage?: string;
+    readonly failedStageId?: string;
+    readonly resumable?: boolean;
     readonly ts: number;
   },
 ): void {
@@ -1020,6 +1065,103 @@ function appendRunEndWhenRecorded(
   appendRunEnd(persistence, payload);
 }
 
+interface RunFailureMetadata {
+  readonly errorMessage: string;
+  readonly failureKind: WorkflowFailureKind;
+  readonly failureMessage: string;
+  readonly failedStageId?: string;
+  readonly resumable: boolean;
+}
+
+function applyFailureToStage(stage: StageSnapshot, failure: WorkflowFailure): void {
+  stage.status = "failed";
+  stage.error = failure.userMessage;
+  stage.failureKind = failure.kind;
+  stage.failureMessage = failure.message;
+}
+
+function runFailureMetadata(err: unknown, stages: readonly StageSnapshot[]): RunFailureMetadata {
+  const classified = classifyWorkflowFailure(err);
+  const failedStage = stages.find((stage) => stage.status === "failed");
+  const failureKind = failedStage?.failureKind ?? classified.kind;
+
+  return {
+    errorMessage: classified.userMessage,
+    failureKind,
+    failureMessage: failedStage?.failureMessage ?? classified.message,
+    ...(failedStage !== undefined ? { failedStageId: failedStage.id } : {}),
+    resumable: classified.resumable,
+  };
+}
+
+type ContinuationReplayDecision =
+  | { readonly kind: "execute"; readonly source?: StageSnapshot }
+  | { readonly kind: "replay"; readonly source: StageSnapshot };
+
+interface ContinuationReplayIndex {
+  decide(stageName: string, parentIds: readonly string[], stageId: string): ContinuationReplayDecision;
+}
+
+function sameStringSet(left: readonly string[], right: readonly string[]): boolean {
+  if (left.length !== right.length) return false;
+  const rightSet = new Set(right);
+  return left.every((value) => rightSet.has(value));
+}
+
+function createContinuationReplayIndex(continuation: RunContinuationOpts | undefined): ContinuationReplayIndex {
+  if (continuation === undefined) return { decide: () => ({ kind: "execute" }) };
+  const resumeStage = continuation.source.stages.find((stage) => stage.id === continuation.resumeFromStageId);
+  if (resumeStage === undefined) {
+    throw new Error(`pi-workflows: insufficient_state: resume stage ${continuation.resumeFromStageId} was not found in source run ${continuation.source.id}`);
+  }
+
+  const stagesByName = new Map<string, StageSnapshot[]>();
+  for (const stage of continuation.source.stages) {
+    const stages = stagesByName.get(stage.name);
+    if (stages === undefined) {
+      stagesByName.set(stage.name, [stage]);
+    } else {
+      stages.push(stage);
+    }
+  }
+
+  const consumedSourceStageIds = new Set<string>();
+  const sourceStageIdByContinuationStageId = new Map<string, string>();
+  return {
+    decide(stageName: string, parentIds: readonly string[], stageId: string): ContinuationReplayDecision {
+      const candidates = stagesByName.get(stageName)?.filter((stage) => !consumedSourceStageIds.has(stage.id)) ?? [];
+      if (candidates.length === 0) return { kind: "execute" };
+
+      const translatedParentIds = parentIds.map((parentId) => sourceStageIdByContinuationStageId.get(parentId));
+      const hasUnmappedParent = translatedParentIds.some((parentId) => parentId === undefined);
+      const matches = hasUnmappedParent
+        ? []
+        : candidates.filter((stage) => sameStringSet(translatedParentIds as string[], stage.parentIds));
+
+      if (matches.length !== 1) {
+        const reason = matches.length === 0 ? "mismatch" : "ambiguous";
+        throw new Error(`pi-workflows: insufficient_state: replay topology ${reason} for stage "${stageName}" in source run ${continuation.source.id}`);
+      }
+
+      const source = matches[0]!;
+      consumedSourceStageIds.add(source.id);
+      sourceStageIdByContinuationStageId.set(stageId, source.id);
+      if (source.status === "completed") return { kind: "replay", source };
+      return { kind: "execute", source };
+    },
+  };
+}
+
+interface ParallelFailFastStage {
+  readonly skip: () => void;
+}
+
+interface ParallelFailFastScope {
+  failed: boolean;
+  firstFailure?: unknown;
+  readonly activeStages: Map<string, ParallelFailFastStage>;
+}
+
 // ---------------------------------------------------------------------------
 // Shared killed finalizer — used for catch-abort and post-body abort check
 // ---------------------------------------------------------------------------
@@ -1031,17 +1173,25 @@ function finalizeKilled(
   persistence: WorkflowPersistencePort | undefined,
   onRunEnd: RunOpts["onRunEnd"],
 ): RunResult {
-  const recorded = activeStore.recordRunEnd(runId, "killed", undefined, "workflow killed");
-  onRunEnd?.(runId, "killed", undefined, "workflow killed");
+  const errorMessage = "workflow killed";
+  const metadata = {
+    failureKind: "cancelled" as const,
+    failureMessage: errorMessage,
+    resumable: false,
+  };
+  const recorded = activeStore.recordRunEnd(runId, "killed", undefined, errorMessage, metadata);
+  onRunEnd?.(runId, "killed", undefined, errorMessage);
   appendRunEndWhenRecorded(persistence, recorded, {
     runId,
     status: "killed",
+    error: errorMessage,
+    ...metadata,
     ts: Date.now(),
   });
   return {
     runId,
     status: "killed",
-    error: "workflow killed",
+    error: errorMessage,
     stages: [...runSnapshot.stages],
   };
 }
@@ -1080,6 +1230,7 @@ export async function run<TInputs extends Record<string, unknown>>(
 
   // 2. Generate runId (or use pre-allocated seam from caller)
   const runId = opts.runId ?? crypto.randomUUID();
+  const replayIndex = createContinuationReplayIndex(opts.continuation);
 
   // 2a. Create own AbortController; forward caller signal if provided
   const ownController = new AbortController();
@@ -1100,6 +1251,10 @@ export async function run<TInputs extends Record<string, unknown>>(
     status: "running",
     stages: [],
     startedAt: Date.now(),
+    ...(opts.continuation !== undefined ? {
+      resumedFromRunId: opts.continuation.source.id,
+      resumeFromStageId: opts.continuation.resumeFromStageId,
+    } : {}),
   };
 
   activeStore.recordRunStart(runSnapshot);
@@ -1120,6 +1275,8 @@ export async function run<TInputs extends Record<string, unknown>>(
       runId,
       name: def.name,
       inputs: resolvedInputs,
+      ...(runSnapshot.resumedFromRunId !== undefined ? { resumedFromRunId: runSnapshot.resumedFromRunId } : {}),
+      ...(runSnapshot.resumeFromStageId !== undefined ? { resumeFromStageId: runSnapshot.resumeFromStageId } : {}),
       ts: runSnapshot.startedAt,
     });
   }
@@ -1147,7 +1304,7 @@ export async function run<TInputs extends Record<string, unknown>>(
   };
 
   const isTerminalStage = (stage: StageSnapshot): boolean =>
-    stage.status === "completed" || stage.status === "failed";
+    stage.status === "completed" || stage.status === "failed" || stage.status === "skipped";
 
   const stageById = (stageId: string): StageSnapshot | undefined =>
     runSnapshot.stages.find((stage) => stage.id === stageId);
@@ -1280,7 +1437,7 @@ export async function run<TInputs extends Record<string, unknown>>(
     inputs: resolvedInputs as TInputs,
     ui: opts.ui ?? makeUnavailableUIContext(),
 
-    stage(name: string, options?: StageOptions) {
+    stage(name: string, options?: StageOptions, stageFailFastScope?: ParallelFailFastScope) {
       // a. Generate stageId
       const stageId = crypto.randomUUID();
 
@@ -1288,12 +1445,24 @@ export async function run<TInputs extends Record<string, unknown>>(
       const parentIds = tracker.onSpawn(stageId, name);
 
       // c. Create StageSnapshot as "pending"
+      const replayDecision = replayIndex.decide(name, parentIds, stageId);
+      const replaySource = replayDecision.kind === "replay" ? replayDecision.source : undefined;
+      const shouldReplay = replaySource !== undefined;
+
       const stageSnapshot: StageSnapshot = {
         id: stageId,
         name,
-        status: "pending",
+        status: shouldReplay ? "completed" : "pending",
         parentIds: Object.freeze(parentIds),
         toolEvents: [],
+        ...(shouldReplay ? {
+          startedAt: Date.now(),
+          endedAt: Date.now(),
+          durationMs: 0,
+          ...(replaySource.result !== undefined ? { result: replaySource.result } : {}),
+          replayedFromStageId: replaySource.id,
+          replayed: true,
+        } : {}),
         // Store mcp scope options on snapshot when provided
         ...(options?.mcp !== undefined
           ? { mcpScope: { allow: options.mcp.allow ?? null, deny: options.mcp.deny ?? null } }
@@ -1301,9 +1470,87 @@ export async function run<TInputs extends Record<string, unknown>>(
         // Mark attachable up-front: the live stage handle is registered
         // below before the first onStageStart fires, so consumers that
         // hook onStageStart see `attachable: true` for the pending stage.
-        attachable: true,
+        attachable: !shouldReplay,
       };
 
+      let stageStartEntryAppended = false;
+      const appendStageStartOnce = (): void => {
+        if (!opts.persistence || stageStartEntryAppended) return;
+        stageStartEntryAppended = true;
+        appendStageStart(opts.persistence, {
+          runId,
+          stageId,
+          name,
+          parentIds: stageSnapshot.parentIds,
+          ...(stageSnapshot.replayedFromStageId !== undefined ? { replayedFromStageId: stageSnapshot.replayedFromStageId } : {}),
+          ...(stageSnapshot.replayed !== undefined ? { replayed: stageSnapshot.replayed } : {}),
+          ts: stageSnapshot.startedAt ?? Date.now(),
+        });
+      };
+
+      if (shouldReplay) {
+        activeStore.recordStageStart(runId, stageSnapshot);
+        opts.onStageStart?.(runId, stageSnapshot);
+        appendStageStartOnce();
+        let replayFinalized = false;
+        const finalizeReplayStage = (): void => {
+          if (replayFinalized) return;
+          replayFinalized = true;
+          activeStore.recordStageEnd(runId, stageSnapshot);
+          opts.onStageEnd?.(runId, stageSnapshot);
+          if (opts.persistence) {
+            appendStageEnd(opts.persistence, {
+              runId,
+              stageId,
+              status: "completed",
+              durationMs: 0,
+              ...(stageSnapshot.result !== undefined ? { summary: stageSnapshot.result } : {}),
+              replayedFromStageId: replaySource.id,
+              replayed: true,
+            });
+          }
+          tracker.onSettle(stageId);
+        };
+        const replayResult = replaySource.result ?? "";
+        const replayText = async (): Promise<string> => {
+          await Promise.resolve();
+          finalizeReplayStage();
+          return replayResult;
+        };
+        const rejectReplayMutation = (action: string): never => {
+          throw new Error(`pi-workflows: replayed stage "${name}" cannot ${action}`);
+        };
+        const replayContext: StageContext & Pick<InternalStageContext, "__modelFallbackMeta"> = {
+          name,
+          prompt: replayText,
+          complete: replayText,
+          steer: async () => rejectReplayMutation("steer"),
+          followUp: async () => rejectReplayMutation("follow up"),
+          subscribe: () => () => {},
+          get sessionFile() { return replaySource.sessionFile; },
+          get sessionId() { return replaySource.sessionId ?? ""; },
+          setModel: async () => rejectReplayMutation("set model"),
+          setThinkingLevel: () => rejectReplayMutation("set thinking level"),
+          cycleModel: async () => rejectReplayMutation("cycle model"),
+          cycleThinkingLevel: () => rejectReplayMutation("cycle thinking level"),
+          get agent() { return undefined as never; },
+          get model() { return replaySource.model as never; },
+          get thinkingLevel() { return undefined as never; },
+          get messages() { return [] as never; },
+          get isStreaming() { return false; },
+          navigateTree: async () => rejectReplayMutation("navigate conversation tree"),
+          compact: async () => rejectReplayMutation("compact"),
+          abortCompaction: () => rejectReplayMutation("abort compaction"),
+          abort: async () => rejectReplayMutation("abort"),
+          __modelFallbackMeta: () => ({
+            ...(replaySource.model !== undefined ? { model: replaySource.model } : {}),
+            ...(replaySource.attemptedModels !== undefined ? { attemptedModels: replaySource.attemptedModels } : {}),
+            ...(replaySource.modelAttempts !== undefined ? { modelAttempts: replaySource.modelAttempts } : {}),
+          }),
+        };
+        return replayContext;
+      }
+
       // d. Create inner AgentSession-like StageContext (raw, without lifecycle wrapping).
       //    Must come before the registry registration because the handle
       //    delegates to it for every operation.
@@ -1349,6 +1596,47 @@ export async function run<TInputs extends Record<string, unknown>>(
         unregisterStageHandle();
         await disposeInnerContext();
       };
+      const hasQueuedLiveWork = (): boolean =>
+        innerCtx.isStreaming || innerCtx.__pendingMessageCount() > 0 || activeAskUserQuestionCalls.size > 0;
+      const releaseLiveHandleWhenIdle = async (): Promise<void> => {
+        dropStageControlHandle();
+        if (!hasQueuedLiveWork()) {
+          await releaseLiveHandle();
+          return;
+        }
+
+        // The queued-work branch installs asynchronous cleanup and returns once
+        // the release watcher is armed. Inner-context events normally trigger
+        // the subscription when streaming/pending-message counters change, but
+        // SDK prompt/tool cleanup can also drain after the stage has stopped
+        // emitting workflow-visible events. The unref'd 250 ms interval is a
+        // fallback for that silent drain path and is cleared as soon as the
+        // handle becomes idle.
+        let unsubscribe = (): void => {};
+        let pollTimer: ReturnType<typeof setInterval> | undefined;
+        const cleanupWatcher = (): void => {
+          unsubscribe();
+          if (pollTimer !== undefined) {
+            clearInterval(pollTimer);
+            pollTimer = undefined;
+          }
+        };
+        const releaseIfIdle = (): void => {
+          if (liveHandleReleased) {
+            cleanupWatcher();
+            return;
+          }
+          if (hasQueuedLiveWork()) return;
+          cleanupWatcher();
+          void releaseLiveHandle().catch((error: unknown) => {
+            console.debug("pi-workflows: failed to release idle stage handle", error);
+          });
+        };
+        unsubscribe = innerCtx.subscribe(() => queueMicrotask(releaseIfIdle));
+        pollTimer = setInterval(releaseIfIdle, 250);
+        pollTimer.unref?.();
+        releaseIfIdle();
+      };
 
       // e. Register a live stage-control handle so attached panes can
       //    prompt/steer/pause/resume the underlying Pi session lazily.
@@ -1372,6 +1660,9 @@ export async function run<TInputs extends Record<string, unknown>>(
         get isStreaming() {
           return innerCtx.isStreaming;
         },
+        get isDisposed() {
+          return liveHandleReleased;
+        },
         get messages() {
           return innerCtx.messages;
         },
@@ -1421,6 +1712,59 @@ export async function run<TInputs extends Record<string, unknown>>(
           return innerCtx.subscribe(listener);
         },
       };
+      let stageFinalized = false;
+      const finalizeStageSnapshot = (): boolean => {
+        if (stageFinalized) return false;
+        stageFinalized = true;
+        stageSnapshot.endedAt = Date.now();
+        stageSnapshot.durationMs = elapsedStageMs(stageSnapshot, stageSnapshot.endedAt);
+
+        const finalModelMeta = innerCtx.__modelFallbackMeta();
+        if (finalModelMeta.model !== undefined) stageSnapshot.model = finalModelMeta.model;
+        if (finalModelMeta.attemptedModels !== undefined) stageSnapshot.attemptedModels = finalModelMeta.attemptedModels;
+        if (finalModelMeta.modelAttempts !== undefined) stageSnapshot.modelAttempts = finalModelMeta.modelAttempts;
+
+        activeStore.recordStageEnd(runId, stageSnapshot);
+        opts.onStageEnd?.(runId, stageSnapshot);
+
+        if (opts.persistence) {
+          appendStageStartOnce();
+          appendStageEnd(opts.persistence, {
+            runId,
+            stageId,
+            status: stageSnapshot.status,
+            durationMs: stageSnapshot.durationMs,
+            ...(stageSnapshot.error !== undefined ? { error: stageSnapshot.error } : {}),
+            ...(stageSnapshot.failureKind !== undefined ? { failureKind: stageSnapshot.failureKind } : {}),
+            ...(stageSnapshot.failureMessage !== undefined ? { failureMessage: stageSnapshot.failureMessage } : {}),
+            ...(stageSnapshot.skippedReason !== undefined ? { skippedReason: stageSnapshot.skippedReason } : {}),
+            ...(stageSnapshot.result !== undefined && stageSnapshot.status === "completed" ? { summary: stageSnapshot.result } : {}),
+            ...(stageSnapshot.replayedFromStageId !== undefined ? { replayedFromStageId: stageSnapshot.replayedFromStageId } : {}),
+            ...(stageSnapshot.replayed !== undefined ? { replayed: stageSnapshot.replayed } : {}),
+          });
+        }
+
+        stageFailFastScope?.activeStages.delete(stageId);
+        tracker.onSettle(stageId);
+        return true;
+      };
+      let skippedForParallelFailFast = false;
+      const markSkippedForParallelFailFast = (): void => {
+        skippedForParallelFailFast = true;
+        stageSnapshot.status = "skipped";
+        stageSnapshot.skippedReason = "fail-fast";
+      };
+      const parallelFailFastError = (): unknown =>
+        stageFailFastScope?.firstFailure ?? new Error("pi-workflows: skipped after parallel fail-fast");
+      const skipForParallelFailFast = (): void => {
+        if (isTerminalStage(stageSnapshot)) return;
+        markSkippedForParallelFailFast();
+        finalizeStageSnapshot();
+        void innerCtx.abort().catch(() => {});
+        void releaseLiveHandleWhenIdle().catch(() => {});
+      };
+      stageFailFastScope?.activeStages.set(stageId, { skip: skipForParallelFailFast });
+
       let stageControlDropped = false;
       dropStageControlHandle = (): void => {
         if (stageControlDropped) return;
@@ -1454,12 +1798,18 @@ export async function run<TInputs extends Record<string, unknown>>(
 
       const runTrackedStageCall = async (call: () => Promise<string>): Promise<string> => {
         await waitForStageRelease();
+        if (stageFinalized) {
+          throw parallelFailFastError();
+        }
 
         // Block here until a concurrency slot is available for this run.
         await limiter.acquire();
 
         try {
           await waitForStageRelease();
+          if (stageFinalized) {
+            throw parallelFailFastError();
+          }
         } catch (err) {
           limiter.release();
           throw err;
@@ -1470,15 +1820,7 @@ export async function run<TInputs extends Record<string, unknown>>(
         activeStore.recordStageStart(runId, stageSnapshot);
 
         // Persistence: append stage.start entry
-        if (opts.persistence) {
-          appendStageStart(opts.persistence, {
-            runId,
-            stageId,
-            name,
-            parentIds: stageSnapshot.parentIds,
-            ts: stageSnapshot.startedAt,
-          });
-        }
+        appendStageStartOnce();
 
         const mcpAllow = options?.mcp?.allow ?? null;
         const mcpDeny = options?.mcp?.deny ?? null;
@@ -1513,6 +1855,13 @@ export async function run<TInputs extends Record<string, unknown>>(
             if (modelMeta.attemptedModels !== undefined) stageSnapshot.attemptedModels = modelMeta.attemptedModels;
             if (modelMeta.modelAttempts !== undefined) stageSnapshot.modelAttempts = modelMeta.modelAttempts;
           }
+          if (stageFailFastScope?.failed === true && stageFailFastScope.activeStages.has(stageId)) {
+            markSkippedForParallelFailFast();
+            throw parallelFailFastError();
+          }
+          if (stageFinalized) {
+            throw parallelFailFastError();
+          }
           stageSnapshot.status = "completed";
           const assistantText = innerCtx.__getLastAssistantText();
           if (assistantText !== undefined) {
@@ -1520,77 +1869,24 @@ export async function run<TInputs extends Record<string, unknown>>(
           }
           return result;
         } catch (err) {
-          if (!ownController.signal.aborted) {
-            stageSnapshot.status = "failed";
-            stageSnapshot.error = err instanceof Error ? err.message : String(err);
+          if (!ownController.signal.aborted && !skippedForParallelFailFast) {
+            applyFailureToStage(stageSnapshot, classifyWorkflowFailure(err));
           }
           throw err;
         } finally {
-          stageSnapshot.endedAt = Date.now();
-          stageSnapshot.durationMs = elapsedStageMs(stageSnapshot, stageSnapshot.endedAt);
-
-          const finalModelMeta = innerCtx.__modelFallbackMeta();
-          if (finalModelMeta.model !== undefined) stageSnapshot.model = finalModelMeta.model;
-          if (finalModelMeta.attemptedModels !== undefined) stageSnapshot.attemptedModels = finalModelMeta.attemptedModels;
-          if (finalModelMeta.modelAttempts !== undefined) stageSnapshot.modelAttempts = finalModelMeta.modelAttempts;
-
           if (opts.mcp && hasMcpScope) {
             opts.mcp.clearScope(stageId);
           }
 
-          activeStore.recordStageEnd(runId, stageSnapshot);
-          opts.onStageEnd?.(runId, stageSnapshot);
-
-          // Persistence: append stage.end entry
-          if (opts.persistence) {
-            appendStageEnd(opts.persistence, {
-              runId,
-              stageId,
-              status: stageSnapshot.status,
-              durationMs: stageSnapshot.durationMs,
-            });
-          }
-
-          tracker.onSettle(stageId);
+          finalizeStageSnapshot();
           // The stage has finished participating in workflow scheduling. Drop it
-          // from run-level pause/resume and cascade-pause lookups immediately,
-          // but keep the SDK session alive below while an attached chat pane is
-          // still using its direct handle.
-          dropStageControlHandle();
-          if (stageSnapshot.attached === true) {
-            let unsubscribeDetach: (() => void) | undefined;
-            let abortListener: (() => void) | undefined;
-            const releaseWhenDetached = (force = false): void => {
-              const currentRun = activeStore.runs().find((r) => r.id === runId);
-              const currentStage = currentRun?.stages.find((s) => s.id === stageId);
-              if (!force && currentStage?.attached === true) return;
-              unsubscribeDetach?.();
-              unsubscribeDetach = undefined;
-              if (abortListener) {
-                ownController.signal.removeEventListener("abort", abortListener);
-                abortListener = undefined;
-              }
-              void releaseLiveHandle().catch(() => {});
-            };
-            unsubscribeDetach = activeStore.subscribe(() => releaseWhenDetached());
-            abortListener = () => releaseWhenDetached(true);
-            if (ownController.signal.aborted) releaseWhenDetached(true);
-            else {
-              ownController.signal.addEventListener(
-                "abort",
-                abortListener,
-                { once: true },
-              );
-            }
-            releaseWhenDetached();
-            limiter.release();
-          } else {
-            try {
-              await releaseLiveHandle();
-            } finally {
-              limiter.release();
-            }
-          }
+          // from run-level pause/resume and cascade-pause lookups immediately.
+          // If no SDK queue/active input remains, release the live chat handle so
+          // the node reopens as a read-only archived session. Queued messages keep
+          // the direct handle alive only until the SDK reports that the queue has
+          // drained.
+          await releaseLiveHandleWhenIdle().catch(() => {});
+          limiter.release();
         }
       };
 
@@ -1678,9 +1974,13 @@ export async function run<TInputs extends Record<string, unknown>>(
       return stageContext;
     },
 
-    async task(name: string, options: WorkflowTaskOptions): Promise<WorkflowTaskResult> {
+    async task(name: string, options: WorkflowTaskOptions, stageFailFastScope?: ParallelFailFastScope): Promise<WorkflowTaskResult> {
       const runTaskOnce = async (taskOptions: WorkflowTaskOptions): Promise<WorkflowTaskResult> => {
-        const stage = ctx.stage(name, taskStageOptions(taskOptions));
+        const stage = (ctx.stage as typeof ctx.stage & ((stageName: string, stageOptions?: StageOptions, scope?: ParallelFailFastScope) => StageContext))(
+          name,
+          taskStageOptions(taskOptions),
+          stageFailFastScope,
+        );
         const rawText = await stage.prompt(
           applyTaskContext(`${taskReadInstruction(taskOptions)}${taskPrompt(taskOptions)}`, taskPrevious(taskOptions)),
           taskPromptOptions(taskOptions),
@@ -1744,12 +2044,23 @@ export async function run<TInputs extends Record<string, unknown>>(
 
     async parallel(steps: readonly WorkflowTaskStep[], options: WorkflowParallelOptions = {}): Promise<WorkflowTaskResult[]> {
       const fallback = parallelFallbackTask(steps, options);
+      const failFastScope: ParallelFailFastScope | undefined = options.failFast === false
+        ? undefined
+        : { failed: false, activeStages: new Map<string, ParallelFailFastStage>() };
       return mapParallelSteps(steps, options.concurrency, options.failFast, async (step) => {
         const prompt = replaceTaskPlaceholder(step.prompt ?? step.task ?? fallback, options.task ?? fallback);
-        return ctx.task(
+        return await (ctx.task as typeof ctx.task & ((taskName: string, taskOptions: WorkflowTaskOptions, scope?: ParallelFailFastScope) => Promise<WorkflowTaskResult>))(
           step.name,
           taskWithSharedDefaults(taskOptionsFromStep(step, prompt, taskPrevious(step)), options),
+          failFastScope,
         );
+      }, (error) => {
+        if (failFastScope === undefined) return;
+        failFastScope.failed = true;
+        failFastScope.firstFailure = error;
+        for (const stage of failFastScope.activeStages.values()) {
+          stage.skip();
+        }
       });
     },
   };
@@ -1792,21 +2103,25 @@ export async function run<TInputs extends Record<string, unknown>>(
       return finalizeKilled(runId, runSnapshot, activeStore, opts.persistence, opts.onRunEnd);
     }
 
-    const errorMessage = err instanceof Error ? err.message : String(err);
-
-    const recorded = activeStore.recordRunEnd(runId, "failed", undefined, errorMessage);
-    opts.onRunEnd?.(runId, "failed", undefined, errorMessage);
+    const metadata = runFailureMetadata(err, runSnapshot.stages);
+    const recorded = activeStore.recordRunEnd(runId, "failed", undefined, metadata.errorMessage, metadata);
+    opts.onRunEnd?.(runId, "failed", undefined, metadata.errorMessage);
 
     appendRunEndWhenRecorded(opts.persistence, recorded, {
       runId,
       status: "failed",
+      error: metadata.errorMessage,
+      failureKind: metadata.failureKind,
+      failureMessage: metadata.failureMessage,
+      ...(metadata.failedStageId !== undefined ? { failedStageId: metadata.failedStageId } : {}),
+      resumable: metadata.resumable,
       ts: Date.now(),
     });
 
     return {
       runId,
       status: "failed",
-      error: errorMessage,
+      error: metadata.errorMessage,
       stages: [...runSnapshot.stages],
     };
   } finally {