@rivetkit/engine-runner 2.0.21 → 2.0.22-rc.2

package/src/mod.ts

@@ -8,7 +8,8 @@ import { importWebSocket } from "./websocket.js";
 import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
 
 const KV_EXPIRE: number = 30_000;
-const PROTOCOL_VERSION: number = 1;
+const PROTOCOL_VERSION: number = 2;
+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;
@@ -43,17 +44,19 @@ export interface RunnerConfig {
 	prepopulateActorNames: Record<string, { metadata: Record<string, any> }>;
 	metadata?: Record<string, any>;
 	onConnected: () => void;
-	onDisconnected: () => void;
+	onDisconnected: (code: number, reason: string) => void;
 	onShutdown: () => void;
 	fetch: (
 		runner: Runner,
 		actorId: string,
+		requestId: protocol.RequestId,
 		request: Request,
 	) => Promise<Response>;
 	websocket?: (
 		runner: Runner,
 		actorId: string,
 		ws: any,
+		requestId: protocol.RequestId,
 		request: Request,
 	) => Promise<void>;
 	onActorStart: (
@@ -62,9 +65,19 @@ export interface RunnerConfig {
 		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;
@@ -155,26 +168,21 @@ export class Runner {
 		const actor = this.#removeActor(actorId, generation);
 		if (!actor) return;
 
-		// Unregister actor from tunnel
-		this.#tunnel?.unregisterActor(actor);
-
 		// If onActorStop times out, Pegboard will handle this timeout with ACTOR_STOP_THRESHOLD_DURATION_MS
+		//
+		// If we receive a request while onActorStop is running, a Service
+		// Unavailable error will be returned to Guard and the request will be
+		// retried
 		try {
 			await this.#config.onActorStop(actorId, actor.generation);
 		} catch (err) {
 			console.error(`Error in onActorStop for actor ${actorId}:`, err);
 		}
 
-		this.#sendActorStateUpdate(actorId, actor.generation, "stopped");
+		// Close requests after onActorStop so you can send messages over the tunnel
+		this.#tunnel?.closeActiveRequests(actor);
 
-		this.#config.onActorStop(actorId, actor.generation).catch((err) => {
-			logger()?.error({
-				msg: "error in onactorstop for actor",
-				runnerId: this.runnerId,
-				actorId,
-				err,
-			});
-		});
+		this.#sendActorStateUpdate(actorId, actor.generation, "stopped");
 	}
 
 	#stopAllActors() {
@@ -221,6 +229,7 @@ export class Runner {
 		);
 	}
 
+	// IMPORTANT: Make sure to call stopActiveRequests if calling #removeActor
 	#removeActor(
 		actorId: string,
 		generation?: number,
@@ -246,24 +255,6 @@ export class Runner {
 
 		this.#actors.delete(actorId);
 
-		// Close all WebSocket connections for this actor
-		const actorWebSockets = this.#actorWebSockets.get(actorId);
-		if (actorWebSockets) {
-			for (const ws of actorWebSockets) {
-				try {
-					ws.close(1000, "Actor stopped");
-				} catch (err) {
-					logger()?.error({
-						msg: "error closing websocket for actor",
-						runnerId: this.runnerId,
-						actorId,
-						err,
-					});
-				}
-			}
-			this.#actorWebSockets.delete(actorId);
-		}
-
 		return actor;
 	}
 
@@ -458,17 +449,24 @@ export class Runner {
 	}
 
 	// MARK: Networking
+	get pegboardEndpoint() {
+		return this.#config.pegboardEndpoint || this.#config.endpoint;
+	}
 	get pegboardUrl() {
-		const endpoint = this.#config.pegboardEndpoint || this.#config.endpoint;
-		const wsEndpoint = endpoint
+		const wsEndpoint = this.pegboardEndpoint
 			.replace("http://", "ws://")
 			.replace("https://", "wss://");
-		return `${wsEndpoint}?protocol_version=${PROTOCOL_VERSION}&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#config.runnerKey)}`;
+
+		// Ensure the endpoint ends with /runners/connect
+		const baseUrl = wsEndpoint.endsWith("/")
+			? wsEndpoint.slice(0, -1)
+			: wsEndpoint;
+		return `${baseUrl}/runners/connect?protocol_version=${PROTOCOL_VERSION}&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#config.runnerKey)}`;
 	}
 
 	// MARK: Runner protocol
 	async #openPegboardWebSocket() {
