rivetkit 2.3.0 → 2.3.1

package/src/workflow/context.ts

@@ -1,5 +1,3 @@
-// @ts-nocheck
-
 import type {
 	BranchConfig,
 	BranchOutput,
@@ -57,7 +55,177 @@ type WorkflowActorQueueNextBatchOptions<
 type WorkflowActorQueueNextBatchOptionsFallback<TCompletable extends boolean> =
 	Omit<QueueNextBatchOptions<string, TCompletable>, "signal">;
 
-type ActorWorkflowLoopConfig<
+// Step run callbacks receive a WorkflowStepContext, which is the only place
+// actor data (state/db/vars/client) may be touched.
+type WorkflowStepRun<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = (
+	step: WorkflowStepContext<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>,
+) => Promise<T>;
+
+// Step rollback callbacks compensate a committed step, so they also run with a
+// WorkflowStepContext to mutate actor data.
+type WorkflowStepRollback<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = (
+	step: WorkflowStepContext<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>,
+	output: T,
+) => Promise<void>;
+
+// Orchestration callbacks (try/loop/race/join) receive a WorkflowContext,
+// because inside them you sequence further steps rather than touch actor data.
+type WorkflowContextRun<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = (
+	ctx: WorkflowContext<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>,
+) => Promise<T>;
+
+export type WorkflowStepConfig<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = Omit<StepConfig<T>, "run" | "rollback"> & {
+	run: WorkflowStepRun<
+		T,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+	rollback?: WorkflowStepRollback<
+		T,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+};
+
+export type WorkflowTryStepConfig<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = Omit<TryStepConfig<T>, "run" | "rollback"> & {
+	run: WorkflowStepRun<
+		T,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+	rollback?: WorkflowStepRollback<
+		T,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+};
+
+export type WorkflowTryConfig<
+	T,
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig,
+	TQueues extends QueueSchemaConfig,
+> = Omit<TryBlockConfig<T>, "run"> & {
+	run: WorkflowContextRun<
+		T,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+};
+
+export type WorkflowLoopConfig<
 	S,
 	T,
 	TState,
@@ -70,7 +238,7 @@ type ActorWorkflowLoopConfig<
 	TQueues extends QueueSchemaConfig,
 > = Omit<LoopConfig<S, T>, "run"> & {
 	run: (
-		ctx: ActorWorkflowContext<
+		ctx: WorkflowContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -86,7 +254,7 @@ type ActorWorkflowLoopConfig<
 	>;
 };
 
-type ActorWorkflowBranchConfig<
+export type WorkflowBranchConfig<
 	TOutput,
 	TState,
 	TConnParams,
@@ -97,8 +265,54 @@ type ActorWorkflowBranchConfig<
 	TEvents extends EventSchemaConfig,
 	TQueues extends QueueSchemaConfig,
 > = {
-	run: (
-		ctx: ActorWorkflowContext<
+	run: WorkflowContextRun<
+		TOutput,
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+};
+
+// Marks a step context inactive once its step has finished. Module-private so it
+// never appears on the public surface.
+const DEACTIVATE_STEP = Symbol("workflow.step.deactivate");
+
+/**
+ * The context handed to a workflow step (`step` / `tryStep` callbacks). This is
+ * the only scope where actor data (state, vars, db, client) and side effects
+ * (broadcast, queue.send) are reachable. It is valid only while its step is
+ * executing; using it after the step resolves throws.
+ */
+export class WorkflowStepContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+	TEvents extends EventSchemaConfig = Record<never, never>,
+	TQueues extends QueueSchemaConfig = Record<never, never>,
+> {
+	#runCtx: RunContext<
+		TState,
+		TConnParams,
+		TConnState,
+		TVars,
+		TInput,
+		TDatabase,
+		TEvents,
+		TQueues
+	>;
+	#active = true;
+	#onGuardViolation: () => void;
+
+	constructor(
+		runCtx: RunContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -108,10 +322,133 @@ type ActorWorkflowBranchConfig<
 			TEvents,
 			TQueues
 		>,
-	) => Promise<TOutput>;
-};
+		onGuardViolation: () => void,
+	) {
+		this.#runCtx = runCtx;
+		this.#onGuardViolation = onGuardViolation;
+	}
 
-export class ActorWorkflowContext<
+	[DEACTIVATE_STEP](): void {
+		this.#active = false;
+	}
+
+	#ensureActive(feature: string): void {
+		if (!this.#active) {
+			this.#onGuardViolation();
+			throw new Error(
+				`${feature} is only available inside workflow steps`,
+			);
+		}
+	}
+
+	get actorId(): string {
+		return this.#runCtx.actorId;
+	}
+
+	get name(): string {
+		return this.#runCtx.name;
+	}
+
+	get key(): string[] {
+		return this.#runCtx.key;
+	}
+
+	get log() {
+		return this.#runCtx.log;
+	}
+
+	get abortSignal(): AbortSignal {
+		return this.#runCtx.abortSignal;
+	}
+
+	get state(): TState extends never ? never : TState {
+		this.#ensureActive("state");
+		return this.#runCtx.state as TState extends never ? never : TState;
+	}
+
+	get vars(): TVars extends never ? never : TVars {
+		this.#ensureActive("vars");
+		return this.#runCtx.vars as TVars extends never ? never : TVars;
+	}
+
+	get db(): TDatabase extends never ? never : InferDatabaseClient<TDatabase> {
+		this.#ensureActive("db");
+		return this.#runCtx.db as TDatabase extends never
+			? never
+			: InferDatabaseClient<TDatabase>;
+	}
+
+	client<R extends Registry<any> = Registry<any>>(): Client<R> {
+		this.#ensureActive("client");
+		return this.#runCtx.client<R>();
+	}
+
+	broadcast<K extends keyof TEvents & string>(
+		name: K,
+		...args: InferEventArgs<InferSchemaMap<TEvents>[K]>
+	): void;
+	broadcast(
+		name: keyof TEvents extends never ? string : never,
+		...args: Array<unknown>
+	): void;
+	broadcast(name: string, ...args: Array<unknown>): void {
+		this.#ensureActive("broadcast");
+		this.#runCtx.broadcast(
+			name as never,
+			...(args as unknown[] as never[]),
+		);
+	}
+
+	get queue() {
+		const self = this;
+		function send<K extends keyof TQueues & string>(
+			name: K,
+			body: InferSchemaMap<TQueues>[K],
+		): Promise<void>;
+		function send(
+			name: keyof TQueues extends never ? string : never,
+			body: unknown,
+		): Promise<void>;
+		async function send(name: string, body: unknown): Promise<void> {
+			self.#ensureActive("queue.send");
+			await self.#runCtx.queue.send(name as never, body as never);
+		}
+		return { send };
+	}
+
+	/**
+	 * Holds the actor awake for the duration of the provided promise. The actor
+	 * cannot idle-sleep or finalize the sleep grace period until the promise
+	 * settles.
+	 */
+	keepAwake<T>(promise: Promise<T>): Promise<T> {
+		this.#ensureActive("keepAwake");
+		return this.#runCtx.keepAwake(promise);
+	}
+
+	/**
+	 * Registers a promise that the sleep grace period will wait on. Use this for
+	 * best-effort flush/cleanup work that may complete inside the grace window.
+	 */
+	waitUntil(promise: Promise<void>): void {
+		this.#ensureActive("waitUntil");
+		this.#runCtx.waitUntil(promise);
+	}
+
+	destroy(): void {
+		this.#ensureActive("destroy");
+		this.#runCtx.destroy();
+	}
+}
+
+/**
+ * The context handed to the workflow function and to orchestration callbacks
+ * (`try` / `loop` / `race` / `join`). It exposes the deterministic, replayable
+ * workflow primitives (step, sleep, queue waits, control flow). It deliberately
+ * does NOT expose actor data; reach `state`, `db`, `vars`, `client`, and
+ * `broadcast` through the step context passed to `step` / `tryStep`.
+ */
+export class WorkflowContext<
 	TState,
 	TConnParams,
 	TConnState,
