chore(rivetkit): implement new hibernating ws protocol

engine/sdks/typescript/runner/src/mod.ts

@@ -10,7 +10,7 @@
 	unreachable,
 } from "./utils";
 import { importWebSocket } from "./websocket.js";
 import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
 
 const KV_EXPIRE: number = 30_000;
 const PROTOCOL_VERSION: number = 3;
@@ -57,30 +57,40 @@
 		requestId: protocol.RequestId,
 		request: Request,
 	) => Promise<Response>;
-	websocket?: (
+	websocket: (
 		runner: Runner,
 		actorId: string,
 		ws: any,
 		requestId: protocol.RequestId,
 		request: Request,
-	) => Promise<void>;
+	) => void;
 	onActorStart: (
 		actorId: string,
 		generation: number,
 		config: ActorConfig,
 	) => Promise<void>;
+	/**
+	 * Called on actor start.
+	 *
+	 * Returns the message indices for hibernating request IDs.
+	 *
+	 * Disconnect any connections that are not in the request ID list.
+	 */
+	restoreHibernativeRequests(
+		actorId: string,
+		requestId: protocol.RequestId[],
+	): number[];
 	onActorStop: (actorId: string, generation: number) => Promise<void>;
-	getActorHibernationConfig: (
+	canWebSocketHibernate: (
 		actorId: string,
 		requestId: ArrayBuffer,
 		request: Request,
-	) => HibernationConfig;
+	) => boolean;
 	noAutoShutdown?: boolean;
 }
 
 export interface HibernationConfig {
-	enabled: boolean;
-	lastMsgIndex: number | undefined;
+	lastMsgIndex: number;
 }
 
 export interface KvListOptions {
@@ -105,7 +115,6 @@
 	}
 
 	#actors: Map<string, ActorInstance> = new Map();
-	#actorWebSockets: Map<string, Set<WebSocketTunnelAdapter>> = new Map();
 
 	// WebSocket
 	#pegboardWebSocket?: WebSocket;
@@ -831,6 +840,11 @@
 			webSockets: new Set(),
 		};
 
+		this.#tunnel.restoreHibernatingRequests(
+			actorId,
+			startCommand.hibernatingRequestIds,
+		);
+
 		this.#actors.set(actorId, instance);
 
 		this.#sendActorStateUpdate(actorId, generation, "running");
@@ -1427,8 +1441,10 @@
 		}
 	}
 
-	sendWebsocketMessageAck(requestId: ArrayBuffer, index: number) {
-		this.#tunnel?.__ackWebsocketMessage(requestId, index);
+	sendHibernatableWebSocketMessageAck(requestId: ArrayBuffer, index: number) {
+		if (!this.#tunnel)
+			throw new Error("missing tunnel to send message ack");
+		this.#tunnel.sendHibernatableWebSocketMessageAck(requestId, index);
 	}
 
 	getServerlessInitPacket(): string | undefined {