fix(runner): do not send websocket close on tunnel shutdown since pegboard-runner takes care of this

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

@@ -2,14 +2,14 @@
 import type { MessageId, RequestId } from "@rivetkit/engine-runner-protocol";
 import type { Logger } from "pino";
 import { stringify as uuidstringify, v4 as uuidv4 } from "uuid";
 import { logger } from "./log";
 import type { ActorInstance, Runner } from "./mod";
 import { unreachable } from "./utils";
 import { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
 
 const GC_INTERVAL = 60000; // 60 seconds
 const MESSAGE_ACK_TIMEOUT = 5000; // 5 seconds
 const WEBSOCKET_STATE_PERSIST_TIMEOUT = 30000; // 30 seconds
 
 interface PendingRequest {
 	resolve: (response: Response) => void;
@@ -23,6 +23,12 @@
 	requestIdStr: string;
 }
 
+class RunnerShutdownError extends Error {
+	constructor() {
+		super("Runner shut down");
+	}
+}
+
 export class Tunnel {
 	#runner: Runner;
 
@@ -49,20 +55,31 @@
 	}
 
 	shutdown() {
+		// NOTE: Pegboard WS already closed at this point, cannot send
+		// anything. All teardown logic is handled by pegboard-runner.
+
 		if (this.#gcInterval) {
 			clearInterval(this.#gcInterval);
 			this.#gcInterval = undefined;
 		}
 
 		// Reject all pending requests
+		//
+		// RunnerShutdownError will be explicitly ignored
 		for (const [_, request] of this.#actorPendingRequests) {
-			request.reject(new Error("Tunnel shutting down"));
+			request.reject(new RunnerShutdownError());
 		}
 		this.#actorPendingRequests.clear();
 
 		// Close all WebSockets
+		//
+		// The WebSocket close event with retry is automatically sent when the
+		// runner WS closes, so we only need to notify the client that the WS
+		// closed:
+		// https://github.com/rivet-dev/rivet/blob/00d4f6a22da178a6f8115e5db50d96c6f8387c2e/engine/packages/pegboard-runner/src/lib.rs#L157
 		for (const [_, ws] of this.#actorWebSockets) {
-			ws.__closeWithRetry();
+			// TODO: Trigger close event, but do not send anything over the tunnel
+			ws.__closeWithoutCallback(1000, "ws.tunnel_shutdown");
 		}
 		this.#actorWebSockets.clear();
 	}
@@ -407,8 +424,16 @@
 				await this.#sendResponse(requestId, response);
 			}
 		} catch (error) {
-			this.log?.error({ msg: "error handling request", error });
-			this.#sendResponseError(requestId, 500, "Internal Server Error");
+			if (error instanceof RunnerShutdownError) {
+				this.log?.debug({ msg: "catught runner shutdown error" });
+			} else {
+				this.log?.error({ msg: "error handling request", error });
+				this.#sendResponseError(
+					requestId,
+					500,
+					"Internal Server Error",
+				);
+			}
 		} finally {
 			// Clean up request tracking
 			const actor = this.#runner.getActor(req.actorId);
@@ -604,8 +629,8 @@
 					headerInit[k] = v;
 				}
 			}
 			headerInit["Upgrade"] = "websocket";
 			headerInit["Connection"] = "Upgrade";
 
 			const request = new Request(`http://localhost${open.path}`, {
 				method: "GET",

engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts

@@ -4,7 +4,7 @@
 import { logger } from "./log";
 
 export class WebSocketTunnelAdapter {
 	#webSocketId: string;
 	#readyState: number = 0; // CONNECTING
 	#eventListeners: Map<string, Set<(event: any) => void>> = new Map();
 	#onopen: ((this: any, ev: any) => any) | null = null;
@@ -190,14 +190,23 @@
 	}
 
 	close(code?: number, reason?: string): void {
-		this.closeInner(code, reason);
+		this.closeInner(code, reason, false, true);
 	}
 
 	__closeWithRetry(code?: number, reason?: string): void {
-		this.closeInner(code, reason, true);
+		this.closeInner(code, reason, true, true);
 	}
 
-	closeInner(code?: number, reason?: string, retry: boolean = false): void {
+	__closeWithoutCallback(code?: number, reason?: string): void {
+		this.closeInner(code, reason, false, false);
+	}
+
+	closeInner(
+		code: number | undefined,
+		reason: string | undefined,
+		retry: boolean,
+		callback: boolean,
+	): void {
 		if (
 			this.#readyState === 2 || // CLOSING
 			this.#readyState === 3 // CLOSED
@@ -208,7 +217,9 @@
 		this.#readyState = 2; // CLOSING
 
 		// Send close through tunnel
-		this.#closeCallback(code, reason, retry);
+		if (callback) {
+			this.#closeCallback(code, reason, retry);
+		}
 
 		// Update state and fire event
 		this.#readyState = 3; // CLOSED
@@ -227,7 +238,7 @@
 	addEventListener(
 		type: string,
 		listener: (event: any) => void,
 		options?: boolean | any,
 	): void {
 		if (typeof listener === "function") {
 			let listeners = this.#eventListeners.get(type);
@@ -245,7 +256,7 @@
 	removeEventListener(
 		type: string,
 		listener: (event: any) => void,
 		options?: boolean | any,
 	): void {
 		if (typeof listener === "function") {
 			const listeners = this.#eventListeners.get(type);
@@ -255,7 +266,7 @@
 		}
 	}
 
 	dispatchEvent(event: any): boolean {
 		// Simple implementation
 		return true;
 	}