fix(rivetkit): properly handle msgIndex for hibernatable websocket reconnection

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit/schemas/actor-persist/v2.bare

@@ -51,5 +51,5 @@ type PersistedActor struct {
 	state: data
 	connections: list<PersistedConnection>
 	scheduledEvents: list<PersistedScheduleEvent>
-	hibernatableWebSocket: list<PersistedHibernatableWebSocket>
+	hibernatableWebSockets: list<PersistedHibernatableWebSocket>
 }

rivetkit-typescript/packages/rivetkit/src/actor/conn.ts

@@ -137,7 +137,7 @@ export class Conn<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			return false;
 		}
 		return (
-			this.#actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex((x) =>
+			this.#actor[PERSIST_SYMBOL].hibernatableWebSockets.findIndex((x) =>
 				arrayBuffersEqual(
 					x.requestId,
 					this.__persist.hibernatableRequestId!,

rivetkit-typescript/packages/rivetkit/src/actor/instance.ts

@@ -822,7 +822,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 				msg: "actor restoring",
 				connections: persistData.connections.length,
 				hibernatableWebSockets:
-					persistData.hibernatableWebSocket.length,
+					persistData.hibernatableWebSockets.length,
 			});
 
 			// Set initial state
@@ -1222,7 +1222,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 		// Check if this connection is for a hibernatable websocket
 		if (socket.requestIdBuf) {
 			const isHibernatable =
-				this.#persist.hibernatableWebSocket.findIndex((ws) =>
+				this.#persist.hibernatableWebSockets.findIndex((ws) =>
 					arrayBuffersEqual(ws.requestId, socket.requestIdBuf!),
 				) !== -1;
 
@@ -1654,7 +1654,7 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 	/**
 	 * Handles raw WebSocket connections to the actor.
 	 */
-	async handleWebSocket(
+	async handleRawWebSocket(
 		websocket: UniversalWebSocket,
 		opts: { request: Request },
 	): Promise<void> {
@@ -1672,116 +1672,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			this.#activeRawWebSockets.add(websocket);
 			this.#resetSleepTimer();
 
-			// Track hibernatable WebSockets
-			let rivetRequestId: ArrayBuffer | undefined;
-			let persistedHibernatableWebSocket:
-				| PersistedHibernatableWebSocket
-				| undefined;
-
-			const onSocketOpened = (event: any) => {
-				rivetRequestId = event?.rivetRequestId;
-
-				// Find hibernatable WS
-				if (rivetRequestId) {
-					const rivetRequestIdLocal = rivetRequestId;
-					persistedHibernatableWebSocket =
-						this.#persist.hibernatableWebSocket.find((ws) =>
-							arrayBuffersEqual(
-								ws.requestId,
-								rivetRequestIdLocal,
-							),
-						);
-
-					if (persistedHibernatableWebSocket) {
-						persistedHibernatableWebSocket.lastSeenTimestamp =
-							BigInt(Date.now());
-					}
-				}
-
-				this.#rLog.debug({
-					msg: "actor instance onSocketOpened",
-					rivetRequestId,
-					isHibernatable: !!persistedHibernatableWebSocket,
-					hibernationMsgIndex:
-						persistedHibernatableWebSocket?.msgIndex,
-				});
-			};
-
-			const onSocketMessage = (event: any) => {
-				// Update state of hibernatable WS
-				if (persistedHibernatableWebSocket) {
-					persistedHibernatableWebSocket.lastSeenTimestamp = BigInt(
-						Date.now(),
-					);
-					persistedHibernatableWebSocket.msgIndex = BigInt(
-						event.rivetMessageIndex,
-					);
-				}
-
-				this.#rLog.debug({
-					msg: "actor instance onSocketMessage",
-					rivetRequestId,
-					isHibernatable: !!persistedHibernatableWebSocket,
-					hibernationMsgIndex:
-						persistedHibernatableWebSocket?.msgIndex,
-				});
-			};
-
 			const onSocketClosed = (_event: any) => {
-				// Remove hibernatable WS
-				if (rivetRequestId) {
-					const rivetRequestIdLocal = rivetRequestId;
-					const wsIndex =
-						this.#persist.hibernatableWebSocket.findIndex((ws) =>
-							arrayBuffersEqual(
-								ws.requestId,
-								rivetRequestIdLocal,
-							),
-						);
-
-					const removed = this.#persist.hibernatableWebSocket.splice(
-						wsIndex,
-						1,
-					);
-					if (removed.length > 0) {
-						this.#rLog.debug({
-							msg: "removed hibernatable websocket",
-							rivetRequestId,
-							hibernationMsgIndex:
-								persistedHibernatableWebSocket?.msgIndex,
-						});
-					} else {
-						this.#rLog.warn({
-							msg: "could not find hibernatable websocket to remove",
-							rivetRequestId,
-							hibernationMsgIndex:
-								persistedHibernatableWebSocket?.msgIndex,
-						});
-					}
-				}
-
-				this.#rLog.debug({
-					msg: "actor instance onSocketMessage",
-					rivetRequestId,
-					isHibernatable: !!persistedHibernatableWebSocket,
-					hibernatableWebSocketCount:
-						this.#persist.hibernatableWebSocket.length,
-				});
-
-				// Remove listener and socket from tracking
-				try {
-					websocket.removeEventListener("open", onSocketOpened);
-					websocket.removeEventListener("message", onSocketMessage);
-					websocket.removeEventListener("close", onSocketClosed);
-					websocket.removeEventListener("error", onSocketClosed);
-				} catch {}
+				// Remove socket from tracking
 				this.#activeRawWebSockets.delete(websocket);
 				this.#resetSleepTimer();
 			};
 
 			try {
-				websocket.addEventListener("open", onSocketOpened);
-				websocket.addEventListener("message", onSocketMessage);
 				websocket.addEventListener("close", onSocketClosed);
 				websocket.addEventListener("error", onSocketClosed);
 			} catch {}
