fix(rivetkit): stall stop handler until start completes

Author: MasterPtato

frontend/packages/example-registry/src/_gen.ts

@@ -68,6 +68,7 @@ import {
 	workflowQueueActor,
 	workflowSleepActor,
 } from "./workflow";
+import { startStopRaceActor, lifecycleObserver } from "./start-stop-race";
 
 // Consolidated setup with all actors
 export const registry = setup({
@@ -160,5 +161,8 @@ export const registry = setup({
 		dbActorDrizzle,
 		// From stateless.ts
 		statelessActor,
+		// From start-stop-race.ts
+		startStopRaceActor,
+		lifecycleObserver,
 	},
 });

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

@@ -0,0 +1,71 @@
+import { actor } from "rivetkit";
+
+/**
+ * Actor designed to test start/stop race conditions.
+ * Has a slow initialization to make race conditions easier to trigger.
+ */
+export const startStopRaceActor = actor({
+	state: {
+		initialized: false,
+		startTime: 0,
+		destroyCalled: false,
+		startCompleted: false,
+	},
+	onWake: async (c) => {
+		c.state.startTime = Date.now();
+
+		// Simulate slow initialization to create window for race condition
+		await new Promise((resolve) => setTimeout(resolve, 100));
+
+		c.state.initialized = true;
+		c.state.startCompleted = true;
+	},
+	onDestroy: (c) => {
+		c.state.destroyCalled = true;
+		// Don't save state here - the actor framework will save it automatically
+	},
+	actions: {
+		getState: (c) => {
+			return {
+				initialized: c.state.initialized,
+				startTime: c.state.startTime,
+				destroyCalled: c.state.destroyCalled,
+				startCompleted: c.state.startCompleted,
+			};
+		},
+		ping: (c) => {
+			return "pong";
+		},
+		destroy: (c) => {
+			c.destroy();
+		},
+	},
+});
+
+/**
+ * Observer actor to track lifecycle events from other actors
+ */
+export const lifecycleObserver = actor({
+	state: {
+		events: [] as Array<{
+			actorKey: string;
+			event: string;
+			timestamp: number;
+		}>,
+	},
+	actions: {
+		recordEvent: (c, params: { actorKey: string; event: string }) => {
+			c.state.events.push({
+				actorKey: params.actorKey,
+				event: params.event,
+				timestamp: Date.now(),
+			});
+		},
+		getEvents: (c) => {
+			return c.state.events;
+		},
+		clearEvents: (c) => {
+			c.state.events = [];
+		},
+	},
+});

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/start-stop-race.ts

@@ -1,5 +1,6 @@
 import { describe } from "vitest";
 import type { DriverTestConfig } from "../mod";
+import { runActorLifecycleTests } from "./actor-lifecycle";
 import { runActorScheduleTests } from "./actor-schedule";
 import { runActorSleepTests } from "./actor-sleep";
 import { runActorStateTests } from "./actor-state";
@@ -14,5 +15,8 @@ export function runActorDriverTests(driverTestConfig: DriverTestConfig) {
 
 		// Run actor sleep tests
 		runActorSleepTests(driverTestConfig);
+
+		// Run actor lifecycle tests
+		runActorLifecycleTests(driverTestConfig);
 	});
 }

rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-driver.ts

