@rivetkit/engine-runner 2.0.23 → 2.0.24

package/src/tunnel.ts

@@ -1,30 +1,44 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
-import type { MessageId, RequestId } from "@rivetkit/engine-runner-protocol";
+import type {
+	GatewayId,
+	MessageId,
+	RequestId,
+} from "@rivetkit/engine-runner-protocol";
 import type { Logger } from "pino";
-import { stringify as uuidstringify, v4 as uuidv4 } from "uuid";
-import { logger } from "./log";
-import type { ActorInstance, Runner } from "./mod";
+import {
+	parse as uuidparse,
+	stringify as uuidstringify,
+	v4 as uuidv4,
+} from "uuid";
+import type { Runner, RunnerActor } from "./mod";
 import {
 	stringifyToClientTunnelMessageKind,
 	stringifyToServerTunnelMessageKind,
 } from "./stringify";
-import { unreachable } from "./utils";
-import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
-
-const GC_INTERVAL = 60000; // 60 seconds
-const MESSAGE_ACK_TIMEOUT = 5000; // 5 seconds
-const WEBSOCKET_STATE_PERSIST_TIMEOUT = 30000; // 30 seconds
+import { arraysEqual, idToStr, unreachable } from "./utils";
+import {
+	HIBERNATABLE_SYMBOL,
+	WebSocketTunnelAdapter,
+} from "./websocket-tunnel-adapter";
 
