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 via `wrangler workflows instances terminate --rollback`, including local mode.
+
+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

@@ -1169,6 +1169,11 @@ const config = {
 												},
 											},
 										},
+										rollback: {
+											type: "boolean",
+											description:
+												"Run rollback before terminating. Only valid when action is terminate.",
+										},
 									},
 								},
 							},

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

@@ -1747,6 +1747,10 @@ export type WorkflowsChangeInstanceStatusData = {
 			 */
 			type?: "do" | "sleep" | "waitForEvent";
 		};
+		/**
+		 * Run rollback before terminating. Only valid when action is terminate.
+		 */
+		rollback?: boolean;
 	};
 	path: {
 		workflow_name: WorkflowsWorkflowName;

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

@@ -1090,6 +1090,7 @@ export const zWorkflowsChangeInstanceStatusData = z.object({
 				type: z.enum(["do", "sleep", "waitForEvent"]).optional(),
 			})
 			.optional(),
+		rollback: z.boolean().optional(),
 	}),
 	path: z.object({
 		workflow_name: zWorkflowsWorkflowName,

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

@@ -1810,6 +1810,10 @@
 												"description": "The step type. Defaults to do."
 											}
 										}
+									},
+									"rollback": {
+										"type": "boolean",
+										"description": "Run rollback before terminating. Only valid when action is terminate."
 									}
 								}
 							}

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

@@ -11,7 +11,10 @@ import type {
 	WorkflowsWorkflow,
 } from "../generated";
 import type { zWorkflowsListInstancesData } from "../generated/zod.gen";
-import type { RestartFromStep } from "@cloudflare/workflows-shared/src/binding";
+import type {
+	RestartFromStep,
+	WorkflowInstanceTerminateOptions,
+} from "@cloudflare/workflows-shared/src/binding";
 import type { z } from "zod";
 
 // ============================================================================
@@ -35,7 +38,7 @@ interface WorkflowHandle {
 	pause(): Promise<void>;
 	resume(): Promise<void>;
 	restart(options?: { from?: RestartFromStep }): Promise<void>;
-	terminate(): Promise<void>;
+	terminate(options?: WorkflowInstanceTerminateOptions): Promise<void>;
 	sendEvent(args: { payload: unknown; type: string }): Promise<void>;
 	status(): Promise<{ status: string; output?: unknown; error?: unknown }>;
 }
@@ -886,6 +889,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) {
@@ -909,7 +920,9 @@ export async function changeWorkflowInstanceStatus(
 				break;
 			}
 			case "terminate":
-				await handle.terminate();
+				await (handle as unknown as WorkflowHandle).terminate(
+					body.rollback === true ? { rollback: true } : undefined
+				);
 				break;
 		}
 

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

@@ -1,6 +1,7 @@
 import type {
 	WorkflowBinding,
 	WorkflowInstanceRestartOptions,
+	WorkflowInstanceTerminateOptions,
 } from "@cloudflare/workflows-shared/src/binding";
 
 class WorkflowImpl implements Workflow {
@@ -84,9 +85,14 @@ class InstanceImpl implements WorkflowInstance {
 		await instance.resume();
 	}
 
-	public async terminate(): Promise<void> {
-		using instance = await this.getInstance();
-		await instance.terminate();
+	public async terminate(
+		options?: WorkflowInstanceTerminateOptions
+	): Promise<void> {
+		using instance = (await this.binding.get(this.id)) as Awaited<
+			ReturnType<WorkflowBinding["get"]>
+		> &
+			Disposable;
+		await instance.terminate(options);
 	}
 
 	public async restart(

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());