-		const protocols = ["rivet", `rivet_target.runner`];
+		const protocols = ["rivet"];
 		if (this.config.token)
 			protocols.push(`rivet_token.${this.config.token}`);
 
@@ -476,8 +474,16 @@ export class Runner {
 		const ws = new WS(this.pegboardUrl, protocols) as any as WebSocket;
 		this.#pegboardWebSocket = ws;
 
+		logger()?.info({
+			msg: "connecting",
+			endpoint: this.pegboardEndpoint,
+			namespace: this.#config.namespace,
+			runnerKey: this.#config.runnerKey,
+			hasToken: !!this.config.token,
+		});
+
 		ws.addEventListener("open", () => {
-			logger()?.info({ msg: "Connected" });
+			logger()?.info({ msg: "connected" });
 
 			// Reset reconnect attempt counter on successful connection
 			this.#reconnectAttempt = 0;
@@ -523,7 +529,6 @@ export class Runner {
 			this.#processUnsentKvRequests();
 
 			// Start ping interval
-			const pingInterval = 1000;
 			const pingLoop = setInterval(() => {
 				if (ws.readyState === 1) {
 					this.__sendToServer({
@@ -539,7 +544,7 @@ export class Runner {
 						runnerId: this.runnerId,
 					});
 				}
-			}, pingInterval);
+			}, RUNNER_PING_INTERVAL);
 			this.#pingLoop = pingLoop;
 
 			// Start command acknowledgment interval (5 minutes)
@@ -652,7 +657,7 @@ export class Runner {
 				reason: ev.reason.toString(),
 			});
 
-			this.#config.onDisconnected();
+			this.#config.onDisconnected(ev.code, ev.reason);
 
 			if (ev.reason.toString().startsWith("ws.eviction")) {
 				logger()?.info({
@@ -1390,6 +1395,10 @@ export class Runner {
 		}
 	}
 
+	sendWebsocketMessageAck(requestId: ArrayBuffer, index: number) {
+		this.#tunnel?.__ackWebsocketMessage(requestId, index);
+	}
+
 	getServerlessInitPacket(): string | undefined {
 		if (!this.runnerId) return undefined;
 

package/src/tunnel.ts

@@ -1,6 +1,6 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
 import type { MessageId, RequestId } from "@rivetkit/engine-runner-protocol";
-import { v4 as uuidv4 } from "uuid";
+import { stringify as uuidstringify, v4 as uuidv4 } from "uuid";
 import { logger } from "./log";
 import type { ActorInstance, Runner } from "./mod";
 import { unreachable } from "./utils";
@@ -8,6 +8,7 @@ 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
 
 interface PendingRequest {
 	resolve: (response: Response) => void;
@@ -56,7 +57,7 @@ export class Tunnel {
 
 		// Close all WebSockets
 		for (const [_, ws] of this.#actorWebSockets) {
-			ws.close();
+			ws.__closeWithRetry();
 		}
 		this.#actorWebSockets.clear();
 	}
@@ -76,12 +77,20 @@ export class Tunnel {
 		// Build message
 		const messageId = generateUuidBuffer();
 
-		const requestIdStr = bufferToString(requestId);
-		this.#pendingTunnelMessages.set(bufferToString(messageId), {
+		const requestIdStr = idToStr(requestId);
+		const messageIdStr = idToStr(messageId);
+		this.#pendingTunnelMessages.set(messageIdStr, {
 			sentAt: Date.now(),
 			requestIdStr,
 		});
 
+		logger()?.debug({
+			msg: "send tunnel msg",
+			requestId: requestIdStr,
+			messageId: messageIdStr,
+			message: messageKind,
+		});
+
 		// Send message
 		const message: protocol.ToServer = {
 			tag: "ToServerTunnelMessage",
@@ -108,6 +117,12 @@ export class Tunnel {
 			},
 		};
 
+		logger()?.debug({
+			msg: "ack tunnel msg",
+			requestId: idToStr(requestId),
+			messageId: idToStr(messageId),
+		});
+
 		this.#runner.__sendToServer(message);
 	}
 
@@ -156,7 +171,10 @@ export class Tunnel {
 				const webSocket = this.#actorWebSockets.get(requestIdStr);
 				if (webSocket) {
 					// Close the WebSocket connection
-					webSocket.close(1000, "Message acknowledgment timeout");
+					webSocket.__closeWithRetry(
+						1000,
+						"Message acknowledgment timeout",
+					);
 
 					// Clean up from actorWebSockets map
 					this.#actorWebSockets.delete(requestIdStr);
@@ -176,7 +194,7 @@ export class Tunnel {
 		}
 	}
 
-	unregisterActor(actor: ActorInstance) {
+	closeActiveRequests(actor: ActorInstance) {
 		const actorId = actor.actorId;
 
 		// Terminate all requests for this actor
@@ -189,30 +207,43 @@ export class Tunnel {
 		}
 		actor.requests.clear();
 
-		// Close all WebSockets for this actor
+		// Flush acks and close all WebSockets for this actor
 		for (const webSocketId of actor.webSockets) {
 			const ws = this.#actorWebSockets.get(webSocketId);
 			if (ws) {
-				ws.close(1000, "Actor stopped");
+				ws.__closeWithRetry(1000, "Actor stopped");
 				this.#actorWebSockets.delete(webSocketId);
 			}
 		}
 		actor.webSockets.clear();
 	}
 
-	async #fetch(actorId: string, request: Request): Promise<Response> {
+	async #fetch(
+		actorId: string,
+		requestId: protocol.RequestId,
+		request: Request,
+	): Promise<Response> {
 		// Validate actor exists
 		if (!this.#runner.hasActor(actorId)) {
 			logger()?.warn({
 				msg: "ignoring request for unknown actor",
 				actorId,
 			});
-			return new Response("Actor not found", { status: 404 });
+
+			// NOTE: This is a special response that will cause Guard to retry the request
+			//
+			// See should_retry_request_inner
+			// https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/guard-core/src/proxy_service.rs#L2458
+			return new Response("Actor not found", {
+				status: 503,
+				headers: { "x-rivet-error": "runner.actor_not_found" },
+			});
 		}
 
 		const fetchHandler = this.#runner.config.fetch(
 			this.#runner,
 			actorId,
+			requestId,
 			request,
 		);
 
@@ -224,44 +255,72 @@ export class Tunnel {
 	}
 
 	async handleTunnelMessage(message: protocol.ToClientTunnelMessage) {
+		const requestIdStr = idToStr(message.requestId);
+		const messageIdStr = idToStr(message.messageId);
+		logger()?.debug({
+			msg: "receive tunnel msg",
+			requestId: requestIdStr,
+			messageId: messageIdStr,
+			message: message.messageKind,
+		});
+
 		if (message.messageKind.tag === "TunnelAck") {
 			// Mark pending message as acknowledged and remove it
-			const msgIdStr = bufferToString(message.messageId);
-			const pending = this.#pendingTunnelMessages.get(msgIdStr);
+			const pending = this.#pendingTunnelMessages.get(messageIdStr);
 			if (pending) {
-				this.#pendingTunnelMessages.delete(msgIdStr);
+				const didDelete =
+					this.#pendingTunnelMessages.delete(messageIdStr);
+				if (!didDelete) {
+					logger()?.warn({
+						msg: "received tunnel ack for nonexistent message",
+						requestId: requestIdStr,
+						messageId: messageIdStr,
+					});
+				}
 			}
 		} else {
-			this.#sendAck(message.requestId, message.messageId);
 			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":
-					await this.#handleWebSocketMessage(
+				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,
@@ -278,7 +337,7 @@ export class Tunnel {
 		req: protocol.ToClientRequestStart,
 	) {
 		// Track this request for the actor
-		const requestIdStr = bufferToString(requestId);
+		const requestIdStr = idToStr(requestId);
 		const actor = this.#runner.getActor(req.actorId);
 		if (actor) {
 			actor.requests.add(requestIdStr);
@@ -329,12 +388,17 @@ export class Tunnel {
 				// Call fetch handler with validation
 				const response = await this.#fetch(
 					req.actorId,
+					requestId,
 					streamingRequest,
 				);
 				await this.#sendResponse(requestId, response);
 			} else {
 				// Non-streaming request
-				const response = await this.#fetch(req.actorId, request);
+				const response = await this.#fetch(
+					req.actorId,
+					requestId,
+					request,
+				);
 				await this.#sendResponse(requestId, response);
 			}
 		} catch (error) {
@@ -353,7 +417,7 @@ export class Tunnel {
 		requestId: ArrayBuffer,
 		chunk: protocol.ToClientRequestChunk,
 	) {
-		const requestIdStr = bufferToString(requestId);
+		const requestIdStr = idToStr(requestId);
 		const pending = this.#actorPendingRequests.get(requestIdStr);
 		if (pending?.streamController) {
 			pending.streamController.enqueue(new Uint8Array(chunk.body));
@@ -365,7 +429,7 @@ export class Tunnel {
 	}
 
 	async #handleRequestAbort(requestId: ArrayBuffer) {
-		const requestIdStr = bufferToString(requestId);
+		const requestIdStr = idToStr(requestId);
 		const pending = this.#actorPendingRequests.get(requestIdStr);
 		if (pending?.streamController) {
 			pending.streamController.error(new Error("Request aborted"));
@@ -426,10 +490,10 @@ export class Tunnel {
 	}
 
 	async #handleWebSocketOpen(
-		requestId: ArrayBuffer,
+		requestId: protocol.RequestId,
 		open: protocol.ToClientWebSocketOpen,
 	) {
-		const webSocketId = bufferToString(requestId);
+		const webSocketId = idToStr(requestId);
 		// Validate actor exists
 		const actor = this.#runner.getActor(open.actorId);
 		if (!actor) {
@@ -437,12 +501,18 @@ export class Tunnel {
 				msg: "ignoring websocket for unknown actor",
 				actorId: open.actorId,
 			});
-			// Send close immediately
+
+			// NOTE: Closing a WebSocket before open is equivalent to a Service
+			// Unavailable error and will cause Guard to retry the request
+			//
+			// See
+			// https://github.com/rivet-dev/rivet/blob/222dae87e3efccaffa2b503de40ecf8afd4e31eb/engine/packages/pegboard-gateway/src/lib.rs#L238
 			this.#sendMessage(requestId, {
 				tag: "ToServerWebSocketClose",
 				val: {
 					code: 1011,
 					reason: "Actor not found",
+					retry: false,
 				},
 			});
 			return;
@@ -460,6 +530,7 @@ export class Tunnel {
 				val: {
 					code: 1011,
 					reason: "Not Implemented",
+					retry: false,
 				},
 			});
 			return;
@@ -490,13 +561,14 @@ export class Tunnel {
 						},
 					});
 				},
-				(code?: number, reason?: string) => {
+				(code?: number, reason?: string, retry: boolean = false) => {
 					// Send close through tunnel
 					this.#sendMessage(requestId, {
 						tag: "ToServerWebSocketClose",
 						val: {
 							code: code || null,
 							reason: reason || null,
+							retry,
 						},
 					});
 
@@ -513,17 +585,10 @@ export class Tunnel {
 			// Store adapter
 			this.#actorWebSockets.set(webSocketId, adapter);
 
-			// Send open confirmation
-			this.#sendMessage(requestId, {
-				tag: "ToServerWebSocketOpen",
-				val: null,
-			});
-
-			// Notify adapter that connection is open
-			adapter._handleOpen();
-
-			// Create a minimal request object for the websocket handler
-			// Include original headers from the open message
+			// 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<
@@ -533,7 +598,6 @@ export class Tunnel {
 					headerInit[k] = v;
 				}
 			}
-			// Ensure websocket upgrade headers are present
 			headerInit["Upgrade"] = "websocket";
 			headerInit["Connection"] = "Upgrade";
 
@@ -542,11 +606,30 @@ export class Tunnel {
 				headers: headerInit,
 			});
 
+			// Send open confirmation
+			const hibernationConfig =
+				this.#runner.config.getActorHibernationConfig(
+					actor.actorId,
+					requestId,
+					request,
+				);
+			this.#sendMessage(requestId, {
+				tag: "ToServerWebSocketOpen",
+				val: {
+					canHibernate: hibernationConfig.enabled,
+					lastMsgIndex: BigInt(hibernationConfig.lastMsgIndex ?? -1),
+				},
+			});
+
+			// Notify adapter that connection is open
+			adapter._handleOpen(requestId);
+
 			// Call websocket handler
 			await websocketHandler(
 				this.#runner,
 				open.actorId,
 				adapter,
+				requestId,
 				request,
 			);
 		} catch (error) {
@@ -557,6 +640,7 @@ export class Tunnel {
 				val: {
 					code: 1011,
 					reason: "Server Error",
+					retry: false,
 				},
 			});
 
