@rivetkit/engine-runner 2.0.24-rc.1 → 2.0.25-rc.1

package/src/mod.ts

@@ -1,39 +1,28 @@
 import * as protocol from "@rivetkit/engine-runner-protocol";
 import type { Logger } from "pino";
 import type WebSocket from "ws";
+import { type ActorConfig, RunnerActor } from "./actor";
 import { logger, setLogger } from "./log.js";
-import { stringifyCommandWrapper, stringifyEvent } from "./stringify";
-import { Tunnel } from "./tunnel";
+import { stringifyToClient, stringifyToServer } from "./stringify";
+import { type HibernatingWebSocketMetadata, Tunnel } from "./tunnel";
 import {
 	calculateBackoff,
 	parseWebSocketCloseReason,
 	unreachable,
 } from "./utils";
 import { importWebSocket } from "./websocket.js";
-import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
+
+export type { HibernatingWebSocketMetadata };
+export { RunnerActor, type ActorConfig };
+export { idToStr } from "./utils";
 
 const KV_EXPIRE: number = 30_000;
-const PROTOCOL_VERSION: number = 2;
+const PROTOCOL_VERSION: number = 3;
 const RUNNER_PING_INTERVAL = 3_000;
 
 /** Warn once the backlog significantly exceeds the server's ack batch size. */
 const EVENT_BACKLOG_WARN_THRESHOLD = 10_000;