@@ -2288,11 +2185,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					},
 				},
 			})),
-			hibernatableWebSocket: persist.hibernatableWebSocket.map((ws) => ({
-				requestId: ws.requestId,
-				lastSeenTimestamp: ws.lastSeenTimestamp,
-				msgIndex: ws.msgIndex,
-			})),
+			hibernatableWebSockets: persist.hibernatableWebSockets.map(
+				(ws) => ({
+					requestId: ws.requestId,
+					lastSeenTimestamp: ws.lastSeenTimestamp,
+					msgIndex: ws.msgIndex,
+				}),
+			),
 		};
 	}
 
@@ -2326,11 +2225,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 					},
 				},
 			})),
-			hibernatableWebSocket: bareData.hibernatableWebSocket.map((ws) => ({
-				requestId: ws.requestId,
-				lastSeenTimestamp: ws.lastSeenTimestamp,
-				msgIndex: ws.msgIndex,
-			})),
+			hibernatableWebSockets: bareData.hibernatableWebSockets.map(
+				(ws) => ({
+					requestId: ws.requestId,
+					lastSeenTimestamp: ws.lastSeenTimestamp,
+					msgIndex: ws.msgIndex,
+				}),
+			),
 		};
 	}
 }

rivetkit-typescript/packages/rivetkit/src/actor/persisted.ts

@@ -5,7 +5,7 @@ export interface PersistedActor<S, CP, CS, I> {
 	state: S;
 	connections: PersistedConn<CP, CS>[];
 	scheduledEvents: PersistedScheduleEvent[];
-	hibernatableWebSocket: PersistedHibernatableWebSocket[];
+	hibernatableWebSockets: PersistedHibernatableWebSocket[];
 }
 
 /** Object representing connection that gets persisted to storage. */

rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts

@@ -188,7 +188,7 @@ export async function handleWebSocketConnect(
 					// Check if this is a hibernatable websocket
 					const isHibernatable =
 						!!requestIdBuf &&
-						actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex(
+						actor[PERSIST_SYMBOL].hibernatableWebSockets.findIndex(
 							(ws) =>
 								arrayBuffersEqual(ws.requestId, requestIdBuf),
 						) !== -1;
@@ -616,7 +616,7 @@ export async function handleRawWebSocketHandler(
 			// Extract rivetRequestId provided by engine runner
 			const rivetRequestId = evt?.rivetRequestId;
 			const isHibernatable =
-				actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex((ws) =>
+				actor[PERSIST_SYMBOL].hibernatableWebSockets.findIndex((ws) =>
 					arrayBuffersEqual(ws.requestId, rivetRequestId),
 				) !== -1;
 
@@ -647,7 +647,7 @@ export async function handleRawWebSocketHandler(
 			});
 
 			// Call the actor's onWebSocket handler with the adapted WebSocket
-			actor.handleWebSocket(adapter, {
+			actor.handleRawWebSocket(adapter, {
 				request: newRequest,
 			});
 		},

rivetkit-typescript/packages/rivetkit/src/driver-helpers/utils.ts

@@ -15,7 +15,7 @@ export function serializeEmptyPersistData(
 		state: bufferToArrayBuffer(cbor.encode(undefined)),
 		connections: [],
 		scheduledEvents: [],
-		hibernatableWebSocket: [],
+		hibernatableWebSockets: [],
 	};
 	return PERSISTED_ACTOR_VERSIONED.serializeWithEmbeddedVersion(persistData);
 }