fix(runner): skip closing hibernatable websockets on shutdown (#3397)

Author: NathanFlurry

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

@@ -82,8 +82,11 @@ export class Tunnel {
 		// closed:
 		// https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157
 		for (const [_, ws] of this.#actorWebSockets) {
-			// TODO: Trigger close event, but do not send anything over the tunnel
-			ws.__closeWithoutCallback(1000, "ws.tunnel_shutdown");
+			// Only close non-hibernatable websockets to prevent sending
+			// unnecessary close messages for websockets that will be hibernated
+			if (!ws.canHibernate) {
+				ws.__closeWithoutCallback(1000, "ws.tunnel_shutdown");
+			}
 		}
 		this.#actorWebSockets.clear();
 	}
@@ -668,6 +671,8 @@ export class Tunnel {
 					requestId,
 					request,
 				);
+			adapter.canHibernate = hibernationConfig.enabled;
+
 			this.#sendMessage(requestId, {
 				tag: "ToServerWebSocketOpen",
 				val: {

engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts

@@ -18,6 +18,7 @@ export class WebSocketTunnelAdapter {
 	#url = "";
 	#sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void;
 	#closeCallback: (code?: number, reason?: string, retry?: boolean) => void;
+	#canHibernate: boolean = false;
 
 	// Event buffering for events fired before listeners are attached
 	#bufferedEvents: Array<{
@@ -73,6 +74,16 @@ export class WebSocketTunnelAdapter {
 		return this.#url;
 	}
 
+	/** @experimental */
+	get canHibernate(): boolean {
+		return this.#canHibernate;
+	}
+
+	/** @experimental */
+	set canHibernate(value: boolean) {
+		this.#canHibernate = value;
+	}
+
 	get onopen(): ((this: any, ev: any) => any) | null {
 		return this.#onopen;
 	}

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

@@ -2089,10 +2089,13 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
 			}
 		}
 
-		// Disconnect existing connections
 		const promises: Promise<unknown>[] = [];
+
+		// Disconnect existing non-hibernatable connections
 		for (const connection of this.#connections.values()) {
-			promises.push(connection.disconnect());
+			if (!connection.isHibernatable) {
+				promises.push(connection.disconnect());
+			}
 
 			// TODO: Figure out how to abort HTTP requests on shutdown. This
 			// might already be handled by the engine runner tunnel shutdown.