rivetkit 2.0.21 → 2.0.22-rc.2

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

@@ -1,6 +1,7 @@
 import type {
 	ActorConfig as EngineActorConfig,
 	RunnerConfig as EngineRunnerConfig,
+	HibernationConfig,
 } from "@rivetkit/engine-runner";
 import { Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
@@ -9,12 +10,14 @@ import { streamSSE } from "hono/streaming";
 import { WSContext } from "hono/ws";
 import invariant from "invariant";
 import { lookupInRegistry } from "@/actor/definition";
+import { PERSIST_SYMBOL } from "@/actor/instance";
 import { deserializeActorKey } from "@/actor/keys";
 import { EncodingSchema } from "@/actor/protocol/serde";
 import { type ActorRouter, createActorRouter } from "@/actor/router";
 import {
 	handleRawWebSocketHandler,
 	handleWebSocketConnect,
+	truncateRawWebSocketPathPrefix,
 } from "@/actor/router-endpoints";
 import type { Client } from "@/client/client";
 import {
@@ -26,7 +29,11 @@ import {
 } from "@/common/actor-router-consts";
 import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { getLogger } from "@/common/log";
-import type { UniversalWebSocket } from "@/common/websocket-interface";
+import type {
+	RivetEvent,
+	RivetMessageEvent,
+	UniversalWebSocket,
+} from "@/common/websocket-interface";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
@@ -37,13 +44,18 @@ import { buildActorNames, type RegistryConfig } from "@/registry/config";
 import type { RunnerConfig } from "@/registry/run-config";
 import { getEndpoint } from "@/remote-manager-driver/api-utils";
 import {
+	arrayBuffersEqual,
+	idToStr,
 	type LongTimeoutHandle,
 	promiseWithResolvers,
 	setLongTimeout,
+	stringifyError,
 } from "@/utils";
 import { KEYS } from "./kv";
 import { logger } from "./log";
 
+const RUNNER_SSE_PING_INTERVAL = 1000;
+
 interface ActorHandler {
 	actor?: AnyActorInstance;
 	actorStartPromise?: ReturnType<typeof promiseWithResolvers<void>>;
@@ -65,6 +77,14 @@ export class EngineActorDriver implements ActorDriver {
 
 	#runnerStarted: PromiseWithResolvers<undefined> = promiseWithResolvers();
 	#runnerStopped: PromiseWithResolvers<undefined> = promiseWithResolvers();
+	#isRunnerStopped: boolean = false;
+
+	// WebSocket message acknowledgment debouncing
+	#wsAckQueue: Map<
+		string,
+		{ requestIdBuf: ArrayBuffer; messageIndex: number }
+	> = new Map();
+	#wsAckFlushInterval?: NodeJS.Timeout;
 
 	constructor(
 		registryConfig: RegistryConfig,
@@ -79,7 +99,7 @@ export class EngineActorDriver implements ActorDriver {
 
 		// HACK: Override inspector token (which are likely to be
 		// removed later on) with token from x-rivet-token header
-		const token = runConfig.token ?? runConfig.token;
+		const token = runConfig.token;
 		if (token && runConfig.inspector && runConfig.inspector.enabled) {
 			runConfig.inspector.token = () => token;
 		}
@@ -96,10 +116,10 @@ export class EngineActorDriver implements ActorDriver {
 			version: this.#version,
 			endpoint: getEndpoint(runConfig),
 			token,
-			namespace: runConfig.namespace ?? runConfig.namespace,
-			totalSlots: runConfig.totalSlots ?? runConfig.totalSlots,
-			runnerName: runConfig.runnerName ?? runConfig.runnerName,
-			runnerKey: runConfig.runnerKey,
+			namespace: runConfig.namespace,
+			totalSlots: runConfig.totalSlots,
+			runnerName: runConfig.runnerName,
+			runnerKey: runConfig.runnerKey ?? crypto.randomUUID(),
 			metadata: {
 				inspectorToken: this.#runConfig.inspector.token(),
 			},
@@ -121,22 +141,149 @@ export class EngineActorDriver implements ActorDriver {
 
 				this.#runnerStarted.resolve(undefined);
 			},
-			onDisconnected: () => {
+			onDisconnected: (code, reason) => {
 				logger().warn({
 					msg: "runner disconnected",
 					namespace: this.#runConfig.namespace,
 					runnerName: this.#runConfig.runnerName,
+					code,
+					reason,
 				});
 				hasDisconnected = true;
 			},
 			onShutdown: () => {
 				this.#runnerStopped.resolve(undefined);
+				this.#isRunnerStopped = true;
 			},
 			fetch: this.#runnerFetch.bind(this),
 			websocket: this.#runnerWebSocket.bind(this),
 			onActorStart: this.#runnerOnActorStart.bind(this),
 			onActorStop: this.#runnerOnActorStop.bind(this),
 			logger: getLogger("engine-runner"),
+			getActorHibernationConfig: (
+				actorId: string,
+				requestId: ArrayBuffer,
+				request: Request,
+			): HibernationConfig => {
+				const url = new URL(request.url);
+				const path = url.pathname;
+
+				// Get actor instance from runner to access actor name
+				const actorInstance = this.#runner.getActor(actorId);
+				if (!actorInstance) {
+					logger().warn({
+						msg: "actor not found in getActorHibernationConfig",
+						actorId,
+					});
+					return { enabled: false, lastMsgIndex: undefined };
+				}
+
+				// Load actor handler to access persisted data
+				const handler = this.#actors.get(actorId);
+				if (!handler) {
+					logger().warn({
+						msg: "actor handler not found in getActorHibernationConfig",
+						actorId,
+					});
+					return { enabled: false, lastMsgIndex: undefined };
+				}
+				if (!handler.actor) {
+					logger().warn({
+						msg: "actor not found in getActorHibernationConfig",
+						actorId,
+					});
+					return { enabled: false, lastMsgIndex: undefined };
+				}
+
+				// Check for existing WS
+				const existingWs = handler.actor[
+					PERSIST_SYMBOL
+				].hibernatableWebSocket.find((ws) =>
+					arrayBuffersEqual(ws.requestId, requestId),
+				);
+
+				// Determine configuration for new WS
+				let hibernationConfig: HibernationConfig;
+				if (existingWs) {
+					hibernationConfig = {
+						enabled: true,
+						lastMsgIndex: Number(existingWs.msgIndex),
+					};
+				} else {
+					if (path === PATH_CONNECT_WEBSOCKET) {
+						hibernationConfig = {
+							enabled: true,
+							lastMsgIndex: undefined,
+						};
+					} else if (path.startsWith(PATH_RAW_WEBSOCKET_PREFIX)) {
+						// Find actor config
+						const definition = lookupInRegistry(
+							this.#registryConfig,
+							actorInstance.config.name,
+						);
+
+						// Check if can hibernate
+						const canHibernatWebSocket =
+							definition.config.options?.canHibernatWebSocket;
+						if (canHibernatWebSocket === true) {
+							hibernationConfig = {
+								enabled: true,
+								lastMsgIndex: undefined,
+							};
+						} else if (typeof canHibernatWebSocket === "function") {
+							try {
+								// Truncate the path to match the behavior on onRawWebSocket
+								const newPath = truncateRawWebSocketPathPrefix(
+									url.pathname,
+								);
+								const truncatedRequest = new Request(
+									`http://actor${newPath}`,
+									request,
+								);
+
+								const canHibernate =
+									canHibernatWebSocket(truncatedRequest);
+								hibernationConfig = {
+									enabled: canHibernate,
+									lastMsgIndex: undefined,
+								};
+							} catch (error) {
+								logger().error({
+									msg: "error calling canHibernatWebSocket",
+									error,
+								});
+								hibernationConfig = {
+									enabled: false,
+									lastMsgIndex: undefined,
+								};
+							}
+						} else {
+							hibernationConfig = {
+								enabled: false,
+								lastMsgIndex: undefined,
+							};
+						}
+					} else {
+						logger().warn({
+							msg: "unexpected path for getActorHibernationConfig",
+							path,
+						});
+						hibernationConfig = {
+							enabled: false,
+							lastMsgIndex: undefined,
+						};
+					}
+				}
+
+				// Save hibernatable WebSocket
+				handler.actor[PERSIST_SYMBOL].hibernatableWebSocket.push({
+					requestId,
+					lastSeenTimestamp: BigInt(Date.now()),
+					msgIndex: -1n,
+				});
+
+				return hibernationConfig;
+			},
 		};
 
 		// Create and start runner
@@ -148,6 +295,15 @@ export class EngineActorDriver implements ActorDriver {
 			namespace: runConfig.namespace,
 			runnerName: runConfig.runnerName,
 		});
+
+		// Start WebSocket ack flush interval
+		//
+		// Decreasing this reduces the amount of buffered messages on the
+		// gateway
+		//
+		// Gateway timeout configured to 30s
+		// https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/shared_state.rs#L17
+		this.#wsAckFlushInterval = setInterval(() => this.#flushWsAcks(), 1000);
 	}
 
 	async #loadActorHandler(actorId: string): Promise<ActorHandler> {
@@ -166,6 +322,19 @@ export class EngineActorDriver implements ActorDriver {
 		return handler.actor;
 	}
 
+	#flushWsAcks(): void {
+		if (this.#wsAckQueue.size === 0) return;
+
+		for (const {
+			requestIdBuf: requestId,
+			messageIndex: index,
+		} of this.#wsAckQueue.values()) {
+			this.#runner.sendWebsocketMessageAck(requestId, index);
+		}
+
+		this.#wsAckQueue.clear();
+	}
+
 	getContext(actorId: string): DriverContext {
 		return {};
 	}
@@ -296,7 +465,14 @@ export class EngineActorDriver implements ActorDriver {
 
 		const handler = this.#actors.get(actorId);
 		if (handler?.actor) {
-			await handler.actor._stop();
+			try {
+				await handler.actor._onStop();
+			} catch (err) {
+				logger().error({
+					msg: "error in _onStop, proceeding with removing actor",
+					err: stringifyError(err),
+				});
+			}
 			this.#actors.delete(actorId);
 		}
 
@@ -304,8 +480,9 @@ export class EngineActorDriver implements ActorDriver {
 	}
 
 	async #runnerFetch(
-		runner: Runner,
+		_runner: Runner,
 		actorId: string,
+		_requestIdBuf: ArrayBuffer,
 		request: Request,
 	): Promise<Response> {
 		logger().debug({
@@ -318,12 +495,14 @@ export class EngineActorDriver implements ActorDriver {
 	}
 
 	async #runnerWebSocket(
-		runner: Runner,
+		_runner: Runner,
 		actorId: string,
 		websocketRaw: any,
+		requestIdBuf: ArrayBuffer,
 		request: Request,
 	): Promise<void> {
 		const websocket = websocketRaw as UniversalWebSocket;
+		const requestId = idToStr(requestIdBuf);
 
 		logger().debug({ msg: "runner websocket", actorId, url: request.url });
 
@@ -367,6 +546,7 @@ export class EngineActorDriver implements ActorDriver {
 				actorId,
 				encoding,
 				connParams,
+				requestId,
 				// Extract connId and connToken from protocols if needed
 				undefined,
 				undefined,
@@ -402,11 +582,37 @@ export class EngineActorDriver implements ActorDriver {
 			});
 		}
 
-		websocket.addEventListener("message", (event) => {
+		websocket.addEventListener("message", (event: RivetMessageEvent) => {
 			wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
+
+			invariant(event.rivetRequestId, "missing rivetRequestId");
+			invariant(event.rivetMessageIndex, "missing rivetMessageIndex");
+
+			// Track only the highest seen message index per request
+			// Convert ArrayBuffer to string for Map key
+			const currentEntry = this.#wsAckQueue.get(requestId);
+			if (currentEntry) {
+				if (event.rivetMessageIndex > currentEntry.messageIndex) {
+					currentEntry.messageIndex = event.rivetMessageIndex;
+				} else {
+					logger().warn({
+						msg: "received lower index than ack queue for message",
+						requestId,
+						queuedMessageIndex: currentEntry,
+						eventMessageIndex: event.rivetMessageIndex,
+					});
+				}
+			} else {
+				this.#wsAckQueue.set(requestId, {
+					requestIdBuf,
+					messageIndex: event.rivetMessageIndex,
+				});
+			}
 		});
 
 		websocket.addEventListener("close", (event) => {
+			// Flush any pending acks before closing
+			this.#flushWsAcks();
 			wsHandlerPromise.then((x) => x.onClose?.(event, wsContext));
 		});
 
@@ -415,12 +621,22 @@ export class EngineActorDriver implements ActorDriver {
 		});
 	}
 
-	async sleep(actorId: string) {
+	startSleep(actorId: string) {
 		this.#runner.sleepActor(actorId);
 	}
 
-	async shutdown(immediate: boolean): Promise<void> {
+	async shutdownRunner(immediate: boolean): Promise<void> {
 		logger().info({ msg: "stopping engine actor driver" });
+
+		// Clear the ack flush interval
+		if (this.#wsAckFlushInterval) {
+			clearInterval(this.#wsAckFlushInterval);
+			this.#wsAckFlushInterval = undefined;
+		}
+
+		// Flush any remaining acks
+		this.#flushWsAcks();
+
 		await this.#runner.shutdown(immediate);
 	}
 
@@ -430,7 +646,11 @@ export class EngineActorDriver implements ActorDriver {
 			stream.onAbort(() => {});
 			c.req.raw.signal.addEventListener("abort", () => {
 				logger().debug("SSE aborted, shutting down runner");
-				this.shutdown(true);
+
+				// We cannot assume that the request will always be closed gracefully by Rivet. We always proceed with a graceful shutdown in case the request was terminated for any other reason.
+				//
+				// If we did not use a graceful shutdown, the runner would
+				this.shutdownRunner(false);
 			});
 
 			await this.#runnerStarted.promise;
@@ -440,7 +660,34 @@ export class EngineActorDriver implements ActorDriver {
 			invariant(payload, "runnerId not set");
 			await stream.writeSSE({ data: payload });
 
+			// Send ping every second to keep the connection alive
+			while (true) {
+				if (this.#isRunnerStopped) {
+					logger().debug({
+						msg: "runner is stopped",
+					});
+					break;
+				}
+
+				if (stream.closed || stream.aborted) {
+					logger().debug({
+						msg: "runner sse stream closed",
+						closed: stream.closed,
+						aborted: stream.aborted,
+					});
+					break;
+				}
+
+				await stream.writeSSE({ event: "ping", data: "" });
+				await stream.sleep(RUNNER_SSE_PING_INTERVAL);
+			}
+
+			// Wait for the runner to stop if the SSE stream aborted early for any reason
 			await this.#runnerStopped.promise;
 		});
 	}
+
+	getExtraActorLogParams(): Record<string, string> {
+		return { runnerId: this.#runner.runnerId ?? "-" };
+	}
 }

package/src/drivers/engine/config.ts

@@ -7,10 +7,8 @@ export const EngingConfigSchema = z
 		/** Unique key for this runner. Runners connecting a given key will replace any other runner connected with the same key. */
 		runnerKey: z
 			.string()
-			.default(
-				() =>
-					getEnvUniversal("RIVET_RUNNER_KEY") ?? crypto.randomUUID(),
-			),
+			.optional()
+			.transform((x) => x ?? getEnvUniversal("RIVET_RUNNER_KEY")),
 
 		/** How many actors this runner can run. */
 		totalSlots: z.number().default(100_000),

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

@@ -79,7 +79,8 @@ export class FileSystemActorDriver implements ActorDriver {
 		return this.#state.createDatabase(actorId);
 	}
 
-	sleep(actorId: string): Promise<void> {
-		return this.#state.sleepActor(actorId);
+	startSleep(actorId: string): void {
+		// Spawns the sleepActor promise
+		this.#state.sleepActor(actorId);
 	}
 }

package/src/drivers/file-system/global-state.ts

@@ -325,7 +325,7 @@ export class FileSystemGlobalState {
 
 		// Stop actor
 		invariant(actor.actor, "actor should be loaded");
-		await actor.actor._stop();
+		await actor.actor._onStop();
 
 		// Remove from map after stop is complete
 		this.#actors.delete(actorId);

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

@@ -1,5 +1,6 @@
 import type { Context as HonoContext } from "hono";
 import invariant from "invariant";
+import { generateConnRequestId } from "@/actor/conn";
 import { type ActorRouter, createActorRouter } from "@/actor/router";
 import {
 	handleRawWebSocketHandler,
@@ -173,6 +174,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				params,
+				generateConnRequestId(),
 				connId,
 				connToken,
 			);
@@ -231,6 +233,7 @@ export class FileSystemManagerDriver implements ManagerDriver {
 				actorId,
 				encoding,
 				connParams,
+				generateConnRequestId(),
 				connId,
 				connToken,
 			);

package/src/engine-process/mod.ts

@@ -8,6 +8,7 @@ import {
 	ensureDirectoryExists,
 	getStoragePath,
 } from "@/drivers/file-system/utils";
+import { EXTRA_ERROR_LOG } from "@/utils";
 import { logger } from "./log";
 
 export const ENGINE_PORT = 6420;
@@ -84,6 +85,25 @@ export async function ensureEngineProcess(
 		stdio: ["inherit", "pipe", "pipe"],
 		env: {
 			...process.env,
+			// In development, runners can be terminated without a graceful
+			// shutdown (i.e. SIGKILL instead of SIGTERM). This is treated as a
+			// crash by Rivet Engine in production and implements a backoff for
+			// rescheduling actors in case of a crash loop.
+			//
+			// This is problematic in development since this will cause actors
+			// to become unresponsive if frequently killing your dev server.
+			//
+			// We reduce the timeouts for resetting a runner as healthy in
+			// order to account for this.
+			RIVET__PEGBOARD__RETRY_RESET_DURATION: "100",
+			RIVET__PEGBOARD__BASE_RETRY_TIMEOUT: "100",
+			// Set max exponent to 1 to have a maximum of base_retry_timeout
+			RIVET__PEGBOARD__RESCHEDULE_BACKOFF_MAX_EXPONENT: "1",
+			// Reduce thresholds for faster development iteration
+			//
+			// Default ping interval is 3s, this gives a 2s & 4s grace
+			RIVET__PEGBOARD__RUNNER_ELIGIBLE_THRESHOLD: "5000",
+			RIVET__PEGBOARD__RUNNER_LOST_THRESHOLD: "7000",
 		},
 	});
 
@@ -110,8 +130,7 @@ export async function ensureEngineProcess(
 			msg: "engine process exited, please report this error",
 			code,
 			signal,
-			issues: "https://github.com/rivet-dev/rivetkit/issues",
-			support: "https://rivet.dev/discord",
+			...EXTRA_ERROR_LOG,
 		});
 		// Clean up log streams
 		stdoutStream.end();
@@ -228,8 +247,7 @@ async function downloadEngineBinaryIfNeeded(
 			msg: "engine download failed, please report this error",
 			tempPath,
 			error,
-			issues: "https://github.com/rivet-dev/rivetkit/issues",
-			support: "https://rivet.dev/discord",
+			...EXTRA_ERROR_LOG,
 		});
 		try {
 			await fs.unlink(tempPath);

package/src/inspector/config.ts

@@ -1,10 +1,6 @@
-import type { cors } from "hono/cors";
 import { z } from "zod";
-import { HEADER_ACTOR_QUERY } from "@/driver-helpers/mod";
 import { getEnvUniversal } from "@/utils";
 
-type CorsOptions = NonNullable<Parameters<typeof cors>[0]>;
-
 const defaultTokenFn = () => {
 	const envToken = getEnvUniversal("RIVETKIT_INSPECTOR_TOKEN");
 
@@ -22,41 +18,6 @@ const defaultEnabled = () => {
 	);
 };
 
-const defaultInspectorOrigins = [
-	"http://localhost:43708",
-	"http://localhost:43709",
-	"https://studio.rivet.gg",
-	"https://inspect.rivet.dev",
-	"https://dashboard.rivet.dev",
-	"https://dashboard.staging.rivet.dev",
-];
-
-const defaultCors: CorsOptions = {
-	origin: (origin) => {
-		if (
-			defaultInspectorOrigins.includes(origin) ||
-			(origin.startsWith("https://") &&
-				origin.endsWith("rivet-dev.vercel.app"))
-		) {
-			return origin;
-		} else {
-			return null;
-		}
-	},
-	allowMethods: ["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
-	allowHeaders: [
-		"Authorization",
-		"Content-Type",
-		"User-Agent",
-		"baggage",
-		"sentry-trace",
-		"x-rivet-actor",
-		"x-rivet-target",
-	],
-	maxAge: 3600,
-	credentials: true,
-};
-
 export const InspectorConfigSchema = z
 	.object({
 		enabled: z
@@ -69,11 +30,6 @@ export const InspectorConfigSchema = z
 			)
 			.optional()
 			.default(defaultEnabled),
-		/** CORS configuration for the router. Uses Hono's CORS middleware options. */
-		cors: z
-			.custom<CorsOptions>()
-			.optional()
-			.default(() => defaultCors),
 
 		/**
 		 * Token used to access the Inspector.
@@ -95,6 +51,5 @@ export const InspectorConfigSchema = z
 	.default(() => ({
 		enabled: defaultEnabled(),
 		token: defaultTokenFn,
-		cors: defaultCors,
 	}));
 export type InspectorConfig = z.infer<typeof InspectorConfigSchema>;