[workflows] adding restart-from-step support to the sdk (#13565)

Author: vaishnav-mk

.changeset/restart-from-step.md

@@ -0,0 +1,10 @@
+---
+"@cloudflare/workflows-shared": minor
+"miniflare": minor
+---
+
+Add restart from step support for local Workflows development
+
+Workflow instances can now be restarted from a specific step in local development. When restarting from a step, all earlier steps preserve their cached results and replay instantly, while the target step and everything after it re-execute.
+
+The `WorkflowInstance.restart()` method now accepts an optional `{ from: { name, count?, type? } }` parameter to specify which step to restart from.

fixtures/vitest-pool-workers-examples/workflows/test/integration.test.ts

@@ -140,6 +140,8 @@ describe("workflow instance lifecycle methods", () => {
 		// DISPOSE: ensured by `await using`
 	});
 
+	// TODO(vaish): add restart-from-step test here once @cloudflare/workers-types ships restart options
+
 	it("should pause a workflow instance", async ({ expect }) => {
 		// CONFIG:
 		await using introspector = await introspectWorkflow(env.MODERATOR);

packages/miniflare/src/workers/local-explorer/resources/workflows.ts

@@ -8,6 +8,7 @@ import type { AppContext } from "../common";
 import type { Env } from "../explorer.worker";
 import type { WorkflowsWorkflow } from "../generated";
 import type { zWorkflowsListInstancesData } from "../generated/zod.gen";
+import type { RestartFromStep } from "@cloudflare/workflows-shared/src/binding";
 import type { z } from "zod";
 
 // ============================================================================
@@ -30,7 +31,7 @@ interface DirectoryEntry {
 interface WorkflowHandle {
 	pause(): Promise<void>;
 	resume(): Promise<void>;
-	restart(): Promise<void>;
+	restart(options?: { from?: RestartFromStep }): Promise<void>;
 	terminate(): Promise<void>;
 	sendEvent(args: { payload: unknown; type: string }): Promise<void>;
 	status(): Promise<{ status: string; output?: unknown; error?: unknown }>;
@@ -871,7 +872,10 @@ export async function changeWorkflowInstanceStatus(
 	}
 
 	try {
-		const body = (await c.req.json()) as { action: string };
+		const body = (await c.req.json()) as {
+			action: string;
+			from?: RestartFromStep;
+		};
 		const { action } = body;
 
 		if (!["pause", "resume", "restart", "terminate"].includes(action)) {
@@ -891,9 +895,19 @@ export async function changeWorkflowInstanceStatus(
 			case "resume":
 				await handle.resume();
 				break;
-			case "restart":
-				await handle.restart();
+			case "restart": {
+				if (body.from && !body.from.name) {
+					return errorResponse(
+						400,
+						10001,
+						"'from.name' is required when restarting from a specific step."
+					);
+				}
+				const opts = body.from ? { from: body.from } : undefined;
+				// TODO(vaish): remove cast once @cloudflare/workers-types ships restart options
+				await (handle as unknown as WorkflowHandle).restart(opts);
 				break;
+			}
 			case "terminate":
 				await handle.terminate();
 				break;

packages/miniflare/src/workers/workflows/wrapped-binding.worker.ts

@@ -1,4 +1,7 @@
-import type { WorkflowBinding } from "@cloudflare/workflows-shared/src/binding";
+import type {
+	WorkflowBinding,
+	WorkflowInstanceRestartOptions,
+} from "@cloudflare/workflows-shared/src/binding";
 
 class WorkflowImpl implements Workflow {
 	constructor(private binding: WorkflowBinding) {}
@@ -86,9 +89,13 @@ class InstanceImpl implements WorkflowInstance {
 		await instance.terminate();
 	}
 
-	public async restart(): Promise<void> {
+	public async restart(
+		options?: WorkflowInstanceRestartOptions
+	): Promise<void> {
 		using instance = await this.getInstance();
-		await instance.restart();
+		// TODO(vaish): remove @ts-expect-error once @cloudflare/workers-types ships restart options
+		// @ts-expect-error WorkflowInstance type does not include options yet
+		await instance.restart(options);
 	}
 
 	public async status(): Promise<InstanceStatus> {

packages/workflows-shared/src/binding.ts

@@ -21,6 +21,17 @@ type Env = {
 	BINDING_NAME: string;
 };
 
+// TODO(vaish): import from @cloudflare/workers-types once restart options are published
+export interface RestartFromStep {
+	name: string;
+	count?: number;
+	type?: "do" | "sleep" | "waitForEvent";
+}
+
+export interface WorkflowInstanceRestartOptions {
+	from?: RestartFromStep;
+}
+
 // this.env.WORKFLOW is WorkflowBinding
 export class WorkflowBinding extends WorkerEntrypoint<Env> {
 	constructor(ctx: ExecutionContext, env: Env) {
@@ -208,9 +219,11 @@ export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
 		}
 	}
 
-	public async restart(): Promise<void> {
+	public async restart(
+		options?: WorkflowInstanceRestartOptions
+	): Promise<void> {
 		try {
-			await this.stub.changeInstanceStatus("restart");
+			await this.stub.changeInstanceStatus("restart", options?.from);
 		} catch (e) {
 			// restart causes instance abortion
 			if (!isUserTriggeredRestart(e)) {

packages/workflows-shared/src/engine.ts

@@ -15,25 +15,29 @@ import {
 	isAbortError,
 	PreservedNonRetryableError,
 	shouldPreserveNonRetryableError,
+	stepNotFoundError,
 	WorkflowFatalError,
 } from "./lib/errors";
 import {
 	ENGINE_TIMEOUT,
 	GracePeriodSemaphore,
 	startGracePeriod,
 } from "./lib/gracePeriodSemaphore";
+import {
+	readAndClearRestartFromStep,
+	resolveGroupKeysToWipe,
+	storeRestartFromStep,
+	wipeRestartState,
+} from "./lib/restart";
 import {
 	createReplayReadableStream,
 	getInvalidStoredStreamOutputError,
 	getStoredStreamOutputPreview,
 	StreamOutputState,
 } from "./lib/streams";
 import { TimePriorityQueue } from "./lib/timePriorityQueue";
-import {
-	isModifierKey,
-	MODIFIER_KEYS,
-	WorkflowInstanceModifier,
-} from "./modifier";
+import { MODIFIER_KEYS, WorkflowInstanceModifier } from "./modifier";
+import type { RestartFromStep } from "./binding";
 import type { Event } from "./context";
 import type { InstanceMetadata, RawInstanceLog } from "./instance";
 import type { StreamOutputMeta } from "./lib/streams";
@@ -753,7 +757,8 @@ export class Engine extends DurableObject<Env> {
 	}
 
 	async changeInstanceStatus(
-		newStatus: "resume" | "pause" | "terminate" | "restart"
+		newStatus: "resume" | "pause" | "terminate" | "restart",
+		from?: RestartFromStep
 	) {
 		const metadata =
 			await this.ctx.storage.get<InstanceMetadata>(INSTANCE_METADATA);
@@ -802,6 +807,12 @@ export class Engine extends DurableObject<Env> {
 				break;
 			}
 			case "restart":
+				if (from) {
+					if (!resolveGroupKeysToWipe(this.ctx.storage.sql, from)) {
+						throw stepNotFoundError(from.name);
+					}
+					await storeRestartFromStep(this.ctx.storage, from);
+				}
 				await this.userTriggeredRestart();
 				break;
 		}
@@ -889,61 +900,27 @@ export class Engine extends DurableObject<Env> {
 		await this.abort(ABORT_REASONS.USER_RESTART);
 	}
 
-	private getMockedEventMapKeys(allKeys: Map<string, unknown>): Set<string> {
-		const mockEventTypes = new Set<string>();
-		for (const key of allKeys.keys()) {
-			if (key.startsWith(MODIFIER_KEYS.MOCK_EVENT)) {
-				mockEventTypes.add(key.slice(MODIFIER_KEYS.MOCK_EVENT.length));
-			}
-		}
-
-		if (mockEventTypes.size === 0) {
-			return new Set();
-		}
+	async attemptRestart() {
+		const restartFromStep = await readAndClearRestartFromStep(this.ctx.storage);
 
-		const preserved = new Set<string>();
-		for (const key of allKeys.keys()) {
-			if (key.startsWith(`${EVENT_MAP_PREFIX}\n`)) {
-				// EVENT_MAP keys are formatted as "EVENT_MAP\n{type}\n{idx}"
-				const eventType = key.split("\n")[1];
-				if (eventType !== undefined && mockEventTypes.has(eventType)) {
-					preserved.add(key);
-				}
+		let groupKeysToWipe: Set<string> | null = null;
+		if (restartFromStep) {
+			groupKeysToWipe = resolveGroupKeysToWipe(
+				this.ctx.storage.sql,
+				restartFromStep
+			);
+			if (!groupKeysToWipe) {
+				throw stepNotFoundError(restartFromStep.name);
 			}
 		}
 
-		return preserved;
-	}
-
-	async attemptRestart() {
-		this.ctx.storage.sql.exec("DELETE FROM states");
-		this.ctx.storage.sql.exec("DELETE FROM priority_queue");
-		// Only delete non-mock streaming chunks. Mock stream outputs are stored
-		// at attempt=0 (see modifier.ts mockStepResult) and their sentinels
-		// survive restart via isModifierKey(), so the underlying SQL rows must
-		// be preserved too.
-		this.ctx.storage.sql.exec(
-			"DELETE FROM streaming_step_chunks WHERE attempt != 0"
+		await wipeRestartState(
+			this.ctx.storage,
+			ENGINE_STATUS_KEY,
+			PAUSE_DATETIME,
+			groupKeysToWipe
 		);
 
-		const allKeys = await this.ctx.storage.list();
-		const preservedEventMapKeys = this.getMockedEventMapKeys(allKeys);
-
-		// Remove all KV keys except:
-		// - INSTANCE_METADATA (needed to re-run the workflow)
-		// - Modifier/mock keys (so mocks survive restart)
-		// - EVENT_MAP entries for mocked event types
-		for (const key of allKeys.keys()) {
-			if (
-				key === INSTANCE_METADATA ||
-				isModifierKey(key) ||
-				preservedEventMapKeys.has(key)
-			) {
-				continue;
-			}
-			await this.ctx.storage.delete(key);
-		}
-
 		const metadata =
 			await this.ctx.storage.get<InstanceMetadata>(INSTANCE_METADATA);
 
@@ -956,14 +933,16 @@ export class Engine extends DurableObject<Env> {
 
 		const { accountId, workflow, version, instance, event } = metadata;
 
-		this.writeLog(InstanceEvent.WORKFLOW_QUEUED, null, null, {
-			params: event.payload,
-			versionId: version.id,
-			trigger: {
-				source: InstanceTrigger.API,
-			},
-		});
-		this.writeLog(InstanceEvent.WORKFLOW_START, null, null, {});
+		if (!groupKeysToWipe) {
+			this.writeLog(InstanceEvent.WORKFLOW_QUEUED, null, null, {
+				params: event.payload,
+				versionId: version.id,
+				trigger: {
+					source: InstanceTrigger.API,
+				},
+			});
+			this.writeLog(InstanceEvent.WORKFLOW_START, null, null, {});
+		}
 
 		void this.init(accountId, workflow, version, instance, event);
 	}

packages/workflows-shared/src/lib/errors.ts

@@ -114,3 +114,10 @@ function getCompatFlag(name: string): boolean {
 export function shouldPreserveNonRetryableError(): boolean {
 	return getCompatFlag("workflows_preserve_non_retryable_error_message");
 }
+
+export function stepNotFoundError(name: string): WorkflowError {
+	return createWorkflowError(
+		`Step "${name}" not found in execution history`,
+		"instance.cannot_restart"
+	);
+}