1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

/**
 * WebSocket client for the ori backend: parses server messages into store
 * events and reconnects automatically with a capped backoff. The server
 * replays the session history on every connection, and the reducer resets on
 * "hello", so a reconnection rebuilds the exact same thread.
 */

import type { ClientMessage, ServerMessage } from "./protocol";
import type { ChatEvent } from "./reducer";

/** What the UI needs from the connection. */
export interface OriSocket {
	send(message: ClientMessage): void;
	close(): void;
}

interface ConnectOptions {
	url: string;
	dispatch: (event: ChatEvent) => void;
	/** Injectable WebSocket constructor, for tests. Defaults to the browser's. */
	webSocketImpl?: typeof WebSocket;
	/** First reconnection delay in ms; doubles up to 16x. Defaults to 500. */
	initialRetryMs?: number;
}

/**
 * connectSocket opens the connection and keeps it alive until close().
 *
 * @example
 * const socket = connectSocket({ url: wsUrl(), dispatch });
 * socket.send({ type: "prompt", text: "hello" });
 */
export function connectSocket(options: ConnectOptions): OriSocket {
	const Impl = options.webSocketImpl ?? WebSocket;
	const initialRetryMs = options.initialRetryMs ?? 500;

	let socket: WebSocket | null = null;
	let retryMs = initialRetryMs;
	let closedByUser = false;

	const open = () => {
		socket = new Impl(options.url);
		options.dispatch({ type: "connection", status: "connecting" });

		socket.onopen = () => {
			retryMs = initialRetryMs;
		};
		socket.onmessage = (event: MessageEvent) => {
			let message: ServerMessage;
			try {
				message = JSON.parse(String(event.data)) as ServerMessage;
			} catch {
				return; // a malformed frame must not kill the connection handling
			}
			options.dispatch({ type: "server", message });
		};
		socket.onclose = () => {
			if (closedByUser) {
				return;
			}
			options.dispatch({ type: "connection", status: "offline" });
			setTimeout(open, retryMs);
			retryMs = Math.min(retryMs * 2, initialRetryMs * 16);
		};
	};
	open();

	return {
		send(message: ClientMessage) {
			if (socket && socket.readyState === Impl.OPEN) {
				socket.send(JSON.stringify(message));
			}
		},
		close() {
			closedByUser = true;
			socket?.close();
		},
	};
}

/** wsUrl derives the backend WebSocket URL from the page's own origin. */
export function wsUrl(): string {
	const scheme = window.location.protocol === "https:" ? "wss" : "ws";
	return `${scheme}://${window.location.host}/ws`;
}