feat: auto reconnect the terminal (#18796)

**Changes:**
- Use [websocket-ts](https://www.npmjs.com/package/websocket-ts) to have
auto reconnection out of the box 🙏
- Update the disconnected alert message to "Trying to connect..." since
the connection is always trying to reconnect
- Remove `useWithRetry` because it is not necessary anymore

**Other topics:**
- The disconnected alert is displaying a simple message, but we can
include more info such as the number of attemtps
- The reconnection feature is in a good state and adding value. IMO, any
improvement can be done after getting this merged

Closes https://github.com/coder/internal/issues/659
This commit is contained in:
Bruno Quaresma
2025-07-09 15:04:24 -03:00
committed by GitHub
parent 00ba0278d2
commit 5a8a19be70
8 changed files with 50 additions and 453 deletions

View File

@ -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": {

8
site/pnpm-lock.yaml generated
View File

@ -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

View File

@ -3,4 +3,3 @@ export * from "./useClickable";
export * from "./useClickableTableRow";
export * from "./useClipboard";
export * from "./usePagination";
export * from "./useWithRetry";

View File

@ -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<void>((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<void>((_, 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<void>((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);
});
});

View File

@ -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<void>): UseWithRetryResult {
const [state, setState] = useState<RetryState>({
isLoading: false,
nextRetryAt: undefined,
});
const timeoutRef = useRef<number | null>(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<void> => {
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,
};
}

View File

@ -170,14 +170,16 @@ const TerminalAlert: FC<AlertProps> = (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<AlertProps> = (props) => {
return (
<TerminalAlert
{...props}
severity="warning"
severity="info"
actions={<RefreshSessionButton />}
>
Disconnected
Trying to connect...
</TerminalAlert>
);
};

View File

@ -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 () => {

View File

@ -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", () => {
terminal.options.disableStdin = true;
terminal.writeln(
`${Language.websocketErrorMessagePrefix}socket errored`,
);
setConnectionStatus("disconnected");
});
websocket.addEventListener("close", () => {
websocket.addEventListener(WebsocketEvent.error, (_, event) => {
console.error("WebSocket error:", event);
terminal.options.disableStdin = true;
setConnectionStatus("disconnected");
});
websocket.addEventListener("message", (event) => {
websocket.addEventListener(WebsocketEvent.close, () => {
terminal.options.disableStdin = true;
setConnectionStatus("disconnected");
});
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");
});