rivetkit 2.0.24-rc.1 → 2.0.25-rc.1

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

@@ -1,65 +1,74 @@
 import type {
 	ActorConfig as EngineActorConfig,
 	RunnerConfig as EngineRunnerConfig,
-	HibernationConfig,
+	HibernatingWebSocketMetadata,
 } from "@rivetkit/engine-runner";
-import { Runner } from "@rivetkit/engine-runner";
+import { idToStr, Runner } from "@rivetkit/engine-runner";
 import * as cbor from "cbor-x";
 import type { Context as HonoContext } from "hono";
 import { streamSSE } from "hono/streaming";
-import { WSContext } from "hono/ws";
+import { WSContext, type WSContextInit } from "hono/ws";
 import invariant from "invariant";
+import { type AnyConn, CONN_STATE_MANAGER_SYMBOL } from "@/actor/conn/mod";
 import { lookupInRegistry } from "@/actor/definition";
-import { PERSIST_SYMBOL } from "@/actor/instance";
+import { KEYS } from "@/actor/instance/kv";
 import { deserializeActorKey } from "@/actor/keys";
-import { EncodingSchema } from "@/actor/protocol/serde";
+import { getValueLength } from "@/actor/protocol/old";
 import { type ActorRouter, createActorRouter } from "@/actor/router";
 import {
-	handleRawWebSocketHandler,
-	handleWebSocketConnect,
+	parseWebSocketProtocols,
+	routeWebSocket,
 	truncateRawWebSocketPathPrefix,
-} from "@/actor/router-endpoints";
+	type UpgradeWebSocketArgs,
+} from "@/actor/router-websocket-endpoints";
 import type { Client } from "@/client/client";
 import {
-	PATH_CONNECT_WEBSOCKET,
-	PATH_RAW_WEBSOCKET_PREFIX,
-	WS_PROTOCOL_CONN_PARAMS,
-	WS_PROTOCOL_ENCODING,
-	WS_PROTOCOL_TOKEN,
+	PATH_CONNECT,
+	PATH_INSPECTOR_CONNECT,
+	PATH_WEBSOCKET_BASE,
+	PATH_WEBSOCKET_PREFIX,
 } from "@/common/actor-router-consts";
-import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
 import { getLogger } from "@/common/log";
 import type {
-	RivetEvent,
 	RivetMessageEvent,
 	UniversalWebSocket,
 } from "@/common/websocket-interface";
 import {
 	type ActorDriver,
 	type AnyActorInstance,
+	getInitialActorKvState,
 	type ManagerDriver,
-	serializeEmptyPersistData,
 } from "@/driver-helpers/mod";
 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;
 
+// Message ack deadline is 30s on the gateway, but we will ack more frequently
+// in order to minimize the message buffer size on the gateway and to give
+// generous breathing room for the timeout.
+//
+// See engine/packages/pegboard-gateway/src/shared_state.rs
+// (HWS_MESSAGE_ACK_TIMEOUT)
+const CONN_MESSAGE_ACK_DEADLINE = 5_000;
+
+// Force saveState when cumulative message size reaches this threshold (0.5 MB)
+//
+// See engine/packages/pegboard-gateway/src/shared_state.rs
+// (HWS_MAX_PENDING_MSGS_SIZE_PER_REQ)
+const CONN_BUFFERED_MESSAGE_SIZE_THRESHOLD = 500_000;
+
 interface ActorHandler {
 	actor?: AnyActorInstance;
 	actorStartPromise?: ReturnType<typeof promiseWithResolvers<void>>;
-	persistedData?: Uint8Array;
 }
 
 export type DriverContext = {};
