fix(desktop): refresh relay JWT on WS reconnect and stop retrying definitive 403s

apps/desktop/src/renderer/lib/auth-client.ts

@@ -8,6 +8,7 @@ import {
 } from "better-auth/client/plugins";
 import { createAuthClient } from "better-auth/react";
 import { env } from "renderer/env.renderer";
+import { decodeJwtExpiresAtMs } from "renderer/lib/jwt-expiry";
 
 let authToken: string | null = null;
 
@@ -20,15 +21,70 @@ export function getAuthToken(): string | null {
 }
 
 let jwt: string | null = null;
+let jwtExpiresAtMs: number | null = null;
+let jwtGeneration = 0;
+let jwtRefreshInFlight: Promise<string | null> | null = null;
+
+// Refresh ahead of expiry so a token handed to a WS URL is still valid by the
+// time the relay verifies it.
+const JWT_REFRESH_LEEWAY_MS = 60_000;
 
 export function setJwt(token: string | null) {
 	jwt = token;
+	jwtGeneration++;
+	jwtExpiresAtMs = token ? decodeJwtExpiresAtMs(token) : null;
+}
+
+function jwtIsFresh(): boolean {
+	if (!jwt) return false;
+	if (jwtExpiresAtMs === null) return true;
+	return Date.now() < jwtExpiresAtMs - JWT_REFRESH_LEEWAY_MS;
 }
 
 export function getJwt(): string | null {
+	// Relay JWTs rotate hourly, but this cache only updates when some API
+	// response happens to carry `set-auth-jwt`. Sync callers (WS URL builders,
+	// reconnect loops) can't await a refresh, so kick one off in the background
+	// and let their next attempt pick up the fresh token.
+	if (jwt && !jwtIsFresh()) void ensureFreshJwt();
 	return jwt;
 }
 
+/**
+ * Returns the cached JWT if it's still valid, otherwise mints a fresh one from
+ * better-auth's `/token` endpoint (deduped across concurrent callers). Falls
+ * back to the stale cached token if the refresh fails.
+ */
+export async function ensureFreshJwt(): Promise<string | null> {
+	if (jwtIsFresh()) return jwt;
+	if (!jwtRefreshInFlight) {
+		const generationAtStart = jwtGeneration;
+		jwtRefreshInFlight = authClient
+			.$fetch<{ token?: string }>("/token")
+			.then((res) => {
+				const token = res.data?.token;
+				// Apply only if nothing (logout, a set-auth-jwt response header)
+				// replaced the cached token while this request was in flight.
+				if (
+					typeof token === "string" &&
+					token &&
+					jwtGeneration === generationAtStart
+				) {
+					setJwt(token);
+				}
+				return jwt;
+			})
+			.catch((err) => {
+				console.warn("[auth] JWT refresh failed:", err);
+				return jwt;
+			})
+			.finally(() => {
+				jwtRefreshInFlight = null;
+			});
+	}
+	return jwtRefreshInFlight;
+}
+
 /**
  * Better Auth client for Electron desktop app.
  *

apps/desktop/src/renderer/lib/jwt-expiry.test.ts

@@ -0,0 +1,28 @@
+import { describe, expect, it } from "bun:test";
+import { decodeJwtExpiresAtMs } from "./jwt-expiry";
+
+function makeJwt(payload: object): string {
+	const encode = (obj: object) =>
+		btoa(JSON.stringify(obj))
+			.replace(/\+/g, "-")
+			.replace(/\//g, "_")
+			.replace(/=+$/, "");
+	return `${encode({ alg: "EdDSA" })}.${encode(payload)}.sig`;
+}
+
+describe("decodeJwtExpiresAtMs", () => {
+	it("returns exp in epoch milliseconds", () => {
+		expect(decodeJwtExpiresAtMs(makeJwt({ exp: 1_800_000_000 }))).toBe(
+			1_800_000_000_000,
+		);
+	});
+
+	it("returns null when exp is missing", () => {
+		expect(decodeJwtExpiresAtMs(makeJwt({ sub: "u1" }))).toBeNull();
+	});
+
+	it("returns null for malformed tokens", () => {
+		expect(decodeJwtExpiresAtMs("not-a-jwt")).toBeNull();
+		expect(decodeJwtExpiresAtMs("a.%%%.c")).toBeNull();
+	});
+});

apps/desktop/src/renderer/lib/jwt-expiry.ts

@@ -0,0 +1,13 @@
+/** Reads the `exp` claim (as epoch ms) from a JWT without verifying it. */
+export function decodeJwtExpiresAtMs(token: string): number | null {
+	try {
+		const payloadPart = token.split(".")[1];
+		if (!payloadPart) return null;
+		const payload = JSON.parse(
+			atob(payloadPart.replace(/-/g, "+").replace(/_/g, "/")),
+		) as { exp?: unknown };
+		return typeof payload.exp === "number" ? payload.exp * 1000 : null;
+	} catch {
+		return null;
+	}
+}

apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts

@@ -3,6 +3,7 @@ import {
 	type RelayAffinityProbe,
 } from "@superset/workspace-client";
 import type { Terminal as XTerm } from "@xterm/xterm";
+import { ensureFreshJwt } from "renderer/lib/auth-client";
 import { posthog } from "renderer/lib/posthog";
 import { classifyTerminalFailure } from "./terminalConnectionDiagnostics";
 import { createWriteCoalescer, type WriteCoalescer } from "./write-coalescer";
@@ -78,6 +79,10 @@ export interface TerminalTransport {
 	_writeCoalescer: WriteCoalescer | null;
 	/** Internal: last `_whoowns` probe, used to explain a failed connection. */
 	_lastProbe: RelayAffinityProbe | null;
+	/** Internal: bumped per connect() so async preflight callbacks from a
+	 * superseded attempt (reconnectNow re-entering the same URL) can't open a
+	 * duplicate socket or fatally terminate the newer attempt. */
+	_connectEpoch: number;
 }
 
 const MAX_LOG_ENTRIES = 200;
@@ -179,6 +184,7 @@ export function createTransport(): TerminalTransport {
 		_resumeListener: null,
 		_writeCoalescer: null,
 		_lastProbe: null,
+		_connectEpoch: 0,
 	};
 }
 
@@ -361,6 +367,7 @@ export function connect(
 	}
 
 	cancelReconnect(transport);
+	const epoch = ++transport._connectEpoch;
 	transport.currentUrl = wsUrl;
 	transport._terminal = terminal;
 	// Recreate per connect so the coalescer always targets the current