-interface PendingRequest {
+export interface PendingRequest {
 	resolve: (response: Response) => void;
 	reject: (error: Error) => void;
 	streamController?: ReadableStreamDefaultController<Uint8Array>;
 	actorId?: string;
+	gatewayId?: GatewayId;
+	requestId?: RequestId;
+	clientMessageIndex: number;
 }
 
-interface PendingTunnelMessage {
-	sentAt: number;
-	requestIdStr: string;
+export interface HibernatingWebSocketMetadata {
+	gatewayId: GatewayId;
+	requestId: RequestId;
+	clientMessageIndex: number;
+	serverMessageIndex: number;
+
+	path: string;
+	headers: Record<string, string>;
 }
 
 class RunnerShutdownError extends Error {
@@ -36,15 +50,19 @@ class RunnerShutdownError extends Error {
 export class Tunnel {
 	#runner: Runner;
 
-	/** Requests over the tunnel to the actor that are in flight. */
-	#actorPendingRequests: Map<string, PendingRequest> = new Map();
-	/** WebSockets over the tunnel to the actor that are in flight. */
-	#actorWebSockets: Map<string, WebSocketTunnelAdapter> = new Map();
-
-	/** Messages sent from the actor over the tunnel that have not been acked by the gateway. */
-	#pendingTunnelMessages: Map<string, PendingTunnelMessage> = new Map();
+	/** Maps request IDs to actor IDs for lookup */
+	#requestToActor: Array<{
+		gatewayId: GatewayId;
+		requestId: RequestId;
+		actorId: string;
+	}> = [];
 
-	#gcInterval?: NodeJS.Timeout;
+	/** Buffer for messages when not connected */
+	#bufferedMessages: Array<{
+		gatewayId: GatewayId;
+		requestId: RequestId;
+		messageKind: protocol.ToServerTunnelMessageKind;
+	}> = [];
 
 	get log(): Logger | undefined {
 		return this.#runner.log;
@@ -55,70 +73,473 @@ export class Tunnel {
 	}
 
 	start(): void {
-		this.#startGarbageCollector();
+		// No-op - kept for compatibility
+	}
+
+	resendBufferedEvents(): void {
+		if (this.#bufferedMessages.length === 0) {
+			return;
+		}
+
+		this.log?.info({
+			msg: "resending buffered tunnel messages",
+			count: this.#bufferedMessages.length,
+		});
+
+		const messages = this.#bufferedMessages;
+		this.#bufferedMessages = [];
+
+		for (const { gatewayId, requestId, messageKind } of messages) {
+			this.#sendMessage(gatewayId, requestId, messageKind);
+		}
 	}
 
 	shutdown() {
 		// NOTE: Pegboard WS already closed at this point, cannot send
 		// anything. All teardown logic is handled by pegboard-runner.
 
-		if (this.#gcInterval) {
-			clearInterval(this.#gcInterval);
-			this.#gcInterval = undefined;
+		// Reject all pending requests and close all WebSockets for all actors
+		// RunnerShutdownError will be explicitly ignored
+		for (const [_actorId, actor] of this.#runner.actors) {
+			// Reject all pending requests for this actor
+			for (const entry of actor.pendingRequests) {
+				entry.request.reject(new RunnerShutdownError());
+			}
+			actor.pendingRequests = [];
+
+			// Close all WebSockets for this actor
+			// The WebSocket close event with retry is automatically sent when the
+			// runner WS closes, so we only need to notify the client that the WS
+			// closed:
+			// https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157
+			for (const entry of actor.webSockets) {
+				// Only close non-hibernatable websockets to prevent sending
+				// unnecessary close messages for websockets that will be hibernated
+				if (!entry.ws[HIBERNATABLE_SYMBOL]) {
+					entry.ws._closeWithoutCallback(1000, "ws.tunnel_shutdown");
+				}
+			}
+			actor.webSockets = [];
 		}
 
-		// Reject all pending requests
-		//
-		// RunnerShutdownError will be explicitly ignored
-		for (const [_, request] of this.#actorPendingRequests) {
-			request.reject(new RunnerShutdownError());
+		// Clear the request-to-actor mapping
+		this.#requestToActor = [];
+	}
+
+	async restoreHibernatingRequests(
+		actorId: string,
+		metaEntries: HibernatingWebSocketMetadata[],
+	) {
+		const actor = this.#runner.getActor(actorId);
+		if (!actor) {
+			throw new Error(
+				`Actor ${actorId} not found for restoring hibernating requests`,
+			);
 		}
-		this.#actorPendingRequests.clear();
 
-		// Close all WebSockets
-		//
-		// The WebSocket close event with retry is automatically sent when the
-		// runner WS closes, so we only need to notify the client that the WS
-		// closed:
-		// https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157
-		for (const [_, ws] of this.#actorWebSockets) {
-			// 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");
+		if (actor.hibernationRestored) {
+			throw new Error(
+				`Actor ${actorId} already restored hibernating requests`,
+			);
+		}
+
+		this.log?.debug({
+			msg: "restoring hibernating requests",
+			actorId,
+			requests: actor.hibernatingRequests.length,
+		});
+
+		// Track all background operations
+		const backgroundOperations: Promise<void>[] = [];
+
+		// Process connected WebSockets
+		let connectedButNotLoadedCount = 0;
+		let restoredCount = 0;
+		for (const { gatewayId, requestId } of actor.hibernatingRequests) {
+			const requestIdStr = idToStr(requestId);
+			const meta = metaEntries.find(
+				(entry) =>
+					arraysEqual(entry.gatewayId, gatewayId) &&
+					arraysEqual(entry.requestId, requestId),
+			);
+
+			if (!meta) {
+				// Connected but not loaded (not persisted) - close it
+				//
+				// This may happen if the metadata was not successfully persisted
+				this.log?.warn({
+					msg: "closing websocket that is not persisted",
+					requestId: requestIdStr,
+				});
+
+				this.#sendMessage(gatewayId, requestId, {
+					tag: "ToServerWebSocketClose",
+					val: {
+						code: 1000,
+						reason: "ws.meta_not_found_during_restore",
+						hibernate: false,
+					},
+				});
+
+				connectedButNotLoadedCount++;
+			} else {
+				// Both connected and persisted - restore it
+				const request = buildRequestForWebSocket(
+					meta.path,
+					meta.headers,
+				);
+
+				// This will call `runner.config.websocket` under the hood to
+				// attach the event listeners to the WebSocket.
+				// Track this operation to ensure it completes
+				const restoreOperation = this.#createWebSocket(
+					actorId,
+					gatewayId,
+					requestId,
+					requestIdStr,
+					meta.serverMessageIndex,
+					true,
+					true,
+					request,
+					meta.path,
+					meta.headers,
+					false,
+				)
+					.then(() => {
+						// Create a PendingRequest entry to track the message index
+						const actor = this.#runner.getActor(actorId);
+						if (actor) {
+							actor.createPendingRequest(
+								gatewayId,
+								requestId,
+								meta.clientMessageIndex,
+							);
+						}
+
+						this.log?.info({
+							msg: "connection successfully restored",
+							actorId,
+							requestId: requestIdStr,
+						});
+					})
+					.catch((err) => {
+						this.log?.error({
+							msg: "error creating websocket during restore",
+							requestId: requestIdStr,
+							err,
+						});
+
+						// Close the WebSocket on error
+						this.#sendMessage(gatewayId, requestId, {
+							tag: "ToServerWebSocketClose",
+							val: {
+								code: 1011,
+								reason: "ws.restore_error",
+								hibernate: false,
+							},
+						});
+					});
+
+				backgroundOperations.push(restoreOperation);
+				restoredCount++;
 			}
 		}
-		this.#actorWebSockets.clear();
+
+		// Process loaded but not connected (stale) - remove them
+		let loadedButNotConnectedCount = 0;
+		for (const meta of metaEntries) {
+			const requestIdStr = idToStr(meta.requestId);
+			const isConnected = actor.hibernatingRequests.some(
+				(req) =>
+					arraysEqual(req.gatewayId, meta.gatewayId) &&
+					arraysEqual(req.requestId, meta.requestId),
+			);
+			if (!isConnected) {
+				this.log?.warn({
+					msg: "removing stale persisted websocket",
+					requestId: requestIdStr,
+				});
+
+				const request = buildRequestForWebSocket(
+					meta.path,
+					meta.headers,
+				);
+
+				// Create adapter to register user's event listeners.
+				// Pass engineAlreadyClosed=true so close callback won't send tunnel message.
+				// Track this operation to ensure it completes
+				const cleanupOperation = this.#createWebSocket(
+					actorId,
+					meta.gatewayId,
+					meta.requestId,
+					requestIdStr,
+					meta.serverMessageIndex,
+					true,
+					true,
+					request,
+					meta.path,
+					meta.headers,
+					true,
+				)
+					.then((adapter) => {
+						// Close the adapter normally - this will fire user's close event handler
+						// (which should clean up persistence) and trigger the close callback
+						// (which will clean up maps but skip sending tunnel message)
+						adapter.close(1000, "ws.stale_metadata");
+					})
+					.catch((err) => {
+						this.log?.error({
+							msg: "error creating stale websocket during restore",
+							requestId: requestIdStr,
+							err,
+						});
+					});
+
+				backgroundOperations.push(cleanupOperation);
+				loadedButNotConnectedCount++;
+			}
+		}
+
+		// Wait for all background operations to complete before finishing
+		await Promise.allSettled(backgroundOperations);
+
+		// Mark restoration as complete
+		actor.hibernationRestored = true;
+
+		this.log?.info({
+			msg: "restored hibernatable websockets",
+			actorId,
+			restoredCount,
+			connectedButNotLoadedCount,
+			loadedButNotConnectedCount,
+		});
+	}
+
+	/**
+	 * Called from WebSocketOpen message and when restoring hibernatable WebSockets.
+	 *
+	 * engineAlreadyClosed will be true if this is only being called to trigger
+	 * the close callback and not to send a close message to the server. This
+	 * is used specifically to clean up zombie WebSocket connections.
+	 */
+	async #createWebSocket(
+		actorId: string,
+		gatewayId: GatewayId,
+		requestId: RequestId,
+		requestIdStr: string,
+		serverMessageIndex: number,
+		isHibernatable: boolean,
+		isRestoringHibernatable: boolean,
+		request: Request,
+		path: string,
+		headers: Record<string, string>,
+		engineAlreadyClosed: boolean,
+	): Promise<WebSocketTunnelAdapter> {
+		this.log?.debug({
+			msg: "createWebSocket creating adapter",
+			actorId,
+			requestIdStr,
+			isHibernatable,
+			path,
+		});
+		// Create WebSocket adapter
+		const adapter = new WebSocketTunnelAdapter(
+			this,
+			actorId,
+			requestIdStr,
+			serverMessageIndex,
+			isHibernatable,
+			isRestoringHibernatable,
+			request,
+			(data: ArrayBuffer | string, isBinary: boolean) => {
+				// Send message through tunnel
+				const dataBuffer =
+					typeof data === "string"
+						? (new TextEncoder().encode(data).buffer as ArrayBuffer)
+						: data;
+
+				this.#sendMessage(gatewayId, requestId, {
+					tag: "ToServerWebSocketMessage",
+					val: {
+						data: dataBuffer,
+						binary: isBinary,
+					},
+				});
+			},
+			(code?: number, reason?: string) => {
+				// Send close through tunnel if engine doesn't already know it's closed
+				if (!engineAlreadyClosed) {
+					this.#sendMessage(gatewayId, requestId, {
+						tag: "ToServerWebSocketClose",
+						val: {
+							code: code || null,
+							reason: reason || null,
+							hibernate: false,
+						},
+					});
+				}
+
+				// Clean up actor tracking
+				const actor = this.#runner.getActor(actorId);
+				if (actor) {
+					actor.deleteWebSocket(gatewayId, requestId);
+					actor.deletePendingRequest(gatewayId, requestId);
+				}
+
+				// Clean up request-to-actor mapping
+				this.#removeRequestToActor(gatewayId, requestId);
+			},
+		);
+
+		// Get actor and add websocket to it
+		const actor = this.#runner.getActor(actorId);
+		if (!actor) {
+			throw new Error(`Actor ${actorId} not found`);
+		}
+
+		actor.setWebSocket(gatewayId, requestId, adapter);
+		this.addRequestToActor(gatewayId, requestId, actorId);
+
+		// Call WebSocket handler. This handler will add event listeners
+		// for `open`, etc.
+		await this.#runner.config.websocket(
+			this.#runner,
+			actorId,
+			adapter,
+			gatewayId,
+			requestId,
+			request,
+			path,
+			headers,
+			isHibernatable,
+			isRestoringHibernatable,
+		);
+
+		return adapter;
+	}
+
+	addRequestToActor(
+		gatewayId: GatewayId,
+		requestId: RequestId,
+		actorId: string,
+	) {
+		this.#requestToActor.push({ gatewayId, requestId, actorId });
+	}
+
+	#removeRequestToActor(gatewayId: GatewayId, requestId: RequestId) {
+		const index = this.#requestToActor.findIndex(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		);
+		if (index !== -1) {
+			this.#requestToActor.splice(index, 1);
+		}
+	}
+
+	getRequestActor(
+		gatewayId: GatewayId,
+		requestId: RequestId,
+	): RunnerActor | undefined {
+		const entry = this.#requestToActor.find(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		);
+
+		if (!entry) {
+			this.log?.warn({
+				msg: "missing requestToActor entry",
+				requestId: idToStr(requestId),
+			});
+			return undefined;
+		}
+
+		const actor = this.#runner.getActor(entry.actorId);
+		if (!actor) {
+			this.log?.warn({
+				msg: "missing actor for requestToActor lookup",
+				requestId: idToStr(requestId),
+				actorId: entry.actorId,
+			});
+			return undefined;
+		}
+
+		return actor;
+	}
+
+	async getAndWaitForRequestActor(
+		gatewayId: GatewayId,
+		requestId: RequestId,
+	): Promise<RunnerActor | undefined> {
+		const actor = this.getRequestActor(gatewayId, requestId);
+		if (!actor) return;
+		await actor.actorStartPromise.promise;
+		return actor;
 	}
 
 	#sendMessage(
