test(rivetkit): expect workflow step rollback

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts

@@ -166,6 +166,7 @@ import {
 	workflowSleepActor,
 	workflowSpawnChildActor,
 	workflowSpawnParentActor,
+	workflowStepRollbackActor,
 	workflowStopTeardownActor,
 	workflowTryActor,
 } from "./workflow";
@@ -322,6 +323,7 @@ export const registry = setup({
 		workflowRunningStepActor,
 		workflowReplayActor,
 		workflowSleepActor,
+		workflowStepRollbackActor,
 		workflowTryActor,
 		warmupActor,
 		workflowStopTeardownActor,

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/workflow.ts

@@ -455,12 +455,17 @@ export const workflowTryActor = actor({
 		} | null,
 		tryJoinFailure: null as string | null,
 	},
+	vars: {
+		innerWrites: 0,
+		recoveryWrites: 0,
+	},
 	run: workflow(async (ctx) => {
 		const stepResult = await ctx.tryStep({
 			name: "charge-card",
 			maxRetries: 0,
 			run: async () => {
 				ctx.state.innerWrites += 1;
+				ctx.vars.innerWrites += 1;
 				throw new Error("card declined");
 			},
 		});
@@ -479,6 +484,7 @@ export const workflowTryActor = actor({
 		});
 
 		await ctx.step("store-try-results", async () => {
+			ctx.vars.recoveryWrites += 1;
 			if (!stepResult.ok) {
 				ctx.state.tryStepFailure = {
 					kind: stepResult.failure.kind,
@@ -492,7 +498,43 @@ export const workflowTryActor = actor({
 		});
 	}),
 	actions: {
-		getState: (c) => c.state,
+		getState: (c) => ({ ...c.state, vars: c.vars }),
+	},
+	options: {
+		sleepTimeout: 50,
+	},
+});
+
+export const workflowStepRollbackActor = actor({
+	state: {
+		failedStateWrites: 0,
+		recoveryStateWrites: 0,
+		failureCaught: false,
+	},
+	vars: {
+		failedVarsWrites: 0,
+		recoveryVarsWrites: 0,
+	},
+	run: workflow(async (ctx) => {
+		const stepResult = await ctx.try(
+			"recover-failed-step",
+			async (tryCtx) => {
+				await tryCtx.step("failing-step", async () => {
+					tryCtx.state.failedStateWrites += 1;
+					tryCtx.vars.failedVarsWrites += 1;
+					throw new Error("step rollback");
+				});
+			},
+		);
+
+		await ctx.step("record-recovery", async () => {
+			ctx.state.recoveryStateWrites += 1;
+			ctx.vars.recoveryVarsWrites += 1;
+			ctx.state.failureCaught = !stepResult.ok;
+		});
+	}),
+	actions: {
+		getSnapshot: (c) => ({ state: c.state, vars: c.vars }),
 	},
 	options: {
 		sleepTimeout: 50,

rivetkit-typescript/packages/rivetkit/tests/driver/actor-workflow.test.ts

@@ -231,7 +231,10 @@ describeDriverMatrix("Actor Workflow", (driverTestConfig) => {
 				state = await actor.getState();
 			}
 
-			expect(state.innerWrites).toBe(1);
+			// Failed workflow steps roll back actor state mutations.
+			expect(state.innerWrites).toBe(0);
+			expect(state.vars.innerWrites).toBe(0);
+			expect(state.vars.recoveryWrites).toBe(1);
 			expect(state.tryStepFailure).toEqual({
 				kind: "exhausted",
 				message: "card declined",
@@ -240,6 +243,29 @@ describeDriverMatrix("Actor Workflow", (driverTestConfig) => {
 			expect(state.tryJoinFailure).toBe("join:parallel");
 		});
 
+		test("failed workflow step rolls back state and vars", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+			const actor = client.workflowStepRollbackActor.getOrCreate([
+				"workflow-step-rollback",
+			]);
+
+			let snapshot = await actor.getSnapshot();
+			for (let i = 0; i < 40 && !snapshot.state.failureCaught; i++) {
+				await waitFor(driverTestConfig, 50);
+				snapshot = await actor.getSnapshot();
+			}
+
+			expect(snapshot.state).toMatchObject({
+				failedStateWrites: 0,
+				recoveryStateWrites: 1,
+				failureCaught: true,
+			});
+			expect(snapshot.vars).toMatchObject({
+				failedVarsWrites: 0,
+				recoveryVarsWrites: 1,
+			});
+		});
+
 		test("sleeps and resumes between ticks", async (c) => {
 			const { client } = await setupDriverTest(c, driverTestConfig);
 			const actor = client.workflowSleepActor.getOrCreate([