diff --git a/canvas/src/components/tabs/DisplayTab.tsx b/canvas/src/components/tabs/DisplayTab.tsx index 94f50730c..38b9a7c0d 100644 --- a/canvas/src/components/tabs/DisplayTab.tsx +++ b/canvas/src/components/tabs/DisplayTab.tsx @@ -1,6 +1,6 @@ "use client"; -import { useEffect, useRef, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import { api } from "@/lib/api"; import type RFB from "@novnc/novnc"; @@ -33,6 +33,11 @@ export function DisplayTab({ workspaceId }: Props) { const [controlBusy, setControlBusy] = useState(false); const [sessionUrl, setSessionUrl] = useState(null); const requestGeneration = useRef(0); + // Freshest signed session URL (token bound to the lease's expires_at). The + // renewal timer keeps this current WITHOUT swapping the live stream's + // sessionUrl (which would needlessly reconnect the desktop); the stream uses + // it only when it has to reconnect after an unclean drop. + const latestSessionUrlRef = useRef(null); useEffect(() => { const generation = requestGeneration.current + 1; @@ -41,6 +46,7 @@ export function DisplayTab({ workspaceId }: Props) { setStatus(null); setControl(null); setSessionUrl(null); + latestSessionUrlRef.current = null; setError(null); setControlError(null); setControlBusy(false); @@ -69,6 +75,41 @@ export function DisplayTab({ workspaceId }: Props) { }; }, [workspaceId]); + // Acquire (or re-acquire) the display-control lease as the current holder. + // Re-acquiring extends the 300s server-side lock AND returns a freshly-signed + // session URL (token bound to the new expires_at). Used both to renew the + // lease on a timer and to mint a non-stale token for each reconnect — a + // cached URL can be past its ~300s expiry, which would make a reconnect 401. + const reacquireSession = useCallback(async (): Promise => { + const generation = requestGeneration.current; + try { + const next = await api.post( + `/workspaces/${workspaceId}/display/control/acquire`, + { controller: "user", ttl_seconds: 300 }, + ); + if (requestGeneration.current !== generation) return null; + setControl(next); + if (next.session_url) latestSessionUrlRef.current = next.session_url; + return next.session_url ?? null; + } catch { + // Transient failure, or another holder took over: the live stream keeps + // running on its existing connection; a reconnect re-evaluates control. + return null; + } + }, [workspaceId]); + + // Renew the lease while we hold it. The lock is a 300s lease with no + // server-side auto-renewal, so without this the control (and the session + // token) silently expire mid-session — the user appears "kicked" every ~5 + // minutes. We renew well inside the TTL and do not touch the live stream. + useEffect(() => { + if (!sessionUrl) return; + const timer = setInterval(() => { + void reacquireSession(); + }, 120_000); + return () => clearInterval(timer); + }, [sessionUrl, reacquireSession]); + const acquireControl = async () => { const generation = requestGeneration.current; const controlPath = `/workspaces/${workspaceId}/display/control`; @@ -82,6 +123,7 @@ export function DisplayTab({ workspaceId }: Props) { if (requestGeneration.current !== generation) return; setControl(next); setSessionUrl(next.session_url || null); + latestSessionUrlRef.current = next.session_url || null; } catch (err) { if (requestGeneration.current !== generation) return; setControlError("Failed to take control"); @@ -108,6 +150,7 @@ export function DisplayTab({ workspaceId }: Props) { if (requestGeneration.current !== generation) return; setControl(next); setSessionUrl(null); + latestSessionUrlRef.current = null; } catch (err) { if (requestGeneration.current !== generation) return; setControlError("Failed to release control"); @@ -235,7 +278,11 @@ export function DisplayTab({ workspaceId }: Props) { /> {sessionUrl ? ( - + ) : (
@@ -311,7 +358,15 @@ function DisplayControlBar({ ); } -function DesktopStream({ sessionUrl }: { sessionUrl: string }) { +function DesktopStream({ + sessionUrl, + latestSessionUrlRef, + reacquireSession, +}: { + sessionUrl: string; + latestSessionUrlRef: { current: string | null }; + reacquireSession: () => Promise; +}) { const containerRef = useRef(null); const rfbRef = useRef(null); const [streamError, setStreamError] = useState(null); @@ -329,20 +384,37 @@ function DesktopStream({ sessionUrl }: { sessionUrl: string }) { clipboardTimer = setTimeout(() => setClipboardStatus(null), 2500); }; - async function connect() { + let attempts = 0; + let retryTimer: ReturnType | null = null; + const maxAttempts = 10; + + async function connect(reacquire = false) { setStreamError(null); try { + // On a reconnect, mint a fresh lease + token first — the original token + // is only ~300s, so a cached URL can be expired and would 401. The + // initial connect already holds a fresh token from acquireControl. + if (reacquire) await reacquireSession(); const mod = await import("@novnc/novnc"); if (cancelled || !containerRef.current) return; - const stream = displayWebSocketConnection(sessionUrl); + const stream = displayWebSocketConnection(latestSessionUrlRef.current || sessionUrl); rfb = new mod.default(containerRef.current, stream.url, { wsProtocols: ["binary", `molecule-display-token.${stream.token}`], }); rfbRef.current = rfb; rfb.scaleViewport = true; - rfb.resizeSession = true; + // Do NOT request a server-side resize: the workspace display runs a + // fixed Xorg modeline and x11vnc rejects SetDesktopSize ("Resize is + // administratively prohibited"), which spams the console on every + // (re)connect. scaleViewport already fits the fixed framebuffer to the + // container client-side, so we don't need the server to resize. + rfb.resizeSession = false; rfb.focusOnClick = true; rfb.focus({ preventScroll: true }); + rfb.addEventListener("connect", () => { + attempts = 0; + if (!cancelled) setStreamError(null); + }); rfb.addEventListener("clipboard", (event: Event) => { const text = (event as CustomEvent<{ text?: string }>).detail?.text ?? ""; if (!text) return; @@ -353,7 +425,20 @@ function DesktopStream({ sessionUrl }: { sessionUrl: string }) { }); rfb.addEventListener("disconnect", (event: Event) => { const detail = (event as CustomEvent<{ clean?: boolean }>).detail; - if (!cancelled && !detail?.clean) setStreamError("Desktop stream disconnected."); + rfbRef.current = null; + if (cancelled || detail?.clean) return; + // Auto-reconnect after an unclean drop (idle/network blip, brief + // agent hiccup); bounded backoff so a genuinely-dead session still + // surfaces an error instead of looping forever. + if (attempts < maxAttempts) { + attempts += 1; + setStreamError(`Reconnecting to desktop… (attempt ${attempts})`); + retryTimer = setTimeout(() => { + if (!cancelled) void connect(true); + }, Math.min(1000 * attempts, 5000)); + } else { + setStreamError("Desktop stream disconnected."); + } }); } catch { if (!cancelled) setStreamError("Desktop stream could not be opened."); @@ -363,11 +448,12 @@ function DesktopStream({ sessionUrl }: { sessionUrl: string }) { connect(); return () => { cancelled = true; + if (retryTimer) clearTimeout(retryTimer); if (clipboardTimer) clearTimeout(clipboardTimer); rfbRef.current = null; rfb?.disconnect(); }; - }, [sessionUrl]); + }, [sessionUrl, reacquireSession, latestSessionUrlRef]); useEffect(() => { const onPaste = (event: ClipboardEvent) => { diff --git a/canvas/src/components/tabs/__tests__/DisplayTab.test.tsx b/canvas/src/components/tabs/__tests__/DisplayTab.test.tsx index afd84361c..30b83ad99 100644 --- a/canvas/src/components/tabs/__tests__/DisplayTab.test.tsx +++ b/canvas/src/components/tabs/__tests__/DisplayTab.test.tsx @@ -2,12 +2,13 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react"; -const { mockGet, mockPost, mockRFBConstructor, mockRFBClipboardPasteFrom, mockRFBFocus } = vi.hoisted(() => ({ +const { mockGet, mockPost, mockRFBConstructor, mockRFBClipboardPasteFrom, mockRFBFocus, rfbInstances } = vi.hoisted(() => ({ mockGet: vi.fn(), mockPost: vi.fn(), mockRFBConstructor: vi.fn(), mockRFBClipboardPasteFrom: vi.fn(), mockRFBFocus: vi.fn(), + rfbInstances: [] as EventTarget[], })); vi.mock("@/lib/api", () => ({ @@ -31,6 +32,7 @@ vi.mock("@novnc/novnc", () => ({ this.url = url; this.options = options; mockRFBConstructor(target, url, options); + rfbInstances.push(this); } clipboardPasteFrom(text: string) { mockRFBClipboardPasteFrom(text); @@ -52,6 +54,7 @@ describe("DisplayTab", () => { mockRFBConstructor.mockReset(); mockRFBClipboardPasteFrom.mockReset(); mockRFBFocus.mockReset(); + rfbInstances.length = 0; }); it("renders unavailable state for non-display workspaces", async () => { @@ -400,6 +403,62 @@ describe("DisplayTab", () => { }); expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy(); }); + + it("auto-reconnects the desktop stream after an unclean disconnect but not a clean one", async () => { + mockGet + .mockResolvedValueOnce({ + available: true, + mode: "desktop-control", + protocol: "novnc", + width: 1920, + height: 1080, + }) + .mockResolvedValueOnce({ controller: "none" }); + // Initial acquire returns token "signed"; the reconnect re-acquire mints a + // FRESH token "signed2" (the lock/token is only ~300s — reconnecting with a + // cached, possibly-expired token would 401 and never recover). + mockPost + .mockResolvedValueOnce({ + controller: "user", + controlled_by: "admin-token", + expires_at: "2026-05-23T08:48:27Z", + session_url: "/workspaces/ws-display/display/session/websockify#token=signed", + }) + .mockResolvedValue({ + controller: "user", + controlled_by: "admin-token", + expires_at: "2026-05-23T08:53:27Z", + session_url: "/workspaces/ws-display/display/session/websockify#token=signed2", + }); + + render(); + await waitFor(() => { + expect(screen.getByRole("button", { name: "Take control" })).toBeTruthy(); + }); + fireEvent.click(screen.getByRole("button", { name: "Take control" })); + await waitFor(() => { + expect(rfbInstances.length).toBe(1); + }); + expect(mockRFBConstructor.mock.calls[0][2].wsProtocols).toContain("molecule-display-token.signed"); + + // An idle/network drop closes the websocket uncleanly. The client must + // re-acquire a fresh token and reconnect instead of giving up — this is the + // "disconnects every ~5 min and stays dead" report. + rfbInstances[0].dispatchEvent(new CustomEvent("disconnect", { detail: { clean: false } })); + await waitFor( + () => { + expect(rfbInstances.length).toBe(2); + }, + { timeout: 3000 }, + ); + // Reconnect dialed with the FRESH token, not the stale original. + expect(mockRFBConstructor.mock.calls[1][2].wsProtocols).toContain("molecule-display-token.signed2"); + + // A clean disconnect (the user released control) must NOT reconnect. + rfbInstances[1].dispatchEvent(new CustomEvent("disconnect", { detail: { clean: true } })); + await new Promise((resolve) => setTimeout(resolve, 1100)); + expect(rfbInstances.length).toBe(2); + }); }); function deferred() {