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

package/src/stringify.ts

@@ -0,0 +1,184 @@
+import type * as protocol from "@rivetkit/engine-runner-protocol";
+
+/**
+ * Helper function to stringify ArrayBuffer for logging
+ */
+function stringifyArrayBuffer(buffer: ArrayBuffer): string {
+	return `ArrayBuffer(${buffer.byteLength})`;
+}
+
+/**
+ * Helper function to stringify bigint for logging
+ */
+function stringifyBigInt(value: bigint): string {
+	return `${value}n`;
+}
+
+/**
+ * Helper function to stringify Map for logging
+ */
+function stringifyMap(map: ReadonlyMap<string, string>): string {
+	const entries = Array.from(map.entries())
+		.map(([k, v]) => `"${k}": "${v}"`)
+		.join(", ");
+	return `Map(${map.size}){${entries}}`;
+}
+
+/**
+ * Stringify ToServerTunnelMessageKind for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyToServerTunnelMessageKind(
+	kind: protocol.ToServerTunnelMessageKind,
+): string {
+	switch (kind.tag) {
+		case "TunnelAck":
+			return "TunnelAck";
+		case "ToServerResponseStart": {
+			const { status, headers, body, stream } = kind.val;
+			const bodyStr = body === null ? "null" : stringifyArrayBuffer(body);
+			return `ToServerResponseStart{status: ${status}, headers: ${stringifyMap(headers)}, body: ${bodyStr}, stream: ${stream}}`;
+		}
+		case "ToServerResponseChunk": {
+			const { body, finish } = kind.val;
+			return `ToServerResponseChunk{body: ${stringifyArrayBuffer(body)}, finish: ${finish}}`;
+		}
+		case "ToServerResponseAbort":
+			return "ToServerResponseAbort";
+		case "ToServerWebSocketOpen": {
+			const { canHibernate, lastMsgIndex } = kind.val;
+			return `ToServerWebSocketOpen{canHibernate: ${canHibernate}, lastMsgIndex: ${stringifyBigInt(lastMsgIndex)}}`;
+		}
+		case "ToServerWebSocketMessage": {
+			const { data, binary } = kind.val;
+			return `ToServerWebSocketMessage{data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`;
+		}
+		case "ToServerWebSocketMessageAck": {
+			const { index } = kind.val;
+			return `ToServerWebSocketMessageAck{index: ${index}}`;
+		}
+		case "ToServerWebSocketClose": {
+			const { code, reason, retry } = kind.val;
+			const codeStr = code === null ? "null" : code.toString();
+			const reasonStr = reason === null ? "null" : `"${reason}"`;
+			return `ToServerWebSocketClose{code: ${codeStr}, reason: ${reasonStr}, retry: ${retry}}`;
+		}
+	}
+}
+
+/**
+ * Stringify ToClientTunnelMessageKind for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyToClientTunnelMessageKind(
+	kind: protocol.ToClientTunnelMessageKind,
+): string {
+	switch (kind.tag) {
+		case "TunnelAck":
+			return "TunnelAck";
+		case "ToClientRequestStart": {
+			const { actorId, method, path, headers, body, stream } = kind.val;
+			const bodyStr = body === null ? "null" : stringifyArrayBuffer(body);
+			return `ToClientRequestStart{actorId: "${actorId}", method: "${method}", path: "${path}", headers: ${stringifyMap(headers)}, body: ${bodyStr}, stream: ${stream}}`;
+		}
+		case "ToClientRequestChunk": {
+			const { body, finish } = kind.val;
+			return `ToClientRequestChunk{body: ${stringifyArrayBuffer(body)}, finish: ${finish}}`;
+		}
+		case "ToClientRequestAbort":
+			return "ToClientRequestAbort";
+		case "ToClientWebSocketOpen": {
+			const { actorId, path, headers } = kind.val;
+			return `ToClientWebSocketOpen{actorId: "${actorId}", path: "${path}", headers: ${stringifyMap(headers)}}`;
+		}
+		case "ToClientWebSocketMessage": {
+			const { index, data, binary } = kind.val;
+			return `ToClientWebSocketMessage{index: ${index}, data: ${stringifyArrayBuffer(data)}, binary: ${binary}}`;
+		}
+		case "ToClientWebSocketClose": {
+			const { code, reason } = kind.val;
+			const codeStr = code === null ? "null" : code.toString();
+			const reasonStr = reason === null ? "null" : `"${reason}"`;
+			return `ToClientWebSocketClose{code: ${codeStr}, reason: ${reasonStr}}`;
+		}
+	}
+}
+
+/**
+ * Stringify Command for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyCommand(command: protocol.Command): string {
+	switch (command.tag) {
+		case "CommandStartActor": {
+			const { actorId, generation, config } = command.val;
+			const keyStr = config.key === null ? "null" : `"${config.key}"`;
+			const inputStr =
+				config.input === null
+					? "null"
+					: stringifyArrayBuffer(config.input);
+			return `CommandStartActor{actorId: "${actorId}", generation: ${generation}, config: {name: "${config.name}", key: ${keyStr}, createTs: ${stringifyBigInt(config.createTs)}, input: ${inputStr}}}`;
+		}
+		case "CommandStopActor": {
+			const { actorId, generation } = command.val;
+			return `CommandStopActor{actorId: "${actorId}", generation: ${generation}}`;
+		}
+	}
+}
+
+/**
+ * Stringify CommandWrapper for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyCommandWrapper(
+	wrapper: protocol.CommandWrapper,
+): string {
+	return `CommandWrapper{index: ${stringifyBigInt(wrapper.index)}, inner: ${stringifyCommand(wrapper.inner)}}`;
+}
+
+/**
+ * Stringify Event for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyEvent(event: protocol.Event): string {
+	switch (event.tag) {
+		case "EventActorIntent": {
+			const { actorId, generation, intent } = event.val;
+			const intentStr =
+				intent.tag === "ActorIntentSleep"
+					? "Sleep"
+					: intent.tag === "ActorIntentStop"
+						? "Stop"
+						: "Unknown";
+			return `EventActorIntent{actorId: "${actorId}", generation: ${generation}, intent: ${intentStr}}`;
+		}
+		case "EventActorStateUpdate": {
+			const { actorId, generation, state } = event.val;
+			let stateStr: string;
+			if (state.tag === "ActorStateRunning") {
+				stateStr = "Running";
+			} else if (state.tag === "ActorStateStopped") {
+				const { code, message } = state.val;
+				const messageStr = message === null ? "null" : `"${message}"`;
+				stateStr = `Stopped{code: ${code}, message: ${messageStr}}`;
+			} else {
+				stateStr = "Unknown";
+			}
+			return `EventActorStateUpdate{actorId: "${actorId}", generation: ${generation}, state: ${stateStr}}`;
+		}
+		case "EventActorSetAlarm": {
+			const { actorId, generation, alarmTs } = event.val;
+			const alarmTsStr =
+				alarmTs === null ? "null" : stringifyBigInt(alarmTs);
+			return `EventActorSetAlarm{actorId: "${actorId}", generation: ${generation}, alarmTs: ${alarmTsStr}}`;
+		}
+	}
+}
+
+/**
+ * Stringify EventWrapper for logging
+ * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
+ */
+export function stringifyEventWrapper(wrapper: protocol.EventWrapper): string {
+	return `EventWrapper{index: ${stringifyBigInt(wrapper.index)}, inner: ${stringifyEvent(wrapper.inner)}}`;
+}

