@bastani/atomic 0.9.16-alpha.2 → 0.9.16-alpha.4

package/dist/builtin/workflows/src/runs/foreground/executor-stage-control.ts

@@ -1,225 +0,0 @@
-import { isTerminalStage } from "./executor-scheduler.js";
-import type { LiveStageRuntime } from "./executor-stage-types.js";
-import type { AgentSessionEventListener, StageControlHandle } from "./stage-control-registry.js";
-import { StageQueuedUserMessageBuffer } from "./stage-queued-user-messages.js";
-import type { StageUserMessageDeliveryAction, StageUserMessagePreparation } from "./stage-runner-types.js";
-import { StageToolExecutionBuffer } from "./stage-tool-execution-buffer.js";
-
-export function createStageControlHandle(runtime: LiveStageRuntime): StageControlHandle {
-	const messagePreparation = (): StageUserMessagePreparation => {
-		const meta = runtime.innerCtx.__sessionMeta();
-		if (meta.sessionId !== undefined || meta.sessionFile !== undefined) {
-			return { beforePreparation: runtime.throwIfStageMutationBlocked };
-		}
-		if (runtime.stageSnapshot.sessionFile !== undefined) {
-			return {
-				sessionFile: runtime.stageSnapshot.sessionFile,
-				beforePreparation: runtime.throwIfStageMutationBlocked,
-			};
-		}
-		if (isTerminalStage(runtime.stageSnapshot)) {
-			throw new Error(
-				`atomic-workflows: cannot message stage "${runtime.name}" because no retained session metadata is available.`,
-			);
-		}
-		return { beforePreparation: runtime.throwIfStageMutationBlocked };
-	};
-	const ensureMessagingSession = async (): Promise<void> => {
-		const sessionFile = messagePreparation().sessionFile;
-		if (sessionFile === undefined) return;
-		await runtime.innerCtx.__ensureSessionFromFile(sessionFile);
-		await runtime.captureStageSessionMeta({ awaitDurable: true });
-	};
-	const toolExecutions = new StageToolExecutionBuffer();
-	const queuedUserMessages = new StageQueuedUserMessageBuffer();
-	// One lifetime subscription feeds both runtime projections, so the queue stays
-	// current while no stage chat is mounted.
-	const unsubscribeStageEvents = runtime.innerCtx.subscribe((event) => {
-		toolExecutions.record(event);
-		queuedUserMessages.record(event);
-	});
-
-	return {
-		runId: runtime.runId,
-		stageId: runtime.stageId,
-		stageName: runtime.name,
-		get status() {
-			return runtime.stageSnapshot.status;
-		},
-		get sessionId() {
-			return runtime.innerCtx.__sessionMeta().sessionId ?? runtime.stageSnapshot.sessionId;
-		},
-		get sessionFile() {
-			return runtime.innerCtx.__sessionMeta().sessionFile ?? runtime.stageSnapshot.sessionFile;
-		},
-		get isStreaming() {
-			return runtime.innerCtx.isStreaming;
-		},
-		get isDisposed() {
-			return runtime.state.liveHandleReleased;
-		},
-		get messages() {
-			return runtime.innerCtx.messages;
-		},
-		get agentSession() {
-			return runtime.innerCtx.__agentSession();
-		},
-		pendingToolExecutionEvents() {
-			return toolExecutions.replayEvents();
-		},
-		queuedUserMessages() {
-			return queuedUserMessages.snapshot();
-		},
-		async ensureAttached() {
-			runtime.throwIfStageMutationBlocked();
-			await ensureMessagingSession();
-			await runtime.innerCtx.__ensureSession();
-			runtime.throwIfStageMutationBlocked();
-			await runtime.captureStageSessionMeta({ awaitDurable: true });
-		},
-		async sendUserMessage(text, options, beforeDelivery) {
-			runtime.throwIfStageMutationBlocked();
-			const preparation = messagePreparation();
-			const admitDelivery = (): void => {
-				runtime.throwIfStageMutationBlocked();
-				beforeDelivery?.();
-			};
-			try {
-				const action = await runtime.innerCtx.__sendUserMessage(text, options, admitDelivery, preparation);
-				if (action === "steer" || action === "followUp") {
-					runtime.state.resumeContinuationPending = "queued-user-message";
-				}
-				return action;
-			} finally {
-				void runtime.captureStageSessionMeta();
-			}
-		},
-		async prompt(text: string) {
-			runtime.throwIfStageMutationBlocked();
-			const preparation = messagePreparation();
-			let action: StageUserMessageDeliveryAction | undefined;
-			try {
-				action = await runtime.innerCtx.__sendUserMessage(
-					text,
-					undefined,
-					runtime.throwIfStageMutationBlocked,
-					preparation,
-				);
-			} finally {
-				void runtime.captureStageSessionMeta();
-			}
-			if (action !== "handled") runtime.throwIfStageMutationBlocked();
-		},
-		async steer(text: string) {
-			runtime.throwIfStageMutationBlocked();
-			await ensureMessagingSession();
-			runtime.throwIfStageMutationBlocked();
-			// A user message queued into an in-flight turn should nudge the stage
-			// back to its objective once that turn ends: arm the pending flag so
-			// drainResumeContinuations injects RESUME_CONTINUATION_PROMPT after the
-			// tracked call resolves. Idle deliveries start a fresh user turn and
-			// need no continuation nudge, so only arm while streaming.
-			const queuedIntoInFlightTurn = runtime.innerCtx.isStreaming;
-			try {
-				await runtime.innerCtx.steer(text);
-				if (queuedIntoInFlightTurn) runtime.state.resumeContinuationPending = "queued-user-message";
-			} finally {
-				void runtime.captureStageSessionMeta();
-			}
-		},
-		async followUp(text: string) {
-			runtime.throwIfStageMutationBlocked();
-			await ensureMessagingSession();
-			runtime.throwIfStageMutationBlocked();
-			// Same in-flight continuation arming as steer(): see comment above.
-			const queuedIntoInFlightTurn = runtime.innerCtx.isStreaming;
-			try {
-				await runtime.innerCtx.followUp(text);
-				if (queuedIntoInFlightTurn) runtime.state.resumeContinuationPending = "queued-user-message";
-			} finally {
-				void runtime.captureStageSessionMeta();
-			}
-		},
-		async pause() {
-			runtime.throwIfStageMutationBlocked();
-			const statusBeforePause = runtime.stageSnapshot.status;
-			if (statusBeforePause === "pending" || statusBeforePause === "running" || runtime.innerCtx.isStreaming) {
-				await runtime.innerCtx.__requestPause();
-			}
-			const changed = runtime.activeStore.recordStagePaused(runtime.runId, runtime.stageId);
-			if (changed) {
-				runtime.scheduler.ensureReleaseBarrier(runtime.stageId);
-				await runtime.scheduler.cascadePauseFrom(runtime.stageId);
-				const run = runtime.activeStore.runs().find((candidate) => candidate.id === runtime.runId);
-				const stillActive =
-					run?.stages.some((stage) => stage.status === "running" && stage.id !== runtime.stageId) ?? false;
-				if (!stillActive) runtime.activeStore.recordRunPaused(runtime.runId);
-			}
-			// Graceful pause/quit is an exact durability boundary. Force the latest
-			// pause-adjusted stage elapsed time even inside the normal 30s bucket.
-			await runtime.captureStageSessionMeta({ forceDurable: true });
-		},
-		async resume(message?: string, beforeResume?: () => void) {
-			runtime.throwIfStageMutationBlocked();
-			await ensureMessagingSession();
-			runtime.throwIfStageMutationBlocked();
-			const wasPausedBeforeResume = runtime.innerCtx.__isPaused();
-			const resumesIdleStageChat = wasPausedBeforeResume && runtime.state.waitingForStageChatTurn;
-			const hasMessage = typeof message === "string" && message.trim().length > 0;
-			const resumeMessage = hasMessage ? message : undefined;
-			const queuedResumeContinuation = wasPausedBeforeResume && !resumesIdleStageChat;
-			const previousResumeContinuation = runtime.state.resumeContinuationPending;
-			let resumeContinuationChanged = false;
-			let wakeReleasedIdleStageChat = false;
-			try {
-				await runtime.innerCtx.__resume(
-					resumesIdleStageChat ? undefined : resumeMessage,
-					({ releasedQueuedMessages, runnerOwnedDeliveryPending }) => {
-						if (runnerOwnedDeliveryPending) return;
-						if (!releasedQueuedMessages && !queuedResumeContinuation) return;
-						resumeContinuationChanged = true;
-						runtime.state.resumeContinuationPending = releasedQueuedMessages
-							? "paused-queued-user-message"
-							: previousResumeContinuation === false
-								? "resume"
-								: previousResumeContinuation;
-						wakeReleasedIdleStageChat = resumesIdleStageChat && releasedQueuedMessages;
-					},
-					beforeResume,
-				);
-				const changed = runtime.activeStore.recordStageResumed(runtime.runId, runtime.stageId);
-				if (changed) {
-					runtime.scheduler.releaseStageBarrier(runtime.stageId);
-					await runtime.scheduler.cascadeResumeFrom(runtime.stageId);
-					// Preserve manual per-stage semantics: once this acknowledged
-					// resume succeeds, the run is active even if sibling stages remain paused.
-					runtime.activeStore.recordRunResumed(runtime.runId, undefined, { source: "stage_control" });
-				}
-				if (wakeReleasedIdleStageChat) runtime.state.wakeWaitingForStageChatTurn?.();
-				if (resumesIdleStageChat && hasMessage) {
-					runtime.throwIfStageMutationBlocked();
-					return await runtime.innerCtx.__sendUserMessage(message, undefined, runtime.throwIfStageMutationBlocked);
-				}
-			} catch (err) {
-				if (resumeContinuationChanged) {
-					runtime.state.resumeContinuationPending = previousResumeContinuation;
-				}
-				throw err;
-			} finally {
-				void runtime.captureStageSessionMeta();
-			}
-		},
-		subscribe(listener: AgentSessionEventListener) {
-			return runtime.innerCtx.subscribe(listener);
-		},
-		subscribeDeliveryActivity(listener) {
-			return runtime.innerCtx.__subscribeDeliveryActivity(listener);
-		},
-		async dispose() {
-			unsubscribeStageEvents();
-			toolExecutions.clear();
-			queuedUserMessages.clear();
-			await runtime.releaseLiveHandle();
-		},
-	};
-}