@@ -0,0 +1,157 @@
+import { describe, expect, test } from "vitest";
+import type { DriverTestConfig } from "../mod";
+import { setupDriverTest } from "../utils";
+
+export function runActorLifecycleTests(driverTestConfig: DriverTestConfig) {
+	describe("Actor Lifecycle Tests", () => {
+		test("actor stop during start waits for start to complete", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			const actorKey = `test-stop-during-start-${Date.now()}`;
+
+			// Create actor - this starts the actor
+			const actor = client.startStopRaceActor.getOrCreate([actorKey]);
+
+			// Immediately try to call an action and then destroy
+			// This creates a race where the actor might not be fully started yet
+			const pingPromise = actor.ping();
+
+			// Get actor ID
+			const actorId = await actor.resolve();
+
+			// Destroy immediately while start might still be in progress
+			await actor.destroy();
+
+			// The ping should still complete successfully because destroy waits for start
+			const result = await pingPromise;
+			expect(result).toBe("pong");
+
+			// Verify actor was actually destroyed
+			let destroyed = false;
+			try {
+				await client.startStopRaceActor.getForId(actorId).ping();
+			} catch (err: any) {
+				destroyed = true;
+				expect(err.group).toBe("actor");
+				expect(err.code).toBe("not_found");
+			}
+			expect(destroyed).toBe(true);
+		});
+
+		test("actor stop before actor instantiation completes cleans up handler", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			const actorKey = `test-stop-before-instantiation-${Date.now()}`;
+
+			// Create multiple actors rapidly to increase chance of race
+			const actors = Array.from({ length: 5 }, (_, i) =>
+				client.startStopRaceActor.getOrCreate([
+					`${actorKey}-${i}`,
+				]),
+			);
+
+			// Resolve all actor IDs (this triggers start)
+			const ids = await Promise.all(actors.map((a) => a.resolve()));
+
+			// Immediately destroy all actors
+			await Promise.all(actors.map((a) => a.destroy()));
+
+			// Verify all actors were cleaned up
+			for (const id of ids) {
+				let destroyed = false;
+				try {
+					await client.startStopRaceActor.getForId(id).ping();
+				} catch (err: any) {
+					destroyed = true;
+					expect(err.group).toBe("actor");
+					expect(err.code).toBe("not_found");
+				}
+				expect(destroyed, `actor ${id} should be destroyed`).toBe(
+					true,
+				);
+			}
+		});
+
+		test("onBeforeActorStart completes before stop proceeds", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			const actorKey = `test-before-actor-start-${Date.now()}`;
+
+			// Create actor
+			const actor = client.startStopRaceActor.getOrCreate([actorKey]);
+
+			// Call action to ensure actor is starting
+			const statePromise = actor.getState();
+
+			// Destroy immediately
+			await actor.destroy();
+
+			// State should be initialized because onBeforeActorStart must complete
+			const state = await statePromise;
+			expect(state.initialized).toBe(true);
+			expect(state.startCompleted).toBe(true);
+		});
+
+		test("multiple rapid create/destroy cycles handle race correctly", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			// Perform multiple rapid create/destroy cycles
+			for (let i = 0; i < 10; i++) {
+				const actorKey = `test-rapid-cycle-${Date.now()}-${i}`;
+				const actor = client.startStopRaceActor.getOrCreate([
+					actorKey,
+				]);
+
+				// Trigger start
+				const resolvePromise = actor.resolve();
+
+				// Immediately destroy
+				const destroyPromise = actor.destroy();
+
+				// Both should complete without errors
+				await Promise.all([resolvePromise, destroyPromise]);
+			}
+
+			// If we get here without errors, the race condition is handled correctly
+			expect(true).toBe(true);
+		});
+
+		test("actor stop called with no actor instance cleans up handler", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			const actorKey = `test-cleanup-no-instance-${Date.now()}`;
+
+			// Create and immediately destroy
+			const actor = client.startStopRaceActor.getOrCreate([actorKey]);
+			const id = await actor.resolve();
+			await actor.destroy();
+
+			// Try to recreate with same key - should work without issues
+			const newActor = client.startStopRaceActor.getOrCreate([
+				actorKey,
+			]);
+			const result = await newActor.ping();
+			expect(result).toBe("pong");
+
+			// Clean up
+			await newActor.destroy();
+		});
+
+		test("onDestroy is called even when actor is destroyed during start", async (c) => {
+			const { client } = await setupDriverTest(c, driverTestConfig);
+
+			const actorKey = `test-ondestroy-during-start-${Date.now()}`;
+
+			// Create actor
+			const actor = client.startStopRaceActor.getOrCreate([actorKey]);
+
+			// Start and immediately destroy
+			const statePromise = actor.getState();
+			await actor.destroy();
+
+			// Verify onDestroy was called (requires actor to be started)
+			const state = await statePromise;
+			expect(state.destroyCalled).toBe(true);
+		});
+	});
+}

rivetkit-typescript/packages/rivetkit/src/driver-test-suite/tests/actor-lifecycle.ts

@@ -153,7 +153,7 @@ export class EngineActorDriver implements ActorDriver {
 			onConnected: () => {
 				this.#runnerStarted.resolve(undefined);
 			},
-			onDisconnected: (_code, _reason) => {},
+			onDisconnected: (_code, _reason) => { },
 			onShutdown: () => {
 				this.#runnerStopped.resolve(undefined);
 				this.#isRunnerStopped = true;
@@ -356,7 +356,7 @@ export class EngineActorDriver implements ActorDriver {
 	async serverlessHandleStart(c: HonoContext): Promise<Response> {
 		return streamSSE(c, async (stream) => {
 			// NOTE: onAbort does not work reliably
-			stream.onAbort(() => {});
+			stream.onAbort(() => { });
 			c.req.raw.signal.addEventListener("abort", () => {
 				logger().debug("SSE aborted, shutting down runner");
 
@@ -498,7 +498,26 @@ export class EngineActorDriver implements ActorDriver {
 		this.#actorStopIntent.delete(actorId);
 
 		const handler = this.#actors.get(actorId);
-		if (handler?.actor) {
+		if (!handler) {
+			logger().debug({ msg: "no runner actor handler to stop", actorId, reason });
+			return;
+		}
+
+		if (handler.actorStartPromise) {
+			try {
+				logger().debug({ msg: "runner actor stopping before it started, waiting", actorId, generation });
+				await handler.actorStartPromise.promise;
+			} catch (err) {
+				// Start failed, but we still want to clean up the handler
+				logger().debug({
+					msg: "actor start failed during stop, cleaning up handler",
+					actorId,
+					err: stringifyError(err),
+				});
+			}
+		}
+
+		if (handler.actor) {
 			try {
 				await handler.actor.onStop(reason);
 			} catch (err) {
@@ -507,9 +526,10 @@ export class EngineActorDriver implements ActorDriver {
 					err: stringifyError(err),
 				});
 			}
-			this.#actors.delete(actorId);
 		}
 
+		this.#actors.delete(actorId);
+
 		logger().debug({ msg: "runner actor stopped", actorId, reason });
 	}