package/src/tunnel.ts

@@ -1,8 +1,13 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
 import type { MessageId, RequestId } from "@rivetkit/engine-runner-protocol";
+import type { Logger } from "pino";
 import { stringify as uuidstringify, v4 as uuidv4 } from "uuid";
 import { logger } from "./log";
 import type { ActorInstance, Runner } from "./mod";
+import {
+	stringifyToClientTunnelMessageKind,
+	stringifyToServerTunnelMessageKind,
+} from "./stringify";
 import { unreachable } from "./utils";
 import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
 
@@ -22,6 +27,12 @@ interface PendingTunnelMessage {
 	requestIdStr: string;
 }
 
+class RunnerShutdownError extends Error {
+	constructor() {
+		super("Runner shut down");
+	}
+}
+
 export class Tunnel {
 	#runner: Runner;
 
@@ -35,6 +46,10 @@ export class Tunnel {
 
 	#gcInterval?: NodeJS.Timeout;
 
+	get log(): Logger | undefined {
+		return this.#runner.log;
+	}
+
 	constructor(runner: Runner) {
 		this.#runner = runner;
 	}
@@ -44,20 +59,34 @@ export class Tunnel {
 	}
 
 	shutdown() {
+		// NOTE: Pegboard WS already closed at this point, cannot send
+		// anything. All teardown logic is handled by pegboard-runner.
+
 		if (this.#gcInterval) {
 			clearInterval(this.#gcInterval);
 			this.#gcInterval = undefined;
 		}
 
 		// Reject all pending requests
+		//
+		// RunnerShutdownError will be explicitly ignored
 		for (const [_, request] of this.#actorPendingRequests) {
-			request.reject(new Error("Tunnel shutting down"));
+			request.reject(new RunnerShutdownError());
 		}
 		this.#actorPendingRequests.clear();
 
 		// Close all WebSockets
+		//
+		// The WebSocket close event with retry is automatically sent when the
+		// runner WS closes, so we only need to notify the client that the WS
+		// closed:
+		// https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157
 		for (const [_, ws] of this.#actorWebSockets) {
-			ws.__closeWithRetry();
+			// Only close non-hibernatable websockets to prevent sending
+			// unnecessary close messages for websockets that will be hibernated
+			if (!ws.canHibernate) {
+				ws.__closeWithoutCallback(1000, "ws.tunnel_shutdown");
+			}
 		}
 		this.#actorWebSockets.clear();
 	}
@@ -68,9 +97,11 @@ export class Tunnel {
 	) {
 		// TODO: Switch this with runner WS
 		if (!this.#runner.__webSocketReady()) {
-			logger()?.warn(
-				"cannot send tunnel message, socket not connected to engine",
-			);
+			this.log?.warn({
+				msg: "cannot send tunnel message, socket not connected to engine",
+				requestId: idToStr(requestId),
+				message: stringifyToServerTunnelMessageKind(messageKind),
+			});
 			return;
 		}
 
@@ -84,11 +115,11 @@ export class Tunnel {
 			requestIdStr,
 		});
 
-		logger()?.debug({
+		this.log?.debug({
 			msg: "send tunnel msg",
 			requestId: requestIdStr,
 			messageId: messageIdStr,
-			message: messageKind,
+			message: stringifyToServerTunnelMessageKind(messageKind),
 		});
 
 		// Send message
@@ -117,7 +148,7 @@ export class Tunnel {
 			},
 		};
 
-		logger()?.debug({
+		this.log?.debug({
 			msg: "ack tunnel msg",
 			requestId: idToStr(requestId),
 			messageId: idToStr(messageId),
@@ -184,7 +215,7 @@ export class Tunnel {
 
 		// Remove timed out messages
 		if (messagesToDelete.length > 0) {
-			logger()?.warn({
+			this.log?.warn({
 				msg: "purging unacked tunnel messages, this indicates that the Rivet Engine is disconnected or not responding",
 				count: messagesToDelete.length,
 			});
@@ -208,11 +239,11 @@ export class Tunnel {
 		actor.requests.clear();
 
 		// Flush acks and close all WebSockets for this actor
-		for (const webSocketId of actor.webSockets) {
-			const ws = this.#actorWebSockets.get(webSocketId);
+		for (const requestIdStr of actor.webSockets) {
+			const ws = this.#actorWebSockets.get(requestIdStr);
 			if (ws) {
 				ws.__closeWithRetry(1000, "Actor stopped");
-				this.#actorWebSockets.delete(webSocketId);
+				this.#actorWebSockets.delete(requestIdStr);
 			}
 		}
 		actor.webSockets.clear();
@@ -225,7 +256,7 @@ export class Tunnel {
 	): Promise<Response> {
 		// Validate actor exists
 		if (!this.#runner.hasActor(actorId)) {
-			logger()?.warn({
+			this.log?.warn({
 				msg: "ignoring request for unknown actor",
 				actorId,
 			});
@@ -257,11 +288,11 @@ export class Tunnel {
 	async handleTunnelMessage(message: protocol.ToClientTunnelMessage) {
 		const requestIdStr = idToStr(message.requestId);
 		const messageIdStr = idToStr(message.messageId);
-		logger()?.debug({
+		this.log?.debug({
 			msg: "receive tunnel msg",
 			requestId: requestIdStr,
 			messageId: messageIdStr,
-			message: message.messageKind,
+			message: stringifyToClientTunnelMessageKind(message.messageKind),
 		});
 
 		if (message.messageKind.tag === "TunnelAck") {
@@ -271,7 +302,7 @@ export class Tunnel {
 				const didDelete =
 					this.#pendingTunnelMessages.delete(messageIdStr);
 				if (!didDelete) {
-					logger()?.warn({
+					this.log?.warn({
 						msg: "received tunnel ack for nonexistent message",
 						requestId: requestIdStr,
 						messageId: messageIdStr,
@@ -402,8 +433,16 @@ export class Tunnel {
 				await this.#sendResponse(requestId, response);
 			}
 		} catch (error) {
-			logger()?.error({ msg: "error handling request", error });
-			this.#sendResponseError(requestId, 500, "Internal Server Error");
+			if (error instanceof RunnerShutdownError) {
+				this.log?.debug({ msg: "catught runner shutdown error" });
+			} else {
+				this.log?.error({ msg: "error handling request", error });
+				this.#sendResponseError(
+					requestId,
+					500,
+					"Internal Server Error",
+				);
+			}
 		} finally {
 			// Clean up request tracking
 			const actor = this.#runner.getActor(req.actorId);
@@ -493,11 +532,12 @@ export class Tunnel {
 		requestId: protocol.RequestId,
 		open: protocol.ToClientWebSocketOpen,
 	) {
-		const webSocketId = idToStr(requestId);
+		const requestIdStr = idToStr(requestId);
+
 		// Validate actor exists
 		const actor = this.#runner.getActor(open.actorId);
 		if (!actor) {
-			logger()?.warn({
+			this.log?.warn({
 				msg: "ignoring websocket for unknown actor",
 				actorId: open.actorId,
 			});
@@ -521,7 +561,7 @@ export class Tunnel {
 		const websocketHandler = this.#runner.config.websocket;
 
 		if (!websocketHandler) {
-			logger()?.error({
+			this.log?.error({
 				msg: "no websocket handler configured for tunnel",
 			});
 			// Send close immediately
@@ -536,15 +576,33 @@ export class Tunnel {
 			return;
 		}
 
+		// Close existing WebSocket if one already exists for this request ID.
+		// There should always be a close message sent before another open
+		// message for the same message ID.
+		//
+		// This should never occur if all is functioning correctly, but this
+		// prevents any edge case that would result in duplicate WebSockets for
+		// the same request.
+		const existingAdapter = this.#actorWebSockets.get(requestIdStr);
+		if (existingAdapter) {
+			this.log?.warn({
+				msg: "closing existing websocket for duplicate open event for the same request id",
+				requestId: requestIdStr,
+			});
+			// Close without sending a message through the tunnel since the server
+			// already knows about the new connection
+			existingAdapter.__closeWithoutCallback(1000, "ws.duplicate_open");
+		}
+
 		// Track this WebSocket for the actor
 		if (actor) {
-			actor.webSockets.add(webSocketId);
+			actor.webSockets.add(requestIdStr);
 		}
 
 		try {
 			// Create WebSocket adapter
 			const adapter = new WebSocketTunnelAdapter(
-				webSocketId,
+				requestIdStr,
 				(data: ArrayBuffer | string, isBinary: boolean) => {
 					// Send message through tunnel
 					const dataBuffer =
@@ -573,17 +631,17 @@ export class Tunnel {
 					});
 
 					// Remove from map
-					this.#actorWebSockets.delete(webSocketId);
+					this.#actorWebSockets.delete(requestIdStr);
 
 					// Clean up actor tracking
 					if (actor) {
-						actor.webSockets.delete(webSocketId);
+						actor.webSockets.delete(requestIdStr);
 					}
 				},
 			);
 
 			// Store adapter
-			this.#actorWebSockets.set(webSocketId, adapter);
+			this.#actorWebSockets.set(requestIdStr, adapter);
 
 			// Convert headers to map
 			//
@@ -613,6 +671,8 @@ export class Tunnel {
 					requestId,
 					request,
 				);
+			adapter.canHibernate = hibernationConfig.enabled;
+
 			this.#sendMessage(requestId, {
 				tag: "ToServerWebSocketOpen",
 				val: {
@@ -633,7 +693,7 @@ export class Tunnel {
 				request,
 			);
 		} catch (error) {
-			logger()?.error({ msg: "error handling websocket open", error });
+			this.log?.error({ msg: "error handling websocket open", error });
 			// Send close on error
 			this.#sendMessage(requestId, {
 				tag: "ToServerWebSocketClose",
@@ -644,11 +704,11 @@ export class Tunnel {
 				},
 			});
 
-			this.#actorWebSockets.delete(webSocketId);
+			this.#actorWebSockets.delete(requestIdStr);
 
 			// Clean up actor tracking
 			if (actor) {
-				actor.webSockets.delete(webSocketId);
+				actor.webSockets.delete(requestIdStr);
 			}
 		}
 	}
@@ -658,8 +718,8 @@ export class Tunnel {
 		requestId: ArrayBuffer,
 		msg: protocol.ToClientWebSocketMessage,
 	): Promise<boolean> {
-		const webSocketId = idToStr(requestId);
-		const adapter = this.#actorWebSockets.get(webSocketId);
+		const requestIdStr = idToStr(requestId);
+		const adapter = this.#actorWebSockets.get(requestIdStr);
 		if (adapter) {
 			const data = msg.binary
 				? new Uint8Array(msg.data)
@@ -677,7 +737,7 @@ export class Tunnel {
 	}
 
 	__ackWebsocketMessage(requestId: ArrayBuffer, index: number) {
-		logger()?.debug({
+		this.log?.debug({
 			msg: "ack ws msg",
 			requestId: idToStr(requestId),
 			index,

package/src/utils.ts

@@ -1,3 +1,5 @@
+import { logger } from "./log";
+
 export function unreachable(x: never): never {
 	throw `Unreachable: ${x}`;
 }
@@ -29,3 +31,36 @@ export function calculateBackoff(
 
 	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,
+	};
+}

package/src/websocket-tunnel-adapter.ts

@@ -18,6 +18,7 @@ export class WebSocketTunnelAdapter {
 	#url = "";
 	#sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void;
 	#closeCallback: (code?: number, reason?: string, retry?: boolean) => void;
+	#canHibernate: boolean = false;
 
 	// Event buffering for events fired before listeners are attached
 	#bufferedEvents: Array<{
@@ -73,6 +74,16 @@ export class WebSocketTunnelAdapter {
 		return this.#url;
 	}
 
+	/** @experimental */
+	get canHibernate(): boolean {
+		return this.#canHibernate;
+	}
+
+	/** @experimental */
+	set canHibernate(value: boolean) {
+		this.#canHibernate = value;
+	}
+
 	get onopen(): ((this: any, ev: any) => any) | null {
 		return this.#onopen;
 	}
@@ -190,14 +201,23 @@ export class WebSocketTunnelAdapter {
 	}
 
 	close(code?: number, reason?: string): void {
-		this.closeInner(code, reason);
+		this.closeInner(code, reason, false, true);
 	}
 
 	__closeWithRetry(code?: number, reason?: string): void {
-		this.closeInner(code, reason, true);
+		this.closeInner(code, reason, true, true);
 	}
 
-	closeInner(code?: number, reason?: string, retry: boolean = false): void {
+	__closeWithoutCallback(code?: number, reason?: string): void {
+		this.closeInner(code, reason, false, false);
+	}
+
+	closeInner(
+		code: number | undefined,
+		reason: string | undefined,
+		retry: boolean,
+		callback: boolean,
+	): void {
 		if (
 			this.#readyState === 2 || // CLOSING
 			this.#readyState === 3 // CLOSED
@@ -208,7 +228,9 @@ export class WebSocketTunnelAdapter {
 		this.#readyState = 2; // CLOSING
 
 		// Send close through tunnel
-		this.#closeCallback(code, reason, retry);
+		if (callback) {
+			this.#closeCallback(code, reason, retry);
+		}
 
 		// Update state and fire event
 		this.#readyState = 3; // CLOSED