@rivetkit/engine 1.0.0

package/sdks/typescript/runner/src/utils.ts

@@ -0,0 +1,159 @@
+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("");
+}

package/sdks/typescript/runner/src/websocket-tunnel-adapter.ts

@@ -0,0 +1,567 @@
+// WebSocket-like adapter for tunneled connections
+// Implements a subset of the WebSocket interface for compatibility with runner code
+
+import type { Logger } from "pino";
+import { logger } from "./log";
+import type { Tunnel } from "./tunnel";
+import { wrappingAddU16, wrappingLteU16, wrappingSubU16 } from "./utils";
+
+export const HIBERNATABLE_SYMBOL = Symbol("hibernatable");
+
+export class WebSocketTunnelAdapter {
+	// MARK: - WebSocket Compat Variables
+	#readyState: number = 0; // CONNECTING
+	#eventListeners: Map<string, Set<(event: any) => void>> = new Map();
+	#onopen: ((this: any, ev: any) => any) | null = null;
+	#onclose: ((this: any, ev: any) => any) | null = null;
+	#onerror: ((this: any, ev: any) => any) | null = null;
+	#onmessage: ((this: any, ev: any) => any) | null = null;
+	#bufferedAmount = 0;
+	#binaryType: "nodebuffer" | "arraybuffer" | "blob" = "nodebuffer";
+	#extensions = "";
+	#protocol = "";
+	#url = "";
+
+	// mARK: - Internal State
+	#tunnel: Tunnel;
+	#actorId: string;
+	#requestId: string;
+	#hibernatable: boolean;
+	#serverMessageIndex: number;
+
+	get [HIBERNATABLE_SYMBOL](): boolean {
+		return this.#hibernatable;
+	}
+
+	/**
+	 * Called when sending a message from this WebSocket.
+	 *
+	 * Used to send a tunnel message from Tunnel.
+	 */
+	#sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void;
+
+	/**
+	 * Called when closing this WebSocket.
+	 *
+	 * Used to send a tunnel message from Tunnel
+	 */
+	#closeCallback: (
+		code?: number,
+		reason?: string,
+		hibernate?: boolean,
+	) => void;
+
+	get #log(): Logger | undefined {
+		return this.#tunnel.log;
+	}
+
+	constructor(
+		tunnel: Tunnel,
+		actorId: string,
+		requestId: string,
+		serverMessageIndex: number,
+		hibernatable: boolean,
+		isRestoringHibernatable: boolean,
+		/** @experimental */
+		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;
+
+		// For restored WebSockets, immediately set to OPEN state
+		if (isRestoringHibernatable) {
+			this.#log?.debug({
+				msg: "setting WebSocket to OPEN state for restored connection",
+				actorId: this.#actorId,
+				requestId: this.#requestId,
+				hibernatable: this.#hibernatable,
+			});
+			this.#readyState = 1; // OPEN
+		}
+	}
+
+	// MARK: - Lifecycle
+	get bufferedAmount(): number {
+		return this.#bufferedAmount;
+	}
+
+	_handleOpen(requestId: ArrayBuffer): void {
+		if (this.#readyState !== 0) {
+			// CONNECTING
+			return;
+		}
+
+		this.#readyState = 1; // OPEN
+
+		const event = {
+			type: "open",
+			rivetRequestId: requestId,
+			target: this,
+		};
+
+		this.#fireEvent("open", event);
+	}
+
+	_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,
+				expectedReadyState: 1,
+				serverMessageIndex: serverMessageIndex,
+				hibernatable: this.#hibernatable,
+			});
+			return true;
+		}
+
+		// Validate message index
+		if (this.#hibernatable) {
+			const previousIndex = this.#serverMessageIndex;
+
+			// Ignore duplicate old messages
+			//
+			// This should only happen if something goes wrong
+			// between persisting the previous index and acking the
+			// message index to the gateway. If the ack is never
+			// received by the gateway (due to a crash or network
+			// issue), the gateway will resend all messages from
+			// the last ack on reconnect.
+			if (wrappingLteU16(serverMessageIndex, previousIndex)) {
+				this.#log?.info({
+					msg: "received duplicate hibernating websocket message, this indicates the actor failed to ack the message index before restarting",
+					requestId,
+					actorId: this.#actorId,
+					previousIndex,
+					expectedIndex: wrappingAddU16(previousIndex, 1),
+					receivedIndex: serverMessageIndex,
+				});
+
+				return true;
+			}
+
+			// Close message if skipped message in sequence
+			//
+			// There is no scenario where this should ever happen
+			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,
+					),
+				});
+
+				// Close the WebSocket and skip processing
+				this.close(1008, closeReason);
+
+				return true;
+			}
+
+			// Update to the next index
+			this.#serverMessageIndex = serverMessageIndex;
+		}
+
+		// Dispatch event
+		let messageData: any;
+		if (isBinary) {
+			// Handle binary data based on binaryType
+			if (this.#binaryType === "nodebuffer") {
+				// Convert to Buffer for Node.js compatibility
+				messageData = Buffer.from(data as Uint8Array);
+			} else if (this.#binaryType === "arraybuffer") {
+				// Convert to ArrayBuffer
+				if (data instanceof Uint8Array) {
+					messageData = data.buffer.slice(
+						data.byteOffset,
+						data.byteOffset + data.byteLength,
+					);
+				} else {
+					messageData = data;
+				}
+			} else {
+				// Blob type - not commonly used in Node.js
+				throw new Error(
+					"Blob binaryType not supported in tunnel adapter",
+				);
+			}
+		} else {
+			messageData = data;
+		}
+
+		const event = {
+			type: "message",
+			data: messageData,
+			rivetRequestId: requestId,
+			rivetMessageIndex: serverMessageIndex,
+			target: this,
+		};
+
+		this.#fireEvent("message", event);
+
+		return false;
+	}
+
+	_handleClose(
+		_requestId: ArrayBuffer,
+		code?: number,
+		reason?: string,
+	): void {
+		this.#closeInner(code, reason, true);
+	}
+
+	_handleError(error: Error): void {
+		const event = {
+			type: "error",
+			target: this,
+			error,
+		};
+
+		this.#fireEvent("error", event);
+	}
+
+	_closeWithoutCallback(code?: number, reason?: string): void {
+		this.#closeInner(code, reason, false);
+	}
+
+	#fireEvent(type: string, event: any): void {
+		// Call all registered event listeners
+		const listeners = this.#eventListeners.get(type);
+
+		if (listeners && listeners.size > 0) {
+			for (const listener of listeners) {
+				try {
+					listener.call(this, event);
+				} catch (error) {
+					logger()?.error({
+						msg: "error in websocket event listener",
+						error,
+						type,
+					});
+				}
+			}
+		}
+
+		// Call the onX property if set
+		switch (type) {
+			case "open":
+				if (this.#onopen) {
+					try {
+						this.#onopen.call(this, event);
+					} catch (error) {
+						logger()?.error({
+							msg: "error in onopen handler",
+							error,
+						});
+					}
+				}
+				break;
+			case "close":
+				if (this.#onclose) {
+					try {
+						this.#onclose.call(this, event);
+					} catch (error) {
+						logger()?.error({
+							msg: "error in onclose handler",
+							error,
+						});
+					}
+				}
+				break;
+			case "error":
+				if (this.#onerror) {
+					try {
+						this.#onerror.call(this, event);
+					} catch (error) {
+						logger()?.error({
+							msg: "error in onerror handler",
+							error,
+						});
+					}
+				}
+				break;
+			case "message":
+				if (this.#onmessage) {
+					try {
+						this.#onmessage.call(this, event);
+					} catch (error) {
+						logger()?.error({
+							msg: "error in onmessage handler",
+							error,
+						});
+					}
+				}
+				break;
+		}
+	}
+
+	#closeInner(
+		code: number | undefined,
+		reason: string | undefined,
+		callback: boolean,
+	): void {
+		if (
+			this.#readyState === 2 || // CLOSING
+			this.#readyState === 3 // CLOSED
+		) {
+			return;
+		}
+
+		this.#readyState = 2; // CLOSING
+
+		// Send close through tunnel
+		if (callback) {
+			this.#closeCallback(code, reason);
+		}
+
+		// Update state and fire event
+		this.#readyState = 3; // CLOSED
+
+		const closeEvent = {
+			wasClean: true,
+			code: code || 1000,
+			reason: reason || "",
+			type: "close",
+			target: this,
+		};
+
+		this.#fireEvent("close", closeEvent);
+	}
+
+	// MARK: - WebSocket Compatible API
+	get readyState(): number {
+		return this.#readyState;
+	}
+
+	get binaryType(): string {
+		return this.#binaryType;
+	}
+
+	set binaryType(value: string) {
+		if (
+			value === "nodebuffer" ||
+			value === "arraybuffer" ||
+			value === "blob"
+		) {
+			this.#binaryType = value;
+		}
+	}
+
+	get extensions(): string {
+		return this.#extensions;
+	}
+
+	get protocol(): string {
+		return this.#protocol;
+	}
+
+	get url(): string {
+		return this.#url;
+	}
+
+	get onopen(): ((this: any, ev: any) => any) | null {
+		return this.#onopen;
+	}
+
+	set onopen(value: ((this: any, ev: any) => any) | null) {
+		this.#onopen = value;
+	}
+
+	get onclose(): ((this: any, ev: any) => any) | null {
+		return this.#onclose;
+	}
+
+	set onclose(value: ((this: any, ev: any) => any) | null) {
+		this.#onclose = value;
+	}
+
+	get onerror(): ((this: any, ev: any) => any) | null {
+		return this.#onerror;
+	}
+
+	set onerror(value: ((this: any, ev: any) => any) | null) {
+		this.#onerror = value;
+	}
+
+	get onmessage(): ((this: any, ev: any) => any) | null {
+		return this.#onmessage;
+	}
+
+	set onmessage(value: ((this: any, ev: any) => any) | null) {
+		this.#onmessage = value;
+	}
+
+	send(data: string | ArrayBuffer | ArrayBufferView | Blob | Buffer): void {
+		// Handle different ready states
+		if (this.#readyState === 0) {
+			// CONNECTING
+			throw new DOMException(
+				"WebSocket is still in CONNECTING state",
+				"InvalidStateError",
+			);
+		}
+
+		if (this.#readyState === 2 || this.#readyState === 3) {
+			// CLOSING or CLOSED - silently ignore
+			return;
+		}
+
+		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;
+			// Convert ArrayBufferView to ArrayBuffer
+			const view = data as ArrayBufferView;
+			// Check if it's a SharedArrayBuffer
+			if (view.buffer instanceof SharedArrayBuffer) {
+				// Copy SharedArrayBuffer to regular ArrayBuffer
+				const bytes = new Uint8Array(
+					view.buffer,
+					view.byteOffset,
+					view.byteLength,
+				);
+				messageData = bytes.buffer.slice(
+					bytes.byteOffset,
+					bytes.byteOffset + bytes.byteLength,
+				) as unknown as ArrayBuffer;
+			} else {
+				messageData = view.buffer.slice(
+					view.byteOffset,
+					view.byteOffset + view.byteLength,
+				) as ArrayBuffer;
+			}
+		} else if (data instanceof Blob) {
+			throw new Error("Blob sending not implemented in tunnel adapter");
+		} else if (typeof Buffer !== "undefined" && Buffer.isBuffer(data)) {
+			isBinary = true;
+			// Convert Buffer to ArrayBuffer
+			const buf = data as Buffer;
+			// Check if it's a SharedArrayBuffer
+			if (buf.buffer instanceof SharedArrayBuffer) {
+				// Copy SharedArrayBuffer to regular ArrayBuffer
+				const bytes = new Uint8Array(
+					buf.buffer,
+					buf.byteOffset,
+					buf.byteLength,
+				);
+				messageData = bytes.buffer.slice(
+					bytes.byteOffset,
+					bytes.byteOffset + bytes.byteLength,
+				) as unknown as ArrayBuffer;
+			} else {
+				messageData = buf.buffer.slice(
+					buf.byteOffset,
+					buf.byteOffset + buf.byteLength,
+				) as ArrayBuffer;
+			}
+		} else {
+			throw new Error("Invalid data type");
+		}
+
+		// Send through tunnel
+		this.#sendCallback(messageData, isBinary);
+	}
+
+	close(code?: number, reason?: string): void {
+		this.#closeInner(code, reason, true);
+	}
+
+	addEventListener(
+		type: string,
+		listener: (event: any) => void,
+		options?: boolean | any,
+	): void {
+		if (typeof listener === "function") {
+			let listeners = this.#eventListeners.get(type);
+			if (!listeners) {
+				listeners = new Set();
+				this.#eventListeners.set(type, listeners);
+			}
+			listeners.add(listener);
+		}
+	}
+
+	removeEventListener(
+		type: string,
+		listener: (event: any) => void,
+		options?: boolean | any,
+	): void {
+		if (typeof listener === "function") {
+			const listeners = this.#eventListeners.get(type);
+			if (listeners) {
+				listeners.delete(listener);
+			}
+		}
+	}
+
+	dispatchEvent(event: any): boolean {
+		// TODO:
+		return true;
+	}
+
+	static readonly CONNECTING = 0;
+	static readonly OPEN = 1;
+	static readonly CLOSING = 2;
+	static readonly CLOSED = 3;
+
+	readonly CONNECTING = 0;
+	readonly OPEN = 1;
+	readonly CLOSING = 2;
+	readonly CLOSED = 3;
+
+	// Additional methods for compatibility
+	ping(data?: any, mask?: boolean, cb?: (err: Error) => void): void {
+		// Not implemented for tunnel - could be added if needed
+		if (cb) cb(new Error("Ping not supported in tunnel adapter"));
+	}
+
+	pong(data?: any, mask?: boolean, cb?: (err: Error) => void): void {
+		// Not implemented for tunnel - could be added if needed
+		if (cb) cb(new Error("Pong not supported in tunnel adapter"));
+	}
+
+	/** @experimental */
+	terminate(): void {
+		// Immediate close without close frame
+		this.#readyState = 3; // CLOSED
+		this.#closeCallback(1006, "Abnormal Closure");
+
+		const event = {
+			wasClean: false,
+			code: 1006,
+			reason: "Abnormal Closure",
+			type: "close",
+			target: this,
+		};
+
+		this.#fireEvent("close", event);
+	}
+}