@bastani/atomic 0.8.14-0 → 0.8.15-0

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

@@ -2,6 +2,7 @@
  * Main DAG executor: run(def, inputs, opts) → RunResult
  */
 
+import { createHash } from "node:crypto";
 import { mkdir, writeFile } from "node:fs/promises";
 import { basename, dirname, extname, isAbsolute, join, resolve } from "node:path";
 import { CONFIG_DIR_NAME } from "@bastani/atomic";
@@ -40,6 +41,8 @@ import type {
   RunSnapshot,
   WorkflowOverlayAdapter,
   WorkflowFailureKind,
+  PendingPrompt,
+  PromptKind,
 } from "../../shared/store-types.js";
 import type { StageControlHandle, StageControlRegistry, AgentSessionEventListener } from "./stage-control-registry.js";
 import type { Store } from "../../shared/store.js";
@@ -68,6 +71,7 @@ import {
 import { validateWorkflowModels } from "../shared/model-fallback.js";
 import type { WorkflowFailure } from "../../shared/workflow-failures.js";
 import { classifyWorkflowFailure } from "../../shared/workflow-failures.js";
+import { selectPromptCallsiteFrame } from "../shared/prompt-callsite.js";
 
 export interface ResolvedInputs extends Record<string, unknown> {}
 
@@ -80,6 +84,8 @@ export interface RunOpts {
   adapters?: StageAdapters;
   /** HIL adapter injected by the pi runtime or test harness. */
   ui?: WorkflowUIAdapter;
+  /** Internal detached-run mode: surface ctx.ui.* as node-local workflow prompt stages. */
+  usePromptNodesForUi?: boolean;
   /** Store override (for testing; defaults to singleton store) */
   store?: Store;
   /** Persistence port for writing session entries (run.start, stage.start, etc.). */
@@ -184,6 +190,69 @@ function resolveInputConcurrency(
 // HIL unavailable fallback — rejects with precise per-primitive error
 // ---------------------------------------------------------------------------
 
+interface PromptDescriptor {
+  readonly kind: PromptKind;
+  readonly message: string;
+  readonly choices?: readonly string[];
+  readonly initial?: string;
+}
+
+function fallbackForPromptDescriptor(descriptor: PromptDescriptor): unknown {
+  switch (descriptor.kind) {
+    case "input":
+    case "editor":
+      return descriptor.initial ?? "";
+    case "confirm":
+      return false;
+    case "select":
+      return descriptor.choices?.[0] ?? "";
+  }
+}
+
+function makePrompt(descriptor: PromptDescriptor): PendingPrompt {
+  return {
+    id: `hil-${crypto.randomUUID()}`,
+    kind: descriptor.kind,
+    message: descriptor.message,
+    ...(descriptor.choices !== undefined ? { choices: descriptor.choices } : {}),
+    ...(descriptor.initial !== undefined ? { initial: descriptor.initial } : {}),
+    createdAt: Date.now(),
+  };
+}
+
+function stableHash(value: unknown): string {
+  // 128 bits is plenty for replay-key identity while keeping graph labels compact.
+  return createHash("sha256").update(JSON.stringify(value)).digest("hex").slice(0, 32);
+}
+
+function promptDescriptorHash(descriptor: PromptDescriptor): string {
+  return stableHash({
+    kind: descriptor.kind,
+    message: descriptor.message,
+    choices: descriptor.choices ?? [],
+    // Include input/editor initial text because it is visible prompt context;
+    // changing it should not replay a stale answer from the same callsite.
+    initial: descriptor.initial ?? null,
+  });
+}
+
+function promptReplayKey(descriptor: PromptDescriptor): string {
+  return `prompt:${descriptor.kind}:${promptDescriptorHash(descriptor)}:${promptCallsiteHash()}`;
+}
+
+function promptCallsiteHash(): string {
+  // Capturing an Error stack is intentional here: HIL prompts are an
+  // interactive slow path, and the author callsite is part of the replay key.
+  const frame = selectPromptCallsiteFrame(new Error().stack ?? "") ?? "unknown";
+  return stableHash(frame);
+}
+
+function hilAbortError(signal: AbortSignal): Error {
+  return signal.reason instanceof Error
+    ? signal.reason
+    : new Error("pi-workflows: HIL aborted");
+}
+
 function makeUnavailableUIContext(): WorkflowUIContext {
   const msg = (primitive: string): string =>
     `pi-workflows: HIL ctx.ui.${primitive} is unavailable because pi runtime did not provide a UI adapter`;
@@ -196,8 +265,8 @@ function makeUnavailableUIContext(): WorkflowUIContext {
 }
 
 type AskUserQuestionToolEvent =
-  | { phase: "start"; callId: string }
-  | { phase: "end"; callId: string; nameMatched: boolean };
+  | { phase: "start"; callId?: string }
+  | { phase: "end"; callId?: string; nameMatched: boolean };
 
 function stringField(value: Record<string, unknown>, keys: readonly string[]): string | undefined {
   for (const key of keys) {
@@ -217,7 +286,7 @@ function askUserQuestionToolEvent(event: unknown): AskUserQuestionToolEvent | un
   const record = event as Record<string, unknown>;
   const type = typeof record["type"] === "string" ? record["type"] : "";
   const toolName = stringField(record, ["toolName", "tool_name", "name"]);
-  const callId = stringField(record, ["toolCallId", "tool_call_id", "toolUseId", "tool_use_id", "id"]) ?? "__ask_user_question__";
+  const callId = stringField(record, ["toolCallId", "tool_call_id", "toolUseId", "tool_use_id", "id"]);
 
   if (type === "tool_execution_start" && isAskUserQuestionToolName(toolName)) {
     return { phase: "start", callId };
@@ -1094,12 +1163,52 @@ function runFailureMetadata(err: unknown, stages: readonly StageSnapshot[]): Run
   };
 }
 
+function stageReplayFields(stage: StageSnapshot): Partial<Pick<StageSnapshot, "replayKey" | "replayedFromStageId" | "replayed">> {
+  return {
+    ...(stage.replayKey !== undefined ? { replayKey: stage.replayKey } : {}),
+    ...(stage.replayedFromStageId !== undefined ? { replayedFromStageId: stage.replayedFromStageId } : {}),
+    ...(stage.replayed !== undefined ? { replayed: stage.replayed } : {}),
+  };
+}
+
+type PromptAnswerReplaySafety = "allowed" | "unavailable" | "ambiguous";
+
+function getPromptAnswerState(
+  hasReplayAnswer: boolean,
+  replaySourceId: string | undefined,
+  answerReplay: PromptAnswerReplaySafety,
+): StageSnapshot["promptAnswerState"] {
+  if (replaySourceId === undefined) return undefined;
+  if (hasReplayAnswer) return "available";
+  if (answerReplay === "ambiguous") return "ambiguous";
+  return "unavailable";
+}
+
 type ContinuationReplayDecision =
-  | { readonly kind: "execute"; readonly source?: StageSnapshot }
-  | { readonly kind: "replay"; readonly source: StageSnapshot };
+  | {
+      readonly kind: "execute";
+      readonly source?: StageSnapshot;
+      readonly parentIds: readonly string[];
+      readonly answerReplay: PromptAnswerReplaySafety;
+    }
+  | {
+      readonly kind: "replay";
+      readonly source: StageSnapshot;
+      readonly parentIds: readonly string[];
+      readonly answerReplay: PromptAnswerReplaySafety;
+    };
+
+interface ContinuationReplayInput {
+  readonly displayName: string;
+  readonly replayKey: string;
+  readonly parentIds: readonly string[];
+  readonly stageId: string;
+  readonly kind: "stage" | "prompt";
+}
 
 interface ContinuationReplayIndex {
-  decide(stageName: string, parentIds: readonly string[], stageId: string): ContinuationReplayDecision;
+  decide(input: ContinuationReplayInput): ContinuationReplayDecision;
+  markPromptAnswerReplayed(stageId: string): void;
 }
 
 function sameStringSet(left: readonly string[], right: readonly string[]): boolean {
@@ -1108,46 +1217,135 @@ function sameStringSet(left: readonly string[], right: readonly string[]): boole
   return left.every((value) => rightSet.has(value));
 }
 
+function sortedIdentity(values: readonly string[]): string {
+  return [...values].sort().join("\u0000");
+}
+
 function createContinuationReplayIndex(continuation: RunContinuationOpts | undefined): ContinuationReplayIndex {
-  if (continuation === undefined) return { decide: () => ({ kind: "execute" }) };
+  if (continuation === undefined) {
+    return {
+      decide: (input) => ({
+        kind: "execute",
+        parentIds: input.parentIds,
+        answerReplay: "unavailable",
+      }),
+      markPromptAnswerReplayed: () => {},
+    };
+  }
   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[]>();
+  const stagesByReplayIdentity = new Map<string, StageSnapshot[]>();
+  const promptDuplicateCounts = new Map<string, number>();
   for (const stage of continuation.source.stages) {
-    const stages = stagesByName.get(stage.name);
+    const identity = stage.replayKey ?? stage.name;
+    const stages = stagesByReplayIdentity.get(identity);
     if (stages === undefined) {
-      stagesByName.set(stage.name, [stage]);
+      stagesByReplayIdentity.set(identity, [stage]);
     } else {
       stages.push(stage);
     }
+    const duplicateKey = `${identity}\u0001${sortedIdentity(stage.parentIds)}`;
+    promptDuplicateCounts.set(duplicateKey, (promptDuplicateCounts.get(duplicateKey) ?? 0) + 1);
   }
 
   const consumedSourceStageIds = new Set<string>();
-  const sourceStageIdByContinuationStageId = new Map<string, string>();
+  const continuationStageIdBySourceStageId = new Map<string, string>();
+  const replayablePromptContinuationStageIds = new Set<string>();
+
+  const failTopology = (displayName: string, replayKey: string, reason: "mismatch" | "ambiguous"): never => {
+    throw new Error(`pi-workflows: insufficient_state: replay topology ${reason} for stage "${displayName}" (replayKey "${replayKey}") in source run ${continuation.source.id}`);
+  };
+
+  const translateSourceParents = (source: StageSnapshot): string[] | undefined => {
+    const parentIds: string[] = [];
+    for (const sourceParentId of source.parentIds) {
+      const continuationParentId = continuationStageIdBySourceStageId.get(sourceParentId);
+      if (continuationParentId === undefined) return undefined;
+      parentIds.push(continuationParentId);
+    }
+    return parentIds;
+  };
+
+  const allSameParentSet = (candidates: readonly { readonly parentIds: readonly string[] }[]): boolean => {
+    const first = candidates[0]?.parentIds;
+    if (first === undefined) return false;
+    return candidates.every((candidate) => sameStringSet(candidate.parentIds, first));
+  };
+
+  const hasOnlyReplayablePromptParentDrift = (
+    sourceParentIds: readonly string[],
+    provisionalParentIds: readonly string[],
+  ): boolean => {
+    const sourceParentSet = new Set(sourceParentIds);
+    const provisionalParentSet = new Set(provisionalParentIds);
+    const driftParentIds = [
+      ...sourceParentIds.filter((parentId) => !provisionalParentSet.has(parentId)),
+      ...provisionalParentIds.filter((parentId) => !sourceParentSet.has(parentId)),
+    ];
+    return driftParentIds.length > 0 && driftParentIds.every((parentId) => replayablePromptContinuationStageIds.has(parentId));
+  };
+
   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}`);
+    markPromptAnswerReplayed(stageId: string): void {
+      replayablePromptContinuationStageIds.add(stageId);
+    },
+
+    decide(input: ContinuationReplayInput): ContinuationReplayDecision {
+      const { displayName, replayKey, parentIds, stageId, kind } = input;
+      let identity = replayKey;
+      let candidates = stagesByReplayIdentity.get(replayKey)?.filter((stage) => !consumedSourceStageIds.has(stage.id)) ?? [];
+      if (candidates.length === 0) {
+        // Legacy snapshots created before replayKey existed can only be matched
+        // by display name. Current stage and prompt nodes always carry replayKey.
+        identity = displayName;
+        candidates = stagesByReplayIdentity.get(displayName)?.filter((stage) => !consumedSourceStageIds.has(stage.id) && stage.replayKey === undefined) ?? [];
       }
+      if (candidates.length === 0) {
+        return { kind: "execute", parentIds, answerReplay: "unavailable" };
+      }
+
+      const mappedCandidates = candidates
+        .map((source) => ({ source, parentIds: translateSourceParents(source) }))
+        .filter((candidate): candidate is { readonly source: StageSnapshot; readonly parentIds: string[] } => candidate.parentIds !== undefined);
 
-      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 };
+      if (mappedCandidates.length === 0) {
+        failTopology(displayName, replayKey, "mismatch");
+      }
+
+      const provisionalMatches = mappedCandidates.filter((candidate) => sameStringSet(candidate.parentIds, parentIds));
+      const hasPromptDriftMatch = kind === "prompt" &&
+        allSameParentSet(mappedCandidates) &&
+        hasOnlyReplayablePromptParentDrift(mappedCandidates[0]!.parentIds, parentIds);
+      let matches: typeof mappedCandidates | undefined;
+      if (provisionalMatches.length > 0) {
+        matches = provisionalMatches;
+      } else if (hasPromptDriftMatch) {
+        matches = mappedCandidates;
+      }
+      if (matches === undefined) {
+        return failTopology(displayName, replayKey, "mismatch");
+      }
+      if (matches.length > 1 && (kind !== "prompt" || !allSameParentSet(matches))) {
+        failTopology(displayName, replayKey, "ambiguous");
+      }
+
+      const selected = matches[0]!;
+      const duplicateKey = `${identity}\u0001${sortedIdentity(selected.source.parentIds)}`;
+      const ambiguousPromptAnswer = kind === "prompt" && (promptDuplicateCounts.get(duplicateKey) ?? 0) > 1;
+      const answerReplay: PromptAnswerReplaySafety = ambiguousPromptAnswer
+        ? "ambiguous"
+        : selected.source.status === "completed"
+          ? "allowed"
+          : "unavailable";
+      consumedSourceStageIds.add(selected.source.id);
+      continuationStageIdBySourceStageId.set(selected.source.id, stageId);
+      if (selected.source.status === "completed" && answerReplay === "allowed") {
+        return { kind: "replay", source: selected.source, parentIds: selected.parentIds, answerReplay };
+      }
+      return { kind: "execute", source: selected.source, parentIds: selected.parentIds, answerReplay };
     },
   };
 }
@@ -1211,6 +1409,9 @@ export async function run<TInputs extends Record<string, unknown>>(
 ): Promise<RunResult> {
   const activeStore = opts.store ?? defaultStore;
   const adapters = opts.adapters ?? {};
+  if (opts.usePromptNodesForUi === true && opts.ui !== undefined) {
+    console.warn("pi-workflows: usePromptNodesForUi ignores the provided RunOpts.ui adapter");
+  }
 
   // 0. maxDepth guard — reject before any store/persistence side effects.
   const depth = opts.depth ?? 0;
@@ -1309,6 +1510,12 @@ export async function run<TInputs extends Record<string, unknown>>(
   const stageById = (stageId: string): StageSnapshot | undefined =>
     runSnapshot.stages.find((stage) => stage.id === stageId);
 
+  const setStageParentIds = (stage: StageSnapshot, parentIds: readonly string[]): void => {
+    // Keep tracker and snapshot parent arrays in sync when topology is refreshed;
+    // consumers should not cache the old parentIds reference across updates.
+    stage.parentIds = Object.freeze([...parentIds]);
+  };
+
   const hasAncestor = (stage: StageSnapshot, ancestorId: string): boolean => {
     const queue = [...stage.parentIds];
     const seen = new Set<string>();
@@ -1432,26 +1639,217 @@ export async function run<TInputs extends Record<string, unknown>>(
     { once: true },
   );
 
+  const buildPromptNodeUiAdapter = (): WorkflowUIAdapter => {
+    const ask = async (descriptor: PromptDescriptor): Promise<unknown> => {
+      if (ownController.signal.aborted) {
+        return fallbackForPromptDescriptor(descriptor);
+      }
+
+      const prompt = makePrompt(descriptor);
+      const stageId = crypto.randomUUID();
+      const provisionalParentIds = tracker.onSpawn(stageId, descriptor.kind);
+      const replayKey = promptReplayKey(descriptor);
+      const replayDecision = replayIndex.decide({
+        displayName: descriptor.kind,
+        replayKey,
+        parentIds: provisionalParentIds,
+        stageId,
+        kind: "prompt",
+      });
+      const parentIds = replayDecision.parentIds;
+      if (!sameStringSet(parentIds, provisionalParentIds)) {
+        tracker.replaceParents(stageId, parentIds);
+      }
+      const replaySource = replayDecision.source;
+      const replayAnswer = replayDecision.kind === "replay"
+        // Replay decisions are only produced when continuation is present.
+        ? activeStore.getStagePromptAnswer(opts.continuation!.source.id, replayDecision.source.id)
+        : undefined;
+      const shouldReplay = replayAnswer !== undefined;
+      if (shouldReplay) {
+        replayIndex.markPromptAnswerReplayed(stageId);
+      }
+      const replaySourceId = replaySource?.id;
+      const promptAnswerStatus = getPromptAnswerState(shouldReplay, replaySourceId, replayDecision.answerReplay);
+      const stageSnapshot: StageSnapshot = {
+        id: stageId,
+        name: descriptor.kind,
+        replayKey,
+        status: shouldReplay ? "completed" : "running",
+        parentIds: Object.freeze(parentIds),
+        startedAt: prompt.createdAt,
+        promptFootprint: { ...prompt },
+        toolEvents: [],
+        attachable: !shouldReplay,
+        ...(shouldReplay ? {
+          endedAt: prompt.createdAt,
+          durationMs: 0,
+          promptAnswerState: promptAnswerStatus,
+          replayedFromStageId: replaySourceId,
+          replayed: true,
+        } : replaySourceId !== undefined ? {
+          promptAnswerState: promptAnswerStatus,
+          replayedFromStageId: replaySourceId,
+          replayed: false,
+        } : {}),
+      };
+      let finalized = false;
+      const finalizePromptStage = (status: "completed" | "failed" | "skipped"): void => {
+        if (finalized) return;
+        finalized = true;
+        stageSnapshot.status = status;
+        stageSnapshot.endedAt = Date.now();
+        stageSnapshot.durationMs = elapsedStageMs(stageSnapshot, stageSnapshot.endedAt);
+        activeStore.recordStageAttachable(runId, stageId, false);
+        activeStore.recordStageEnd(runId, stageSnapshot);
+        opts.onStageEnd?.(runId, stageSnapshot);
+        if (opts.persistence) {
+          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 } : {}),
+            ...stageReplayFields(stageSnapshot),
+          });
+        }
+        tracker.onSettle(stageId);
+      };
+
+      activeStore.recordStageStart(runId, stageSnapshot);
+      opts.onStageStart?.(runId, stageSnapshot);
+      if (opts.persistence) {
+        appendStageStart(opts.persistence, {
+          runId,
+          stageId,
+          name: stageSnapshot.name,
+          parentIds: stageSnapshot.parentIds,
+          ...stageReplayFields(stageSnapshot),
+          ts: prompt.createdAt,
+        });
+      }
+      if (shouldReplay) {
+        await Promise.resolve();
+        finalizePromptStage("completed");
+        return replayAnswer.value;
+      }
+      const accepted = activeStore.recordStagePendingPrompt(runId, stageId, prompt);
+      if (!accepted) {
+        stageSnapshot.skippedReason = "prompt-unavailable";
+        finalizePromptStage("skipped");
+        return fallbackForPromptDescriptor(descriptor);
+      }
+
+      const waiter = activeStore.awaitStagePendingPrompt(runId, stageId, prompt.id);
+      try {
+        const response = await new Promise<unknown>((resolve, reject) => {
+          const onAbort = (): void => {
+            activeStore.resolveStagePendingPrompt(
+              runId,
+              stageId,
+              prompt.id,
+              fallbackForPromptDescriptor(descriptor),
+              { recordAnswer: false },
+            );
+            reject(hilAbortError(ownController.signal));
+          };
+          if (ownController.signal.aborted) {
+            onAbort();
+            return;
+          }
+          ownController.signal.addEventListener("abort", onAbort, { once: true });
+          waiter.then(
+            (value) => {
+              ownController.signal.removeEventListener("abort", onAbort);
+              resolve(value);
+            },
+            (err: unknown) => {
+              ownController.signal.removeEventListener("abort", onAbort);
+              reject(err);
+            },
+          );
+        });
+        finalizePromptStage("completed");
+        return response;
+      } catch (err) {
+        if (ownController.signal.aborted) {
+          stageSnapshot.skippedReason = "run-aborted";
+          finalizePromptStage("skipped");
+        } else {
+          applyFailureToStage(stageSnapshot, classifyWorkflowFailure(err));
+          finalizePromptStage("failed");
+        }
+        throw err;
+      }
+    };
+
+    return {
+      async input(promptText: string): Promise<string> {
+        const response = await ask({ kind: "input", message: promptText });
+        return typeof response === "string" ? response : String(response ?? "");
+      },
+      async confirm(message: string): Promise<boolean> {
+        const response = await ask({ kind: "confirm", message });
+        return response === true;
+      },
+      async select<T extends string>(message: string, options: readonly T[]): Promise<T> {
+        if (options.length === 0) {
+          throw new Error("pi-workflows: ctx.ui.select requires at least one option");
+        }
+        const response = await ask({ kind: "select", message, choices: options });
+        if (typeof response === "string" && (options as readonly string[]).includes(response)) {
+          return response as T;
+        }
+        return options[0]!;
+      },
+      async editor(initial?: string): Promise<string> {
+        const response = await ask({
+          kind: "editor",
+          message: "Edit and save to continue.",
+          initial,
+        });
+        return typeof response === "string" ? response : initial ?? "";
+      },
+    };
+  };
+
   // 5. Build WorkflowRunContext
   const ctx: WorkflowRunContext<TInputs> = {
     inputs: resolvedInputs as TInputs,
-    ui: opts.ui ?? makeUnavailableUIContext(),
+    // Prompt nodes and caller-provided UI adapters are mutually exclusive;
+    // executor-owned prompt nodes intentionally take precedence when enabled.
+    ui: opts.usePromptNodesForUi === true ? buildPromptNodeUiAdapter() : opts.ui ?? makeUnavailableUIContext(),
 
     stage(name: string, options?: StageOptions, stageFailFastScope?: ParallelFailFastScope) {
       // a. Generate stageId
       const stageId = crypto.randomUUID();
 
-      // b. tracker.onSpawn → parentIds
-      const parentIds = tracker.onSpawn(stageId, name);
+      // b. tracker.onSpawn → provisional parentIds
+      const provisionalParentIds = tracker.onSpawn(stageId, name);
 
       // c. Create StageSnapshot as "pending"
-      const replayDecision = replayIndex.decide(name, parentIds, stageId);
+      const replayKey = `stage:${name}`;
+      const replayDecision = replayIndex.decide({
+        displayName: name,
+        replayKey,
+        parentIds: provisionalParentIds,
+        stageId,
+        kind: "stage",
+      });
+      const parentIds = replayDecision.parentIds;
+      if (!sameStringSet(parentIds, provisionalParentIds)) {
+        tracker.replaceParents(stageId, parentIds);
+      }
       const replaySource = replayDecision.kind === "replay" ? replayDecision.source : undefined;
       const shouldReplay = replaySource !== undefined;
 
       const stageSnapshot: StageSnapshot = {
         id: stageId,
         name,
+        replayKey,
         status: shouldReplay ? "completed" : "pending",
         parentIds: Object.freeze(parentIds),
         toolEvents: [],
@@ -1482,8 +1880,7 @@ export async function run<TInputs extends Record<string, unknown>>(
           stageId,
           name,
           parentIds: stageSnapshot.parentIds,
-          ...(stageSnapshot.replayedFromStageId !== undefined ? { replayedFromStageId: stageSnapshot.replayedFromStageId } : {}),
-          ...(stageSnapshot.replayed !== undefined ? { replayed: stageSnapshot.replayed } : {}),
+          ...stageReplayFields(stageSnapshot),
           ts: stageSnapshot.startedAt ?? Date.now(),
         });
       };
@@ -1505,8 +1902,7 @@ export async function run<TInputs extends Record<string, unknown>>(
               status: "completed",
               durationMs: 0,
               ...(stageSnapshot.result !== undefined ? { summary: stageSnapshot.result } : {}),
-              replayedFromStageId: replaySource.id,
-              replayed: true,
+              ...stageReplayFields(stageSnapshot),
             });
           }
           tracker.onSettle(stageId);
@@ -1520,7 +1916,7 @@ export async function run<TInputs extends Record<string, unknown>>(
         const rejectReplayMutation = (action: string): never => {
           throw new Error(`pi-workflows: replayed stage "${name}" cannot ${action}`);
         };
-        const replayContext: StageContext & Pick<InternalStageContext, "__modelFallbackMeta"> = {
+        const replayContext: InternalStageContext = {
           name,
           prompt: replayText,
           complete: replayText,
@@ -1542,11 +1938,24 @@ export async function run<TInputs extends Record<string, unknown>>(
           compact: async () => rejectReplayMutation("compact"),
           abortCompaction: () => rejectReplayMutation("abort compaction"),
           abort: async () => rejectReplayMutation("abort"),
+          __dispose: async () => {},
+          __getLastAssistantText: () => replayResult,
+          getLastAssistantText: () => replayResult,
+          __ensureSession: async () => {},
+          __sessionMeta: () => ({
+            sessionId: replaySource.sessionId,
+            sessionFile: replaySource.sessionFile,
+          }),
+          __agentSession: () => undefined,
+          __pendingMessageCount: () => 0,
           __modelFallbackMeta: () => ({
             ...(replaySource.model !== undefined ? { model: replaySource.model } : {}),
             ...(replaySource.attemptedModels !== undefined ? { attemptedModels: replaySource.attemptedModels } : {}),
             ...(replaySource.modelAttempts !== undefined ? { modelAttempts: replaySource.modelAttempts } : {}),
           }),
+          __requestPause: async () => rejectReplayMutation("pause"),
+          __resume: async () => rejectReplayMutation("resume"),
+          __isPaused: () => false,
         };
         return replayContext;
       }
@@ -1564,25 +1973,35 @@ export async function run<TInputs extends Record<string, unknown>>(
         models: opts.models,
       });
       const activeAskUserQuestionCalls = new Set<string>();
+      let activeAskUserQuestionAnonymousCalls = 0;
+      const hasActiveAskUserQuestion = (): boolean =>
+        activeAskUserQuestionCalls.size > 0 || activeAskUserQuestionAnonymousCalls > 0;
       const unsubscribeAskUserQuestionWatcher = innerCtx.subscribe((event) => {
         const toolEvent = askUserQuestionToolEvent(event);
         if (!toolEvent) return;
         if (toolEvent.phase === "start") {
-          activeAskUserQuestionCalls.add(toolEvent.callId);
+          if (toolEvent.callId !== undefined) activeAskUserQuestionCalls.add(toolEvent.callId);
+          else activeAskUserQuestionAnonymousCalls += 1;
           activeStore.recordStageAwaitingInput(runId, stageId, true);
           return;
         }
 
-        if (toolEvent.nameMatched || activeAskUserQuestionCalls.has(toolEvent.callId)) {
+        if (toolEvent.callId !== undefined && activeAskUserQuestionCalls.has(toolEvent.callId)) {
           activeAskUserQuestionCalls.delete(toolEvent.callId);
-          if (activeAskUserQuestionCalls.size === 0) {
-            activeStore.recordStageAwaitingInput(runId, stageId, false);
-          }
+        } else if (toolEvent.callId === undefined && toolEvent.nameMatched) {
+          activeAskUserQuestionAnonymousCalls = Math.max(0, activeAskUserQuestionAnonymousCalls - 1);
+        } else {
+          return;
+        }
+
+        if (!hasActiveAskUserQuestion()) {
+          activeStore.recordStageAwaitingInput(runId, stageId, false);
         }
       });
       const disposeInnerContext = async (): Promise<void> => {
         unsubscribeAskUserQuestionWatcher();
         activeAskUserQuestionCalls.clear();
+        activeAskUserQuestionAnonymousCalls = 0;
         activeStore.recordStageAwaitingInput(runId, stageId, false);
         await innerCtx.__dispose();
       };
@@ -1596,46 +2015,13 @@ 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> => {
+      const dropStageControlForCompletion = async (): Promise<void> => {
+        // Completion removes the stage from workflow-level pause/resume and
+        // dependency cascades, but must not turn the attached/reopenable chat
+        // into a read-only archive. Keep the direct live handle registered for
+        // post-completion follow-ups until the registry/store is explicitly
+        // cleared by the host.
         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
@@ -1711,6 +2097,9 @@ export async function run<TInputs extends Record<string, unknown>>(
         subscribe(listener: AgentSessionEventListener) {
           return innerCtx.subscribe(listener);
         },
+        async dispose() {
+          await releaseLiveHandle();
+        },
       };
       let stageFinalized = false;
       const finalizeStageSnapshot = (): boolean => {
@@ -1739,8 +2128,7 @@ export async function run<TInputs extends Record<string, unknown>>(
             ...(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 } : {}),
+            ...stageReplayFields(stageSnapshot),
           });
         }
 
@@ -1761,7 +2149,7 @@ export async function run<TInputs extends Record<string, unknown>>(
         markSkippedForParallelFailFast();
         finalizeStageSnapshot();
         void innerCtx.abort().catch(() => {});
-        void releaseLiveHandleWhenIdle().catch(() => {});
+        void dropStageControlForCompletion().catch(() => {});
       };
       stageFailFastScope?.activeStages.set(stageId, { skip: skipForParallelFailFast });
 
@@ -1815,6 +2203,13 @@ export async function run<TInputs extends Record<string, unknown>>(
           throw err;
         }
 
+        if (opts.continuation === undefined && stageSnapshot.startedAt === undefined) {
+          const actualParentIds = tracker.currentParents();
+          if (!sameStringSet(actualParentIds, stageSnapshot.parentIds)) {
+            tracker.replaceParents(stageId, actualParentIds);
+            setStageParentIds(stageSnapshot, actualParentIds);
+          }
+        }
         stageSnapshot.status = "running";
         stageSnapshot.startedAt = Date.now();
         activeStore.recordStageStart(runId, stageSnapshot);
@@ -1880,12 +2275,10 @@ export async function run<TInputs extends Record<string, unknown>>(
 
           finalizeStageSnapshot();
           // The stage has finished participating in workflow scheduling. Drop it
-          // 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(() => {});
+          // from run-level pause/resume and cascade-pause lookups immediately,
+          // while retaining the direct chat handle so completed nodes can be
+          // reopened and continued instead of becoming read-only archives.
+          await dropStageControlForCompletion().catch(() => {});
           limiter.release();
         }
       };