@@ -569,45 +653,72 @@ export class Tunnel {
 		}
 	}
 
+	/// Returns false if the message was sent off
 	async #handleWebSocketMessage(
 		requestId: ArrayBuffer,
-		msg: protocol.ToServerWebSocketMessage,
-	) {
-		const webSocketId = bufferToString(requestId);
+		msg: protocol.ToClientWebSocketMessage,
+	): Promise<boolean> {
+		const webSocketId = idToStr(requestId);
 		const adapter = this.#actorWebSockets.get(webSocketId);
 		if (adapter) {
 			const data = msg.binary
 				? new Uint8Array(msg.data)
 				: new TextDecoder().decode(new Uint8Array(msg.data));
 
-			adapter._handleMessage(data, msg.binary);
+			return adapter._handleMessage(
+				requestId,
+				data,
+				msg.index,
+				msg.binary,
+			);
+		} else {
+			return true;
 		}
 	}
 
+	__ackWebsocketMessage(requestId: ArrayBuffer, index: number) {
+		logger()?.debug({
+			msg: "ack ws msg",
+			requestId: idToStr(requestId),
+			index,
+		});
+
+		if (index < 0 || index > 65535)
+			throw new Error("invalid websocket ack index");
+
+		// Send the ack message
+		this.#sendMessage(requestId, {
+			tag: "ToServerWebSocketMessageAck",
+			val: {
+				index,
+			},
+		});
+	}
+
 	async #handleWebSocketClose(
 		requestId: ArrayBuffer,
-		close: protocol.ToServerWebSocketClose,
+		close: protocol.ToClientWebSocketClose,
 	) {
-		const webSocketId = bufferToString(requestId);
-		const adapter = this.#actorWebSockets.get(webSocketId);
+		const requestIdStr = idToStr(requestId);
+		const adapter = this.#actorWebSockets.get(requestIdStr);
 		if (adapter) {
 			adapter._handleClose(
+				requestId,
 				close.code || undefined,
 				close.reason || undefined,
 			);
-			this.#actorWebSockets.delete(webSocketId);
+			this.#actorWebSockets.delete(requestIdStr);
 		}
 	}
 }
 
-/** Converts a buffer to a string. Used for storing strings in a lookup map. */
-function bufferToString(buffer: ArrayBuffer): string {
-	return Buffer.from(buffer).toString("base64");
-}
-
 /** Generates a UUID as bytes. */
 function generateUuidBuffer(): ArrayBuffer {
 	const buffer = new Uint8Array(16);
 	uuidv4(undefined, buffer);
 	return buffer.buffer;
 }
+
+function idToStr(id: ArrayBuffer): string {
+	return uuidstringify(new Uint8Array(id));
+}