@@ -79,12 +88,32 @@ export class EngineActorDriver implements ActorDriver {
 	#runnerStopped: PromiseWithResolvers<undefined> = promiseWithResolvers();
 	#isRunnerStopped: boolean = false;
 
-	// WebSocket message acknowledgment debouncing
-	#wsAckQueue: Map<
+	// HACK: Track actor stop intent locally since the runner protocol doesn't
+	// pass the stop reason to onActorStop. This will be fixed when the runner
+	// protocol is updated to send the intent directly (see RVT-5284)
+	#actorStopIntent: Map<string, "sleep" | "destroy"> = new Map();
+
+	// Map of conn IDs to message index waiting to be persisted before sending
+	// an ack
+	//
+	// serverMessageIndex is updated and pendingAck is flagged in needed in
+	// onBeforePersistConnect, then the HWS ack message is sent in
+	// onAfterPersistConn. This allows us to track what's about to be written
+	// to storage to prevent race conditions with the serverMessageIndex being
+	// updated while writing the existing state.
+	//
+	// bufferedMessageSize tracks the total bytes received since last persist
+	// to force a saveState when threshold is reached. This is the amount of
+	// data currently buffered on the gateway.
+	#hwsMessageIndex = new Map<
 		string,
-		{ requestIdBuf: ArrayBuffer; messageIndex: number }
-	> = new Map();
-	#wsAckFlushInterval?: NodeJS.Timeout;
+		{
+			serverMessageIndex: number;
+			bufferedMessageSize: number;
+			pendingAckFromMessageIndex: boolean;
+			pendingAckFromBufferSize: boolean;
+		}
+	>();
 
 	constructor(
 		registryConfig: RegistryConfig,
@@ -133,159 +162,12 @@ export class EngineActorDriver implements ActorDriver {
 			},
 			fetch: this.#runnerFetch.bind(this),
 			websocket: this.#runnerWebSocket.bind(this),
+			hibernatableWebSocket: {
+				canHibernate: this.#hwsCanHibernate.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 hibernatableArray =
-					handler.actor[PERSIST_SYMBOL].hibernatableWebSocket;
-				logger().debug({
-					msg: "checking hibernatable websockets",
-					requestId: idToStr(requestId),
-					existingHibernatableWebSockets: hibernatableArray.length,
-				});
-				const existingWs = hibernatableArray.find((ws) =>
-					arrayBuffersEqual(ws.requestId, requestId),
-				);
-
-				// Determine configuration for new WS
-				let hibernationConfig: HibernationConfig;
-				if (existingWs) {
-					logger().debug({
-						msg: "found existing hibernatable websocket",
-						requestId: idToStr(requestId),
-						lastMsgIndex: existingWs.msgIndex,
-					});
-					hibernationConfig = {
-						enabled: true,
-						lastMsgIndex: Number(existingWs.msgIndex),
-					};
-				} else {
-					logger().debug({
-						msg: "no existing hibernatable websocket found",
-						requestId: idToStr(requestId),
-					});
-					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 or update hibernatable WebSocket
-				if (existingWs) {
-					logger().debug({
-						msg: "updated existing hibernatable websocket timestamp",
-						requestId: idToStr(requestId),
-					});
-					existingWs.lastSeenTimestamp = BigInt(Date.now());
-				} else {
-					logger().debug({
-						msg: "created new hibernatable websocket entry",
-						requestId: idToStr(requestId),
-					});
-					handler.actor[PERSIST_SYMBOL].hibernatableWebSocket.push({
-						requestId,
-						lastSeenTimestamp: BigInt(Date.now()),
-						msgIndex: -1n,
-					});
-				}
-
-				return hibernationConfig;
-			},
 		};
 
 		// Create and start runner
@@ -297,15 +179,10 @@ 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);
+	getExtraActorLogParams(): Record<string, string> {
+		return { runnerId: this.#runner.runnerId ?? "-" };
 	}
 
 	async #loadActorHandler(actorId: string): Promise<ActorHandler> {
@@ -318,52 +195,10 @@ export class EngineActorDriver implements ActorDriver {
 		return handler;
 	}
 
-	async loadActor(actorId: string): Promise<AnyActorInstance> {
-		const handler = await this.#loadActorHandler(actorId);
-		if (!handler.actor) throw new Error(`Actor ${actorId} failed to load`);
-		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 {};
 	}
 
-	async readPersistedData(actorId: string): Promise<Uint8Array | undefined> {
-		const handler = this.#actors.get(actorId);
-		if (!handler) throw new Error(`Actor ${actorId} not loaded`);
-
-		// This was loaded during actor startup
-		return handler.persistedData;
-	}
-
-	async writePersistedData(actorId: string, data: Uint8Array): Promise<void> {
-		const handler = this.#actors.get(actorId);
-		if (!handler) throw new Error(`Actor ${actorId} not loaded`);
-
-		handler.persistedData = data;
-
-		logger().debug({
-			msg: "writing persisted data for actor",
-			actorId,
-			dataSize: data.byteLength,
-		});
-
-		await this.#runner.kvPut(actorId, [[KEYS.PERSIST_DATA, data]]);
-	}
-
 	async setAlarm(actor: AnyActorInstance, timestamp: number): Promise<void> {
 		// Clear prev timeout
 		if (this.#alarmTimeout) {
@@ -374,7 +209,7 @@ export class EngineActorDriver implements ActorDriver {
 		// Set alarm
 		const delay = Math.max(0, timestamp - Date.now());
 		this.#alarmTimeout = setLongTimeout(() => {
-			actor._onAlarm();
+			actor.onAlarm();
 			this.#alarmTimeout = undefined;
 		}, delay);
 
@@ -385,7 +220,7 @@ export class EngineActorDriver implements ActorDriver {
 		// Instead, it just wakes the actor on the alarm (if not
 		// already awake).
 		//
-		// _onAlarm is automatically called on `ActorInstance.start` when waking
+		// onAlarm is automatically called on `ActorInstance.start` when waking
 		// again.
 		this.#runner.setAlarm(actor.id, timestamp);
 	}
@@ -394,7 +229,172 @@ export class EngineActorDriver implements ActorDriver {
 		return undefined;
 	}
 
-	// Runner lifecycle callbacks
+	// MARK: - Batch KV operations
+	async kvBatchPut(
+		actorId: string,
+		entries: [Uint8Array, Uint8Array][],
+	): Promise<void> {
+		await this.#runner.kvPut(actorId, entries);
+	}
+
+	async kvBatchGet(
+		actorId: string,
+		keys: Uint8Array[],
+	): Promise<(Uint8Array | null)[]> {
+		return await this.#runner.kvGet(actorId, keys);
+	}
+
+	async kvBatchDelete(actorId: string, keys: Uint8Array[]): Promise<void> {
+		await this.#runner.kvDelete(actorId, keys);
+	}
+
+	async kvList(actorId: string): Promise<Uint8Array[]> {
+		const entries = await this.#runner.kvListPrefix(
+			actorId,
+			new Uint8Array(),
+		);
+		const keys = entries.map(([key]) => key);
+		logger().info({
+			msg: "kvList called",
+			actorId,
+			keysCount: keys.length,
+			keys: keys.map((k) => new TextDecoder().decode(k)),
+		});
+		return keys;
+	}
+
+	async kvListPrefix(
+		actorId: string,
+		prefix: Uint8Array,
+	): Promise<[Uint8Array, Uint8Array][]> {
+		const result = await this.#runner.kvListPrefix(actorId, prefix);
+		logger().info({
+			msg: "kvListPrefix called",
+			actorId,
+			prefixStr: new TextDecoder().decode(prefix),
+			entriesCount: result.length,
+			keys: result.map(([key]) => new TextDecoder().decode(key)),
+		});
+		return result;
+	}
+
+	// MARK: - Actor Lifecycle
+	async loadActor(actorId: string): Promise<AnyActorInstance> {
+		const handler = await this.#loadActorHandler(actorId);
+		if (!handler.actor) throw new Error(`Actor ${actorId} failed to load`);
+		return handler.actor;
+	}
+
+	startSleep(actorId: string) {
+		// HACK: Track intent for onActorStop (see RVT-5284)
+		this.#actorStopIntent.set(actorId, "sleep");
+		this.#runner.sleepActor(actorId);
+	}
+
+	startDestroy(actorId: string) {
+		// HACK: Track intent for onActorStop (see RVT-5284)
+		this.#actorStopIntent.set(actorId, "destroy");
+		this.#runner.stopActor(actorId);
+	}
+
+	async shutdownRunner(immediate: boolean): Promise<void> {
+		logger().info({ msg: "stopping engine actor driver", immediate });
+
+		// TODO: We need to update the runner to have a draining state so:
+		// 1. Send ToServerDraining
+		//		- This causes Pegboard to stop allocating actors to this runner
+		// 2. Pegboard sends ToClientStopActor for all actors on this runner which handles the graceful migration of each actor independently
+		// 3. Send ToServerStopping once all actors have successfully stopped
+		//
+		// What's happening right now is:
+		// 1. All actors enter stopped state
+		// 2. Actors still respond to requests because only RivetKit knows it's
+		//    stopping, this causes all requests to issue errors that the actor is
+		//    stopping. (This will NOT return a 503 bc the runner has no idea the
+		//    actors are stopping.)
+		// 3. Once the last actor stops, then the runner finally stops + actors
+		//    reschedule
+		//
+		// This means that:
+		// - All actors on this runner are bricked until the slowest onStop finishes
+		// - Guard will not gracefully handle requests bc it's not receiving a 503
+		// - Actors can still be scheduled to this runner while the other
+		//   actors are stopping, meaning that those actors will NOT get onStop
+		//   and will potentiall corrupt their state
+		//
+		// HACK: Stop all actors to allow state to be saved
+		// NOTE: onStop is only supposed to be called by the runner, we're
+		// abusing it here
+		logger().debug({
+			msg: "stopping all actors before shutdown",
+			actorCount: this.#actors.size,
+		});
+		const stopPromises: Promise<void>[] = [];
+		for (const [_actorId, handler] of this.#actors.entries()) {
+			if (handler.actor) {
+				stopPromises.push(
+					handler.actor.onStop("sleep").catch((err) => {
+						handler.actor?.rLog.error({
+							msg: "onStop errored",
+							error: stringifyError(err),
+						});
+					}),
+				);
+			}
+		}
+		await Promise.all(stopPromises);
+		logger().debug({ msg: "all actors stopped" });
+
+		await this.#runner.shutdown(immediate);
+	}
+
+	async serverlessHandleStart(c: HonoContext): Promise<Response> {
+		return streamSSE(c, async (stream) => {
+			// NOTE: onAbort does not work reliably
+			stream.onAbort(() => {});
+			c.req.raw.signal.addEventListener("abort", () => {
+				logger().debug("SSE aborted, shutting down runner");
+
+				// 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;
+
+			// Runner id should be set if the runner started
+			const payload = this.#runner.getServerlessInitPacket();
+			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;
+		});
+	}
+
 	async #runnerOnActorStart(
 		actorId: string,
 		generation: number,
@@ -422,32 +422,33 @@ export class EngineActorDriver implements ActorDriver {
 			// create the same handler simultaneously.
 			handler = {
 				actorStartPromise: promiseWithResolvers(),
-				persistedData: undefined,
 			};
 			this.#actors.set(actorId, handler);
+		}
 
-			// Load persisted data from storage
-			const [persistedValue] = await this.#runner.kvGet(actorId, [
-				KEYS.PERSIST_DATA,
-			]);
-
-			handler.persistedData =
-				persistedValue !== null
-					? persistedValue
-					: serializeEmptyPersistData(input);
+		const name = actorConfig.name as string;
+		invariant(actorConfig.key, "actor should have a key");
+		const key = deserializeActorKey(actorConfig.key);
 
+		// Initialize storage
+		const [persistDataBuffer] = await this.#runner.kvGet(actorId, [
+			KEYS.PERSIST_DATA,
+		]);
+		if (persistDataBuffer === null) {
+			const initialKvState = getInitialActorKvState(input);
+			await this.#runner.kvPut(actorId, initialKvState);
+			logger().debug({
+				msg: "initialized persist data for new actor",
+				actorId,
+			});
+		} else {
 			logger().debug({
-				msg: "loaded persisted data for actor",
+				msg: "found existing persist data for actor",
 				actorId,
-				dataSize: handler.persistedData?.byteLength,
-				wasInStorage: persistedValue !== null,
+				dataSize: persistDataBuffer.byteLength,
 			});
 		}
 