+		gatewayId: GatewayId,
 		requestId: RequestId,
 		messageKind: protocol.ToServerTunnelMessageKind,
 	) {
-		// TODO: Switch this with runner WS
+		// Buffer message if not connected
 		if (!this.#runner.__webSocketReady()) {
-			this.log?.warn({
-				msg: "cannot send tunnel message, socket not connected to engine",
+			this.log?.debug({
+				msg: "buffering tunnel message, socket not connected to engine",
 				requestId: idToStr(requestId),
 				message: stringifyToServerTunnelMessageKind(messageKind),
 			});
+			this.#bufferedMessages.push({ gatewayId, requestId, messageKind });
 			return;
 		}
 
-		// Build message
-		const messageId = generateUuidBuffer();
-
+		// Get or initialize message index for this request
+		//
+		// We don't have to wait for the actor to start since we're not calling
+		// any callbacks on the actor
+		const gatewayIdStr = idToStr(gatewayId);
 		const requestIdStr = idToStr(requestId);
-		const messageIdStr = idToStr(messageId);
-		this.#pendingTunnelMessages.set(messageIdStr, {
-			sentAt: Date.now(),
-			requestIdStr,
-		});
+		const actor = this.getRequestActor(gatewayId, requestId);
+		if (!actor) {
+			this.log?.warn({
+				msg: "cannot send tunnel message, actor not found",
+				gatewayId: gatewayIdStr,
+				requestId: requestIdStr,
+			});
+			return;
+		}
+
+		// Get message index from pending request
+		let clientMessageIndex: number;
+		const pending = actor.getPendingRequest(gatewayId, requestId);
+		if (pending) {
+			clientMessageIndex = pending.clientMessageIndex;
+			pending.clientMessageIndex++;
+		} else {
+			// No pending request
+			this.log?.warn({
+				msg: "missing pending request for send message, defaulting to message index 0",
+				gatewayId: gatewayIdStr,
+				requestId: requestIdStr,
+			});
+			clientMessageIndex = 0;
+		}
+
+		// Build message ID from gatewayId + requestId + messageIndex
+		const messageId: protocol.MessageId = {
+			gatewayId,
+			requestId,
+			messageIndex: clientMessageIndex,
+		};
+		const messageIdStr = `${idToStr(messageId.gatewayId)}-${idToStr(messageId.requestId)}-${messageId.messageIndex}`;
 
 		this.log?.debug({
-			msg: "send tunnel msg",
-			requestId: requestIdStr,
+			msg: "sending tunnel msg",
 			messageId: messageIdStr,
+			gatewayId: gatewayIdStr,
+			requestId: requestIdStr,
+			messageIndex: clientMessageIndex,
 			message: stringifyToServerTunnelMessageKind(messageKind),
 		});
 
@@ -126,7 +547,6 @@ export class Tunnel {
 		const message: protocol.ToServer = {
 			tag: "ToServerTunnelMessage",
 			val: {
-				requestId,
 				messageId,
 				messageKind,
 			},
@@ -134,123 +554,33 @@ export class Tunnel {
 		this.#runner.__sendToServer(message);
 	}
 
-	#sendAck(requestId: RequestId, messageId: MessageId) {
-		if (!this.#runner.__webSocketReady()) {
-			return;
-		}
-
-		const message: protocol.ToServer = {
-			tag: "ToServerTunnelMessage",
-			val: {
-				requestId,
-				messageId,
-				messageKind: { tag: "TunnelAck", val: null },
-			},
-		};
-
-		this.log?.debug({
-			msg: "ack tunnel msg",
-			requestId: idToStr(requestId),
-			messageId: idToStr(messageId),
-		});
-
-		this.#runner.__sendToServer(message);
-	}
-
-	#startGarbageCollector() {
-		if (this.#gcInterval) {
-			clearInterval(this.#gcInterval);
-		}
-
-		this.#gcInterval = setInterval(() => {
-			this.#gc();
-		}, GC_INTERVAL);
-	}
-
-	#gc() {
-		const now = Date.now();
-		const messagesToDelete: string[] = [];
-
-		for (const [messageId, pendingMessage] of this.#pendingTunnelMessages) {
-			// Check if message is older than timeout
-			if (now - pendingMessage.sentAt > MESSAGE_ACK_TIMEOUT) {
-				messagesToDelete.push(messageId);
-
-				const requestIdStr = pendingMessage.requestIdStr;
-
-				// Check if this is an HTTP request
-				const pendingRequest =
-					this.#actorPendingRequests.get(requestIdStr);
-				if (pendingRequest) {
-					// Reject the pending HTTP request
-					pendingRequest.reject(
-						new Error("Message acknowledgment timeout"),
-					);
-
-					// Close stream controller if it exists
-					if (pendingRequest.streamController) {
-						pendingRequest.streamController.error(
-							new Error("Message acknowledgment timeout"),
-						);
-					}
-
-					// Clean up from actorPendingRequests map
-					this.#actorPendingRequests.delete(requestIdStr);
-				}
-
-				// Check if this is a WebSocket
-				const webSocket = this.#actorWebSockets.get(requestIdStr);
-				if (webSocket) {
-					// Close the WebSocket connection
-					webSocket.__closeWithRetry(
-						1000,
-						"Message acknowledgment timeout",
-					);
-
-					// Clean up from actorWebSockets map
-					this.#actorWebSockets.delete(requestIdStr);
-				}
-			}
-		}
-
-		// Remove timed out messages
-		if (messagesToDelete.length > 0) {
-			this.log?.warn({
-				msg: "purging unacked tunnel messages, this indicates that the Rivet Engine is disconnected or not responding",
-				count: messagesToDelete.length,
-			});
-			for (const messageId of messagesToDelete) {
-				this.#pendingTunnelMessages.delete(messageId);
-			}
-		}
-	}
-
-	closeActiveRequests(actor: ActorInstance) {
+	closeActiveRequests(actor: RunnerActor) {
 		const actorId = actor.actorId;
 
-		// Terminate all requests for this actor
-		for (const requestId of actor.requests) {
-			const pending = this.#actorPendingRequests.get(requestId);
-			if (pending) {
-				pending.reject(new Error(`Actor ${actorId} stopped`));
-				this.#actorPendingRequests.delete(requestId);
+		// Terminate all requests for this actor. This will no send a
+		// ToServerResponse* message since the actor will no longer be loaded.
+		// The gateway is responsible for closing the request.
+		for (const entry of actor.pendingRequests) {
+			entry.request.reject(new Error(`Actor ${actorId} stopped`));
+			if (entry.gatewayId && entry.requestId) {
+				this.#removeRequestToActor(entry.gatewayId, entry.requestId);
 			}
 		}
-		actor.requests.clear();
 
-		// Flush acks and close all WebSockets for this actor
-		for (const requestIdStr of actor.webSockets) {
-			const ws = this.#actorWebSockets.get(requestIdStr);
-			if (ws) {
-				ws.__closeWithRetry(1000, "Actor stopped");
-				this.#actorWebSockets.delete(requestIdStr);
+		// Close all WebSockets. Only send close event to non-HWS. The gateway is
+		// responsible for hibernating HWS and closing regular WS.
+		for (const entry of actor.webSockets) {
+			const isHibernatable = entry.ws[HIBERNATABLE_SYMBOL];
+			if (!isHibernatable) {
+				entry.ws._closeWithoutCallback(1000, "actor.stopped");
 			}
+			// Note: request-to-actor mapping is cleaned up in the close callback
 		}
-		actor.webSockets.clear();
 	}
 
 	async #fetch(
 		actorId: string,
+		gatewayId: protocol.GatewayId,
 		requestId: protocol.RequestId,
 		request: Request,
 	): Promise<Response> {
@@ -274,6 +604,7 @@ export class Tunnel {
 		const fetchHandler = this.#runner.config.fetch(
 			this.#runner,
 			actorId,
+			gatewayId,
 			requestId,
 			request,
 		);
@@ -286,94 +617,88 @@ export class Tunnel {
 	}
 
 	async handleTunnelMessage(message: protocol.ToClientTunnelMessage) {
-		const requestIdStr = idToStr(message.requestId);
-		const messageIdStr = idToStr(message.messageId);
+		// Parse the gateway ID, request ID, and message index from the messageId
+		const { gatewayId, requestId, messageIndex } = message.messageId;
+
+		const gatewayIdStr = idToStr(gatewayId);
+		const requestIdStr = idToStr(requestId);
 		this.log?.debug({
 			msg: "receive tunnel msg",
+			gatewayId: gatewayIdStr,
 			requestId: requestIdStr,
-			messageId: messageIdStr,
+			messageIndex: message.messageId.messageIndex,
 			message: stringifyToClientTunnelMessageKind(message.messageKind),
 		});
 
-		if (message.messageKind.tag === "TunnelAck") {
-			// Mark pending message as acknowledged and remove it
-			const pending = this.#pendingTunnelMessages.get(messageIdStr);
-			if (pending) {
-				const didDelete =
-					this.#pendingTunnelMessages.delete(messageIdStr);
-				if (!didDelete) {
-					this.log?.warn({
-						msg: "received tunnel ack for nonexistent message",
-						requestId: requestIdStr,
-						messageId: messageIdStr,
-					});
-				}
-			}
-		} else {
-			switch (message.messageKind.tag) {
-				case "ToClientRequestStart":
-					this.#sendAck(message.requestId, message.messageId);
-
-					await this.#handleRequestStart(
-						message.requestId,
-						message.messageKind.val,
-					);
-					break;
-				case "ToClientRequestChunk":
-					this.#sendAck(message.requestId, message.messageId);
-
-					await this.#handleRequestChunk(
-						message.requestId,
-						message.messageKind.val,
-					);
-					break;
-				case "ToClientRequestAbort":
-					this.#sendAck(message.requestId, message.messageId);
-
-					await this.#handleRequestAbort(message.requestId);
-					break;
-				case "ToClientWebSocketOpen":
-					this.#sendAck(message.requestId, message.messageId);
-
-					await this.#handleWebSocketOpen(
-						message.requestId,
-						message.messageKind.val,
-					);
-					break;
-				case "ToClientWebSocketMessage": {
-					this.#sendAck(message.requestId, message.messageId);
-
-					const _unhandled = await this.#handleWebSocketMessage(
-						message.requestId,
-						message.messageKind.val,
-					);
-					break;
-				}
-				case "ToClientWebSocketClose":
-					this.#sendAck(message.requestId, message.messageId);
-
-					await this.#handleWebSocketClose(
-						message.requestId,
-						message.messageKind.val,
-					);
-					break;
-				default:
-					unreachable(message.messageKind);
+		switch (message.messageKind.tag) {
+			case "ToClientRequestStart":
+				await this.#handleRequestStart(
+					gatewayId,
+					requestId,
+					message.messageKind.val,
+				);
+				break;
+			case "ToClientRequestChunk":
+				await this.#handleRequestChunk(
+					gatewayId,
+					requestId,
+					message.messageKind.val,
+				);
+				break;
+			case "ToClientRequestAbort":
+				await this.#handleRequestAbort(gatewayId, requestId);
+				break;
+			case "ToClientWebSocketOpen":
+				await this.#handleWebSocketOpen(
+					gatewayId,
+					requestId,
+					message.messageKind.val,
+				);
+				break;
+			case "ToClientWebSocketMessage": {
+				await this.#handleWebSocketMessage(
+					gatewayId,
+					requestId,
+					messageIndex,
+					message.messageKind.val,
+				);
+				break;
 			}
