diff --git a/site/package.json b/site/package.json index 1512a803b0a96..e3a99b9d8eebf 100644 --- a/site/package.json +++ b/site/package.json @@ -120,6 +120,7 @@ "undici": "6.21.2", "unique-names-generator": "4.7.1", "uuid": "9.0.1", + "websocket-ts": "2.2.1", "yup": "1.6.1" }, "devDependencies": { diff --git a/site/pnpm-lock.yaml b/site/pnpm-lock.yaml index 62cdc6176092a..3c7f5176b5b6b 100644 --- a/site/pnpm-lock.yaml +++ b/site/pnpm-lock.yaml @@ -274,6 +274,9 @@ importers: uuid: specifier: 9.0.1 version: 9.0.1 + websocket-ts: + specifier: 2.2.1 + version: 2.2.1 yup: specifier: 1.6.1 version: 1.6.1 @@ -6344,6 +6347,9 @@ packages: webpack-virtual-modules@0.5.0: resolution: {integrity: sha512-kyDivFZ7ZM0BVOUteVbDFhlRt7Ah/CSPwJdi8hBpkK7QLumUqdLtVfm/PX/hkcnrvr0i77fO5+TjZ94Pe+C9iw==, tarball: https://registry.npmjs.org/webpack-virtual-modules/-/webpack-virtual-modules-0.5.0.tgz} + websocket-ts@2.2.1: + resolution: {integrity: sha512-YKPDfxlK5qOheLZ2bTIiktZO1bpfGdNCPJmTEaPW7G9UXI1GKjDdeacOrsULUS000OPNxDVOyAuKLuIWPqWM0Q==, tarball: https://registry.npmjs.org/websocket-ts/-/websocket-ts-2.2.1.tgz} + whatwg-encoding@2.0.0: resolution: {integrity: sha512-p41ogyeMUrw3jWclHWTQg1k05DSVXPLcVxRTYsXUk+ZooOCZLcoYgPZ/HL/D/N+uQPOtcp1me1WhBEaX02mhWg==, tarball: https://registry.npmjs.org/whatwg-encoding/-/whatwg-encoding-2.0.0.tgz} engines: {node: '>=12'} @@ -13266,6 +13272,8 @@ snapshots: webpack-virtual-modules@0.5.0: {} + websocket-ts@2.2.1: {} + whatwg-encoding@2.0.0: dependencies: iconv-lite: 0.6.3 diff --git a/site/src/hooks/index.ts b/site/src/hooks/index.ts index 4453e36fa4bb4..901fee8a50ded 100644 --- a/site/src/hooks/index.ts +++ b/site/src/hooks/index.ts @@ -3,4 +3,3 @@ export * from "./useClickable"; export * from "./useClickableTableRow"; export * from "./useClipboard"; export * from "./usePagination"; -export * from "./useWithRetry"; diff --git a/site/src/hooks/useWithRetry.test.ts b/site/src/hooks/useWithRetry.test.ts deleted file mode 100644 index 7ed7b4331f21e..0000000000000 --- a/site/src/hooks/useWithRetry.test.ts +++ /dev/null @@ -1,329 +0,0 @@ -import { act, renderHook } from "@testing-library/react"; -import { useWithRetry } from "./useWithRetry"; - -// Mock timers -jest.useFakeTimers(); - -describe("useWithRetry", () => { - let mockFn: jest.Mock; - - beforeEach(() => { - mockFn = jest.fn(); - jest.clearAllTimers(); - }); - - afterEach(() => { - jest.clearAllMocks(); - }); - - it("should initialize with correct default state", () => { - const { result } = renderHook(() => useWithRetry(mockFn)); - - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).toBe(undefined); - }); - - it("should execute function successfully on first attempt", async () => { - mockFn.mockResolvedValue(undefined); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - await act(async () => { - await result.current.call(); - }); - - expect(mockFn).toHaveBeenCalledTimes(1); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).toBe(undefined); - }); - - it("should set isLoading to true during execution", async () => { - let resolvePromise: () => void; - const promise = new Promise((resolve) => { - resolvePromise = resolve; - }); - mockFn.mockReturnValue(promise); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - act(() => { - result.current.call(); - }); - - expect(result.current.isLoading).toBe(true); - - await act(async () => { - resolvePromise!(); - await promise; - }); - - expect(result.current.isLoading).toBe(false); - }); - - it("should retry on failure with exponential backoff", async () => { - mockFn - .mockRejectedValueOnce(new Error("First failure")) - .mockRejectedValueOnce(new Error("Second failure")) - .mockResolvedValueOnce(undefined); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the call - await act(async () => { - await result.current.call(); - }); - - expect(mockFn).toHaveBeenCalledTimes(1); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Fast-forward to first retry (1 second) - await act(async () => { - jest.advanceTimersByTime(1000); - }); - - expect(mockFn).toHaveBeenCalledTimes(2); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Fast-forward to second retry (2 seconds) - await act(async () => { - jest.advanceTimersByTime(2000); - }); - - expect(mockFn).toHaveBeenCalledTimes(3); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).toBe(undefined); - }); - - it("should continue retrying without limit", async () => { - mockFn.mockRejectedValue(new Error("Always fails")); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the call - await act(async () => { - await result.current.call(); - }); - - expect(mockFn).toHaveBeenCalledTimes(1); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Fast-forward through multiple retries to verify it continues - for (let i = 1; i < 15; i++) { - const delay = Math.min(1000 * 2 ** (i - 1), 600000); // exponential backoff with max delay - await act(async () => { - jest.advanceTimersByTime(delay); - }); - expect(mockFn).toHaveBeenCalledTimes(i + 1); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - } - - // Should still be retrying after 15 attempts - expect(result.current.nextRetryAt).not.toBe(null); - }); - - it("should respect max delay of 10 minutes", async () => { - mockFn.mockRejectedValue(new Error("Always fails")); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the call - await act(async () => { - await result.current.call(); - }); - - expect(result.current.isLoading).toBe(false); - - // Fast-forward through several retries to reach max delay - // After attempt 9, delay would be 1000 * 2^9 = 512000ms, which is less than 600000ms (10 min) - // After attempt 10, delay would be 1000 * 2^10 = 1024000ms, which should be capped at 600000ms - - // Skip to attempt 9 (delay calculation: 1000 * 2^8 = 256000ms) - for (let i = 1; i < 9; i++) { - const delay = 1000 * 2 ** (i - 1); - await act(async () => { - jest.advanceTimersByTime(delay); - }); - } - - expect(mockFn).toHaveBeenCalledTimes(9); - expect(result.current.nextRetryAt).not.toBe(null); - - // The 9th retry should use max delay (600000ms = 10 minutes) - await act(async () => { - jest.advanceTimersByTime(600000); - }); - - expect(mockFn).toHaveBeenCalledTimes(10); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Continue with more retries at max delay to verify it continues indefinitely - await act(async () => { - jest.advanceTimersByTime(600000); - }); - - expect(mockFn).toHaveBeenCalledTimes(11); - expect(result.current.nextRetryAt).not.toBe(null); - }); - - it("should cancel previous retry when call is invoked again", async () => { - mockFn - .mockRejectedValueOnce(new Error("First failure")) - .mockResolvedValueOnce(undefined); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the first call - await act(async () => { - await result.current.call(); - }); - - expect(mockFn).toHaveBeenCalledTimes(1); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Call again before retry happens - await act(async () => { - await result.current.call(); - }); - - expect(mockFn).toHaveBeenCalledTimes(2); - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).toBe(undefined); - - // Advance time to ensure previous retry was cancelled - await act(async () => { - jest.advanceTimersByTime(5000); - }); - - expect(mockFn).toHaveBeenCalledTimes(2); // Should not have been called again - }); - - it("should set nextRetryAt when scheduling retry", async () => { - mockFn - .mockRejectedValueOnce(new Error("Failure")) - .mockResolvedValueOnce(undefined); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the call - await act(async () => { - await result.current.call(); - }); - - const nextRetryAt = result.current.nextRetryAt; - expect(nextRetryAt).not.toBe(null); - expect(nextRetryAt).toBeInstanceOf(Date); - - // nextRetryAt should be approximately 1 second in the future - const expectedTime = Date.now() + 1000; - const actualTime = nextRetryAt!.getTime(); - expect(Math.abs(actualTime - expectedTime)).toBeLessThan(100); // Allow 100ms tolerance - - // Advance past retry time - await act(async () => { - jest.advanceTimersByTime(1000); - }); - - expect(result.current.nextRetryAt).toBe(undefined); - }); - - it("should cleanup timer on unmount", async () => { - mockFn.mockRejectedValue(new Error("Failure")); - - const { result, unmount } = renderHook(() => useWithRetry(mockFn)); - - // Start the call to create timer - await act(async () => { - await result.current.call(); - }); - - expect(result.current.isLoading).toBe(false); - expect(result.current.nextRetryAt).not.toBe(null); - - // Unmount should cleanup timer - unmount(); - - // Advance time to ensure timer was cleared - await act(async () => { - jest.advanceTimersByTime(5000); - }); - - // Function should not have been called again - expect(mockFn).toHaveBeenCalledTimes(1); - }); - - it("should prevent scheduling retries when function completes after unmount", async () => { - let rejectPromise: (error: Error) => void; - const promise = new Promise((_, reject) => { - rejectPromise = reject; - }); - mockFn.mockReturnValue(promise); - - const { result, unmount } = renderHook(() => useWithRetry(mockFn)); - - // Start the call - this will make the function in-flight - act(() => { - result.current.call(); - }); - - expect(result.current.isLoading).toBe(true); - - // Unmount while function is still in-flight - unmount(); - - // Function completes with error after unmount - await act(async () => { - rejectPromise!(new Error("Failed after unmount")); - await promise.catch(() => {}); // Suppress unhandled rejection - }); - - // Advance time to ensure no retry timers were scheduled - await act(async () => { - jest.advanceTimersByTime(5000); - }); - - // Function should only have been called once (no retries after unmount) - expect(mockFn).toHaveBeenCalledTimes(1); - }); - - it("should do nothing when call() is invoked while function is already loading", async () => { - let resolvePromise: () => void; - const promise = new Promise((resolve) => { - resolvePromise = resolve; - }); - mockFn.mockReturnValue(promise); - - const { result } = renderHook(() => useWithRetry(mockFn)); - - // Start the first call - this will set isLoading to true - act(() => { - result.current.call(); - }); - - expect(result.current.isLoading).toBe(true); - expect(mockFn).toHaveBeenCalledTimes(1); - - // Try to call again while loading - should do nothing - act(() => { - result.current.call(); - }); - - // Function should not have been called again - expect(mockFn).toHaveBeenCalledTimes(1); - expect(result.current.isLoading).toBe(true); - - // Complete the original promise - await act(async () => { - resolvePromise!(); - await promise; - }); - - expect(result.current.isLoading).toBe(false); - expect(mockFn).toHaveBeenCalledTimes(1); - }); -}); diff --git a/site/src/hooks/useWithRetry.ts b/site/src/hooks/useWithRetry.ts deleted file mode 100644 index 1310da221efc5..0000000000000 --- a/site/src/hooks/useWithRetry.ts +++ /dev/null @@ -1,106 +0,0 @@ -import { useCallback, useEffect, useRef, useState } from "react"; -import { useEffectEvent } from "./hookPolyfills"; - -const DELAY_MS = 1_000; -const MAX_DELAY_MS = 600_000; // 10 minutes -// Determines how much the delay between retry attempts increases after each -// failure. -const MULTIPLIER = 2; - -interface UseWithRetryResult { - call: () => void; - nextRetryAt: Date | undefined; - isLoading: boolean; -} - -interface RetryState { - isLoading: boolean; - nextRetryAt: Date | undefined; -} - -/** - * Hook that wraps a function with automatic retry functionality - * Provides a simple interface for executing functions with exponential backoff retry - */ -export function useWithRetry(fn: () => Promise): UseWithRetryResult { - const [state, setState] = useState({ - isLoading: false, - nextRetryAt: undefined, - }); - - const timeoutRef = useRef(null); - const mountedRef = useRef(true); - - const clearTimeout = useCallback(() => { - if (timeoutRef.current) { - window.clearTimeout(timeoutRef.current); - timeoutRef.current = null; - } - }, []); - - const stableFn = useEffectEvent(fn); - - const call = useCallback(() => { - if (state.isLoading) { - return; - } - - clearTimeout(); - - const executeAttempt = async (attempt = 0): Promise => { - if (!mountedRef.current) { - return; - } - setState({ - isLoading: true, - nextRetryAt: undefined, - }); - - try { - await stableFn(); - if (mountedRef.current) { - setState({ isLoading: false, nextRetryAt: undefined }); - } - } catch (error) { - if (!mountedRef.current) { - return; - } - const delayMs = Math.min( - DELAY_MS * MULTIPLIER ** attempt, - MAX_DELAY_MS, - ); - - setState({ - isLoading: false, - nextRetryAt: new Date(Date.now() + delayMs), - }); - - timeoutRef.current = window.setTimeout(() => { - if (!mountedRef.current) { - return; - } - setState({ - isLoading: false, - nextRetryAt: undefined, - }); - executeAttempt(attempt + 1); - }, delayMs); - } - }; - - executeAttempt(); - }, [state.isLoading, stableFn, clearTimeout]); - - useEffect(() => { - return () => { - mountedRef.current = false; - clearTimeout(); - }; - }, [clearTimeout]); - - return { - call, - nextRetryAt: state.nextRetryAt, - isLoading: state.isLoading, - }; -} diff --git a/site/src/pages/TerminalPage/TerminalAlerts.tsx b/site/src/pages/TerminalPage/TerminalAlerts.tsx index 07740135769f3..6a06a76964128 100644 --- a/site/src/pages/TerminalPage/TerminalAlerts.tsx +++ b/site/src/pages/TerminalPage/TerminalAlerts.tsx @@ -170,14 +170,16 @@ const TerminalAlert: FC = (props) => { ); }; +// Since the terminal connection is always trying to reconnect, we show this +// alert to indicate that the terminal is trying to connect. const DisconnectedAlert: FC = (props) => { return ( } > - Disconnected + Trying to connect... ); }; diff --git a/site/src/pages/TerminalPage/TerminalPage.test.tsx b/site/src/pages/TerminalPage/TerminalPage.test.tsx index 7600fa5257d43..4591190ad9904 100644 --- a/site/src/pages/TerminalPage/TerminalPage.test.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.test.tsx @@ -85,7 +85,7 @@ describe("TerminalPage", () => { await expectTerminalText(container, Language.workspaceErrorMessagePrefix); }); - it("shows an error if the websocket fails", async () => { + it("shows reconnect message when websocket fails", async () => { server.use( http.get("/api/v2/workspaceagents/:agentId/pty", () => { return HttpResponse.json({}, { status: 500 }); @@ -94,7 +94,9 @@ describe("TerminalPage", () => { const { container } = await renderTerminal(); - await expectTerminalText(container, Language.websocketErrorMessagePrefix); + await waitFor(() => { + expect(container.textContent).toContain("Trying to connect..."); + }); }); it("renders data from the backend", async () => { diff --git a/site/src/pages/TerminalPage/TerminalPage.tsx b/site/src/pages/TerminalPage/TerminalPage.tsx index 2023bdb0eeb29..5c13e89c30005 100644 --- a/site/src/pages/TerminalPage/TerminalPage.tsx +++ b/site/src/pages/TerminalPage/TerminalPage.tsx @@ -26,6 +26,13 @@ import { openMaybePortForwardedURL } from "utils/portForward"; import { terminalWebsocketUrl } from "utils/terminal"; import { getMatchingAgentOrFirst } from "utils/workspace"; import { v4 as uuidv4 } from "uuid"; +// Use websocket-ts for better WebSocket handling and auto-reconnection. +import { + ExponentialBackoff, + type Websocket, + WebsocketBuilder, + WebsocketEvent, +} from "websocket-ts"; import { TerminalAlerts } from "./TerminalAlerts"; import type { ConnectionStatus } from "./types"; @@ -221,7 +228,7 @@ const TerminalPage: FC = () => { } // Hook up terminal events to the websocket. - let websocket: WebSocket | null; + let websocket: Websocket | null; const disposers = [ terminal.onData((data) => { websocket?.send( @@ -259,9 +266,11 @@ const TerminalPage: FC = () => { if (disposed) { return; // Unmounted while we waited for the async call. } - websocket = new WebSocket(url); + websocket = new WebsocketBuilder(url) + .withBackoff(new ExponentialBackoff(1000, 6)) + .build(); websocket.binaryType = "arraybuffer"; - websocket.addEventListener("open", () => { + websocket.addEventListener(WebsocketEvent.open, () => { // Now that we are connected, allow user input. terminal.options = { disableStdin: false, @@ -278,18 +287,16 @@ const TerminalPage: FC = () => { ); setConnectionStatus("connected"); }); - websocket.addEventListener("error", () => { + websocket.addEventListener(WebsocketEvent.error, (_, event) => { + console.error("WebSocket error:", event); terminal.options.disableStdin = true; - terminal.writeln( - `${Language.websocketErrorMessagePrefix}socket errored`, - ); setConnectionStatus("disconnected"); }); - websocket.addEventListener("close", () => { + websocket.addEventListener(WebsocketEvent.close, () => { terminal.options.disableStdin = true; setConnectionStatus("disconnected"); }); - websocket.addEventListener("message", (event) => { + websocket.addEventListener(WebsocketEvent.message, (_, event) => { if (typeof event.data === "string") { // This exclusively occurs when testing. // "jest-websocket-mock" doesn't support ArrayBuffer. @@ -298,12 +305,25 @@ const TerminalPage: FC = () => { terminal.write(new Uint8Array(event.data)); } }); + websocket.addEventListener(WebsocketEvent.reconnect, () => { + if (websocket) { + websocket.binaryType = "arraybuffer"; + websocket.send( + new TextEncoder().encode( + JSON.stringify({ + height: terminal.rows, + width: terminal.cols, + }), + ), + ); + } + }); }) .catch((error) => { if (disposed) { return; // Unmounted while we waited for the async call. } - terminal.writeln(Language.websocketErrorMessagePrefix + error.message); + console.error("WebSocket connection failed:", error); setConnectionStatus("disconnected"); });