@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

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

@@ -40,8 +40,6 @@ export function stringifyToServerTunnelMessageKind(
 	kind: protocol.ToServerTunnelMessageKind,
 ): string {
 	switch (kind.tag) {
-		case "DeprecatedTunnelAck":
-			return "DeprecatedTunnelAck";
 		case "ToServerResponseStart": {
 			const { status, headers, body, stream } = kind.val;
 			const bodyStr = body === null ? "null" : stringifyArrayBuffer(body);
@@ -82,8 +80,6 @@ export function stringifyToClientTunnelMessageKind(
 	kind: protocol.ToClientTunnelMessageKind,
 ): string {
 	switch (kind.tag) {
-		case "DeprecatedTunnelAck":
-			return "DeprecatedTunnelAck";
 		case "ToClientRequestStart": {
 			const { actorId, method, path, headers, body, stream } = kind.val;
 			const bodyStr = body === null ? "null" : stringifyArrayBuffer(body);
@@ -119,8 +115,7 @@ export function stringifyToClientTunnelMessageKind(
 export function stringifyCommand(command: protocol.Command): string {
 	switch (command.tag) {
 		case "CommandStartActor": {
-			const { actorId, generation, config, hibernatingRequests } =
-				command.val;
+			const { config, hibernatingRequests } = command.val;
 			const keyStr = config.key === null ? "null" : `"${config.key}"`;
 			const inputStr =
 				config.input === null
@@ -130,11 +125,10 @@ export function stringifyCommand(command: protocol.Command): string {
 				hibernatingRequests.length > 0
 					? `[${hibernatingRequests.map((hr) => `{gatewayId: ${idToStr(hr.gatewayId)}, requestId: ${idToStr(hr.requestId)}}`).join(", ")}]`
 					: "[]";
-			return `CommandStartActor{actorId: "${actorId}", generation: ${generation}, config: {name: "${config.name}", key: ${keyStr}, createTs: ${stringifyBigInt(config.createTs)}, input: ${inputStr}}, hibernatingRequests: ${hibernatingRequestsStr}}`;
+			return `CommandStartActor{config: {name: "${config.name}", key: ${keyStr}, createTs: ${stringifyBigInt(config.createTs)}, input: ${inputStr}}, hibernatingRequests: ${hibernatingRequestsStr}}`;
 		}
 		case "CommandStopActor": {
-			const { actorId, generation } = command.val;
-			return `CommandStopActor{actorId: "${actorId}", generation: ${generation}}`;
+			return `CommandStopActor`;
 		}
 	}
 }
@@ -146,7 +140,7 @@ export function stringifyCommand(command: protocol.Command): string {
 export function stringifyCommandWrapper(
 	wrapper: protocol.CommandWrapper,
 ): string {
-	return `CommandWrapper{index: ${stringifyBigInt(wrapper.index)}, inner: ${stringifyCommand(wrapper.inner)}}`;
+	return `CommandWrapper{actorId: "${wrapper.checkpoint.actorId}", generation: "${wrapper.checkpoint.generation}", index: ${stringifyBigInt(wrapper.checkpoint.index)}, inner: ${stringifyCommand(wrapper.inner)}}`;
 }
 
 /**
@@ -156,17 +150,17 @@ export function stringifyCommandWrapper(
 export function stringifyEvent(event: protocol.Event): string {
 	switch (event.tag) {
 		case "EventActorIntent": {
-			const { actorId, generation, intent } = event.val;
+			const { intent } = event.val;
 			const intentStr =
 				intent.tag === "ActorIntentSleep"
 					? "Sleep"
 					: intent.tag === "ActorIntentStop"
 						? "Stop"
 						: "Unknown";
-			return `EventActorIntent{actorId: "${actorId}", generation: ${generation}, intent: ${intentStr}}`;
+			return `EventActorIntent{intent: ${intentStr}}`;
 		}
 		case "EventActorStateUpdate": {
-			const { actorId, generation, state } = event.val;
+			const { state } = event.val;
 			let stateStr: string;
 			if (state.tag === "ActorStateRunning") {
 				stateStr = "Running";
@@ -177,13 +171,13 @@ export function stringifyEvent(event: protocol.Event): string {
 			} else {
 				stateStr = "Unknown";
 			}
-			return `EventActorStateUpdate{actorId: "${actorId}", generation: ${generation}, state: ${stateStr}}`;
+			return `EventActorStateUpdate{state: ${stateStr}}`;
 		}
 		case "EventActorSetAlarm": {
-			const { actorId, generation, alarmTs } = event.val;
+			const { alarmTs } = event.val;
 			const alarmTsStr =
 				alarmTs === null ? "null" : stringifyBigInt(alarmTs);
-			return `EventActorSetAlarm{actorId: "${actorId}", generation: ${generation}, alarmTs: ${alarmTsStr}}`;
+			return `EventActorSetAlarm{alarmTs: ${alarmTsStr}}`;
 		}
 	}
 }
@@ -193,7 +187,7 @@ export function stringifyEvent(event: protocol.Event): string {
  * 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)}}`;
+	return `EventWrapper{actorId: ${wrapper.checkpoint.actorId}, generation: "${wrapper.checkpoint.generation}", index: ${stringifyBigInt(wrapper.checkpoint.index)}, inner: ${stringifyEvent(wrapper.inner)}}`;
 }
 
 /**
@@ -207,34 +201,33 @@ export function stringifyToServer(message: protocol.ToServer): string {
 				name,
 				version,
 				totalSlots,
-				lastCommandIdx,
 				prepopulateActorNames,
 				metadata,
 			} = message.val;
-			const lastCommandIdxStr =
-				lastCommandIdx === null
-					? "null"
-					: stringifyBigInt(lastCommandIdx);
 			const prepopulateActorNamesStr =
 				prepopulateActorNames === null
 					? "null"
 					: `Map(${prepopulateActorNames.size})`;
 			const metadataStr = metadata === null ? "null" : `"${metadata}"`;
-			return `ToServerInit{name: "${name}", version: ${version}, totalSlots: ${totalSlots}, lastCommandIdx: ${lastCommandIdxStr}, prepopulateActorNames: ${prepopulateActorNamesStr}, metadata: ${metadataStr}}`;
+			return `ToServerInit{name: "${name}", version: ${version}, totalSlots: ${totalSlots}, prepopulateActorNames: ${prepopulateActorNamesStr}, metadata: ${metadataStr}}`;
 		}
 		case "ToServerEvents": {
 			const events = message.val;
 			return `ToServerEvents{count: ${events.length}, events: [${events.map((e) => stringifyEventWrapper(e)).join(", ")}]}`;
 		}
 		case "ToServerAckCommands": {
-			const { lastCommandIdx } = message.val;
-			return `ToServerAckCommands{lastCommandIdx: ${stringifyBigInt(lastCommandIdx)}}`;
+			const { lastCommandCheckpoints } = message.val;
+			const checkpointsStr =
+				lastCommandCheckpoints.length > 0
+					? `[${lastCommandCheckpoints.map((cp) => `{actorId: "${cp.actorId}", index: ${stringifyBigInt(cp.index)}}`).join(", ")}]`
+					: "[]";
+			return `ToServerAckCommands{lastCommandCheckpoints: ${checkpointsStr}}`;
 		}
 		case "ToServerStopping":
 			return "ToServerStopping";
-		case "ToServerPing": {
+		case "ToServerPong": {
 			const { ts } = message.val;
-			return `ToServerPing{ts: ${stringifyBigInt(ts)}}`;
+			return `ToServerPong{ts: ${stringifyBigInt(ts)}}`;
 		}
 		case "ToServerKvRequest": {
 			const { actorId, requestId, data } = message.val;
@@ -255,19 +248,25 @@ export function stringifyToServer(message: protocol.ToServer): string {
 export function stringifyToClient(message: protocol.ToClient): string {
 	switch (message.tag) {
 		case "ToClientInit": {
-			const { runnerId, lastEventIdx, metadata } = message.val;
+			const { runnerId, metadata } = message.val;
 			const metadataStr = `{runnerLostThreshold: ${stringifyBigInt(metadata.runnerLostThreshold)}}`;
-			return `ToClientInit{runnerId: "${runnerId}", lastEventIdx: ${stringifyBigInt(lastEventIdx)}, metadata: ${metadataStr}}`;
+			return `ToClientInit{runnerId: "${runnerId}", metadata: ${metadataStr}}`;
+		}
+		case "ToClientPing": {
+			const { ts } = message.val;
+			return `ToClientPing{ts: ${stringifyBigInt(ts)}}`;
 		}
-		case "ToClientClose":
-			return "ToClientClose";
 		case "ToClientCommands": {
 			const commands = message.val;
 			return `ToClientCommands{count: ${commands.length}, commands: [${commands.map((c) => stringifyCommandWrapper(c)).join(", ")}]}`;
 		}
 		case "ToClientAckEvents": {
-			const { lastEventIdx } = message.val;
-			return `ToClientAckEvents{lastEventIdx: ${stringifyBigInt(lastEventIdx)}}`;
+			const { lastEventCheckpoints } = message.val;
+			const checkpointsStr =
+				lastEventCheckpoints.length > 0
+					? `[${lastEventCheckpoints.map((cp) => `{actorId: "${cp.actorId}", index: ${stringifyBigInt(cp.index)}}`).join(", ")}]`
+					: "[]";
+			return `ToClientAckEvents{lastEventCheckpoints: ${checkpointsStr}}`;
 		}
 		case "ToClientKvResponse": {
 			const { requestId, data } = message.val;
@@ -304,6 +303,10 @@ function stringifyKvRequestData(data: protocol.KvRequestData): string {
 			const { keys } = data.val;
 			return `KvDeleteRequest{keys: ${keys.length}}`;
 		}
+		case "KvDeleteRangeRequest": {
+			const { start, end } = data.val;
+			return `KvDeleteRangeRequest{start: ${stringifyArrayBuffer(start)}, end: ${stringifyArrayBuffer(end)}}`;
+		}
 		case "KvDropRequest":
 			return "KvDropRequest";
 	}

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

@@ -5,17 +5,12 @@ import type {
 	RequestId,
 } from "@rivetkit/engine-runner-protocol";
 import type { Logger } from "pino";
-import {
-	parse as uuidparse,
-	stringify as uuidstringify,
-	v4 as uuidv4,
-} from "uuid";
-import type { Runner, RunnerActor } from "./mod";
+import { type Runner, type RunnerActor, RunnerShutdownError } from "./mod";
 import {
 	stringifyToClientTunnelMessageKind,
 	stringifyToServerTunnelMessageKind,
 } from "./stringify";
-import { arraysEqual, idToStr, unreachable } from "./utils";
+import { arraysEqual, idToStr, MAX_PAYLOAD_SIZE, stringifyError, unreachable } from "./utils";
 import {
 	HIBERNATABLE_SYMBOL,
 	WebSocketTunnelAdapter,
@@ -41,12 +36,6 @@ export interface HibernatingWebSocketMetadata {
 	headers: Record<string, string>;
 }
 
-class RunnerShutdownError extends Error {
-	constructor() {
-		super("Runner shut down");
-	}
-}
-
 export class Tunnel {
 	#runner: Runner;
 
@@ -189,6 +178,15 @@ export class Tunnel {
 					meta.headers,
 				);
 
+				// Restored websocket handlers can synchronously send messages
+				// during onRestore, so the pending request tracking must exist
+				// before the websocket is handed to user code.
+				actor.resetPendingRequestMessageIndex(
+					meta.gatewayId,
+					meta.requestId,
+					meta.clientMessageIndex,
+				);
+
 				// This will call `runner.config.websocket` under the hood to
 				// attach the event listeners to the WebSocket.
 				// Track this operation to ensure it completes
@@ -206,16 +204,6 @@ export class Tunnel {
 					false,
 				)
 					.then(() => {
-						// Create a PendingRequest entry to track the message index
-						const actor = this.#runner.getActor(actorId);
-						if (actor) {
-							actor.createPendingRequest(
-								gatewayId,
-								requestId,
-								meta.clientMessageIndex,
-							);
-						}
-
 						this.log?.info({
 							msg: "connection successfully restored",
 							actorId,
@@ -226,7 +214,7 @@ export class Tunnel {
 						this.log?.error({
 							msg: "error creating websocket during restore",
 							requestId: requestIdStr,
-							err,
+							error: stringifyError(err),
 						});
 
 						// Close the WebSocket on error
@@ -291,7 +279,7 @@ export class Tunnel {
 						this.log?.error({
 							msg: "error creating stale websocket during restore",
 							requestId: requestIdStr,
-							err,
+							error: stringifyError(err),
 						});
 					});
 
@@ -401,11 +389,11 @@ export class Tunnel {
 		this.addRequestToActor(gatewayId, requestId, actorId);
 
 		// Call WebSocket handler. This handler will add event listeners
-		// for `open`, etc.
+		// for `open`, etc. Pass the VirtualWebSocket (not the adapter) to the actor.
 		await this.#runner.config.websocket(
 			this.#runner,
 			actorId,
-			adapter,
+			adapter.websocket,
 			gatewayId,
 			requestId,
 			request,
@@ -484,7 +472,7 @@ export class Tunnel {
 		messageKind: protocol.ToServerTunnelMessageKind,
 	) {
 		// Buffer message if not connected
-		if (!this.#runner.__webSocketReady()) {
+		if (!this.#runner.getPegboardWebSocketIfReady()) {
 			this.log?.debug({
 				msg: "buffering tunnel message, socket not connected to engine",
 				requestId: idToStr(requestId),
@@ -648,13 +636,14 @@ export class Tunnel {
 			case "ToClientRequestAbort":
 				await this.#handleRequestAbort(gatewayId, requestId);
 				break;
-			case "ToClientWebSocketOpen":
-				await this.#handleWebSocketOpen(
-					gatewayId,
-					requestId,
-					message.messageKind.val,
-				);
-				break;
+				case "ToClientWebSocketOpen":
+					await this.#handleWebSocketOpen(
+						gatewayId,
+						requestId,
+						message.messageId.messageIndex,
+						message.messageKind.val,
+					);
+					break;
 			case "ToClientWebSocketMessage": {
 				await this.#handleWebSocketMessage(
 					gatewayId,
@@ -671,9 +660,6 @@ export class Tunnel {
 					message.messageKind.val,
 				);
 				break;
-			case "DeprecatedTunnelAck":
-				// Ignore deprecated tunnel ack messages
-				break;
 			default:
 				unreachable(message.messageKind);
 		}
@@ -864,6 +850,10 @@ export class Tunnel {
 		// Read the body first to get the actual content
 		const body = response.body ? await response.arrayBuffer() : null;
 
+		if (body && body.byteLength > MAX_PAYLOAD_SIZE) {
+			throw new Error("Response body too large");
+		}
+
 		// Convert headers to map and add Content-Length if not present
 		const headers = new Map<string, string>();
 		response.headers.forEach((value, key) => {
@@ -922,11 +912,12 @@ export class Tunnel {
 	async #handleWebSocketOpen(
 		gatewayId: GatewayId,
 		requestId: RequestId,
+		serverMessageIndex: number,
 		open: protocol.ToClientWebSocketOpen,
 	) {
 		// NOTE: This method is safe to be async since we will not receive any
 		// further WebSocket events until we send a ToServerWebSocketOpen
-		// tunnel message. We can do any async logic we need to between thoes two events.
+		// tunnel message. We can do any async logic we need to between those two events.
 		//
 		// Sending a ToServerWebSocketClose will terminate the WebSocket early.
 
@@ -961,13 +952,18 @@ export class Tunnel {
 		// WebSockets from retransmits.
 		const existingAdapter = actor.getWebSocket(gatewayId, requestId);
 		if (existingAdapter) {
-			this.log?.warn({
-				msg: "closing existing websocket for duplicate open event for the same request id",
-				requestId: requestIdStr,
+			const replayAction =
+				existingAdapter._handleOpenReplay(serverMessageIndex);
+			if (replayAction === "reset") {
+				actor.resetPendingRequestMessageIndex(gatewayId, requestId, 0);
+			}
+			this.#sendMessage(gatewayId, requestId, {
+				tag: "ToServerWebSocketOpen",
+				val: {
+					canHibernate: existingAdapter[HIBERNATABLE_SYMBOL],
+				},
 			});
-			// Close without sending a message through the tunnel since the server
-			// already knows about the new connection
-			existingAdapter._closeWithoutCallback(1000, "ws.duplicate_open");
+			return;
 		}
 
 		// Create WebSocket
@@ -989,15 +985,15 @@ export class Tunnel {
 			// hood to add the event listeners for open, etc. If this handler
 			// throws, then the WebSocket will be closed before sending the
 			// open event.
-			const adapter = await this.#createWebSocket(
-				actor.actorId,
-				gatewayId,
-				requestId,
-				requestIdStr,
-				0,
-				canHibernate,
-				false,
-				request,
+				const adapter = await this.#createWebSocket(
+					actor.actorId,
+					gatewayId,
+					requestId,
+					requestIdStr,
+					serverMessageIndex,
+					canHibernate,
+					false,
+					request,
 				open.path,
 				Object.fromEntries(open.headers),
 				false,
@@ -1088,7 +1084,7 @@ export class Tunnel {
 		});
 
 		if (clientMessageIndex < 0 || clientMessageIndex > 65535)
-			throw new Error("invalid websocket ack index");
+			throw new Error("Invalid websocket ack index");
 
 		// Get the actor to find the gatewayId
 		//
@@ -1166,7 +1162,7 @@ function buildRequestForWebSocket(
 	};
 
 	if (!path.startsWith("/")) {
-		throw new Error("path must start with leading slash");
+		throw new Error("Path must start with leading slash");
 	}
 
 	const request = new Request(`http://actor${path}`, {

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

@@ -1,5 +1,8 @@
 import { logger } from "./log";
 
+// 20MiB. Keep in sync with runner_max_response_payload_body_size from engine/packages/config/src/config/pegboard.rs
+export const MAX_PAYLOAD_SIZE = 20 * 1024 * 1024;
+
 export function unreachable(x: never): never {
 	throw `Unreachable: ${x}`;
 }
@@ -157,3 +160,19 @@ export function idToStr(id: ArrayBuffer): string {
 		.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);
+	}
+}