diff --git a/site/src/pages/AgentsPage/DesktopPanel.stories.tsx b/site/src/pages/AgentsPage/DesktopPanel.stories.tsx index e8fa5afe67..27dfb170c4 100644 --- a/site/src/pages/AgentsPage/DesktopPanel.stories.tsx +++ b/site/src/pages/AgentsPage/DesktopPanel.stories.tsx @@ -1,14 +1,17 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; -import { DesktopPanel } from "./DesktopPanel"; -import { mockAttach, mockDesktopConnection } from "./desktopStoryUtils"; +import { fn } from "storybook/test"; +import { DesktopPanelView, type DesktopPanelViewProps } from "./DesktopPanel"; -const meta: Meta = { +const defaults: DesktopPanelViewProps = { + status: "idle", + reconnect: fn(), + attach: fn(), +}; + +const meta: Meta = { title: "pages/AgentsPage/DesktopPanel", - component: DesktopPanel, - args: { - isExpanded: false, - chatId: "test-chat-id", - }, + component: DesktopPanelView, + args: defaults, decorators: [ (Story) => (
@@ -18,35 +21,22 @@ const meta: Meta = { ], }; export default meta; -type Story = StoryObj; +type Story = StoryObj; -export const Connected: Story = { - args: { - connectionOverride: mockDesktopConnection({ - status: "connected", - hasConnected: true, - attach: mockAttach(), - }), - }, -}; +export const Idle: Story = {}; export const Connecting: Story = { - args: { - connectionOverride: mockDesktopConnection({ status: "connecting" }), - }, + args: { status: "connecting" }, }; -export const ErrorState: Story = { - args: { - connectionOverride: mockDesktopConnection({ status: "error" }), - }, +export const Connected: Story = { + args: { status: "connected" }, }; export const Disconnected: Story = { - args: { - connectionOverride: mockDesktopConnection({ - status: "disconnected", - hasConnected: true, - }), - }, + args: { status: "disconnected" }, +}; + +export const ErrorState: Story = { + args: { status: "error" }, }; diff --git a/site/src/pages/AgentsPage/DesktopPanel.tsx b/site/src/pages/AgentsPage/DesktopPanel.tsx index 3f53189a98..10d90c1fb9 100644 --- a/site/src/pages/AgentsPage/DesktopPanel.tsx +++ b/site/src/pages/AgentsPage/DesktopPanel.tsx @@ -1,57 +1,37 @@ import { Button } from "components/Button/Button"; import { Spinner } from "components/Spinner/Spinner"; -import { type FC, useEffect, useRef } from "react"; -import { - type UseDesktopConnectionResult, - useDesktopConnection, -} from "./useDesktopConnection"; +import type { FC } from "react"; +import { useDesktopConnection } from "./useDesktopConnection"; + +type DesktopConnectionStatus = + | "idle" + | "connecting" + | "connected" + | "disconnected" + | "error"; interface DesktopPanelProps { chatId: string; - isExpanded: boolean; - /** Optional override for the desktop connection. Used in stories. */ - connectionOverride?: UseDesktopConnectionResult; } -export const DesktopPanel: FC = ({ - chatId, - isExpanded: _isExpanded, - connectionOverride, +export interface DesktopPanelViewProps { + status: DesktopConnectionStatus; + reconnect: () => void; + attach: (container: HTMLElement) => void; +} + +export const DesktopPanel: FC = ({ chatId }) => { + const { status, reconnect, attach } = useDesktopConnection({ chatId }); + return ( + + ); +}; + +export const DesktopPanelView: FC = ({ + status, + reconnect, + attach, }) => { - // When an override is provided, pass undefined chatId to prevent - // the real hook from attempting any WebSocket connections. - const hookResult = useDesktopConnection({ - chatId: connectionOverride ? undefined : chatId, - }); - const { status, connect, disconnect, attach } = - connectionOverride ?? hookResult; - const containerRef = useRef(null); - - const attachToContainer = (el: HTMLDivElement | null) => { - containerRef.current = el; - if (el) { - attach(el); - } - }; - - // Connect on mount, disconnect on unmount. This drives the - // visibility-based lifecycle: DesktopPanel is only rendered - // when the Desktop tab is active, so mounting/unmounting - // naturally starts and stops the WebSocket connection. - useEffect(() => { - connect(); - return () => { - disconnect(); - }; - }, [connect, disconnect]); - - // Re-attach when status changes to connected (e.g., after reconnect). - useEffect(() => { - if (status === "connected" && containerRef.current) { - attach(containerRef.current); - } - }, [status, attach]); - if (status === "connecting") { return (
@@ -73,17 +53,11 @@ export const DesktopPanel: FC = ({ if (status === "error") { return (
- - Failed to connect to the desktop session. + + Failed to connect to the desktop session. The agent may not be + connected or the desktop environment may not be available. -
@@ -100,5 +74,12 @@ export const DesktopPanel: FC = ({ } // status === "connected" - return
; + return ( +
{ + if (el) attach(el); + }} + className="h-full w-full" + /> + ); }; diff --git a/site/src/pages/AgentsPage/SidebarTabView.stories.tsx b/site/src/pages/AgentsPage/SidebarTabView.stories.tsx index 6105c49ca5..01b20178f2 100644 --- a/site/src/pages/AgentsPage/SidebarTabView.stories.tsx +++ b/site/src/pages/AgentsPage/SidebarTabView.stories.tsx @@ -1,6 +1,5 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { fn } from "storybook/test"; -import { mockAttach, mockDesktopConnection } from "./desktopStoryUtils"; import type { SidebarTab } from "./SidebarTabView"; import { SidebarTabView } from "./SidebarTabView"; @@ -108,31 +107,3 @@ export const NarrowPanel: Story = { ), ], }; - -export const DesktopConnecting: Story = { - args: { - tabs: [], - desktopChatId: "test-chat-id", - desktopConnectionOverride: mockDesktopConnection({ status: "connecting" }), - }, -}; - -export const DesktopConnected: Story = { - args: { - tabs: [], - desktopChatId: "test-chat-id", - desktopConnectionOverride: mockDesktopConnection({ - status: "connected", - hasConnected: true, - attach: mockAttach(), - }), - }, -}; - -export const DesktopError: Story = { - args: { - tabs: [], - desktopChatId: "test-chat-id", - desktopConnectionOverride: mockDesktopConnection({ status: "error" }), - }, -}; diff --git a/site/src/pages/AgentsPage/SidebarTabView.tsx b/site/src/pages/AgentsPage/SidebarTabView.tsx index f1f647bd90..475abdcbc2 100644 --- a/site/src/pages/AgentsPage/SidebarTabView.tsx +++ b/site/src/pages/AgentsPage/SidebarTabView.tsx @@ -11,7 +11,6 @@ import type { ReactNode } from "react"; import { type FC, useEffect, useId, useRef, useState } from "react"; import { cn } from "utils/cn"; import { DesktopPanel } from "./DesktopPanel"; -import type { UseDesktopConnectionResult } from "./useDesktopConnection"; /** A single tab definition for the sidebar panel. */ export interface SidebarTab { @@ -43,8 +42,6 @@ interface SidebarTabViewProps { onClose?: () => void; /** Desktop chat ID. Omitted if desktop is not available. */ desktopChatId?: string; - /** Optional override for the desktop connection. Used in stories. */ - desktopConnectionOverride?: UseDesktopConnectionResult; } /** How far (px) each chevron click scrolls the tab strip. */ @@ -110,7 +107,6 @@ export const SidebarTabView: FC = ({ chatTitle, onClose, desktopChatId, - desktopConnectionOverride, }) => { const tabIdPrefix = useId(); const [activeTabId, setActiveTabId] = useState( @@ -324,11 +320,7 @@ export const SidebarTabView: FC = ({ className="min-h-0 flex-1" > {effectiveTabId === "desktop" && desktopChatId ? ( - + ) : ( activeTab?.content )} diff --git a/site/src/pages/AgentsPage/desktopStoryUtils.ts b/site/src/pages/AgentsPage/desktopStoryUtils.ts deleted file mode 100644 index ee89cb32b5..0000000000 --- a/site/src/pages/AgentsPage/desktopStoryUtils.ts +++ /dev/null @@ -1,46 +0,0 @@ -import { fn } from "storybook/test"; -import type { UseDesktopConnectionResult } from "./useDesktopConnection"; - -/** - * Creates a mock attach function that inserts a placeholder element - * simulating a noVNC canvas so the "connected" state is visible in - * stories. - */ -export function mockAttach(): (container: HTMLElement) => void { - const placeholder = document.createElement("div"); - Object.assign(placeholder.style, { - width: "100%", - height: "100%", - background: - "linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%)", - display: "flex", - alignItems: "center", - justifyContent: "center", - color: "#888", - fontSize: "14px", - fontFamily: "monospace", - }); - placeholder.textContent = "VNC canvas placeholder"; - - const attachFn = fn((container: HTMLElement) => { - if (placeholder.parentElement !== container) { - container.appendChild(placeholder); - } - }); - - return attachFn; -} - -export function mockDesktopConnection( - overrides: Partial = {}, -): UseDesktopConnectionResult { - return { - status: "idle", - hasConnected: false, - connect: fn(), - disconnect: fn(), - attach: fn(), - rfb: null, - ...overrides, - }; -} diff --git a/site/src/pages/AgentsPage/useDesktopConnection.test.ts b/site/src/pages/AgentsPage/useDesktopConnection.test.ts index 1197cfb2dd..b53c887de0 100644 --- a/site/src/pages/AgentsPage/useDesktopConnection.test.ts +++ b/site/src/pages/AgentsPage/useDesktopConnection.test.ts @@ -23,13 +23,20 @@ interface MockRFBInstance { const { FakeRFB, lastInstance } = vi.hoisted(() => { const ref: { current: MockRFBInstance | null } = { current: null }; + // When true, the constructor throws to simulate failures + // like missing WebGL support. + let shouldThrow = false; + + class FakeRFB implements MockRFBInstance { + static set throwOnConstruct(v: boolean) { + shouldThrow = v; + } - class FakeRFB { scaleViewport = false; resizeSession = true; disconnect = vi.fn(); addEventListener = vi.fn((type: string, handler: (ev: unknown) => void) => { - (this as unknown as MockRFBInstance).listeners.set(type, handler); + this.listeners.set(type, handler); }); listeners = new Map void>(); @@ -43,10 +50,12 @@ const { FakeRFB, lastInstance } = vi.hoisted(() => { } constructor() { - ref.current = this as unknown as MockRFBInstance; + if (shouldThrow) { + throw new Error("WebGL not supported"); + } + ref.current = this; } } - return { FakeRFB, lastInstance: ref }; }); @@ -68,7 +77,9 @@ function getLastRFBInstance(): MockRFBInstance { } function createMockSocket(): WebSocket { - return { binaryType: "arraybuffer" } as unknown as WebSocket; + const socket = new WebSocket("ws://localhost"); + vi.spyOn(socket, "close").mockImplementation(() => {}); + return socket; } // ---- tests ----------------------------------------------------------------- @@ -78,15 +89,16 @@ describe("useDesktopConnection", () => { mockWatchChatDesktop.mockReset(); mockWatchChatDesktop.mockReturnValue(createMockSocket()); lastInstance.current = null; + FakeRFB.throwOnConstruct = false; }); afterEach(() => { vi.restoreAllMocks(); }); - it("starts in idle status and does not connect automatically", () => { + it("does nothing when chatId is undefined", () => { const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), + useDesktopConnection({ chatId: undefined }), ); expect(result.current.status).toBe("idle"); @@ -95,29 +107,15 @@ describe("useDesktopConnection", () => { expect(mockWatchChatDesktop).not.toHaveBeenCalled(); }); - it("does nothing when chatId is undefined and connect() is called", () => { - const { result } = renderHook(() => - useDesktopConnection({ chatId: undefined }), - ); - - act(() => result.current.connect()); - - expect(result.current.status).toBe("idle"); - expect(mockWatchChatDesktop).not.toHaveBeenCalled(); - }); - - it("transitions to connecting then connected on connect()", () => { + it("auto-connects on mount when chatId is set", () => { const { result } = renderHook(() => useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); - const rfb = getLastRFBInstance(); - expect(mockWatchChatDesktop).toHaveBeenCalledWith("chat-1"); expect(result.current.status).toBe("connecting"); - expect(result.current.hasConnected).toBe(false); + const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); expect(result.current.status).toBe("connected"); @@ -125,53 +123,18 @@ describe("useDesktopConnection", () => { }); it("sets scaleViewport and resizeSession on the RFB instance", () => { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); - - act(() => result.current.connect()); + renderHook(() => useDesktopConnection({ chatId: "chat-1" })); const rfb = getLastRFBInstance(); expect(rfb.scaleViewport).toBe(true); expect(rfb.resizeSession).toBe(false); }); - it("connect() is a no-op when already connecting", () => { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); - - act(() => result.current.connect()); - expect(result.current.status).toBe("connecting"); - - mockWatchChatDesktop.mockClear(); - - act(() => result.current.connect()); - expect(mockWatchChatDesktop).not.toHaveBeenCalled(); - }); - - it("connect() is a no-op when already connected", () => { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); - - act(() => result.current.connect()); - const rfb = getLastRFBInstance(); - act(() => rfb.simulateEvent("connect")); - expect(result.current.status).toBe("connected"); - - mockWatchChatDesktop.mockClear(); - - act(() => result.current.connect()); - expect(mockWatchChatDesktop).not.toHaveBeenCalled(); - }); - it("transitions to error on securityfailure", () => { const { result } = renderHook(() => useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("securityfailure", { @@ -191,7 +154,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb1 = getLastRFBInstance(); act(() => rfb1.simulateEvent("connect")); expect(result.current.status).toBe("connected"); @@ -207,9 +169,9 @@ describe("useDesktopConnection", () => { expect(mockWatchChatDesktop).toHaveBeenCalledTimes(1); const rfb2 = getLastRFBInstance(); - // Reconnect attempt fails (no "connect" event) but desktop - // was previously reachable, so it retries. + // Reconnect succeeds then drops again. // attempt 1 → 2000ms delay. + act(() => rfb2.simulateEvent("connect")); act(() => rfb2.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("disconnected"); @@ -221,6 +183,7 @@ describe("useDesktopConnection", () => { const rfb3 = getLastRFBInstance(); // attempt 2 → 4000ms delay. + act(() => rfb3.simulateEvent("connect")); act(() => rfb3.simulateEvent("disconnect", { clean: false })); mockWatchChatDesktop.mockClear(); @@ -241,7 +204,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb1 = getLastRFBInstance(); act(() => rfb1.simulateEvent("connect")); @@ -274,22 +236,18 @@ describe("useDesktopConnection", () => { vi.useFakeTimers(); try { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); + renderHook(() => useDesktopConnection({ chatId: "chat-1" })); - act(() => result.current.connect()); let rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); // Burn through attempts: 1s, 2s, 4s, 8s, 16s. - // Reconnect attempts fail (no "connect" event) but the - // desktop was previously reachable so backoff accumulates. const delays = [1000, 2000, 4000, 8000, 16000]; for (const delay of delays) { act(() => rfb.simulateEvent("disconnect", { clean: false })); act(() => vi.advanceTimersByTime(delay)); rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); } // Attempt 5 — should be capped at 30_000, not 32_000. @@ -305,60 +263,14 @@ describe("useDesktopConnection", () => { } }); - it("disconnect() cleans up and resets to idle", () => { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); - - act(() => result.current.connect()); - const rfb = getLastRFBInstance(); - act(() => rfb.simulateEvent("connect")); - expect(result.current.status).toBe("connected"); - - act(() => result.current.disconnect()); - - expect(rfb.disconnect).toHaveBeenCalled(); - expect(result.current.status).toBe("idle"); - }); - - it("disconnect() cancels pending reconnect timers", () => { - vi.useFakeTimers(); - - try { - const { result } = renderHook(() => - useDesktopConnection({ chatId: "chat-1" }), - ); - - act(() => result.current.connect()); - const rfb = getLastRFBInstance(); - act(() => rfb.simulateEvent("connect")); - - // Trigger reconnect timer. - act(() => rfb.simulateEvent("disconnect", { clean: false })); - expect(result.current.status).toBe("disconnected"); - - // Manually disconnect before timer fires. - act(() => result.current.disconnect()); - expect(result.current.status).toBe("idle"); - - // Timer should be cancelled — no reconnect. - mockWatchChatDesktop.mockClear(); - act(() => vi.advanceTimersByTime(60_000)); - expect(mockWatchChatDesktop).not.toHaveBeenCalled(); - } finally { - vi.useRealTimers(); - } - }); - it("cleans up on unmount and does not reconnect", () => { vi.useFakeTimers(); try { - const { result, unmount } = renderHook(() => + const { unmount } = renderHook(() => useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); @@ -374,24 +286,28 @@ describe("useDesktopConnection", () => { } }); - it("resets state when chatId changes", () => { + it("tears down and reconnects when chatId changes", () => { const { result, rerender } = renderHook( ({ chatId }: { chatId: string | undefined }) => useDesktopConnection({ chatId }), { initialProps: { chatId: "chat-aaa" as string | undefined } }, ); - act(() => result.current.connect()); const rfb1 = getLastRFBInstance(); act(() => rfb1.simulateEvent("connect")); expect(result.current.status).toBe("connected"); expect(result.current.hasConnected).toBe(true); + mockWatchChatDesktop.mockClear(); rerender({ chatId: "chat-bbb" }); + // Old RFB was torn down. expect(rfb1.disconnect).toHaveBeenCalled(); - expect(result.current.status).toBe("idle"); expect(result.current.hasConnected).toBe(false); + + // Auto-reconnected with the new chatId. + expect(mockWatchChatDesktop).toHaveBeenCalledWith("chat-bbb"); + expect(result.current.status).toBe("connecting"); }); it("attach() appends the offscreen container to the target", () => { @@ -399,7 +315,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); @@ -415,7 +330,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); @@ -431,7 +345,7 @@ describe("useDesktopConnection", () => { expect(container2.children[0]).toBe(screen); expect(container1.children.length).toBe(0); - // WebSocket was only opened once. + // WebSocket was only opened once (plus the initial auto-connect). expect(mockWatchChatDesktop).toHaveBeenCalledTimes(1); }); @@ -443,11 +357,9 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); - // Disconnect fires before connect — e.g. agent returned 424 - // because portabledesktop is not installed. + // Disconnect fires before connect — e.g. agent returned 424. act(() => rfb.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("error"); @@ -461,7 +373,7 @@ describe("useDesktopConnection", () => { } }); - it("retries when reconnect attempt fails but desktop was previously reachable", () => { + it("does not retry when reconnect attempt fails before handshake completes", () => { vi.useFakeTimers(); try { @@ -469,8 +381,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - // Establish a successful connection first. - act(() => result.current.connect()); const rfb1 = getLastRFBInstance(); act(() => rfb1.simulateEvent("connect")); expect(result.current.status).toBe("connected"); @@ -483,15 +393,14 @@ describe("useDesktopConnection", () => { act(() => vi.advanceTimersByTime(1000)); const rfb2 = getLastRFBInstance(); - // Reconnect attempt fails (disconnect before connect), but - // desktop was previously reachable — should keep retrying. + // Reconnect attempt fails (disconnect before handshake). act(() => rfb2.simulateEvent("disconnect", { clean: false })); - expect(result.current.status).toBe("disconnected"); + expect(result.current.status).toBe("error"); - // Another retry should be scheduled. + // No further retry should be scheduled. mockWatchChatDesktop.mockClear(); - act(() => vi.advanceTimersByTime(2000)); - expect(mockWatchChatDesktop).toHaveBeenCalledTimes(1); + act(() => vi.advanceTimersByTime(60_000)); + expect(mockWatchChatDesktop).not.toHaveBeenCalled(); } finally { vi.useRealTimers(); } @@ -502,7 +411,6 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); const rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); @@ -522,13 +430,10 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); let rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); expect(result.current.status).toBe("connected"); - // Burn through 10 reconnect attempts that all fail before - // the handshake completes. for (let i = 0; i < 10; i++) { act(() => rfb.simulateEvent("disconnect", { clean: false })); const delay = Math.min(1000 * 2 ** i, 30_000); @@ -540,7 +445,6 @@ describe("useDesktopConnection", () => { act(() => rfb.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("error"); - // No more retries. mockWatchChatDesktop.mockClear(); act(() => vi.advanceTimersByTime(60_000)); expect(mockWatchChatDesktop).not.toHaveBeenCalled(); @@ -557,30 +461,26 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); let rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); - // 9 failed reconnects (just under the cap). for (let i = 0; i < 9; i++) { act(() => rfb.simulateEvent("disconnect", { clean: false })); const delay = Math.min(1000 * 2 ** i, 30_000); act(() => vi.advanceTimersByTime(delay)); rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); } - // This reconnect succeeds — counter resets after stability period. - act(() => rfb.simulateEvent("connect")); expect(result.current.status).toBe("connected"); // Advance past the stability period so the counter resets. act(() => vi.advanceTimersByTime(3000)); - // Another drop + failed reconnect should NOT hit the cap - // because the counter was reset. act(() => rfb.simulateEvent("disconnect", { clean: false })); act(() => vi.advanceTimersByTime(1000)); rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); act(() => rfb.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("disconnected"); } finally { @@ -596,28 +496,20 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); let rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); expect(result.current.status).toBe("connected"); - // Drop the connection immediately — before the 3s - // stability window elapses. The counter should NOT - // reset, preventing an infinite 1s reconnect loop. act(() => rfb.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("disconnected"); - // First retry at 1000ms (attempt 0). act(() => vi.advanceTimersByTime(1000)); rfb = getLastRFBInstance(); - // This reconnect also "succeeds" briefly then drops. act(() => rfb.simulateEvent("connect")); act(() => rfb.simulateEvent("disconnect", { clean: false })); - // The next retry should use 2000ms (attempt 1), NOT - // 1000ms. If the counter had been reset on connect, - // this would be 1000ms, creating an infinite loop. + // Should use 2000ms (attempt 1), not 1000ms. mockWatchChatDesktop.mockClear(); act(() => vi.advanceTimersByTime(1999)); expect(mockWatchChatDesktop).not.toHaveBeenCalled(); @@ -636,29 +528,21 @@ describe("useDesktopConnection", () => { useDesktopConnection({ chatId: "chat-1" }), ); - act(() => result.current.connect()); let rfb = getLastRFBInstance(); act(() => rfb.simulateEvent("connect")); expect(result.current.status).toBe("connected"); - // Simulate 10 flapping cycles: each reconnect attempt - // briefly connects then immediately disconnects. Since - // the stability timer never fires, the attempt counter - // keeps incrementing. for (let i = 0; i < 10; i++) { act(() => rfb.simulateEvent("disconnect", { clean: false })); const delay = Math.min(1000 * 2 ** i, 30_000); act(() => vi.advanceTimersByTime(delay)); rfb = getLastRFBInstance(); - // Brief connect then immediate disconnect. act(() => rfb.simulateEvent("connect")); } - // The 11th disconnect should give up. act(() => rfb.simulateEvent("disconnect", { clean: false })); expect(result.current.status).toBe("error"); - // No more retries. mockWatchChatDesktop.mockClear(); act(() => vi.advanceTimersByTime(60_000)); expect(mockWatchChatDesktop).not.toHaveBeenCalled(); @@ -666,4 +550,274 @@ describe("useDesktopConnection", () => { vi.useRealTimers(); } }); + + it("does not retry after securityfailure even if desktop was previously reachable", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + const rfb1 = getLastRFBInstance(); + act(() => rfb1.simulateEvent("connect")); + expect(result.current.status).toBe("connected"); + + act(() => rfb1.simulateEvent("disconnect", { clean: false })); + act(() => vi.advanceTimersByTime(1000)); + const rfb2 = getLastRFBInstance(); + + act(() => + rfb2.simulateEvent("securityfailure", { + status: 1, + reason: "Insufficient resources", + }), + ); + act(() => rfb2.simulateEvent("disconnect", { clean: false })); + + expect(result.current.status).toBe("error"); + + mockWatchChatDesktop.mockClear(); + act(() => vi.advanceTimersByTime(60_000)); + expect(mockWatchChatDesktop).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + // -- Generation counter --------------------------------------------------- + + it("stale event handlers from previous session are ignored after reconnect", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + const rfb1 = getLastRFBInstance(); + act(() => rfb1.simulateEvent("connect")); + expect(result.current.status).toBe("connected"); + + act(() => rfb1.simulateEvent("disconnect", { clean: false })); + expect(result.current.status).toBe("disconnected"); + + // Timer fires, doConnect() runs and bumps generation. + act(() => vi.advanceTimersByTime(1000)); + expect(result.current.status).toBe("connecting"); + + // Old rfb1 fires a late "disconnect" — should be ignored. + act(() => rfb1.simulateEvent("disconnect", { clean: false })); + expect(result.current.status).toBe("connecting"); + } finally { + vi.useRealTimers(); + } + }); + + // -- Connection timeout --------------------------------------------------- + + it("transitions to error when connection times out", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + expect(result.current.status).toBe("connecting"); + + act(() => vi.advanceTimersByTime(29_999)); + expect(result.current.status).toBe("connecting"); + + act(() => vi.advanceTimersByTime(1)); + expect(result.current.status).toBe("error"); + } finally { + vi.useRealTimers(); + } + }); + + it("clears timeout when connection succeeds before deadline", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + const rfb = getLastRFBInstance(); + expect(result.current.status).toBe("connecting"); + + act(() => vi.advanceTimersByTime(15_000)); + expect(result.current.status).toBe("connecting"); + + act(() => rfb.simulateEvent("connect")); + expect(result.current.status).toBe("connected"); + + act(() => vi.advanceTimersByTime(20_000)); + expect(result.current.status).toBe("connected"); + } finally { + vi.useRealTimers(); + } + }); + + // -- Reconnect button ----------------------------------------------------- + + it("reconnect() tears down and restarts the connection", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + const rfb = getLastRFBInstance(); + + // Disconnect before handshake → error. + act(() => rfb.simulateEvent("disconnect", { clean: false })); + expect(result.current.status).toBe("error"); + + // User clicks Reconnect. + mockWatchChatDesktop.mockClear(); + act(() => result.current.reconnect()); + + expect(mockWatchChatDesktop).toHaveBeenCalledWith("chat-1"); + expect(result.current.status).toBe("connecting"); + } finally { + vi.useRealTimers(); + } + }); + + it("reconnect() resets the attempt counter after hitting the cap", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + let rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); + + // Exhaust all 10 reconnect attempts. + for (let i = 0; i < 10; i++) { + act(() => rfb.simulateEvent("disconnect", { clean: false })); + const delay = Math.min(1000 * 2 ** i, 30_000); + act(() => vi.advanceTimersByTime(delay)); + rfb = getLastRFBInstance(); + } + + act(() => rfb.simulateEvent("disconnect", { clean: false })); + expect(result.current.status).toBe("error"); + + // Reconnect should work — counter was reset. + mockWatchChatDesktop.mockClear(); + act(() => result.current.reconnect()); + + expect(result.current.status).toBe("connecting"); + rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); + expect(result.current.status).toBe("connected"); + } finally { + vi.useRealTimers(); + } + }); + + // -- chatId transitions --------------------------------------------------- + + it("resets to idle when chatId becomes undefined", () => { + const { result, rerender } = renderHook( + ({ chatId }: { chatId: string | undefined }) => + useDesktopConnection({ chatId }), + { initialProps: { chatId: "chat-1" as string | undefined } }, + ); + + const rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); + expect(result.current.status).toBe("connected"); + + rerender({ chatId: undefined }); + + expect(rfb.disconnect).toHaveBeenCalled(); + expect(result.current.status).toBe("idle"); + expect(result.current.hasConnected).toBe(false); + }); + + it("cancels pending reconnect timer on chatId change", () => { + vi.useFakeTimers(); + + try { + const { result, rerender } = renderHook( + ({ chatId }: { chatId: string | undefined }) => + useDesktopConnection({ chatId }), + { initialProps: { chatId: "chat-aaa" as string | undefined } }, + ); + + const rfb = getLastRFBInstance(); + act(() => rfb.simulateEvent("connect")); + + // Drop → reconnect timer pending. + act(() => rfb.simulateEvent("disconnect", { clean: false })); + expect(result.current.status).toBe("disconnected"); + + // chatId changes before timer fires. + mockWatchChatDesktop.mockClear(); + rerender({ chatId: "chat-bbb" }); + + // Should connect with new chatId, not fire old timer. + expect(mockWatchChatDesktop).toHaveBeenCalledWith("chat-bbb"); + expect(result.current.status).toBe("connecting"); + + // Old timer should be cancelled — no extra connect. + mockWatchChatDesktop.mockClear(); + act(() => vi.advanceTimersByTime(60_000)); + expect(mockWatchChatDesktop).not.toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + // -- Constructor failure -------------------------------------------------- + + it("closes socket and sets error when RFB constructor throws", () => { + const mockSocket = createMockSocket(); + mockWatchChatDesktop.mockReturnValue(mockSocket); + FakeRFB.throwOnConstruct = true; + + const { result } = renderHook(() => + useDesktopConnection({ chatId: "chat-1" }), + ); + + expect(result.current.status).toBe("error"); + expect(mockSocket.close).toHaveBeenCalled(); + }); + + // -- Stale connect event -------------------------------------------------- + + it("ignores stale connect event from a previous session", () => { + vi.useFakeTimers(); + + try { + const { result, rerender } = renderHook( + ({ chatId }: { chatId: string | undefined }) => + useDesktopConnection({ chatId }), + { initialProps: { chatId: "chat-aaa" as string | undefined } }, + ); + + const rfb1 = getLastRFBInstance(); + // Don't fire connect yet — switch chatId first. + + rerender({ chatId: "chat-bbb" }); + expect(result.current.status).toBe("connecting"); + + // Old rfb1 fires a late "connect" — should be ignored. + act(() => rfb1.simulateEvent("connect")); + + // hasConnected should still be false — the connect was + // from the stale session. + expect(result.current.hasConnected).toBe(false); + expect(result.current.status).toBe("connecting"); + } finally { + vi.useRealTimers(); + } + }); }); diff --git a/site/src/pages/AgentsPage/useDesktopConnection.ts b/site/src/pages/AgentsPage/useDesktopConnection.ts index 3ccfd617b6..9a6771b1f1 100644 --- a/site/src/pages/AgentsPage/useDesktopConnection.ts +++ b/site/src/pages/AgentsPage/useDesktopConnection.ts @@ -13,20 +13,16 @@ type DesktopConnectionStatus = | "disconnected" | "error"; -export interface UseDesktopConnectionResult { +interface UseDesktopConnectionResult { /** Current connection status. */ status: DesktopConnectionStatus; /** Whether the connection has ever been established. */ hasConnected: boolean; /** - * Start the connection. No-op if already connected/connecting. - * Called when the user first opens the Desktop tab. + * Tear down the current connection and start a fresh one. + * Used by the "Reconnect" button after an error. */ - connect: () => void; - /** - * Disconnect and clean up. Called on unmount. - */ - disconnect: () => void; + reconnect: () => void; /** * Attach the noVNC canvas to a container element. Can be called * multiple times (e.g., when the tab is re-selected). The RFB @@ -41,6 +37,7 @@ export interface UseDesktopConnectionResult { const MAX_BACKOFF_MS = 30_000; const MAX_RECONNECT_ATTEMPTS = 10; const STABLE_CONNECTION_MS = 3_000; +const CONNECT_TIMEOUT_MS = 30_000; export function useDesktopConnection({ chatId, @@ -60,154 +57,22 @@ export function useDesktopConnection({ const reconnectStableTimerRef = useRef | null>( null, ); - const disposedRef = useRef(false); - // Track whether connect() has been called at least once. - const connectRequestedRef = useRef(false); - // Ref mirror of hasConnected. Reading the hasConnected *state* - // inside doConnect's event-handler closures would make React - // Compiler track it as a reactive dependency of connect(), - // giving connect a new identity whenever hasConnected changes. - // That would re-fire DesktopPanel's useEffect([connect, - // disconnect]) and tear down a working connection. The ref - // lets event handlers read the latest value without becoming - // a dependency. - const hasConnectedRef = useRef(false); + const connectTimeoutRef = useRef | null>(null); + // Monotonically increasing counter. Incremented at the start of + // every doConnect() and inside the cleanup. Event handlers + // capture the current value and bail when it no longer matches, + // which prevents async noVNC callbacks from a previous session + // (e.g. a phantom "disconnect" event fired after cleanupRfb) + // from scheduling unwanted reconnect timers. + const generationRef = useRef(0); - // Disconnect and clear the current RFB instance. Only reads - // refs and stable setters, so React Compiler memoizes this as - // a singleton — no effect needed. - const cleanupRfb = () => { - if (rfbRef.current) { - try { - rfbRef.current.disconnect(); - } catch { - // Ignore errors during disconnect. - } - rfbRef.current = null; - setRfbInstance(null); - } - }; + // Populated by the lifecycle effect so reconnect() can tear + // down and restart the connection from outside the effect. + // Cleared on cleanup so calls after unmount are no-ops. + const restartRef = useRef<(() => void) | null>(null); - const doConnect = () => { - if (!chatId || disposedRef.current) { - return; - } - - if (reconnectStableTimerRef.current !== null) { - clearTimeout(reconnectStableTimerRef.current); - reconnectStableTimerRef.current = null; - } - - cleanupRfb(); - setStatus("connecting"); - - // Temporary offscreen container for the RFB canvas; moved into - // the visible panel by `attach()`. - offscreenContainerRef.current = document.createElement("div"); - offscreenContainerRef.current.style.width = "100%"; - offscreenContainerRef.current.style.height = "100%"; - - const socket = watchChatDesktop(chatId); - - try { - const rfb = new RFB(offscreenContainerRef.current, socket, { - shared: true, - }); - - rfb.scaleViewport = true; - rfb.resizeSession = false; - - // Track whether this particular RFB instance completed the - // VNC handshake. - let sessionConnected = false; - - rfb.addEventListener("connect", () => { - if (disposedRef.current) return; - sessionConnected = true; - setStatus("connected"); - setHasConnected(true); - hasConnectedRef.current = true; - // Only reset the reconnect counter after the connection - // has been stable for a minimum duration. This prevents - // infinite reconnect loops when the VNC handshake succeeds - // but the connection drops immediately (exit code 1006). - reconnectStableTimerRef.current = setTimeout(() => { - reconnectAttemptRef.current = 0; - }, STABLE_CONNECTION_MS); - }); - - rfb.addEventListener("disconnect", () => { - if (disposedRef.current) return; - if (reconnectStableTimerRef.current !== null) { - clearTimeout(reconnectStableTimerRef.current); - reconnectStableTimerRef.current = null; - } - rfbRef.current = null; - setRfbInstance(null); - - if (!sessionConnected && !hasConnectedRef.current) { - // The VNC handshake never completed and the desktop - // has never been reachable. The endpoint is not - // available (e.g. portabledesktop not installed, - // no workspace, agent down). Don't retry. - setStatus("error"); - return; - } - - const attempt = reconnectAttemptRef.current; - - if (attempt >= MAX_RECONNECT_ATTEMPTS) { - // Too many consecutive failures. Give up. - setStatus("error"); - return; - } - - setStatus("disconnected"); - - // Either this session was connected and dropped, or a - // previous session was connected (transient reconnect - // failure). Retry with exponential backoff. - const delay = Math.min(1000 * 2 ** attempt, MAX_BACKOFF_MS); - reconnectAttemptRef.current = attempt + 1; - reconnectTimerRef.current = setTimeout(doConnect, delay); - }); - - rfb.addEventListener("securityfailure", () => { - if (disposedRef.current) return; - rfbRef.current = null; - setRfbInstance(null); - setStatus("error"); - }); - - rfbRef.current = rfb; - setRfbInstance(rfb); - } catch { - setStatus("error"); - } - }; - - const connect = () => { - if (connectRequestedRef.current) { - return; - } - connectRequestedRef.current = true; - doConnect(); - }; - - const disconnect = () => { - if (reconnectTimerRef.current !== null) { - clearTimeout(reconnectTimerRef.current); - reconnectTimerRef.current = null; - } - if (reconnectStableTimerRef.current !== null) { - clearTimeout(reconnectStableTimerRef.current); - reconnectStableTimerRef.current = null; - } - cleanupRfb(); - offscreenContainerRef.current = null; - setStatus("idle"); - connectRequestedRef.current = false; - reconnectAttemptRef.current = 0; + const reconnect = () => { + restartRef.current?.(); }; const attach = (container: HTMLElement) => { @@ -217,13 +82,28 @@ export function useDesktopConnection({ } }; - // Cleanup on unmount or chatId change. - // biome-ignore lint/correctness/useExhaustiveDependencies: chatId is an intentional trigger to reset state for a new conversation + // Single lifecycle effect that owns the entire connection. + // Connects on mount, tears down on unmount, and resets when + // chatId changes. + // + // All connection logic (doConnect, cleanupRfb, timer helpers) + // is defined inside the effect so the dependency array only + // contains primitives. This avoids any reliance on the React + // Compiler for function identity stability. useEffect(() => { - disposedRef.current = false; + const cleanupRfb = () => { + if (rfbRef.current) { + try { + rfbRef.current.disconnect(); + } catch { + // Ignore errors during disconnect. + } + rfbRef.current = null; + setRfbInstance(null); + } + }; - return () => { - disposedRef.current = true; + const clearAllTimers = () => { if (reconnectTimerRef.current !== null) { clearTimeout(reconnectTimerRef.current); reconnectTimerRef.current = null; @@ -232,21 +112,181 @@ export function useDesktopConnection({ clearTimeout(reconnectStableTimerRef.current); reconnectStableTimerRef.current = null; } + if (connectTimeoutRef.current !== null) { + clearTimeout(connectTimeoutRef.current); + connectTimeoutRef.current = null; + } + }; + + const teardown = () => { + generationRef.current++; + clearAllTimers(); cleanupRfb(); + if (offscreenContainerRef.current?.parentElement) { + offscreenContainerRef.current.remove(); + } offscreenContainerRef.current = null; + reconnectAttemptRef.current = 0; + }; + + const doConnect = () => { + if (!chatId) { + return; + } + + // Bump the generation so any in-flight async callbacks + // from the previous RFB session become stale. + generationRef.current++; + const gen = generationRef.current; + + clearAllTimers(); + cleanupRfb(); + setStatus("connecting"); + + // Remove the previous offscreen container from the DOM + // so reconnect cycles don't leak detached divs. + if (offscreenContainerRef.current?.parentElement) { + offscreenContainerRef.current.remove(); + } + + // Temporary offscreen container for the RFB canvas; + // moved into the visible panel by `attach()`. + offscreenContainerRef.current = document.createElement("div"); + offscreenContainerRef.current.style.width = "100%"; + offscreenContainerRef.current.style.height = "100%"; + + const socket = watchChatDesktop(chatId); + + try { + const rfb = new RFB(offscreenContainerRef.current, socket, { + shared: true, + }); + + rfb.scaleViewport = true; + rfb.resizeSession = false; + + // Per-session flags scoped to this RFB instance. + // NOT refs — each doConnect() gets fresh copies so + // state from a previous session cannot leak. + let sessionConnected = false; + let securityFailed = false; + + // Fail if the VNC handshake doesn't complete within + // a reasonable window. Bump the generation before + // cleanupRfb() so the resulting noVNC "disconnect" + // event is treated as stale and doesn't redundantly + // set status to "error". + connectTimeoutRef.current = setTimeout(() => { + if (gen !== generationRef.current) return; + if (!sessionConnected) { + generationRef.current++; + cleanupRfb(); + setStatus("error"); + } + }, CONNECT_TIMEOUT_MS); + + rfb.addEventListener("connect", () => { + if (gen !== generationRef.current) return; + sessionConnected = true; + if (connectTimeoutRef.current !== null) { + clearTimeout(connectTimeoutRef.current); + connectTimeoutRef.current = null; + } + setStatus("connected"); + setHasConnected(true); + // Only reset the reconnect counter after the + // connection has been stable for a minimum + // duration. This prevents infinite reconnect + // loops when the VNC handshake succeeds but the + // connection drops immediately. + // No gen check needed — clearAllTimers() + // always cancels this before a new session. + reconnectStableTimerRef.current = setTimeout(() => { + reconnectAttemptRef.current = 0; + }, STABLE_CONNECTION_MS); + }); + + rfb.addEventListener("disconnect", () => { + if (gen !== generationRef.current) return; + if (reconnectStableTimerRef.current !== null) { + clearTimeout(reconnectStableTimerRef.current); + reconnectStableTimerRef.current = null; + } + if (connectTimeoutRef.current !== null) { + clearTimeout(connectTimeoutRef.current); + connectTimeoutRef.current = null; + } + rfbRef.current = null; + setRfbInstance(null); + + // Security failures are terminal — the + // securityfailure handler already moved to + // "error". noVNC fires disconnect after + // securityfailure; ignore it so we don't + // accidentally schedule a retry. + if (securityFailed) { + return; + } + + // Only retry if THIS session's VNC handshake + // completed. A previous session having connected + // is irrelevant — the desktop may have become + // permanently unavailable. + if (!sessionConnected) { + setStatus("error"); + return; + } + + const attempt = reconnectAttemptRef.current; + + if (attempt >= MAX_RECONNECT_ATTEMPTS) { + setStatus("error"); + return; + } + + setStatus("disconnected"); + + const delay = Math.min(1000 * 2 ** attempt, MAX_BACKOFF_MS); + reconnectAttemptRef.current = attempt + 1; + reconnectTimerRef.current = setTimeout(doConnect, delay); + }); + + rfb.addEventListener("securityfailure", () => { + if (gen !== generationRef.current) return; + securityFailed = true; + rfbRef.current = null; + setRfbInstance(null); + setStatus("error"); + }); + + rfbRef.current = rfb; + setRfbInstance(rfb); + } catch { + socket.close(); + setStatus("error"); + } + }; + + // Expose a restart handle for the Reconnect button. + restartRef.current = () => { + teardown(); + doConnect(); + }; + + doConnect(); + + return () => { + restartRef.current = null; + teardown(); setStatus("idle"); setHasConnected(false); - hasConnectedRef.current = false; - connectRequestedRef.current = false; - reconnectAttemptRef.current = 0; }; }, [chatId]); return { status, hasConnected, - connect, - disconnect, + reconnect, attach, rfb: rfbInstance, };