From 7a186f675311a0f8088d3024cb7706dc5cbe6cca Mon Sep 17 00:00:00 2001
From: Nathan Flurry <developer@nathanflurry.com>
Date: Wed, 5 Nov 2025 16:35:23 -0800
Subject: [PATCH] fix(runner): do not send websocket close on tunnel shutdown
 since pegboard-runner takes care of this

---
 engine/sdks/typescript/runner/src/tunnel.ts   | 33 ++++++++++++++++---
 .../runner/src/websocket-tunnel-adapter.ts    | 19 ++++++++---
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/engine/sdks/typescript/runner/src/tunnel.ts b/engine/sdks/typescript/runner/src/tunnel.ts
index 29bfbe79d0..5c62b6b5dd 100644
--- a/engine/sdks/typescript/runner/src/tunnel.ts
+++ b/engine/sdks/typescript/runner/src/tunnel.ts
@@ -23,6 +23,12 @@ interface PendingTunnelMessage {
 	requestIdStr: string;
 }
 
+class RunnerShutdownError extends Error {
+	constructor() {
+		super("Runner shut down");
+	}
+}
+
 export class Tunnel {
 	#runner: Runner;
 
@@ -49,20 +55,31 @@ export class Tunnel {
 	}
 
 	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 @@ export class Tunnel {
 				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);
diff --git a/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts b/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts
index 3eee2ce072..00e3a7e8a7 100644
--- a/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts
+++ b/engine/sdks/typescript/runner/src/websocket-tunnel-adapter.ts
@@ -190,14 +190,23 @@ export class WebSocketTunnelAdapter {
 	}
 
 	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 @@ export class WebSocketTunnelAdapter {
 		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