@rivetkit/engine-runner 2.0.24-rc.1 → 2.0.24

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rivetkit/engine-runner",
-  "version": "2.0.24-rc.1",
+  "version": "2.0.24",
   "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.24-rc.1"
+    "@rivetkit/engine-runner-protocol": "2.0.24"
   },
   "devDependencies": {
     "@types/node": "^22.18.1",

package/src/actor.ts

@@ -0,0 +1,196 @@
+import type * as protocol from "@rivetkit/engine-runner-protocol";
+import { logger } from "./log";
+import type { PendingRequest } from "./tunnel";
+import { arraysEqual, idToStr, promiseWithResolvers } from "./utils";
+import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
+
+export interface ActorConfig {
+	name: string;
+	key: string | null;
+	createTs: bigint;
+	input: Uint8Array | null;
+}
+
+export class RunnerActor {
+	actorId: string;
+	generation: number;
+	config: ActorConfig;
+	pendingRequests: Array<{
+		gatewayId: protocol.GatewayId;
+		requestId: protocol.RequestId;
+		request: PendingRequest;
+	}> = [];
+	webSockets: Array<{
+		gatewayId: protocol.GatewayId;
+		requestId: protocol.RequestId;
+		ws: WebSocketTunnelAdapter;
+	}> = [];
+	actorStartPromise: ReturnType<typeof promiseWithResolvers<void>>;
+
+	/**
+	 * If restoreHibernatingRequests has been called. This is used to assert
+	 * that the caller is implemented correctly.
+	 **/
+	hibernationRestored: boolean = false;
+
+	constructor(
+		actorId: string,
+		generation: number,
+		config: ActorConfig,
+		/**
+		 * List of hibernating requests provided by the gateway on actor start.
+		 * This represents the WebSocket connections that the gateway knows about.
+		 **/
+		public hibernatingRequests: readonly protocol.HibernatingRequest[],
+	) {
+		this.actorId = actorId;
+		this.generation = generation;
+		this.config = config;
+		this.actorStartPromise = promiseWithResolvers();
+	}
+
+	// Pending request methods
+	getPendingRequest(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+	): PendingRequest | undefined {
+		return this.pendingRequests.find(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		)?.request;
+	}
+
+	createPendingRequest(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+		clientMessageIndex: number,
+	) {
+		const exists =
+			this.getPendingRequest(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({
+				msg: "attempting to set pending request twice, replacing existing",
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
+			});
+			// Delete existing pending request before adding the new one
+			this.deletePendingRequest(gatewayId, requestId);
+		}
+		this.pendingRequests.push({
+			gatewayId,
+			requestId,
+			request: {
+				resolve: () => {},
+				reject: () => {},
+				actorId: this.actorId,
+				gatewayId: gatewayId,
+				requestId: requestId,
+				clientMessageIndex,
+			},
+		});
+		logger()?.debug({
+			msg: "added pending request",
+			gatewayId: idToStr(gatewayId),
+			requestId: idToStr(requestId),
+			length: this.pendingRequests.length,
+		});
+	}
+
+	createPendingRequestWithStreamController(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+		clientMessageIndex: number,
+		streamController: ReadableStreamDefaultController<Uint8Array>,
+	) {
+		const exists =
+			this.getPendingRequest(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({
+				msg: "attempting to set pending request twice, replacing existing",
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
+			});
+			// Delete existing pending request before adding the new one
+			this.deletePendingRequest(gatewayId, requestId);
+		}
+		this.pendingRequests.push({
+			gatewayId,
+			requestId,
+			request: {
+				resolve: () => {},
+				reject: () => {},
+				actorId: this.actorId,
+				gatewayId: gatewayId,
+				requestId: requestId,
+				clientMessageIndex,
+				streamController,
+			},
+		});
+		logger()?.debug({
+			msg: "added pending request with stream controller",
+			gatewayId: idToStr(gatewayId),
+			requestId: idToStr(requestId),
+			length: this.pendingRequests.length,
+		});
+	}
+
+	deletePendingRequest(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+	) {
+		const index = this.pendingRequests.findIndex(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		);
+		if (index !== -1) {
+			this.pendingRequests.splice(index, 1);
+			logger()?.debug({
+				msg: "removed pending request",
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
+				length: this.pendingRequests.length,
+			});
+		}
+	}
+
+	// WebSocket methods
+	getWebSocket(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+	): WebSocketTunnelAdapter | undefined {
+		return this.webSockets.find(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		)?.ws;
+	}
+
+	setWebSocket(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+		ws: WebSocketTunnelAdapter,
+	) {
+		const exists = this.getWebSocket(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({ msg: "attempting to set websocket twice" });
+			return;
+		}
+		this.webSockets.push({ gatewayId, requestId, ws });
+	}
+
+	deleteWebSocket(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+	) {
+		const index = this.webSockets.findIndex(
+			(entry) =>
+				arraysEqual(entry.gatewayId, gatewayId) &&
+				arraysEqual(entry.requestId, requestId),
+		);
+		if (index !== -1) {
+			this.webSockets.splice(index, 1);
+		}
+	}
+}