@rivetkit/engine-runner 2.0.33 → 2.0.34-rc.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rivetkit/engine-runner",
-  "version": "2.0.33",
+  "version": "2.0.34-rc.2",
   "type": "module",
   "files": [
     "dist",
@@ -21,7 +21,8 @@
     "uuid": "^12.0.0",
     "pino": "^9.9.5",
     "ws": "^8.18.3",
-    "@rivetkit/engine-runner-protocol": "2.0.33"
+    "@rivetkit/virtual-websocket": "2.0.33",
+    "@rivetkit/engine-runner-protocol": "2.0.34-rc.2"
   },
   "devDependencies": {
     "@types/node": "^22.18.1",

package/src/mod.ts

@@ -201,7 +201,7 @@ export class Runner {
 	#actors: Map<string, RunnerActor> = new Map();
 
 	// WebSocket
-	__pegboardWebSocket?: WebSocket;
+	#pegboardWebSocket?: WebSocket;
 	runnerId?: string;
 	#started: boolean = false;
 	#shutdown: boolean = false;
@@ -528,8 +528,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");
@@ -590,7 +590,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,
 			});
 		}
 
@@ -715,8 +715,21 @@ 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",
@@ -1621,7 +1634,7 @@ export class Runner {
 
 			this.#kvRequests.set(requestId, requestEntry);
 
-			if (this.__webSocketReady()) {
+			if (this.getPegboardWebSocketIfReady()) {
 				// Send immediately
 				this.#sendSingleKvRequest(requestId);
 			}
@@ -1654,7 +1667,7 @@ export class Runner {
 	}
 
 	#processUnsentKvRequests() {
-		if (!this.__webSocketReady()) {
+		if (!this.getPegboardWebSocketIfReady()) {
 			return;
 		}
 
@@ -1672,13 +1685,15 @@ export class Runner {
 	}
 
 	/** 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) {
@@ -1688,8 +1703,9 @@ export class Runner {
 		});
 
 		const encoded = protocol.encodeToServer(message);
-		if (this.__webSocketReady()) {
-			this.__pegboardWebSocket.send(encoded);
+		const pegboardWebSocket = this.getPegboardWebSocketIfReady();
+		if (pegboardWebSocket) {
+			pegboardWebSocket.send(encoded);
 		} else {
 			this.log?.error({
 				msg: "WebSocket not available or not open for sending data",
@@ -1781,6 +1797,13 @@ export class Runner {
 			msg: `Scheduling reconnect attempt ${this.#reconnectAttempt + 1} in ${delay}ms`,
 		});
 
+		if (this.#reconnectTimeout) {
+			this.log?.info(
+				"clearing previous reconnect timeout in schedule reconnect",
+			);
+			clearTimeout(this.#reconnectTimeout);
+		}
+
 		this.#reconnectTimeout = setTimeout(() => {
 			if (!this.#shutdown) {
 				this.#reconnectAttempt++;

package/src/tunnel.ts

@@ -395,11 +395,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,
@@ -478,7 +478,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),
@@ -917,7 +917,7 @@ export class Tunnel {
 	) {
 		// 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.