@rivetkit/engine-runner 2.0.4-rc.1

package/src/utils.ts

@@ -0,0 +1,175 @@
+import { logger } from "./log";
+
+export function unreachable(x: never): never {
+	throw `Unreachable: ${x}`;
+}
+
+export interface BackoffOptions {
+	initialDelay?: number;
+	maxDelay?: number;
+	multiplier?: number;
+	jitter?: boolean;
+}
+
+export function calculateBackoff(
+	attempt: number,
+	options: BackoffOptions = {},
+): number {
+	const {
+		initialDelay = 1000,
+		maxDelay = 30000,
+		multiplier = 2,
+		jitter = true,
+	} = options;
+
+	let delay = Math.min(initialDelay * multiplier ** attempt, maxDelay);
+
+	if (jitter) {
+		// Add random jitter between 0% and 25% of the delay
+		delay = delay * (1 + Math.random() * 0.25);
+	}
+
+	return Math.floor(delay);
+}
+
+export interface ParsedCloseReason {
+	group: string;
+	error: string;
+	rayId?: string;
+}
+
+/**
+ * Parses a WebSocket close reason in the format: {group}.{error} or {group}.{error}#{ray_id}
+ *
+ * Examples:
+ *   - "ws.eviction#t1s80so6h3irenp8ymzltfoittcl00"
+ *   - "ws.client_closed"
+ *
+ * Returns undefined if the format is invalid
+ */
+export function parseWebSocketCloseReason(
+	reason: string,
+): ParsedCloseReason | undefined {
+	const [mainPart, rayId] = reason.split("#");
+	const [group, error] = mainPart.split(".");
+
+	if (!group || !error) {
+		logger()?.warn({ msg: "failed to parse close reason", reason });
+		return undefined;
+	}
+
+	return {
+		group,
+		error,
+		rayId,
+	};
+}
+
+const U16_MAX = 65535;
+
+/**
+ * Wrapping greater than comparison for u16 values.
+ * Based on shared_state.rs wrapping_gt implementation.
+ */
+export function wrappingGtU16(a: number, b: number): boolean {
+	return a !== b && wrappingSub(a, b, U16_MAX) < U16_MAX / 2;
+}
+
+/**
+ * Wrapping less than comparison for u16 values.
+ * Based on shared_state.rs wrapping_lt implementation.
+ */
+export function wrappingLtU16(a: number, b: number): boolean {
+	return a !== b && wrappingSub(b, a, U16_MAX) < U16_MAX / 2;
+}
+
+/**
+ * Wrapping greater than or equal comparison for u16 values.
+ */
+export function wrappingGteU16(a: number, b: number): boolean {
+	return a === b || wrappingGtU16(a, b);
+}
+
+/**
+ * Wrapping less than or equal comparison for u16 values.
+ */
+export function wrappingLteU16(a: number, b: number): boolean {
+	return a === b || wrappingLtU16(a, b);
+}
+
+/**
+ * Performs wrapping addition for u16 values.
+ */
+export function wrappingAddU16(a: number, b: number): number {
+	return (a + b) % (U16_MAX + 1);
+}
+
+/**
+ * Performs wrapping subtraction for u16 values.
+ */
+export function wrappingSubU16(a: number, b: number): number {
+	return wrappingSub(a, b, U16_MAX);
+}
+
+/**
+ * Performs wrapping subtraction for unsigned integers.
+ */
+function wrappingSub(a: number, b: number, max: number): number {
+	const result = a - b;
+	if (result < 0) {
+		return result + max + 1;
+	}
+	return result;
+}
+
+export function arraysEqual(a: ArrayBuffer, b: ArrayBuffer): boolean {
+	const ua = new Uint8Array(a);
+	const ub = new Uint8Array(b);
+	if (ua.length !== ub.length) return false;
+	for (let i = 0; i < ua.length; i++) {
+		if (ua[i] !== ub[i]) return false;
+	}
+	return true;
+}
+
+/**
+ * Polyfill for Promise.withResolvers().
+ *
+ * This is specifically for Cloudflare Workers. Their implementation of Promise.withResolvers does not work correctly.
+ */
+export function promiseWithResolvers<T>(): {
+	promise: Promise<T>;
+	resolve: (value: T | PromiseLike<T>) => void;
+	reject: (reason?: any) => void;
+} {
+	let resolve!: (value: T | PromiseLike<T>) => void;
+	let reject!: (reason?: any) => void;
+	const promise = new Promise<T>((res, rej) => {
+		resolve = res;
+		reject = rej;
+	});
+	return { promise, resolve, reject };
+}
+
+export function idToStr(id: ArrayBuffer): string {
+	const bytes = new Uint8Array(id);
+	return Array.from(bytes)
+		.map((byte) => byte.toString(16).padStart(2, "0"))
+		.join("");
+}
+
+export function stringifyError(error: unknown): string {
+	if (error instanceof Error) {
+		return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ""}`;
+	} else if (typeof error === "string") {
+		return error;
+	} else if (typeof error === "object" && error !== null) {
+		try {
+			return `${JSON.stringify(error)}`;
+		} catch {
+			return `[object ${error.constructor?.name || "Object"}]`;
+		}
+	} else {
+		return String(error);
+	}
+}

package/src/websocket-tunnel-adapter.ts

@@ -0,0 +1,201 @@
+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";
+
+export const HIBERNATABLE_SYMBOL = Symbol("hibernatable");
+
+export class WebSocketTunnelAdapter {
+	#readyState: 0 | 1 | 2 | 3 = 0;
+	#binaryType: "nodebuffer" | "arraybuffer" | "blob" = "nodebuffer";
+	#ws: VirtualWebSocket;
+	#tunnel: Tunnel;
+	#actorId: string;
+	#requestId: string;
+	#hibernatable: boolean;
+	#serverMessageIndex: number;
+	#sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void;
+	#closeCallback: (code?: number, reason?: string) => void;
+
+	get [HIBERNATABLE_SYMBOL](): boolean {
+		return this.#hibernatable;
+	}
+
+	get #log(): Logger | undefined {
+		return this.#tunnel.log;
+	}
+
+	constructor(
+		tunnel: Tunnel,
+		actorId: string,
+		requestId: string,
+		serverMessageIndex: number,
+		hibernatable: boolean,
+		isRestoringHibernatable: boolean,
+		public readonly request: Request,
+		sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void,
+		closeCallback: (code?: number, reason?: string) => void,
+	) {
+		this.#tunnel = tunnel;
+		this.#actorId = actorId;
+		this.#requestId = requestId;
+		this.#hibernatable = hibernatable;
+		this.#serverMessageIndex = serverMessageIndex;
+		this.#sendCallback = sendCallback;
+		this.#closeCallback = closeCallback;
+
+		this.#ws = new VirtualWebSocket({
+			getReadyState: () => this.#readyState,
+			onSend: (data) => this.#handleSend(data),
+			onClose: (code, reason) => this.#close(code, reason, true),
+			onTerminate: () => this.#terminate(),
+		});
+
+		if (isRestoringHibernatable) {
+			this.#log?.debug({
+				msg: "setting WebSocket to OPEN state for restored connection",
+				actorId: this.#actorId,
+				requestId: this.#requestId,
+			});
+			this.#readyState = 1;
+		}
+	}
+
+	get websocket(): UniversalWebSocket {
+		return this.#ws;
+	}
+
+	#handleSend(data: string | ArrayBufferLike | Blob | ArrayBufferView): void {
+		let isBinary = false;
+		let messageData: string | ArrayBuffer;
+
+		if (typeof data === "string") {
+			messageData = data;
+		} else if (data instanceof ArrayBuffer) {
+			isBinary = true;
+			messageData = data;
+		} else if (ArrayBuffer.isView(data)) {
+			isBinary = true;
+			const view = data;
+			const buffer = view.buffer instanceof SharedArrayBuffer
+				? new Uint8Array(view.buffer, view.byteOffset, view.byteLength).slice().buffer
+				: view.buffer.slice(view.byteOffset, view.byteOffset + view.byteLength);
+			messageData = buffer as ArrayBuffer;
+		} else {
+			throw new Error("Unsupported data type");
+		}
+
+		this.#sendCallback(messageData, isBinary);
+	}
+
+	// Called by Tunnel when WebSocket is opened
+	_handleOpen(requestId: ArrayBuffer): void {
+		if (this.#readyState !== 0) return;
+		this.#readyState = 1;
+		this.#ws.dispatchEvent({ type: "open", rivetRequestId: requestId, target: this.#ws });
+	}
+
+	// Called by Tunnel when message is received
+	_handleMessage(
+		requestId: ArrayBuffer,
+		data: string | Uint8Array,
+		serverMessageIndex: number,
+		isBinary: boolean,
+	): boolean {
+		if (this.#readyState !== 1) {
+			this.#log?.warn({
+				msg: "WebSocket message ignored - not in OPEN state",
+				requestId: this.#requestId,
+				actorId: this.#actorId,
+				currentReadyState: this.#readyState,
+			});
+			return true;
+		}
+
+		// Validate message index for hibernatable websockets
+		if (this.#hibernatable) {
+			const previousIndex = this.#serverMessageIndex;
+
+			if (wrappingLteU16(serverMessageIndex, previousIndex)) {
+				this.#log?.info({
+					msg: "received duplicate hibernating websocket message",
+					requestId,
+					actorId: this.#actorId,
+					previousIndex,
+					receivedIndex: serverMessageIndex,
+				});
+				return true;
+			}
+
+			const expectedIndex = wrappingAddU16(previousIndex, 1);
+			if (serverMessageIndex !== expectedIndex) {
+				const closeReason = "ws.message_index_skip";
+				this.#log?.warn({
+					msg: "hibernatable websocket message index out of sequence, closing connection",
+					requestId,
+					actorId: this.#actorId,
+					previousIndex,
+					expectedIndex,
+					receivedIndex: serverMessageIndex,
+					closeReason,
+					gap: wrappingSubU16(wrappingSubU16(serverMessageIndex, previousIndex), 1),
+				});
+				this.#close(1008, closeReason, true);
+				return true;
+			}
+
+			this.#serverMessageIndex = serverMessageIndex;
+		}
+
+		// Convert data based on binaryType
+		let messageData: any = data;
+		if (isBinary && data instanceof Uint8Array) {
+			if (this.#binaryType === "nodebuffer") {
+				messageData = Buffer.from(data);
+			} else if (this.#binaryType === "arraybuffer") {
+				messageData = data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength);
+			}
+		}
+
+		this.#ws.dispatchEvent({
+			type: "message",
+			data: messageData,
+			rivetRequestId: requestId,
+			rivetMessageIndex: serverMessageIndex,
+			target: this.#ws,
+		} as RivetMessageEvent);
+
+		return false;
+	}
+
+	// Called by Tunnel when close is received
+	_handleClose(_requestId: ArrayBuffer, code?: number, reason?: string): void {
+		this.#close(code, reason, true);
+	}
+
+	// Close without sending close message to tunnel
+	_closeWithoutCallback(code?: number, reason?: string): void {
+		this.#close(code, reason, false);
+	}
+
+	// Public close method (used by tunnel.ts for stale websocket cleanup)
+	close(code?: number, reason?: string): void {
+		this.#close(code, reason, true);
+	}
+
+	#close(code: number | undefined, reason: string | undefined, sendCallback: boolean): void {
+		if (this.#readyState >= 2) return;
+
+		this.#readyState = 2;
+		if (sendCallback) this.#closeCallback(code, reason);
+		this.#readyState = 3;
+		this.#ws.triggerClose(code ?? 1000, reason ?? "");
+	}
+
+	#terminate(): void {
+		// Immediate close without close frame
+		this.#readyState = 3;
+		this.#closeCallback(1006, "Abnormal Closure");
+		this.#ws.triggerClose(1006, "Abnormal Closure", false);
+	}
+}

package/src/websocket.ts

@@ -0,0 +1,43 @@
+import { logger } from "./log";
+
+// Global singleton promise that will be reused for subsequent calls
+let webSocketPromise: Promise<typeof WebSocket> | null = null;
+
+export async function importWebSocket(): Promise<typeof WebSocket> {
+	// Return existing promise if we already started loading
+	if (webSocketPromise !== null) {
+		return webSocketPromise;
+	}
+
+	// Create and store the promise
+	webSocketPromise = (async () => {
+		let _WebSocket: typeof WebSocket;
+
+		if (typeof WebSocket !== "undefined") {
+			// Native
+			_WebSocket = WebSocket as unknown as typeof WebSocket;
+			logger()?.debug({ msg: "using native websocket" });
+		} else {
+			// Node.js package
+			try {
+				const ws = await import("ws");
+				_WebSocket = ws.default as unknown as typeof WebSocket;
+				logger()?.debug({ msg: "using websocket from npm" });
+			} catch {
+				// WS not available
+				_WebSocket = class MockWebSocket {
+					constructor() {
+						throw new Error(
+							'WebSocket support requires installing the "ws" peer dependency.',
+						);
+					}
+				} as unknown as typeof WebSocket;
+				logger()?.debug({ msg: "using mock websocket" });
+			}
+		}
+
+		return _WebSocket;
+	})();
+
+	return webSocketPromise;
+}