@comprehend/telemetry-browser/src/WebSocketConnection.ts

Integration of comprehend.dev with OpenTelemetry in browser environments.

import { ObservationInputMessage } from "./wire-protocol";

const INGESTION_ENDPOINT = 'wss://ingestion.comprehend.dev';

export class WebSocketConnection {
    private ws?: WebSocket;
    private organization: string;
    private token: string;
    private debug?: (message: string) => void;
    private reconnectAttempts = 0;
    private maxReconnectAttempts = 5;
    private reconnectDelay = 1000;
    private messageQueue: ObservationInputMessage[] = [];
    private isConnected = false;

    constructor(organization: string, token: string, debug?: (message: string) => void) {
        this.organization = organization;
        this.token = token;
        this.debug = debug;
        this.connect();
    }

    private connect(): void {
        try {
            const url = `${INGESTION_ENDPOINT}/${this.organization}/observations`;
            this.ws = new WebSocket(url);

            this.ws.onopen = () => {
                this.debug?.('WebSocket connected');
                this.isConnected = true;
                this.reconnectAttempts = 0;
                
                // Send init message
                this.sendMessage({
                    event: "init",
                    protocolVersion: 1,
                    token: this.token
                });

                // Send queued messages
                while (this.messageQueue.length > 0) {
                    const message = this.messageQueue.shift();
                    if (message) {
                        this.sendMessage(message);
                    }
                }
            };

            this.ws.onmessage = (event) => {
                try {
                    const message = JSON.parse(event.data);
                    this.debug?.(`Received: ${JSON.stringify(message)}`);
                } catch (error) {
                    this.debug?.(`Failed to parse message: ${error}`);
                }
            };

            this.ws.onclose = (event) => {
                this.debug?.(`WebSocket closed: ${event.code} ${event.reason}`);
                this.isConnected = false;
                this.scheduleReconnect();
            };

            this.ws.onerror = (event) => {
                this.debug?.(`WebSocket error: ${event}`);
                this.isConnected = false;
            };

        } catch (error) {
            this.debug?.(`Failed to connect: ${error}`);
            this.scheduleReconnect();
        }
    }

    private scheduleReconnect(): void {
        if (this.reconnectAttempts < this.maxReconnectAttempts) {
            this.reconnectAttempts++;
            const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);
            this.debug?.(`Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})`);
            
            setTimeout(() => {
                this.connect();
            }, delay);
        } else {
            this.debug?.('Max reconnection attempts reached');
        }
    }

    sendMessage(message: ObservationInputMessage): void {
        if (!this.isConnected || !this.ws || this.ws.readyState !== WebSocket.OPEN) {
            this.debug?.('WebSocket not ready, queuing message');
            this.messageQueue.push(message);
            return;
        }

        try {
            const messageStr = JSON.stringify(message);
            this.ws.send(messageStr);
            this.debug?.(`Sent: ${messageStr}`);
        } catch (error) {
            this.debug?.(`Failed to send message: ${error}`);
            this.messageQueue.push(message);
        }
    }

    close(): void {
        if (this.ws) {
            this.ws.close();
        }
    }
}