From ecc28a66508e24d6e4750a59a362567579753b3b Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Sat, 21 Mar 2026 13:34:00 +0000 Subject: [PATCH] fix(site): prevent infinite desktop reconnect loop on exit code 1006 (#23401) --- site/src/pages/AgentsPage/DesktopPanel.tsx | 9 +- .../AgentsPage/useDesktopConnection.test.ts | 89 ++++++++++++++++++- .../pages/AgentsPage/useDesktopConnection.ts | 75 ++++++++++++---- 3 files changed, 151 insertions(+), 22 deletions(-) diff --git a/site/src/pages/AgentsPage/DesktopPanel.tsx b/site/src/pages/AgentsPage/DesktopPanel.tsx index 53ba6b4a2c..3f53189a98 100644 --- a/site/src/pages/AgentsPage/DesktopPanel.tsx +++ b/site/src/pages/AgentsPage/DesktopPanel.tsx @@ -76,7 +76,14 @@ export const DesktopPanel: FC = ({ Failed to connect to the desktop session. - diff --git a/site/src/pages/AgentsPage/useDesktopConnection.test.ts b/site/src/pages/AgentsPage/useDesktopConnection.test.ts index 03849352be..1197cfb2dd 100644 --- a/site/src/pages/AgentsPage/useDesktopConnection.test.ts +++ b/site/src/pages/AgentsPage/useDesktopConnection.test.ts @@ -250,10 +250,13 @@ describe("useDesktopConnection", () => { act(() => vi.advanceTimersByTime(1000)); const rfb2 = getLastRFBInstance(); - // Reconnect succeeds — counter should reset. + // Reconnect succeeds — counter resets after stability period. act(() => rfb2.simulateEvent("connect")); expect(result.current.status).toBe("connected"); + // Advance past the stability period so the counter resets. + act(() => vi.advanceTimersByTime(3000)); + // Next disconnect should use 1000ms again (not 2000ms). act(() => rfb2.simulateEvent("disconnect", { clean: false })); @@ -566,10 +569,13 @@ describe("useDesktopConnection", () => { rfb = getLastRFBInstance(); } - // This reconnect succeeds — counter resets. + // 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 })); @@ -581,4 +587,83 @@ describe("useDesktopConnection", () => { vi.useRealTimers(); } }); + + it("does not reset reconnect counter when connection drops before stability period", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + 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. + mockWatchChatDesktop.mockClear(); + act(() => vi.advanceTimersByTime(1999)); + expect(mockWatchChatDesktop).not.toHaveBeenCalled(); + act(() => vi.advanceTimersByTime(1)); + expect(mockWatchChatDesktop).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }); + + it("gives up when connection keeps flapping (connect then immediate disconnect)", () => { + vi.useFakeTimers(); + + try { + const { result } = renderHook(() => + 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(); + } finally { + vi.useRealTimers(); + } + }); }); diff --git a/site/src/pages/AgentsPage/useDesktopConnection.ts b/site/src/pages/AgentsPage/useDesktopConnection.ts index 84fa2f8449..3ccfd617b6 100644 --- a/site/src/pages/AgentsPage/useDesktopConnection.ts +++ b/site/src/pages/AgentsPage/useDesktopConnection.ts @@ -40,6 +40,7 @@ export interface UseDesktopConnectionResult { const MAX_BACKOFF_MS = 30_000; const MAX_RECONNECT_ATTEMPTS = 10; +const STABLE_CONNECTION_MS = 3_000; export function useDesktopConnection({ chatId, @@ -47,39 +48,57 @@ export function useDesktopConnection({ const [status, setStatus] = useState("idle"); const [hasConnected, setHasConnected] = useState(false); + // rfbRef provides synchronous access for cleanup and event + // handlers. rfbInstance (state) provides reactivity so consumers + // re-render when the RFB instance changes. const [rfbInstance, setRfbInstance] = useState(null); const rfbRef = useRef(null); + const offscreenContainerRef = useRef(null); const reconnectAttemptRef = useRef(0); const reconnectTimerRef = useRef | null>(null); + 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 so disconnect handlers can read - // the latest value without stale closures. + // 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 cleanupRfbRef = useRef(() => {}); - useEffect(() => { - cleanupRfbRef.current = () => { - if (rfbRef.current) { - try { - rfbRef.current.disconnect(); - } catch { - // Ignore errors during disconnect. - } - rfbRef.current = null; - setRfbInstance(null); + // 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); + } + }; const doConnect = () => { if (!chatId || disposedRef.current) { return; } - cleanupRfbRef.current(); + if (reconnectStableTimerRef.current !== null) { + clearTimeout(reconnectStableTimerRef.current); + reconnectStableTimerRef.current = null; + } + + cleanupRfb(); setStatus("connecting"); // Temporary offscreen container for the RFB canvas; moved into @@ -108,11 +127,21 @@ export function useDesktopConnection({ setStatus("connected"); setHasConnected(true); hasConnectedRef.current = true; - reconnectAttemptRef.current = 0; + // 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); @@ -170,7 +199,11 @@ export function useDesktopConnection({ clearTimeout(reconnectTimerRef.current); reconnectTimerRef.current = null; } - cleanupRfbRef.current(); + if (reconnectStableTimerRef.current !== null) { + clearTimeout(reconnectStableTimerRef.current); + reconnectStableTimerRef.current = null; + } + cleanupRfb(); offscreenContainerRef.current = null; setStatus("idle"); connectRequestedRef.current = false; @@ -195,7 +228,11 @@ export function useDesktopConnection({ clearTimeout(reconnectTimerRef.current); reconnectTimerRef.current = null; } - cleanupRfbRef.current(); + if (reconnectStableTimerRef.current !== null) { + clearTimeout(reconnectStableTimerRef.current); + reconnectStableTimerRef.current = null; + } + cleanupRfb(); offscreenContainerRef.current = null; setStatus("idle"); setHasConnected(false);