package/dist/builtin/workflows/src/runs/foreground/executor-stage-factory.ts

@@ -1,473 +0,0 @@
-import { runCallback, runSynchronousCallback } from "@bastani/atomic";
-import type { GraphFrontierTracker } from "../../engine/graph-inference.js";
-import type { EngineStageRuntimeOptions } from "../../engine/options.js";
-import type { RunBudgetController } from "../../engine/run-budget.js";
-import { appendStageEnd, appendStageStart } from "../../shared/persistence-session-entries.js";
-import { buildStagePromptAdapter } from "../../shared/stage-prompt.js";
-import { stageUiBroker } from "../../shared/stage-ui-broker.js";
-import type { Store } from "../../shared/store.js";
-import type { StageSnapshot } from "../../shared/store-types.js";
-import { elapsedStageMs } from "../../shared/timing.js";
-import type { StageOptions } from "../../shared/types.js";
-import type { WorkflowFailure } from "../../shared/workflow-failures.js";
-import type { ConcurrencyLimiter } from "../shared/concurrency.js";
-import type { ContinuationReplayIndex } from "./executor-continuation.js";
-import { sameStringSet } from "./executor-continuation.js";
-import {
-	type GitWorktreeSetupCache,
-	stageOptionsWithGitWorktree,
-	stageOptionsWithInputDefaults,
-} from "./executor-direct-helpers.js";
-import type { WorkflowExitManager } from "./executor-exit-manager.js";
-import { askUserQuestionToolEvent, toolResultHasChatAnswer } from "./executor-hil.js";
-import { stageReplayFields } from "./executor-lifecycle.js";
-import { createQueuedUserMessageConsumptionWatcher } from "./executor-queued-user-message.js";
-import type { StageScheduler } from "./executor-scheduler.js";
-import { isTerminalStage } from "./executor-scheduler.js";
-import { createTrackedStageCaller } from "./executor-stage-call.js";
-import { createStageContext } from "./executor-stage-context.js";
-import { createStageControlHandle } from "./executor-stage-control.js";
-import { createReplayStageContext } from "./executor-stage-replay.js";
-import type {
-	LiveStageMutableState,
-	LiveStageRuntime,
-	StageContextWithMeta,
-	StageMcpScope,
-} from "./executor-stage-types.js";
-import type { ParallelFailFastScope, StageSessionCheckpointOptions } from "./executor-types.js";
-import type { StageControlRegistry } from "./stage-control-registry.js";
-import {
-	createStageContext as createInnerStageContext,
-	type InternalStageContext,
-	type StageAdapters,
-} from "./stage-runner.js";
-import { createStageSessionHeartbeat } from "./stage-session-heartbeat.js";
-
-export function createWorkflowStageFactory(input: {
-	readonly runId: string;
-	readonly workflowIntercomGroup: string;
-	readonly activeStore: Store;
-	readonly opts: EngineStageRuntimeOptions;
-	readonly adapters: StageAdapters;
-	readonly signal: AbortSignal;
-	readonly tracker: GraphFrontierTracker;
-	readonly scheduler: StageScheduler;
-	readonly replayIndex: ContinuationReplayIndex;
-	readonly limiter: ConcurrencyLimiter;
-	readonly inputRuntimeDefaults: Partial<StageOptions>;
-	readonly workflowInvocationCwd: string;
-	readonly gitWorktreeSetupCache: GitWorktreeSetupCache;
-	readonly stageRegistry: StageControlRegistry;
-	readonly exit: WorkflowExitManager;
-	readonly budget: RunBudgetController;
-	readonly classifyExecutorFailure: (error: unknown) => WorkflowFailure;
-	readonly createMcpScope: (stageId: string, options: StageOptions | undefined) => StageMcpScope;
-}): (name: string, options?: StageOptions, stageFailFastScope?: ParallelFailFastScope) => StageContextWithMeta {
-	return (name: string, options?: StageOptions, stageFailFastScope?: ParallelFailFastScope): StageContextWithMeta => {
-		input.exit.throwIfWorkflowExitSelected();
-		options = stageOptionsWithGitWorktree(
-			stageOptionsWithInputDefaults(options, input.inputRuntimeDefaults),
-			input.workflowInvocationCwd,
-			input.gitWorktreeSetupCache,
-		);
-		const stageId = options?.durableStageId ?? crypto.randomUUID();
-		const provisionalParentIds = input.tracker.onSpawn(stageId, name);
-		const scopedParentIds =
-			options?.durableParentIds ??
-			(input.opts.continuation === undefined ? stageFailFastScope?.parentIds : undefined);
-		const initialParentIds = scopedParentIds === undefined ? provisionalParentIds : [...scopedParentIds];
-		if (scopedParentIds !== undefined && !sameStringSet(scopedParentIds, provisionalParentIds)) {
-			input.tracker.replaceParents(stageId, scopedParentIds);
-		}
-
-		const replayKey = options?.durableReplayKey ?? `stage:${name}`;
-		const replayDecision = input.replayIndex.decide({
-			displayName: name,
-			replayKey,
-			parentIds: initialParentIds,
-			stageId,
-			kind: "stage",
-		});
-		const parentIds = options?.durableParentIds ?? replayDecision.parentIds;
-		if (!sameStringSet(parentIds, provisionalParentIds)) input.tracker.replaceParents(stageId, parentIds);
-		const replaySource = replayDecision.kind === "replay" ? replayDecision.source : undefined;
-		const executeReplaySource = replayDecision.kind === "execute" ? replayDecision.source : undefined;
-		const shouldReplay = replaySource !== undefined;
-
-		const stageSnapshot: StageSnapshot = {
-			id: stageId,
-			name,
-			replayKey,
-			status: shouldReplay ? "completed" : "pending",
-			parentIds: Object.freeze(parentIds),
-			toolEvents: [],
-			...(shouldReplay
-				? {
-						startedAt: Date.now(),
-						endedAt: Date.now(),
-						durationMs: 0,
-						...(replaySource.result !== undefined ? { result: replaySource.result } : {}),
-						...(replaySource.sessionId !== undefined ? { sessionId: replaySource.sessionId } : {}),
-						...(replaySource.sessionFile !== undefined ? { sessionFile: replaySource.sessionFile } : {}),
-						replayedFromStageId: replaySource.id,
-						replayed: true,
-					}
-				: {}),
-			...(options?.mcp !== undefined
-				? { mcpScope: { allow: options.mcp.allow ?? null, deny: options.mcp.deny ?? null } }
-				: {}),
-			attachable: !shouldReplay,
-		};
-
-		if (shouldReplay) {
-			return createReplayStageContext({
-				runId: input.runId,
-				name,
-				stageId,
-				stageSnapshot,
-				replaySource,
-				activeStore: input.activeStore,
-				opts: input.opts,
-				tracker: input.tracker,
-				registerWorkflowExitCleanup: input.exit.registerWorkflowExitCleanup,
-				workflowExitSkippedReason: input.exit.workflowExitSkippedReason,
-				throwIfWorkflowExitSelected: input.exit.throwIfWorkflowExitSelected,
-			}) as StageContextWithMeta;
-		}
-
-		const stageOptionsForContext: StageOptions | undefined =
-			executeReplaySource?.sessionFile === undefined
-				? options
-				: {
-						...(options ?? {}),
-						context: options?.context ?? "fork",
-						forkFromSessionFile: options?.forkFromSessionFile ?? executeReplaySource.sessionFile,
-					};
-
-		const applyModelFallbackMeta = (meta: ReturnType<InternalStageContext["__modelFallbackMeta"]>): void => {
-			if (meta.model !== undefined) stageSnapshot.model = meta.model;
-			if (meta.fastMode !== undefined) {
-				if (meta.fastMode) stageSnapshot.fastMode = true;
-				else delete stageSnapshot.fastMode;
-			}
-			if (meta.attemptedModels !== undefined) stageSnapshot.attemptedModels = meta.attemptedModels;
-			if (meta.modelAttempts !== undefined) stageSnapshot.modelAttempts = meta.modelAttempts;
-		};
-
-		let runtime!: LiveStageRuntime;
-		const stageSessionHeartbeat = createStageSessionHeartbeat(async () => {
-			try {
-				await runtime.captureStageSessionMeta({ awaitDurable: true });
-			} catch (error) {
-				await runtime.innerCtx.abort().catch(() => {});
-				throw error;
-			}
-		});
-		const innerCtx = createInnerStageContext({
-			stageId,
-			stageName: name,
-			adapters: input.adapters,
-			runId: input.runId,
-			workflowIntercomGroup: input.workflowIntercomGroup,
-			signal: input.signal,
-			stageOptions: stageOptionsForContext,
-			models: input.opts.models,
-			executionMode: input.opts.executionMode,
-			defaultSessionDir: input.opts.defaultSessionDir,
-			onModelFallbackMetaChange(meta) {
-				applyModelFallbackMeta(meta);
-				if (stageSnapshot.status === "running") input.activeStore.recordStageStart(input.runId, stageSnapshot);
-			},
-			async onSessionReady() {
-				if (stageSnapshot.status !== "running") return;
-				await runtime.captureStageSessionMeta({ awaitDurable: true });
-				runtime.startStageSessionHeartbeat();
-			},
-		});
-
-		const state: LiveStageMutableState = {
-			activeAskUserQuestionAnonymousCalls: 0,
-			askUserQuestionObservedThisTurn: false,
-			chatAnswerObservedThisTurn: false,
-			resumeContinuationPending: false,
-			suppressQueuedUserMessageContinuation: false,
-			waitingForStageChatTurn: false,
-			wakeWaitingForStageChatTurn: undefined,
-			liveHandleReleased: false,
-			stageClosedByWorkflowExit: false,
-			stageFinalized: false,
-			skippedForParallelFailFast: false,
-			stageControlDropped: false,
-		};
-		const activeAskUserQuestionCalls = new Set<string>();
-		const hasActiveAskUserQuestion = (): boolean =>
-			activeAskUserQuestionCalls.size > 0 || state.activeAskUserQuestionAnonymousCalls > 0;
-		const unsubscribeAskUserQuestionWatcher = innerCtx.subscribe((event) => {
-			const toolEvent = askUserQuestionToolEvent(event);
-			if (!toolEvent) return;
-			if (toolEvent.phase === "start") {
-				state.askUserQuestionObservedThisTurn = true;
-				if (toolEvent.callId !== undefined) activeAskUserQuestionCalls.add(toolEvent.callId);
-				else state.activeAskUserQuestionAnonymousCalls += 1;
-				const adapter = buildStagePromptAdapter(
-					toolEvent.callId ?? `ask-user-question-${stageId}`,
-					"ask_user_question",
-					toolEvent.args,
-					Date.now(),
-				);
-				if (adapter) stageUiBroker.provideStagePrompt(input.runId, stageId, adapter);
-				input.activeStore.recordStageAwaitingInput(input.runId, stageId, true);
-				return;
-			}
-			if (toolEvent.callId !== undefined && activeAskUserQuestionCalls.has(toolEvent.callId)) {
-				activeAskUserQuestionCalls.delete(toolEvent.callId);
-			} else if (toolEvent.callId === undefined && toolEvent.nameMatched) {
-				state.activeAskUserQuestionAnonymousCalls = Math.max(0, state.activeAskUserQuestionAnonymousCalls - 1);
-			} else {
-				return;
-			}
-			if (toolResultHasChatAnswer((event as Record<string, unknown>).result))
-				state.chatAnswerObservedThisTurn = true;
-			if (!hasActiveAskUserQuestion()) {
-				input.activeStore.recordStageAwaitingInput(input.runId, stageId, false);
-				stageUiBroker.clearStagePrompt(input.runId, stageId);
-			}
-		});
-		const unsubscribeQueuedUserMessageWatcher = innerCtx.subscribe(
-			createQueuedUserMessageConsumptionWatcher(() => {
-				if (!state.suppressQueuedUserMessageContinuation) {
-					state.resumeContinuationPending = "queued-user-message";
-				}
-			}),
-		);
-
-		const disposeInnerContext = async (): Promise<void> => {
-			unsubscribeAskUserQuestionWatcher();
-			unsubscribeQueuedUserMessageWatcher();
-			activeAskUserQuestionCalls.clear();
-			state.activeAskUserQuestionAnonymousCalls = 0;
-			input.activeStore.recordStageAwaitingInput(input.runId, stageId, false);
-			stageUiBroker.clearStagePrompt(input.runId, stageId);
-			await innerCtx.__dispose();
-		};
-
-		const captureStageSessionMeta = async (checkpointOptions?: StageSessionCheckpointOptions): Promise<void> => {
-			const meta = innerCtx.__sessionMeta();
-			if (meta.sessionId !== undefined) stageSnapshot.sessionId = meta.sessionId;
-			if (meta.sessionFile !== undefined) stageSnapshot.sessionFile = meta.sessionFile;
-			if (meta.sessionId !== undefined || meta.sessionFile !== undefined)
-				input.activeStore.recordStageSession(input.runId, stageId, meta);
-			const pending = input.opts.onStageSession?.(input.runId, stageSnapshot, checkpointOptions);
-			if (checkpointOptions?.forceDurable === true || checkpointOptions?.awaitDurable === true) {
-				await pending;
-				return;
-			}
-			void Promise.resolve(pending).catch(() => {});
-		};
-		const releaseLiveHandle = async (): Promise<void> => {
-			if (state.liveHandleReleased) return;
-			state.liveHandleReleased = true;
-			stageSessionHeartbeat.stop();
-			runtime.dropStageControlHandle();
-			runtime.unregisterStageHandle();
-			await disposeInnerContext();
-		};
-		const dropStageControlForCompletion = async (): Promise<void> => {
-			runtime.dropStageControlHandle();
-			// A registry/session clear can release a confirmed-paused handle before
-			// its owner resumes. The handle is then ownerless, so normal completion
-			// must make the second disposal call that drains the controller's deferred
-			// cleanup instead of retaining the session and its listeners forever.
-			if (state.liveHandleReleased) await disposeInnerContext();
-		};
-		const throwIfStageMutationBlocked = (): void => {
-			if (state.stageClosedByWorkflowExit) {
-				input.exit.throwIfWorkflowExitSelected();
-				throw new Error(`atomic-workflows: stage "${name}" skipped by workflow exit`);
-			}
-			input.exit.throwIfWorkflowExitSelected();
-			if (input.signal.aborted) {
-				throw input.signal.reason ?? new DOMException("workflow killed", "AbortError");
-			}
-		};
-
-		let stageStartEntryAppended = false;
-		const appendStageStartOnce = (): void => {
-			if (!input.opts.persistence || stageStartEntryAppended) return;
-			stageStartEntryAppended = true;
-			appendStageStart(input.opts.persistence, {
-				runId: input.runId,
-				stageId,
-				name,
-				parentIds: stageSnapshot.parentIds,
-				...stageReplayFields(stageSnapshot),
-				ts: stageSnapshot.startedAt ?? Date.now(),
-			});
-		};
-
-		const finalizeStageSnapshot = async (): Promise<boolean> => {
-			stageSessionHeartbeat.stop();
-			if (state.stageFinalized) return false;
-			if (stageSnapshot.endedAt !== undefined && isTerminalStage(stageSnapshot)) {
-				state.stageFinalized = true;
-				runtime.unregisterWorkflowExitCleanup();
-				stageFailFastScope?.activeStages.delete(stageId);
-				input.tracker.onSettle(stageId);
-				return false;
-			}
-			if (state.skippedForParallelFailFast) {
-				stageSnapshot.status = "skipped";
-				stageSnapshot.skippedReason = "fail-fast";
-			}
-			state.stageFinalized = true;
-			runtime.unregisterWorkflowExitCleanup();
-			stageSnapshot.endedAt = Date.now();
-			stageSnapshot.durationMs = elapsedStageMs(stageSnapshot, stageSnapshot.endedAt);
-			applyModelFallbackMeta(innerCtx.__modelFallbackMeta());
-			input.activeStore.recordStageEnd(input.runId, stageSnapshot);
-			stageUiBroker.cancelStagePrompt(
-				input.runId,
-				stageId,
-				new Error(`atomic-workflows: stage ${stageId} completed with pending custom UI`),
-			);
-			if (input.opts.onStageEnd) {
-				await runCallback(
-					{ kind: "workflow.stage_adapter", name: `onStageEnd:${name}`, runId: input.runId, stageId },
-					() => input.opts.onStageEnd!(input.runId, stageSnapshot),
-				);
-			}
-			if (input.opts.persistence) {
-				appendStageStartOnce();
-				appendStageEnd(input.opts.persistence, {
-					runId: input.runId,
-					stageId,
-					status: stageSnapshot.status,
-					durationMs: stageSnapshot.durationMs,
-					...(stageSnapshot.error !== undefined ? { error: stageSnapshot.error } : {}),
-					...(stageSnapshot.failureKind !== undefined ? { failureKind: stageSnapshot.failureKind } : {}),
-					...(stageSnapshot.failureCode !== undefined ? { failureCode: stageSnapshot.failureCode } : {}),
-					...(stageSnapshot.failureRecoverability !== undefined
-						? { failureRecoverability: stageSnapshot.failureRecoverability }
-						: {}),
-					...(stageSnapshot.failureDisposition !== undefined
-						? { failureDisposition: stageSnapshot.failureDisposition }
-						: {}),
-					...(stageSnapshot.failureMessage !== undefined ? { failureMessage: stageSnapshot.failureMessage } : {}),
-					...(stageSnapshot.retryAfterMs !== undefined ? { retryAfterMs: stageSnapshot.retryAfterMs } : {}),
-					...(stageSnapshot.skippedReason !== undefined ? { skippedReason: stageSnapshot.skippedReason } : {}),
-					...(stageSnapshot.sessionId !== undefined ? { sessionId: stageSnapshot.sessionId } : {}),
-					...(stageSnapshot.sessionFile !== undefined ? { sessionFile: stageSnapshot.sessionFile } : {}),
-					...(stageSnapshot.modelAttempts !== undefined ? { modelAttempts: stageSnapshot.modelAttempts } : {}),
-					...(stageSnapshot.result !== undefined && stageSnapshot.status === "completed"
-						? { summary: stageSnapshot.result }
-						: {}),
-					...stageReplayFields(stageSnapshot),
-				});
-			}
-			stageFailFastScope?.activeStages.delete(stageId);
-			input.tracker.onSettle(stageId);
-			return true;
-		};
-
-		const markSkippedForParallelFailFast = (): void => {
-			// Defer graph-visible fields until `finalizeStageSnapshot`, directly
-			// before its version-bumping `recordStageEnd` call.
-			state.skippedForParallelFailFast = true;
-		};
-		const parallelFailFastError = (): unknown =>
-			stageFailFastScope?.firstFailure ?? new Error("atomic-workflows: skipped after parallel fail-fast");
-
-		runtime = {
-			runId: input.runId,
-			stageId,
-			name,
-			stageSnapshot,
-			innerCtx,
-			activeStore: input.activeStore,
-			opts: input.opts,
-			stageRegistry: input.stageRegistry,
-			scheduler: input.scheduler,
-			signal: input.signal,
-			exit: input.exit,
-			budget: input.budget,
-			classifyExecutorFailure: input.classifyExecutorFailure,
-			mcpScope: input.createMcpScope(stageId, options),
-			...(stageFailFastScope !== undefined ? { stageFailFastScope } : {}),
-			state,
-			unregisterStageHandle: () => {},
-			dropStageControlHandle: () => {},
-			unregisterWorkflowExitCleanup: () => {},
-			captureStageSessionMeta,
-			startStageSessionHeartbeat: stageSessionHeartbeat.start,
-			drainStageSessionHeartbeat: stageSessionHeartbeat.drain,
-			raceStageSessionHeartbeat: stageSessionHeartbeat.race,
-			applyModelFallbackMeta,
-			appendStageStartOnce,
-			finalizeStageSnapshot,
-			releaseLiveHandle,
-			dropStageControlForCompletion,
-			markSkippedForParallelFailFast,
-			parallelFailFastError,
-			throwIfStageMutationBlocked,
-		};
-
-		const handle = createStageControlHandle(runtime);
-		runtime.dropStageControlHandle = (): void => {
-			if (state.stageControlDropped) return;
-			state.stageControlDropped = true;
-			input.activeStore.recordStageAttachable(input.runId, stageId, false);
-			input.stageRegistry.detachControl(input.runId, stageId, handle);
-		};
-		runtime.unregisterStageHandle = input.stageRegistry.register(handle);
-
-		input.activeStore.recordStageStart(input.runId, stageSnapshot);
-		if (input.opts.onStageStart) {
-			runSynchronousCallback(
-				{ kind: "workflow.stage_adapter", name: `onStageStart:${name}`, runId: input.runId, stageId },
-				() => input.opts.onStageStart!(input.runId, stageSnapshot),
-			);
-		}
-		const blockedBy = input.scheduler.blockingAncestorFor(stageSnapshot);
-		if (blockedBy !== undefined) input.scheduler.blockStageUntilCascadeRelease(stageSnapshot, blockedBy);
-
-		// Parallel fail-fast and workflow-exit cleanup can both target a live stage.
-		// The first terminal path owns the snapshot: finalization unregisters
-		// workflow-exit cleanup and removes the stage from the fail-fast active set.
-		// Later paths must not overwrite the terminal skippedReason; they only abort
-		// and release idempotent live handles.
-		const skipForParallelFailFast = async (): Promise<void> => {
-			if (isTerminalStage(stageSnapshot)) return;
-			markSkippedForParallelFailFast();
-			innerCtx.__sealGeneration();
-			await innerCtx.abort().catch(() => {});
-			await innerCtx.__closeGeneration();
-			await finalizeStageSnapshot();
-			await dropStageControlForCompletion().catch(() => {});
-		};
-		stageFailFastScope?.activeStages.set(stageId, { skip: skipForParallelFailFast });
-		runtime.unregisterWorkflowExitCleanup = input.exit.registerWorkflowExitCleanup(stageId, {
-			async skipForWorkflowExit(reason?: string): Promise<void> {
-				state.stageClosedByWorkflowExit = true;
-				innerCtx.__sealGeneration();
-				await innerCtx.abort().catch(() => {});
-				await innerCtx.__closeGeneration();
-				if (!isTerminalStage(stageSnapshot)) {
-					stageSnapshot.status = "skipped";
-					stageSnapshot.skippedReason = input.exit.workflowExitSkippedReason(reason);
-					await finalizeStageSnapshot();
-				}
-				await releaseLiveHandle().catch(() => {});
-			},
-		});
-
-		const runTrackedStageCall = createTrackedStageCaller({
-			runtime,
-			limiter: input.limiter,
-			options,
-			adapters: input.adapters,
-			hasContinuation: input.opts.continuation !== undefined,
-			hasScopedParents: stageFailFastScope?.parentIds !== undefined,
-		});
-		return createStageContext({ runtime, runTrackedStageCall });
-	};
-}