@@ -120,8 +457,7 @@ export class ActorWorkflowContext<
 	TDatabase extends AnyDatabaseProvider,
 	TEvents extends EventSchemaConfig = Record<never, never>,
 	TQueues extends QueueSchemaConfig = Record<never, never>,
-> implements WorkflowContextInterface
-{
+> {
 	#inner: WorkflowContextInterface;
 	#runCtx: RunContext<
 		TState,
@@ -133,9 +469,6 @@ export class ActorWorkflowContext<
 		TEvents,
 		TQueues
 	>;
-	#actorAccessDepth = 0;
-	#allowActorAccess = false;
-	#guardViolation = false;
 
 	constructor(
 		inner: WorkflowContextInterface,
@@ -162,6 +495,22 @@ export class ActorWorkflowContext<
 		return this.#inner.abortSignal;
 	}
 
+	get actorId(): string {
+		return this.#runCtx.actorId;
+	}
+
+	get name(): string {
+		return this.#runCtx.name;
+	}
+
+	get key(): string[] {
+		return this.#runCtx.key;
+	}
+
+	get log() {
+		return this.#runCtx.log;
+	}
+
 	get queue() {
 		const self = this;
 		function next<
@@ -181,10 +530,9 @@ export class ActorWorkflowContext<
 				TCompletable
 			>
 		>;
-		async function next(
-			name: string,
-			opts?: WorkflowActorQueueNextOptions<string, boolean>,
-		): Promise<WorkflowQueueMessage<unknown>> {
+		// The implementation signature stays broad so the schema-typed public
+		// overloads above remain compatible with it.
+		async function next(name: string, opts?: any): Promise<any> {
 			const message = await self.#inner.queue.next(name, opts);
 			return self.#toActorQueueMessage(message);
 		}
@@ -210,106 +558,232 @@ export class ActorWorkflowContext<
 				>
 			>
 		>;
-		async function nextBatch(
-			name: string,
-			opts?: WorkflowActorQueueNextBatchOptions<string, boolean>,
-		): Promise<Array<WorkflowQueueMessage<unknown>>> {
+		async function nextBatch(name: string, opts?: any): Promise<any> {
 			const messages = await self.#inner.queue.nextBatch(name, opts);
 			return messages.map((message) =>
 				self.#toActorQueueMessage(message),
 			);
 		}
 
-		function send<K extends keyof TQueues & string>(
-			name: K,
-			body: InferSchemaMap<TQueues>[K],
-		): Promise<void>;
-		function send(
-			name: keyof TQueues extends never ? string : never,
-			body: unknown,
-		): Promise<void>;
-		async function send(name: string, body: unknown): Promise<void> {
-			self.#ensureActorAccess("queue.send");
-			await self.#runCtx.queue.send(name as never, body as never);
-		}
-
 		return {
 			next,
 			nextBatch,
-			send,
 		};
 	}
 
+	step<T>(
+		name: string,
+		run: WorkflowStepRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<T>;
+	step<T>(
+		config: WorkflowStepConfig<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<T>;
 	async step<T>(
-		nameOrConfig: string | StepConfig<T>,
-		run?: () => Promise<T>,
+		nameOrConfig:
+			| string
+			| WorkflowStepConfig<
+					T,
+					TState,
+					TConnParams,
+					TConnState,
+					TVars,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+			  >,
+		run?: WorkflowStepRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
 	): Promise<T> {
 		if (typeof nameOrConfig === "string") {
 			if (!run) {
 				throw new Error("Step run function missing");
 			}
+			const stepRun = run;
 			return await this.#wrapActive(() =>
-				this.#inner.step(nameOrConfig, () =>
-					this.#withActorAccessAndStateRollback(run),
-				),
+				this.#inner.step(nameOrConfig, () => this.#runStep(stepRun)),
 			);
 		}
-		const stepConfig = nameOrConfig as StepConfig<T>;
+		const stepConfig = nameOrConfig;
+		const rollback = stepConfig.rollback;
 		const config: StepConfig<T> = {
 			...stepConfig,
-			run: () => this.#withActorAccessAndStateRollback(stepConfig.run),
+			run: () => this.#runStep(stepConfig.run),
+			rollback: rollback
+				? (_ctx, output) => this.#runRollback(rollback, output)
+				: undefined,
 		};
 		return await this.#wrapActive(() => this.#inner.step(config));
 	}
 
+	tryStep<T>(
+		name: string,
+		run: WorkflowStepRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<TryStepResult<T>>;
+	tryStep<T>(
+		config: WorkflowTryStepConfig<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<TryStepResult<T>>;
 	async tryStep<T>(
-		nameOrConfig: string | TryStepConfig<T>,
-		run?: () => Promise<T>,
+		nameOrConfig:
+			| string
+			| WorkflowTryStepConfig<
+					T,
+					TState,
+					TConnParams,
+					TConnState,
+					TVars,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+			  >,
+		run?: WorkflowStepRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
 	): Promise<TryStepResult<T>> {
 		if (typeof nameOrConfig === "string") {
 			if (!run) {
 				throw new Error("Step run function missing");
 			}
+			const stepRun = run;
 			return await this.#wrapActive(() =>
-				this.#inner.tryStep(nameOrConfig, () =>
-					this.#withActorAccessAndStateRollback(run),
-				),
+				this.#inner.tryStep(nameOrConfig, () => this.#runStep(stepRun)),
 			);
 		}
-		const stepConfig = nameOrConfig as TryStepConfig<T>;
+		const stepConfig = nameOrConfig;
+		const rollback = stepConfig.rollback;
 		const config: TryStepConfig<T> = {
 			...stepConfig,
-			run: () => this.#withActorAccessAndStateRollback(stepConfig.run),
+			run: () => this.#runStep(stepConfig.run),
+			rollback: rollback
+				? (_ctx, output) => this.#runRollback(rollback, output)
+				: undefined,
 		};
 		return await this.#wrapActive(() => this.#inner.tryStep(config));
 	}
 
+	try<T>(
+		name: string,
+		run: WorkflowContextRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<TryBlockResult<T>>;
+	try<T>(
+		config: WorkflowTryConfig<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+	): Promise<TryBlockResult<T>>;
 	async try<T>(
-		nameOrConfig: string | Parameters<WorkflowContextInterface["try"]>[0],
-		run?: (
-			ctx: ActorWorkflowContext<
-				TState,
-				TConnParams,
-				TConnState,
-				TVars,
-				TInput,
-				TDatabase,
-				TEvents,
-				TQueues
-			>,
-		) => Promise<T>,
+		nameOrConfig:
+			| string
+			| WorkflowTryConfig<
+					T,
+					TState,
+					TConnParams,
+					TConnState,
+					TVars,
+					TInput,
+					TDatabase,
+					TEvents,
+					TQueues
+			  >,
+		run?: WorkflowContextRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
 	): Promise<TryBlockResult<T>> {
 		if (typeof nameOrConfig === "string") {
 			if (!run) {
 				throw new Error("Try run function missing");
 			}
+			const tryRun = run;
 			return await this.#wrapActive(() =>
 				this.#inner.try(nameOrConfig, async (ctx) =>
-					run(this.#createChildContext(ctx)),
+					tryRun(this.#createChildContext(ctx)),
 				),
 			);
 		}
-		const tryConfig = nameOrConfig as TryBlockConfig<T>;
+		const tryConfig = nameOrConfig;
 		const config: TryBlockConfig<T> = {
 			...tryConfig,
 			run: async (ctx) => tryConfig.run(this.#createChildContext(ctx)),
@@ -317,10 +791,10 @@ export class ActorWorkflowContext<
 		return await this.#wrapActive(() => this.#inner.try(config));
 	}
 
-	async loop<T>(
+	loop<T>(
 		name: string,
 		run: (
-			ctx: ActorWorkflowContext<
+			ctx: WorkflowContext<
 				TState,
 				TConnParams,
 				TConnState,
@@ -332,14 +806,8 @@ export class ActorWorkflowContext<
 			>,
 		) => Promise<LoopResult<undefined, T> | undefined | void>,
 	): Promise<T>;
-	async loop<T>(
-		name: string,
-		run: (
-			ctx: WorkflowContextInterface,
-		) => Promise<LoopResult<undefined, T> | undefined | void>,
-	): Promise<T>;
-	async loop<S, T>(
-		config: ActorWorkflowLoopConfig<
+	loop<S, T>(
+		config: WorkflowLoopConfig<
 			S,
 			T,
 			TState,
@@ -352,12 +820,10 @@ export class ActorWorkflowContext<
 			TQueues
 		>,
 	): Promise<T>;
-	async loop<S, T>(config: LoopConfig<S, T>): Promise<T>;
 	async loop(
 		nameOrConfig:
 			| string
-			| LoopConfig<any, any>
-			| ActorWorkflowLoopConfig<
+			| WorkflowLoopConfig<
 					any,
 					any,
 					TState,
@@ -370,7 +836,7 @@ export class ActorWorkflowContext<
 					TQueues
 			  >,
 		run?: (
-			ctx: ActorWorkflowContext<
+			ctx: WorkflowContext<
 				TState,
 				TConnParams,
 				TConnState,
@@ -386,16 +852,26 @@ export class ActorWorkflowContext<
 			if (!run) {
 				throw new Error("Loop run function missing");
 			}
+			const loopRun = run;
 			return await this.#wrapActive(() =>
-				this.#inner.loop(nameOrConfig, async (ctx) =>
-					run(this.#createChildContext(ctx)),
+				this.#inner.loop(
+					nameOrConfig,
+					// A void return (no explicit Loop result) is undefined at
+					// runtime, which the engine treats as continue.
+					async (
+						ctx,
+					): Promise<LoopResult<undefined, any> | undefined> =>
+						(await loopRun(this.#createChildContext(ctx))) ??
+						undefined,
 				),
 			);
 		}
+		const loopConfig = nameOrConfig;
 		const wrapped: LoopConfig<any, any> = {
-			...nameOrConfig,
-			run: async (ctx, state) =>
-				nameOrConfig.run(this.#createChildContext(ctx), state),
+			...loopConfig,
+			run: (async (ctx, state) =>
+				(await loopConfig.run(this.#createChildContext(ctx), state)) ??
+				undefined) as LoopConfig<any, any>["run"],
 		};
 		return await this.#wrapActive(() => this.#inner.loop(wrapped));
 	}
@@ -408,11 +884,6 @@ export class ActorWorkflowContext<
 		return this.#inner.sleepUntil(name, timestampMs);
 	}
 
-	destroy(): void {
-		this.#ensureActorAccess("destroy");
-		this.#runCtx.destroy();
-	}
-
 	async rollbackCheckpoint(name: string): Promise<void> {
 		await this.#wrapActive(() => this.#inner.rollbackCheckpoint(name));
 	}
@@ -420,7 +891,7 @@ export class ActorWorkflowContext<
 	async join<
 		T extends Record<
 			string,
-			ActorWorkflowBranchConfig<
+			WorkflowBranchConfig<
 				unknown,
 				TState,
 				TConnParams,
@@ -440,7 +911,10 @@ export class ActorWorkflowContext<
 		name: string,
 		branches: T,
 	): Promise<{ [K in keyof T]: BranchOutput<T[K]> }>;
-	async join(name: string, branches: Record<string, BranchConfig<unknown>>) {
+	async join(
+		name: string,
+		branches: Record<string, { run: (ctx: any) => Promise<unknown> }>,
+	) {
 		const wrappedBranches = Object.fromEntries(
 			Object.entries(branches).map(([key, branch]) => [
 				key,
@@ -459,25 +933,17 @@ export class ActorWorkflowContext<
 		name: string,
 		branches: Array<{
 			name: string;
-			run: (
-				ctx: ActorWorkflowContext<
-					TState,
-					TConnParams,
-					TConnState,
-					TVars,
-					TInput,
-					TDatabase,
-					TEvents,
-					TQueues
-				>,
-			) => Promise<T>;
-		}>,
-	): Promise<{ winner: string; value: T }>;
-	async race<T>(
-		name: string,
-		branches: Array<{
-			name: string;
-			run: (ctx: WorkflowContextInterface) => Promise<T>;
+			run: WorkflowContextRun<
+				T,
+				TState,
+				TConnParams,
+				TConnState,
+				TVars,
+				TInput,
+				TDatabase,
+				TEvents,
+				TQueues
+			>;
 		}>,
 	): Promise<{ winner: string; value: T }> {
 		const wrappedBranches = branches.map((branch) => ({
@@ -498,124 +964,33 @@ export class ActorWorkflowContext<
 		return this.#inner.isEvicted();
 	}
 
-	get state(): TState extends never ? never : TState {
-		this.#ensureActorAccess("state");
-		return this.#runCtx.state as TState extends never ? never : TState;
-	}
-
-	get vars(): TVars extends never ? never : TVars {
-		this.#ensureActorAccess("vars");
-		return this.#runCtx.vars as TVars extends never ? never : TVars;
-	}
-
-	client<R extends Registry<any> = Registry<any>>(): Client<R> {
-		this.#ensureActorAccess("client");
-		return this.#runCtx.client<R>();
-	}
-
-	get db(): TDatabase extends never ? never : InferDatabaseClient<TDatabase> {
-		this.#ensureActorAccess("db");
-		return this.#runCtx.db as TDatabase extends never
-			? never
-			: InferDatabaseClient<TDatabase>;
-	}
-
-	get log() {
-		return this.#runCtx.log;
-	}
-
-	/** @deprecated No-op. Use `keepAwake(promise)` or `waitUntil(promise)` instead. */
-	setPreventSleep(_prevent: boolean): void {
-		this.#ensureActorAccess("setPreventSleep");
-	}
-
-	/** @deprecated No-op. Always returns `false`. */
-	get preventSleep(): boolean {
-		this.#ensureActorAccess("preventSleep");
-		return false;
-	}
-
-	/**
-	 * Holds the actor awake for the duration of the provided promise. The
-	 * actor cannot idle-sleep or finalize the sleep grace period until the
-	 * promise settles.
-	 */
-	keepAwake<T>(promise: Promise<T>): Promise<T> {
-		this.#ensureActorAccess("keepAwake");
-		return this.#runCtx.keepAwake(promise);
-	}
-
-	/**
-	 * Registers a promise that the sleep grace period will wait on. Use this
-	 * for best-effort flush/cleanup work that may complete inside the grace
-	 * window. For work the actor must stay running through, prefer
-	 * `c.keepAwake(promise)` which also blocks idle sleep.
-	 */
-	waitUntil(promise: Promise<void>): void {
-		this.#ensureActorAccess("waitUntil");
-		this.#runCtx.waitUntil(promise);
-	}
-
-	get actorId(): string {
-		return this.#runCtx.actorId;
-	}
-
-	broadcast<K extends keyof TEvents & string>(
-		name: K,
-		...args: InferEventArgs<InferSchemaMap<TEvents>[K]>
-	): void;
-	broadcast(
-		name: keyof TEvents extends never ? string : never,
-		...args: Array<unknown>
-	): void;
-	broadcast(name: string, ...args: Array<unknown>): void {
-		this.#ensureActorAccess("broadcast");
-		this.#runCtx.broadcast(
-			name as never,
-			...(args as unknown[] as never[]),
-		);
-	}
-
-	#toActorQueueMessage<T>(
-		message: WorkflowQueueMessage<T>,
-	): WorkflowQueueMessage<T> & { id: bigint } {
-		let id: bigint;
-		try {
-			id = BigInt(message.id);
-		} catch {
-			throw new Error(`Invalid queue message id "${message.id}"`);
-		}
-		return {
-			id,
-			name: message.name,
-			body: message.body,
-			createdAt: message.createdAt,
-			...(message.complete ? { complete: message.complete } : {}),
-		};
-	}
-
-	async #wrapActive<T>(run: () => Promise<T>): Promise<T> {
-		return await this.#runCtx.internalKeepAwake(run);
-	}
-
-	async #withActorAccess<T>(run: () => Promise<T>): Promise<T> {
-		this.#actorAccessDepth++;
-		if (this.#actorAccessDepth === 1) {
-			this.#allowActorAccess = true;
-		}
-		try {
-			return await run();
-		} finally {
-			this.#actorAccessDepth--;
-			if (this.#actorAccessDepth === 0) {
-				this.#allowActorAccess = false;
-			}
-		}
-	}
-
-	async #withActorAccessAndStateRollback<T>(
-		run: () => Promise<T>,
+	// Runs a user step body inside a fresh step context, snapshotting actor
+	// state/vars so a thrown step rolls back its mutations, and deactivating the
+	// step context once the body settles so it cannot be used after the step.
+	async #runStep<T>(
+		run: WorkflowStepRun<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
 	): Promise<T> {
+		const stepCtx = new WorkflowStepContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>(this.#runCtx, () => this.#markGuardTriggered());
+
 		let stateSnapshot: { state: TState } | null = null;
 		try {
 			stateSnapshot = { state: this.#runCtx[RAW_STATE_SYMBOL]() };
@@ -629,33 +1004,77 @@ export class ActorWorkflowContext<
 			stateSnapshot.state = structuredClone(stateSnapshot.state);
 		}
 		const varsSnapshot = structuredClone(this.#runCtx.vars);
+
 		try {
-			return await this.#withActorAccess(run);
+			return await run(stepCtx);
 		} catch (error) {
 			if (stateSnapshot) {
 				this.#runCtx.state = stateSnapshot.state;
 			}
 			this.#runCtx.vars = varsSnapshot;
 			throw error;
+		} finally {
+			stepCtx[DEACTIVATE_STEP]();
 		}
 	}
 
-	#ensureActorAccess(feature: string): void {
-		if (!this.#allowActorAccess) {
-			this.#guardViolation = true;
-			this.#markGuardTriggered();
-			throw new Error(
-				`${feature} is only available inside workflow steps`,
-			);
+	#toActorQueueMessage<T>(
+		message: WorkflowQueueMessage<T>,
+	): WorkflowQueueMessage<T> & { id: bigint } {
+		let id: bigint;
+		try {
+			id = BigInt(message.id);
+		} catch {
+			throw new Error(`Invalid queue message id "${message.id}"`);
+		}
+		return {
+			id,
+			name: message.name,
+			body: message.body,
+			createdAt: message.createdAt,
+			...(message.complete ? { complete: message.complete } : {}),
+		};
+	}
+
+	// Runs a step rollback compensation with an active step context. Rollbacks
+	// intentionally mutate actor state, so their writes are not snapshotted.
+	async #runRollback<T>(
+		rollback: WorkflowStepRollback<
+			T,
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>,
+		output: T,
+	): Promise<void> {
+		const stepCtx = new WorkflowStepContext<
+			TState,
+			TConnParams,
+			TConnState,
+			TVars,
+			TInput,
+			TDatabase,
+			TEvents,
+			TQueues
+		>(this.#runCtx, () => this.#markGuardTriggered());
+		try {
+			await rollback(stepCtx, output);
+		} finally {
+			stepCtx[DEACTIVATE_STEP]();
 		}
 	}
 
-	consumeGuardViolation(): boolean {
-		const violated = this.#guardViolation;
-		this.#guardViolation = false;
-		return violated;
+	async #wrapActive<T>(run: () => Promise<T>): Promise<T> {
+		return await this.#runCtx.internalKeepAwake(run);
 	}
 
+	// Records that a step context was used outside its step. Mirrors the value
+	// onto actor state and a KV flag so callers can observe the violation.
 	#markGuardTriggered(): void {
 		try {
 			const state = this.#runCtx.state as Record<string, unknown>;
@@ -686,7 +1105,7 @@ export class ActorWorkflowContext<
 
 	#createChildContext(
 		ctx: WorkflowContextInterface,
-	): ActorWorkflowContext<
+	): WorkflowContext<
 		TState,
 		TConnParams,
 		TConnState,
@@ -696,7 +1115,7 @@ export class ActorWorkflowContext<
 		TEvents,
 		TQueues
 	> {
-		return new ActorWorkflowContext(ctx, this.#runCtx);
+		return new WorkflowContext(ctx, this.#runCtx);
 	}
 }
 
@@ -712,7 +1131,7 @@ export type WorkflowContextOf<AD extends AnyActorDefinition> =
 		infer Q extends QueueSchemaConfig,
 		any
 	>
-		? ActorWorkflowContext<S, CP, CS, V, I, DB, E, Q>
+		? WorkflowContext<S, CP, CS, V, I, DB, E, Q>
 		: never;
 
 export type WorkflowLoopContextOf<AD extends AnyActorDefinition> =
@@ -722,4 +1141,16 @@ export type WorkflowBranchContextOf<AD extends AnyActorDefinition> =
 	WorkflowContextOf<AD>;
 
 export type WorkflowStepContextOf<AD extends AnyActorDefinition> =
-	WorkflowContextOf<AD>;
+	AD extends BaseActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		infer E extends EventSchemaConfig,
+		infer Q extends QueueSchemaConfig,
+		any
+	>
+		? WorkflowStepContext<S, CP, CS, V, I, DB, E, Q>
+		: never;