@rivetkit/engine-runner 2.1.3 → 2.1.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rivetkit/engine-runner",
-  "version": "2.1.3",
+  "version": "2.1.4",
   "type": "module",
   "files": [
     "dist",
@@ -22,7 +22,7 @@
     "pino": "^9.9.5",
     "ws": "^8.18.3",
     "@rivetkit/virtual-websocket": "2.0.33",
-    "@rivetkit/engine-runner-protocol": "2.1.3"
+    "@rivetkit/engine-runner-protocol": "2.1.4"
   },
   "devDependencies": {
     "@types/node": "^22.18.1",

package/src/mod.ts

@@ -208,8 +208,10 @@ export class Runner {
 	#reconnectAttempt: number = 0;
 	#reconnectTimeout?: NodeJS.Timeout;
 
+	// Protocol metadata
+	#protocolMetadata?: protocol.ProtocolMetadata;
+
 	// Runner lost threshold management
-	#runnerLostThreshold?: number;
 	#runnerLostTimeout?: NodeJS.Timeout;
 
 	// Event storage for resending
@@ -839,14 +841,11 @@ export class Runner {
 					this.#stopAllActors();
 				}
 
-				// Store the runner lost threshold from metadata
-				this.#runnerLostThreshold = init.metadata?.runnerLostThreshold
-					? Number(init.metadata.runnerLostThreshold)
-					: undefined;
+				this.#protocolMetadata = init.metadata;
 
 				this.log?.info({
 					msg: "received init",
-					runnerLostThreshold: this.#runnerLostThreshold,
+					protocolMetadata: this.#protocolMetadata,
 				});
 
 				// Resend pending events
@@ -888,27 +887,7 @@ export class Runner {
 			});
 
 			if (!this.#shutdown) {
-				// Start runner lost timeout if we have a threshold and are not shutting down
-				if (
-					!this.#runnerLostTimeout &&
-					this.#runnerLostThreshold &&
-					this.#runnerLostThreshold > 0
-				) {
-					this.log?.info({
-						msg: "starting runner lost timeout",
-						seconds: this.#runnerLostThreshold / 1000,
-					});
-					this.#runnerLostTimeout = setTimeout(() => {
-						try {
-							this.#handleLost();
-						} catch (err) {
-							this.log?.error({
-								msg: "error handling runner lost",
-								error: stringifyError(err),
-							});
-						}
-					}, this.#runnerLostThreshold);
-				}
+				this.#startRunnerLostTimeout();
 
 				// Attempt to reconnect if not stopped
 				this.#scheduleReconnect();
@@ -944,27 +923,7 @@ export class Runner {
 					this.#ackInterval = undefined;
 				}
 
-				// Start runner lost timeout if we have a threshold and are not shutting down
-				if (
-					!this.#runnerLostTimeout &&
-					this.#runnerLostThreshold &&
-					this.#runnerLostThreshold > 0
-				) {
-					this.log?.info({
-						msg: "starting runner lost timeout",
-						seconds: this.#runnerLostThreshold / 1000,
-					});
-					this.#runnerLostTimeout = setTimeout(() => {
-						try {
-							this.#handleLost();
-						} catch (err) {
-							this.log?.error({
-								msg: "error handling runner lost",
-								error: stringifyError(err),
-							});
-						}
-					}, this.#runnerLostThreshold);
-				}
+				this.#startRunnerLostTimeout();
 
 				// Attempt to reconnect if not stopped
 				this.#scheduleReconnect();
@@ -976,6 +935,30 @@ export class Runner {
 		});
 	}
 
