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

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

@@ -8,22 +8,31 @@ import { type HibernatingWebSocketMetadata, Tunnel } from "./tunnel";
 import {
 	calculateBackoff,
 	parseWebSocketCloseReason,
+	stringifyError,
 	unreachable,
 } from "./utils";
 import { importWebSocket } from "./websocket.js";
+import {
+	v4 as uuidv4,
+} from "uuid";
 
 export type { HibernatingWebSocketMetadata };
 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 = 7;
 
 /** 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;
@@ -34,7 +43,6 @@ export interface RunnerConfig {
 	namespace: string;
 	totalSlots: number;
 	runnerName: string;
-	runnerKey: string;
 	prepopulateActorNames: Record<string, { metadata: Record<string, any> }>;
 	metadata?: Record<string, any>;
 	onConnected: () => void;
@@ -169,6 +177,15 @@ export interface RunnerConfig {
 
 	onActorStop: (actorId: string, generation: number) => Promise<void>;
 	noAutoShutdown?: boolean;
+
+	/**
+	 * Debug option to inject artificial latency (in ms) into WebSocket
+	 * communication. Messages are queued and delivered in order after the
+	 * configured delay.
+	 *
+	 * @experimental For testing only.
+	 */
+	debugLatencyMs?: number;
 }
 
 export interface KvListOptions {
@@ -187,6 +204,7 @@ interface KvRequestEntry {
 
 export class Runner {
 	#config: RunnerConfig;
+	#runnerKey: string = uuidv4();
 
 	get config(): RunnerConfig {
 		return this.#config;
@@ -195,23 +213,21 @@ export class Runner {
 	#actors: Map<string, RunnerActor> = new Map();
 
 	// WebSocket
-	__pegboardWebSocket?: WebSocket;
+	#pegboardWebSocket?: WebSocket;
 	runnerId?: string;
-	#lastCommandIdx: number = -1;
-	#pingLoop?: NodeJS.Timeout;
-	#nextEventIdx: bigint = 0n;
 	#started: boolean = false;
 	#shutdown: boolean = false;
-	#shuttingDown: boolean = false;
+	#draining: boolean = false;
 	#reconnectAttempt: number = 0;
 	#reconnectTimeout?: NodeJS.Timeout;
 
+	// Protocol metadata
+	#protocolMetadata?: protocol.ProtocolMetadata;
+
 	// Runner lost threshold management
-	#runnerLostThreshold?: number;
 	#runnerLostTimeout?: NodeJS.Timeout;
 
 	// Event storage for resending
-	#eventHistory: protocol.EventWrapper[] = [];
 	#eventBacklogWarned: boolean = false;
 
 	// Command acknowledgment
@@ -255,7 +271,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
 	}
 
@@ -281,7 +304,30 @@ export class Runner {
 		// The server will send a StopActor command if it wants to fully stop
 	}
 
+	/**
+	 * Like stopActor but marks the actor for graceful destruction.
+	 * This ensures the engine destroys the actor instead of sleeping it.
+	 *
+	 * NOTE: If a drain (GoingAway) occurs after this is called but before the
+	 * stop completes, the engine's going_away flag overrides graceful_exit and
+	 * the actor will sleep instead of being destroyed. The destroy intent is
+	 * lost in this race. This is acceptable since the actor will be rescheduled
+	 * elsewhere and can be destroyed on the next wake.
+	 */
+	destroyActor(actorId: string, generation?: number) {
+		const actor = this.getActor(actorId, generation);
+		if (!actor) return;
+
+		actor.stopIntentSent = true;
+		this.#sendActorIntent(actorId, actor.generation, "stop");
+	}
+
 	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 +345,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),
+				});
+			});
 		}
 	}
 
@@ -412,6 +474,8 @@ export class Runner {
 			throw error;
 		}
 