-		const name = actorConfig.name as string;
-		invariant(actorConfig.key, "actor should have a key");
-		const key = deserializeActorKey(actorConfig.key);
-
 		// Create actor instance
 		const definition = lookupInRegistry(
 			this.#registryConfig,
@@ -465,10 +466,6 @@ export class EngineActorDriver implements ActorDriver {
 			"unknown", // TODO: Add regions
 		);
 
-		// Resolve promise if waiting
-		handler.actorStartPromise?.resolve();
-		handler.actorStartPromise = undefined;
-
 		logger().debug({ msg: "runner actor started", actorId, name, key });
 	}
 
@@ -478,25 +475,37 @@ export class EngineActorDriver implements ActorDriver {
 	): Promise<void> {
 		logger().debug({ msg: "runner actor stopping", actorId, generation });
 
+		// HACK: Retrieve the stop intent we tracked locally (see RVT-5284)
+		// Default to "sleep" if no intent was recorded (e.g., if the runner
+		// initiated the stop)
+		//
+		// TODO: This will not work if the actor is destroyed from the API
+		// correctly. Currently, it will use the sleep intent, but it's
+		// actually a destroy intent.
+		const reason = this.#actorStopIntent.get(actorId) ?? "sleep";
+		this.#actorStopIntent.delete(actorId);
+
 		const handler = this.#actors.get(actorId);
 		if (handler?.actor) {
 			try {
-				await handler.actor._onStop();
+				await handler.actor.onStop(reason);
 			} catch (err) {
 				logger().error({
-					msg: "error in _onStop, proceeding with removing actor",
+					msg: "error in onStop, proceeding with removing actor",
 					err: stringifyError(err),
 				});
 			}
 			this.#actors.delete(actorId);
 		}
 
-		logger().debug({ msg: "runner actor stopped", actorId });
+		logger().debug({ msg: "runner actor stopped", actorId, reason });
 	}
 
