Support workflow terminate rollback

Author: vaishnav-mk

.changeset/workflows-terminate-rollback.md

@@ -0,0 +1,11 @@
+---
+"@cloudflare/workflows-shared": minor
+"wrangler": minor
+"miniflare": minor
+---
+
+Add rollback support when terminating Workflow instances
+
+`WorkflowInstance.terminate({ rollback: true })` now runs registered rollback handlers before marking a local Workflow instance as terminated. Wrangler also supports this for remote instances via `wrangler workflows instances terminate --rollback`.
+
+The rollback option is only sent for terminate operations and is rejected by the Local Explorer API for pause, resume, and restart actions.

packages/miniflare/scripts/openapi-filter-config.ts

@@ -1146,6 +1146,11 @@ const config = {
 											description:
 												"The action to perform on the workflow instance.",
 										},
+										rollback: {
+											type: "boolean",
+											description:
+												"Run rollback before terminating. Only valid when action is terminate.",
+										},
 										from: {
 											type: "object",
 											description:

packages/miniflare/src/workers/local-explorer/generated/types.gen.ts

@@ -1730,6 +1730,10 @@ export type WorkflowsChangeInstanceStatusData = {
 		 * The action to perform on the workflow instance.
 		 */
 		action: "pause" | "resume" | "restart" | "terminate";
+		/**
+		 * Run rollback before terminating. Only valid when action is terminate.
+		 */
+		rollback?: boolean;
 		/**
 		 * The step to restart the instance from. Only valid when action is restart.
 		 */

packages/miniflare/src/workers/local-explorer/generated/zod.gen.ts

@@ -1083,6 +1083,7 @@ export const zWorkflowsGetInstanceDetailsResponse =
 export const zWorkflowsChangeInstanceStatusData = z.object({
 	body: z.object({
 		action: z.enum(["pause", "resume", "restart", "terminate"]),
+		rollback: z.boolean().optional(),
 		from: z
 			.object({
 				name: z.string(),

packages/miniflare/src/workers/local-explorer/openapi.local.json

@@ -1790,6 +1790,10 @@
 										"enum": ["pause", "resume", "restart", "terminate"],
 										"description": "The action to perform on the workflow instance."
 									},
+									"rollback": {
+										"type": "boolean",
+										"description": "Run rollback before terminating. Only valid when action is terminate."
+									},
 									"from": {
 										"type": "object",
 										"description": "The step to restart the instance from. Only valid when action is restart.",

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

@@ -11,7 +11,7 @@ import type {
 	WorkflowsWorkflow,
 } from "../generated";
 import type { zWorkflowsListInstancesData } from "../generated/zod.gen";
-import type { RestartFromStep } from "@cloudflare/workflows-shared/src/binding";
+import type { WorkflowBinding } from "@cloudflare/workflows-shared/src/binding";
 import type { z } from "zod";
 
 // ============================================================================
@@ -30,16 +30,6 @@ interface DirectoryEntry {
 	birthtimeMs: number;
 }
 
-/** Methods on a WorkflowInstance handle (from workflow.get()). */
-interface WorkflowHandle {
-	pause(): Promise<void>;
-	resume(): 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 }>;
-}
-
 /** RPC methods exposed by the Engine Durable Object. */
 interface EngineStub {
 	getInstanceMetadata(): Promise<{
@@ -173,12 +163,15 @@ function getEngineNamespace(
  * Get the Workflow proxy binding for a workflow.
  * Used for operations that go through the standard Workflow interface (create).
  */
-function getWorkflowBinding(env: Env, workflowName: string): Workflow | null {
+function getWorkflowBinding(
+	env: Env,
+	workflowName: string
+): WorkflowBinding | null {
 	const info = env.LOCAL_EXPLORER_BINDING_MAP.workflows[workflowName];
 	if (!info) {
 		return null;
 	}
-	return env[info.binding] as Workflow;
+	return env[info.binding] as unknown as WorkflowBinding;
 }
 
 function getLocalWorkflows(env: Env): WorkflowsWorkflow[] {
@@ -886,6 +879,14 @@ export async function changeWorkflowInstanceStatus(
 			);
 		}
 
+		if (body.rollback !== undefined && action !== "terminate") {
+			return errorResponse(
+				400,
+				10001,
+				"'rollback' is only valid when terminating."
+			);
+		}
+
 		const handle = await workflow.get(instanceId);
 
 		switch (action) {
@@ -904,12 +905,13 @@ export async function changeWorkflowInstanceStatus(
 					);
 				}
 				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);
+				await handle.restart(opts);
 				break;
 			}
 			case "terminate":
-				await handle.terminate();
+				await handle.terminate(
+					body.rollback === true ? { rollback: true } : undefined
+				);
 				break;
 		}
 
@@ -1044,9 +1046,7 @@ export async function sendWorkflowInstanceEvent(
 			// Empty body is fine — payload is optional
 		}
 
-		const handle = (await workflow.get(
-			instanceId
-		)) as unknown as WorkflowHandle;
+		const handle = await workflow.get(instanceId);
 		await handle.sendEvent({ payload, type: eventType });
 
 		return c.json(wrapResponse({ success: true }));

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

@@ -1,6 +1,8 @@
 import type {
 	WorkflowBinding,
+	WorkflowHandle,
 	WorkflowInstanceRestartOptions,
+	WorkflowInstanceTerminateOptions,
 } from "@cloudflare/workflows-shared/src/binding";
 
 class WorkflowImpl implements Workflow {
@@ -70,8 +72,8 @@ class InstanceImpl implements WorkflowInstance {
 		private binding: WorkflowBinding
 	) {}
 
-	private async getInstance(): Promise<WorkflowInstance & Disposable> {
-		return (await this.binding.get(this.id)) as WorkflowInstance & Disposable;
+	private async getInstance(): Promise<WorkflowHandle & Disposable> {
+		return (await this.binding.get(this.id)) as WorkflowHandle & Disposable;
 	}
 
 	public async pause(): Promise<void> {
@@ -84,9 +86,11 @@ class InstanceImpl implements WorkflowInstance {
 		await instance.resume();
 	}
 
-	public async terminate(): Promise<void> {
+	public async terminate(
+		options?: WorkflowInstanceTerminateOptions
+	): Promise<void> {
 		using instance = await this.getInstance();
-		await instance.terminate();
+		await instance.terminate(options);
 	}
 
 	public async restart(
@@ -98,7 +102,8 @@ class InstanceImpl implements WorkflowInstance {
 
 	public async status(): Promise<InstanceStatus> {
 		using instance = await this.getInstance();
-		using res = (await instance.status()) as InstanceStatus & Disposable;
+		using res = (await instance.status()) as unknown as InstanceStatus &
+			Disposable;
 		return structuredClone(res);
 	}
 

packages/workflows-shared/src/binding.ts

@@ -29,6 +29,13 @@ export interface RestartFromStep {
 	type?: "do" | "sleep" | "waitForEvent";
 }
 
+export interface WorkflowInstanceTerminateOptions {
+	/**
+	 * If true, run registered rollback handlers before terminating the instance.
+	 */
+	rollback?: boolean;
+}
+
 export interface WorkflowInstanceRestartOptions {
 	from?: RestartFromStep;
 }
@@ -92,7 +99,7 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> {
 		};
 	}
 
-	public async get(id: string): Promise<WorkflowInstance> {
+	public async get(id: string): Promise<WorkflowHandle> {
 		const stubId = this.env.ENGINE.idFromName(id);
 		const stub = this.env.ENGINE.get(stubId);
 
@@ -210,9 +217,11 @@ export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
 		await this.stub.changeInstanceStatus("resume");
 	}
 
-	public async terminate(): Promise<void> {
+	public async terminate(
+		options?: WorkflowInstanceTerminateOptions
+	): Promise<void> {
 		try {
-			await this.stub.changeInstanceStatus("terminate");
+			await this.stub.changeInstanceStatus("terminate", undefined, options);
 		} catch (e) {
 			// terminate causes instance abortion
 			if (!isUserTriggeredTerminate(e)) {

packages/workflows-shared/src/engine.ts

@@ -38,7 +38,10 @@ import {
 } from "./lib/streams";
 import { TimePriorityQueue } from "./lib/timePriorityQueue";
 import { MODIFIER_KEYS, WorkflowInstanceModifier } from "./modifier";
-import type { RestartFromStep } from "./binding";
+import type {
+	RestartFromStep,
+	WorkflowInstanceTerminateOptions,
+} from "./binding";
 import type { Event } from "./context";
 import type { InstanceMetadata, RawInstanceLog } from "./instance";
 import type { RollbackRegistryEntry } from "./lib/rollback";
@@ -804,8 +807,9 @@ export class Engine extends DurableObject<Env> {
 
 	async changeInstanceStatus(
 		newStatus: "resume" | "pause" | "terminate" | "restart",
-		from?: RestartFromStep
-	) {
+		from?: RestartFromStep,
+		terminateOptions?: WorkflowInstanceTerminateOptions
+	): Promise<void> {
 		const metadata =
 			await this.ctx.storage.get<InstanceMetadata>(INSTANCE_METADATA);
 
@@ -849,7 +853,7 @@ export class Engine extends DurableObject<Env> {
 						"instance.cannot_terminate"
 					);
 				}
-				await this.userTriggeredTerminate();
+				await this.userTriggeredTerminate(terminateOptions);
 				break;
 			}
 			case "restart":
@@ -864,7 +868,7 @@ export class Engine extends DurableObject<Env> {
 		}
 	}
 
-	async userTriggeredTerminate() {
+	async userTriggeredTerminate(options?: WorkflowInstanceTerminateOptions) {
 		const metadata =
 			await this.ctx.storage.get<InstanceMetadata>(INSTANCE_METADATA);
 
@@ -875,6 +879,16 @@ export class Engine extends DurableObject<Env> {
 			);
 		}
 
+		if (options?.rollback === true && this.rollbackRegistry.size > 0) {
+			const error = new Error("Instance terminated during rollback");
+			error.name = "Terminated";
+			try {
+				await executeRollbacks(this, error);
+			} catch (rollbackErr) {
+				console.error("Rollback execution failed:", rollbackErr);
+			}
+		}
+
 		this.writeLog(InstanceEvent.WORKFLOW_TERMINATED, null, null, {
 			trigger: {
 				source: InstanceTrigger.API,

packages/workflows-shared/tests/engine.test.ts

@@ -1673,6 +1673,52 @@ describe("Rollback", () => {
 		expect(countOf(logs, InstanceEvent.ROLLBACK_COMPLETE)).toBe(0);
 	});
 
+	it("runs rollback when terminate requests it", async ({ expect }) => {
+		const instanceId = "RB-TERMINATE";
+		const engineStub = await runWorkflow(instanceId, async (_e, step) => {
+			await doWithRollback(
+				step,
+				"setup-resource",
+				async () => "resource-id",
+				rollbackOptions()
+			);
+			await step.sleep("wait-forever", "1 hour");
+		});
+
+		await readLogsAfter(engineStub, (logs) =>
+			logs.logs.some(
+				(log) =>
+					log.event === InstanceEvent.STEP_SUCCESS &&
+					log.target === "setup-resource-1"
+			)
+		);
+
+		try {
+			await runInDurableObject(engineStub, async (engine) => {
+				await engine.changeInstanceStatus("terminate", undefined, {
+					rollback: true,
+				});
+			});
+		} catch (error) {
+			if (
+				!(error instanceof Error) ||
+				!error.message.startsWith("Aborting engine:")
+			) {
+				throw error;
+			}
+		}
+
+		const logs = await readLogsAfter(
+			env.ENGINE.get(env.ENGINE.idFromName(instanceId)),
+			(logs) =>
+				logs.logs.some((log) => log.event === InstanceEvent.WORKFLOW_TERMINATED)
+		);
+		expect(targetsOf(logs, InstanceEvent.ROLLBACK_STEP_SUCCESS)).toEqual([
+			"setup-resource-1",
+		]);
+		expect(countOf(logs, InstanceEvent.ROLLBACK_COMPLETE)).toBe(1);
+	});
+
 	it("does not run rollback when workflow succeeds", async ({ expect }) => {
 		const stub = await runWorkflowAndAwait("RB-NOOP", async (_e, step) => {
 			await doWithRollback(step, "a", async () => "ok", rollbackOptions());