+			case "ToClientWebSocketClose":
+				await this.#handleWebSocketClose(
+					gatewayId,
+					requestId,
+					message.messageKind.val,
+				);
+				break;
+			case "DeprecatedTunnelAck":
+				// Ignore deprecated tunnel ack messages
+				break;
+			default:
+				unreachable(message.messageKind);
 		}
 	}
 
 	async #handleRequestStart(
-		requestId: ArrayBuffer,
+		gatewayId: GatewayId,
+		requestId: RequestId,
 		req: protocol.ToClientRequestStart,
 	) {
 		// Track this request for the actor
 		const requestIdStr = idToStr(requestId);
-		const actor = this.#runner.getActor(req.actorId);
-		if (actor) {
-			actor.requests.add(requestIdStr);
+		const actor = await this.#runner.getAndWaitForActor(req.actorId);
+		if (!actor) {
+			this.log?.warn({
+				msg: "actor does not exist in handleRequestStart, request will leak",
+				actorId: req.actorId,
+				requestId: requestIdStr,
+			});
+			return;
 		}
 
+		// Add to request-to-actor mapping
+		this.addRequestToActor(gatewayId, requestId, req.actorId);
+
 		try {
 			// Convert headers map to Headers object
 			const headers = new Headers();
@@ -394,18 +719,22 @@ export class Tunnel {
 				const stream = new ReadableStream<Uint8Array>({
 					start: (controller) => {
 						// Store controller for chunks
-						const existing =
-							this.#actorPendingRequests.get(requestIdStr);
+						const existing = actor.getPendingRequest(
+							gatewayId,
+							requestId,
+						);
 						if (existing) {
 							existing.streamController = controller;
 							existing.actorId = req.actorId;
+							existing.gatewayId = gatewayId;
+							existing.requestId = requestId;
 						} else {
-							this.#actorPendingRequests.set(requestIdStr, {
-								resolve: () => {},
-								reject: () => {},
-								streamController: controller,
-								actorId: req.actorId,
-							});
+							actor.createPendingRequestWithStreamController(
+								gatewayId,
+								requestId,
+								0,
+								controller,
+							);
 						}
 					},
 				});
@@ -419,18 +748,35 @@ export class Tunnel {
 				// Call fetch handler with validation
 				const response = await this.#fetch(
 					req.actorId,
+					gatewayId,
 					requestId,
 					streamingRequest,
 				);
-				await this.#sendResponse(requestId, response);
+				await this.#sendResponse(
+					actor.actorId,
+					actor.generation,
+					gatewayId,
+					requestId,
+					response,
+				);
 			} else {
 				// Non-streaming request
+				// Create a pending request entry to track messageIndex for the response
+				actor.createPendingRequest(gatewayId, requestId, 0);
+
 				const response = await this.#fetch(
 					req.actorId,
+					gatewayId,
 					requestId,
 					request,
 				);
