rivetkit 2.0.7-rc.1 → 2.0.8

package/src/driver-test-suite/tests/actor-schedule.ts

@@ -6,6 +6,8 @@ export function runActorScheduleTests(driverTestConfig: DriverTestConfig) {
 	describe.skipIf(driverTestConfig.skip?.schedule)(
 		"Actor Schedule Tests",
 		() => {
+			// See alarm + actor sleeping test in actor-sleep.ts
+
 			describe("Scheduled Alarms", () => {
 				test("executes c.schedule.at() with specific timestamp", async (c) => {
 					const { client } = await setupDriverTest(c, driverTestConfig);
@@ -13,12 +15,12 @@ export function runActorScheduleTests(driverTestConfig: DriverTestConfig) {
 					// Create instance
 					const scheduled = client.scheduled.getOrCreate();
 
-					// Schedule a task to run in 100ms using timestamp
-					const timestamp = Date.now() + 100;
+					// Schedule a task to run using timestamp
+					const timestamp = Date.now() + 250;
 					await scheduled.scheduleTaskAt(timestamp);
 
 					// Wait for longer than the scheduled time
-					await waitFor(driverTestConfig, 200);
+					await waitFor(driverTestConfig, 500);
 
 					// Verify the scheduled task ran
 					const lastRun = await scheduled.getLastRun();
@@ -34,11 +36,11 @@ export function runActorScheduleTests(driverTestConfig: DriverTestConfig) {
 					// Create instance
 					const scheduled = client.scheduled.getOrCreate();
 
-					// Schedule a task to run in 100ms using delay
-					await scheduled.scheduleTaskAfter(100);
+					// Schedule a task to run using delay
+					await scheduled.scheduleTaskAfter(250);
 
 					// Wait for longer than the scheduled time
-					await waitFor(driverTestConfig, 200);
+					await waitFor(driverTestConfig, 500);
 
 					// Verify the scheduled task ran
 					const lastRun = await scheduled.getLastRun();
@@ -48,31 +50,6 @@ export function runActorScheduleTests(driverTestConfig: DriverTestConfig) {
 					expect(scheduledCount).toBe(1);
 				});
 
-				test("scheduled tasks persist across actor restarts", async (c) => {
-					const { client } = await setupDriverTest(c, driverTestConfig);
-
-					// Create instance and schedule
-					const scheduled = client.scheduled.getOrCreate();
-					await scheduled.scheduleTaskAfter(200);
-
-					// Wait a little so the schedule is stored but hasn't triggered yet
-					await waitFor(driverTestConfig, 100);
-
-					// Get a new reference to simulate actor restart
-					const newInstance = client.scheduled.getOrCreate();
-
-					// Verify the schedule still exists but hasn't run yet
-					const initialCount = await newInstance.getScheduledCount();
-					expect(initialCount).toBe(0);
-
-					// Wait for the scheduled task to execute
-					await waitFor(driverTestConfig, 200);
-
-					// Verify the scheduled task ran after "restart"
-					const scheduledCount = await newInstance.getScheduledCount();
-					expect(scheduledCount).toBe(1);
-				});
-
 				test("multiple scheduled tasks execute in order", async (c) => {
 					const { client } = await setupDriverTest(c, driverTestConfig);
 
@@ -83,22 +60,22 @@ export function runActorScheduleTests(driverTestConfig: DriverTestConfig) {
 					await scheduled.clearHistory();
 
 					// Schedule multiple tasks with different delays
-					await scheduled.scheduleTaskAfterWithId("first", 100);
-					await scheduled.scheduleTaskAfterWithId("second", 300);
-					await scheduled.scheduleTaskAfterWithId("third", 500);
+					await scheduled.scheduleTaskAfterWithId("first", 250);
+					await scheduled.scheduleTaskAfterWithId("second", 750);
+					await scheduled.scheduleTaskAfterWithId("third", 1250);
 
 					// Wait for first task only
-					await waitFor(driverTestConfig, 200);
+					await waitFor(driverTestConfig, 500);
 					const history1 = await scheduled.getTaskHistory();
 					expect(history1).toEqual(["first"]);
 
 					// Wait for second task
-					await waitFor(driverTestConfig, 200);
+					await waitFor(driverTestConfig, 500);
 					const history2 = await scheduled.getTaskHistory();
 					expect(history2).toEqual(["first", "second"]);
 
 					// Wait for third task
-					await waitFor(driverTestConfig, 200);
+					await waitFor(driverTestConfig, 500);
 					const history3 = await scheduled.getTaskHistory();
 					expect(history3).toEqual(["first", "second", "third"]);
 				});

package/src/driver-test-suite/tests/actor-sleep.ts

@@ -29,7 +29,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			await sleepActor.triggerSleep();
 
 			// HACK: Wait for sleep to finish in background
-			await waitFor(driverTestConfig, 100);
+			await waitFor(driverTestConfig, 250);
 
 			// Get sleep count after restore
 			{
@@ -59,7 +59,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			await sleepActor.dispose();
 
 			// HACK: Wait for sleep to finish in background
-			await waitFor(driverTestConfig, 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Reconnect to get sleep count after restore
 			const sleepActor2 = client.sleep.getOrCreate();
@@ -84,7 +84,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait for sleep
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Get sleep count after restore
 			{
@@ -111,7 +111,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			await sleepActor.dispose();
 
 			// Wait for sleep
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Reconnect to get sleep count after restore
 			const sleepActor2 = client.sleep.getOrCreate();
@@ -136,7 +136,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait almost until sleep timeout, then make RPC call
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT - 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT - 250);
 
 			// RPC call should reset the sleep timer
 			{
@@ -146,7 +146,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait another partial timeout period - actor should still be awake
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT - 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT - 250);
 
 			// Actor should still be awake
 			{
@@ -156,7 +156,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Now wait for full timeout without any RPC calls
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should have slept and restarted
 			{
@@ -180,10 +180,10 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Set an alarm to keep the actor awake
-			await sleepActor.setAlarm(SLEEP_TIMEOUT - 100);
+			await sleepActor.setAlarm(SLEEP_TIMEOUT - 250);
 
 			// Wait until after SLEEPT_IMEOUT to validate the actor did not sleep
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should not have slept
 			{
@@ -207,7 +207,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Set an alarm to keep the actor awake
-			await sleepActor.setAlarm(SLEEP_TIMEOUT + 100);
+			await sleepActor.setAlarm(SLEEP_TIMEOUT + 250);
 
 			// Wait until after SLEEPT_IMEOUT to validate the actor did not sleep
 			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 200);
@@ -239,7 +239,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			);
 			const longRunningPromise = sleepActor.longRunningRpc();
 			await waitPromise;
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 			await sleepActor.finishLongRunningRpc();
 			await longRunningPromise;
 
@@ -252,7 +252,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			await sleepActor.dispose();
 
 			// Now wait for the sleep timeout
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should have slept after the timeout
 			const sleepActor2 = client.sleepWithLongRpc.getOrCreate();
@@ -298,7 +298,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			});
 
 			// Wait longer than sleep timeout while keeping WebSocket connected
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Send a message to check if actor is still alive
 			ws.send(JSON.stringify({ type: "getCounts" }));
@@ -320,7 +320,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			ws.close();
 
 			// Wait for sleep timeout after WebSocket closed
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should have slept after WebSocket closed
 			{
@@ -347,7 +347,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Start a long-running fetch request
-			const fetchDuration = SLEEP_TIMEOUT + 100;
+			const fetchDuration = SLEEP_TIMEOUT + 250;
 			const fetchPromise = sleepActor.fetch(
 				`long-request?duration=${fetchDuration}`,
 			);
@@ -365,7 +365,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait for sleep timeout
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should have slept after timeout
 			{
@@ -389,7 +389,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait longer than sleep timeout
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should NOT have slept due to noSleep option
 			{
@@ -399,7 +399,7 @@ export function runActorSleepTests(driverTestConfig: DriverTestConfig) {
 			}
 
 			// Wait even longer to be sure
-			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 100);
+			await waitFor(driverTestConfig, SLEEP_TIMEOUT + 250);
 
 			// Actor should still not have slept
 			{

package/src/drivers/engine/actor-driver.ts

@@ -35,7 +35,11 @@ import {
 } from "@/driver-helpers/mod";
 import type { RegistryConfig } from "@/registry/config";
 import type { RunConfig } from "@/registry/run-config";
-import { promiseWithResolvers } from "@/utils";
+import {
+	type LongTimeoutHandle,
+	promiseWithResolvers,
+	setLongTimeout,
+} from "@/utils";
 import type { Config } from "./config";
 import { KEYS } from "./kv";
 import { logger } from "./log";
@@ -58,6 +62,7 @@ export class EngineActorDriver implements ActorDriver {
 	#actors: Map<string, ActorHandler> = new Map();
 	#actorRouter: ActorRouter;
 	#version: number = 1; // Version for the runner protocol
+	#alarmTimeout?: LongTimeoutHandle;
 
 	#runnerStarted: PromiseWithResolvers<undefined> = Promise.withResolvers();
 	#runnerStopped: PromiseWithResolvers<undefined> = Promise.withResolvers();
@@ -74,17 +79,21 @@ export class EngineActorDriver implements ActorDriver {
 		this.#managerDriver = managerDriver;
 		this.#inlineClient = inlineClient;
 		this.#config = config;
-		this.#actorRouter = createActorRouter(runConfig, this);
+		this.#actorRouter = createActorRouter(
+			runConfig,
+			this,
+			registryConfig.test.enabled,
+		);
 
 		// Create runner configuration
 		let hasDisconnected = false;
 		const runnerConfig: RunnerConfig = {
 			version: this.#version,
 			endpoint: config.endpoint,
-			token: config.token,
+			token: runConfig.token ?? config.token,
 			pegboardEndpoint: config.pegboardEndpoint,
 			namespace: config.namespace,
-			totalSlots: config.totalSlots,
+			totalSlots: runConfig.totalSlots ?? config.totalSlots,
 			runnerName: config.runnerName,
 			runnerKey: config.runnerKey,
 			metadata: {
@@ -186,9 +195,29 @@ export class EngineActorDriver implements ActorDriver {
 	}
 
 	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
-		// TODO: Set timeout
-		// TODO: Use alarm on sleep
-		// TODO: Send alarm to runner
+		// Clear prev timeout
+		if (this.#alarmTimeout) {
+			this.#alarmTimeout.abort();
+			this.#alarmTimeout = undefined;
+		}
+
+		// Set alarm
+		const delay = Math.max(0, timestamp - Date.now());
+		this.#alarmTimeout = setLongTimeout(() => {
+			actor._onAlarm();
+			this.#alarmTimeout = undefined;
+		}, delay);
+
+		// TODO: This call may not be needed on ActorInstance.start, but it does help ensure that the local state is synced with the alarm state
+		// Set alarm on Rivet
+		//
+		// This does not call an "alarm" event like Durable Objects.
+		// Instead, it just wakes the actor on the alarm (if not
+		// already awake).
+		//
+		// _onAlarm is automatically called on `ActorInstance.start` when waking
+		// again.
+		this.#runner.setAlarm(actor.id, timestamp);
 	}
 
 	async getDatabase(_actorId: string): Promise<unknown | undefined> {
@@ -382,9 +411,9 @@ export class EngineActorDriver implements ActorDriver {
 
 		return streamSSE(c, async (stream) => {
 			// Runner id should be set if the runner started
-			const runnerId = this.#runner.runnerId;
-			invariant(runnerId, "runnerId not set");
-			stream.writeSSE({ data: runnerId });
+			const payload = this.#runner.getServerlessInitPacket();
+			invariant(payload, "runnerId not set");
+			stream.writeSSE({ data: payload });
 
 			return this.#runnerStopped.promise;
 		});

package/src/drivers/file-system/manager.ts

@@ -123,7 +123,11 @@ export class FileSystemManagerDriver implements ManagerDriver {
 			this,
 			inlineClient,
 		);
-		this.#actorRouter = createActorRouter(this.#runConfig, this.#actorDriver);
+		this.#actorRouter = createActorRouter(
+			this.#runConfig,
+			this.#actorDriver,
+			registryConfig.test.enabled,
+		);
 	}
 
 	async sendRequest(actorId: string, actorRequest: Request): Promise<Response> {

package/src/manager/router.ts

@@ -105,7 +105,10 @@ export function createManagerRouter(
 }
 
 function addServerlessRoutes(
-	serverlessActorDriverBuilder: () => ActorDriver,
+	serverlessActorDriverBuilder: (
+		token: string | undefined,
+		totalSlots: number | undefined,
+	) => ActorDriver,
 	router: OpenAPIHono,
 	cors: MiddlewareHandler,
 ) {
@@ -118,7 +121,13 @@ function addServerlessRoutes(
 
 	// Serverless start endpoint
 	router.get("/start", cors, async (c) => {
-		const actorDriver = serverlessActorDriverBuilder();
+		const token = c.req.header("x-rivet-token");
+		let totalSlots: number | undefined = parseInt(
+			c.req.header("x-rivetkit-total-slots") as any,
+		);
+		if (isNaN(totalSlots)) totalSlots = undefined;
+
+		const actorDriver = serverlessActorDriverBuilder(token, totalSlots);
 		invariant(
 			actorDriver.serverlessHandleStart,
 			"missing serverlessHandleStart on ActorDriver",

package/src/registry/mod.ts

@@ -215,7 +215,16 @@ export class Registry<A extends RegistryActors> {
 			console.log();
 		}
 
-		let serverlessActorDriverBuilder: (() => ActorDriver) | undefined = () => {
+		let serverlessActorDriverBuilder:
+			| ((token?: string, totalSlots?: number) => ActorDriver)
+			| undefined = (
+			token: string | undefined,
+			totalSlots: number | undefined,
+		) => {
+			// Override config
+			if (token) config.token = token;
+			if (totalSlots) config.totalSlots = totalSlots;
+
 			return driver.actor(this.#config, config, managerDriver, client);
 		};