rivetkit 2.3.23 → 2.3.24-rc.1

package/src/registry/run-handler-coordinator.test.ts

@@ -1,5 +1,6 @@
 import { describe, expect, test, vi } from "vitest";
 import { defineRunHandler, type RunControl } from "@/actor/config";
+import { workflow } from "@/workflow/mod";
 import { RunHandlerCoordinator } from "./run-handler-coordinator";
 
 class Deferred<T = void> {
@@ -51,6 +52,107 @@ function createCoordinator() {
 }
 
 describe("RunHandlerCoordinator", () => {
+	test("workflow inspector registrations stay per generation", () => {
+		const coordinator = new RunHandlerCoordinator(workflow(async () => {}));
+		const restart = vi.fn();
+
+		const oldInspector = coordinator.getInspector(
+			"actor-1",
+			restart,
+			1,
+		)?.workflow;
+		const newInspector = coordinator.getInspector(
+			"actor-1",
+			restart,
+			2,
+		)?.workflow;
+		expect(oldInspector).toBeDefined();
+		expect(newInspector).toBeDefined();
+		expect(newInspector).not.toBe(oldInspector);
+
+		// A superseded generation cannot re-initialize or rebind the current generation's
+		// workflow registration.
+		expect(coordinator.getInspector("actor-1", restart, 1)).toBeUndefined();
+		expect(coordinator.getInspector("actor-1", restart, 2)?.workflow).toBe(
+			newInspector,
+		);
+	});
+
+	test("a superseded generation's run is rejected before it touches the inspector", async () => {
+		const coordinator = new RunHandlerCoordinator(workflow(async () => {}));
+		const restart = vi.fn();
+		const newInspector = coordinator.getInspector(
+			"actor-1",
+			restart,
+			2,
+		)?.workflow;
+
+		await expect(
+			coordinator.run("actor-1", restart, async () => {}, 1),
+		).resolves.toBe("closed");
+		expect(coordinator.getInspector("actor-1", restart, 2)?.workflow).toBe(
+			newInspector,
+		);
+	});
+
+	test("a newer generation does not wait on an older generation's stuck run", async () => {
+		const coordinator = new RunHandlerCoordinator(
+			defineRunHandler(async () => {}, {
+				inspectorKind: "workflow",
+				createInspector: () => ({
+					inspector: {
+						workflow: {
+							getHistory: () => null,
+							getState: async () => null,
+							onHistoryUpdated: () => () => {},
+							replayFromStep: async () => null,
+						},
+					},
+				}),
+			}),
+		);
+		const oldRun = new Deferred();
+		const oldOutcome = coordinator.run(
+			"actor-1",
+			() => {},
+			() => oldRun.promise,
+			1,
+		);
+
+		// The lost generation's JS callback cannot be cancelled, so it stays active.
+		await expect(
+			coordinator.run(
+				"actor-1",
+				() => {},
+				async () => {},
+				2,
+			),
+		).resolves.toBe("ran");
+
+		// The superseded generation cannot start new runs, and its late destroy leaves the
+		// newer generation's state in place.
+		await expect(
+			coordinator.run(
+				"actor-1",
+				() => {},
+				async () => {},
+				1,
+			),
+		).resolves.toBe("closed");
+		coordinator.destroy("actor-1", 1);
+		await expect(
+			coordinator.run(
+				"actor-1",
+				() => {},
+				async () => {},
+				2,
+			),
+		).resolves.toBe("ran");
+
+		oldRun.resolve();
+		await expect(oldOutcome).resolves.toBe("ran");
+	});
+
 	test("fails explicitly when a JavaScript factory omits required workflow controls", async () => {
 		const coordinator = new RunHandlerCoordinator(
 			defineRunHandler(async () => {}, {

package/src/registry/run-handler-coordinator.ts

@@ -11,6 +11,8 @@ import { RivetError } from "@/actor/errors";
 type RunHandler = ((...args: any[]) => any) | { run: (...args: any[]) => any };
 
 interface ActorRunState {
+	/** Actor generation that owns this state, when the runtime reports one. */
+	actorGeneration: number | undefined;
 	active: boolean;
 	closed: boolean;
 	exclusive: boolean;
@@ -46,6 +48,32 @@ function notifyStateChanged(state: ActorRunState): void {
 	}
 }
 
+function createRunState(actorGeneration: number | undefined): ActorRunState {
+	return {
+		actorGeneration,
+		active: false,
+		closed: false,
+		exclusive: false,
+		exclusiveGeneration: 0,
+		inspectorInitialized: false,
+		queued: 0,
+		waiters: new Set(),
+	};
+}
+
+/**
+ * Closes run state. The inspector disposer is user code, so it is skipped for a generation
+ * that was declared lost.
+ */
+function closeRunState(state: ActorRunState, disposeInspector = true): void {
+	state.closed = true;
+	if (disposeInspector) {
+		state.inspector?.dispose?.();
+	}
+	state.inspector = undefined;
+	notifyStateChanged(state);
+}
+
 function waitForStateChange(state: ActorRunState): Promise<void> {
 	return new Promise((resolve) => {
 		state.waiters.add(resolve);
@@ -69,8 +97,14 @@ export class RunHandlerCoordinator {
 		actorId: string,
 		restart: () => void | Promise<void>,
 		callback: () => void | Promise<void>,
+		actorGeneration?: number,
 	): Promise<RunHandlerOutcome> {
-		const state = this.#getOrCreate(actorId);
+		const state = this.#getOrCreate(actorId, actorGeneration);
+		// Closed state belongs to a superseded generation. It must not initialize an inspector,
+		// which would rebind the workflow registration of the current generation.
+		if (state.closed) {
+			return "closed";
+		}
 		state.restart = restart;
 		this.#initializeInspector(actorId, state);
 		state.queued += 1;
@@ -110,21 +144,37 @@ export class RunHandlerCoordinator {
 	getInspector(
 		actorId: string,
 		restart: () => void | Promise<void>,
+		actorGeneration?: number,
 	): RunInspectorConfig | undefined {
-		const state = this.#getOrCreate(actorId);
+		const state = this.#getOrCreate(actorId, actorGeneration);
+		if (state.closed) {
+			return undefined;
+		}
 		state.restart = restart;
 		this.#initializeInspector(actorId, state);
 		return state.inspector?.inspector;
 	}
 
-	destroy(actorId: string): void {
+	/**
+	 * Closes the run state for `actorId`. With `actorGeneration`, only state owned by that
+	 * generation is closed, so an older generation's late cleanup cannot close a newer one.
+	 */
+	destroy(
+		actorId: string,
+		actorGeneration?: number,
+		options: { disposeInspector?: boolean } = {},
+	): void {
 		const state = this.#states.get(actorId);
 		if (!state) return;
+		if (
+			actorGeneration !== undefined &&
+			state.actorGeneration !== undefined &&
+			state.actorGeneration !== actorGeneration
+		) {
+			return;
+		}
 
-		state.closed = true;
-		state.inspector?.dispose?.();
-		state.inspector = undefined;
-		notifyStateChanged(state);
+		closeRunState(state, options.disposeInspector ?? true);
 		this.#states.delete(actorId);
 	}
 
@@ -173,18 +223,34 @@ export class RunHandlerCoordinator {
 		};
 	}
 
-	#getOrCreate(actorId: string): ActorRunState {
+	#getOrCreate(
+		actorId: string,
+		actorGeneration: number | undefined,
+	): ActorRunState {
 		let state = this.#states.get(actorId);
+		if (
+			state &&
+			actorGeneration !== undefined &&
+			state.actorGeneration !== undefined &&
+			state.actorGeneration !== actorGeneration
+		) {
+			if (actorGeneration < state.actorGeneration) {
+				// A superseded generation gets closed, detached state so it cannot run or
+				// block the current generation.
+				const detached = createRunState(actorGeneration);
+				detached.closed = true;
+				return detached;
+			}
+			// A newer generation replaces state left by an older one, which may still be
+			// marked active because its JS callback cannot be cancelled. That generation did not
+			// clean up after itself, which only happens when it was lost, so its user inspector
+			// disposer is not called.
+			closeRunState(state, false);
+			this.#states.delete(actorId);
+			state = undefined;
+		}
 		if (!state) {
-			state = {
-				active: false,
-				closed: false,
-				exclusive: false,
-				exclusiveGeneration: 0,
-				inspectorInitialized: false,
-				queued: 0,
-				waiters: new Set(),
-			};
+			state = createRunState(actorGeneration);
 			this.#states.set(actorId, state);
 		}
 		return state;
@@ -194,6 +260,7 @@ export class RunHandlerCoordinator {
 		if (state.inspectorInitialized) return;
 		const inspector = createRunInspector(this.#run, {
 			actorId,
+			actorGeneration: state.actorGeneration,
 			control: this.#control(actorId, state),
 		});
 		if (this.inspectorKind === "workflow") {

package/src/registry/runtime.ts

@@ -583,6 +583,8 @@ export interface CoreRuntime {
 		writes: RuntimeWorkflowKvWrite[],
 	): Promise<void>;
 	actorId(ctx: ActorContextHandle): string;
+	actorGeneration(ctx: ActorContextHandle): number | undefined;
+	actorIsLost(ctx: ActorContextHandle): boolean;
 	/**
 	 * Runs one actor callback with `ctx` as the current invocation: operations
 	 * on retained handles for the same actor resolve to it, and its Core span

package/src/registry/wasm-runtime.ts

@@ -538,6 +538,14 @@ export class WasmCoreRuntime implements CoreRuntime {
 		return callHandle(asWasmActorContext(ctx), "actorId");
 	}
 
+	actorGeneration(ctx: ActorContextHandle): number | undefined {
+		return callHandle(asWasmActorContext(ctx), "generation") ?? undefined;
+	}
+
+	actorIsLost(ctx: ActorContextHandle): boolean {
+		return callHandle(asWasmActorContext(ctx), "isLost");
+	}
+
 	runWithActorInvocationContext<T>(
 		_ctx: ActorContextHandle,
 		run: () => T,

package/src/workflow/context.ts

@@ -12,6 +12,7 @@ import type {
 	WorkflowContextInterface,
 	WorkflowQueueMessage,
 } from "@rivetkit/workflow-engine";
+import { EvictedError } from "@rivetkit/workflow-engine";
 import type {
 	QueueFilterName,
 	QueueNextBatchOptions,
@@ -19,6 +20,7 @@ import type {
 	QueueResultMessageForName,
 	RunContext,
 } from "@/actor/config";
+import { ACTOR_CONTEXT_INTERNAL_SYMBOL } from "@/actor/config";
 import type {
 	AnyActorDefinition,
 	BaseActorDefinition,
@@ -282,6 +284,20 @@ export type WorkflowBranchConfig<
 // never appears on the public surface.
 const DEACTIVATE_STEP = Symbol("workflow.step.deactivate");
 
+/**
+ * Stops the workflow before it starts more user code once its actor generation was declared
+ * lost. The engine treats the eviction like a sleep and replays from history in the next
+ * generation.
+ */
+export function throwIfGenerationLost(runCtx: unknown): void {
+	const actor = (
+		runCtx as { [ACTOR_CONTEXT_INTERNAL_SYMBOL]?: { isLost?: boolean } }
+	)[ACTOR_CONTEXT_INTERNAL_SYMBOL];
+	if (actor?.isLost === true) {
+		throw new EvictedError();
+	}
+}
+
 /**
  * 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
@@ -994,6 +1010,7 @@ export class WorkflowContext<
 			TQueues
 		>(this.#runCtx, () => this.#markGuardTriggered());
 
+		throwIfGenerationLost(this.#runCtx);
 		try {
 			return await run(stepCtx);
 		} finally {
@@ -1044,6 +1061,7 @@ export class WorkflowContext<
 			TEvents,
 			TQueues
 		>(this.#runCtx, () => this.#markGuardTriggered());
+		throwIfGenerationLost(this.#runCtx);
 		try {
 			await rollback(stepCtx, output);
 		} finally {
@@ -1097,6 +1115,8 @@ export class WorkflowContext<
 		TEvents,
 		TQueues
 	> {
+		// Every try, loop, join, and race body runs through a child context.
+		throwIfGenerationLost(this.#runCtx);
 		return new WorkflowContext(ctx, this.#runCtx);
 	}
 }

package/src/workflow/driver.ts

@@ -10,6 +10,7 @@ import type { RunContext } from "@/actor/config";
 import type { AnyStaticActorInstance } from "@/actor/definition";
 import { makeWorkflowKey, workflowStoragePrefix } from "@/actor/keys";
 import type { RawAccess } from "@/common/database/config";
+import { throwIfGenerationLost } from "./context";
 
 const WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix();
 // Keep workflow flushes below depot's 320-dirty-page commit ceiling. The
@@ -69,6 +70,26 @@ function normalizeSqlBlob(value: unknown): Uint8Array {
 	throw new Error("workflow sqlite value was not a blob");
 }
 
+function guardRawAccessAgainstLost(
+	access: RawAccess,
+	runCtx: RunContext<any, any, any, any, any, any, any, any> | undefined,
+): RawAccess {
+	return {
+		...access,
+		execute: ((...args: Parameters<RawAccess["execute"]>) => {
+			throwIfGenerationLost(runCtx);
+			return access.execute(...args);
+		}) as RawAccess["execute"],
+		transaction: (callback, options) => {
+			throwIfGenerationLost(runCtx);
+			return access.transaction(
+				(tx) => callback(guardRawAccessAgainstLost(tx, runCtx)),
+				options,
+			);
+		},
+	};
+}
+
 function runtimeDbFromContext(
 	runCtx?: RunContext<any, any, any, any, any, any, any, any>,
 ): RawAccess | undefined {
@@ -79,7 +100,10 @@ function runtimeDbFromContext(
 		"execute" in db &&
 		"transaction" in db
 	) {
-		return db as RawAccess;
+		// The client can come from a user database provider. Workflow storage calls it after
+		// awaits, including inside its own transactions, so every call checks that the
+		// generation was not lost first.
+		return guardRawAccessAgainstLost(db as RawAccess, runCtx);
 	}
 	return undefined;
 }

package/src/workflow/mod.ts

@@ -24,7 +24,7 @@ import { isActorAbortedError, RivetError } from "@/actor/errors";
 import type { EventSchemaConfig, QueueSchemaConfig } from "@/actor/schema";
 import type { AnyDatabaseProvider } from "@/common/database/config";
 import { stringifyError } from "@/utils";
-import { WorkflowContext } from "./context";
+import { throwIfGenerationLost, WorkflowContext } from "./context";
 import { ActorWorkflowControlDriver, ActorWorkflowDriver } from "./driver";
 import { createWorkflowInspectorAdapter } from "./inspector";
 
@@ -178,13 +178,28 @@ export function workflow<
 	type WorkflowInspectorRegistration = ReturnType<
 		typeof createWorkflowInspectorAdapter
 	> & { control?: RunControl };
+	// Keyed by actor id and generation so a lost generation that is still running cannot rebind
+	// the control or state of the generation that replaced it.
 	const workflowInspectors = new Map<string, WorkflowInspectorRegistration>();
 
-	function getWorkflowInspector(actorId: string) {
-		let workflowInspector = workflowInspectors.get(actorId);
+	function workflowInspectorKey(
+		actorId: string,
+		actorGeneration: number | undefined,
+	): string {
+		return actorGeneration === undefined
+			? actorId
+			: `${actorId}#${actorGeneration}`;
+	}
+
+	function getWorkflowInspector(
+		actorId: string,
+		actorGeneration: number | undefined,
+	) {
+		const key = workflowInspectorKey(actorId, actorGeneration);
+		let workflowInspector = workflowInspectors.get(key);
 		if (!workflowInspector) {
 			workflowInspector = createWorkflowInspectorAdapter();
-			workflowInspectors.set(actorId, workflowInspector);
+			workflowInspectors.set(key, workflowInspector);
 		}
 		return workflowInspector;
 	}
@@ -207,7 +222,11 @@ export function workflow<
 			}
 		)[ACTOR_CONTEXT_INTERNAL_SYMBOL];
 		invariant(actor, "workflow() requires an actor instance");
-		const workflowInspector = getWorkflowInspector(actor.id);
+		const actorGeneration = (actor as { generation?: number }).generation;
+		const workflowInspector = getWorkflowInspector(
+			actor.id,
+			actorGeneration,
+		);
 
 		const driver = new ActorWorkflowDriver(actor, runCtx);
 		const controlDriver = new ActorWorkflowControlDriver(actor, runCtx);
@@ -246,7 +265,10 @@ export function workflow<
 
 		const handle = runWorkflow(
 			actor.id,
-			async (ctx) => await fn(new WorkflowContext(ctx, runCtx)),
+			async (ctx) => {
+				throwIfGenerationLost(runCtx);
+				return await fn(new WorkflowContext(ctx, runCtx));
+			},
 			undefined,
 			driver,
 			{
@@ -256,7 +278,15 @@ export function workflow<
 				logger: runCtx.log as RunWorkflowOptions["logger"],
 				onHistoryUpdated: workflowInspector.update,
 				onError: onError
-					? async (event) => await onError(runCtx, event)
+					? async (event) => {
+							// A lost generation starts no new user hook.
+							if (
+								(actor as { isLost?: boolean }).isLost === true
+							) {
+								return;
+							}
+							await onError(runCtx, event);
+						}
 					: undefined,
 			},
 		);
@@ -302,16 +332,20 @@ export function workflow<
 	return defineRunHandler(run, {
 		icon: "diagram-project",
 		inspectorKind: "workflow",
-		createInspector: ({ actorId, control }) => {
-			const workflowInspector = getWorkflowInspector(actorId);
+		createInspector: ({ actorId, actorGeneration, control }) => {
+			const key = workflowInspectorKey(actorId, actorGeneration);
+			const workflowInspector = getWorkflowInspector(
+				actorId,
+				actorGeneration,
+			);
 			workflowInspector.control = control;
 			return {
 				inspector: { workflow: workflowInspector.adapter },
 				dispose: () => {
 					// Do not let a stale disposer remove a newly-created adapter for
-					// the same actor id.
-					if (workflowInspectors.get(actorId) === workflowInspector) {
-						workflowInspectors.delete(actorId);
+					// the same actor generation.
+					if (workflowInspectors.get(key) === workflowInspector) {
+						workflowInspectors.delete(key);
 					}
 				},
 			};