@rivetkit/engine-runner 2.0.27 → 2.0.28

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rivetkit/engine-runner",
-  "version": "2.0.27",
+  "version": "2.0.28",
   "type": "module",
   "exports": {
     "import": {
@@ -16,7 +16,7 @@
     "uuid": "^12.0.0",
     "pino": "^9.9.5",
     "ws": "^8.18.3",
-    "@rivetkit/engine-runner-protocol": "2.0.27"
+    "@rivetkit/engine-runner-protocol": "2.0.28"
   },
   "devDependencies": {
     "@types/node": "^22.18.1",

package/src/actor.ts

@@ -27,6 +27,10 @@ export class RunnerActor {
 	}> = [];
 	actorStartPromise: ReturnType<typeof promiseWithResolvers<void>>;
 
+	lastCommandIdx: bigint = -1n;
+	nextEventIdx: bigint = 0n;
+	eventHistory: protocol.EventWrapper[] = [];
+
 	/**
 	 * If restoreHibernatingRequests has been called. This is used to assert
 	 * that the caller is implemented correctly.
@@ -81,8 +85,8 @@ export class RunnerActor {
 			gatewayId,
 			requestId,
 			request: {
-				resolve: () => {},
-				reject: () => {},
+				resolve: () => { },
+				reject: () => { },
 				actorId: this.actorId,
 				gatewayId: gatewayId,
 				requestId: requestId,
@@ -118,8 +122,8 @@ export class RunnerActor {
 			gatewayId,
 			requestId,
 			request: {
-				resolve: () => {},
-				reject: () => {},
+				resolve: () => { },
+				reject: () => { },
 				actorId: this.actorId,
 				gatewayId: gatewayId,
 				requestId: requestId,
@@ -193,4 +197,14 @@ export class RunnerActor {
 			this.webSockets.splice(index, 1);
 		}
 	}
+
+	handleAckEvents(lastEventIdx: bigint) {
+		this.eventHistory = this.eventHistory.filter(
+			(event) => event.checkpoint.index > lastEventIdx,
+		);
+	}
+
+	recordEvent(eventWrapper: protocol.EventWrapper) {
+		this.eventHistory.push(eventWrapper);
+	}
 }

package/src/mod.ts

@@ -8,6 +8,7 @@ import { type HibernatingWebSocketMetadata, Tunnel } from "./tunnel";
 import {
 	calculateBackoff,
 	parseWebSocketCloseReason,
+	stringifyError,
 	unreachable,
 } from "./utils";
 import { importWebSocket } from "./websocket.js";
@@ -17,13 +18,18 @@ export { RunnerActor, type ActorConfig };
 export { idToStr } from "./utils";
 
 const KV_EXPIRE: number = 30_000;
-const PROTOCOL_VERSION: number = 3;
-const RUNNER_PING_INTERVAL = 3_000;
+const PROTOCOL_VERSION: number = 5;
 
 /** Warn once the backlog significantly exceeds the server's ack batch size. */
 const EVENT_BACKLOG_WARN_THRESHOLD = 10_000;
 const SIGNAL_HANDLERS: (() => void | Promise<void>)[] = [];
 
+export class RunnerShutdownError extends Error {
+	constructor() {
+		super("Runner shut down");
+	}
+}
+
 export interface RunnerConfig {
 	logger?: Logger;
 	version: number;
@@ -197,9 +203,6 @@ export class Runner {
 	// WebSocket
 	__pegboardWebSocket?: WebSocket;
 	runnerId?: string;
-	#lastCommandIdx: number = -1;
-	#pingLoop?: NodeJS.Timeout;
-	#nextEventIdx: bigint = 0n;
 	#started: boolean = false;
 	#shutdown: boolean = false;
 	#shuttingDown: boolean = false;
@@ -211,7 +214,6 @@ export class Runner {
 	#runnerLostTimeout?: NodeJS.Timeout;
 
 	// Event storage for resending
-	#eventHistory: protocol.EventWrapper[] = [];
 	#eventBacklogWarned: boolean = false;
 
 	// Command acknowledgment
@@ -255,7 +257,14 @@ export class Runner {
 
 		// Start cleaning up old unsent KV requests every 15 seconds
 		this.#kvCleanupInterval = setInterval(() => {
-			this.#cleanupOldKvRequests();
+			try {
+				this.#cleanupOldKvRequests();
+			} catch (err) {
+				this.log?.error({
+					msg: "error cleaning up kv requests",
+					error: stringifyError(err),
+				});
+			}
 		}, 15000); // Run every 15 seconds
 	}
 
@@ -282,6 +291,11 @@ export class Runner {
 	}
 
 	async forceStopActor(actorId: string, generation?: number) {
+		this.log?.debug({
+			msg: "force stopping actor",
+			actorId,
+		});
+
 		const actor = this.getActor(actorId, generation);
 		if (!actor) return;
 
@@ -299,22 +313,38 @@ export class Runner {
 		// Close requests after onActorStop so you can send messages over the tunnel
 		this.#tunnel?.closeActiveRequests(actor);
 
+		this.#sendActorStateUpdate(actorId, actor.generation, "stopped");
+
 		// Remove actor after stopping in order to ensure that we can still
-		// call actions on the runner. Do this before sending stopped update in
-		// order to ensure we don't have duplicate actors.
+		// call actions on the runner
 		this.#removeActor(actorId, generation);
-
-		this.#sendActorStateUpdate(actorId, actor.generation, "stopped");
 	}
 
-	#stopAllActors() {
+	#handleLost() {
 		this.log?.info({
-			msg: "stopping all actors due to runner lost threshold exceeded",
+			msg: "stopping all actors due to runner lost threshold",
 		});
 
+		// Remove all remaining kv requests
+		for (const [_, request] of this.#kvRequests.entries()) {
+			request.reject(new RunnerShutdownError());
+		}
+
+		this.#kvRequests.clear();
+
+		this.#stopAllActors();
+	}
+
+	#stopAllActors() {
 		const actorIds = Array.from(this.#actors.keys());
 		for (const actorId of actorIds) {
-			this.forceStopActor(actorId);
+			this.forceStopActor(actorId).catch((err) => {
+				this.log?.error({
+					msg: "error stopping actor",
+					actorId,
+					error: stringifyError(err),
+				});
+			});
 		}
 	}
 
@@ -477,12 +507,6 @@ export class Runner {
 			this.#runnerLostTimeout = undefined;
 		}
 
-		// Clear ping loop
-		if (this.#pingLoop) {
-			clearInterval(this.#pingLoop);
-			this.#pingLoop = undefined;
-		}
-
 		// Clear ack interval
 		if (this.#ackInterval) {
 			clearInterval(this.#ackInterval);
@@ -738,10 +762,6 @@ export class Runner {
 				name: this.#config.runnerName,
 				version: this.#config.version,
 				totalSlots: this.#config.totalSlots,
-				lastCommandIdx:
-					this.#lastCommandIdx >= 0
-						? BigInt(this.#lastCommandIdx)
-						: null,
 				prepopulateActorNames: new Map(
 					Object.entries(this.#config.prepopulateActorNames).map(
 						([name, data]) => [
@@ -758,33 +778,22 @@ export class Runner {
 				val: init,
 			});
 
-			// Start ping interval
-			const pingLoop = setInterval(() => {
-				if (ws.readyState === 1) {
-					this.__sendToServer({
-						tag: "ToServerPing",
-						val: {
-							ts: BigInt(Date.now()),
-						},
-					});
-				} else {
-					clearInterval(pingLoop);
-					this.log?.info({
-						msg: "WebSocket not open, stopping ping loop",
-					});
-				}
-			}, RUNNER_PING_INTERVAL);
-			this.#pingLoop = pingLoop;
-
 			// Start command acknowledgment interval (5 minutes)
 			const ackInterval = 5 * 60 * 1000; // 5 minutes in milliseconds
 			const ackLoop = setInterval(() => {
-				if (ws.readyState === 1) {
-					this.#sendCommandAcknowledgment();
-				} else {
-					clearInterval(ackLoop);
-					this.log?.info({
-						msg: "WebSocket not open, stopping ack loop",
+				try {
+					if (ws.readyState === 1) {
+						this.#sendCommandAcknowledgment();
+					} else {
+						clearInterval(ackLoop);
+						this.log?.info({
+							msg: "WebSocket not open, stopping ack loop",
+						});
+					}
+				} catch (err) {
+					this.log?.error({
+						msg: "error in command acknowledgment loop",
+						error: stringifyError(err),
 					});
 				}
 			}, ackInterval);
@@ -815,8 +824,8 @@ export class Runner {
 				if (this.runnerId !== init.runnerId) {
 					this.runnerId = init.runnerId;
 
-					// Clear history if runner id changed
-					this.#eventHistory.length = 0;
+					// Clear actors if runner id changed
+					this.#stopAllActors();
 				}
 
 				// Store the runner lost threshold from metadata
@@ -826,13 +835,12 @@ export class Runner {
 
 				this.log?.info({
 					msg: "received init",
-					lastEventIdx: init.lastEventIdx,
 					runnerLostThreshold: this.#runnerLostThreshold,
 				});
 
 				// Resend pending events
 				this.#processUnsentKvRequests();
-				this.#resendUnacknowledgedEvents(init.lastEventIdx);
+				this.#resendUnacknowledgedEvents();
 				this.#tunnel?.resendBufferedEvents();
 
 				this.#config.onConnected();
@@ -845,10 +853,19 @@ export class Runner {
 				const kvResponse = message.val;
 				this.#handleKvResponse(kvResponse);
 			} else if (message.tag === "ToClientTunnelMessage") {
-				this.#tunnel?.handleTunnelMessage(message.val);
-			} else if (message.tag === "ToClientClose") {
-				this.#tunnel?.shutdown();
-				ws.close(1000, "manual closure");
+				this.#tunnel?.handleTunnelMessage(message.val).catch((err) => {
+					this.log?.error({
+						msg: "error handling tunnel message",
+						error: stringifyError(err),
+					});
+				});
+			} else if (message.tag === "ToClientPing") {
+				this.__sendToServer({
+					tag: "ToServerPong",
+					val: {
+						ts: message.val.ts,
+					},
+				});
 			} else {
 				unreachable(message);
 			}
@@ -871,7 +888,14 @@ export class Runner {
 						seconds: this.#runnerLostThreshold / 1000,
 					});
 					this.#runnerLostTimeout = setTimeout(() => {
-						this.#stopAllActors();
+						try {
+							this.#handleLost();
+						} catch (err) {
+							this.log?.error({
+								msg: "error handling runner lost",
+								error: stringifyError(err),
+							});
+						}
 					}, this.#runnerLostThreshold);
 				}
 
@@ -909,12 +933,6 @@ export class Runner {
 				this.#config.onDisconnected(ev.code, ev.reason);
 			}
 
-			// Clear ping loop on close
-			if (this.#pingLoop) {
-				clearInterval(this.#pingLoop);
-				this.#pingLoop = undefined;
-			}
-
 			// Clear ack interval on close
 			if (this.#ackInterval) {
 				clearInterval(this.#ackInterval);
@@ -933,7 +951,14 @@ export class Runner {
 						seconds: this.#runnerLostThreshold / 1000,
 					});
 					this.#runnerLostTimeout = setTimeout(() => {
-						this.#stopAllActors();
+						try {
+							this.#handleLost();
+						} catch (err) {
+							this.log?.error({
+								msg: "error handling runner lost",
+								error: stringifyError(err),
+							});
+						}
 					}, this.#runnerLostThreshold);
 				}
 
@@ -952,52 +977,88 @@ export class Runner {
 		for (const commandWrapper of commands) {
 			if (commandWrapper.inner.tag === "CommandStartActor") {
 				// Spawn background promise
-				this.#handleCommandStartActor(commandWrapper);
+				this.#handleCommandStartActor(commandWrapper).catch((err) => {
+					this.log?.error({
+						msg: "error handling start actor command",
+						actorId: commandWrapper.checkpoint.actorId,
+						error: stringifyError(err),
+					});
+				});
+
+				// NOTE: We don't do this for CommandStopActor because the actor will be removed by that call
+				// so we cant update the checkpoint
+				const actor = this.getActor(
+					commandWrapper.checkpoint.actorId,
+					commandWrapper.checkpoint.generation,
+				);
+				if (actor) actor.lastCommandIdx = commandWrapper.checkpoint.index;
 			} else if (commandWrapper.inner.tag === "CommandStopActor") {
 				// Spawn background promise
-				this.#handleCommandStopActor(commandWrapper);
+				this.#handleCommandStopActor(commandWrapper).catch((err) => {
+					this.log?.error({
+						msg: "error handling stop actor command",
+						actorId: commandWrapper.checkpoint.actorId,
+						error: stringifyError(err),
+					});
+				});
 			} else {
 				unreachable(commandWrapper.inner);
 			}
-
-			this.#lastCommandIdx = Number(commandWrapper.index);
 		}
 	}
 
 	#handleAckEvents(ack: protocol.ToClientAckEvents) {
-		const lastAckedIdx = ack.lastEventIdx;
+		let originalTotalEvents = Array.from(this.#actors).reduce(
+			(s, [_, actor]) => s + actor.eventHistory.length,
+			0,
+		);
+
+		for (const [_, actor] of this.#actors) {
+			let checkpoint = ack.lastEventCheckpoints.find(
+				(x) => x.actorId == actor.actorId,
+			);
 
-		const originalLength = this.#eventHistory.length;
-		this.#eventHistory = this.#eventHistory.filter(
-			(event) => event.index > lastAckedIdx,
+			if (checkpoint) actor.handleAckEvents(checkpoint.index);
+		}
+
+		const totalEvents = Array.from(this.#actors).reduce(
+			(s, [_, actor]) => s + actor.eventHistory.length,
+			0,
 		);
+		const prunedCount = originalTotalEvents - totalEvents;
 
-		const prunedCount = originalLength - this.#eventHistory.length;
 		if (prunedCount > 0) {
 			this.log?.info({
 				msg: "pruned acknowledged events",
-				lastAckedIdx: lastAckedIdx.toString(),
 				prunedCount,
 			});
 		}
 
-		if (this.#eventHistory.length <= EVENT_BACKLOG_WARN_THRESHOLD) {
+		if (totalEvents <= EVENT_BACKLOG_WARN_THRESHOLD) {
 			this.#eventBacklogWarned = false;
 		}
 	}
 
 	/** Track events to send to the server in case we need to resend it on disconnect. */
 	#recordEvent(eventWrapper: protocol.EventWrapper) {
-		this.#eventHistory.push(eventWrapper);
+		const actor = this.getActor(eventWrapper.checkpoint.actorId);
+		if (!actor) return;
+
+		actor.recordEvent(eventWrapper);
+
+		let totalEvents = Array.from(this.#actors).reduce(
+			(s, [_, actor]) => s + actor.eventHistory.length,
+			0,
+		);
 
 		if (
-			this.#eventHistory.length > EVENT_BACKLOG_WARN_THRESHOLD &&
+			totalEvents > EVENT_BACKLOG_WARN_THRESHOLD &&
 			!this.#eventBacklogWarned
 		) {
 			this.#eventBacklogWarned = true;
 			this.log?.warn({
 				msg: "unacknowledged event backlog exceeds threshold",
-				backlogSize: this.#eventHistory.length,
+				backlogSize: totalEvents,
 				threshold: EVENT_BACKLOG_WARN_THRESHOLD,
 			});
 		}
@@ -1013,8 +1074,8 @@ export class Runner {
 		const startCommand = commandWrapper.inner
 			.val as protocol.CommandStartActor;
 
-		const actorId = startCommand.actorId;
-		const generation = startCommand.generation;
+		const actorId = commandWrapper.checkpoint.actorId;
+		const generation = commandWrapper.checkpoint.generation;
 		const config = startCommand.config;
 
 		const actorConfig: ActorConfig = {
@@ -1094,8 +1155,8 @@ export class Runner {
 		const stopCommand = commandWrapper.inner
 			.val as protocol.CommandStopActor;
 
-		const actorId = stopCommand.actorId;
-		const generation = stopCommand.generation;
+		const actorId = commandWrapper.checkpoint.actorId;
+		const generation = commandWrapper.checkpoint.generation;
 
 		await this.forceStopActor(actorId, generation);
 	}
@@ -1105,6 +1166,9 @@ export class Runner {
 		generation: number,
 		intentType: "sleep" | "stop",
 	) {
+		const actor = this.getActor(actorId, generation);
+		if (!actor) return;
+
 		let actorIntent: protocol.ActorIntent;
 
 		if (intentType === "sleep") {
@@ -1119,14 +1183,15 @@ export class Runner {
 		}
 
 		const intentEvent: protocol.EventActorIntent = {
-			actorId,
-			generation,
 			intent: actorIntent,
 		};
 
-		const eventIndex = this.#nextEventIdx++;
 		const eventWrapper: protocol.EventWrapper = {
-			index: eventIndex,
+			checkpoint: {
+				actorId,
+				generation,
+				index: actor.nextEventIdx++,
+			},
 			inner: {
 				tag: "EventActorIntent",
 				val: intentEvent,
@@ -1146,6 +1211,9 @@ export class Runner {
 		generation: number,
 		stateType: "running" | "stopped",
 	) {
+		const actor = this.getActor(actorId, generation);
+		if (!actor) return;
+
 		let actorState: protocol.ActorState;
 
 		if (stateType === "running") {
@@ -1163,14 +1231,15 @@ export class Runner {
 		}
 
 		const stateUpdateEvent: protocol.EventActorStateUpdate = {
-			actorId,
-			generation,
 			state: actorState,
 		};
 
-		const eventIndex = this.#nextEventIdx++;
 		const eventWrapper: protocol.EventWrapper = {
-			index: eventIndex,
+			checkpoint: {
+				actorId,
+				generation,
+				index: actor.nextEventIdx++,
+			},
 			inner: {
 				tag: "EventActorStateUpdate",
 				val: stateUpdateEvent,
@@ -1186,9 +1255,19 @@ export class Runner {
 	}
 
 	#sendCommandAcknowledgment() {
-		if (this.#lastCommandIdx < 0) {
-			// No commands received yet, nothing to acknowledge
-			return;
+		const lastCommandCheckpoints = [];
+
+		for (const [_, actor] of this.#actors) {
+			if (actor.lastCommandIdx < 0) {
+				// No commands received yet, nothing to acknowledge
+				continue;
+			}
+
+			lastCommandCheckpoints.push({
+				actorId: actor.actorId,
+				generation: actor.generation,
+				index: actor.lastCommandIdx,
+			});
 		}
 
 		//this.#log?.log("Sending command acknowledgment", this.#lastCommandIdx);
@@ -1196,7 +1275,7 @@ export class Runner {
 		this.__sendToServer({
 			tag: "ToServerAckCommands",
 			val: {
-				lastCommandIdx: BigInt(this.#lastCommandIdx),
+				lastCommandCheckpoints,
 			},
 		});
 	}
@@ -1495,14 +1574,15 @@ export class Runner {
 		if (!actor) return;
 
 		const alarmEvent: protocol.EventActorSetAlarm = {
-			actorId,
-			generation: actor.generation,
 			alarmTs: alarmTs !== null ? BigInt(alarmTs) : null,
 		};
 
-		const eventIndex = this.#nextEventIdx++;
 		const eventWrapper: protocol.EventWrapper = {
-			index: eventIndex,
+			checkpoint: {
+				actorId,
+				generation: actor.generation,
+				index: actor.nextEventIdx++,
+			},
 			inner: {
 				tag: "EventActorSetAlarm",
 				val: alarmEvent,
@@ -1669,6 +1749,7 @@ export class Runner {
 			tag: "ToServerlessServerInit",
 			val: {
 				runnerId: this.runnerId,
+				runnerProtocolVersion: PROTOCOL_VERSION,
 			},
 		});
 
@@ -1699,27 +1780,34 @@ export class Runner {
 			msg: `Scheduling reconnect attempt ${this.#reconnectAttempt + 1} in ${delay}ms`,
 		});
 
-		this.#reconnectTimeout = setTimeout(async () => {
+		this.#reconnectTimeout = setTimeout(() => {
 			if (!this.#shutdown) {
 				this.#reconnectAttempt++;
 				this.log?.debug({
 					msg: `Attempting to reconnect (attempt ${this.#reconnectAttempt})...`,
 				});
-				await this.#openPegboardWebSocket();
+				this.#openPegboardWebSocket().catch((err) => {
+					this.log?.error({
+						msg: "error during websocket reconnection",
+						error: stringifyError(err),
+					});
+				});
 			}
 		}, delay);
 	}
 
-	#resendUnacknowledgedEvents(lastEventIdx: bigint) {
-		const eventsToResend = this.#eventHistory.filter(
-			(event) => event.index > lastEventIdx,
-		);
+	#resendUnacknowledgedEvents() {
+		const eventsToResend = [];
+
+		for (const [_, actor] of this.#actors) {
+			eventsToResend.push(...actor.eventHistory);
+		}
 
 		if (eventsToResend.length === 0) return;
 
 		this.log?.info({
 			msg: "resending unacknowledged events",
-			fromIndex: lastEventIdx + 1n,
+			count: eventsToResend.length,
 		});
 
 		// Resend events in batches