-const SIGNAL_HANDLERS: (() => void)[] = [];
-
-export interface ActorInstance {
-	actorId: string;
-	generation: number;
-	config: ActorConfig;
-	requests: Set<string>; // Track active request IDs
-	webSockets: Set<string>; // Track active WebSocket IDs
-}
-
-export interface ActorConfig {
-	name: string;
-	key: string | null;
-	createTs: bigint;
-	input: Uint8Array | null;
-}
+const SIGNAL_HANDLERS: (() => void | Promise<void>)[] = [];
 
 export interface RunnerConfig {
 	logger?: Logger;
@@ -51,38 +40,137 @@ export interface RunnerConfig {
 	onConnected: () => void;
 	onDisconnected: (code: number, reason: string) => void;
 	onShutdown: () => void;
+
+	/** Called when receiving a network request. */
 	fetch: (
 		runner: Runner,
 		actorId: string,
+		gatewayId: protocol.GatewayId,
 		requestId: protocol.RequestId,
 		request: Request,
 	) => Promise<Response>;
-	websocket?: (
+
+	/**
+	 * Called when receiving a WebSocket connection.
+	 *
+	 * All event listeners must be added synchronously inside this function or
+	 * else events may be missed. The open event will fire immediately after
+	 * this function finishes.
+	 *
+	 * Any errors thrown here will disconnect the WebSocket immediately.
+	 *
+	 * While `path` and `headers` are partially redundant to the data in the
+	 * `Request`, they may vary slightly from the actual content of `Request`.
+	 * Prefer to persist the `path` and `headers` properties instead of the
+	 * `Request` itself.
+	 *
+	 * ## Hibernating Web Sockets
+	 *
+	 * ### Implementation Requirements
+	 *
+	 * **Requirement 1: Persist HWS Immediately**
+	 *
+	 * This is responsible for persisting hibernatable WebSockets immediately
+	 * (do not wait for open event). It is not time sensitive to flush the
+	 * connection state. If this fails to persist the HWS, the client's
+	 * WebSocket will be disconnected on next wake in the call to
+	 * `Tunnel::restoreHibernatingRequests` since the connection entry will not
+	 * exist.
+	 *
+	 * **Requirement 2: Persist Message Index On `message`**
+	 *
+	 * In the `message` event listener, this handler must persist the message
+	 * index from the event. The request ID is available at
+	 * `event.rivetRequestId` and message index at `event.rivetMessageIndex`.
+	 *
+	 * The message index should not be flushed immediately. Instead, this
+	 * should:
+	 *
+	 * - Debounce calls to persist the message index
+	 * - After each persist, call
+	 *   `Runner::sendHibernatableWebSocketMessageAck` to acknowledge the
+	 *   message
+	 *
+	 * This mechanism allows us to buffer messages on the gateway so we can
+	 * batch-persist events on our end on a given interval.
+	 *
+	 * If this fails to persist, then the gateway will replay unacked
+	 * messages when the actor starts again.
+	 *
+	 * **Requirement 3: Remove HWS From Storage On `close`**
+	 *
+	 * This handler should add an event listener for `close` to remove the
+	 * connection from storage.
+	 *
+	 * If the connection remove fails to persist, the close event will be
+	 * called again on the next actor start in
+	 * `Tunnel::restoreHibernatingRequests` since there will be no request for
+	 * the given connection.
+	 *
+	 * ### Restoring Connections
+	 *
+	 * The user of this library is responsible for:
+	 * 1. Loading all persisted hibernatable WebSocket metadata for an actor
+	 * 2. Calling `Runner::restoreHibernatingRequests` with this metadata at
+	 *    the end of `onActorStart`
+	 *
+	 * `restoreHibernatingRequests` will restore all connections and attach
+	 * the appropriate event listeners.
+	 *
+	 * ### No Open Event On Restoration
+	 *
+	 * When restoring a HWS, the open event will not be called again. It will
+	 * go straight to the message or close event.
+	 */
+	websocket: (
 		runner: Runner,
 		actorId: string,
 		ws: any,
+		gatewayId: protocol.GatewayId,
 		requestId: protocol.RequestId,
 		request: Request,
+		path: string,
+		headers: Record<string, string>,
+		isHibernatable: boolean,
+		isRestoringHibernatable: boolean,
 	) => Promise<void>;
+
+	hibernatableWebSocket: {
+		/**
+		 * Determines if a WebSocket can continue to live while an actor goes to
+		 * sleep.
+		 */
+		canHibernate: (
+			actorId: string,
+			gatewayId: ArrayBuffer,
+			requestId: ArrayBuffer,
+			request: Request,
+		) => boolean;
+	};
+
+	/**
+	 * Called when an actor starts.
+	 *
+	 * This callback is responsible for:
+	 * 1. Initializing the actor instance
+	 * 2. Loading all persisted hibernatable WebSocket metadata for this actor
+	 * 3. Calling `Runner::restoreHibernatingRequests` with the loaded metadata
+	 *    to restore hibernatable WebSocket connections
+	 *
+	 * The actor should not be marked as "ready" until after
+	 * `restoreHibernatingRequests` completes to ensure all hibernatable
+	 * connections are fully restored before the actor processes new requests.
+	 */
 	onActorStart: (
 		actorId: string,
 		generation: number,
 		config: ActorConfig,
 	) => Promise<void>;
+
 	onActorStop: (actorId: string, generation: number) => Promise<void>;
-	getActorHibernationConfig: (
-		actorId: string,
-		requestId: ArrayBuffer,
-		request: Request,
-	) => HibernationConfig;
 	noAutoShutdown?: boolean;
 }
 
-export interface HibernationConfig {
-	enabled: boolean;
-	lastMsgIndex: number | undefined;
-}
-
 export interface KvListOptions {
 	reverse?: boolean;
 	limit?: number;
@@ -104,17 +192,17 @@ export class Runner {
 		return this.#config;
 	}
 
-	#actors: Map<string, ActorInstance> = new Map();
-	#actorWebSockets: Map<string, Set<WebSocketTunnelAdapter>> = new Map();
+	#actors: Map<string, RunnerActor> = new Map();
 
 	// WebSocket
-	#pegboardWebSocket?: WebSocket;
+	__pegboardWebSocket?: WebSocket;
 	runnerId?: string;
 	#lastCommandIdx: number = -1;
 	#pingLoop?: NodeJS.Timeout;
 	#nextEventIdx: bigint = 0n;
 	#started: boolean = false;
 	#shutdown: boolean = false;
+	#shuttingDown: boolean = false;
 	#reconnectAttempt: number = 0;
 	#reconnectTimeout?: NodeJS.Timeout;
 
@@ -130,7 +218,7 @@ export class Runner {
 	#ackInterval?: NodeJS.Timeout;
 
 	// KV operations
-	#nextRequestId: number = 0;
+	#nextKvRequestId: number = 0;
 	#kvRequests: Map<number, KvRequestEntry> = new Map();
 	#kvCleanupInterval?: NodeJS.Timeout;
 
@@ -173,13 +261,6 @@ export class Runner {
 
 	// MARK: Manage actors
 	sleepActor(actorId: string, generation?: number) {
-		if (this.#shutdown) {
-			this.log?.warn({
-				msg: "runner is shut down, cannot sleep actor",
-			});
-			return;
-		}
-
 		const actor = this.getActor(actorId, generation);
 		if (!actor) return;
 
@@ -201,7 +282,7 @@ export class Runner {
 	}
 
 	async forceStopActor(actorId: string, generation?: number) {
-		const actor = this.#removeActor(actorId, generation);
+		const actor = this.getActor(actorId, generation);
 		if (!actor) return;
 
 		// If onActorStop times out, Pegboard will handle this timeout with ACTOR_STOP_THRESHOLD_DURATION_MS
@@ -218,6 +299,11 @@ export class Runner {
 		// Close requests after onActorStop so you can send messages over the tunnel
 		this.#tunnel?.closeActiveRequests(actor);
 
+		// Remove actor after stopping in order to ensure that we can still
+		// call actions on the runner. Do this before sending stopped update in
+		// order to ensure we don't have duplicate actors.
+		this.#removeActor(actorId, generation);
+
 		this.#sendActorStateUpdate(actorId, actor.generation, "stopped");
 	}
 
@@ -232,17 +318,17 @@ export class Runner {
 		}
 	}
 
-	getActor(actorId: string, generation?: number): ActorInstance | undefined {
+	getActor(actorId: string, generation?: number): RunnerActor | undefined {
 		const actor = this.#actors.get(actorId);
 		if (!actor) {
-			this.log?.error({
+			this.log?.warn({
 				msg: "actor not found",
 				actorId,
 			});
 			return undefined;
 		}
 		if (generation !== undefined && actor.generation !== generation) {
-			this.log?.error({
+			this.log?.warn({
 				msg: "actor generation mismatch",
 				actorId,
 				generation,
@@ -253,6 +339,16 @@ export class Runner {
 		return actor;
 	}
 
+	async getAndWaitForActor(
+		actorId: string,
+		generation?: number,
+	): Promise<RunnerActor | undefined> {
+		const actor = this.getActor(actorId, generation);
+		if (!actor) return;
+		await actor.actorStartPromise.promise;
+		return actor;
+	}
+
 	hasActor(actorId: string, generation?: number): boolean {
 		const actor = this.#actors.get(actorId);
 
@@ -262,11 +358,15 @@ export class Runner {
 		);
 	}
 
+	get actors() {
+		return this.#actors;
+	}
+
 	// IMPORTANT: Make sure to call stopActiveRequests if calling #removeActor
 	#removeActor(
 		actorId: string,
 		generation?: number,
-	): ActorInstance | undefined {
+	): RunnerActor | undefined {
 		const actor = this.#actors.get(actorId);
 		if (!actor) {
 			this.log?.error({
@@ -286,6 +386,12 @@ export class Runner {
 
 		this.#actors.delete(actorId);
 
+		this.log?.info({
+			msg: "removed actor",
+			actorId,
+			actors: this.#actors.size,
+		});
+
 		return actor;
 	}
 
@@ -308,23 +414,25 @@ export class Runner {
 
 		if (!this.#config.noAutoShutdown) {
 			if (!SIGNAL_HANDLERS.length) {
-				process.on("SIGTERM", () => {
+				process.on("SIGTERM", async () => {
 					this.log?.debug("received SIGTERM");
 
 					for (const handler of SIGNAL_HANDLERS) {
-						handler();
+						await handler();
 					}
 
-					process.exit(0);
+					// TODO: Add back
+					// process.exit(0);
 				});
-				process.on("SIGINT", () => {
+				process.on("SIGINT", async () => {
 					this.log?.debug("received SIGINT");
 
 					for (const handler of SIGNAL_HANDLERS) {
-						handler();
+						await handler();
 					}
 
-					process.exit(0);
+					// TODO: Add back
+					// process.exit(0);
 				});
 
 				this.log?.debug({
@@ -332,15 +440,24 @@ export class Runner {
 				});
 			}
 
-			SIGNAL_HANDLERS.push(() => {
+			SIGNAL_HANDLERS.push(async () => {
 				const weak = new WeakRef(this);
-				weak.deref()?.shutdown(false, false);
+				await weak.deref()?.shutdown(false, false);
 			});
 		}
 	}
 
 	// MARK: Shutdown
 	async shutdown(immediate: boolean, exit: boolean = false) {
+		// Prevent concurrent shutdowns
+		if (this.#shuttingDown) {
+			this.log?.debug({
+				msg: "shutdown already in progress, ignoring",
+			});
+			return;
+		}
+		this.#shuttingDown = true;
+
 		this.log?.info({
 			msg: "starting shutdown",
 			immediate,
@@ -387,11 +504,8 @@ export class Runner {
 		this.#kvRequests.clear();
 
 		// Close WebSocket
-		if (
-			this.#pegboardWebSocket &&
-			this.#pegboardWebSocket.readyState === 1
-		) {
-			const pegboardWebSocket = this.#pegboardWebSocket;
+		if (this.__webSocketReady()) {
+			const pegboardWebSocket = this.__pegboardWebSocket;
 			if (immediate) {
 				// Stop immediately
 				pegboardWebSocket.close(1000, "pegboard.runner_shutdown");
@@ -403,22 +517,14 @@ export class Runner {
 						readyState: pegboardWebSocket.readyState,
 					});
 
-					// NOTE: We don't use #sendToServer here because that function checks if the runner is
-					// shut down
-					const encoded = protocol.encodeToServer({
+					// Start stopping
+					//
+					// The runner workflow will send StopActor commands for all
+					// actors
+					this.__sendToServer({
 						tag: "ToServerStopping",
 						val: null,
 					});
-					if (
-						this.#pegboardWebSocket &&
-						this.#pegboardWebSocket.readyState === 1
-					) {
-						this.#pegboardWebSocket.send(encoded);
-					} else {
-						this.log?.error(
-							"WebSocket not available or not open for sending data",
-						);
-					}
 
 					const closePromise = new Promise<void>((resolve) => {
 						if (!pegboardWebSocket)
@@ -434,7 +540,8 @@ export class Runner {
 						});
 					});
 
-					// TODO: Wait for all actors to stop before closing ws
+					// Wait for all actors to stop before closing ws
+					await this.#waitForActorsToStop(pegboardWebSocket);
 
 					this.log?.info({
 						msg: "closing WebSocket",
@@ -459,7 +566,7 @@ export class Runner {
 			// the runner has already shut down
 			this.log?.debug({
 				msg: "no runner WebSocket to shutdown or already closed",
-				readyState: this.#pegboardWebSocket?.readyState,
+				readyState: this.__pegboardWebSocket?.readyState,
 			});
 		}
 
@@ -469,9 +576,96 @@ export class Runner {
 			this.#tunnel = undefined;
 		}
 
+		this.#config.onShutdown();
+
 		if (exit) process.exit(0);
+	}
 
-		this.#config.onShutdown();
+	/**
+	 * Wait for all actors to stop before proceeding with shutdown.
+	 *
+	 * This method polls every 100ms to check if all actors have been stopped.
+	 *
+	 * It will resolve early if:
+	 * - All actors are stopped
+	 * - The WebSocket connection is closed
+	 * - The shutdown timeout is reached (120 seconds)
+	 */
+	async #waitForActorsToStop(ws: WebSocket): Promise<void> {
+		const shutdownTimeout = 120_000; // 120 seconds
+		const shutdownCheckInterval = 100; // Check every 100ms
+		const progressLogInterval = 5_000; // Log progress every 5 seconds
+		const shutdownStartTs = Date.now();
+		let lastProgressLogTs = 0; // Ensure first log happens immediately
+
+		return new Promise<void>((resolve) => {
+			const checkActors = () => {
+				const now = Date.now();
+				const elapsed = now - shutdownStartTs;
+				const wsIsClosed = ws.readyState === 2 || ws.readyState === 3;
+
+				if (this.#actors.size === 0) {
+					this.log?.info({
+						msg: "all actors stopped",
+						elapsed,
+					});
+					return true;
+				} else if (wsIsClosed) {
+					this.log?.warn({
+						msg: "websocket closed before all actors stopped",
+						remainingActors: this.#actors.size,
+						elapsed,
+					});
+					return true;
+				} else if (elapsed >= shutdownTimeout) {
+					this.log?.warn({
+						msg: "shutdown timeout reached, forcing close",
+						remainingActors: this.#actors.size,
+						elapsed,
+					});
+					return true;
+				} else {
+					// Log progress every 5 seconds
+					if (now - lastProgressLogTs >= progressLogInterval) {
+						this.log?.info({
+							msg: "waiting for actors to stop",
+							remainingActors: this.#actors.size,
+							elapsed,
+						});
+						lastProgressLogTs = now;
+					}
+					return false;
+				}
+			};
+
+			// Check immediately first
+			if (checkActors()) {
+				this.log?.debug({
+					msg: "actors check completed immediately",
+				});
+				resolve();
+				return;
+			}
+
+			this.log?.debug({
+				msg: "starting actor wait interval",
+				checkInterval: shutdownCheckInterval,
+			});
+
+			const interval = setInterval(() => {
+				this.log?.debug({
+					msg: "actor wait interval tick",
+					actorCount: this.#actors.size,
+				});
+				if (checkActors()) {
+					this.log?.debug({
+						msg: "actors check completed, clearing interval",
+					});
+					clearInterval(interval);
+					resolve();
+				}
+			}, shutdownCheckInterval);
+		});
 	}
 
 	// MARK: Networking
@@ -498,7 +692,7 @@ export class Runner {
 
 		const WS = await importWebSocket();
 		const ws = new WS(this.pegboardUrl, protocols) as any as WebSocket;
-		this.#pegboardWebSocket = ws;
+		this.__pegboardWebSocket = ws;
 
 		this.log?.info({
 			msg: "connecting",
@@ -564,9 +758,6 @@ export class Runner {
 				val: init,
 			});
 
-			// Process unsent KV requests
-			this.#processUnsentKvRequests();
-
 			// Start ping interval
 			const pingLoop = setInterval(() => {
 				if (ws.readyState === 1) {
@@ -612,6 +803,10 @@ export class Runner {
 
 			// Parse message
 			const message = protocol.decodeToClient(buf);
+			this.log?.debug({
+				msg: "received runner message",
+				data: stringifyToClient(message),
+			});
 
 			// Handle message
 			if (message.tag === "ToClientInit") {
@@ -635,8 +830,10 @@ export class Runner {
 					runnerLostThreshold: this.#runnerLostThreshold,
 				});
 
-				// Resend events that haven't been acknowledged
+				// Resend pending events
+				this.#processUnsentKvRequests();
 				this.#resendUnacknowledgedEvents(init.lastEventIdx);
+				this.#tunnel?.resendBufferedEvents();
 
 				this.#config.onConnected();
 			} else if (message.tag === "ToClientCommands") {
@@ -753,13 +950,11 @@ export class Runner {
 		});
 
 		for (const commandWrapper of commands) {
-			this.log?.info({
-				msg: "received command",
-				command: stringifyCommandWrapper(commandWrapper),
-			});
 			if (commandWrapper.inner.tag === "CommandStartActor") {
+				// Spawn background promise
 				this.#handleCommandStartActor(commandWrapper);
 			} else if (commandWrapper.inner.tag === "CommandStopActor") {
+				// Spawn background promise
 				this.#handleCommandStopActor(commandWrapper);
 			} else {
 				unreachable(commandWrapper.inner);
@@ -808,7 +1003,13 @@ export class Runner {
 		}
 	}
 
-	#handleCommandStartActor(commandWrapper: protocol.CommandWrapper) {
+	async #handleCommandStartActor(commandWrapper: protocol.CommandWrapper) {
+		// IMPORTANT: Make sure no async code runs before inserting #actors and
+		// calling addRequestToActor in order to prevent race conditions with
+		// subsequence commands
+
+		if (!this.#tunnel) throw new Error("missing tunnel on actor start");
+
 		const startCommand = commandWrapper.inner
 			.val as protocol.CommandStartActor;
 
@@ -823,43 +1024,80 @@ export class Runner {
 			input: config.input ? new Uint8Array(config.input) : null,
 		};
 
-		const instance: ActorInstance = {
+		const instance = new RunnerActor(
 			actorId,
 			generation,
-			config: actorConfig,
-			requests: new Set(),
-			webSockets: new Set(),
-		};
+			actorConfig,
+			startCommand.hibernatingRequests,
+		);
+
+		const existingActor = this.#actors.get(actorId);
+		if (existingActor) {
+			this.log?.warn({
+				msg: "replacing existing actor in actors map",
+				actorId,
+				existingGeneration: existingActor.generation,
+				newGeneration: generation,
+				existingPendingRequests: existingActor.pendingRequests.length,
+			});
+		}
 
 		this.#actors.set(actorId, instance);
 
+		// NOTE: We have to populate the requestToActor map BEFORE running any
+		// async code in order for incoming tunnel messages to wait for
+		// instance.actorStartPromise before processing messages
+		// TODO: Where is this GC'd if something fails?
+		for (const hr of startCommand.hibernatingRequests) {
+			this.#tunnel.addRequestToActor(hr.gatewayId, hr.requestId, actorId);
+		}
+
+		this.log?.info({
+			msg: "created actor",
+			actors: this.#actors.size,
+			actorId,
+			name: config.name,
+			key: config.key,
+			generation,
+			hibernatingRequests: startCommand.hibernatingRequests.length,
+		});
+
 		this.#sendActorStateUpdate(actorId, generation, "running");
 
-		// TODO: Add timeout to onActorStart
-		// Call onActorStart asynchronously and handle errors
-		this.#config
-			.onActorStart(actorId, generation, actorConfig)
-			.catch((err) => {
-				this.log?.error({
-					msg: "error in onactorstart for actor",
-					actorId,
-					err,
-				});
+		try {
+			// TODO: Add timeout to onActorStart
+			// Call onActorStart asynchronously and handle errors
+			this.log?.debug({
+				msg: "calling onActorStart",
+				actorId,
+				generation,
+			});
+			await this.#config.onActorStart(actorId, generation, actorConfig);
 
-				// TODO: Mark as crashed
-				// Send stopped state update if start failed
-				this.forceStopActor(actorId, generation);
+			instance.actorStartPromise.resolve();
+		} catch (err) {
+			this.log?.error({
+				msg: "error starting runner actor",
+				actorId,
+				err,
 			});
+
+			instance.actorStartPromise.reject(err);
+
+			// TODO: Mark as crashed
+			// Send stopped state update if start failed
+			await this.forceStopActor(actorId, generation);
+		}
 	}
 
-	#handleCommandStopActor(commandWrapper: protocol.CommandWrapper) {
+	async #handleCommandStopActor(commandWrapper: protocol.CommandWrapper) {
 		const stopCommand = commandWrapper.inner
 			.val as protocol.CommandStopActor;
 
 		const actorId = stopCommand.actorId;
 		const generation = stopCommand.generation;
 
-		this.forceStopActor(actorId, generation);
+		await this.forceStopActor(actorId, generation);
 	}
 
 	#sendActorIntent(
@@ -867,13 +1105,6 @@ export class Runner {
 		generation: number,
 		intentType: "sleep" | "stop",
 	) {
-		if (this.#shutdown) {
-			console.trace("send actor intent", actorId, intentType);
-			this.log?.warn({
-				msg: "Runner is shut down, cannot send actor intent",
-			});
-			return;
-		}
 		let actorIntent: protocol.ActorIntent;
 
 		if (intentType === "sleep") {
@@ -904,12 +1135,6 @@ export class Runner {
 
 		this.#recordEvent(eventWrapper);
 
-		this.log?.info({
-			msg: "sending event to server",
-			event: stringifyEvent(eventWrapper.inner),
-			index: eventWrapper.index.toString(),
-		});
-
 		this.__sendToServer({
 			tag: "ToServerEvents",
 			val: [eventWrapper],
@@ -921,12 +1146,6 @@ export class Runner {
 		generation: number,
 		stateType: "running" | "stopped",
 	) {
-		if (this.#shutdown) {
-			this.log?.warn({
-				msg: "Runner is shut down, cannot send actor state update",
-			});
-			return;
-		}
 		let actorState: protocol.ActorState;
 
 		if (stateType === "running") {
@@ -960,12 +1179,6 @@ export class Runner {
 
 		this.#recordEvent(eventWrapper);
 
-		this.log?.info({
-			msg: "sending event to server",
-			event: stringifyEvent(eventWrapper.inner),
-			index: eventWrapper.index.toString(),
-		});
-
 		this.__sendToServer({
 			tag: "ToServerEvents",
 			val: [eventWrapper],
@@ -973,13 +1186,6 @@ export class Runner {
 	}
 
 	#sendCommandAcknowledgment() {
-		if (this.#shutdown) {
-			this.log?.warn({
-				msg: "Runner is shut down, cannot send command acknowledgment",
-			});
-			return;
-		}
-
 		if (this.#lastCommandIdx < 0) {
 			// No commands received yet, nothing to acknowledge
 			return;
@@ -1288,11 +1494,6 @@ export class Runner {
 		const actor = this.getActor(actorId, generation);
 		if (!actor) return;
 
-		if (this.#shutdown) {
-			console.warn("Runner is shut down, cannot set alarm");
-			return;
-		}
-
 		const alarmEvent: protocol.EventActorSetAlarm = {
 			actorId,
 			generation: actor.generation,
@@ -1325,15 +1526,7 @@ export class Runner {
 		requestData: protocol.KvRequestData,
 	): Promise<any> {
 		return new Promise((resolve, reject) => {
-			if (this.#shutdown) {
-				reject(new Error("Runner is shut down"));
-				return;
-			}
-
-			const requestId = this.#nextRequestId++;
-			const isConnected =
-				this.#pegboardWebSocket &&
-				this.#pegboardWebSocket.readyState === 1;
+			const requestId = this.#nextKvRequestId++;
 
 			// Store the request
 			const requestEntry = {
@@ -1347,7 +1540,7 @@ export class Runner {
 
 			this.#kvRequests.set(requestId, requestEntry);
 
-			if (isConnected) {
+			if (this.__webSocketReady()) {
 				// Send immediately
 				this.#sendSingleKvRequest(requestId);
 			}
@@ -1380,10 +1573,7 @@ export class Runner {
 	}
 
 	#processUnsentKvRequests() {
-		if (
-			!this.#pegboardWebSocket ||
-			this.#pegboardWebSocket.readyState !== 1
-		) {
+		if (!this.__webSocketReady()) {
 			return;
 		}
 
@@ -1400,26 +1590,25 @@ export class Runner {
 		}
 	}
 
-	__webSocketReady(): boolean {
-		return this.#pegboardWebSocket
-			? this.#pegboardWebSocket.readyState === 1
-			: false;
+	/** Asserts WebSocket exists and is ready. */
+	__webSocketReady(): this is this & {
+		__pegboardWebSocket: NonNullable<Runner["__pegboardWebSocket"]>;
+	} {
+		return (
+			!!this.__pegboardWebSocket &&
+			this.__pegboardWebSocket.readyState === 1
+		);
 	}
 
 	__sendToServer(message: protocol.ToServer) {
-		if (this.#shutdown) {
-			this.log?.warn({
-				msg: "Runner is shut down, cannot send message to server",
-			});
-			return;
-		}
+		this.log?.debug({
+			msg: "sending runner message",
+			data: stringifyToServer(message),
+		});
 
 		const encoded = protocol.encodeToServer(message);
-		if (
-			this.#pegboardWebSocket &&
-			this.#pegboardWebSocket.readyState === 1
-		) {
-			this.#pegboardWebSocket.send(encoded);
+		if (this.__webSocketReady()) {
+			this.__pegboardWebSocket.send(encoded);
 		} else {
 			this.log?.error({
 				msg: "WebSocket not available or not open for sending data",
@@ -1427,8 +1616,50 @@ export class Runner {
 		}
 	}
 
-	sendWebsocketMessageAck(requestId: ArrayBuffer, index: number) {
-		this.#tunnel?.__ackWebsocketMessage(requestId, index);
+	sendHibernatableWebSocketMessageAck(
+		gatewayId: ArrayBuffer,
+		requestId: ArrayBuffer,
+		index: number,
+	) {
+		if (!this.#tunnel)
+			throw new Error("missing tunnel to send message ack");
+		this.#tunnel.sendHibernatableWebSocketMessageAck(
+			gatewayId,
+			requestId,
+			index,
+		);
+	}
+
+	/**
+	 * Restores hibernatable WebSocket connections for an actor.
+	 *
+	 * This method should be called at the end of `onActorStart` after the
+	 * actor instance is fully initialized.
+	 *
+	 * This method will:
+	 * - Restore all provided hibernatable WebSocket connections
+	 * - Attach event listeners to the restored WebSockets
+	 * - Close any WebSocket connections that failed to restore
+	 *
+	 * The provided metadata list should include all hibernatable WebSockets
+	 * that were persisted for this actor. The gateway will automatically
+	 * close any connections that are not restored (i.e., not included in
+	 * this list).
+	 *
+	 * **Important:** This method must be called after `onActorStart` completes
+	 * and before marking the actor as "ready" to ensure all hibernatable
+	 * connections are fully restored.
+	 *
+	 * @param actorId - The ID of the actor to restore connections for
+	 * @param metaEntries - Array of hibernatable WebSocket metadata to restore
+	 */
+	async restoreHibernatingRequests(
+		actorId: string,
+		metaEntries: HibernatingWebSocketMetadata[],
+	) {
+		if (!this.#tunnel)
+			throw new Error("missing tunnel to restore hibernating requests");
+		await this.#tunnel.restoreHibernatingRequests(actorId, metaEntries);
 	}
 
 	getServerlessInitPacket(): string | undefined {
@@ -1486,9 +1717,10 @@ export class Runner {
 
 		if (eventsToResend.length === 0) return;
 
-		//this.#log?.log(
-		//	`Resending ${eventsToResend.length} unacknowledged events from index ${Number(lastEventIdx) + 1}`,
-		//);
+		this.log?.info({
+			msg: "resending unacknowledged events",
+			fromIndex: lastEventIdx + 1n,
+		});
 
 		// Resend events in batches
 		this.__sendToServer({