-				await this.#sendResponse(requestId, response);
+				await this.#sendResponse(
+					actor.actorId,
+					actor.generation,
+					gatewayId,
+					requestId,
+					response,
+				);
 			}
 		} catch (error) {
 			if (error instanceof RunnerShutdownError) {
@@ -438,6 +784,9 @@ export class Tunnel {
 			} else {
 				this.log?.error({ msg: "error handling request", error });
 				this.#sendResponseError(
+					actor.actorId,
+					actor.generation,
+					gatewayId,
 					requestId,
 					500,
 					"Internal Server Error",
@@ -445,38 +794,67 @@ export class Tunnel {
 			}
 		} finally {
 			// Clean up request tracking
-			const actor = this.#runner.getActor(req.actorId);
-			if (actor) {
-				actor.requests.delete(requestIdStr);
+			if (this.#runner.hasActor(req.actorId, actor.generation)) {
+				actor.deletePendingRequest(gatewayId, requestId);
+				this.#removeRequestToActor(gatewayId, requestId);
 			}
 		}
 	}
 
 	async #handleRequestChunk(
-		requestId: ArrayBuffer,
+		gatewayId: GatewayId,
+		requestId: RequestId,
 		chunk: protocol.ToClientRequestChunk,
 	) {
-		const requestIdStr = idToStr(requestId);
-		const pending = this.#actorPendingRequests.get(requestIdStr);
-		if (pending?.streamController) {
-			pending.streamController.enqueue(new Uint8Array(chunk.body));
-			if (chunk.finish) {
-				pending.streamController.close();
-				this.#actorPendingRequests.delete(requestIdStr);
+		const actor = await this.getAndWaitForRequestActor(
+			gatewayId,
+			requestId,
+		);
+		if (actor) {
+			const pending = actor.getPendingRequest(gatewayId, requestId);
+			if (pending?.streamController) {
+				pending.streamController.enqueue(new Uint8Array(chunk.body));
+				if (chunk.finish) {
+					pending.streamController.close();
+					actor.deletePendingRequest(gatewayId, requestId);
+					this.#removeRequestToActor(gatewayId, requestId);
+				}
 			}
 		}
 	}
 
-	async #handleRequestAbort(requestId: ArrayBuffer) {
-		const requestIdStr = idToStr(requestId);
-		const pending = this.#actorPendingRequests.get(requestIdStr);
-		if (pending?.streamController) {
-			pending.streamController.error(new Error("Request aborted"));
+	async #handleRequestAbort(gatewayId: GatewayId, requestId: RequestId) {
+		const actor = await this.getAndWaitForRequestActor(
+			gatewayId,
+			requestId,
+		);
+		if (actor) {
+			const pending = actor.getPendingRequest(gatewayId, requestId);
+			if (pending?.streamController) {
+				pending.streamController.error(new Error("Request aborted"));
+			}
+			actor.deletePendingRequest(gatewayId, requestId);
+			this.#removeRequestToActor(gatewayId, requestId);
 		}
-		this.#actorPendingRequests.delete(requestIdStr);
 	}
 
-	async #sendResponse(requestId: ArrayBuffer, response: Response) {
+	async #sendResponse(
+		actorId: string,
+		generation: number,
+		gatewayId: GatewayId,
+		requestId: ArrayBuffer,
+		response: Response,
+	) {
+		if (!this.#runner.hasActor(actorId, generation)) {
+			this.log?.warn({
+				msg: "actor not loaded to send response, assuming gateway has closed request",
+				actorId,
+				generation,
+				requestId,
+			});
+			return;
+		}
+
 		// Always treat responses as non-streaming for now
 		// In the future, we could detect streaming responses based on:
 		// - Transfer-Encoding: chunked
@@ -497,8 +875,8 @@ export class Tunnel {
 			headers.set("content-length", String(body.byteLength));
 		}
 
-		// Send as non-streaming response
-		this.#sendMessage(requestId, {
+		// Send as non-streaming response if actor has not stopped
+		this.#sendMessage(gatewayId, requestId, {
 			tag: "ToServerResponseStart",
 			val: {
 				status: response.status as protocol.u16,
@@ -510,14 +888,27 @@ export class Tunnel {
 	}
 
 	#sendResponseError(
+		actorId: string,
+		generation: number,
+		gatewayId: GatewayId,
 		requestId: ArrayBuffer,
 		status: number,
 		message: string,
 	) {
+		if (!this.#runner.hasActor(actorId, generation)) {
+			this.log?.warn({
+				msg: "actor not loaded to send response, assuming gateway has closed request",
+				actorId,
+				generation,
+				requestId,
+			});
+			return;
+		}
+
 		const headers = new Map<string, string>();
 		headers.set("content-type", "text/plain");
 
-		this.#sendMessage(requestId, {
+		this.#sendMessage(gatewayId, requestId, {
 			tag: "ToServerResponseStart",
 			val: {
 				status: status as protocol.u16,
@@ -529,13 +920,20 @@ export class Tunnel {
 	}
 
 	async #handleWebSocketOpen(
-		requestId: protocol.RequestId,
+		gatewayId: GatewayId,
+		requestId: RequestId,
 		open: protocol.ToClientWebSocketOpen,
 	) {
+		// NOTE: This method is safe to be async since we will not receive any
+		// further WebSocket events until we send a ToServerWebSocketOpen
+		// tunnel message. We can do any async logic we need to between thoes two events.
+		//
+		// Sending a ToServerWebSocketClose will terminate the WebSocket early.
+
 		const requestIdStr = idToStr(requestId);
 
 		// Validate actor exists
-		const actor = this.#runner.getActor(open.actorId);
+		const actor = await this.#runner.getAndWaitForActor(open.actorId);
 		if (!actor) {
 			this.log?.warn({
 				msg: "ignoring websocket for unknown actor",
@@ -547,43 +945,21 @@ export class Tunnel {
 			//
 			// See
 			// https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/lib.rs#L238
-			this.#sendMessage(requestId, {
+			this.#sendMessage(gatewayId, requestId, {
 				tag: "ToServerWebSocketClose",
 				val: {
 					code: 1011,
 					reason: "Actor not found",
-					retry: false,
-				},
-			});
-			return;
-		}
-
-		const websocketHandler = this.#runner.config.websocket;
-
-		if (!websocketHandler) {
-			this.log?.error({
-				msg: "no websocket handler configured for tunnel",
-			});
-			// Send close immediately
-			this.#sendMessage(requestId, {
-				tag: "ToServerWebSocketClose",
-				val: {
-					code: 1011,
-					reason: "Not Implemented",
-					retry: false,
+					hibernate: false,
 				},
 			});
 			return;
 		}
 
 		// Close existing WebSocket if one already exists for this request ID.
-		// There should always be a close message sent before another open
-		// message for the same message ID.
-		//
-		// This should never occur if all is functioning correctly, but this
-		// prevents any edge case that would result in duplicate WebSockets for
-		// the same request.
-		const existingAdapter = this.#actorWebSockets.get(requestIdStr);
+		// This should never happen, but prevents any potential duplicate
+		// WebSockets from retransmits.
+		const existingAdapter = actor.getWebSocket(gatewayId, requestId);
 		if (existingAdapter) {
 			this.log?.warn({
 				msg: "closing existing websocket for duplicate open event for the same request id",
@@ -591,194 +967,212 @@ export class Tunnel {
 			});
 			// Close without sending a message through the tunnel since the server
 			// already knows about the new connection
-			existingAdapter.__closeWithoutCallback(1000, "ws.duplicate_open");
-		}
-
-		// Track this WebSocket for the actor
-		if (actor) {
-			actor.webSockets.add(requestIdStr);
+			existingAdapter._closeWithoutCallback(1000, "ws.duplicate_open");
 		}
 
+		// Create WebSocket
 		try {
-			// Create WebSocket adapter
-			const adapter = new WebSocketTunnelAdapter(
-				requestIdStr,
-				(data: ArrayBuffer | string, isBinary: boolean) => {
-					// Send message through tunnel
-					const dataBuffer =
-						typeof data === "string"
-							? (new TextEncoder().encode(data)
-									.buffer as ArrayBuffer)
-							: data;
-
-					this.#sendMessage(requestId, {
-						tag: "ToServerWebSocketMessage",
-						val: {
-							data: dataBuffer,
-							binary: isBinary,
-						},
-					});
-				},
-				(code?: number, reason?: string, retry: boolean = false) => {
-					// Send close through tunnel
-					this.#sendMessage(requestId, {
-						tag: "ToServerWebSocketClose",
-						val: {
-							code: code || null,
-							reason: reason || null,
-							retry,
-						},
-					});
-
-					// Remove from map
-					this.#actorWebSockets.delete(requestIdStr);
-
-					// Clean up actor tracking
-					if (actor) {
-						actor.webSockets.delete(requestIdStr);
-					}
-				},
+			const request = buildRequestForWebSocket(
+				open.path,
+				Object.fromEntries(open.headers),
 			);
 
-			// Store adapter
-			this.#actorWebSockets.set(requestIdStr, adapter);
-
-			// Convert headers to map
-			//
-			// We need to manually ensure the original Upgrade/Connection WS
-			// headers are present
-			const headerInit: Record<string, string> = {};
-			if (open.headers) {
-				for (const [k, v] of open.headers as ReadonlyMap<
-					string,
-					string
-				>) {
-					headerInit[k] = v;
-				}
-			}
-			headerInit["Upgrade"] = "websocket";
-			headerInit["Connection"] = "Upgrade";
-
-			const request = new Request(`http://localhost${open.path}`, {
-				method: "GET",
-				headers: headerInit,
-			});
-
-			// Send open confirmation
-			const hibernationConfig =
-				this.#runner.config.getActorHibernationConfig(
+			const canHibernate =
+				this.#runner.config.hibernatableWebSocket.canHibernate(
 					actor.actorId,
+					gatewayId,
 					requestId,
 					request,
 				);
-			adapter.canHibernate = hibernationConfig.enabled;
 
-			this.#sendMessage(requestId, {
+			// #createWebSocket will call `runner.config.websocket` under the
+			// hood to add the event listeners for open, etc. If this handler
+			// throws, then the WebSocket will be closed before sending the
+			// open event.
+			const adapter = await this.#createWebSocket(
+				actor.actorId,
+				gatewayId,
+				requestId,
+				requestIdStr,
+				0,
+				canHibernate,
+				false,
+				request,
+				open.path,
+				Object.fromEntries(open.headers),
+				false,
+			);
+
+			// Create a PendingRequest entry to track the message index
+			actor.createPendingRequest(gatewayId, requestId, 0);
+
+			// Open the WebSocket after `config.socket` so (a) the event
+			// handlers can be added and (b) any errors in `config.websocket`
+			// will cause the WebSocket to terminate before the open event.
+			this.#sendMessage(gatewayId, requestId, {
 				tag: "ToServerWebSocketOpen",
 				val: {
-					canHibernate: hibernationConfig.enabled,
-					lastMsgIndex: BigInt(hibernationConfig.lastMsgIndex ?? -1),
+					canHibernate,
 				},
 			});
 