+		// When changing SIGTERM/shutdown behavior, update
+		// website/src/content/docs/actors/versions.mdx (SIGTERM Handling section).
 		if (!this.#config.noAutoShutdown) {
 			if (!SIGNAL_HANDLERS.length) {
 				process.on("SIGTERM", async () => {
@@ -450,20 +514,20 @@ export class Runner {
 	// MARK: Shutdown
 	async shutdown(immediate: boolean, exit: boolean = false) {
 		// Prevent concurrent shutdowns
-		if (this.#shuttingDown) {
+		if (this.#shutdown) {
 			this.log?.debug({
 				msg: "shutdown already in progress, ignoring",
 			});
 			return;
 		}
-		this.#shuttingDown = true;
+		this.#shutdown = true;
+		this.#draining = !immediate;
 
 		this.log?.info({
 			msg: "starting shutdown",
 			immediate,
 			exit,
 		});
-		this.#shutdown = true;
 
 		// Clear reconnect timeout
 		if (this.#reconnectTimeout) {
@@ -477,12 +541,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);
@@ -504,8 +562,8 @@ export class Runner {
 		this.#kvRequests.clear();
 
 		// Close WebSocket
-		if (this.__webSocketReady()) {
-			const pegboardWebSocket = this.__pegboardWebSocket;
+		const pegboardWebSocket = this.getPegboardWebSocketIfReady();
+		if (pegboardWebSocket) {
 			if (immediate) {
 				// Stop immediately
 				pegboardWebSocket.close(1000, "pegboard.runner_shutdown");
@@ -566,7 +624,7 @@ export class Runner {
 			// the runner has already shut down
 			this.log?.debug({
 				msg: "no runner WebSocket to shutdown or already closed",
-				readyState: this.__pegboardWebSocket?.readyState,
+				readyState: this.#pegboardWebSocket?.readyState,
 			});
 		}
 
@@ -590,6 +648,9 @@ export class Runner {
 	 * - All actors are stopped
 	 * - The WebSocket connection is closed
 	 * - The shutdown timeout is reached (120 seconds)
+	 *
+	 * When changing this timeout, update
+	 * website/src/content/docs/actors/versions.mdx (SIGTERM Handling section).
 	 */
 	async #waitForActorsToStop(ws: WebSocket): Promise<void> {
 		const shutdownTimeout = 120_000; // 120 seconds
@@ -681,7 +742,7 @@ export class Runner {
 		const baseUrl = wsEndpoint.endsWith("/")
 			? wsEndpoint.slice(0, -1)
 			: wsEndpoint;
-		return `${baseUrl}/runners/connect?protocol_version=${PROTOCOL_VERSION}&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#config.runnerKey)}`;
+		return `${baseUrl}/runners/connect?protocol_version=${PROTOCOL_VERSION}&namespace=${encodeURIComponent(this.#config.namespace)}&runner_key=${encodeURIComponent(this.#runnerKey)}`;
 	}
 
 	// MARK: Runner protocol
@@ -691,14 +752,27 @@ export class Runner {
 			protocols.push(`rivet_token.${this.config.token}`);
 
 		const WS = await importWebSocket();
+
+		// Assertion to clear previous WebSocket
+		if (
+			this.#pegboardWebSocket &&
+			(this.#pegboardWebSocket.readyState === WS.CONNECTING ||
+				this.#pegboardWebSocket.readyState === WS.OPEN)
+		) {
+			this.log?.error(
+				"found duplicate pegboardWebSocket, closing previous",
+			);
+			this.#pegboardWebSocket.close(1000, "duplicate_websocket");
+		}
+
 		const ws = new WS(this.pegboardUrl, protocols) as any as WebSocket;
-		this.__pegboardWebSocket = ws;
+		this.#pegboardWebSocket = ws;
 
 		this.log?.info({
 			msg: "connecting",
 			endpoint: this.pegboardEndpoint,
 			namespace: this.#config.namespace,
-			runnerKey: this.#config.runnerKey,
+			runnerKey: this.#runnerKey,
 			hasToken: !!this.config.token,
 		});
 
@@ -738,10 +812,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 +828,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);
@@ -792,88 +851,113 @@ export class Runner {
 		});
 
 		ws.addEventListener("message", async (ev) => {
-			let buf: Uint8Array;
-			if (ev.data instanceof Blob) {
-				buf = new Uint8Array(await ev.data.arrayBuffer());
-			} else if (Buffer.isBuffer(ev.data)) {
-				buf = new Uint8Array(ev.data);
-			} else {
-				throw new Error(`expected binary data, got ${typeof ev.data}`);
-			}
+			try {
+				if (ws !== this.#pegboardWebSocket) {
+					this.log?.debug({
+						msg: "ignoring runner message from stale websocket",
+					});
+					return;
+				}
 
-			// Parse message
-			const message = protocol.decodeToClient(buf);
-			this.log?.debug({
-				msg: "received runner message",
-				data: stringifyToClient(message),
-			});
+				let buf: Uint8Array;
+				if (ev.data instanceof Blob) {
+					buf = new Uint8Array(await ev.data.arrayBuffer());
+				} else if (Buffer.isBuffer(ev.data)) {
+					buf = new Uint8Array(ev.data);
+				} else {
+					throw new Error(`expected binary data, got ${typeof ev.data}`);
+				}
+
+				await this.#injectLatency();
 
-			// Handle message
-			if (message.tag === "ToClientInit") {
-				const init = message.val;
+				// Parse message
+				const message = protocol.decodeToClient(buf);
+				this.log?.debug({
+					msg: "received runner message",
+					data: stringifyToClient(message),
+				});
 
-				if (this.runnerId !== init.runnerId) {
-					this.runnerId = init.runnerId;
+				// Handle message
+				if (message.tag === "ToClientInit") {
+					const init = message.val;
 
-					// Clear history if runner id changed
-					this.#eventHistory.length = 0;
-				}
+					if (this.runnerId !== init.runnerId) {
+						this.runnerId = init.runnerId;
 
-				// Store the runner lost threshold from metadata
-				this.#runnerLostThreshold = init.metadata?.runnerLostThreshold
-					? Number(init.metadata.runnerLostThreshold)
-					: undefined;
+						// Clear actors if runner id changed
+						this.#stopAllActors();
+					}
 
-				this.log?.info({
-					msg: "received init",
-					lastEventIdx: init.lastEventIdx,
-					runnerLostThreshold: this.#runnerLostThreshold,
-				});
+					this.#protocolMetadata = init.metadata;
 
-				// Resend pending events
-				this.#processUnsentKvRequests();
-				this.#resendUnacknowledgedEvents(init.lastEventIdx);
-				this.#tunnel?.resendBufferedEvents();
-
-				this.#config.onConnected();
-			} else if (message.tag === "ToClientCommands") {
-				const commands = message.val;
-				this.#handleCommands(commands);
-			} else if (message.tag === "ToClientAckEvents") {
-				this.#handleAckEvents(message.val);
-			} else if (message.tag === "ToClientKvResponse") {
-				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");
-			} else {
-				unreachable(message);
+					this.log?.info({
+						msg: "received init",
+						protocolMetadata: this.#protocolMetadata,
+					});
+
+					// Resend pending events
+					this.#processUnsentKvRequests();
+					this.#resendUnacknowledgedEvents();
+					this.#tunnel?.resendBufferedEvents();
+
+					this.#config.onConnected();
+				} else if (message.tag === "ToClientCommands") {
+					const commands = message.val;
+					this.#handleCommands(commands);
+				} else if (message.tag === "ToClientAckEvents") {
+					this.#handleAckEvents(message.val);
+				} else if (message.tag === "ToClientKvResponse") {
+					const kvResponse = message.val;
+					this.#handleKvResponse(kvResponse);
+				} else if (message.tag === "ToClientTunnelMessage") {
+					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);
+				}
+			} catch (error) {
+				if (this.#shutdown || ws.readyState !== ws.OPEN) {
+					this.log?.debug({
+						msg: "ignoring runner websocket message during shutdown",
+						error: stringifyError(error),
+						readyState: ws.readyState,
+					});
+					return;
+				}
+
+				this.log?.error({
+					msg: "failed to decode runner websocket message",
+					error: stringifyError(error),
+				});
+				try {
+					ws.close(1011, "runner.invalid_frame");
+				} catch (closeError) {
+					this.log?.debug({
+						msg: "failed closing runner websocket after decode error",
+						error: stringifyError(closeError),
+					});
+				}
 			}
 		});
 
 		ws.addEventListener("error", (ev) => {
 			this.log?.error({
-				msg: `WebSocket error: ${ev.error}`,
+				msg: `WebSocket error: ${stringifyError(ev.error)}`,
 			});
 
 			if (!this.#shutdown) {
-				// Start runner lost timeout if we have a threshold and are not shutting down
-				if (
-					!this.#runnerLostTimeout &&
-					this.#runnerLostThreshold &&
-					this.#runnerLostThreshold > 0
-				) {
-					this.log?.info({
-						msg: "starting runner lost timeout",
-						seconds: this.#runnerLostThreshold / 1000,
-					});
-					this.#runnerLostTimeout = setTimeout(() => {
-						this.#stopAllActors();
-					}, this.#runnerLostThreshold);
-				}
+				this.#startRunnerLostTimeout();
 
 				// Attempt to reconnect if not stopped
 				this.#scheduleReconnect();
@@ -881,22 +965,20 @@ export class Runner {
 		});
 
 		ws.addEventListener("close", async (ev) => {
-			const closeError = parseWebSocketCloseReason(ev.reason);
-			if (
-				closeError?.group === "ws" &&
-				closeError?.error === "eviction"
-			) {
-				this.log?.info("runner websocket evicted");
-
-				this.#config.onDisconnected(ev.code, ev.reason);
-
-				await this.shutdown(true);
-			} else {
+			if (this.#pegboardWebSocket === ws) {
+				this.#pegboardWebSocket = undefined;
+			}
+			if (!this.#shutdown) {
+				const closeError = parseWebSocketCloseReason(ev.reason);
 				if (
-					closeError?.group === "pegboard" &&
-					closeError?.error === "runner_shutdown"
+					closeError?.group === "ws" &&
+					closeError?.error === "eviction"
 				) {
-					this.log?.info("runner shutdown");
+					this.log?.info("runner websocket evicted");
+
+					this.#config.onDisconnected(ev.code, ev.reason);
+
+					await this.shutdown(true);
 				} else {
 					this.log?.warn({
 						msg: "runner disconnected",
@@ -904,45 +986,52 @@ export class Runner {
 						reason: ev.reason.toString(),
 						closeError,
 					});
-				}
 
-				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);
-				this.#ackInterval = undefined;
-			}
+					this.#config.onDisconnected(ev.code, ev.reason);
+				}
 
-			if (!this.#shutdown) {
-				// Start runner lost timeout if we have a threshold and are not shutting down
-				if (
-					!this.#runnerLostTimeout &&
-					this.#runnerLostThreshold &&
-					this.#runnerLostThreshold > 0
-				) {
-					this.log?.info({
-						msg: "starting runner lost timeout",
-						seconds: this.#runnerLostThreshold / 1000,
-					});
-					this.#runnerLostTimeout = setTimeout(() => {
-						this.#stopAllActors();
-					}, this.#runnerLostThreshold);
+				// Clear ack interval on close
+				if (this.#ackInterval) {
+					clearInterval(this.#ackInterval);
+					this.#ackInterval = undefined;
 				}
 
+				this.#startRunnerLostTimeout();
+
 				// Attempt to reconnect if not stopped
 				this.#scheduleReconnect();
+			} else {
+				this.log?.info("websocket closed");
+
+				this.#config.onDisconnected(ev.code, ev.reason);
 			}
 		});
 	}
 
+	#startRunnerLostTimeout() {
+		// Start runner lost timeout if we have a threshold and are not shutting down
+		if (
+			!this.#runnerLostTimeout &&
+			this.#protocolMetadata &&
+			this.#protocolMetadata.runnerLostThreshold > 0
+		) {
+			this.log?.info({
+				msg: "starting runner lost timeout",
+				seconds: this.#protocolMetadata.runnerLostThreshold / 1000n,
+			});
+			this.#runnerLostTimeout = setTimeout(() => {
+				try {
+					this.#handleLost();
+				} catch (err) {
+					this.log?.error({
+						msg: "error handling runner lost",
+						error: stringifyError(err),
+					});
+				}
+			}, Number(this.#protocolMetadata.runnerLostThreshold));
+		}
+	}
+
 	#handleCommands(commands: protocol.ToClientCommands) {
 		this.log?.info({
 			msg: "received commands",
@@ -952,52 +1041,89 @@ 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;
+		const originalTotalEvents = Array.from(this.#actors).reduce(
+			(s, [_, actor]) => s + actor.eventHistory.length,
+			0,
+		);
+
+		for (const [_, actor] of this.#actors) {
+			const 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);
+
+		const 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 +1139,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 +1220,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 +1231,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 +1248,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 +1276,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") {
@@ -1154,7 +1287,9 @@ export class Runner {
 			actorState = {
 				tag: "ActorStateStopped",
 				val: {
-					code: protocol.StopCode.Ok,
+					code: actor.stopIntentSent || this.#draining
+						? protocol.StopCode.Ok
+						: protocol.StopCode.Error,
 					message: null,
 				},
 			};
@@ -1163,14 +1298,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 +1322,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 +1342,7 @@ export class Runner {
 		this.__sendToServer({
 			tag: "ToServerAckCommands",
 			val: {
-				lastCommandIdx: BigInt(this.#lastCommandIdx),
+				lastCommandCheckpoints,
 			},
 		});
 	}
@@ -1480,6 +1626,31 @@ export class Runner {
 		await this.#sendKvRequest(actorId, requestData);
 	}
 
+	async kvDeleteRange(
+		actorId: string,
+		start: Uint8Array,
+		end: Uint8Array,
+	): Promise<void> {
+		const startKey: protocol.KvKey = start.buffer.slice(
+			start.byteOffset,
+			start.byteOffset + start.byteLength,
+		) as ArrayBuffer;
+		const endKey: protocol.KvKey = end.buffer.slice(
+			end.byteOffset,
+			end.byteOffset + end.byteLength,
+		) as ArrayBuffer;
+
+		const requestData: protocol.KvRequestData = {
+			tag: "KvDeleteRangeRequest",
+			val: {
+				start: startKey,
+				end: endKey,
+			},
+		};
+
+		await this.#sendKvRequest(actorId, requestData);
+	}
+
 	async kvDrop(actorId: string): Promise<void> {
 		const requestData: protocol.KvRequestData = {
 			tag: "KvDropRequest",
@@ -1495,14 +1666,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,
@@ -1540,7 +1712,7 @@ export class Runner {
 
 			this.#kvRequests.set(requestId, requestEntry);
 
-			if (this.__webSocketReady()) {
+			if (this.getPegboardWebSocketIfReady()) {
 				// Send immediately
 				this.#sendSingleKvRequest(requestId);
 			}
@@ -1573,7 +1745,7 @@ export class Runner {
 	}
 
 	#processUnsentKvRequests() {
-		if (!this.__webSocketReady()) {
+		if (!this.getPegboardWebSocketIfReady()) {
 			return;
 		}
 
@@ -1590,14 +1762,23 @@ export class Runner {
 		}
 	}
 
+	/** Resolves after the configured debug latency, or immediately if none. */
+	#injectLatency(): Promise<void> {
+		const ms = this.#config.debugLatencyMs;
+		if (!ms) return Promise.resolve();
+		return new Promise((resolve) => setTimeout(resolve, ms));
+	}
+
 	/** Asserts WebSocket exists and is ready. */
-	__webSocketReady(): this is this & {
-		__pegboardWebSocket: NonNullable<Runner["__pegboardWebSocket"]>;
-	} {
-		return (
-			!!this.__pegboardWebSocket &&
-			this.__pegboardWebSocket.readyState === 1
-		);
+	getPegboardWebSocketIfReady(): WebSocket | undefined {
+		if (
+			!!this.#pegboardWebSocket &&
+			this.#pegboardWebSocket.readyState === 1
+		) {
+			return this.#pegboardWebSocket;
+		} else {
+			return undefined;
+		}
 	}
 
 	__sendToServer(message: protocol.ToServer) {
@@ -1607,13 +1788,19 @@ export class Runner {
 		});
 
 		const encoded = protocol.encodeToServer(message);
-		if (this.__webSocketReady()) {
-			this.__pegboardWebSocket.send(encoded);
-		} else {
-			this.log?.error({
-				msg: "WebSocket not available or not open for sending data",
-			});
-		}
+
+		// Normally synchronous. When debugLatencyMs is set, the send is
+		// deferred but message order is preserved.
+		this.#injectLatency().then(() => {
+			const pegboardWebSocket = this.getPegboardWebSocketIfReady();
+			if (pegboardWebSocket) {
+				pegboardWebSocket.send(encoded);
+			} else {
+				this.log?.error({
+					msg: "WebSocket not available or not open for sending data",
+				});
+			}
+		});
 	}
 
 	sendHibernatableWebSocketMessageAck(
@@ -1669,6 +1856,7 @@ export class Runner {
 			tag: "ToServerlessServerInit",
 			val: {
 				runnerId: this.runnerId,
+				runnerProtocolVersion: PROTOCOL_VERSION,
 			},
 		});
 
@@ -1699,27 +1887,41 @@ export class Runner {
 			msg: `Scheduling reconnect attempt ${this.#reconnectAttempt + 1} in ${delay}ms`,
 		});
 
-		this.#reconnectTimeout = setTimeout(async () => {
+		if (this.#reconnectTimeout) {
+			this.log?.info(
+				"clearing previous reconnect timeout in schedule reconnect",
+			);
+			clearTimeout(this.#reconnectTimeout);
+		}
+
+		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
@@ -1752,4 +1954,8 @@ export class Runner {
 			//this.#log?.log(`Cleaned up ${toDelete.length} expired KV requests`);
 		}
 	}
+
+	getProtocolMetadata(): protocol.ProtocolMetadata | undefined {
+		return this.#protocolMetadata;
+	}
 }