@@ -376,23 +383,25 @@ export function connect(
 		? appendQueryParam(wsUrl, "replay", "0")
 		: wsUrl;
 
-	const openSocket = () => {
-		// Bail if the transport raced into a different URL or was disconnected
-		// while the pre-flight was in flight.
+	const openSocket = (targetUrl: string) => {
+		// Bail if the transport raced into a different URL, was disconnected, or
+		// a newer connect() superseded this attempt while the pre-flight was in
+		// flight.
 		if (
+			transport._connectEpoch !== epoch ||
 			transport.currentUrl !== wsUrl ||
 			transport.connectionState !== "connecting"
 		) {
 			return;
 		}
 		let socket: WebSocket;
 		try {
-			socket = new WebSocket(actualUrl);
+			socket = new WebSocket(targetUrl);
 		} catch (err) {
 			pushLog(
 				transport,
 				"error",
-				`WebSocket construction failed for ${formatWsEndpoint(actualUrl)}: ${
+				`WebSocket construction failed for ${formatWsEndpoint(targetUrl)}: ${
 					err instanceof Error ? err.message : String(err)
 				}`,
 			);
@@ -409,16 +418,53 @@ export function connect(
 	};
 
 	// Pre-flight `_whoowns` to pin fly edge affinity before the WS upgrade (see
-	// primeRelayAffinity); skip for non-/hosts URLs. Probe `wsUrl`, not the
-	// `replay`-carrying `actualUrl`, and keep the result to explain a failure.
+	// primeRelayAffinity); skip for non-/hosts URLs. Keep the result to explain
+	// a failure.
 	if (isRelayHostUrl(wsUrl)) {
 		transport._lastProbe = null;
-		void primeRelayAffinity(wsUrl).then((probe) => {
-			transport._lastProbe = probe;
-			openSocket();
+		// Relay URLs carry the user JWT, which rotates hourly. The URL was signed
+		// at render time, so reconnect attempts (this function re-entered from
+		// scheduleReconnect with the same currentUrl) would otherwise redial with
+		// an expired token forever — re-sign every attempt with a fresh one.
+		void ensureFreshJwt().then((token) => {
+			const signedUrl = token
+				? appendQueryParam(actualUrl, "token", token)
+				: actualUrl;
+			return primeRelayAffinity(signedUrl).then((probe) => {
+				transport._lastProbe = probe;
+				// 403 is a definitive access denial (the token is fresh), not a
+				// transient failure — retrying would hammer the relay with the same
+				// answer. Stop and surface it; a manual reconnect re-enters connect().
+				if (probe?.status === 403) {
+					if (
+						transport._connectEpoch !== epoch ||
+						transport.currentUrl !== wsUrl ||
+						transport.connectionState !== "connecting"
+					) {
+						return;
+					}
+					const diagnosis = classifyTerminalFailure(probe, true);
+					pushLog(
+						transport,
+						"error",
+						`Connection refused for ${formatWsEndpoint(wsUrl)}: ${diagnosis.message} Not retrying.`,
+					);
+					posthog.capture("terminal_connect_failed", {
+						endpoint: formatWsEndpoint(wsUrl),
+						preflight_status: probe.status,
+						tunnel_region: probe.region,
+						reconnect_attempts: transport._reconnectAttempt,
+						category: diagnosis.category,
+					});
+					transport._terminated = true;
+					setConnectionState(transport, "closed");
+					return;
+				}
+				openSocket(signedUrl);
+			});
 		});
 	} else {
-		openSocket();
+		openSocket(actualUrl);
 	}
 }
 

apps/relay/src/index.ts

@@ -293,9 +293,12 @@ app.get(
 
 // ── Pre-flight for WS replay (host hits this once before opening WS to a host) ─
 
-// Pre-flight for WS upgrade routing. Requires a valid JWT (no checkHostAccess —
-// the destination machine still authorizes) so we don't leak tunnel-presence
-// or fly topology to unauthenticated probers.
+// Pre-flight for WS upgrade routing. Requires a valid JWT so we don't leak
+// tunnel-presence or fly topology to unauthenticated probers. When this
+// machine owns the tunnel, it also runs the (cached) access check: the WS
+// upgrade's 403 is invisible to browser clients (they only see a 1006 close),
+// so this is the one place a definitive denial can surface — clients use it
+// to stop reconnect-looping against hosts they'll never be allowed to reach.
 app.get("/hosts/:hostId/_whoowns", async (c) => {
 	const token = extractToken(c);
 	if (!token) return c.json({ error: "Unauthorized" }, 401);
@@ -305,9 +308,19 @@ app.get("/hosts/:hostId/_whoowns", async (c) => {
 	const hostId = c.req.param("hostId");
 	const replay = await maybeReplay(hostId);
 	if (!replay) {
-		return tunnelManager.hasTunnel(hostId)
-			? c.json({ ok: true, region: env.FLY_REGION })
-			: c.json({ error: "Host not connected" }, 503);
+		if (!tunnelManager.hasTunnel(hostId)) {
+			return c.json({ error: "Host not connected" }, 503);
+		}
+		const access = await checkHostAccess(auth, token, hostId);
+		if (!access.ok) {
+			const detail = `Forbidden: ${accessDenialMessage(access.reason)}`;
+			// "error" means the access check itself failed (API unreachable), not
+			// a denial — don't 403, or clients would stop retrying permanently.
+			return access.reason === "error"
+				? c.json({ error: detail }, 500)
+				: c.json({ error: detail }, 403);
+		}
+		return c.json({ ok: true, region: env.FLY_REGION });
 	}
 	return c.body(null, 200, replay.header);
 });

packages/workspace-client/src/lib/eventBus.ts

@@ -92,6 +92,10 @@ interface ListenerEntry {
 
 const RECONNECT_BASE_MS = 1_000;
 const RECONNECT_MAX_MS = 30_000;
+// Definitive access denial (preflight 403): the relay will keep saying no, so
+// exponential 1-30s retries just hammer it. Poll slowly instead of stopping
+// outright so access granted later (host sharing) is picked up eventually.
+const ACCESS_DENIED_RETRY_MS = 5 * 60_000;
 
 interface ConnectionState {
 	socket: WebSocket | null;
@@ -215,8 +219,12 @@ function connect(
 	// machine before the WS upgrade. fly-replay isn't transparent to all WS
 	// clients on the upgrade itself, but is on plain HTTP, so a quick GET
 	// avoids the connect → 1006 close → reconnect flicker.
-	void primeRelayAffinity(wsUrl).then(() => {
+	void primeRelayAffinity(wsUrl).then((probe) => {
 		if (state.disposed || state.socket) return;
+		if (probe?.status === 403) {
+			scheduleReconnect(state, hostUrl, getWsToken, ACCESS_DENIED_RETRY_MS);
+			return;
+		}
 		let socket: WebSocket;
 		try {
 			socket = new WebSocket(wsUrl);
@@ -255,13 +263,16 @@ function scheduleReconnect(
 	state: ConnectionState,
 	hostUrl: string,
 	getWsToken: () => string | null,
+	fixedDelayMs?: number,
 ): void {
 	if (state.disposed || state.reconnectTimer) return;
 
-	const delay = Math.min(
-		RECONNECT_BASE_MS * 2 ** state.reconnectAttempts,
-		RECONNECT_MAX_MS,
-	);
+	const delay =
+		fixedDelayMs ??
+		Math.min(
+			RECONNECT_BASE_MS * 2 ** state.reconnectAttempts,
+			RECONNECT_MAX_MS,
+		);
 	state.reconnectAttempts++;
 
 	state.reconnectTimer = setTimeout(() => {

plans/20260712-partysocket-ws-consolidation.md

@@ -0,0 +1,42 @@
+# Consolidate reconnecting WebSocket clients on partysocket
+
+Ticket: SUPER-1469 · Prereq: PR #5628 (merged — relay `_whoowns` now returns 403 on definitive denial)
+
+## Problem
+
+Three hand-rolled reconnecting-WS clients re-solve backoff, buffering, and token freshness:
+
+| Transport | Lines | Consumers | Token handling today |
+|---|---|---|---|
+| `packages/workspace-client/src/lib/eventBus.ts` | ~400 | 9 files | sync `getWsToken()` per attempt |
+| `apps/web/.../WebTerminal/TerminalConnection.ts` | ~240 | 1 | async TTL-cached token per attempt (correct) |
+| `apps/desktop/.../terminal-ws-transport.ts` | ~665 | via registry | fresh-token re-sign per attempt (fixed in #5628) |
+
+The desktop copy's drift caused the July 2026 401 reconnect-loop bug (~25 users/day). One library removes the bug class instead of patching each copy.
+
+## Decision
+
+Adopt `partysocket` (Cloudflare/PartyKit's maintained reconnecting-websocket fork; generic `WebSocket` export, no PartyKit coupling). Key primitive: `url` accepts `() => Promise<string>`, evaluated before **every** connection attempt — fresh token per dial for free. Also: `maxRetries`, delay bounds + grow factor, `connectionTimeout`, `minUptime`, `maxEnqueuedMessages`, `binaryType`, permanent `close()`.
+
+## Design: shared wrapper in `packages/workspace-client`
+
+One `createRelaySocket(opts)` (~150 lines) wrapping partysocket's generic export:
+
+- **URL provider**: async, calls the caller's token getter (desktop `ensureFreshJwt`, web `getAuthToken`) and signs the URL per attempt.
+- **Preflight**: runs `primeRelayAffinity` (`_whoowns`) inside the provider before returning the URL — pins Fly edge affinity and surfaces the real HTTP status the WS API hides.
+- **Fatal denial**: preflight **403** (only — 401 means expired token and self-heals via the provider) → permanent `close()` + `onAccessDenied` callback. Non-terminal consumers may instead pass `accessDeniedRetryMs` (eventBus: 5 min).
+- **Bounded buffering**: `maxEnqueuedMessages` mapped from each transport's current cap.
+
+## Phases
+
+1. **Spike (blocking)**: partysocket's behavior when the async `url` provider *rejects* is undocumented — verify it schedules a retry (vs wedging/unhandled rejection) before building on it. Also confirm `event-target-polyfill` behaves in the Electron renderer.
+2. **eventBus** (~½–1d): replace `connect`/`scheduleReconnect` internals; `getEventBus` API and fs:watch resend unchanged → zero consumer changes.
+3. **Web TerminalConnection** (~½–1d): dedupe (not a fix — web already refreshes correctly); keep attach/replay semantics.
+4. **Desktop terminal transport**: deferred. Only ~200/665 lines are dial/backoff; the rest (sleep/wake watchdog, write coalescer, replay flags, epoch guards, telemetry) must be re-hung on partysocket's silent socket-swapping model — high regression risk on just-verified code for mostly aesthetic gain. Revisit when the transport needs touching anyway.
+
+## Risks / constraints
+
+- partysocket swaps the underlying socket silently; code that compares socket identity (desktop transport does) needs restructuring — main reason phase 4 is deferred.
+- React Native: cloudflare/partykit#401 (Hermes lacks `MessageEvent`) — check before any apps/mobile use.
+- Relay semantics (post-#5628): 401 = expired/invalid JWT (retry with fresh token), 403 = definitive access denial (stop/slow-poll). The wrapper must never treat 401 as fatal.
+- New WS consumers should use the wrapper from day one.