diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index af10f8fe0d..1f83c8a6f7 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -372,6 +372,7 @@ export declare class ActorContext { sleep(): void destroy(): void destroyRequested(): boolean + generation(): number | null waitForDestroyCompletion(): Promise setPreventSleep(preventSleep: boolean): void preventSleep(): boolean diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index 842f6c71d3..e1cc6a99f7 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -607,6 +607,11 @@ impl ActorContext { self.inner.is_destroy_requested() } + #[napi] + pub fn generation(&self) -> Option { + self.inner.generation() + } + #[napi] pub async fn wait_for_destroy_completion(&self) { self.inner.wait_for_destroy_completion_public().await; diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index 57b8eebe37..d6062e3974 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -1281,6 +1281,11 @@ impl WasmActorContext { self.inner.actor_id().to_owned() } + #[wasm_bindgen(js_name = generation)] + pub fn generation(&self) -> Option { + self.inner.generation() + } + #[wasm_bindgen] pub fn name(&self) -> String { self.inner.name().to_owned() diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts index d33ac26ed1..b836fefc26 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/config.ts @@ -1079,6 +1079,8 @@ export interface RunControl { /** @experimental */ export interface RunInspectorFactoryContext { actorId: string; + /** Actor generation the inspector belongs to, when the runtime reports one. */ + actorGeneration?: number; control: RunControl; } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index f822e70a0e..05c834aeff 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -600,6 +600,10 @@ export class NapiCoreRuntime implements CoreRuntime { return asNativeActorContext(ctx).actorId(); } + actorGeneration(ctx: ActorContextHandle): number | undefined { + return asNativeActorContext(ctx).generation() ?? undefined; + } + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { return this.#runAs(asNativeActorContext(ctx), run); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 0c5c55648e..ce4a5af4bd 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -2963,6 +2963,10 @@ export class ActorContextHandleAdapter { return callNativeSync(() => this.#runtime.actorId(this.#ctx)); } + get actorGeneration(): number | undefined { + return callNativeSync(() => this.#runtime.actorGeneration(this.#ctx)); + } + get name(): string { return callNativeSync(() => this.#runtime.actorName(this.#ctx)); } @@ -3658,6 +3662,14 @@ class NativeWorkflowRuntimeAdapter { #onMessagesReceived: () => void; #completions = new Map Promise>(); + /** + * Actor generation, used to keep workflow registrations of different generations apart. Read + * lazily because most contexts never run a workflow. + */ + get generation(): number | undefined { + return this.#ctx.actorGeneration; + } + readonly id: string; readonly driver: { kvBatchGet: ( @@ -4104,9 +4116,14 @@ export function buildNativeFactory( const actorId = callNativeSync(() => runtime.actorId(ctx)); const restart = () => callNativeSync(() => runtime.actorRestartRunHandler(ctx)); - return (runHandlerCoordinator?.getInspector(actorId, restart) - ?.workflow ?? - getRunInspectorConfig(config.run, actorId)?.workflow) as + const actorGeneration = callNativeSync(() => + runtime.actorGeneration(ctx), + ); + return (runHandlerCoordinator?.getInspector( + actorId, + restart, + actorGeneration, + )?.workflow ?? getRunInspectorConfig(config.run, actorId)?.workflow) as | NativeWorkflowInspectorConfig | undefined; }; @@ -4824,7 +4841,10 @@ export function buildNativeFactory( saveActorState, ); } finally { - runHandlerCoordinator?.destroy(actorId); + runHandlerCoordinator?.destroy( + actorId, + callNativeSync(() => runtime.actorGeneration(ctx)), + ); disposeRunInspector(config.run, actorId); await actorCtx.dispose(); } @@ -4838,7 +4858,10 @@ export function buildNativeFactory( const actorId = callNativeSync(() => runtime.actorId(ctx)); // Close run control before user cleanup so replay cannot race actor // destruction. Recreating this actor id receives a fresh controller. - runHandlerCoordinator?.destroy(actorId); + runHandlerCoordinator?.destroy( + actorId, + callNativeSync(() => runtime.actorGeneration(ctx)), + ); disposeRunInspector(config.run, actorId); try { if (typeof config.onDestroy === "function") { @@ -5349,6 +5372,7 @@ export function buildNativeFactory( runtime.actorRestartRunHandler(ctx), ), executeRun, + callNativeSync(() => runtime.actorGeneration(ctx)), ); // Only a start suppressed behind a failed replay must preserve a // consumed durable wake. Closed actors must remain closed. diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts index 9a888d3f7b..31686e2090 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.test.ts +++ b/rivetkit-typescript/packages/rivetkit/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 { @@ -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 () => {}, { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts index e15a65f2de..1cc8d391dd 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts +++ b/rivetkit-typescript/packages/rivetkit/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,26 @@ 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(), + }; +} + +function closeRunState(state: ActorRunState): void { + state.closed = true; + state.inspector?.dispose?.(); + state.inspector = undefined; + notifyStateChanged(state); +} + function waitForStateChange(state: ActorRunState): Promise { return new Promise((resolve) => { state.waiters.add(resolve); @@ -69,8 +91,14 @@ export class RunHandlerCoordinator { actorId: string, restart: () => void | Promise, callback: () => void | Promise, + actorGeneration?: number, ): Promise { - 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 +138,33 @@ export class RunHandlerCoordinator { getInspector( actorId: string, restart: () => void | Promise, + 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): 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); this.#states.delete(actorId); } @@ -173,18 +213,32 @@ 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. + closeRunState(state); + 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 +248,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") { diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index 286868ff45..dfa85e8ba4 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -583,6 +583,7 @@ export interface CoreRuntime { writes: RuntimeWorkflowKvWrite[], ): Promise; actorId(ctx: ActorContextHandle): string; + actorGeneration(ctx: ActorContextHandle): number | undefined; /** * 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 diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index 052c0fb591..9384f17d17 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -538,6 +538,10 @@ export class WasmCoreRuntime implements CoreRuntime { return callHandle(asWasmActorContext(ctx), "actorId"); } + actorGeneration(ctx: ActorContextHandle): number | undefined { + return callHandle(asWasmActorContext(ctx), "generation") ?? undefined; + } + runWithActorInvocationContext( _ctx: ActorContextHandle, run: () => T, diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts index 7a41686ea3..5e024b495c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts @@ -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(); - 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); @@ -302,16 +321,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); } }, }; diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index 867a9a6743..afeef468f5 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -196,6 +196,10 @@ class FakeActorContext { }; } + generation(): number | null { + return null; + } + restartRunHandler(): void { this.scenario.runRestarts += 1; }