-			// Notify adapter that connection is open
+			// Dispatch open event
 			adapter._handleOpen(requestId);
-
-			// Call websocket handler
-			await websocketHandler(
-				this.#runner,
-				open.actorId,
-				adapter,
-				requestId,
-				request,
-			);
 		} catch (error) {
 			this.log?.error({ msg: "error handling websocket open", error });
+
+			// TODO: Call close event on adapter if needed
+
 			// Send close on error
-			this.#sendMessage(requestId, {
+			this.#sendMessage(gatewayId, requestId, {
 				tag: "ToServerWebSocketClose",
 				val: {
 					code: 1011,
 					reason: "Server Error",
-					retry: false,
+					hibernate: false,
 				},
 			});
 
-			this.#actorWebSockets.delete(requestIdStr);
-
 			// Clean up actor tracking
-			if (actor) {
-				actor.webSockets.delete(requestIdStr);
-			}
+			actor.deleteWebSocket(gatewayId, requestId);
+			actor.deletePendingRequest(gatewayId, requestId);
+			this.#removeRequestToActor(gatewayId, requestId);
 		}
 	}
 
-	/// Returns false if the message was sent off
 	async #handleWebSocketMessage(
-		requestId: ArrayBuffer,
+		gatewayId: GatewayId,
+		requestId: RequestId,
+		serverMessageIndex: number,
 		msg: protocol.ToClientWebSocketMessage,
