js-sdk/src/_shared/websocket.ts

import WebSocket from "ws";

import { TimeoutError } from "../exceptions.js";

/** Options used to establish a Wrenn WebSocket connection. */
export interface WsConnectionOpts {
	/** HTTP(S) API origin. Converted to WS(S) for the socket URL. */
	baseUrl: string;
	/** WebSocket path relative to the base URL. */
	path: string;
	/** API key sent as `X-API-Key`. */
	apiKey?: string;
	/** Host token sent as `X-Host-Token`. */
	hostToken?: string;
	/** Callback invoked for each JSON message or raw text payload. */
	onMessage: (data: unknown) => void;
	/** Callback invoked for socket errors after connection establishment. */
	onError?: (error: Error) => void;
	/** Callback invoked when the socket closes after connection establishment. */
	onClose?: (code: number, reason: string) => void;
	/** Connection timeout in milliseconds. Defaults to 30 seconds. */
	timeoutMs?: number;
}

/** Minimal WebSocket wrapper for JSON-oriented Wrenn streaming endpoints. */
export class WsConnection {
	private ws: WebSocket;
	private closed = false;

	private constructor(ws: WebSocket) {
		this.ws = ws;
	}

	/** Sends a JSON-encoded message over the open WebSocket. */
	send(data: unknown): void {
		if (this.closed || this.ws.readyState !== WebSocket.OPEN) {
			throw new Error("WebSocket is not open");
		}
		this.ws.send(JSON.stringify(data));
	}

	/** Closes the WebSocket connection if it is still open. */
	close(): void {
		if (this.closed) return;
		this.closed = true;
		this.ws.close();
	}

	/** Indicates whether the connection has closed or failed. */
	get isClosed(): boolean {
		return this.closed;
	}

	/**
	 * Opens a WebSocket connection and resolves once the socket is ready.
	 *
	 * @param opts - Connection URL, authentication, callbacks, and timeout.
	 * @returns An established WebSocket connection wrapper.
	 * @throws TimeoutError When the connection is not established before timeout.
	 */
	static connect(opts: WsConnectionOpts): Promise<WsConnection> {
		return new Promise((resolve, reject) => {
			const url = new URL(`${opts.baseUrl}${opts.path}`);
			const protocol = url.protocol === "https:" ? "wss:" : "ws:";
			url.protocol = protocol;

			const headers: Record<string, string> = {};
			if (opts.apiKey) {
				headers["X-API-Key"] = opts.apiKey;
			}
			if (opts.hostToken) {
				headers["X-Host-Token"] = opts.hostToken;
			}

			const ws = new WebSocket(url.toString(), {
				headers,
			});

			const timeout = opts.timeoutMs ?? 30_000;
			let settled = false;
			let timeoutHandle: ReturnType<typeof setTimeout> | undefined;

			const cleanup = () => {
				if (timeoutHandle) clearTimeout(timeoutHandle);
				settled = true;
			};

			timeoutHandle = setTimeout(() => {
				if (!settled) {
					cleanup();
					ws.terminate();
					reject(
						new TimeoutError(
							`WebSocket connection timed out after ${timeout}ms`,
						),
					);
				}
			}, timeout);

			ws.on("open", () => {
				if (settled) return;
				cleanup();
				const conn = new WsConnection(ws);
				ws.on("message", (raw) => {
					try {
						const data = JSON.parse(raw.toString());
						opts.onMessage(data);
					} catch {
						opts.onMessage(raw.toString());
					}
				});
				ws.on("error", (err) => {
					conn.closed = true;
					opts.onError?.(err);
				});
				ws.on("close", (code, reason) => {
					conn.closed = true;
					opts.onClose?.(code, reason.toString());
				});
				resolve(conn);
			});

			ws.on("error", (err) => {
				if (settled) return;
				cleanup();
				reject(err);
			});

			ws.on("close", (code, reason) => {
				if (settled) return;
				cleanup();
				reject(
					new Error(
						`WebSocket closed before opening (${code}): ${reason.toString()}`,
					),
				);
			});
		});
	}
}