@bastani/atomic 0.9.13-alpha.1 → 0.9.13-alpha.3

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

@@ -107,7 +107,9 @@ export function createTrackedStageCaller(input: {
 			const suppressQueuedContinuation = reason === "paused-queued-user-message";
 			if (suppressQueuedContinuation) runtime.state.suppressQueuedUserMessageContinuation = true;
 			try {
-				result = (await raceAbort(runtime.innerCtx.prompt(RESUME_CONTINUATION_PROMPT), runtime.signal)) as T;
+				result = (await runtime.raceStageSessionHeartbeat(
+					raceAbort(runtime.innerCtx.prompt(RESUME_CONTINUATION_PROMPT), runtime.signal),
+				)) as T;
 			} finally {
 				if (suppressQueuedContinuation) runtime.state.suppressQueuedUserMessageContinuation = false;
 			}
@@ -134,61 +136,69 @@ export function createTrackedStageCaller(input: {
 		}
 
 		const trackStageLifecycle = !runtime.state.stageFinalized;
-		let refreshedParentIds: readonly string[] | undefined;
-		if (
-			trackStageLifecycle &&
-			!input.hasContinuation &&
-			runtime.stageSnapshot.startedAt === undefined &&
-			!input.hasScopedParents
-		) {
-			const actualParentIds = runtime.scheduler.tracker.currentParents();
-			const sameParents =
-				actualParentIds.length === runtime.stageSnapshot.parentIds.length &&
-				actualParentIds.every((value) => runtime.stageSnapshot.parentIds.includes(value));
-			if (!sameParents) {
-				runtime.scheduler.tracker.replaceParents(runtime.stageId, actualParentIds);
-				refreshedParentIds = actualParentIds;
-			}
-		}
-		if (trackStageLifecycle) {
-			const now = Date.now();
-			const hasNoExplicitModelConfig =
-				input.options?.model === undefined && input.options?.fallbackModels === undefined;
-			const promptAdapterHandlesInitialPrompt = input.adapters.prompt !== undefined;
+		let applyTerminalStageState: (() => void) | undefined;
+		// Set only alongside the successful "completed" terminal state, so a failed
+		// final durable checkpoint can replace it without touching a real
+		// failure/skip classification recorded by the catch block.
+		let terminalStateIsSuccess = false;
+		try {
+			let refreshedParentIds: readonly string[] | undefined;
 			if (
-				callOptions.eagerSession &&
-				!promptAdapterHandlesInitialPrompt &&
-				(hasNoExplicitModelConfig ||
-					(await hasExplicitFastModeCandidate({
-						model: input.options?.model,
-						fallbackModels: input.options?.fallbackModels,
-						models: runtime.opts.models,
-					})))
+				trackStageLifecycle &&
+				!input.hasContinuation &&
+				runtime.stageSnapshot.startedAt === undefined &&
+				!input.hasScopedParents
 			) {
-				try {
-					await runtime.innerCtx.__ensureSession();
-					runtime.captureStageSessionMeta();
-				} catch (err) {
-					if (!(err instanceof Error && err.message.includes("prompt adapter not configured"))) throw err;
+				const actualParentIds = runtime.scheduler.tracker.currentParents();
+				const sameParents =
+					actualParentIds.length === runtime.stageSnapshot.parentIds.length &&
+					actualParentIds.every((value) => runtime.stageSnapshot.parentIds.includes(value));
+				if (!sameParents) {
+					runtime.scheduler.tracker.replaceParents(runtime.stageId, actualParentIds);
+					refreshedParentIds = actualParentIds;
 				}
 			}
-			if (refreshedParentIds !== undefined) {
-				runtime.scheduler.setStageParentIds(runtime.stageSnapshot, refreshedParentIds);
+			if (trackStageLifecycle) {
+				const now = Date.now();
+				const hasNoExplicitModelConfig =
+					input.options?.model === undefined && input.options?.fallbackModels === undefined;
+				const promptAdapterHandlesInitialPrompt = input.adapters.prompt !== undefined;
+				if (
+					callOptions.eagerSession &&
+					!promptAdapterHandlesInitialPrompt &&
+					(hasNoExplicitModelConfig ||
+						(await hasExplicitFastModeCandidate({
+							model: input.options?.model,
+							fallbackModels: input.options?.fallbackModels,
+							models: runtime.opts.models,
+						})))
+				) {
+					try {
+						await runtime.innerCtx.__ensureSession();
+					} catch (err) {
+						if (!(err instanceof Error && err.message.includes("prompt adapter not configured"))) throw err;
+					}
+				}
+				if (refreshedParentIds !== undefined) {
+					runtime.scheduler.setStageParentIds(runtime.stageSnapshot, refreshedParentIds);
+				}
+				runtime.stageSnapshot.status = "running";
+				runtime.stageSnapshot.startedAt ??= rebasedStageStartedAt(input.options?.durableAccumulatedDurationMs, now);
+				runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
+				// Publish every graph-visible live write before this task can yield.
+				runtime.activeStore.recordStageStart(runtime.runId, runtime.stageSnapshot);
+				runtime.appendStageStartOnce();
+				const sessionMeta = runtime.innerCtx.__sessionMeta();
+				if (sessionMeta.sessionId !== undefined || sessionMeta.sessionFile !== undefined) {
+					await runtime.captureStageSessionMeta({ awaitDurable: true });
+				}
+				runtime.startStageSessionHeartbeat();
+			} else {
+				runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
 			}
-			runtime.stageSnapshot.status = "running";
-			runtime.stageSnapshot.startedAt ??= rebasedStageStartedAt(input.options?.durableAccumulatedDurationMs, now);
-			runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
-			// Publish every graph-visible live write before this task can yield.
-			runtime.activeStore.recordStageStart(runtime.runId, runtime.stageSnapshot);
-			runtime.appendStageStartOnce();
-		} else {
-			runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
-		}
 
-		runtime.mcpScope.apply();
+			runtime.mcpScope.apply();
 
-		let applyTerminalStageState: (() => void) | undefined;
-		try {
 			const abortSession = (): void => {
 				void runtime.innerCtx.abort().catch(() => {});
 			};
@@ -198,7 +208,7 @@ export function createTrackedStageCaller(input: {
 			try {
 				runtime.state.askUserQuestionObservedThisTurn = false;
 				runtime.state.chatAnswerObservedThisTurn = false;
-				result = await raceAbort(call(), runtime.signal);
+				result = await runtime.raceStageSessionHeartbeat(raceAbort(call(), runtime.signal));
 				const initialDrain = await drainResumeContinuations(result);
 				result = initialDrain.result;
 				let repeatReadinessAfterChatTurn = initialDrain.chatAnswerObserved;
@@ -226,12 +236,14 @@ export function createTrackedStageCaller(input: {
 							runtime.state.chatAnswerObservedThisTurn = false;
 							runtime.state.waitingForStageChatTurn = true;
 							try {
-								await raceAbort(
-									new Promise<void>((resolve) => {
-										resolveNextTurnEnd = resolve;
-										runtime.state.wakeWaitingForStageChatTurn = resolve;
-									}),
-									runtime.signal,
+								await runtime.raceStageSessionHeartbeat(
+									raceAbort(
+										new Promise<void>((resolve) => {
+											resolveNextTurnEnd = resolve;
+											runtime.state.wakeWaitingForStageChatTurn = resolve;
+										}),
+										runtime.signal,
+									),
 								);
 							} finally {
 								runtime.state.wakeWaitingForStageChatTurn = undefined;
@@ -254,7 +266,7 @@ export function createTrackedStageCaller(input: {
 				runtime.signal.removeEventListener("abort", abortSession);
 			}
 			await runtime.innerCtx.__closeGeneration();
-			runtime.captureStageSessionMeta();
+			await runtime.captureStageSessionMeta({ awaitDurable: true });
 			runtime.applyModelFallbackMeta(runtime.innerCtx.__modelFallbackMeta());
 			if (
 				trackStageLifecycle &&
@@ -267,6 +279,7 @@ export function createTrackedStageCaller(input: {
 			if (trackStageLifecycle && runtime.state.stageFinalized) throw runtime.parallelFailFastError();
 			if (trackStageLifecycle) {
 				const assistantText = runtime.innerCtx.__getLastAssistantText();
+				terminalStateIsSuccess = true;
 				applyTerminalStageState = () => {
 					runtime.stageSnapshot.status = "completed";
 					if (assistantText !== undefined) runtime.stageSnapshot.result = assistantText;
@@ -294,15 +307,36 @@ export function createTrackedStageCaller(input: {
 			// If finalizeStageSnapshot() throws, the limiter must still be released
 			// so the concurrency semaphore is not leaked.
 			// cross-ref: issue #1498 — durable finalization failures must not leak the stage limiter.
-			runtime.mcpScope.clear();
 			let finalizationError: { readonly thrown: true; readonly error: unknown } | undefined;
+			// Tracked separately from `finalizationError`: only a failed final durable
+			// checkpoint may rewrite an otherwise-successful terminal stage state.
+			let durableCheckpointError: { readonly thrown: true; readonly error: unknown } | undefined;
 			try {
-				await runtime.innerCtx.__closeGeneration();
-				runtime.captureStageSessionMeta();
+				runtime.mcpScope.clear();
 			} catch (err) {
 				finalizationError = { thrown: true, error: err };
 			}
+			try {
+				// Drain rather than stop: a checkpoint still in flight must be awaited
+				// here, or its late rejection is discarded with the cancelled timer.
+				await runtime.drainStageSessionHeartbeat();
+			} catch (err) {
+				durableCheckpointError = { thrown: true, error: err };
+			}
+			try {
+				await runtime.innerCtx.__closeGeneration();
+				await runtime.captureStageSessionMeta({ awaitDurable: true });
+			} catch (err) {
+				durableCheckpointError ??= { thrown: true, error: err };
+			}
+			finalizationError ??= durableCheckpointError;
 			if (trackStageLifecycle) {
+				// A stage whose final durability checkpoint failed must never be
+				// persisted as `completed` while its caller receives that rejection.
+				if (durableCheckpointError !== undefined && terminalStateIsSuccess) {
+					const failure = runtime.classifyExecutorFailure(durableCheckpointError.error);
+					applyTerminalStageState = () => applyFailureToStage(runtime.stageSnapshot, failure);
+				}
 				if (!runtime.state.stageFinalized) applyTerminalStageState?.();
 				// `finalizeStageSnapshot` calls the version-bumping `recordStageEnd`
 				// before its first yield, so the live write cannot outpace the cache.

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

@@ -52,7 +52,7 @@ export function createStageContext(input: {
 			try {
 				runtime.mcpScope.clear();
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 				runtime.applyModelFallbackMeta(innerCtx.__modelFallbackMeta());
 			}
 		}

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

@@ -28,7 +28,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 		const sessionFile = messagePreparation().sessionFile;
 		if (sessionFile === undefined) return;
 		await runtime.innerCtx.__ensureSessionFromFile(sessionFile);
-		runtime.captureStageSessionMeta();
+		await runtime.captureStageSessionMeta({ awaitDurable: true });
 	};
 	const toolExecutions = new StageToolExecutionBuffer();
 	const queuedUserMessages = new StageQueuedUserMessageBuffer();
@@ -75,7 +75,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 			await ensureMessagingSession();
 			await runtime.innerCtx.__ensureSession();
 			runtime.throwIfStageMutationBlocked();
-			runtime.captureStageSessionMeta();
+			await runtime.captureStageSessionMeta({ awaitDurable: true });
 		},
 		async sendUserMessage(text, options, beforeDelivery) {
 			runtime.throwIfStageMutationBlocked();
@@ -91,7 +91,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 				}
 				return action;
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 			}
 		},
 		async prompt(text: string) {
@@ -106,7 +106,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 					preparation,
 				);
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 			}
 			if (action !== "handled") runtime.throwIfStageMutationBlocked();
 		},
@@ -124,7 +124,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 				await runtime.innerCtx.steer(text);
 				if (queuedIntoInFlightTurn) runtime.state.resumeContinuationPending = "queued-user-message";
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 			}
 		},
 		async followUp(text: string) {
@@ -137,7 +137,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 				await runtime.innerCtx.followUp(text);
 				if (queuedIntoInFlightTurn) runtime.state.resumeContinuationPending = "queued-user-message";
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 			}
 		},
 		async pause() {
@@ -206,7 +206,7 @@ export function createStageControlHandle(runtime: LiveStageRuntime): StageContro
 				}
 				throw err;
 			} finally {
-				runtime.captureStageSessionMeta();
+				void runtime.captureStageSessionMeta();
 			}
 		},
 		subscribe(listener: AgentSessionEventListener) {

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

@@ -40,6 +40,7 @@ import {
 	type InternalStageContext,
 	type StageAdapters,
 } from "./stage-runner.js";
+import { createStageSessionHeartbeat } from "./stage-session-heartbeat.js";
 
 export function createWorkflowStageFactory(input: {
 	readonly runId: string;
@@ -151,6 +152,15 @@ export function createWorkflowStageFactory(input: {
 			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,
@@ -166,6 +176,11 @@ export function createWorkflowStageFactory(input: {
 				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 = {
@@ -234,21 +249,23 @@ export function createWorkflowStageFactory(input: {
 			await innerCtx.__dispose();
 		};
 
-		let runtime: LiveStageRuntime;
-		const captureStageSessionMeta = (checkpointOptions?: StageSessionCheckpointOptions): unknown => {
+		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) return pending;
+			if (checkpointOptions?.forceDurable === true || checkpointOptions?.awaitDurable === true) {
+				await pending;
+				return;
+			}
 			void Promise.resolve(pending).catch(() => {});
-			return undefined;
 		};
 		const releaseLiveHandle = async (): Promise<void> => {
 			if (state.liveHandleReleased) return;
 			state.liveHandleReleased = true;
+			stageSessionHeartbeat.stop();
 			runtime.dropStageControlHandle();
 			runtime.unregisterStageHandle();
 			await disposeInnerContext();
@@ -282,6 +299,7 @@ export function createWorkflowStageFactory(input: {
 		};
 
 		const finalizeStageSnapshot = async (): Promise<boolean> => {
+			stageSessionHeartbeat.stop();
 			if (state.stageFinalized) return false;
 			if (stageSnapshot.endedAt !== undefined && isTerminalStage(stageSnapshot)) {
 				state.stageFinalized = true;
@@ -371,6 +389,9 @@ export function createWorkflowStageFactory(input: {
 			dropStageControlHandle: () => {},
 			unregisterWorkflowExitCleanup: () => {},
 			captureStageSessionMeta,
+			startStageSessionHeartbeat: stageSessionHeartbeat.start,
+			drainStageSessionHeartbeat: stageSessionHeartbeat.drain,
+			raceStageSessionHeartbeat: stageSessionHeartbeat.race,
 			applyModelFallbackMeta,
 			appendStageStartOnce,
 			finalizeStageSnapshot,

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

@@ -51,7 +51,10 @@ export interface LiveStageRuntime {
 	unregisterStageHandle: () => void;
 	dropStageControlHandle: () => void;
 	unregisterWorkflowExitCleanup: () => void;
-	readonly captureStageSessionMeta: (options?: StageSessionCheckpointOptions) => unknown;
+	readonly captureStageSessionMeta: (options?: StageSessionCheckpointOptions) => Promise<void>;
+	readonly startStageSessionHeartbeat: () => void;
+	readonly drainStageSessionHeartbeat: () => Promise<void>;
+	readonly raceStageSessionHeartbeat: <T>(work: Promise<T>) => Promise<T>;
 	readonly applyModelFallbackMeta: (meta: ReturnType<InternalStageContext["__modelFallbackMeta"]>) => void;
 	readonly appendStageStartOnce: () => void;
 	readonly finalizeStageSnapshot: () => Promise<boolean>;

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

@@ -34,6 +34,8 @@ export interface RunContinuationOpts {
 
 export interface StageSessionCheckpointOptions {
 	readonly forceDurable?: boolean;
+	/** Await and surface the persistence result without bypassing duration bucketing. */
+	readonly awaitDurable?: boolean;
 }
 export interface RunOpts
 	extends Omit<

package/dist/builtin/workflows/src/runs/foreground/stage-runner-controller.ts

@@ -14,6 +14,7 @@ import type {
 	StageUserMessageContent,
 	WorkflowModelAttempt,
 	WorkflowModelCatalogPort,
+	WorkflowModelUsage,
 } from "../../shared/types.js";
 import {
 	buildModelCandidatesFromCatalog,
@@ -67,6 +68,58 @@ import {
 
 type RetryPauseResume = NonNullable<ReturnType<StageSessionPause["currentResume"]>>;
 
+type StageMessage = StageSessionRuntime["messages"][number];
+type StageAssistantMessage = Extract<StageMessage, { role: "assistant" }>;
+
+type AttemptUsageTotals = {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cost: number;
+	turns: number;
+};
+
+/**
+ * A present usage object is meaningful when any token bucket, the total token
+ * count, or the provider-reported total cost is positive. Reject the entire
+ * provider record when any aggregated value is non-finite or negative so
+ * malformed telemetry cannot poison a durable stage checkpoint.
+ */
+function hasMeaningfulUsage(usage: StageAssistantMessage["usage"] | undefined): boolean {
+	if (usage === undefined) return false;
+	const values = [usage.input, usage.output, usage.cacheRead, usage.cacheWrite, usage.totalTokens, usage.cost.total];
+	return values.every((value) => Number.isFinite(value) && value >= 0) && values.some((value) => value > 0);
+}
+
+/**
+ * Aggregate the assistant-message usage admitted since the latched attempt
+ * boundary, plus usage accrued from error assistants that a retry restoration
+ * removed from the live transcript. Returns undefined when nothing meaningful
+ * was seen, so no `usage: undefined` enters the serializable attempt record.
+ */
+function usageForAttempt(
+	messages: readonly StageMessage[],
+	start: number,
+	discarded: AttemptUsageTotals,
+): WorkflowModelUsage | undefined {
+	const totals = { ...discarded };
+	let meaningful = totals.turns > 0;
+
+	for (const message of messages.slice(start)) {
+		if (message.role !== "assistant" || !hasMeaningfulUsage(message.usage)) continue;
+		totals.input += message.usage.input;
+		totals.output += message.usage.output;
+		totals.cacheRead += message.usage.cacheRead;
+		totals.cacheWrite += message.usage.cacheWrite;
+		totals.cost += message.usage.cost.total;
+		totals.turns += 1;
+		meaningful = true;
+	}
+
+	return meaningful ? totals : undefined;
+}
+
 interface ThrownErrorRetryState {
 	readonly controller: AbortController;
 	pauseResume?: RetryPauseResume;
@@ -159,6 +212,17 @@ export class StageSessionController {
 	private sessionPromise: Promise<StageSessionRuntime> | undefined;
 	private reattachSessionFile: string | undefined;
 	private lastPromptStartIndex: number | undefined;
+	/** Message index latched once per high-level candidate prompt; never re-read later. */
+	private attemptUsageStartIndex: number | undefined;
+	/** Usage accrued from error assistants removed by retry restoration. */
+	private discardedAttemptUsage: AttemptUsageTotals = {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		cost: 0,
+		turns: 0,
+	};
 	private readonly terminatingToolCallIds = new Set<string>();
 	private latestStructuredOutputToolErrorValue: string | undefined;
 	private unsubscribeTerminateWatcher: (() => void) | undefined;
@@ -418,6 +482,7 @@ export class StageSessionController {
 				this.activeCandidateIndex = index;
 				this.selectedModel = candidate.id;
 				this.notifyModelFallbackMetaChange();
+				this.beginAttemptUsage(activeSession);
 				const { terminalScanStartIndex } = await this.promptWithThrownErrorRetry(
 					activeSession,
 					promptText,
@@ -563,6 +628,49 @@ export class StageSessionController {
 		}
 		return undefined;
 	}
+	/**
+	 * Latch a fresh high-level usage window against the active session. The
+	 * boundary is the message length before the candidate prompt, so retries,
+	 * pause/re-prompts, and recording never re-read a later index.
+	 */
+	private beginAttemptUsage(session: StageSessionRuntime): void {
+		this.attemptUsageStartIndex = session.messages.length;
+		this.discardedAttemptUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 };
+	}
+
+	/**
+	 * Consume the current window exactly once. A session-creation failure has no
+	 * window: clear the latch instead of reading a previous session's messages.
+	 */
+	private takeAttemptUsage(session: StageSessionRuntime | undefined): WorkflowModelUsage | undefined {
+		const start = this.attemptUsageStartIndex;
+		const usage =
+			session === undefined || start === undefined
+				? undefined
+				: usageForAttempt(session.messages, start, this.discardedAttemptUsage);
+		this.clearAttemptUsage();
+		return usage;
+	}
+
+	private clearAttemptUsage(): void {
+		this.attemptUsageStartIndex = undefined;
+		this.discardedAttemptUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 };
+	}
+
+	/**
+	 * Keep the usage of an error assistant that a retry restoration is about to
+	 * remove from the live transcript, so the final attempt total still counts it.
+	 */
+	private accrueDiscardedAttemptUsage(message: StageAssistantMessage): void {
+		if (!hasMeaningfulUsage(message.usage)) return;
+		this.discardedAttemptUsage.input += message.usage.input;
+		this.discardedAttemptUsage.output += message.usage.output;
+		this.discardedAttemptUsage.cacheRead += message.usage.cacheRead;
+		this.discardedAttemptUsage.cacheWrite += message.usage.cacheWrite;
+		this.discardedAttemptUsage.cost += message.usage.cost.total;
+		this.discardedAttemptUsage.turns += 1;
+	}
+
 	/**
 	 * Drop this attempt's failed input from live state before a same-candidate
 	 * retry, keeping unrelated concurrent messages. The durable transcript keeps
@@ -581,6 +689,12 @@ export class StageSessionController {
 	): StageSessionRuntime["messages"][number] | undefined {
 		const snapshotMessages = new Set(snapshot);
 		const admitted = session.messages.filter((message) => !snapshotMessages.has(message));
+		// Error assistants the filter below removes still carry real provider
+		// usage; accrue it before the splice so the retry total keeps it.
+		for (const message of admitted) {
+			if (message.role !== "assistant" || message.stopReason !== "error") continue;
+			this.accrueDiscardedAttemptUsage(message);
+		}
 		const failedAssistantIndex = admitted.findLastIndex(
 			(message) => message.role === "assistant" && message.stopReason === "error",
 		);
@@ -1054,9 +1168,12 @@ export class StageSessionController {
 				throw new Error(`atomic-workflows: stage "${this.opts.stageName}" session has been disposed`);
 			throw this.staleCreationReason(startGeneration);
 		}
-		return attachCreatedStageSession(created, this.disposed, this.opts.stageName, (result) =>
+		const session = attachCreatedStageSession(created, this.disposed, this.opts.stageName, (result) =>
 			this.attachSession(result),
 		);
+		if (session instanceof Promise) return await session;
+		await this.opts.onSessionReady?.();
+		return session;
 	}
 
 	private attachSession(created: StageSessionRuntime | StageSessionCreateResult): StageSessionRuntime {
@@ -1175,11 +1292,17 @@ export class StageSessionController {
 		const resumedSession = this.session;
 		const resumedLabel = this.selectedModel ?? workflowModelId(resumedSession.model) ?? candidates[0]!.id;
 		this.notifyModelFallbackMetaChange();
+		this.beginAttemptUsage(resumedSession);
 		try {
 			const { terminalScanStartIndex } = await this.promptWithThrownErrorRetry(resumedSession, text, sdkOptions);
 			const terminalFailure = latestTerminalAssistantFailureSince(resumedSession.messages, terminalScanStartIndex);
 			if (terminalFailure === undefined || this.capturedStructuredOutputForAttempt()) {
-				this.modelAttempts.push({ model: resumedLabel, success: true });
+				const usage = this.takeAttemptUsage(resumedSession);
+				this.modelAttempts.push({
+					model: resumedLabel,
+					success: true,
+					...(usage === undefined ? {} : { usage }),
+				});
 				this.pendingFallbackWarnings.length = 0;
 				this.resumeCurrentSession = true;
 				return true;
@@ -1187,13 +1310,24 @@ export class StageSessionController {
 			throw new WorkflowPromptModelFailure(terminalFailure);
 		} catch (err) {
 			if (this.capturedStructuredOutputForAttempt() && isRetryableModelFailure(err)) {
-				this.modelAttempts.push({ model: resumedLabel, success: true });
+				const usage = this.takeAttemptUsage(resumedSession);
+				this.modelAttempts.push({
+					model: resumedLabel,
+					success: true,
+					...(usage === undefined ? {} : { usage }),
+				});
 				this.pendingFallbackWarnings.length = 0;
 				this.resumeCurrentSession = true;
 				return true;
 			}
 			const message = errorMessage(err);
-			this.modelAttempts.push({ model: resumedLabel, success: false, error: message });
+			const usage = this.takeAttemptUsage(resumedSession);
+			this.modelAttempts.push({
+				model: resumedLabel,
+				success: false,
+				...(usage === undefined ? {} : { usage }),
+				error: message,
+			});
 			if (this.opts.signal?.aborted || !isRetryableModelFailure(err)) {
 				this.modelWarnings.push(...this.pendingFallbackWarnings);
 				this.pendingFallbackWarnings.length = 0;
@@ -1233,10 +1367,12 @@ export class StageSessionController {
 			this.recordSuccessfulAttempt(candidate);
 			return "handled";
 		}
+		const usage = this.takeAttemptUsage(this.session);
 		this.modelAttempts.push({
 			model: candidate.id,
 			success: false,
 			...modelAttemptReasoning(candidate, this.effectiveStageOptions?.thinkingLevel),
+			...(usage === undefined ? {} : { usage }),
 			error: message,
 		});
 		if (this.opts.signal?.aborted || !isRetryableModelFailure(err) || index === candidates.length - 1) {
@@ -1258,10 +1394,12 @@ export class StageSessionController {
 	}
 
 	private recordSuccessfulAttempt(candidate: WorkflowResolvedModelCandidate): void {
+		const usage = this.takeAttemptUsage(this.session);
 		this.modelAttempts.push({
 			model: candidate.id,
 			success: true,
 			...modelAttemptReasoning(candidate, this.effectiveStageOptions?.thinkingLevel),
+			...(usage === undefined ? {} : { usage }),
 		});
 		this.pendingFallbackWarnings.length = 0;
 		this.resumeCurrentSession = true;

package/dist/builtin/workflows/src/runs/foreground/stage-runner-types.ts

@@ -147,6 +147,8 @@ export interface StageRunnerOpts {
 	defaultSessionDir?: string;
 	/** Internal: notifies the executor when an in-flight fallback changes model/fast metadata. */
 	onModelFallbackMetaChange?: (meta: StageModelFallbackMeta) => void;
+	/** Internal: persist stage-session identity once the SDK has created its path. */
+	onSessionReady?: () => void | Promise<void>;
 }
 
 export interface InternalStageContext extends StageContext {