-	): Promise<boolean> {
-		const requestIdStr = idToStr(requestId);
-		const adapter = this.#actorWebSockets.get(requestIdStr);
-		if (adapter) {
-			const data = msg.binary
-				? new Uint8Array(msg.data)
-				: new TextDecoder().decode(new Uint8Array(msg.data));
+	) {
+		const actor = await this.getAndWaitForRequestActor(
+			gatewayId,
+			requestId,
+		);
+		if (actor) {
+			const adapter = actor.getWebSocket(gatewayId, requestId);
+			if (adapter) {
+				const data = msg.binary
+					? new Uint8Array(msg.data)
+					: new TextDecoder().decode(new Uint8Array(msg.data));
 
-			return adapter._handleMessage(
-				requestId,
-				data,
-				msg.index,
-				msg.binary,
-			);
-		} else {
-			return true;
+				adapter._handleMessage(
+					requestId,
+					data,
+					serverMessageIndex,
+					msg.binary,
+				);
+				return;
+			}
 		}
+
+		// TODO: This will never retransmit the socket and the socket will close
+		this.log?.warn({
+			msg: "missing websocket for incoming websocket message, this may indicate the actor stopped before processing a message",
+			requestId,
+		});
 	}
 
-	__ackWebsocketMessage(requestId: ArrayBuffer, index: number) {
+	sendHibernatableWebSocketMessageAck(
+		gatewayId: ArrayBuffer,
+		requestId: ArrayBuffer,
+		clientMessageIndex: number,
+	) {
+		const requestIdStr = idToStr(requestId);
+
 		this.log?.debug({
 			msg: "ack ws msg",
-			requestId: idToStr(requestId),
-			index,
+			requestId: requestIdStr,
+			index: clientMessageIndex,
 		});
 
-		if (index < 0 || index > 65535)
+		if (clientMessageIndex < 0 || clientMessageIndex > 65535)
 			throw new Error("invalid websocket ack index");
 
+		// Get the actor to find the gatewayId
+		//
+		// We don't have to wait for the actor to start since we're not calling
+		// any callbacks on the actor
+		const actor = this.getRequestActor(gatewayId, requestId);
+		if (!actor) {
+			this.log?.warn({
+				msg: "cannot send websocket ack, actor not found",
+				requestId: requestIdStr,
+			});
+			return;
+		}
+
+		// Get gatewayId from the pending request
+		const pending = actor.getPendingRequest(gatewayId, requestId);
+		if (!pending?.gatewayId) {
+			this.log?.warn({
+				msg: "cannot send websocket ack, gatewayId not found in pending request",
+				requestId: requestIdStr,
+			});
+			return;
+		}
+
 		// Send the ack message
-		this.#sendMessage(requestId, {
+		this.#sendMessage(pending.gatewayId, requestId, {
 			tag: "ToServerWebSocketMessageAck",
 			val: {
-				index,
+				index: clientMessageIndex,
 			},
 		});
 	}
 
 	async #handleWebSocketClose(
