@rivetkit/engine-runner 2.0.33 → 2.0.34-rc.2

package/src/websocket-tunnel-adapter.ts

@@ -1,56 +1,26 @@
-// 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 { 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 {
-	// 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;
+	#readyState: 0 | 1 | 2 | 3 = 0;
 	#binaryType: "nodebuffer" | "arraybuffer" | "blob" = "nodebuffer";
-	#extensions = "";
-	#protocol = "";
-	#url = "";
-
-	// mARK: - Internal State
+	#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;
 	}
 
-	/**
-	 * 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;
 	}
@@ -62,7 +32,6 @@ export class WebSocketTunnelAdapter {
 		serverMessageIndex: number,
 		hibernatable: boolean,
 		isRestoringHibernatable: boolean,
-		/** @experimental */
 		public readonly request: Request,
 		sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void,
 		closeCallback: (code?: number, reason?: string) => void,
@@ -75,40 +44,58 @@ export class WebSocketTunnelAdapter {
 		this.#sendCallback = sendCallback;
 		this.#closeCallback = closeCallback;
 
-		// For restored WebSockets, immediately set to OPEN state
+		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,
-				hibernatable: this.#hibernatable,
 			});
-			this.#readyState = 1; // OPEN
+			this.#readyState = 1;
 		}
 	}
 
-	// MARK: - Lifecycle
-	get bufferedAmount(): number {
-		return this.#bufferedAmount;
+	get websocket(): UniversalWebSocket {
+		return this.#ws;
 	}
 
-	_handleOpen(requestId: ArrayBuffer): void {
-		if (this.#readyState !== 0) {
-			// CONNECTING
-			return;
-		}
+	#handleSend(data: string | ArrayBufferLike | Blob | ArrayBufferView): void {
+		let isBinary = false;
+		let messageData: string | ArrayBuffer;
 
-		this.#readyState = 1; // OPEN
+		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");
+		}
 
-		const event = {
-			type: "open",
-			rivetRequestId: requestId,
-			target: this,
-		};
+		this.#sendCallback(messageData, isBinary);
+	}
 
-		this.#fireEvent("open", event);
+	// 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,
@@ -121,45 +108,28 @@ export class WebSocketTunnelAdapter {
 				requestId: this.#requestId,
 				actorId: this.#actorId,
 				currentReadyState: this.#readyState,
-				expectedReadyState: 1,
-				serverMessageIndex: serverMessageIndex,
-				hibernatable: this.#hibernatable,
 			});
 			return true;
 		}
 
-		// Validate message index
+		// Validate message index for hibernatable websockets
 		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",
+					msg: "received duplicate hibernating websocket message",
 					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,
@@ -168,400 +138,64 @@ export class WebSocketTunnelAdapter {
 					expectedIndex,
 					receivedIndex: serverMessageIndex,
 					closeReason,
-					gap: wrappingSubU16(
-						wrappingSubU16(serverMessageIndex, previousIndex),
-						1,
-					),
+					gap: wrappingSubU16(wrappingSubU16(serverMessageIndex, previousIndex), 1),
 				});
-
-				// Close the WebSocket and skip processing
-				this.close(1008, closeReason);
-
+				this.#close(1008, closeReason, true);
 				return true;
 			}
 
-			// Update to the next index
 			this.#serverMessageIndex = serverMessageIndex;
 		}
 
-		// Dispatch event
-		let messageData: any;
-		if (isBinary) {
-			// Handle binary data based on binaryType
+		// Convert data based on binaryType
+		let messageData: any = data;
+		if (isBinary && data instanceof Uint8Array) {
 			if (this.#binaryType === "nodebuffer") {
-				// Convert to Buffer for Node.js compatibility
-				messageData = Buffer.from(data as Uint8Array);
+				messageData = Buffer.from(data);
 			} 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",
-				);
+				messageData = data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength);
 			}
-		} else {
-			messageData = data;
 		}
 
-		const event = {
+		this.#ws.dispatchEvent({
 			type: "message",
 			data: messageData,
 			rivetRequestId: requestId,
 			rivetMessageIndex: serverMessageIndex,
-			target: this,
-		};
-
-		this.#fireEvent("message", event);
+			target: this.#ws,
+		} as RivetMessageEvent);
 
 		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);
+	// 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.#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);
+		this.#close(code, reason, false);
 	}
 
+	// Public close method (used by tunnel.ts for stale websocket cleanup)
 	close(code?: number, reason?: string): void {
-		this.#closeInner(code, reason, true);
+		this.#close(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;
+	#close(code: number | undefined, reason: string | undefined, sendCallback: boolean): void {
+		if (this.#readyState >= 2) return;
 
-	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"));
+		this.#readyState = 2;
+		if (sendCallback) this.#closeCallback(code, reason);
+		this.#readyState = 3;
+		this.#ws.triggerClose(code ?? 1000, reason ?? "");
 	}
 
-	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 {
+	#terminate(): void {
 		// Immediate close without close frame
-		this.#readyState = 3; // CLOSED
+		this.#readyState = 3;
 		this.#closeCallback(1006, "Abnormal Closure");
-
-		const event = {
-			wasClean: false,
-			code: 1006,
-			reason: "Abnormal Closure",
-			type: "close",
-			target: this,
-		};
-
-		this.#fireEvent("close", event);
+		this.#ws.triggerClose(1006, "Abnormal Closure", false);
 	}
 }