dathagerty.com
An AI agent built to do Ralph loops - plan mode for planning and ralph mode for implementing.
1import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
2import type { WsEvent } from '../types';
3import { WsConnection } from './websocket';
4
5// Mock WebSocket
6class MockWebSocket {
7 readyState = 0; // CONNECTING
8 onopen: ((event: Event) => void) | null = null;
9 onclose: ((event: Event) => void) | null = null;
10 onerror: ((event: Event) => void) | null = null;
11 onmessage: ((event: MessageEvent) => void) | null = null;
12
13 constructor(private url: string) {
14 // Store for access in tests
15 MockWebSocket.lastInstance = this;
16 }
17
18 send(data: string) {
19 // No-op for testing
20 }
21
22 close() {
23 this.readyState = 3; // CLOSED
24 }
25
26 simulateOpen() {
27 this.readyState = 1; // OPEN
28 if (this.onopen) {
29 this.onopen(new Event('open'));
30 }
31 }
32
33 simulateClose() {
34 this.readyState = 3; // CLOSED
35 if (this.onclose) {
36 this.onclose(new Event('close'));
37 }
38 }
39
40 simulateMessage(data: string) {
41 if (this.onmessage) {
42 this.onmessage(new MessageEvent('message', { data }));
43 }
44 }
45
46 static lastInstance: MockWebSocket | undefined;
47}
48
49// Replace global WebSocket with mock
50const originalWebSocket = (globalThis as Record<string, unknown>).WebSocket;
51beforeEach(() => {
52 (globalThis as Record<string, unknown>).WebSocket = MockWebSocket;
53 vi.useFakeTimers();
54});
55
56afterEach(() => {
57 (globalThis as Record<string, unknown>).WebSocket = originalWebSocket;
58 vi.useRealTimers();
59 MockWebSocket.lastInstance = undefined;
60});
61
62describe('WsConnection', () => {
63 it('connects to provided URL', () => {
64 const url = 'ws://localhost:7400/ws';
65 const conn = new WsConnection(url);
66 conn.connect();
67
68 expect(MockWebSocket.lastInstance).toBeDefined();
69 });
70
71 it('auto-detects URL from window.location when not provided', () => {
72 // Create a connection without explicit URL
73 // Note: In the real implementation, it auto-detects from window.location
74 // For testing, we'll test with explicit URL
75 const url = 'ws://localhost:7400/ws';
76 const conn = new WsConnection(url);
77 expect(conn['wsUrl']).toBe(url);
78 });
79
80 it('dispatches events to registered listeners', () => {
81 const url = 'ws://localhost:7400/ws';
82 const conn = new WsConnection(url);
83 const listener = vi.fn();
84
85 conn.onEvent(listener);
86 conn.connect();
87
88 // Simulate connection opening
89 MockWebSocket.lastInstance?.simulateOpen();
90
91 // Simulate an agent_spawned event
92 const event: WsEvent = {
93 type: 'agent_spawned',
94 agent_id: 'agent-123',
95 profile: 'coder',
96 goal_id: 'goal-456',
97 };
98 MockWebSocket.lastInstance?.simulateMessage(JSON.stringify(event));
99
100 expect(listener).toHaveBeenCalledWith(event);
101 });
102
103 it('handles multiple listeners', () => {
104 const url = 'ws://localhost:7400/ws';
105 const conn = new WsConnection(url);
106 const listener1 = vi.fn();
107 const listener2 = vi.fn();
108
109 conn.onEvent(listener1);
110 conn.onEvent(listener2);
111 conn.connect();
112 MockWebSocket.lastInstance?.simulateOpen();
113
114 const event: WsEvent = {
115 type: 'agent_spawned',
116 agent_id: 'agent-123',
117 profile: 'coder',
118 goal_id: 'goal-456',
119 };
120 MockWebSocket.lastInstance?.simulateMessage(JSON.stringify(event));
121
122 expect(listener1).toHaveBeenCalledWith(event);
123 expect(listener2).toHaveBeenCalledWith(event);
124 });
125
126 it('removes listeners with offEvent', () => {
127 const url = 'ws://localhost:7400/ws';
128 const conn = new WsConnection(url);
129 const listener = vi.fn();
130
131 conn.onEvent(listener);
132 conn.offEvent(listener);
133 conn.connect();
134 MockWebSocket.lastInstance?.simulateOpen();
135
136 const event: WsEvent = {
137 type: 'agent_spawned',
138 agent_id: 'agent-123',
139 profile: 'coder',
140 goal_id: 'goal-456',
141 };
142 MockWebSocket.lastInstance?.simulateMessage(JSON.stringify(event));
143
144 expect(listener).not.toHaveBeenCalled();
145 });
146
147 it('sets connected to true when WebSocket opens', () => {
148 const url = 'ws://localhost:7400/ws';
149 const conn = new WsConnection(url);
150 expect(conn.connected).toBe(false);
151
152 conn.connect();
153 MockWebSocket.lastInstance?.simulateOpen();
154
155 expect(conn.connected).toBe(true);
156 });
157
158 it('sets connected to false when WebSocket closes', () => {
159 const url = 'ws://localhost:7400/ws';
160 const conn = new WsConnection(url);
161
162 conn.connect();
163 MockWebSocket.lastInstance?.simulateOpen();
164 expect(conn.connected).toBe(true);
165
166 MockWebSocket.lastInstance?.simulateClose();
167 expect(conn.connected).toBe(false);
168 });
169
170 it('handles malformed JSON gracefully', () => {
171 const url = 'ws://localhost:7400/ws';
172 const conn = new WsConnection(url);
173 const listener = vi.fn();
174 const consoleWarnSpy = vi.spyOn(console, 'warn');
175
176 conn.onEvent(listener);
177 conn.connect();
178 MockWebSocket.lastInstance?.simulateOpen();
179
180 MockWebSocket.lastInstance?.simulateMessage('invalid json');
181
182 expect(consoleWarnSpy).toHaveBeenCalled();
183 expect(listener).not.toHaveBeenCalled();
184
185 consoleWarnSpy.mockRestore();
186 });
187
188 it('handles unknown event types gracefully', () => {
189 const url = 'ws://localhost:7400/ws';
190 const conn = new WsConnection(url);
191 const listener = vi.fn();
192 const consoleWarnSpy = vi.spyOn(console, 'warn');
193
194 conn.onEvent(listener);
195 conn.connect();
196 MockWebSocket.lastInstance?.simulateOpen();
197
198 MockWebSocket.lastInstance?.simulateMessage(JSON.stringify({ type: 'unknown_type' }));
199
200 expect(consoleWarnSpy).toHaveBeenCalled();
201 expect(listener).not.toHaveBeenCalled();
202
203 consoleWarnSpy.mockRestore();
204 });
205
206 it('reconnects with exponential backoff on close', async () => {
207 const url = 'ws://localhost:7400/ws';
208 const conn = new WsConnection(url);
209
210 conn.connect();
211 MockWebSocket.lastInstance?.simulateOpen();
212 const firstInstance = MockWebSocket.lastInstance;
213
214 // Close the connection
215 MockWebSocket.lastInstance?.simulateClose();
216
217 // After 1 second, it should reconnect
218 vi.advanceTimersByTime(1000);
219 expect(MockWebSocket.lastInstance).not.toBe(firstInstance);
220
221 // Next reconnect should be after 2 seconds
222 MockWebSocket.lastInstance?.simulateClose();
223 vi.advanceTimersByTime(2000);
224 const thirdInstance = MockWebSocket.lastInstance;
225
226 // Next should be 4 seconds
227 MockWebSocket.lastInstance?.simulateClose();
228 vi.advanceTimersByTime(4000);
229 const fourthInstance = MockWebSocket.lastInstance;
230
231 expect(fourthInstance).not.toBe(thirdInstance);
232 });
233
234 it('resets backoff on successful connection', () => {
235 const url = 'ws://localhost:7400/ws';
236 const conn = new WsConnection(url);
237
238 // First connection
239 conn.connect();
240 MockWebSocket.lastInstance?.simulateOpen();
241 const firstInstance = MockWebSocket.lastInstance;
242
243 // Close and reconnect quickly
244 MockWebSocket.lastInstance?.simulateClose();
245 vi.advanceTimersByTime(1000); // First backoff
246 MockWebSocket.lastInstance?.simulateOpen();
247
248 // Now close again - should use 1s backoff again, not 2s
249 MockWebSocket.lastInstance?.simulateClose();
250 const beforeReconnect = MockWebSocket.lastInstance;
251 vi.advanceTimersByTime(1000);
252 const afterReconnect = MockWebSocket.lastInstance;
253
254 // Should have reconnected after 1s (reset backoff)
255 expect(afterReconnect).not.toBe(beforeReconnect);
256 });
257
258 it('stops reconnection attempts when disconnect() is called', () => {
259 const url = 'ws://localhost:7400/ws';
260 const conn = new WsConnection(url);
261
262 conn.connect();
263 MockWebSocket.lastInstance?.simulateOpen();
264 const firstInstance = MockWebSocket.lastInstance;
265
266 MockWebSocket.lastInstance?.simulateClose();
267 conn.disconnect();
268
269 vi.advanceTimersByTime(2000);
270 // Should still be the closed instance, no reconnect
271 expect(MockWebSocket.lastInstance).toBe(firstInstance);
272 });
273
274 it('caps backoff at 30 seconds', () => {
275 const url = 'ws://localhost:7400/ws';
276 const conn = new WsConnection(url);
277
278 conn.connect();
279 MockWebSocket.lastInstance?.simulateOpen();
280
281 // Trigger multiple reconnections: 1s, 2s, 4s, 8s, 16s, 32s → capped at 30s
282 for (let i = 0; i < 6; i++) {
283 MockWebSocket.lastInstance?.simulateClose();
284 vi.advanceTimersByTime(32000); // Advance beyond any possible backoff
285 }
286
287 // The last reconnect should use 30s max
288 MockWebSocket.lastInstance?.simulateClose();
289 const beforeAdvance = MockWebSocket.lastInstance;
290 vi.advanceTimersByTime(29000); // Less than 30s
291 expect(MockWebSocket.lastInstance).toBe(beforeAdvance);
292
293 vi.advanceTimersByTime(1000); // Now we're past 30s
294 expect(MockWebSocket.lastInstance).not.toBe(beforeAdvance);
295 });
296
297 it('dispatches all event types correctly', () => {
298 const url = 'ws://localhost:7400/ws';
299 const conn = new WsConnection(url);
300 const listener = vi.fn();
301
302 conn.onEvent(listener);
303 conn.connect();
304 MockWebSocket.lastInstance?.simulateOpen();
305
306 const events: Array<WsEvent> = [
307 {
308 type: 'agent_spawned',
309 agent_id: 'agent-1',
310 profile: 'coder',
311 goal_id: 'goal-1',
312 },
313 {
314 type: 'agent_progress',
315 agent_id: 'agent-1',
316 turn: 1,
317 summary: 'working',
318 },
319 {
320 type: 'agent_completed',
321 agent_id: 'agent-1',
322 outcome_type: 'success',
323 summary: 'completed',
324 tokens_used: 1000,
325 },
326 {
327 type: 'node_created',
328 parent_id: null,
329 id: 'node-1',
330 project_id: 'proj-1',
331 node_type: 'task',
332 title: 'New Task',
333 description: 'Task description',
334 status: 'pending',
335 priority: 'high',
336 assigned_to: null,
337 created_by: null,
338 labels: [],
339 created_at: '2026-02-11T00:00:00Z',
340 started_at: null,
341 completed_at: null,
342 blocked_reason: null,
343 metadata: {},
344 },
345 {
346 type: 'node_status_changed',
347 node_id: 'node-1',
348 node_type: 'task',
349 old_status: 'pending',
350 new_status: 'completed',
351 },
352 {
353 type: 'edge_created',
354 id: 'edge-1',
355 edge_type: 'contains',
356 from_node: 'goal-1',
357 to_node: 'node-1',
358 label: null,
359 created_at: '2026-02-11T00:00:00Z',
360 },
361 {
362 type: 'session_ended',
363 session_id: 'sess-1',
364 handoff_notes: 'notes',
365 },
366 {
367 type: 'tool_execution',
368 agent_id: 'agent-1',
369 tool: 'read_file',
370 args: { path: '/test' },
371 result: 'content',
372 },
373 {
374 type: 'orchestrator_state_changed',
375 goal_id: 'goal-1',
376 state: 'completed',
377 },
378 ];
379
380 for (const event of events) {
381 MockWebSocket.lastInstance?.simulateMessage(JSON.stringify(event));
382 }
383
384 expect(listener).toHaveBeenCalledTimes(events.length);
385 for (let i = 0; i < events.length; i++) {
386 expect(listener).toHaveBeenNthCalledWith(i + 1, expect.objectContaining({ type: events[i].type }));
387 }
388 });
389});