+	#startRunnerLostTimeout() {
+		// Start runner lost timeout if we have a threshold and are not shutting down
+		if (
+			!this.#runnerLostTimeout &&
+			this.#protocolMetadata &&
+			this.#protocolMetadata.runnerLostThreshold > 0
+		) {
+			this.log?.info({
+				msg: "starting runner lost timeout",
+				seconds: this.#protocolMetadata.runnerLostThreshold / 1000n,
+			});
+			this.#runnerLostTimeout = setTimeout(() => {
+				try {
+					this.#handleLost();
+				} catch (err) {
+					this.log?.error({
+						msg: "error handling runner lost",
+						error: stringifyError(err),
+					});
+				}
+			}, Number(this.#protocolMetadata.runnerLostThreshold));
+		}
+	}
+
 	#handleCommands(commands: protocol.ToClientCommands) {
 		this.log?.info({
 			msg: "received commands",
@@ -1859,4 +1842,8 @@ export class Runner {
 			//this.#log?.log(`Cleaned up ${toDelete.length} expired KV requests`);
 		}
 	}
+
+	getProtocolMetadata(): protocol.ProtocolMetadata | undefined {
+		return this.#protocolMetadata;
+	}
 }

package/src/tunnel.ts

@@ -15,7 +15,7 @@ import {
 	stringifyToClientTunnelMessageKind,
 	stringifyToServerTunnelMessageKind,
 } from "./stringify";
-import { arraysEqual, idToStr, stringifyError, unreachable } from "./utils";
+import { arraysEqual, idToStr, MAX_PAYLOAD_SIZE, stringifyError, unreachable } from "./utils";
 import {
 	HIBERNATABLE_SYMBOL,
 	WebSocketTunnelAdapter,
@@ -855,6 +855,10 @@ export class Tunnel {
 		// Read the body first to get the actual content
 		const body = response.body ? await response.arrayBuffer() : null;
 
+		if (body && body.byteLength > MAX_PAYLOAD_SIZE) {
+			throw new Error("Response body too large");
+		}
+
 		// Convert headers to map and add Content-Length if not present
 		const headers = new Map<string, string>();
 		response.headers.forEach((value, key) => {
@@ -1079,7 +1083,7 @@ export class Tunnel {
 		});
 
 		if (clientMessageIndex < 0 || clientMessageIndex > 65535)
-			throw new Error("invalid websocket ack index");
+			throw new Error("Invalid websocket ack index");
 
 		// Get the actor to find the gatewayId
 		//
@@ -1157,7 +1161,7 @@ function buildRequestForWebSocket(
 	};
 
 	if (!path.startsWith("/")) {
-		throw new Error("path must start with leading slash");
+		throw new Error("Path must start with leading slash");
 	}
 
 	const request = new Request(`http://actor${path}`, {

package/src/utils.ts

@@ -1,5 +1,8 @@
 import { logger } from "./log";
 
+// 20MiB. Keep in sync with runner_max_response_payload_body_size from engine/packages/config/src/config/pegboard.rs
+export const MAX_PAYLOAD_SIZE = 20 * 1024 * 1024;
+
 export function unreachable(x: never): never {
 	throw `Unreachable: ${x}`;
 }

package/src/websocket-tunnel-adapter.ts

@@ -1,7 +1,7 @@
 import type { Logger } from "pino";
 import { VirtualWebSocket, type UniversalWebSocket, type RivetMessageEvent } from "@rivetkit/virtual-websocket";
 import type { Tunnel } from "./tunnel";
-import { wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils";
+import { MAX_PAYLOAD_SIZE, wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils";
 
 export const HIBERNATABLE_SYMBOL = Symbol("hibernatable");
 
@@ -70,11 +70,20 @@ export class WebSocketTunnelAdapter {
 		let messageData: string | ArrayBuffer;
 
 		if (typeof data === "string") {
+			const encoder = new TextEncoder();
+			if (encoder.encode(data).byteLength > MAX_PAYLOAD_SIZE) {
+				throw new Error("WebSocket message too large");
+			}
+
 			messageData = data;
 		} else if (data instanceof ArrayBuffer) {
+			if (data.byteLength > MAX_PAYLOAD_SIZE) throw new Error("WebSocket message too large");
+
 			isBinary = true;
 			messageData = data;
 		} else if (ArrayBuffer.isView(data)) {
+			if (data.byteLength > MAX_PAYLOAD_SIZE) throw new Error("WebSocket message too large");
+
 			isBinary = true;
 			const view = data;
 			const buffer = view.buffer instanceof SharedArrayBuffer