-		requestId: ArrayBuffer,
+		gatewayId: GatewayId,
+		requestId: RequestId,
 		close: protocol.ToClientWebSocketClose,
 	) {
-		const requestIdStr = idToStr(requestId);
-		const adapter = this.#actorWebSockets.get(requestIdStr);
-		if (adapter) {
-			adapter._handleClose(
-				requestId,
-				close.code || undefined,
-				close.reason || undefined,
-			);
-			this.#actorWebSockets.delete(requestIdStr);
+		const actor = await this.getAndWaitForRequestActor(
+			gatewayId,
+			requestId,
+		);
+		if (actor) {
+			const adapter = actor.getWebSocket(gatewayId, requestId);
+			if (adapter) {
+				// We don't need to send a close response
+				adapter._handleClose(
+					requestId,
+					close.code || undefined,
+					close.reason || undefined,
+				);
+				actor.deleteWebSocket(gatewayId, requestId);
+				actor.deletePendingRequest(gatewayId, requestId);
+				this.#removeRequestToActor(gatewayId, requestId);
+			}
 		}
 	}
 }
 
-/** Generates a UUID as bytes. */
-function generateUuidBuffer(): ArrayBuffer {
-	const buffer = new Uint8Array(16);
-	uuidv4(undefined, buffer);
-	return buffer.buffer;
-}
+/**
+ * Builds a request that represents the incoming request for a given WebSocket.
+ *
+ * This request is not a real request and will never be sent. It's used to be passed to the actor to behave like a real incoming request.
+ */
+function buildRequestForWebSocket(
+	path: string,
+	headers: Record<string, string>,
+): Request {
+	// We need to manually ensure the original Upgrade/Connection WS
+	// headers are present
+	const fullHeaders = {
+		...headers,
+		Upgrade: "websocket",
+		Connection: "Upgrade",
+	};
+
+	if (!path.startsWith("/")) {
+		throw new Error("path must start with leading slash");
+	}
+
+	const request = new Request(`http://actor${path}`, {
+		method: "GET",
+		headers: fullHeaders,
+	});
 
-function idToStr(id: ArrayBuffer): string {
-	return uuidstringify(new Uint8Array(id));
+	return request;
 }