opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(site): fix desktop reconnect loop by moving connection lifecycle into hook (#23404)

Browse Source
This commit is contained in:
Danielle Maywood
2026-03-22 02:01:33 +00:00
committed by GitHub
parent bf702cc3b9
commit a08b6848f2
7 changed files with 580 additions and 498 deletions
Download Patch File Download Diff File Expand all files Collapse all files
site/src/pages/AgentsPage/DesktopPanel.stories.tsx
+21 -31
View File
@@ -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<typeof DesktopPanel> = {
const defaults: DesktopPanelViewProps = {
status: "idle",
reconnect: fn(),
attach: fn(),
};
const meta: Meta<typeof DesktopPanelView> = {
title: "pages/AgentsPage/DesktopPanel",
component: DesktopPanel,
args: {
isExpanded: false,
chatId: "test-chat-id",
},
component: DesktopPanelView,
args: defaults,
decorators: [
(Story) => (
<div style={{ height: 400, width: 480, border: "1px solid #333" }}>
@@ -18,35 +21,22 @@ const meta: Meta<typeof DesktopPanel> = {
],
};
export default meta;
type Story = StoryObj<typeof DesktopPanel>;
type Story = StoryObj<typeof DesktopPanelView>;
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" },
};
site/src/pages/AgentsPage/DesktopPanel.tsx
+38 -57
View File
@@ -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<DesktopPanelProps> = ({
chatId,
isExpanded: _isExpanded,
connectionOverride,
export interface DesktopPanelViewProps {
status: DesktopConnectionStatus;
reconnect: () => void;
attach: (container: HTMLElement) => void;
}
export const DesktopPanel: FC<DesktopPanelProps> = ({ chatId }) => {
const { status, reconnect, attach } = useDesktopConnection({ chatId });
return (
<DesktopPanelView status={status} reconnect={reconnect} attach={attach} />
);
};
export const DesktopPanelView: FC<DesktopPanelViewProps> = ({
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<HTMLDivElement | null>(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 (
<div className="flex h-full flex-col items-center justify-center gap-2 text-content-secondary">
@@ -73,17 +53,11 @@ export const DesktopPanel: FC<DesktopPanelProps> = ({
if (status === "error") {
return (
<div className="flex h-full flex-col items-center justify-center gap-3 text-content-secondary">
<span className="text-sm">
Failed to connect to the desktop session.
<span className="text-center text-sm">
Failed to connect to the desktop session. The agent may not be
connected or the desktop environment may not be available.
</span>
<Button
variant="outline"
size="sm"
onClick={() => {
disconnect();
connect();
}}
>
<Button variant="outline" size="sm" onClick={reconnect}>
Reconnect
</Button>
</div>
@@ -100,5 +74,12 @@ export const DesktopPanel: FC<DesktopPanelProps> = ({
}
// status === "connected"
return <div ref={attachToContainer} className="h-full w-full" />;
return (
<div
ref={(el) => {
if (el) attach(el);
}}
className="h-full w-full"
/>
);
};
site/src/pages/AgentsPage/SidebarTabView.stories.tsx
-29
View File
@@ -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" }),
},
};
site/src/pages/AgentsPage/SidebarTabView.tsx
+1 -9
View File
@@ -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<SidebarTabViewProps> = ({
chatTitle,
onClose,
desktopChatId,
desktopConnectionOverride,
}) => {
const tabIdPrefix = useId();
const [activeTabId, setActiveTabId] = useState<string | null>(
@@ -324,11 +320,7 @@ export const SidebarTabView: FC<SidebarTabViewProps> = ({
className="min-h-0 flex-1"
>
{effectiveTabId === "desktop" && desktopChatId ? (
<DesktopPanel
chatId={desktopChatId}
isExpanded={isExpanded}
connectionOverride={desktopConnectionOverride}
/>
<DesktopPanel chatId={desktopChatId} />
) : (
activeTab?.content
)}
site/src/pages/AgentsPage/desktopStoryUtils.ts
-46
View File
@@ -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> = {},
): UseDesktopConnectionResult {
return {
status: "idle",
hasConnected: false,
connect: fn(),
disconnect: fn(),
attach: fn(),
rfb: null,
...overrides,
};
}
site/src/pages/AgentsPage/useDesktopConnection.test.ts
+316 -162
View File
@@ -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<string, (ev: unknown) => 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();
}
});
});
site/src/pages/AgentsPage/useDesktopConnection.ts
+204 -164
View File
@@ -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<ReturnType<typeof setTimeout> | 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<ReturnType<typeof setTimeout> | 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,
};