Skip to content

Commit 92fd55f

Browse files
ThomasK33ThomasK33
authored
🤖 fix: auto-reconnect WebSocket on server restart (#1071)
## Problem When running `make dev-server` and the backend restarts via nodemon, the frontend incorrectly shows "Authentication Required" modal instead of automatically reconnecting. ## Root Cause - The `error` handler set `auth_required` even when no token existed - The `close` handler only handled auth codes (1008/4401), ignoring normal disconnection codes (1001, 1006) - No reconnection logic existed ## Solution Implement auto-reconnect with exponential backoff: - **New `reconnecting` status** - allows the app to stay visible while reconnecting - **Track connection history** - distinguish "first connect failed" vs "was connected, lost connection" - **Exponential backoff** - 100ms base delay, doubling up to 10s max, 10 attempts before giving up - **Smart error handling**: - Auth error codes (1008, 4401) → show auth modal - Normal disconnects after being connected → reconnect silently - First connection failure → show auth modal (server might require auth) ## Testing Added 5 unit tests covering: - Reconnects on close without showing auth_required when previously connected - Shows auth_required on close with auth error codes (4401, 1008) - Shows auth_required on first connection error without token - Reconnects on error when previously connected --- <details> <summary>📋 Implementation Plan</summary> # Plan: Fix WebSocket Reconnection on Server Restart ## Problem Statement When running `make dev-server` and the backend restarts via nodemon, the frontend incorrectly shows "Authentication Required" modal instead of automatically reconnecting. ## Root Cause Analysis In `src/browser/contexts/API.tsx`: 1. **`error` handler (lines 152-160)** - Sets `auth_required` even when no token exists: ```typescript } else { setState({ status: "auth_required" }); // Wrong: assumes auth needed } ``` 2. **`close` handler (lines 162-168)** - Only handles auth codes 1008/4401, ignores normal disconnection codes (1001, 1006) 3. **No reconnection logic** - Frontend never attempts to reconnect after disconnection ## Solution: Auto-reconnect with exponential backoff ### Changes to `src/browser/contexts/API.tsx` #### 1. Add new connection state for reconnecting ```typescript export type APIState = | { status: "connecting"; api: null; error: null } | { status: "connected"; api: APIClient; error: null } | { status: "reconnecting"; api: null; error: null; attempt: number } // NEW | { status: "auth_required"; api: null; error: string | null } | { status: "error"; api: null; error: string }; ``` #### 2. Track "has ever connected" state Add a ref to track whether we've successfully connected at least once: ```typescript const hasConnectedRef = useRef(false); ``` Set to `true` when connection succeeds; used to distinguish "first connect failed" vs "reconnect needed". #### 3. Add reconnection logic with exponential backoff ```typescript const reconnectAttemptRef = useRef(0); const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null); const MAX_RECONNECT_ATTEMPTS = 10; const BASE_DELAY_MS = 100; const MAX_DELAY_MS = 10000; const scheduleReconnect = useCallback(() => { const attempt = reconnectAttemptRef.current; if (attempt >= MAX_RECONNECT_ATTEMPTS) { setState({ status: "error", error: "Connection lost. Please refresh the page." }); return; } const delay = Math.min(BASE_DELAY_MS * Math.pow(2, attempt), MAX_DELAY_MS); reconnectAttemptRef.current = attempt + 1; setState({ status: "reconnecting", api: null, error: null, attempt: attempt + 1 }); reconnectTimeoutRef.current = setTimeout(() => { connect(authToken); }, delay); }, [authToken, connect]); ``` #### 4. Update error handler ```typescript ws.addEventListener("error", () => { cleanup(); if (hasConnectedRef.current) { // Was connected before - try to reconnect scheduleReconnect(); } else if (token) { // First connection failed with token - might be invalid clearStoredAuthToken(); setState({ status: "auth_required", error: "Connection failed - invalid token?" }); } else { // First connection failed without token - server might require auth setState({ status: "auth_required" }); } }); ``` #### 5. Update close handler ```typescript ws.addEventListener("close", (event) => { // Auth-specific close codes if (event.code === 1008 || event.code === 4401) { cleanup(); clearStoredAuthToken(); hasConnectedRef.current = false; // Reset - need fresh auth setState({ status: "auth_required", error: "Authentication required" }); return; } // Normal disconnection (server restart, network issue) if (hasConnectedRef.current) { cleanup(); scheduleReconnect(); } }); ``` #### 6. Reset reconnect state on successful connection In the `open` handler's success path: ```typescript .then(() => { hasConnectedRef.current = true; reconnectAttemptRef.current = 0; // Reset backoff // ... existing code }) ``` #### 7. Cleanup on unmount ```typescript useEffect(() => { connect(authToken); return () => { cleanupRef.current?.(); if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); } }; }, []); ``` ### Changes to `src/browser/App.tsx` Update the status check to handle `reconnecting` state - no modal, just let the app show while reconnecting: ```typescript // Show auth modal if authentication is required if (status === "auth_required") { return <AuthTokenModal isOpen={true} onSubmit={authenticate} error={error} />; } // Reconnecting is handled gracefully - app stays visible // Optional: could add a subtle toast/indicator later ``` ## Files to Modify 1. `src/browser/contexts/API.tsx` - Main changes (reconnection logic) 2. `src/browser/App.tsx` - Handle `reconnecting` state (minor, just update type handling) ## Testing ### Unit Test: `src/browser/contexts/API.test.tsx` Test the reconnection behavior without showing auth modal on disconnect: ```typescript import { act, cleanup, render, waitFor } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { GlobalWindow } from "happy-dom"; import { APIProvider, useAPI } from "./API"; // Mock WebSocket that we can control class MockWebSocket { static instances: MockWebSocket[] = []; url: string; readyState = 0; // CONNECTING onopen: (() => void) | null = null; onclose: ((event: { code: number }) => void) | null = null; onerror: (() => void) | null = null; eventListeners: Map<string, Function[]> = new Map(); constructor(url: string) { this.url = url; MockWebSocket.instances.push(this); } addEventListener(event: string, handler: Function) { const handlers = this.eventListeners.get(event) || []; handlers.push(handler); this.eventListeners.set(event, handlers); } close() { this.readyState = 3; // CLOSED } // Test helpers simulateOpen() { this.readyState = 1; // OPEN this.eventListeners.get("open")?.forEach((h) => h()); } simulateClose(code: number) { this.readyState = 3; this.eventListeners.get("close")?.forEach((h) => h({ code })); } simulateError() { this.eventListeners.get("error")?.forEach((h) => h()); } static reset() { MockWebSocket.instances = []; } static lastInstance() { return MockWebSocket.instances[MockWebSocket.instances.length - 1]; } } // Test component to observe API state function APIStateObserver({ onState }: { onState: (state: ReturnType<typeof useAPI>) => void }) { const apiState = useAPI(); onState(apiState); return null; } describe("API reconnection", () => { beforeEach(() => { const window = new GlobalWindow(); globalThis.window = window as unknown as Window & typeof globalThis; globalThis.document = window.document; globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket; MockWebSocket.reset(); }); afterEach(() => { cleanup(); MockWebSocket.reset(); }); test("reconnects on close without showing auth_required when previously connected", async () => { const states: string[] = []; render( <APIProvider> <APIStateObserver onState={(s) => states.push(s.status)} /> </APIProvider> ); // Initial connection const ws1 = MockWebSocket.lastInstance(); expect(ws1).toBeDefined(); // Simulate successful connection await act(async () => { ws1.simulateOpen(); // Mock the ping succeeding (need to mock orpc client) }); // Simulate server restart (close code 1006 = abnormal closure) await act(async () => { ws1.simulateClose(1006); }); // Should be "reconnecting", NOT "auth_required" await waitFor(() => { expect(states).toContain("reconnecting"); expect(states.filter((s) => s === "auth_required")).toHaveLength(0); }); // New WebSocket should be created for reconnect attempt expect(MockWebSocket.instances.length).toBeGreaterThan(1); }); test("shows auth_required on close with auth error codes (1008, 4401)", async () => { const states: string[] = []; render( <APIProvider> <APIStateObserver onState={(s) => states.push(s.status)} /> </APIProvider> ); const ws1 = MockWebSocket.lastInstance(); // Simulate successful connection then auth rejection await act(async () => { ws1.simulateOpen(); ws1.simulateClose(4401); // Auth required code }); await waitFor(() => { expect(states).toContain("auth_required"); }); }); test("shows error after max reconnect attempts exhausted", async () => { // Fast-forward timers for this test const states: string[] = []; render( <APIProvider> <APIStateObserver onState={(s) => states.push(s.status)} /> </APIProvider> ); // Simulate 10 failed reconnection attempts for (let i = 0; i < 11; i++) { const ws = MockWebSocket.lastInstance(); await act(async () => { ws.simulateError(); }); // Advance timers for backoff delay } await waitFor(() => { expect(states).toContain("error"); }); }); }); ``` ### Key Test Cases | Scenario | Expected Status | Auth Modal? | |----------|-----------------|-------------| | Server restart (close code 1006) after connected | `reconnecting` | No | | Server restart (close code 1001) after connected | `reconnecting` | No | | Auth rejection (close code 4401) | `auth_required` | Yes | | Auth rejection (close code 1008) | `auth_required` | Yes | | First connection fails, no token | `auth_required` | Yes | | First connection fails, with token | `auth_required` | Yes | | Max reconnect attempts (10) exhausted | `error` | No | ### Manual Testing 1. Run `make dev-server` 2. Open frontend in browser 3. Make a code change to trigger nodemon restart 4. Verify frontend reconnects automatically without showing auth modal 5. Test with auth token to ensure auth errors still show modal appropriately </details> --- _Generated with `mux`_ --------- Signed-off-by: Thomas Kosiewski <[email protected]>
1 parent 89e52e0 commit 92fd55f

File tree

2 files changed

+345
-13
lines changed

2 files changed

+345
-13
lines changed

‎src/browser/contexts/API.test.tsx‎

Lines changed: 267 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,267 @@
1+
import { act, cleanup, render, waitFor } from "@testing-library/react";
2+
import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
3+
import { GlobalWindow } from "happy-dom";
4+
5+
// Mock WebSocket that we can control
6+
class MockWebSocket {
7+
static instances: MockWebSocket[] = [];
8+
url: string;
9+
readyState = 0; // CONNECTING
10+
eventListeners = new Map<string, Array<(event?: unknown) => void>>();
11+
12+
constructor(url: string) {
13+
this.url = url;
14+
MockWebSocket.instances.push(this);
15+
}
16+
17+
addEventListener(event: string, handler: (event?: unknown) => void) {
18+
const handlers = this.eventListeners.get(event) ?? [];
19+
handlers.push(handler);
20+
this.eventListeners.set(event, handlers);
21+
}
22+
23+
close() {
24+
this.readyState = 3; // CLOSED
25+
}
26+
27+
// Test helpers
28+
simulateOpen() {
29+
this.readyState = 1; // OPEN
30+
this.eventListeners.get("open")?.forEach((h) => h());
31+
}
32+
33+
simulateClose(code: number) {
34+
this.readyState = 3;
35+
this.eventListeners.get("close")?.forEach((h) => h({ code }));
36+
}
37+
38+
simulateError() {
39+
this.eventListeners.get("error")?.forEach((h) => h());
40+
}
41+
42+
static reset() {
43+
MockWebSocket.instances = [];
44+
}
45+
46+
static lastInstance(): MockWebSocket | undefined {
47+
return MockWebSocket.instances[MockWebSocket.instances.length - 1];
48+
}
49+
}
50+
51+
// Mock orpc client
52+
void mock.module("@/common/orpc/client", () => ({
53+
createClient: () => ({
54+
general: {
55+
ping: () => Promise.resolve("pong"),
56+
},
57+
}),
58+
}));
59+
60+
void mock.module("@orpc/client/websocket", () => ({
61+
RPCLink: class {},
62+
}));
63+
64+
void mock.module("@orpc/client/message-port", () => ({
65+
RPCLink: class {},
66+
}));
67+
68+
void mock.module("@/browser/components/AuthTokenModal", () => ({
69+
getStoredAuthToken: () => null,
70+
// eslint-disable-next-line @typescript-eslint/no-empty-function
71+
clearStoredAuthToken: () => {},
72+
}));
73+
74+
// Import the real API module types (not the mocked version)
75+
import type { UseAPIResult as _UseAPIResult, APIProvider as APIProviderType } from "./API";
76+
77+
// IMPORTANT: Other test files mock @/browser/contexts/API with a fake APIProvider.
78+
// Module mocks leak between test files in bun (https://github.com/oven-sh/bun/issues/12823).
79+
// The query string creates a distinct module cache key, bypassing any mocked version.
80+
/* eslint-disable @typescript-eslint/no-require-imports, @typescript-eslint/no-unsafe-assignment */
81+
const RealAPIModule: {
82+
APIProvider: typeof APIProviderType;
83+
useAPI: () => _UseAPIResult;
84+
} = require("./API?real=1");
85+
/* eslint-enable @typescript-eslint/no-require-imports, @typescript-eslint/no-unsafe-assignment */
86+
const { APIProvider, useAPI } = RealAPIModule;
87+
type UseAPIResult = _UseAPIResult;
88+
89+
// Test component to observe API state
90+
function APIStateObserver(props: { onState: (state: UseAPIResult) => void }) {
91+
const apiState = useAPI();
92+
props.onState(apiState);
93+
return null;
94+
}
95+
96+
// Factory that creates MockWebSocket instances (injected via prop)
97+
const createMockWebSocket = (url: string) => new MockWebSocket(url) as unknown as WebSocket;
98+
99+
describe("API reconnection", () => {
100+
beforeEach(() => {
101+
// Minimal DOM setup required by @testing-library/react
102+
const happyWindow = new GlobalWindow();
103+
globalThis.window = happyWindow as unknown as Window & typeof globalThis;
104+
globalThis.document = happyWindow.document as unknown as Document;
105+
MockWebSocket.reset();
106+
});
107+
108+
afterEach(() => {
109+
cleanup();
110+
MockWebSocket.reset();
111+
globalThis.window = undefined as unknown as Window & typeof globalThis;
112+
globalThis.document = undefined as unknown as Document;
113+
});
114+
115+
test("reconnects on close without showing auth_required when previously connected", async () => {
116+
const states: string[] = [];
117+
118+
render(
119+
<APIProvider createWebSocket={createMockWebSocket}>
120+
<APIStateObserver onState={(s) => states.push(s.status)} />
121+
</APIProvider>
122+
);
123+
124+
const ws1 = MockWebSocket.lastInstance();
125+
expect(ws1).toBeDefined();
126+
127+
// Simulate successful connection (open + ping success)
128+
await act(async () => {
129+
ws1!.simulateOpen();
130+
// Wait for ping promise to resolve
131+
await new Promise((r) => setTimeout(r, 10));
132+
});
133+
134+
expect(states).toContain("connected");
135+
136+
// Simulate server restart (close code 1006 = abnormal closure)
137+
act(() => {
138+
ws1!.simulateClose(1006);
139+
});
140+
141+
// Should be "reconnecting", NOT "auth_required"
142+
await waitFor(() => {
143+
expect(states).toContain("reconnecting");
144+
});
145+
146+
expect(states.filter((s) => s === "auth_required")).toHaveLength(0);
147+
148+
// New WebSocket should be created for reconnect attempt (after delay)
149+
await waitFor(() => {
150+
expect(MockWebSocket.instances.length).toBeGreaterThan(1);
151+
});
152+
});
153+
154+
test("shows auth_required on close with auth error codes (4401)", async () => {
155+
const states: string[] = [];
156+
157+
render(
158+
<APIProvider createWebSocket={createMockWebSocket}>
159+
<APIStateObserver onState={(s) => states.push(s.status)} />
160+
</APIProvider>
161+
);
162+
163+
const ws1 = MockWebSocket.lastInstance();
164+
expect(ws1).toBeDefined();
165+
166+
await act(async () => {
167+
ws1!.simulateOpen();
168+
await new Promise((r) => setTimeout(r, 10));
169+
});
170+
171+
expect(states).toContain("connected");
172+
173+
act(() => {
174+
ws1!.simulateClose(4401);
175+
});
176+
177+
await waitFor(() => {
178+
expect(states).toContain("auth_required");
179+
});
180+
});
181+
182+
test("shows auth_required on close with auth error codes (1008)", async () => {
183+
const states: string[] = [];
184+
185+
render(
186+
<APIProvider createWebSocket={createMockWebSocket}>
187+
<APIStateObserver onState={(s) => states.push(s.status)} />
188+
</APIProvider>
189+
);
190+
191+
const ws1 = MockWebSocket.lastInstance();
192+
expect(ws1).toBeDefined();
193+
194+
await act(async () => {
195+
ws1!.simulateOpen();
196+
await new Promise((r) => setTimeout(r, 10));
197+
});
198+
199+
expect(states).toContain("connected");
200+
201+
act(() => {
202+
ws1!.simulateClose(1008);
203+
});
204+
205+
await waitFor(() => {
206+
expect(states).toContain("auth_required");
207+
});
208+
});
209+
210+
test("shows auth_required on first connection failure without token", async () => {
211+
const states: string[] = [];
212+
213+
render(
214+
<APIProvider createWebSocket={createMockWebSocket}>
215+
<APIStateObserver onState={(s) => states.push(s.status)} />
216+
</APIProvider>
217+
);
218+
219+
const ws1 = MockWebSocket.lastInstance();
220+
expect(ws1).toBeDefined();
221+
222+
// First connection fails - browser fires error then close
223+
act(() => {
224+
ws1!.simulateError();
225+
ws1!.simulateClose(1006);
226+
});
227+
228+
await waitFor(() => {
229+
expect(states).toContain("auth_required");
230+
});
231+
232+
expect(states.filter((s) => s === "reconnecting")).toHaveLength(0);
233+
});
234+
235+
test("reconnects on connection loss when previously connected", async () => {
236+
const states: string[] = [];
237+
238+
render(
239+
<APIProvider createWebSocket={createMockWebSocket}>
240+
<APIStateObserver onState={(s) => states.push(s.status)} />
241+
</APIProvider>
242+
);
243+
244+
const ws1 = MockWebSocket.lastInstance();
245+
expect(ws1).toBeDefined();
246+
247+
await act(async () => {
248+
ws1!.simulateOpen();
249+
await new Promise((r) => setTimeout(r, 10));
250+
});
251+
252+
expect(states).toContain("connected");
253+
254+
// Connection lost after being connected
255+
act(() => {
256+
ws1!.simulateError();
257+
ws1!.simulateClose(1006);
258+
});
259+
260+
await waitFor(() => {
261+
expect(states).toContain("reconnecting");
262+
});
263+
264+
const authRequiredAfterConnected = states.slice(states.indexOf("connected") + 1);
265+
expect(authRequiredAfterConnected.filter((s) => s === "auth_required")).toHaveLength(0);
266+
});
267+
});

0 commit comments

Comments
 (0)