From 8a51dde3928f2c744f37e69cafa3b7ef78b75a48 Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Sat, 11 Jul 2026 16:18:17 -0700 Subject: [PATCH 1/4] fix(desktop): refresh relay JWT on WS reconnect and stop retrying definitive 403s Relay JWTs rotate hourly, but the renderer's cached JWT only updated when an API response happened to carry set-auth-jwt, and terminal reconnects reused the URL signed at render time. Any terminal older than ~1h that hit a reconnect (sleep/wake, host restart) redialed with an expired token and looped on 401 until app restart (~25 users/day in PostHog). - auth-client: track token exp; ensureFreshJwt() mints from /api/auth/token (deduped); getJwt() kicks a background refresh when near expiry - terminal transport: re-sign the WS URL with a fresh token on every dial; treat a preflight 403 as fatal (surface "no access", emit telemetry, stop the reconnect loop) - eventBus: retry preflight 403 every 5 minutes instead of the 1-30s exponential loop, so unauthorized teammate-hosts stop hammering the relay --- apps/desktop/src/renderer/lib/auth-client.ts | 42 ++++++++++++++ .../src/renderer/lib/jwt-expiry.test.ts | 28 ++++++++++ apps/desktop/src/renderer/lib/jwt-expiry.ts | 13 +++++ .../lib/terminal/terminal-ws-transport.ts | 55 ++++++++++++++++--- packages/workspace-client/src/lib/eventBus.ts | 21 +++++-- 5 files changed, 145 insertions(+), 14 deletions(-) create mode 100644 apps/desktop/src/renderer/lib/jwt-expiry.test.ts create mode 100644 apps/desktop/src/renderer/lib/jwt-expiry.ts diff --git a/apps/desktop/src/renderer/lib/auth-client.ts b/apps/desktop/src/renderer/lib/auth-client.ts index 68d76577c34..0385b1d37ae 100644 --- a/apps/desktop/src/renderer/lib/auth-client.ts +++ b/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,56 @@ export function getAuthToken(): string | null { } let jwt: string | null = null; +let jwtExpiresAtMs: number | null = null; +let jwtRefreshInFlight: Promise | 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; + 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 { + if (jwtIsFresh()) return jwt; + if (!jwtRefreshInFlight) { + jwtRefreshInFlight = authClient + .$fetch<{ token?: string }>("/token") + .then((res) => { + const token = res.data?.token; + if (typeof token === "string" && token) setJwt(token); + return jwt; + }) + .catch(() => jwt) + .finally(() => { + jwtRefreshInFlight = null; + }); + } + return jwtRefreshInFlight; +} + /** * Better Auth client for Electron desktop app. * diff --git a/apps/desktop/src/renderer/lib/jwt-expiry.test.ts b/apps/desktop/src/renderer/lib/jwt-expiry.test.ts new file mode 100644 index 00000000000..eb19c9eac66 --- /dev/null +++ b/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(); + }); +}); diff --git a/apps/desktop/src/renderer/lib/jwt-expiry.ts b/apps/desktop/src/renderer/lib/jwt-expiry.ts new file mode 100644 index 00000000000..78b2644a84c --- /dev/null +++ b/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; + } +} diff --git a/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts b/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts index b1fb38e0828..9062e4ee939 100644 --- a/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts +++ b/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"; @@ -376,7 +377,7 @@ export function connect( ? appendQueryParam(wsUrl, "replay", "0") : wsUrl; - const openSocket = () => { + const openSocket = (targetUrl: string) => { // Bail if the transport raced into a different URL or was disconnected // while the pre-flight was in flight. if ( @@ -387,12 +388,12 @@ export function connect( } 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 +410,52 @@ 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.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); } } diff --git a/packages/workspace-client/src/lib/eventBus.ts b/packages/workspace-client/src/lib/eventBus.ts index c368b03aa2b..59202043c0b 100644 --- a/packages/workspace-client/src/lib/eventBus.ts +++ b/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(() => { From b76609d176a7fec949ae95446043c929f5e80d70 Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Sat, 11 Jul 2026 16:34:26 -0700 Subject: [PATCH 2/4] fix(relay): surface access denials on the _whoowns preflight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The WS upgrade's 403 is invisible to browser clients (1006 close), and _whoowns skipped the access check entirely, so a denied client could never distinguish "no access" from a routing hiccup — which is why denied clients reconnect-loop forever. Run the cached checkHostAccess when this machine owns the tunnel and return 403 on denial (500 when the check itself errored, so a transient API outage doesn't read as a permanent denial). --- apps/relay/src/index.ts | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/apps/relay/src/index.ts b/apps/relay/src/index.ts index d7ac6684bfb..c398cf68bfe 100644 --- a/apps/relay/src/index.ts +++ b/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); }); From 4ab51e68deb41260c0a7d159e092136a2053e3fb Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Sat, 11 Jul 2026 16:37:17 -0700 Subject: [PATCH 3/4] fix(desktop): guard JWT refresh and terminal preflight against races Review follow-ups: generation-check the /token result so a refresh resolving after logout (or a newer set-auth-jwt header) can't repopulate a stale token; log refresh failures; epoch-guard connect() preflight callbacks so a superseded attempt can't open a duplicate socket or fatally terminate the newer one. --- apps/desktop/src/renderer/lib/auth-client.ts | 18 ++++++++++++++++-- .../lib/terminal/terminal-ws-transport.ts | 13 +++++++++++-- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/apps/desktop/src/renderer/lib/auth-client.ts b/apps/desktop/src/renderer/lib/auth-client.ts index 0385b1d37ae..72cd04c7f4a 100644 --- a/apps/desktop/src/renderer/lib/auth-client.ts +++ b/apps/desktop/src/renderer/lib/auth-client.ts @@ -22,6 +22,7 @@ export function getAuthToken(): string | null { let jwt: string | null = null; let jwtExpiresAtMs: number | null = null; +let jwtGeneration = 0; let jwtRefreshInFlight: Promise | null = null; // Refresh ahead of expiry so a token handed to a WS URL is still valid by the @@ -30,6 +31,7 @@ const JWT_REFRESH_LEEWAY_MS = 60_000; export function setJwt(token: string | null) { jwt = token; + jwtGeneration++; jwtExpiresAtMs = token ? decodeJwtExpiresAtMs(token) : null; } @@ -56,14 +58,26 @@ export function getJwt(): string | null { export async function ensureFreshJwt(): Promise { if (jwtIsFresh()) return jwt; if (!jwtRefreshInFlight) { + const generationAtStart = jwtGeneration; jwtRefreshInFlight = authClient .$fetch<{ token?: string }>("/token") .then((res) => { const token = res.data?.token; - if (typeof token === "string" && token) setJwt(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; }) - .catch(() => jwt) .finally(() => { jwtRefreshInFlight = null; }); diff --git a/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts b/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts index 9062e4ee939..d49af0a1efc 100644 --- a/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts +++ b/apps/desktop/src/renderer/lib/terminal/terminal-ws-transport.ts @@ -79,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; @@ -180,6 +184,7 @@ export function createTransport(): TerminalTransport { _resumeListener: null, _writeCoalescer: null, _lastProbe: null, + _connectEpoch: 0, }; } @@ -362,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 @@ -378,9 +384,11 @@ export function connect( : wsUrl; const openSocket = (targetUrl: string) => { - // Bail if the transport raced into a different URL or was disconnected - // while the pre-flight was in flight. + // 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" ) { @@ -429,6 +437,7 @@ export function connect( // 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" ) { From ca76fc130fc3bd5c8654ffaa7a0c967d9f8b70b1 Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Sun, 12 Jul 2026 11:09:19 -0700 Subject: [PATCH 4/4] docs(plans): partysocket WS-client consolidation design (SUPER-1469) --- .../20260712-partysocket-ws-consolidation.md | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 plans/20260712-partysocket-ws-consolidation.md diff --git a/plans/20260712-partysocket-ws-consolidation.md b/plans/20260712-partysocket-ws-consolidation.md new file mode 100644 index 00000000000..c94b567f88e --- /dev/null +++ b/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`, 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.