+	// MARK: - Runner Networking
 	async #runnerFetch(
 		_runner: Runner,
 		actorId: string,
+		_gatewayIdBuf: ArrayBuffer,
 		_requestIdBuf: ArrayBuffer,
 		request: Request,
 	): Promise<Response> {
@@ -513,243 +522,424 @@ export class EngineActorDriver implements ActorDriver {
 		_runner: Runner,
 		actorId: string,
 		websocketRaw: any,
+		gatewayIdBuf: ArrayBuffer,
 		requestIdBuf: ArrayBuffer,
 		request: Request,
+		requestPath: string,
+		requestHeaders: Record<string, string>,
+		isHibernatable: boolean,
+		isRestoringHibernatable: boolean,
 	): Promise<void> {
 		const websocket = websocketRaw as UniversalWebSocket;
-		const requestId = idToStr(requestIdBuf);
 
-		logger().debug({ msg: "runner websocket", actorId, url: request.url });
+		// Add a unique ID to track this WebSocket object
+		const wsUniqueId = `ws_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
+		(websocket as any).__rivet_ws_id = wsUniqueId;
 
-		const url = new URL(request.url);
+		logger().debug({
+			msg: "runner websocket",
+			actorId,
+			url: request.url,
+			isRestoringHibernatable,
+			websocketObjectId: websocketRaw
+				? Object.prototype.toString.call(websocketRaw)
+				: "null",
+			websocketType: websocketRaw?.constructor?.name,
+			wsUniqueId,
+			websocketProps: websocketRaw
+				? Object.keys(websocketRaw).join(", ")
+				: "null",
+		});
 
 		// Parse configuration from Sec-WebSocket-Protocol header (optional for path-based routing)
 		const protocols = request.headers.get("sec-websocket-protocol");
-
-		let encodingRaw: string | undefined;
-		let connParamsRaw: string | undefined;
-
-		if (protocols) {
-			const protocolList = protocols.split(",").map((p) => p.trim());
-			for (const protocol of protocolList) {
-				if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
-					encodingRaw = protocol.substring(
-						WS_PROTOCOL_ENCODING.length,
-					);
-				} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
-					connParamsRaw = decodeURIComponent(
-						protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
-					);
-				}
-			}
-		}
-
-		const encoding = EncodingSchema.parse(encodingRaw);
-		const connParams = connParamsRaw
-			? JSON.parse(connParamsRaw)
-			: undefined;
+		const { encoding, connParams } = parseWebSocketProtocols(protocols);
 
 		// Fetch WS handler
 		//
 		// We store the promise since we need to add WebSocket event listeners immediately that will wait for the promise to resolve
-		let wsHandlerPromise: Promise<UpgradeWebSocketArgs>;
-		if (url.pathname === PATH_CONNECT_WEBSOCKET) {
-			wsHandlerPromise = handleWebSocketConnect(
+		let wsHandler: UpgradeWebSocketArgs;
+		try {
+			wsHandler = await routeWebSocket(
 				request,
+				requestPath,
+				requestHeaders,
 				this.#runConfig,
 				this,
 				actorId,
 				encoding,
 				connParams,
-				requestId,
-				requestIdBuf,
-				// Extract connId and connToken from protocols if needed
-				undefined,
-				undefined,
-			);
-		} else if (url.pathname.startsWith(PATH_RAW_WEBSOCKET_PREFIX)) {
-			wsHandlerPromise = handleRawWebSocketHandler(
-				request,
-				url.pathname + url.search,
-				this,
-				actorId,
+				gatewayIdBuf,
 				requestIdBuf,
+				isHibernatable,
+				isRestoringHibernatable,
 			);
-		} else {
-			throw new Error(`Unreachable path: ${url.pathname}`);
+		} catch (err) {
+			logger().error({ msg: "building websocket handlers errored", err });
+			websocketRaw.close(1011, "ws.route_error");
+			return;
 		}
 
-		// TODO: Add close
-
 		// Connect the Hono WS hook to the adapter
+		//
+		// We need to assign to `raw` in order for WSContext to expose it on
+		// `ws.raw`
+		(websocket as WSContextInit).raw = websocket;
 		const wsContext = new WSContext(websocket);
 
-		wsHandlerPromise.catch((err) => {
-			logger().error({ msg: "building websocket handlers errored", err });
-			wsContext.close(1011, `${err}`);
+		// Get connection and actor from wsHandler (may be undefined for inspector endpoint)
+		const conn = wsHandler.conn;
+		const actor = wsHandler.actor;
+		const connStateManager = conn?.[CONN_STATE_MANAGER_SYMBOL];
+
+		// Bind event listeners to Hono WebSocket handlers
+		//
+		// We update the HWS data after calling handlers in order to ensure
+		// that the handler ran successfully. By doing this, we ensure at least
+		// once delivery of events to the event handlers.
+
+		// Log when attaching event listeners
+		logger().debug({
+			msg: "attaching websocket event listeners",
+			actorId,
+			connId: conn?.id,
+			wsUniqueId: (websocket as any).__rivet_ws_id,
+			isRestoringHibernatable,
+			websocketType: websocket?.constructor?.name,
 		});
 
-		if (websocket.readyState === 1) {
-			wsHandlerPromise.then((x) =>
-				x.onOpen?.(new Event("open"), wsContext),
-			);
-		} else {
-			websocket.addEventListener("open", (event) => {
-				wsHandlerPromise.then((x) => x.onOpen?.(event, wsContext));
-			});
+		if (isRestoringHibernatable) {
+			wsHandler.onRestore?.(wsContext);
 		}
 
+		websocket.addEventListener("open", (event) => {
+			wsHandler.onOpen(event, wsContext);
+		});
+
 		websocket.addEventListener("message", (event: RivetMessageEvent) => {
-			wsHandlerPromise.then((x) => x.onMessage?.(event, wsContext));
+			logger().debug({
+				msg: "websocket message event listener triggered",
+				connId: conn?.id,
+				actorId: actor?.id,
+				messageIndex: event.rivetMessageIndex,
+				hasWsHandler: !!wsHandler,
+				hasOnMessage: !!wsHandler?.onMessage,
+				actorIsStopping: actor?.isStopping,
+				websocketType: websocket?.constructor?.name,
+				wsUniqueId: (websocket as any).__rivet_ws_id,
+				eventTargetWsId: (event.target as any)?.__rivet_ws_id,
+			});
 
-			invariant(event.rivetRequestId, "missing rivetRequestId");
-			invariant(event.rivetMessageIndex, "missing rivetMessageIndex");
+			// Check if actor is stopping - if so, don't process new messages.
+			// These messages will be reprocessed when the actor wakes up from hibernation.
+			// TODO: This will never retransmit the socket and the socket will close
+			if (actor?.isStopping) {
+				logger().debug({
+					msg: "ignoring ws message, actor is stopping",
+					connId: conn?.id,
+					actorId: actor?.id,
+					messageIndex: event.rivetMessageIndex,
+				});
+				return;
+			}
 
-			// 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;
+			// Process message
+			logger().debug({
+				msg: "calling wsHandler.onMessage",
+				connId: conn?.id,
+				messageIndex: event.rivetMessageIndex,
+			});
+			wsHandler.onMessage(event, wsContext);
+
+			// Persist message index for hibernatable connections
+			const hibernate = connStateManager?.hibernatableData;
+
+			if (hibernate && conn && actor) {
+				invariant(
+					typeof event.rivetMessageIndex === "number",
+					"missing event.rivetMessageIndex",
+				);
+
+				// Persist message index
+				const previousMsgIndex = hibernate.serverMessageIndex;
+				hibernate.serverMessageIndex = event.rivetMessageIndex;
+				logger().info({
+					msg: "persisting message index",
+					connId: conn.id,
+					previousMsgIndex,
+					newMsgIndex: event.rivetMessageIndex,
+				});
+
+				// Calculate message size and track cumulative size
+				const entry = this.#hwsMessageIndex.get(conn.id);
+				if (entry) {
+					// Track message length
+					const messageLength = getValueLength(event.data);
+					entry.bufferedMessageSize += messageLength;
+
+					if (
+						entry.bufferedMessageSize >=
+						CONN_BUFFERED_MESSAGE_SIZE_THRESHOLD
+					) {
+						// Reset buffered message size immeidatley (instead
+						// of waiting for onAfterPersistConn) since we may
+						// receive more messages before onAfterPersistConn
+						// is called, which would called saveState
+						// immediate multiple times
+						entry.bufferedMessageSize = 0;
+						entry.pendingAckFromBufferSize = true;
+
+						// Save state immediately if approaching buffer threshold
+						actor.stateManager.saveState({
+							immediate: true,
+						});
+					} else {
+						// Save message index. The maxWait is set to the ack deadline
+						// since we ack the message immediately after persisting the index.
+						// If cumulative size exceeds threshold, force immediate persist.
+						//
+						// This will call EngineActorDriver.onAfterPersistConn after
+						// persist to send the ack to the gateway.
+						actor.stateManager.saveState({
+							maxWait: CONN_MESSAGE_ACK_DEADLINE,
+						});
+					}
 				} else {
-					logger().warn({
-						msg: "received lower index than ack queue for message",
-						requestId,
-						queuedMessageIndex: currentEntry,
-						eventMessageIndex: event.rivetMessageIndex,
+					// Fallback if entry missing
+					actor.stateManager.saveState({
+						maxWait: CONN_MESSAGE_ACK_DEADLINE,
 					});
 				}
-			} 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));
+			wsHandler.onClose(event, wsContext);
+
+			// NOTE: Persisted connection is removed when `conn.disconnect`
+			// is called by the WebSocket route
 		});
 
 		websocket.addEventListener("error", (event) => {
-			wsHandlerPromise.then((x) => x.onError?.(event, wsContext));
+			wsHandler.onError(event, wsContext);
 		});
-	}
 
-	startSleep(actorId: string) {
-		this.#runner.sleepActor(actorId);
+		// Log event listener attachment for restored connections
+		if (isRestoringHibernatable) {
+			logger().info({
+				msg: "event listeners attached to restored websocket",
+				actorId,
+				connId: conn?.id,
+				gatewayId: idToStr(gatewayIdBuf),
+				requestId: idToStr(requestIdBuf),
+				websocketType: websocket?.constructor?.name,
+				hasMessageListener: !!websocket.addEventListener,
+			});
+		}
 	}
 
-	async shutdownRunner(immediate: boolean): Promise<void> {
-		logger().info({ msg: "stopping engine actor driver", immediate });
+	// MARK: - Hibernating WebSockets
+	#hwsCanHibernate(
+		actorId: string,
+		gatewayId: ArrayBuffer,
+		requestId: ArrayBuffer,
+		request: Request,
+	): boolean {
+		const url = new URL(request.url);
+		const path = url.pathname;
 
-		// TODO: We need to update the runner to have a draining state so:
-		// 1. Send ToServerDraining
-		//		- This causes Pegboard to stop allocating actors to this runner
-		// 2. Pegboard sends ToClientStopActor for all actors on this runner which handles the graceful migration of each actor independently
-		// 3. Send ToServerStopping once all actors have successfully stopped
-		//
-		// What's happening right now is:
-		// 1. All actors enter stopped state
-		// 2. Actors still respond to requests because only RivetKit knows it's
-		//    stopping, this causes all requests to issue errors that the actor is
-		//    stopping. (This will NOT return a 503 bc the runner has no idea the
-		//    actors are stopping.)
-		// 3. Once the last actor stops, then the runner finally stops + actors
-		//    reschedule
-		//
-		// This means that:
-		// - All actors on this runner are bricked until the slowest _onStop finishes
-		// - Guard will not gracefully handle requests bc it's not receiving a 503
-		// - Actors can still be scheduled to this runner while the other
-		//   actors are stopping, meaning that those actors will NOT get _onStop
-		//   and will potentiall corrupt their state
-		//
-		// HACK: Stop all actors to allow state to be saved
-		// NOTE: _onStop is only supposed to be called by the runner, we're
-		// abusing it here
+		// Get actor instance from runner to access actor name
+		const actorInstance = this.#runner.getActor(actorId);
+		if (!actorInstance) {
+			logger().warn({
+				msg: "actor not found in #hwsCanHibernate",
+				actorId,
+			});
+			return false;
+		}
+
+		// Load actor handler to access persisted data
+		const handler = this.#actors.get(actorId);
+		if (!handler) {
+			logger().warn({
+				msg: "actor handler not found in #hwsCanHibernate",
+				actorId,
+			});
+			return false;
+		}
+		if (!handler.actor) {
+			logger().warn({
+				msg: "actor not found in #hwsCanHibernate",
+				actorId,
+			});
+			return false;
+		}
+
+		// Determine configuration for new WS
 		logger().debug({
-			msg: "stopping all actors before shutdown",
-			actorCount: this.#actors.size,
+			msg: "no existing hibernatable websocket found",
+			gatewayId: idToStr(gatewayId),
+			requestId: idToStr(requestId),
 		});
-		const stopPromises: Promise<void>[] = [];
-		for (const [_actorId, handler] of this.#actors.entries()) {
-			if (handler.actor) {
-				stopPromises.push(
-					handler.actor._onStop().catch((err) => {
-						handler.actor?.rLog.error({
-							msg: "_onStop errored",
-							error: stringifyError(err),
-						});
-					}),
-				);
+		if (path === PATH_CONNECT) {
+			return true;
+		} else if (
+			path === PATH_WEBSOCKET_BASE ||
+			path.startsWith(PATH_WEBSOCKET_PREFIX)
+		) {
+			// Find actor config
+			const definition = lookupInRegistry(
+				this.#registryConfig,
+				actorInstance.config.name,
+			);
+
+			// Check if can hibernate
+			const canHibernateWebSocket =
+				definition.config.options?.canHibernateWebSocket;
+			if (canHibernateWebSocket === true) {
+				return true;
+			} else if (typeof canHibernateWebSocket === "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 =
+						canHibernateWebSocket(truncatedRequest);
+					return canHibernate;
+				} catch (error) {
+					logger().error({
+						msg: "error calling canHibernateWebSocket",
+						error,
+					});
+					return false;
+				}
+			} else {
+				return false;
 			}
+		} else if (path === PATH_INSPECTOR_CONNECT) {
+			return false;
+		} else {
+			logger().warn({
+				msg: "unexpected path for getActorHibernationConfig",
+				path,
+			});
+			return false;
 		}
-		await Promise.all(stopPromises);
-		logger().debug({ msg: "all actors stopped" });
+	}
 
-		// Clear the ack flush interval
-		if (this.#wsAckFlushInterval) {
-			clearInterval(this.#wsAckFlushInterval);
-			this.#wsAckFlushInterval = undefined;
-		}
+	async #hwsLoadAll(
+		actorId: string,
+	): Promise<HibernatingWebSocketMetadata[]> {
+		const actor = await this.loadActor(actorId);
+		return actor.conns
+			.values()
+			.map((conn) => {
+				const connStateManager = conn[CONN_STATE_MANAGER_SYMBOL];
+				const hibernatable = connStateManager.hibernatableData;
+				if (!hibernatable) return undefined;
+				return {
+					gatewayId: hibernatable.gatewayId,
+					requestId: hibernatable.requestId,
+					serverMessageIndex: hibernatable.serverMessageIndex,
+					clientMessageIndex: hibernatable.clientMessageIndex,
+					path: hibernatable.requestPath,
+					headers: hibernatable.requestHeaders,
+				} satisfies HibernatingWebSocketMetadata;
+			})
+			.filter((x) => x !== undefined)
+			.toArray();
+	}
 
-		// Flush any remaining acks
-		this.#flushWsAcks();
+	async onBeforeActorStart(actor: AnyActorInstance): Promise<void> {
+		// Resolve promise if waiting
+		const handler = this.#actors.get(actor.id);
+		invariant(handler, "missing actor handler in onBeforeActorReady");
+		handler.actorStartPromise?.resolve();
+		handler.actorStartPromise = undefined;
 
-		await this.#runner.shutdown(immediate);
+		// Restore hibernating requests
+		const metaEntries = await this.#hwsLoadAll(actor.id);
+		await this.#runner.restoreHibernatingRequests(actor.id, metaEntries);
 	}
 
-	async serverlessHandleStart(c: HonoContext): Promise<Response> {
-		return streamSSE(c, async (stream) => {
-			// NOTE: onAbort does not work reliably
-			stream.onAbort(() => {});
-			c.req.raw.signal.addEventListener("abort", () => {
-				logger().debug("SSE aborted, shutting down runner");
+	onCreateConn(conn: AnyConn) {
+		const hibernatable = conn[CONN_STATE_MANAGER_SYMBOL].hibernatableData;
+		if (!hibernatable) return;
 
-				// 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);
-			});
+		this.#hwsMessageIndex.set(conn.id, {
+			serverMessageIndex: hibernatable.serverMessageIndex,
+			bufferedMessageSize: 0,
+			pendingAckFromMessageIndex: false,
+			pendingAckFromBufferSize: false,
+		});
 
-			await this.#runnerStarted.promise;
+		logger().debug({
+			msg: "created #hwsMessageIndex entry",
+			connId: conn.id,
+			serverMessageIndex: hibernatable.serverMessageIndex,
+		});
+	}
 
-			// Runner id should be set if the runner started
-			const payload = this.#runner.getServerlessInitPacket();
-			invariant(payload, "runnerId not set");
-			await stream.writeSSE({ data: payload });
+	onDestroyConn(conn: AnyConn) {
+		this.#hwsMessageIndex.delete(conn.id);
 
-			// Send ping every second to keep the connection alive
-			while (true) {
-				if (this.#isRunnerStopped) {
-					logger().debug({
-						msg: "runner is stopped",
-					});
-					break;
-				}
+		logger().debug({
+			msg: "removed #hwsMessageIndex entry",
+			connId: conn.id,
+		});
+	}
 
-				if (stream.closed || stream.aborted) {
-					logger().debug({
-						msg: "runner sse stream closed",
-						closed: stream.closed,
-						aborted: stream.aborted,
-					});
-					break;
-				}
+	onBeforePersistConn(conn: AnyConn) {
+		const stateManager = conn[CONN_STATE_MANAGER_SYMBOL];
+		const hibernatable = stateManager.hibernatableDataOrError();
 
-				await stream.writeSSE({ event: "ping", data: "" });
-				await stream.sleep(RUNNER_SSE_PING_INTERVAL);
-			}
+		const entry = this.#hwsMessageIndex.get(conn.id);
+		if (!entry) {
+			logger().warn({
+				msg: "missing EngineActorDriver.#hwsMessageIndex entry for conn",
+				connId: conn.id,
+			});
+			return;
+		}
 
-			// Wait for the runner to stop if the SSE stream aborted early for any reason
-			await this.#runnerStopped.promise;
-		});
+		// There is a newer message index
+		entry.pendingAckFromMessageIndex =
+			hibernatable.serverMessageIndex > entry.serverMessageIndex;
+		entry.serverMessageIndex = hibernatable.serverMessageIndex;
 	}
 
-	getExtraActorLogParams(): Record<string, string> {
-		return { runnerId: this.#runner.runnerId ?? "-" };
+	onAfterPersistConn(conn: AnyConn) {
+		const stateManager = conn[CONN_STATE_MANAGER_SYMBOL];
+		const hibernatable = stateManager.hibernatableDataOrError();
+
+		const entry = this.#hwsMessageIndex.get(conn.id);
+		if (!entry) {
+			logger().warn({
+				msg: "missing EngineActorDriver.#hwsMessageIndex entry for conn",
+				connId: conn.id,
+			});
+			return;
+		}
+
+		// Ack entry
+		if (
+			entry.pendingAckFromMessageIndex ||
+			entry.pendingAckFromBufferSize
+		) {
+			this.#runner.sendHibernatableWebSocketMessageAck(
+				hibernatable.gatewayId,
+				hibernatable.requestId,
+				entry.serverMessageIndex,
+			);
+			entry.pendingAckFromMessageIndex = false;
+			entry.pendingAckFromBufferSize = false;
+			entry.bufferedMessageSize = 0;
+		}
 	}
 }