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

package/src/websocket-tunnel-adapter.ts

@@ -17,7 +17,7 @@ export class WebSocketTunnelAdapter {
 	#protocol = "";
 	#url = "";
 	#sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void;
-	#closeCallback: (code?: number, reason?: string) => void;
+	#closeCallback: (code?: number, reason?: string, retry?: boolean) => void;
 
 	// Event buffering for events fired before listeners are attached
 	#bufferedEvents: Array<{
@@ -28,7 +28,11 @@ export class WebSocketTunnelAdapter {
 	constructor(
 		webSocketId: string,
 		sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void,
-		closeCallback: (code?: number, reason?: string) => void,
+		closeCallback: (
+			code?: number,
+			reason?: string,
+			retry?: boolean,
+		) => void,
 	) {
 		this.#webSocketId = webSocketId;
 		this.#sendCallback = sendCallback;
@@ -186,6 +190,14 @@ export class WebSocketTunnelAdapter {
 	}
 
 	close(code?: number, reason?: string): void {
+		this.closeInner(code, reason);
+	}
+
+	__closeWithRetry(code?: number, reason?: string): void {
+		this.closeInner(code, reason, true);
+	}
+
+	closeInner(code?: number, reason?: string, retry: boolean = false): void {
 		if (
 			this.#readyState === 2 || // CLOSING
 			this.#readyState === 3 // CLOSED
@@ -196,7 +208,7 @@ export class WebSocketTunnelAdapter {
 		this.#readyState = 2; // CLOSING
 
 		// Send close through tunnel
-		this.#closeCallback(code, reason);
+		this.#closeCallback(code, reason, retry);
 
 		// Update state and fire event
 		this.#readyState = 3; // CLOSED
@@ -410,7 +422,7 @@ export class WebSocketTunnelAdapter {
 	}
 
 	// Internal methods called by the Tunnel class
-	_handleOpen(): void {
+	_handleOpen(requestId: ArrayBuffer): void {
 		if (this.#readyState !== 0) {
 			// CONNECTING
 			return;
@@ -420,16 +432,23 @@ export class WebSocketTunnelAdapter {
 
 		const event = {
 			type: "open",
+			rivetRequestId: requestId,
 			target: this,
 		};
 
 		this.#fireEvent("open", event);
 	}
 
-	_handleMessage(data: string | Uint8Array, isBinary: boolean): void {
+	/// Returns false if the message was sent off.
+	_handleMessage(
+		requestId: ArrayBuffer,
+		data: string | Uint8Array,
+		index: number,
+		isBinary: boolean,
+	): boolean {
 		if (this.#readyState !== 1) {
 			// OPEN
-			return;
+			return true;
 		}
 
 		let messageData: any;
@@ -460,15 +479,19 @@ export class WebSocketTunnelAdapter {
 		}
 
 		const event = {
-			data: messageData,
 			type: "message",
+			data: messageData,
+			rivetRequestId: requestId,
+			rivetMessageIndex: index,
 			target: this,
 		};
 
 		this.#fireEvent("message", event);
+
+		return false;
 	}
 
-	_handleClose(code?: number, reason?: string): void {
+	_handleClose(requestId: ArrayBuffer, code?: number, reason?: string): void {
 		if (this.#readyState === 3) {
 			// CLOSED
 			return;
@@ -477,10 +500,11 @@ export class WebSocketTunnelAdapter {
 		this.#readyState = 3; // CLOSED
 
 		const event = {
+			type: "close",
 			wasClean: true,
 			code: code || 1000,
 			reason: reason || "",
-			type: "close",
+			rivetRequestId: requestId,
 			target: this,
 		};